A Framework of SO(3)-equivariant Non-linear Representation Learning and its Application to Electronic-Structure Hamiltonian Prediction

Authors: Shi Yin, Xinyang Pan, Fengyan Wang, Lixin He
Primary area: applications to physical sciences (physics, chemistry, biology, etc.)
Keywords: SO(3)-equivariant representation learning; Non-linear expressiveness; Electronic-structure Hamiltonian prediction

Abstract:
We propose both a theoretical and a methodological framework to address a critical challenge in applying deep learning to physical systems: the reconciliation of non-linear expressiveness with SO(3)-equivariance in predictions of SO(3)-equivariant quantities, such as the electronic-structure Hamiltonians. Inspired by covariant theory in physics, we present a solution by  exploring the mathematical relationships between SO(3)-invariant and SO(3)-equivariant quantities and their representations. We first construct theoretical SO(3)-invariant quantities derived from the SO(3)-equivariant regression targets, and use these invariant quantities as supervisory labels to guide the learning of high-quality SO(3)-invariant features. Given that SO(3)-invariance is preserved under non-linear operations, the encoding process for invariant features can extensively utilize non-linear mappings, thereby fully capturing the non-linear patterns inherent in physical systems. Building on this, we propose a gradient-based mechanism to induce SO(3)-equivariant encodings of various degrees from the learned SO(3)-invariant features. This mechanism can incorporate non-linear expressive capabilities into SO(3)-equivariant representations, while theoretically preserving their equivariant properties as we prove, establishing a strong foundation for regressing complex SO(3)-equivariant targets. We apply our theory and method to the electronic-structure Hamiltonian prediction tasks, experimental results on eight benchmark databases covering multiple types of systems and challenging scenarios show substantial improvements on the state-of-the-art prediction accuracy of deep learning paradigm. Our method boosts Hamiltonian prediction accuracy by up to 40\% and enhances downstream physical quantities, such as occupied orbital energy, by a maximum of 76\%. Our method also significantly promotes the acceleration performance for the convergence of traditional Density Functional Theory methods.

Decision: reject
Reviewer scores: 6, 6, 6, 6

Reviews and discussion:
<thread>
creation date:
1729744227048

thread_type:
official_review

<thread_object>
creator: reviewer_7q7d

<content>
summary:
The paper introduces a method combining SO(3)-equivariance with non-linear expressiveness, and applied it for Hamiltonian prediction. This method has a theoretical foundation, and experiments demonstrate the effectiveness of the approach.

strengths:
- Theoretical analysis is provided to ensure equivariance is maintained alongside non-linear expressiveness.
- Experimental results indicate significant improvement in Hamiltonian prediction accuracy, outperforming QHNet and DeepH-E3.
- Detailed explanations are provided for the experimental setup, showcasing a comprehensive evaluation.
- The proposed method is simple and can be easily combined with existing approaches.

weaknesses:
- To show whether the acceleration in DFT outweighs the additional computation, it would be helpful if the paper provided a comparison of the wall time cost for accelerating traditional DFT algorithms. 
- Since SO(3)-equivariance is important in many contexts, could the authors provide an example of how this model might be applied to another prediction task?

questions:
see above
</content>
</thread_object>

</thread>
<thread>
creation date:
1729944664712

thread_type:
official_review

<thread_object>
creator: reviewer_g9ts

<content>
summary:
The authors present a novel architecture and loss function designed for learning targets in SO(3) symmetry. Their focus is on the challenge of predicting the Hamiltonian within Kohn-Sham Density Functional Theory (DFT). Their contribution is twofold: First, they suggest fitting the Forbenius norm of the target alongside the direct fitting of the equivariant targets. Second, they propose a method for differentiating the predicted Forbenius norm of the Hamiltonian, utilizing invariant representations to produce equivariant embeddings that capture non-linear dependencies within SO(3)-equivariant embeddings.

strengths:
* The proposed TraceGrad method shows significant improvements over previous approaches.
* The use of differentiation to achieve higher-order SO(3)-equivariant embeddings is a novel concept, as it has not been extensively explored before, apart from its application to atomic forces.

weaknesses:
* The paper suffers from unclear writing, particularly in the abstract and introduction, which fail to clearly convey the contributions and include overly verbose sections (e.g., lines 66-72).
* Although the paper claims to present a general framework for SO(3) equivariance, it is heavily focused on Hamiltonian prediction, which limits its broader applicability.
* The method is inherently designed to learn conservative vector fields on the embeddings, a point that is not addressed but could significantly influence its empirical performance.
* l.57-61, SO(3)-equivariant models include non-linearities through gating or activating the norm of a vector.
* l.129,130, "Nevertheless, multiplying SO(3)-equivariant features with linear coefficients may not fundamentally improve their non-linear expressiveness." is a claim without evidence.
* l.43, In DFT, the Hamiltonian is not predicted but computed or solved for.
* l.46, Hamiltonian prediction cannot work in O(N) given that the matrix itself is O(N²).

questions:
* Are the models in the experiments adjusted for size and compute time? The extra computations through TraceGrad could equally be spent on larger networks, how does this compare?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730549668947

thread_type:
official_review

<thread_object>
creator: reviewer_7k3e

<content>
summary:
This paper addresses the challenge of incorporating non-linear expressiveness into SO(3)-equivariant deep learning frameworks for physical systems. The authors propose a theoretical framework that constructs SO(3)-invariant quantities tr(Q·Q†) from SO(3)-equivariant targets and utilizes gradient operations to induce SO(3)-equivariant features while maintaining non-linear expressiveness. The method is applied to electronic-structure Hamiltonian prediction tasks and evaluated on eight benchmark databases, showing improvements in prediction accuracy and downstream physical quantities. The method demonstrates generalization ability across various scenarios including thermal motions and bilayer twists, and aids in accelerating traditional DFT calculations through better initialization. The main contribution is providing a framework that attempts to combine strict SO(3)-equivariance with non-linear expressiveness in deep learning models for physical systems.

strengths:
Originality:
- Proposes a novel perspective on bridging SO(3)-equivariance and non-linear expressiveness through invariant quantities and gradient operations
- Creates a theoretical connection between invariant and equivariant representations for l>1.
- Introduces a gradient-based mechanism for inducing equivariant features while preserving non-linear information

Quality:
- Provides rigorous mathematical proofs for the theoretical framework
- Conducts comprehensive experiments across several different benchmark databases

Clarity:
- Well-structured presentation of the theoretical framework
- Clear mathematical derivations with detailed proofs in appendices
- Thorough experimental analysis with multiple evaluation metrics

weaknesses:
While proposing a new method, it fails to provide detailed comparisons of parameter counts, memory consumption, training time, and convergence behavior between the baseline and proposed models. The claimed superiority of gradient-based transfer over traditional gating activation mechanisms is not convincingly demonstrated through empirical evidence. Moreover, the practical significance of the accuracy improvements is not well justified - the paper does not discuss what constitutes a meaningful accuracy threshold for Hamiltonian prediction tasks, nor does it adequately analyze the trade-offs between these moderate accuracy improvements and their associated computational costs. These omissions make it difficult to fully evaluate the practical value and efficiency of the proposed method relative to existing approaches. A comprehensive analysis of these aspects would significantly strengthen the paper's contributions and help readers better understand the method's practical advantages and limitations.

questions:
1. The paper proposes using gradient operation (∂z/∂f) to transfer nonlinearity from l=0 features to equivariant features, where z = S_nonlin(u) and u = CGDecomp(f⊗f, 0). However, similar to the gated activation mechanisms, this method still essentially applies nonlinearity only to l=0 components. The fundamental advantage of gradient-based transfer over the gated activation mechanism remains unclear - could the authors provide theoretical insights on why gradient operations might be more effective than gated activation for transferring nonlinear information from l=0 to higher-degree features?

2. The paper lacks crucial implementation details comparing the baseline and proposed method, such as parameter counts, training epochs, convergence curves, and computational overhead. More importantly, the reported improvements are moderate (e.g., MAE reduction from 0.274 to 0.226 on BB dataset) - such accuracy gains could potentially be achieved through increasing training epoch, learning rate adjustment, or other hyperparameter tuning of the baseline models. To demonstrate the superiority of the proposed method, one needs to compare the upper limits of model expressiveness after controlling for various training factors (number of epochs, hyperparameters, optimization strategies, random seeds, etc.). The current empirical evidence does not adequately address this comparison.

3. For electronic-structure Hamiltonian prediction tasks, what is the practically meaningful accuracy threshold? If, for instance, 1 meV accuracy is sufficient for most practical applications, how do we justify the value of further accuracy improvements (e.g., from 0.274 to 0.226 meV on BB dataset)? Could the authors discuss the practical significance of such fine-grained improvements in the context of real-world applications, such as materials design and DFT calculations?

4. In line 1116, The paper compares the O(N) complexity of Hamiltonian prediction with the O(N³) complexity of traditional DFT methods. However, this comparison seems unfair: the O(N³) complexity in traditional DFT includes solving the Kohn-Sham equations, while the O(N) complexity here only covers Hamiltonian prediction. To obtain physical properties from the predicted Hamiltonian, one still needs O(N³) operations for diagonalization. Could the authors clarify this comparison and provide a more comprehensive complexity analysis that includes all steps necessary for obtaining the same physical properties in both approaches?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730592354364

thread_type:
official_review

<thread_object>
creator: reviewer_uxwk

<content>
summary:
This paper proposes a new framework that strictly satisfies the SO(3) equivalence and has the potential to enhance expressive power with nonlinearity. The effectiveness of the proposed ansatz is demonstrated with experiments on Hamiltonian prediction tasks and DFT.

strengths:
-	Balancing between physics constraints (i.e. SO3-equivalence in this task) and the expressive power of the network is highly non-trivial. Therefore, an ansatz that both strictly satisfies the equivariance constraint and has the potential to leverage nonlinearity for expressiveness improvement is a noticeable contribution. 
-	The authors test their new framework on various tasks

weaknesses:
The main weakness of this paper is lack of readability. Many important background and detailed introduction to the problem, the task, and technical details are missing, making it hard to follow.

-	L156-160, it is very abrupt when concepts like Q and direct-product state are mentioned without any context. The *problem formulation* section should provide details (with clearly explained equations) on how the related concepts are defined, along with their physical meaning related to the task, which is also not clearly delivered.

-	In appendix A, it is hard to understand why the authors only include well-known concepts like definitions of group and equivariance but omit those more problem-specific and highly relevant ones, e.g. irreducible representation, properties of Wigner-D-matrix, Clebsch-Gordan decomposition, etc. These are crucial for readers to understand their methods. Hiding all these related contexts is not appreciated.

-	The learning task considered in the experiment section is not explicitly described. If the authors find themselves running out of pages, a detailed description on electronic-structure Hamiltonian prediction task should then be added in the appendix, including the physics background, the mapping to be learned, and methods adopted as well as model architecture. Same comments on the DFT experiment.

Also, I didn’t find a comparison on computation cost, e.g. GPU hours / FLOPS for the Hamiltonian prediction task.

questions:
See Weakness.
</content>
</thread_object>

</thread>
<thread>
creation date:
1733125698308

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Letter to Area Chairs

comment:
Dear Area Chairs:

We have completed the rebuttal process and sincerely thank the reviewers for their time and valuable comments, which have significantly helped us improve the quality of our paper.

In the revised paper, we have addressed all of the reviewers' comments and questions and have improved our paper in four aspects.

1.	We have improved the representation of the paper, clarifying the main learning tasks and some key concepts, thereby enhancing its readability.

2.	We have provided a clearer theoretical explanation of the effectiveness of the proposed method, included a more detailed comparison with related studies, and incorporated robust experimental evidence.

3.	We have supplemented more experiments to demonstrate the efficiency and effectiveness of our method.

4.	To demonstrate the generality of our method, we have supplemented our experiments to demonstrate its effectiveness on the additional task, namely energy and force field prediction.

Overall, we have thoroughly addressed the reviewers' concerns and refined our work according to their suggestions. The reviewers have responded positively, and a consensus has been reached regarding the evaluation scores.

**We would like to emphasize that our method demonstrates the significant potential of deep learning in addressing core tasks in quantum physics, representing a cross-disciplinary effort bridging AI and physics research. We believe this work holds substantial value for both the AI and quantum physics communities.**

We sincerely request your consideration of our work and would like to express our deep appreciation for your time and thoughtful oversight throughout the review process. Your support and guidance are invaluable to us.

Best regards,

The authors
</content>
</thread_object>

</thread>


Metareview:
metareview:
The paper proposes a neural network architecture that preserves rotational equivariance for high-order tensors. The novelty is the introduced equivariant component constructed as the gradient of a nonlinear mapping on a scalar feature of the tensor. The architecture is applied to the DFT Hamiltonian prediction task which involves tensors in various orders, and preferred results are demonstrated.

Reviewers raised some concerns and insufficiencies, and after the rebuttal, they appreciated the noticeable efforts in providing more details in the background, more results on hardware and timing, accuracy-efficiency trade-off evaluation, more empirical comparisons with other architectures, demonstration on the effect to accelerate DFT in wall time, and extended results for predicting energy and force. Three out of four of them have raised their total scores as a result, and all scores are leaning towards positive.

Nevertheless, I still feel certain remaining concerns:
* Expressiveness comparison with other methods. If expanding the expression, the equivariant update in each component would be a nonlinear transformation on the scalar $\\mathrm{CGDecomp}(\\mathbf{f} \\otimes \\mathbf{f}, 0)$, times the gradient of $\\mathrm{CGDecomp}(\\mathbf{f} \\otimes \\mathbf{f}, 0)$, which is in turn a specific scalar times the tensor f. In comparison, many other equivariant architectures also follow this form of a nonlinear transformation on the scalar $\\mathrm{CGDecomp}(\\mathbf{f} \\otimes \\mathbf{f}, 0)$ times the tensor feature $\\mathbf{f}$. The authors are expected to provide more discussion/rationale on why the proposed formulation could be more expressive.

* Conservativeness. The gradient of a scalar function makes a conservative vector/tensor field (function), but not all vector/tensor fields are conservative. This is not a problem (even a desired thing) in machine learning force field since force fields are conservative. But the authors are expected to explain why it is still the case for predicting the Hamiltonian. Although the authors did not bing it up in the paper, the DFT Hamiltonian is the gradient of electronic energy w.r.t the ground-state density matrix (in the DFT sense). But the gradient taken in the model is not w.r.t a tensor that holds the semantics of a density matrix, and even if there is, the ground-state density matrix still needs to be inferred from the input conformation. I noticed the authors made a discussion on this point in response to Reviewer G9tS, but it does not seem to reflect my question.

  Moreover, in the domain of machine learning force field, it is an empirical observation that the gradient formulation does not seem to better fit the force than directly predicting the vector, possibly due to that neural networks are not typically designed to produce expressive gradient functions, and may have numerical difficulties for the gradient. The authors may also need to discuss this point when motivating the gradient-based approach.

* On learning the SO(3)-invariant quantity $\\mathbf{T} = \\mathrm{tr}(\\mathbf{Q} \\mathbf{Q}^\\dagger)$ simultaneously. I still have the concern why this could be consistent with letting the gradient learning an arbitrary equivariant target. The argument using $\\mathbf{Q} = \\frac{\\partial \\mathbf{T}}{\\partial \\mathrm{Conj}(\\mathbf{Q})}$ seems vague: the gradient in the model is not taken w.r.t $\\mathrm{Conj}(\\mathbf{Q})$ but w.r.t the $\\mathbf{f}$ equivariant feature, and if the authors were suggesting $\\mathbf{f}$ could learn $\\mathrm{Conj}(\\mathbf{Q})$, then what the gradient produces is just $\\mathbf{Q}$, which may restrict the capability in fitting the equivariant Hamitonian target.

I communicated these points with reviewers, and Reviewer G9tS expressed the same remaining concerns regarding the conservativeness, and the potential conflict between fitting $\\mathbf{T}$ and $\\mathbf{Q}$ simultaneously, as the explanation and derivation presented in the paper seem insufficient, despite the empirical results. (He/She also resolved my question on the linear scaling.)

As the reviewers all rated only a borderline accept, and there still remains concerns, I tend to recommend a reject for the current time, and encourage the authors to present a deeper understanding of their proposed architecture in future versions.