Tensor Decomposition Networks for Accelerating Machine Learning Force Field Computations

We appreciate that Reviewer S6y1 finds our method significant.

The low-rank structure of the tensor M deserves further discussion. While the paper relates the rank R to the maximum angular degree L, a theoretical justification or empirical analysis of the inherent low-rank nature of M is lacking.

Our work mainly approximates the Clebsch Gordan (CG) tensor product via CP decomposition. The CG coefficient tensor is intrinsically low-rank because of the sparsity that most index triples $(\ell_1, \ell_2, \ell_3)$ in the CG tensor violate the selection rules $|\ell_1 - \ell_2| \le \ell_3 \le \ell_1 + \ell_2$ and $m_3 = m_1 + m_2$; their entries are therefore identically zero.

Empirical evidence supports this view. By the upper bound of the tensor rank [1], the maximum rank of a CG tensor is $(L+1)^4 \propto O(L^4)$ while we successfully employ only an $O(L^2)$ rank schedule to approximate the CG tensor. Below we present the ablation study of ranks over n-body system dataset [2], where the TDN is trained with a maximum angular degree $L=2$, 4 interaction blocks, and a hidden dimension of 64 in 5000 epochs. In the table, our adopted schedule uses much smaller ranks yet matches the accuracy obtained with the largest ranks, corresponding to the results of rank schedule selection in Section 4.1 and showing the low rank nature of the CG tensor.

The method for selecting the optimal rank R is not clearly explained. Is R tuned via cross-validation? Clarifying whether this is a learned or predefined hyperparameter is needed.

The method for selecting the optimal rank R is not clearly explained. Is R tuned via cross-validation?

Rank $R$ is a predefined hyperparameter, not a learned quantity. Because computing the exact CP rank of the CG tensor is NP-hard, we instead explore several simple growth rules before network training (Section 4.1) to guide the selection of rank $R$. Our adopted schedule gives a relatively small SO(3)-equivariance errors with low computational cost as demonstrated in Figure 3, and rank-sweep ablations of $L=2$ in the above table confirm that this schedule suffices for downstream tasks.

Although this heuristic meets the needs of our current force‑field tasks, a data‑driven approach, such as learning a rank predictor or applying Bayesian optimization over ranks, could generalize the method and further reduce redundancy; we identify this as an important direction for future work.

The theoretical analysis focuses on the forward computation of the tensor product. However, during training, the gradients with respect to the CP decomposition parameters must also be computed. Therefore, an analysis of the computational complexity of gradient backpropagation should be included.

We present both the theoretical and empirical time complexity analysis of the CP decomposition for the tensor product. Let $A\in \mathbb{R}^{d_3\times R}, B\in \mathbb{R}^{d_1\times R}, C\in \mathbb{R}^{d_2\times R}$ and input $x\in \mathbb{R}^{d_1}, y\in \mathbb{R}^{d_2}$. By Equation (3), the forward map is

$$z = A(B^\top x\odot C^\top y ).$$

Given loss gradient $g_z = \frac{\partial \mathcal L}{\partial z}\in \mathbb{R}^{d_3}$, the backward gradient to $x$ is

$$g_x = g_z\frac{\partial{A(B^\top x\odot C^\top y )}}{\partial x} = g_z\frac{\partial A \text{diag}(C^\top y) B^\top x}{\partial x} = g_z A \text{diag}(C^\top y) B^\top = (g_z A \odot C^\top y) B^\top$$

and similarly the backward gradient to $y$ is

$$g_y = (g_z A \odot B^\top x) C^\top.$$

Each expression requires a time complexity of $O(R(d_1+d_2+d_3))$. Therefore, the total time complexity to compute the backward gradient is $O(R(d_1+d_2+d_3))$, identical to the forward pass.

Furthermore, the computation time of both forward and backward is presented in the below table. Each pass is measured with a batch size of 1000 and a multiplicity of 128. The backward step is consistently $\approx 2-3×$ the forward cost, independent of L, confirming that the same time complexity of forward and backward pass.

Figure 3: It is unclear which dataset this figure refers to. This should be explicitly stated in the caption or text.

Which dataset Figure 3 refers to?

Figure 3 is not tied to any dataset. The curves in Figure 3 report the approximation error and runtime of the CP decomposition for the CG tensor product itself, which only depends on the chosen rank, multiplicities, and maximum angular degree. The low-rank reconstruction can be learned and evaluated without any molecular data. To learn the CP factor matrices in CP decomposition, we minimize the Frobenius norm between the exact and approximate CG tensor product on batches of standard normal inputs; this guards against data bias and is done once before any network training. We will amend the caption to state this explicitly in Figure 3 in the camera‑ready version.

The evaluation of computation time is described in terms of “wall-clock runtime” but it is ambiguous whether this refers to training time, inference time, or both. A more precise description of the measurement setup is needed.

How do you define the “wall-clock runtime” that you evaluated in the experiments? Is it related to computation time during training, testing or both?

The "wall-clock runtime" refers to the inference time over the PubchemQC-Traj-Small dataset. We will revise line 260 to state this explicitly in the camera‑ready version.

In Line 285, a reference to a table is made, but the table number is missing.

The table reference should be "Table 3". We will revise line 285 and update this in our camera-ready version.

[1] Kolda, T. and Bader, B. Tensor decompositions and applications. SIAM review, 51(3):455–500, 2009.

[2] Kipf, T., Fetaya, E., Wang, K.-C., Welling, M., and Zemel, R. Neural relational inference for interacting systems. In International conference on machine learning, pp. 26882697. PMLR, 2018.

Dear Reviewer S6y1,

Thank you again for the positive evaluation and constructive feedback. As we approach the end of the reviewer-author discussion period, we would like to kindly check if there are any new questions or concerns you would like us to address. If our rebuttal has fully addressed your concerns, we would be grateful if you might consider increasing your score to reflect a solid accept.

Best regards,

Authors

We appreciate that Reviewer oXft finds our method significant.

While the approximation method is justified theoretically, the required CP rank R is treated as a hyperparameter. Estimating or learning optimal ranks adaptively could improve generality.

CP Rank Selection is Manual: The required CP rank R is treated as a fixed hyperparameter. Since determining the optimal CP rank is NP-hard, this introduces potential inefficiency. A more adaptive or data-driven approach to rank selection would improve scalability and generality.

Selecting the CP rank $R$ is itself an NP‑hard problem as mentioned in the limitation section, so in practice the rank schedules are empirical heuristics. The adopted one gives a relatively small SO(3)-equivariance error with low computational cost across every schedule in Figure 3 and Section 4.1, and ablation studies in the below table of NequIP confirm that this schedule suffices for downstream tasks while cutting time cost substantially. Note that the NequIP is trained with a maximum angular degree $L=2$, 4 interaction blocks, and a multiplicity of 64 for 100 epochs with our proposed schedule.

Although this heuristic meets the needs of our current force‑field tasks, we agree that a data‑driven approach, such as learning a rank predictor or applying Bayesian optimization over ranks, could generalize the method and further reduce redundancy; we identify this as an important direction for future work.

The approximation is "approximately equivariant". The trade-off between error and efficiency for approximate equivariance has been well-explored in the literature. I think this work could benefit from more discussion in real-world generalization contexts.

We agree that the value of approximate equivariance ultimately rests on how well the model generalizes. We share the concern and therefore measured our CP-decomposition-based approximation on two large-scale, experiment-derived benchmarks: Open Catalyst Project and PubChemQC-Traj. Both datasets contain millions of chemically diverse geometries and are labelled with high-fidelity DFT energies and forces. We believe the performance on these two datasets is a strong indicator of generalization ability.

To push the generalization study further, we plan to train on OMol [1], a newly released molecular dataset that spans even broader compositional space. These additional experiments may help clarify how the approximation behaves on systems that differ more sharply from the training distribution.

Missing some important citation: Brandon Anderson, Truong-Son Hy and Risi Kondor, Cormorant: Covariant molecular neural networks, NeurIPS 2019.

Thank you for raising this issue. We will update the citation at line 25 and line 72 in the camera-ready version.

Approximate Equivariance: The method only maintains approximate SO(3)-equivariance. While the authors provide theoretical bounds on the equivariance error, the practical implications of this approximation, especially in sensitive downstream tasks, are not deeply analyzed.

No Ablation on Universality vs. Accuracy Trade-off: Although the paper proves universality theoretically, it lacks an ablation showing how the approximation degrades or improves as R increases. Practical guidance for choosing R would be valuable.

Thank you for raising these points. To further demonstrate the practical implication of this approximation and strengthen the contribution, we conduct the ablation study of ranks over n-body system dataset [2]. In the below rank‑sweep table for the n-body system experiment, the TDN is trained with a maximum angular degree $L=2$, 4 interaction blocks, and a hidden dimension of 64 in 5000 epochs. In that table, our adopted schedule uses much smaller ranks yet matches the accuracy obtained with the largest ranks, corresponding to the results of rank schedule selection in Section 4.1. Note that the largest rank is deduced from the upper bound in line 164; see [3].

Limited Scope of Decomposition: The proposed CP decomposition is demonstrated only for the Clebsch–Gordan tensor product. It's unclear how well the approach generalizes to other types of equivariant operations (for example, permutation-equivariance).

Tensor product is widely used in SO(3)-equivariant neural networks. Particularly, the Clebsch-Gordan tensor product is the canonical formulation of an SO(3)-equivariant neural network. For permutation equivariance, to our knowledge, it is unknown that tensor product will play a role in it. The permutation-equivariant application is usually built by deep sets or graph neural networks, where bilinear operations are not necessarily used. For other types of equivariant operations, as long as they can be written in bilinear forms, our CP decomposition can be applied to approximate the operations.

The evaluation is mostly limited to molecular force field tasks (energy and force prediction). It would strengthen the paper to show broader utility in other domains where SO(3)-equivariance is important (e.g., protein modeling, 3D object recognition).

We appreciate this suggestion. Our study deliberately targets machine-learning interatomic potentials since accelerating the Clebsch Gordan map yields a great practical benefit there, and the benchmarks (Open Catalyst Project, PubChemQC-Traj) are well established. Extending the approach to other domains, e.g., protein folding and 3-D object recognition, is indeed an exciting next step, and we leave it as future explorations.

The path-weight sharing mechanism drastically reduces parameter count by using a single weight matrix across all paths. While this helps with efficiency, it may limit the expressiveness or fine-grained control across angular degrees, especially at high L.

We agree that path-weight sharing will limit the expressiveness of the tensor product. However, coupling with CP decomposition, the result can match the full Clebsch Gordan (CG) map. To demonstrate this, we further conduct the ablations of CP decomposition and path-weight sharing for both the n-body experiment dataset and a subset of PubchemQC-Traj-Small dataset in the below table. For the n-body experiment, the testing results of TDN are presented in the table, where the TDN is trained with a maximum angular degree $L=2$, 4 interaction blocks, a hidden dimension of 72, and a weight decay of $5\times 10^{-6}$ in 2000 epochs. And for the subset of PubchemQC-Traj-Small dataset, we randomly select one conformer for each trajectory from PubchemQC-Traj-Small to build the training, validation, and testing set. The testing results of TDN over this subset are presented in the table, where the TDN is trained with a maximum angular degree $L=2$, 4 interaction blocks, and a hidden dimension of 64 in 60 epochs. Note that in the table, the n-Body experiment appears under the label "n-Body," while the subset of the PubChemQC-Traj-Small dataset is labeled simply as "Subset."

As shown in the above table, introducing path-weight sharing alone will indeed incur a small accuracy drop. However, introducing both path-weight sharing and CP decomposition produces comparable results to the full CG tensor product while significantly boosting the efficiency. We believe this is because the Clebsch Gordan map is replaced by a low-rank surrogate with rank-1 factors, which act as an implicit regularizer as they filter out high-variance directions that weight sharing alone cannot suppress. We will include this table in the appendix of the camera‑ready version.

[1] Levine, D. S., et al. The open molecules 2025 (omol25) dataset, evaluations, and models.

[2] Kipf, T., et al. Neural relational inference for interacting systems. In International conference on machine learning, pp. 26882697. PMLR, 2018.

[3] Kolda, T. and Bader, B. Tensor decompositions and applications. SIAM review, 51(3):455–500, 2009.

Dear Reviewer oXft,

Thank you again for the positive evaluation and constructive feedback. As we approach the end of the reviewer-author discussion period, we would like to kindly check if there are any new questions or concerns you would like us to address. If our rebuttal has fully addressed your concerns, we would be grateful if you might consider increasing your score to reflect a solid accept.

Best regards,

Authors

We appreciate that Reviewer cmnz finds our method significant.

Why is EQUIFORMER selected as the main method for acceleration? It would be nice to see if the CPD applicability and performance are not only for EQUIFORMER but also for other tensor-product-based methods.

We focused on Equiformer [1] because Equiformer uses the fully-connected Clebsch Gordan (CG) tensor product as the basic equivariant operation. This is the canonical formulation of an SO(3)‑equivariant network and is significant for the machine learning force field prediction task [1-3]. Also, the fully-connected CG tensor product is a cost‑dominating kernel so the acceleration is easy to measure. Furthermore, Equiformer outperforms many other baselines with fully-connected CG tensor products in several molecular tasks [1]. Considering all the above, Equiformer provides a standard and widely recognized baseline.

The CP‑decomposition method itself is not tied to Equiformer. Any architecture whose core step is a bilinear operation, e.g., Tensor Field Networks (TFN) [2], NequIP [3], or other e3nn “TensorProduct” blocks, can be approximated in the same way. We further provide the ablation study over the NequIP architecture [3] with and without our techniques in the below table. The employed dataset is the subset of PubchemQC-Traj-Small dataset, where we randomly select one conformer for each trajectory from PubchemQC-Traj-Small to build the training, validation, and testing set. And the NequIP is trained with a maximum angular degree $L=2$, 4 interaction blocks, and a multiplicity of 64 for 100 epochs. As shown in the below table, the energy and force validation results are similar, but the training and inference time are accelerated by 2 times with our techniques.

Is the GPU runtime reported in Table 2? Can the authors show the CPU runtime as well?

The time reported in Table 2 is the inference time over a single A100-80GB GPU, and the runtime includes both GPU and CPU overhead. The pure CPU implementation runtime over AMD EPYC 9684X 96-Core Processor is detailed in the below table.

We will include this CPU table in the appendix of the camera‑ready version.

Would the proposed method benefit from other tensor network (decomposition) structures? What is the author's opinion?

The current study uses the CP decomposition because it is the simplest: it requires only rank‑1 components, is easy to fit once, and introduces no new hyperparameters beyond the rank schedule. Additionally, the error bound and the uniqueness of CP decomposition are well studied in [4]. Other tensor decomposition methods are promising as well, such as Tucker decomposition [4] and tensor‑train decomposition [5]. These alternatives offer a different trade‑off between parameter count, arithmetic complexity, and the cache locality. We view a systematic comparison of these alternatives, and their effect on training stability and accuracy, as valuable future work.

It would be nice to add relative numbers on parameter counts and processing speed directly to Table 2.

We present the updated Table 2 below, where we include the multiplicative ratio with respect to the Equiformer baseline with the same $L$. We will include this table in the camera‑ready version.

[1] Liao, Y. and Smidt, T. Equiformer: Equivariant graph attention transformer for 3d atomistic graphs. In The Eleventh International Conference on Learning Representations, 2022.

[2] Thomas, N., Smidt, T., Kearnes, S., Yang, L., Li, L., Kohlhoff, K., and Riley, P. Tensor field networks: Rotation-and translation-equivariant neural networks for 3d point clouds. arXiv preprint arXiv:1802.08219, 2018.

[3] Batzner, S., Musaelian, A., Sun, L., Geiger, M., Mailoa, J. P., Kornbluth, M., Molinari, N., Smidt, T. E., and Kozinsky, B. "E(3)-equivariant graph neural networks for data-efficient and accurate interatomic potentials." Nature communications 13, no. 1 (2022): 2453

[4] Kolda, T. and Bader, B. Tensor decompositions and applications. SIAM review, 51(3):455–500, 2009.

[5] Oseledets, I. V. Tensor-train decomposition. SIAM Journal on Scientific Computing, 33(5), 2295-2317.

Dear Reviewer cmnz,

Thank you again for the positive evaluation and constructive feedback. As we approach the end of the reviewer-author discussion period, we would like to kindly check if there are any new questions or concerns you would like us to address. If our rebuttal has fully addressed your concerns, we would be grateful if you might consider increasing your score to reflect a solid accept.

Best regards,

Authors

We appreciate that Reviewer 7Mvd finds our method significant. To further strengthen our contribution, we add n-body experiments [1,2] as an additional benchmark.

In line 233, what is the motivation of the proposed rank schedules?

The rank schedules are empirical heuristics and are corresponded to different complexity growth of the CP decomposition. The adopted one gives a relatively small SO(3)-equivariance error with low computational cost across every schedule in Figure 3 and Section 4.1. Further ablations clarify that this schedule gives a good trade-off between efficiency and approximation error:

• Ablations of ranks at $L=2$ show that using this schedule leaves downstream metrics with negligible impact. In the below rank‑sweep table of the n-body system experiment, the TDN is trained with a maximum angular degree $L=2$, 4 interaction blocks, and a hidden dimension of 64 in 5000 epochs. In that table, our schedule uses much smaller ranks yet matches the accuracy obtained with the largest rank to approximate the Clebsch-Gordan tensor product. Note that the largest rank is deduced from the upper bound in line 164; see [3].

• Likewise, in the ablations of CP-decomposition and path‑weight‑sharing (in the next reply), turning CP decomposition on and off under our schedule has negligible effect on final performance but markedly improves efficiency.

Lack of ablation study on CP decomposition and path-weight sharing.

Below we present the ablation study of CP decomposition and path-weight sharing for both the n-body experiment dataset and the subset of PubchemQC-Traj-Small dataset. For the n-body experiment, the testing results of TDN are presented in the table, where the TDN is trained with a maximum angular degree $L=2$, 4 interaction blocks, a hidden dimension of 72, and a weight decay of $5\times 10^{-6}$ in 2000 epochs. And for the subset of PubchemQC-Traj-Small dataset, we randomly select one conformer for each trajectory from PubchemQC-Traj-Small to build the training, validation, and testing set. The testing results of TDN over this subset are presented in the table, where the TDN is trained with a maximum angular degree $L=2$, 4 interaction blocks, and a hidden dimension of 64 in 60 epochs. Note that in the table, the n-Body experiment appears under the label "n-Body," while the subset of the PubChemQC-Traj-Small dataset is labeled simply as "Subset."

As shown in the above table, introducing both path-weight sharing and CP decomposition does not affect the performance a lot while significantly boosting the efficiency, although we notice that introducing path-weight sharing alone will hurt the performance a bit. We believe this is because the Clebsch Gordan map is replaced by the low-rank surrogate with rank-1 factors, which act as an implicit regularizer as they filter out high-variance directions that weight sharing alone cannot suppress. We will include this table in the appendix of the camera‑ready version.

While CP decomposition improves computational efficiency, it also introduces abstraction: the learned rank-1 components may lack clear physical or geometric interpretations.

Thank you for raising this point. Although the rank‑1 components from the CP decomposition are indeed harder to interpret physically, they do not strip away information contained in the Clebsch Gordan tensor if the rank is sufficiently large, as shown in the empirical ablation study above and Theorem 3.2.

Conceptually, Clebsch Gordan coefficients are coordinates of the SO(3) tensor product in the standard coupled-angular-momentum basis. Switching to a CP basis simply means re-expressing the same sparse linear map in a different compact rank‑1 factor basis; the transformation is analogous to choosing a new set of axes in linear algebra, and information is re‑packaged.

[1] Kipf, T., Fetaya, E., Wang, K.-C., Welling, M., and Zemel, R. Neural relational inference for interacting systems. In International conference on machine learning, pp. 26882697. PMLR, 2018.

[2] Satorras, V. G., Hoogeboom, E., and Welling, M. E(n) equivariant graph neural networks. In International conference on machine learning, pp. 9323–9332. PMLR, 2021.

[3] Kolda, T. and Bader, B. Tensor decompositions and applications. SIAM review, 51(3):455–500, 2009.

Dear Reviewer 7Mvd,

Thank you again for the positive evaluation and constructive feedback. As we approach the end of the reviewer-author discussion period, we would like to kindly check if there are any new questions or concerns you would like us to address. If our rebuttal has fully addressed your concerns, we would be grateful if you might consider increasing your score to reflect a solid accept.

Best regards,

Authors