TraceGrad: a Framework Learning Expressive SO(3)-equivariant Non-linear Representations for Electronic-Structure Hamiltonian Prediction
Authors: Shi Yin
Xinyang Pan
Fengyan Wang
Lixin He
Venue: ICML 2025 poster (2025)
Keywords: Atomic System Modeling; Electronic-structure Hamiltonian prediction; SO(3)-equivariant representation learning; Non-linear expressiveness
Primary area: applications->chemistry_physics_and_earth_sciences

Abstract: We propose a  framework to combine strong non-linear expressiveness with strict SO(3)-equivariance in prediction of the electronic-structure Hamiltonian, by exploring the mathematical relationships between SO(3)-invariant and SO(3)-equivariant quantities and their representations. The proposed framework, called **TraceGrad**,  first constructs theoretical SO(3)-invariant **trace** quantities derived from the Hamiltonian 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 learning of 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 **grad**ient-based mechanism to induce SO(3)-equivariant encodings of various degrees from the learned SO(3)-invariant features. This mechanism can incorporate powerful non-linear expressive capabilities into SO(3)-equivariant features with correspondence of physical dimensions to the regression targets, while theoretically preserving equivariant properties, establishing a strong foundation for predicting electronic-structure Hamiltonian. Experimental results on eight challenging benchmark databases demonstrate that our method achieves state-of-the-art performance in Hamiltonian prediction.

TraceGrad: a Framework Learning Expressive S0(3)-equivarian Non-linear Representations for Electronic-Strcture Hamiltonian Prediction
Shi Yin +!
Xinyang Pan
Fengyan Wang'
f Lixin He
Abstract
1. Introduction
" "Symmetry dictates interaction"
C. N. Yang
We propose a framework to combine strong non-linear expressiveness with strict SO(3)- equivariance prediction of the electronic- structure Hamiltonian, by exploring the math- ematical relationships between 50(3)-invariam) and SO(()-equivariant quantities and their rep- resentations. The proposed framework, called TraceGrad, first constructs theoretical SO(3)- invariant trace quantities derived from the Hamiltonian targets, and use these invariant quantities as supervisory labels guide the learning of high-quality SO(3)-invariant features. Given that non-linear operations, the learning of invariant features can extensively utilize non-linear map- pings, thereby fully capturing the non-linear pat- terns inherent in physical systems. Building on this, we propose gradient- baseed mech- anism to induce -O(3)--quivariant encodings of various degrees from the learned SO(3)- invariant features. This mechanism can incorpo- rate powerful non-linear expressive capabilities into SQuii-equivariant features with correspon- dence of physical dimensions to the regression targets, while theoretically preserving equivari- ant properties, establishing strong foundation for predicting electronic-struce Hamiltonian. Experimental results eight challenging bench- mark databases demonstrate that our method achieves state-of-the-art performance in Hamil- tonian prediction.
Electronic structure calculations provide crucial insights into the behavior of electrons in condensed matter, form- ing the foundation for predicting material properties such as conductivity, magnetism, optical response, and chem- ical reactivity, playing a key role applications ranging from semiconductors to renewable energy and catalysis, ul- timately helping to reshape our lives. The core of clec- tronic structure calculations lies in solving the Hamilto- nian matrices, from which critical physical quantities such as orbital energy, band structure and electronic wavefunc- tion are induced. However, traditional Density Functional Theory (DFT) methods (Hohenberg & Kohn 1964, Kohn & Sham 1965) suffer from the extremely time-consuming self-consistent field (SCF) algorithms used obtain the Hamiltonians. These algorithms involve the exhaustive it- erative diagonalization of large matrices, with each diag- onalization scaling with a high computational complexity of ((N*), where N is the number of atoms in the system. With advantages in computational complexity and gener- alization capabilities, the deep learning paradigm has sig- nificantly advanced physics research (Zhang et al. 2023). In the task of predicting the Hamiltonians, deep learning approaches (Unke al.) 2021 Yu et al.) 2023b Gong
is preserved under
et al., 2023 Zhang et al., 2024 Wang et al., 2024c) by-
pass the SCF process, dramatically reducing the computa- tional complexity of solving Hamiltonians, opening up new possibilities for analyzing extremely large atomic systems, enabling efficient materials simulation and design that was previously unimaginable. See Appendix for more details.
To align with fundamental physical laws, deep learning methods for Hamiltonian prediction must strictly adhere to cquivariance under 3D rotational transformations, i.c., el- ements from the SO(3) group. However, it remains chal- lenging for deep learning methods to simultaneously en- sure strict SO(3)-equivariiace and high numerical accuracy predicting Hamiltonians. The root cause of this prob- lem lies in the conflict between SO(3)-equivariaace and non-linear expressiveness: specifically, directly applying non-linear activation functions on SQuii-cquiaarian fea- tures (with degree 2 1) may lead to the loss of equiv-
'Institute of Artificial Intelligence, Hefei Comprehensive Na- tional Science Center *CAS Key Laboratory of Quantum Infor- mation, University of Science and Technology of China. Cor- respondence to: Shi Yin <shiyin @iai ust.de.te.d...m.. Lixin He Seedx@uutc.cdu.c...
Proceedings of the 1nnd International Conference Machine Learning, Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

TraceGrad
ariance, while bypassing non-linear mappings severely re- stricts the network's expressive capabilities and thereby lowering down the achievable accuracy. This issue is com- monly encountered in physics-oriented oriented machine learning tasks that require both strict equivariance and fine-grained generalization performance, as discussed by Zitnick et al. (2022). Several recent methods, such as the gated mech- anism (Weiler et al.] 2018). spherical channel-based ap- proaches (Zitnick et al. [2022] Liao et al.) 2024 Wang
- representations for predicting the regres- sion target H in the inference phase from high-quality non- linear S0(3) sion of tr(H- H in the training phase. Taking SO(3)- invariant features as a bridge, this mechanism can incorpo- rate powerful non-linear expressive capabilities into SO(3)- equivariant representations while preserving their equivari- ant properties, as we prove. Compared the gated mech- anism (Weiler et al.] 2018), the proposed mechanism nat- urally follows the correspondence of physical dimensions between the 50(3)-equivarian) Hamiltonian and the SO(3)- invariant trace quantity, laying a solid foundation for re- gressing Hamiltonian with the help of supervision signals from the trace quantity.
features learned under the supervi-
et al. 2024c), local coordinate strategies (Li et al. 2022)
and hybrid regression methods (Yin et al. 2024), have
been proposed to address this challenge. However, com- bining strong non-linear expressiveness with strict SO(3)-
equivariance remains challenging for the Hamiltonian pre- diction task, as analyzed in detail in Section 2
Experimental results demonstrate that, our TraceGrad framework significantly improves the state-of-the- per- formance in Hamiltonian prediction on eight databases from the well-known DeepH and QH9 benchmark series (Li et al. 2022 Gong et al. 2023; Yu et al. 2023a). It demonstrates excellent generalization performance to both crystalline and molecular systems, covering challenging scenarios such as thermal motions, bilayer twists, scale variations, and new trajectories. Furthermore, as observed from the experiments on the QH9 benchmark series, our approach also substantially enhances the prediction accu- racy of downstream physical quantities of Hamiltonian in- cluding occupied orbital energy and electronic wavefunc- tion. Moreover, our method also demonstrates superior performance in accelerating the convergence of classical DFT by providing predicted Hamiltonians as initializa- tion matrices. Our leading performance comprehensively demonstrates that our method, while satisfying SO(3)- equivariance, possesses excellent expressive power and generalization performance, providing an effective deep learning tool for efficient and accurate electronic-structure calculations of atomic systems.
To address the equivatianceeexpresiiveeesss dilemma, we make theoretical and methodological explorations on uni- fying strict (0(3)--quivariane) with strong non-linear ex- pressiveness within the realm of deep representation learn- ing for dlectronic-structur Hamiltonian prediction. We are inspired by the insight that invariant quantities in trans- formation (Resnick) 199T) often reflect the mathematical nature of physical laws and can induce other quantities with equivariant properties, and aim to extend the rela- tionship between invariance and equivariance from spe- cific physical quantities to hidden representations of neu- ral networks. From the perspective of deep represen- tation learning, the attribute that invariance is preserved under non-linear operations is a significant advantage, given its compatibility with non-linear expressive capabil- ities. Built upon these insights, we propose solution to the aquivariance-expeesiieneee dilemma by intensively exploring and making use of the intrinsic relationships between SO(3)-invariant and SO(3)-equivariant quantities and representations: we first dedicate efforts to supervis- ing high-quality SO(33-innariint features with ample non- linear expressiveness, and subsequently, we derive SO(3)- equivariant non-linear representations and the target quan- tities from these SO(3)-invariant ones. Specifically:
2. Related Work
The -quiivvanva representation learning paradigm (Thomas et al. 2018 Geiger & Smidi (022) typically developed group theory-based tensor operators, such as linear scaling. element-wise sum, direct products, di- rect sums, Clebsch-Gordan decomposition, and equivari- ant normalization, to encode equivariant features. These operators have been used to construct graph neural net- work architectures for tasks in 3D point cloud analysis (Fuchs et al. 2020), energy and force prediction (Liao & Smidt 2023 as well as Hamiltonian prediction (Schütt
First, we propose a theoretical construct of 50(3)-invariam quantities, namely tr(H. Hi), where tr(-) signifies the trace operation, t denotes the conjugate transpose opera- tion, and H denotes the SQu3)-equivariant regression tar- gets, i.e., the block of Hamiltonian matrix. A significant advantage of these SO(3))invarian) quantities lies in the fact that they are directly derived from the SO(3)-equivarian target labels and can serve as unique supervision labels for the effective learning of informative -nvaaainn fea- tures that capture the intrinsic symmetry properties of the mathematical structure of H without requiring additional labeling resources.
ct al. 2019 Unke et al.] [022]] Gong et al. 2023 Zhong
et al. 2023). However, as non-linear activation func-
tions may result in the loss of equivariance, they are re- stricted when applied to S0(3)-cquivariant features with
Second, we propose gradient- mechanism to induce
2

---

TraceGrad
degree 1 2 1. This restriction severely limits the net- work's capability to model complex non-linear mappings. To alleviate this issue, Hamiltonian prediction methods like DeepH-E3 (Gong et al., (023) and QHNet (Yu et al.) 2023b) introduced a gated activation mechanism (Weiler eet al. 2018) that fed SO(3)-invariant features (1 = 0) into non-linear activation functions and used these features as linear gating coefficients for -quiivequivariann features ( 2 1), aiming to enhance their expressive power while maintaining strict equivariance. However, these methods only indirectly supervised the learning of SO(3)-invariant features through equivaaqii Hamiltonians taking -oui)-equivariant representations as a bridge, and the lack of direct S0(3)-invarian supervision signals may result in insufficient learning of SO(3)-invariant features, limiting their expressiveness assist SO(3)-qqivvarian features. Furthermore, the gated mechanism, which directly multi- plies -quuivaaaa features with SO(3)-invariant fea- tures to yield new S0(3)-equivariant features, lacks con- sideration of the physical dimensional relationship between the actual -quui-aqqi and Snvaaiin physical quantities, making it difficult to achieve perfect coupling of their representations. See Remark 4.3 for a more com- prehensive analysis.
through data-driven methods does not have strict theoreti- cal guarantee, even with rotational data augmentation. In many applications of electronic structure calculations, the demands for symmetry are extremely high. Even minor de- viations from perfect SQui-eqqivaariane can result in in- correct physical results. In this paper, we seek for com- bining strict SO(()--quivariance with the powerful non- linear neural expressiveness to resolve the equivariance- expressiveness dilemma for electronic-structure Hamilto- nian prediction.
3. Preliminary
In order to better understand our theory and methods, we recommend that the readers first refer to the prelimi- nary in Appendices: Appendix A provides the definitions of foundational concepts relevant to the research topic; Appendix [B] describes the application tasks, specifically the clectronic-structur Hamiltonian prediction task; Ap- pendix C formalizes Hamiltonian prediction as an SO(3)- equivariant non-linear representation learning problem.
4. Theory
Theorem 4.1. The quantity T = tr(HH HH) is SO(3)- invariant, where H is the simplified representation (with- out superscripts) of the basic block H' of Hamiltonian matrix defined in Appendix B and i denotes the conjugate transpose operation, (r(-) is the trace operation.
In order to improve non-linear expressiveness, Sitressiiveess,ziiikkk Zitnick et al. 2022) proposed spherical channel methods, which decom- posed -quiiiaqua features into SO(3)-invariant co- efficients of spherical harmonic basis functions. These SO(3)-invariant coefficients were processed by non-linear neural networks enhance expressiveness, with equivari- ance regained by recombining the updated coefficients with the basis functions. In subsequent developments (Passaro & Zitmick 2023) et al., 2024 Wang et al. 2024b
Theorem 4.2. The non-linear neural mapping Inontin (-) defined the following is SO(3) -qqivvaaaa::
Inontin (f)
let al. 2025), this approach has demonstrated a remarkable capacity to fit complex functions. However, as pointed out by existing literature (Zhang et al., 2023). this approach degenerates from continuous to discrete rotational equiv- ariance, losing strict equivariance to continuous rotational transformations due to the decomposition based on inner- product operations with discrete basis functions; Li et al. 2022) proposed a local coordinate strategy, projecting ro- tating global coordinates onto SO(3)-invariam local ones for the non-linear neural network to encode. However, this strategy effective only for global rigid rotations and fails to maintain symmetry under non-rigid perturbations like thermal fluctuations or bilayer twists, as it lacks neu- ral mechanism to enforce equivariance; Yin et al. (2024) proposed a hybrid regression framework, where the non- linear mechanisms showed remarkable capability at learn- ing -quivirquvv from the data with the help of the theoretically SO(3) equivvraaa mechanisms, and released powerful non-linear expressive capabilities to achieve more numerical accuracy. However, the equivariance achieved
Of
subject to z Snotlin (u), u CGDecomp(f & f, 0)
(1)
where is an input S0(3) feature with degree in the direct-sum state, X denotes the direct-product op- eration of tensors, CGDecomp(-- 0) refers to performing a Clebsch-Gordan decomposition of the tensor and return- ing the SO(3)-invariant component corresponding to de- gree O, Snontin (.) represents arbitrary differentiable non- linear neural modules, u and are SO(3)-invarian features, and v is the outputted feature encoded by Inontin (-).
The proofs of the two theorems above are presented in Ap- pendix [
Remark 4.3. In physics, dimensions describe the funda- mental properties of physical quantities. The seven base dimensions are mass [M], time (T). length [4., elec- tric current [4], temperature [0), amount of substance [N], and luminous intensity [J]. Other physical quanti- ties can be derived from these base dimensions. For in- stance, the dimension of energy E is given by dim E
3

---

TraceGrad
The quantities H and T presented in Theo- rem 4.1 correspond the physical dimensions of dim E and dim E²² respectively, where their units can be taken as VeV and (meV)². These two physical dimensions are re- lated through: dim H dim E dim E²
M
posed of recently developed equivariant operators (Thomas) et al. 2018 Geiger & Smidt 2022) like linear scaling. direct products, direct sums, Clebsch-Gordan decomposi- tion, gated activation, equivariant normalization, and etc., to encode the physical as spherical harmonies (Schrodinger 1926). or represen- tations passed from the previous neural layers, as equiv- ariant feature (()) in the current hidden layer. Next, we construct the feature v(8) Imonlin (f**), to achieve suffi- cient non-linear expressiveness while maintaining SO(3)- equivariance, where Inonlia! (.), Snonlin (-) are the non-linear functions defined in Eq. The function Snontin (-) can be implemented as any differentiable non-linear neural net- work module, such as feed-forward layers with non-linear activation functions like SiLU and normalization opera- tions like Layernorm.
dim T As we
's initial representations, such
dim E dim H developed in Section 5 in the feature space, the SO(3)-
invariant feature 2 is used to regress T, and the SO(3)- cquivariant features f and v are used to regress H. Our based mechanism v = af
0= in Theorem 4.2 nat-
urally preserves the dimensional correspondence between H and T. On the other hand, the existing gated activa- tion function (Weiler et al. 2018), which can be expressed as f. does not enforce such a dimensional cor- respondence and could not offer a strong guarantee of a balanced coupling between the SO(3)-invariant and SO(3)- equivariant representations. In contrast, our approach, with adherence to the dimensional correspondence between H and T, ensures that the learned features are more physi- cally grounded and expressive, effectively capturing the un- derlying relationships between sov3)ainvariant and SO(3)- equivariant quantities. See Appendix H for an empirical study comparing between the muddient-based mechanism and the gated mechanism.
For an expressive representation of a complex physical system, p()) is usually not a feature of single degree but a direct-sum concatenation of series of components (10)), p())a, p(A)s ,..} at multiple degrees, i.c., (()) = (l, l2, ...), where some components share the same de- gree while others differ. In this case, it becomes necessary to extend the decomposition operator, i.e., CGDecomp(.) in Eq. [ to accommodate various components of de- grees. Moreover, in the context of neural networks, intro- ducing learnable parameters W and more feature channels C may improve the model capacity. Based on these con- siderations, when constructing the encoding module, the CGDecomp(-) operation in Eq. ① can be expanded as CGDecompent ()), and its outputs can be expanded as u((), as shown in Eq. 2
5. Method
As shown in Fig. [ taking 2 in Eq. (T) serving as the bridge between Theorem 4.1 and Theorem 4.2, we propose method for learning expressive non-linear representations that satisfy the SO(3)-equivariaace property outlined in Eq. (TO). The core of our method can be abstractly referred to as TraceGrad. At the label level, it incorporates the SO(3)-invariant trace quantity tr(H- Hi) introduced in Theorem as a supervisory signal for learning z, i.e., the SO(3)-invariant internal representations of gnonlin (-) in The- orem 4.2 Given the attribute that invariance is preserved under non-linear operations, is encoded by (-)ntin()) to obtain non-linear expressive capabilities. At the repre- sentation level, by taking the gradient of = with respect to f. the non-linear expressiveness of 2 is transferred to the equivariant feature v, while maintaining strict SO(3)- cquivariance as we prove. Subsequently, merging v and f. and applying Inontin(-) (-) in a stacked manner can yield rich SO(3) -qquivaauv non-linear representations for inferring the (0(3))00 electronic-structure Hamiltonians.
u (k) = [u] U
(K)
(k)
C
( CGDecomp ext
x (k) W)
Σ
WE Ct (K),, @ pAN 0)
(2)
where WS represents learnable parameters, we is the th channel of (k),, CGDecomp (-) adds learnable parame- ters to CGDecomp(.) and expands its output from single channel to multiple channels. To further enhance the model capacity, we also expand the output of smonlin (.) to multi- (k)
ple channels as follows: ((k) [2(())
We define (!)) = aFūe and construct the new features (k) by (k)) = Ee v(()) is evident that these exten- sions maintain the -ova)-inarrianee of u(k) and (*),, as well as the -quivaarine of (k)) For different fea- ture components, such as p p(k) and (k)4 in Eq. (2). as long as they share the same degree (l, = [1,),, they can be used to construct SOv3)-invariant and squii-equivarian features through the aforementioned operations, even if their mag- nitudes and directions differ. These operations ultimately
(x(*)
5.1. Encoding and Decoding Framework
Our encoding framework corresponds to total of K en- coding modules, which are sequentially connected to form a deep encoding framework. For the k-th module (1 < k < K), it first introduces an SO(3)-equivariant back- bone encoder, e.g., the encoders from DecpH-E3 (Gong et al. 2 or QHNet (Yu et al. 2023b) which is com-
4

---

TraceGrad
Input Physical System
Superviee Tratning
The & th Encoding Module
Initial Fearunes
Backbone cquivariant Feature Encoder
Sojjiaaaiarian Features F
SOL3--
500))-quivariant Regresion Target H
equivariantt
Decoder
Direct Product & Gorbac-GGordan Decomposition
SOoppiiiivaiian Fostures (V"") with Non-linas lineer Expressiveness
Returning Componcuts Combining
with Mahti-chanels
Learnable Parameters
SOD) invatiant
8f
Festures 1"
S0(3) inveriant Testures (x with Non-linear
Non-linrar
Neural Modules
S03) Invariant Decader
30X3)-navarinn
Quantity
T= H
Sapervised Trainins
Figure I: The proposed framework for learning -Quui-equivariant representations with strong non-linear expressiveness regress the eqqivvarinn clectronic-structure Hamiltonian.
enable the encoding of new features with flexible directions and magnitudes, thereby enhancing the expressive power of the resulting representations. (kk is combined with f() in
ficient to regularize the relative scale of the two loss terms:
((lossH,lossp) = À No.Grad loss
(4)
a residual manner like (ak) (1k) + v(k) as the output of
losST
the & th encoding module.
where No Grad(-) denotes gradient discarding when cal- culating such coefficient, as this coefficient is only used for adjusting the relative scale between the two loss terms and should not itself be source of training gradients, other- wise it would counteract the gradients from lossy in Eq. 3 All of the encoding and decoding modules are trained
We follow previous literature (Gong al., 2023 Yu et al.) 2023b) to send the features from the last layer of the encoder, i.e., (K) in our framework, into the SO(3)- cquivariant decoder to regress the predictions of H. For this, we can directly utilize the mature design of the SO(3)- equivariant decoders in DeepH-E3 or QHNet. The new part the decoding phase introduced by our method is the SO(3)-invariant decoder, consisting of feed-forward layers taking = (z()), T.
jointly by Eq. 3
6. Experiments
z (k) z [K)] the input to predict
6.1. Experimental Conditions
We apply our theory and method to the electromic-structure Hamiltonian prediction task, and collect results eight benchmark databases, i.e., Monolayer Graphene (MG), Monolayer MoS2 (MM), Bilayer Graphene (BG), Bi- layer Bismuthene (BB), Bilayer Bi2Te3 (BT), Bilayer Bi2Se3 (BS), Qt9-stable (QS), and QH9-dynamic (QD). The first six databases, consisting of periodic crystalline systems with elements like C, Mo. S, Bi, Te and Se, are from the DeepH benchmark series (Li al. 2022 Gong et al. 2023). The last two databases are from the QH9 benchmark series (Yu et al. 2023a). composed of molec- ular systems with elements like C. H. O, N and F. These databases present diverse and complex challenges to a re- gression model. Regarding MG, MM, and QD, as their samples are prepared from an temperature environment
5.2. Training
The training loss function is shown the following:
min loss loss + p[lossH losst) lOSST,
loss Error H.H') H"), losst Error T
(3)
where A denotes all of the learnable parameters of our framework, H, T and H', T respectively denote the pre- dictions and labels of H and T. In order to prevent the nu- merical disparity between the two loss terms and stabilize the training for both of the sQuii-equi and SO(3)- invariant branches, we apply p(lossH, lossT), i.e., a coef-
5

---

TraceGrad
Table I: Experimental results measured by the M M Graphene (MG), Monolayer MoS2 (MM), Bilayer Graphene (BG), Bilayer Bismuthene (BB), Bilayer Bi2Te3 (BT) and Bilayer Bi2Se3 (BS) databases, where the superseripts nt and respectively denote the non-twisted and twisted subsets.
allt M MAE cha and M AA b metrics (meV) on the Monolayer
MG
MM
Methods
MAE (4)
MAEH
MAEH cha.s
MAEH
MAEH
M MAEH
MAEH
cha.b
cha.s
cha..
DeepH-E3 (Baseline)
0.251
0.357
0.362
0.406
0.574
1.103
Deepll-E3+Tracerrad
0.175
0.257
0.228
0.285
0.412
0.808
Methods
BG
BG
DeepH-E3 (Baseline)
0,389
0.453
0.644
0.264
0.429
0.609
TrepH-E3+TraeGGrad
0.291
0.323
0.430
0.198
0.372
0.406
Methods
B
BB
DeepH-E3 (Baseline)
0.274
0.304
1.042
0.468
0.602
2.399
0.226
0.256
0.740
0.384
0.503
1.284
Methods
BT
BT*
DeepH-E3 (Baseline)
0.447
0.480
1.387
0.831
0.850
4.572
DeepHl-F3+Traccrrrd
0.295
0.312
0.718
0.735
0.755
4.418
Methods
BSa
BS¹
DeepH-E3 (Baseline)
0.397
0.424
0.867
0.370
0.390
0.875
DeepH-E3+ Trarcerra
0.300
0.332
0.644
0.291
0.302
0.674
three-hundred Kelvin, the thermal motions lead to complex non-rigid deformations, increasing the difficulty of Hamil- tonian prediction. For BG, BB, BT, and BS, the twisted structures, with interplay of SQuua-equivariant effects and van der Waals (vdW) force variations bring signifi- cant generalization challenges, which are further exacer- bated by the absence of any twisted samples in the training sets. Besides, BB, BT, and BS exhibit strong spin-orbit coupling (SOC) effects, which further increase the com- plexity Hamiltonian modeling. For the QS database, the "ood' strategy from the official settings is used split the training, validation, and testing sets, ensuring that the atom number of samples do not overlap across the three subsets. For the QD database, the 'mol' strategy provided by Yu et al. (2023a) is applied to split the training, validation, and testing sets, ensuring that there are no thermal motion samples from the same temporal trajectory across the three subsets. The 'mol' and 'ood' strategies aim to assess the regression model's extrapolation capability with respect to the number of atoms and the temporal trajectories, respec- tively. Detailed statistic information of these databases can be found in the Appendix E
We use a comprehensive set of metrics to deeply evalu- ate the accuracy performance of deep learning clectronic- structure Hamiltonian prediction models. On the databases from the DeepH benchmark series (Li al. 2022 Gong et al., 2023), set of Mean Absolute Error (MAE) metrics between pre- dicted and ground truth Hamiltonians, including MAE", for measuring average MAE all samples and matrix elements, MAE for measuring the MAE of chal- lenging samples where the baseline model performs the worst, M MAE for measuring the MAE of different ba- sie blocks in the Hamiltonian matrix, and M MAFH for measuring the MAE on the most challenging Hamiltonian block where the baseline model shows the poorest perfor- mance (with the largest M AAEHH.)). These metrics com- prehensively reflect the accuracy performance, covering not only the average accuracy but also the accuracy on difficult samples and challenging blocks of the Hamilto- nian matrices. On the two databases from the QH9 bench- mark series, adopt the metrics introduced by their orig- inal paper YYu et al. 2023a), including MAE of Hamilto- nian matrices, which are further subdivided into MAEH for measuring average MAE, MAEH diog for measuring MAE of Hamiltonian matrix formed by an atom with it- self, and M MEH Ao for measuring MAE of Hamilto-
follow Yin et al. (2024 to adopt
Implementation details of our method for experiments on these databases are presented in Appendix
6

---

TraceGrad
Table 2: Experimental results measured by the M AA M MAE M MA non diag' M AES and Sim( metrics on the -t9-st (QS) and QH9-dynamic (QD) databases respectively using 'ood' and 'mol' split strategies YYu et al. (023a).
+ means lower values correspond to better accuracy, while f means higher values correspond to better performance. The units of MAE metries are meV, while Sim({) is the cosine similarity which dimensionless.
QS
Methods
MAE (L)
Sim(t) (f)
M MAR
MAE
M MAR
MAF
diag
dond
QHNet (Baseline)
1.962
3.040
1.902
17.528
0.937
QrNeteTraceGra
1.191
2.125
1.139
8.579
0.948
Methods
QD
QHNet (Baseline)
4.733
11.347
4.182
264.483
0.792
QHNet+TraceGral
2.819
6.844
2.497
63.375
0.927
nian matrices formed by different atoms; well as the MAE (M AE of occupied orbital energies e induced by the predicted Hamiltonians and compared to the ground truth ones, and the cosine similarity (Sim(15) between the electronic wavefunctions vi induced by the predicted and ground truth Hamiltonians. é and & are crucial downstream physical quantities for determining multiple properties of the atomic systems as well as their dynamies, highly re- flecting the application values of the Hamiltonian regres- sion model.
peated experiments. Regarding the metric of M block for every Hamiltonian block, due its large data volume, we just present its values for all databases from the DeepH series in Appendix G We use the same fixed random seed as adopted by DeepH-E3 for all random processes in exper- iments on these six databases. As a result, for each Hamil- tonian MAE result presented in Table [] the standard de- viation across independent repeated experiments does not exceed 0.007 meVV and is negligible.
From results presented in Tables D we could find that the proposed TraceGrad method dramatically enhances the accuracy performance of the baseline method DeepH-B3, both on average and for challenging samples and blocks, both on the non-twisted samples and the twisted sam- ples. Specifically, on the corresponding datasets, Trace- Grad lowers down the M MAFH all and M A of DeepH- E3 with relative ratios of up to 34% and 35%, respectively, Furthermore, from the results included in Appendix G TraceGrad significantly improves the performance for the vast majority of basic blocks. Particularly, for the blocks where DeepH-EB perform the worst, TraceGrad reduces the MAE (MAE by a maximum of 48%. The lead- ing performance on the MG and M M databases prepared at hureeehundeed Kelvin temperature demonstrates the bustness of our method against thermal motion. The high accuracy on the BB, BT,, and BS databases, which have strong SOC effects, indicates our method's strong capabil- ity to model such effects. The excellent performance on the BG BB*, BT and BS subsets showcases the method's superior generalization twisted structures, which are present in the training data. The outstanding performance such samples highlights the good potential for study- ing twist-related phenomena, a hot research topic that may bring new electrical and transport properties (Cao et al.. 2018 Wang et al.) 2024a He et al.) 2004).. Additionally, the BG BB', BT and BS¹ subsets contain significantly
6.2. Results and Analysis
We compare experimental results from two setups: the first one is the experimental results of the baseline SO(3)- equivariant regression model (Gong et al.] 2023) Yu et al.) 20236) for Hamiltonian prediction, and the second one is the experimental results of extending the architecture and pipeline of the baseline model through the proposed TraceGrad method, which incorporates non-linear expres- siveness into the SO(3)-equivariant features of the base- line model with the gradient operations of SO(3)-invariant non-linear features learned under the supervision of the trace targets. We choose DeepH-E3 (Gong et al.) 2023) as the baseline model for databases from the DeepH bench- mark series (Li et al.) 2022] Gong et al. 2023); and we choose QHNet (Yu et al. 2023b) as the baseline model for databases from the QH9 benchmark series (Yu et al. 2023a). They the respective state-of-the-art (SOTA) methods with strict -quivaaquee the correspond- ing databases.
We list the results of DeepH-E3 and DeepH-E3+ TraceGrrad in Table [ for databases from the DeepH benchmark se- ries, reporting the values of M A AF
M M cha.s' and M Al a. The results of DeepH-E3 to be compared are copied from Yin al.] (2024). The results of DeepH- are the average from 10 independent re-

---

TraceGrad
larger unit cells compared to the training set (see Appendix for statistics of their sizes), yet method still excels these subsets as measured by the multiple MAE metrics, demonstrating its good scalability on the sizes of atomic systems it handles.
7. Summary of Appendices
Due to the page limit of the main manuscript, we have to provide some important contents in the Appendices, sum- marized as follows:
In Table 2 we present the results QHNet and QH- Net+TraceGrad under the metrics of M MAE alt- M MAE!! diog' M AE diag' MAE', and (im( for the QS and QD databases. The results of QHNet to be compared are taken from their original paper (Yu et al., 2023a). and for the unification of MAE units, we convert the units of MAE from 10-6 Hartree (En) in the original paper to meV (LE, 27211.4 mcV). The results of QHNet+TraceGnin are the average from 10 independent repeated experiments. To ensure reproducibility, we use the same fixed random seeds as employed in QHNet for all random processes in the experiments on the QS and QD databases. As a re- sult, for each Hamiltonian MAE result presented in Table 2 the standard deviation across repeated experiments is greater than 0.009 V and is also negligible.
Appendix A provides the definitions of foundational concepts relevant to the problem addressed in this pa- per.
Appendix B describes the application task in this pa- per, i.c., the electronic-structus Hamiltonian predic- tion task.
Appendix C formalizes Hamiltonian prediction as an S0(3)-equivarian non-linear representation learning problem.
Appendix D provides the proofs for all of the proposed theorems.
Appendix experimental databases.
presents the detailed information of the
The results presented in Table 2 demonstrate that the pro- posed TraceGrad method significantly enhances the accu- racy of the baseline QHNet model across all metrics on the QS and QD databases. Specifically, TraceGrad reduces M ME'H M MAE Aooo.ddag M A deag AAE,, and Sim(y) of QHNet with relative reductions by maximum of 40%, 39%, 40%, 76%, and 17%, respectively, on the correspond- ing databases. The significant accuracy improvements the QS database, partitioned using the 'ood' split strat- egy (Yu et al.) 2023a) without scale overlapping among the training, validation, and testing sets, once again demon- strate the method's strong generalization capabilities across different scales of atomic systems. Meanwhile, perfor- mance the QD database under the 'mol' strategy (Yu et al. 2023a). which partitions the training, validation, and testing sets with samples from completely different ther- mal motion trajectories, highlight our method's robustness in generalizing to new thermal motion sequences. Further- more, the substantial improvement in the prediction accu- racy of e, i.e., occupied orbital energies crucial for deter- mining electronic properties such as optical characteristics and conductivity in atomic systems, and tb, i.e., the elec- tronic wavefunctions essential for understanding electron distribution and interactions, underscores the potential val- ues of our method for applications like material design, molecular pharmacology, and quantum computing.
Appendix presents the implementation details of the experiments.
non_diag
Appendix compares the block-level MAE statis- tics (M A A Bock and MAF for DeepH-E3 and CeepH-EE+Tracedde
Appendix H reports the results of ablation study for each individual mechanism of our framework and the quantitative comparison between the pro- posed gradient- based mechanism (Grad) with the ex- isting gated activation mechanism (Gate). Experi- mental results indicate that each individual mecha- nism of our method can contribute individually to the performance. Moreover, their combination pro- vides even better performance. Experimental results also demonstrate better accuracy performance of the proposed gradient-based mechanism compared the gated mechanism.
Appendix[[proviae a theoretical analysis of the com- putational complexity advantage of our method over traditional DFT calculations.
Appendix | makes joint discussion on GPU time costs and performance gains brought by our Trace- Grad method. It underscores the superiority of the TraceGrad method in improving accuracy perfor- mance while maintaining time efficiency.
We also evaluate the acceleration performance of our method for DFT computations and report the results in Ap- pendix K. Experimental results show that while combining TraceGrad introduces only slight increase from the infer- ence time of QHNet, it delivers significant improvements in accelerating the convergence of DFT methods.
Appendix K reports the acceleration performance of our method for traditional DFT algorithms.
8

---

TraceGrad
Appendix corresponds the synergy of method with approximately -Que)-eqiivarian methodol- ogy
ten lack transparency, hindering comprehensive under- standing of the learning outcomes and limiting our ability to gain deeper insights. We believe it is important in- vestigate the interpretability of such models, particularly in terms of how they apply physical knowledge in a com- prehensible way. Additionally, is crucial to continually improve the correctness and fairness of deep learning mod- els on this area. Ensuring high-quality and diverse training data, implementing sound model designs, and performing ongoing validation and refinement are necessary to guar- antee model accuracy and the broad applicability of their results.
Yin et al. 2024).
Appendix M reports the experimental results of ap- plying the proposed TraceGrad method to energy and force field prediction tasks. The results show that our approach improves the prediction aceuracy for both quantities, demonstrating its generality.
The codes of this work available in this link
8. Conclusion
References
We propose a theoretical and methodological framework to tackle the issue of reconciling powerful non-linear expressiveness with SO(3)--quivariance deep learning paradigm for electronic-struccurre Hamiltonian prediction, through deeply investigating the mathematical connections between SO(3)-invariant and SO(3) -quiivaaain quantities, as well as their representations. We first constructs SO(3)- invariant quantities from SO(3)-equivariant regression tar- gets, using them to train informative SO(3)-invariaan non- linear representations. From these, -qqiivaanva features are derived with gradient operations, achieving non-linear expressiveness while maintaining strict SO(3)- cquivariance and physical dimensional correspondence to the target labels. Experimental results on eight benchmark databases covering multiple types of systems and challeng- ing scenarios show substantial improvements on the state- of-the-art prediction accuracy of deep learning paradigm on electromic-structure Hamiltonian prediction. Our method also significantly promotes the acceleration performance for the convergence of traditional DFT methods.
Cao, Y., Fatemi, V., Fang, S., Watanabe, K., Taniguchi, T., Kaxiras, E., and Jarillo-Herrero, P. Unconventional su- perconductivity in magic-angle graphene superlattices. Nature, 556(7699):43-50) 2018.
Chmiela, S., Tkatchenko, A., Sauceda, H. E., Poltavsky, I., Schütt,, K. T., and Müller, K.-R. Machine learning ac- curate energy-conserving molecular force fields. Science Advances, 3(5):e1603015, 2017.
Chmiela, S., Sauceda, H. E., Müller, K.-R., and Tkatchenko, A. Towards exact molecular dynamics sim- ulations with machine-learned force fields. Nature com- munications, 91):38877, 2018.
Dresselhaus, M. S., Dresselhaus, G., and Jorio, A. Group theory: application to the physics of condensed matter. Springer Science & Business Media, 2007.
Fuchs, F., Worrall, D. E., Fischer, V., and Welling, M. Se(3)-transformers: 3d roto-translation equivariant at- tention networks. In NeurIPS, 2020.
Acknowledgements
Geiger, M. and Smidt, T. E. e3nn: Euclidean neural net- works. CoRR, 0is/2207.094555 2022.
This work was supported by the the Strategic Priority Re- search Program of Chinese Academy of Sciences (Grant Number XDB0500201) and the National Natural Science Foundation of China (Grant Number 12134012).
Gong, X., Li, H., Zou, N., Xu, R., Duan, W., and Xu, Y.
General framework for (3)-equivariant neural network representation of density functional theory hamiltonian. Nature Communications, 14(1):2848, 2023.
Impact Statement
He, M., Cai, J., Zheng, H., Seewald, E., Taniguchi, T. Watanabe, K., Yan, J., Yankowitz, M.. Pasupathy, A., Yao, W., et al. Dynamically tunable moiré exciton ry- dberg states in monolayer semiconductor on twisted bilayer graphene. Nature Materials, 2024.
This paper presents work whose goal is to advance the field of machine learning driven by physics research. We recog- nize that, although our research field has not shown direct negative social ethical consequences, there are several important issues that require our vigilance and attention. Currently, although our method yields accurate predictions, the decision decision-making processes of deep learning systems of-
Hohenberg, P. and Kohn, W. Inhomogeneous electron gas. Physical Review, 136(3B):B864, 1964.
Jones, R. O. Density functional theory: Its origins, rise to prominence, and future. Reviews of modern physics, 87 (3):877, 2015.
KVWrrcy IKWrrccy_fIx800555 wwtp://asxxxsseseoeoeeeeee file/d/ usp-sharing

---

TraceGrad
Kohn, W. and Sham, L. J. Self-consistent equations includ- ing exchange and correlation effects. Physical Review, 140(4A);A1133, 1965.
Sun, Q.. Berkelbach, T. C.,, Blunt, N. Booth, G. H., Guo, S., Li, Z.,, Liu, J., McClain, J. D., Sayfutyarova, E. R., Sharma, S., et al. PySCF: the Python-based simulations of chemistry framework. Wiley Interdisciplinary Re- views: Computational Molecular Science, 8(1):e134)) 2018.
Lehoucq, R. B., Sorensen, D. C., and Yang, C. ARPACK Users' Guide: Solution of Large-Scale Eigenvalue Problems with Implicitly Restarted Arnoldi Methods. SIAM, Philadelphia, PA, USA, 1998. ISBN 978- 0898714074, URL https:/ /ttp:://ppppppp/ siam. org/ doi /ta//bok//a0 10. .T..11177.78888811988888
Thomas, N., Smidt, T., Kearnes, S., Yang. L., Li, L., Kohlhoff, K., and Riley, P. Tensor field networks: Rotation-and cannslation-quivariinn neural networks for 3d point clouds. arXiv preprint arXiv: 802.08219. 2018.
Li, H., Wang, Z., Zou, N., Ye, M., Xu, R., Gong, X., Duan, W., and Xu, Y. Deep-learning density functional theory hamiltonian for efficient ab initio electronic-struccure calculation. Nature Computational Science, 2(6):367- 377, 2022.
Unke, O., Bogojeski. M., Gastegger, M., Geiger, M., Smidt, T., and Müller, K.-R. eqqiivrran predic- tion of molecular wavefunctions and electronic densities. In NeurlPS, volume 34, pp. 14434-14447, 2021.
Li, Y., Xia, Z.,, Huang, I.., Wei, X., Harshe, S., Yang. H., Luo, E., Wang, Z., Zhang, J., Liu, C., al. Enhancing the scalability and upplicability of kohn-sham hamiltoni- ans for molecular systems. In ICLR, 2025.
Wang, C., Zhang, X.-W., Liu, X., He, Y., Xu, X., Ran, Y., Cao, T., and Xiao, D. Fractional chern insulator in twisted bilayer mote 2. Physical Review Letters, 132(3), 2024a.
Liao, Y. and Smidt, E. Equiformer: Equivariant graph attention transformer for 3d atomistic graphs. In ICLR, 2023.
Wang, Y., Li, H., Tang, Z., Tao, H., Wang, Y., Yuan, Z., Chen, Z., Duan, W., and Xu, Y. Deeph-2: En- hancing deep-learning electronic structure via an equiv- ariant local-coordinate transformer. arXiv preprint 2rXhv:2401.17015 2024b.
Liao, Y., Wood, B. M., Das, A., and Smidt, T. E. Equiformerv2: Improved equivariant transformer for scaling to higher-degree representations. In ICLR, 2024.
Wang. Y. Li, Y., Tang, Z., Li, H., Yuan, Z., Tao, H., Zou, N., Bao, T., Liang, X., Chen, 7., et al. Universal ma- terials model of deep-learning density functional theory hamiltonian. Science Bulletin, 2024c.
Lin, P., Ren, X., Liu, X., and He, L. Ab initio elec- tronic structure calculations based on numerical atomic orbitals: Basic fomalisms and recent progresses. WIREs Computational Molecular Science, 14(n/a)=e1687, 2023.
Weiler, M., Geiger, M., Welling, M., Boomsma, W., and Cohen, T. S. 3d steerable cnns: Learning rotationally equivariant features in volumetric data. NeurlPS, 2018.
Nagy, Á. Density functional. theory and application to atoms and molecules. Physics Reports, 298(1):1-79, 1998.
Yin, S., Pan, X., Zhu, X., Gao, T. Zhang, H., Wu, F,, and He, L. Towards harmonization of so ())-equivariance and expressiveness: a hybrid deep learning framework for electronic-structure hamiltonian prediction. Machine Learning: Science and Technology, 5(4):045038, 2024.
Passaro, S. and Zitnick, C. L. Reducing SO(3) convolutions to SO(2) for efficient equivariant gnns. In ICML, pp. 27420-27438, 2023.
Resnick, R. Introduction to special relativity. 1991.
Yu, H., Liu, M., Luo, Y., Strasser, A., Qian, X., Qian, X., and Ji, S. QH9: A quantum hamiltonian prediction benchmark for QM9 molecules. In NeurlPS, 2023a.
Schrodinger, E. Quantisierung als eigenwertproblem. An- nalen der Physik, 384(4):361-376, 1926.
Yu, H., Xu, Z., Qian, X., Qian, X., and Ji, Efficient and equivariant graph networks for predicting quantum hamiltonian. In ICML, pp. 40412-40424 2023b.
Schütt, K. T., Arbabzadah, F., Chmicla, S., Müller, K. R., and Tkatchenko, A. Quantum-chemica insights from deep tensor neural networks. Nature communications, (1):13890,, 2017.
Zhang, H., Liu, C., Wang, Z., Wei, X., Liu, S., Zheng, N., Shao, B., and Liu, T. Self-consistency training for censity-functionall ICML, 2024.
Schütt, K. T., Gastegger, M., Tkatchenko, A., Müller, K.-R., and Maurer, R. J. Unifying machine learning and quantum chemistry with a deep neural network for molecular wavefunctions. Nature Communications, 10 (1):5024, 2019.
hamiltonian prediction. In
Zhang, X., Wang, L., Helwig, J., Luo, Y., Fu, C., Xie, Y., Liu, M., Lin, Y., Xu, Z., Yan, K., Adams, K., Weiler,
10

---

TraceGrad
M., Li, X., Fu, T., Wang, Y., Yu, H., Xie, Y., Fu, X., Strasser, A., Xu, S., Liu, Y., Du, Y., Saxton, A., Ling, H., Lawrence, H., Stärk, H., Gui, S., Edwards, C., Gao, N., Ladera, A., Wu, T., Hofgard, E. F., Tehrani, A. M., Wang, R., Daigavane, A., Bohde, M., Kurtin, J., Huang, Q. Phung, T., Xu, M., Joshi, C. K., Mathis, S. V., Az- izzadenesheli, K., Fang, A., Aspuru-Guzik, A., Bekkers, E., Bronstein, M. Zitnik, M., Anandkumar, A., Ermon, S., Liò, P., Yu, R., Günnemann, S., Leskovec, J., Ji, H., Sun, J., Barzilay, R., Jaakkola, T., Coley, C. W., Qian, X., Qian, X., Smidt, T., and Ji, S. Artificial intelligence for science in quantum, atomistic, and continuum systems. arXiv preprint arXiv:2307.08423; 2023.
Zhong, Y., Yu, H., Su, M., Gong, X., and Xiang, H. Trans- ferable equivariant graph neural networks for the hamil- tomians of molecules and solids. npj Computational Ma- terials, 11): 2023.
Zitmick, L., Das. A., Kolluru, A., Lan, J., Shuaibi, M., Sri- ram, A., Ulissi, Z. W., and Wood, B. M. Spherical chan- nels for modeling atomic interactions. In NeurlPS, 2022.

---

TraceGrad
Appendices
A. Definition of Concepts
This section provides definitions of foundational concepts relevant to the problem addressed this paper. For additional background, please refer to the book by Dresselhaus et al. (2007).
Definition .1. Group. set G, denoted as {.... 9.... }. equipped with binary operation is called group if it satisfies the following four axioms:
1. Closure: For any f.g E G, the result of their operation f 9 is also an element of G: f.g...g. E G.
2. Associativity: For all f,g, h E G, the operation satisfies (f - g) h f (g-h).
3. Identity Element: There exists a unique element G (called the identity) such that for all f =
4. Inverse Element: For each f E G. there exists unique element f-1 G (called the inverse) such that f f
f-1 f=c.
Definition .2. SO(3) Group. The special orthogonal group SO(3) is the group of all 3 x 3 orthogonal matrices with determinant I. Formally, it is defined as:
SO(3) (R E 3
R I, (R) = =];
where R. denotes the transpose of R, and / is the 3 identity matrix. The elements of SO(3) represent rotations in three-dimensional Euclidean space.
Definition .3. Group Representation. representation of group G on a tensor space T(V) homomorphism p from G the general linear group GL(T(V)), the group of all invertible linear transformations on T'(V). Here, T(V) represents the tensor space associated with the vector space V, encompassing all tensors that can be formed from elements of V. Formally, the homomorphism is defined as:
P:G P
such that for all 92 E G.
p(g192) p(g))
and p(e)= I, where is the identity element of G, and / is the identity transformation on T(V).
Definition .4. Irreducible Representation. representation p - GL(V) of a group G on vector space V is said be irreducible if there no proper subspace W V such that P(g) C W for all g E G. In other words, the only invariant subspaces under the group action the trivial subspaces (0) and V itself. If such nontrivial invariant subspace exists, the representation is said to be reducible.
Definition .s. SO(3) Group Representation. A representation of the special orthogonal group SO(3) on a vector space V is homomorphism
p P:SO(3) + GL(V),
where GL(V) denotes the group of all invertible linear transformations on V. The irreducible representation of SO(3), each labeled by a degree I, often corresponds to the quantum number for angular momentum in quantum physics.
Definition .6. Wigner-D Matrices. One of the irreducible representations of the rotation group SO(3) is given by the Wigner-D) matrices D' (R), where is the degree of the representation, typically associated with angular momentum quan- tum number in quantum mechanics.
The Wigner-D D matrices are defined as:
DJ D'mmm m (R) = (1, m |R] m),
12

---

TraceGrad
where m) are the eigenstates of the angular momentum operator 1,2 and its z-component Ly with m and m' being the magnetic quantum numbers corresponding the eigenvalue of f.g. The Wigner-D matrices D' (R) encode how the angular momentum states ,1, m) transform under the action of rotation.
Definition .7. Equivariance with Respect to a Group. Let G be group, and let Pr(v) : G - GL(T(V)) and pr(w) : G GL(T(W)) be representations of G tensor spaces T(V) and T(W), respectively. A map T(V) T(W) is said to be equivariant with respect to the group G if the following condition holds:
I (PF (g) v) - (r( f(v) for all E T(V) and g E G.
where generally denotes the operation defined on the tensor space.
Definition .8. Invariance with Respect to a Group. Let G be group, and let Pr(v) G + GL(((V)) be a representa- tion of G on tensor space TVV). A function f: T(V) + T() said to be invariant under the group if the following condition holds:
{(p(())9)) u) = f(v) for all v T(V) and g E G.
This definition indicates that the function f remains unchanged under the action of the group G.
Definition .9. Direct-Product State. Let V1 and V2 be two vector spaces, and let (uz) V1 and (12) e V2 be arbitrary elements of these spaces. The direct-product state of (v1) and (u2) is defined as the element (v1) @ (02) in the direct-product space Vi @ V2. The direct-product state represents all possible combinations of the elements of V11 and V22 and forms new vector space that captures the joint state of two systems. The action of a group on the direct-produet state is defined by the direct product of the individual actions on (u1) and (12), i.c., the effect of the group action on the direct- prrect-product space is the direct product of the actions on each individual space:
g (lu1) (02)) (g (11)) & (g (o2)), G.
Definition .10. Physical Quantity Formed by the Direct-Product of Two Degrees. In quantum mechanics, direct- product spaces are used to describe the combined states of systems with distinct degrees of freedom, such as angular momentum. The combined state captures both the independent action of each degree and their joint transformation under rotations governed by the SO(3) group.
Let degrees 4p and lg represent the angular momentum quantum numbers of two systems. A direct- prrect-product state Q Qla04e (21, (1), combines these degrees and transforms under the SO(3) group according to:
Q = - Q (D'
where D's (R) E (2) and D (R) E (2) are the Wigner-D matrices for the degrees lp and by- respectively, and denotes the conjugate transpose, ensuring unitary transformations.
Definition .11. Direct-Sum State. Let Vi and V22 be two vector spaces, and let (wn) E Vi and (02) E V/2 be arbitrary elements of these spaces. The direct-sum state of (u)) and (u2) is defined the element (u1) @ (o2) in the direct-sum space Vi e V22
The direct-sum space V1 @ V22 consists of ordered pairs of elements, where each element is drawn from one of the original spaces. The operations of vector addition and scalar multiplication in V @ V2 are defined component-wise:
(a1, 02) + (bi, b2) = (ar + b1, a2 + b2), c (ai, 02) = (c- a,, c a2);
for a2), C.
02), (bi, bz) Vi ④ V/2 and F, where F is field, typically either the real numbers I the complex numbers
The direct-sum state (w1) + (22) represents a combination where the components remain in their respective vector spaces and do not interact with each other. The action of a group G on the direct-sum state is defined by its independent action on each component:
g (111) e (o2)) = (g (11)) @ (g 112)), Vg E G.
13

---

TraceGrad
14
1/
11 61
a 1
1/
11
1 1 1
:
basic block H
1
81
Chattr of alomic painstso Hannikanian matrices paritionull from the Hamilicaion matire the whole gystem
The atomic pairwise Hamiltonian matrix Hidefined by atomic pair (ij), where is within the cutoff radius
Figure 2: Illustration of atomic pairwise Hamiltonian matrices partitioned from the complete Hamiltonian matrix of the whole system. Each Hamiltonian matrix contains multiple basic blocks, denoted as H4
Definition .12. Clebsch-Gordan Decomposition for SO(3) Group.
In quantum mechanics, the Clebsch-Gordan decomposition is used to describe the coupling and decoupling of two angular momentum systems. For two angular momentum quantum numbers lp and l,p their combined state can be decomposed into series components in the direct-sum state. Taking Q', defined in Definition as an example, it can be decomposed as follows:
De 0l))
q' d'm
Σ Mp..n
Chipddy my
= L= -bpl
where CGDecomp()) denotes the Clebsch-Gordan decomposition operator, Q'a the elements of the quantity Q4181y in the direct-produc state, Cm,mp.m, are the Clebsch-Gordan coefficients, and q' denotes the decomposed results, with oi representing the component under quantum numbers / and m.
B. Application Task Description: Flletronic-strucctuee Hamiltonian Calculation
Density Functional Theory (DFT) (Hohenberg Kohn 1964 Kohn Sham 1965) has become a cornerstone modern electronic structure theory, playing a pivotal role condensed matter physics, quantum chemistry, and materials science. Introduced in the 1960s through the foundational work of Hohenberg, Kohn, and Sham, DFT provides a framework for studying many-electron systems by replacing the computationally expensive many-body wavefunction with the electron density p(r) as the fundamental variable. This reformulation significantly reduces computational complexity while preserv- ing essential quantum mechanical effects, enabling researchers to investigate systems of practical interest with managcable computational resources. Over the decades, DFT has proven to be instrumental in calculating orbital energies and band structures, optimizing structural geometries, and exploring a wide range of material properties, underscoring its versatility and importance across diverse scientific disciplines.
14

---

TraceGrad
At the heart of DFT lies the Kohn-Sham equation (Kohn & Sham set of ainle-particle equations. These equations are expressed as:
1965), which simplifies the many-body problem into
(r)(r) = (r), subject to H
V2 + Van(r) + VHxc[ol(r)
(5)
2m
where H is the Hamiltonian operator, the Hartree-exchange corcchange-corrlattoon potential (r) = (r) + Vxclpl(r) is a
functional of the electron density p(r). The electron density obtained from the Kohn-Sham orbitals (r) as:
M
p(r) = Σ (x))1
where M represents the total number of electrons.
Atomic orbitals offer computationally efficient basis for electronic structure caleulations because they require fewer basis functions compared to other types of basis sets (Lin et al.) 2023). Additionally, their inherently localized nature makes them particularly advantageous for large systems. When expressed in an atomic orbital basis set, the Kohn-Sham equations can be formulated into as generalized cigenvalue problem:
HC = eSC,
(6)
where H is the Hamiltonian matrix incorporating the contributions from Vex and Vixc. S the overlap matrix, C contains the orbital coefficients, and represents the eigenvalues of the system.
The atomic orbitals used the basis functions typically take the form:
(r) (r(rr) (m,,), (0,
where (r) is the radial part of the wavefunction, and m,,, (6, Φ) are the spherical harmonics describing the angular dependence. These orbitals are localized around their respective atomic centers, decay rapidly as r increases, and are truncated to zero beyond a cutoff radius re.
The matrix elements of the Hamiltonian and overlap matrices, HE and Spe, are defined as:
H pv - (200))))))))))
Sp - (T)))))) dr.
The localized nature of atomic orbitals ensures that matrix elements are non-zero only when the orbitals o and Φ. overlap according their cutoff radius. Beyond this range, the corresponding matrix elements are treated as zero, leading to sparsity in the Hamiltonian and overlap matrices. This sparsity significantly reduces computational effort and memory usage, especially for large systems.
Solving Hav involves an iterative self-consistent process: starting with an initial guess for ( (r). the potentials VHXc[p] (r) are updated, and the equations are solved repeatedly until p(r) converges. This process can be illustrated as (r)) + vi(1) (r)) H 12 (1) (r) (2)
((T) (r), until the charge density p(r) converges. At that point, the Hamiltonian matrix is outputted by (T) - H from which downstream physical properties such as orbital energies and band structures are induced, determining the electronic, magnetic, and transport characteristics of the electron system.
Despite the remarkable success of Kohn-Sham DFT in advancing fields such as materials science, energy, and biomedicine over recent decades (Nagy) 1998 Jones 2015). the challenge of high computational complexity remains unresolved. The main computational bottleneck lies in iterative solving the Kohn-Sham equation, especially in the matrix diagonalization of Eq. (6). which has complexity of O(N)), where N is the number of atoms in the system. To address this challenge, recent approaches (Li et al. 2022; Yu et al. 2023b) Gong et al. 2023) have applied the deep graph learning paradigm to predict the self-consistent Hamiltonian. These methods use the Hamiltonian matrix (T),, calculated from traditional DFT methods, as labels. This matrix can be partitioned into a series atomic pairwise Hamiltonian matrices { {H j E 2(1)), as shown in Fig. 2 where and j represent two atoms the system. Assuming that the angular momentum quantum numbers of the orbitals of atom are is E I2,,, }, and those of atom j are 13 E -1/ 12 ..., 12). ). the basic block of Hd can be expressed as H4, which is formed by the direct product between the 15. and
15

---

TraceGrad
These methods train an efficient graph neural network to directly predict each H from the 3D structure of the atomic system, thereby circumventing the extremely time-consuming self-consistent iterations. During the inference phase, these methods successfully reduced the computational complexity of calculating the Hamiltonian matrices to O(N), while show- ing good potential in generalizing to larger atomic systems, even though the training set, constrained by the expensive DFT labels, only includes smaller systems. Once the Hamiltonian matrices are obtained, many downstream physical properties can be efficiently calculated with (IV) complexity. A detailed analysis of the computational complexity is provided in Ap- pendix | This paper aims to tackle the challenge of predicting electronic Hamiltonians with high accuracy and reliability, while rigorously maintaining -quiviaquiv
C. Formulation of Hamiltonian Prediction as an SO(3)-Equivariant Non-linear Representation Learning Problem
As shown in Fig. 2 a basic block of Hamiltonian can be denoted as H which is formed by the direct product between degrees i and tg. It obeys the following S0(3) law:
HR)) = D')(R) H S0 (D)(())*
where i denotes the conjugate transpose operation. R E R333 is the rotational matrix, D (R) E (((++)(2)+))) and DVd (R) E +22 are the Wigner-D matrices of degrees I and l respectively; H(R)S0 (
denotes the transformed results of H H E g2 121))+)) through the rotational operation by R.
s in the firret-produect state can be further decomposed into series of direct-sum state components, i.e., h (11 1/ S 1 s i + Hg), which follows SO(3) -qquivaaara law mathematically equivalent to Eq. (7) but with a simpler form:
h hRR) = (R) h 11, -
1 s +
(8)
where h E R22+1 and hRR)² E g2++1 respectively denote the components with degree / before and after the rotational operation by R.
For ease of processing, the internal representations of -Quui-equivariant neural networks (Gong al., 2023 Liao & Smidt 2023) are typically in the direct-sum form. To obey the equivariant law of Eq. (8) for regressing Hamiltonian, these hidden representations must also satisfy the same form of equivariance:
f(R) (+)) = DR( (
where (()) E 22++ and ((R) E 2R respectively denote one channel of hidden representations with degree 7 before and after the rotational operation by R, at the k th hidden layer.
Due to the intrinsic complexity and non-linearity of Hamiltonians, the regression neural networks are supposed to equip with powerful non-linear mappings as feature encoders to fully capture the intrinsic patterns of the Hamiltonians, which crucial for precise and generalizable prediction performance. Meanwhile, the non-linear mappings, denoted gnentin (-), must also preserve -ouia-equi,, which is expressed as:
((R)(++NN= =D(RR) p( subject to p(k+1y Inontin (f(())
(10)
However, directly implementing (n)ntin(.) as neural network module with non-linear activation functions, such as Sigmoid, Softmar and SiLU, may result in the destruction of strict equivariance. How make (nonlin both the- oretically -quiiaaqui and capable of powerful non-linear expressiveness, and effectively apply it to the prediction of SO(3--cquivaaia complex physical quantities, i.c., electronic-stmueture Hamiltonians in this context, is the core problem this paper aims at solving.
D. Proofs of Theorems
Proof of Theorem 1. Under an SO(3) rotation represented by the rotational matrix R. H H is transformed as HR)):
H(R) D (R) H (R)),
16

---

TraceGrad
where D' (R) and (R) are the Wigner-D matrices for the degrees of i and 1/. respectively, corresponding to the rotation R.
The conjugate transpose of the transformed quantity is:
H(R) = D H (R))...
Using the cyclic property of the trace, which states that the trace of product of matrices remains unchanged under cyclic permutations (i.e., tr (ABCC) tr(BCA) tr(CA)),, and combining the properties that (R)(R) DD (R)) I and
(R) (R)) = I, we can rearrange the terms inside the trace as follows:
T(R) tr (H(R) H(R))) = tr (D' (R) H (R))) (D R H D (R)))
tr (D' H. D' (R)) tr (H-HH H & (R) (R) = tr(H = TT
Therefore, T (H H is invariant under SO(3) transformations, its SO(3))invarians proved.
Proof of Theorem 2. Under the given condition, the input feature in direct-sum state is 50(3)-cqiivarian, meaning that under an SO(3) rotation represented by R, it transforms as follows:
((R) = D'(R)
where D'(R) is the Wigner-D matrix corresponding to degree
First, according group theory, u CGDecomp(f & 0) is an SO(3)-invariant scalar the degree-zero component from the Clebsch-Gordan decomposition invariant under rotations. Since applying non-linear operation an SO(3)-
invariant quantity does not change its invariance, Sponlin (u) is also SO(3) -nvvarinn,,,,,,t, independent the specific form of Saealin (.). It formally holds that:
R =
(11)
Next, we apply the chain rule in Jacobian form. Considering f(R) is in the form of a column vector, to facilitate the application of the chain rule in vector form, we first transpose into a row vector f(R)T, then differentiate:
Oz ( OF (R)
Oz
Oz
OFT Oz ØhT
Oz
D'(R)
D'(R)T
(12)
OF (R)
ØFT
Here we utilize the property that D'(R)-1 = D' (R) Since the representations of neural networks are generally real numbers, the corresponding Wigner-D matrix is also real unitary.
Finally, we transpose the result back to a column vector:
(x(R) TT ( orT
Oz
dz = D((R) -'(R)))"DD(RR.-- Ør = D'(R)
v vR) Inontin (f(R)
(13)
This proves that (neatin.) (-) is SO(3)-equivariant non-linear operator: when applying its non-linearity to SO(3)- equivariant feature f, the output feature v remains SO(3)-equivariant,
E. Information of Experimental Databases
In this part, we provide detailed information about the experimental databases, including the statistical information of the six databases from the DeepH benchmark series (Li et al. 2022 Gong et al. (023) and the two databases from the QH9
benchmark series Yuu et al. 2023a). listed in Table 3 and Table respectively. Additionally, we visualize two types of challenging testing samples; samples with non-rigid deformation from thermal motions, as well as the bilayer samples with interlayer twists, which are shown in Fig. and Fig. respectively.
In Theorem and Theorem 2, the Wigner-D matrices are in the complex and real fields, respectively, since the target quantity may be complex, whereas the internal representations of neural networks are typically in the real field. Nonetheless, neural network representations in the real field can still predict complex-valued targets with SO(()--quivariance. Previous literature (Gong al. 2023 has provided mechanisms for converting the network outputs in real field into regression targets with real and imaginary parts.
17

---

TraceGrad
F. Implementation Details
The hardware environment for our experiments is a server cluster equipped with Nvidia RTX A6000 GPUs, each with 48 GiB of memory. Other experimental details may differ across the DeepH and QH9 benchmark series, which we will describe separately.
F.1. IMPLEMENTATION DETAILS ON THE DEEPH BENCHMARK SERIES
The software environment used is Pytorch 2.0.1 for experiments on the six crystalline databases from the DeepH benchmark series. When combining the proposed TraceGrad method with the DeepH-E3 architecture, the implementation of DeepH- by (Gon et al. (2023), keeping the architecture and model hyperparameters consistent
E3 is based on the projecr@roovidee with their setup. In our framework, we use the same number of encoding modules K as DeepH-E3, which is set to 3. In each encoding module, we apply the I-ontin (+) module proposed in Section 4 and 5 to each -Quu)--quivarian edge feature, enabling non-linear expressiveness. We set the number of channels for the SO(3)-invariant feature u**) (1 < k < 3) to 1024. The neural network module Snontin (.) within Irontin is implemented as three-layers fully- connected module: the input size is set to 1024, consistent with ( the hidden layer size is also 1024, with SiLU as the non-linear activation function and LayerNorm as the normalization mechanism; and the output layer size (i.e., the dimensionality of ((()) is set to be equal to the number of basic blocks for a Hamiltonian matrix, which 25 for MG and BG, 49 for M,, and 196 for BB, BT, and BS. It is worth noting that, while (-)ndan()) can be implemented as any differentiable neural network module, we here implement it as a simple fully-connected module. This decision is made to avoid adding significant computational burden to the whole network. Meanwhile, as DeepH-E3 already incorporates complex graph network mechanisms for information aggregation and message- message-passing, there is no need for Snontin (-) to be overly complex. Its role is focusing on to filling in the gaps left by the existing equivariant mechanisms in DeepH-E3: to introduce a non-linear mapping mechanism that maintains equivariance, thereby activating and unleashing the expressive power of the overall network architecture through non-linearity. The SQuiiaequivarrant decoder we adopt is the same as that of DeepH-E3; The SO(3)-invarian) decoder we adopt is four-layers fully- fully-connected module: the input size is 3 (K) times of the dimensionality of z², e.g., 75 for MG; the hidden layers have 1024 neurons with SiLU as the non-linear activation function and LayerNorm as the normalization mechanism; the size of the output layer is the number of basic blocks for an atomic pairwise Hamiltonian matrix. Since each basic block of the Hamiltonian matrix can compute a trace, the total number of trace variables corresponds to the number of basic matrix blocks. Regarding the eгror metric in the loss function Eq. (3) for the first term, we follow DeepH-F3 to use MSE (Mean Squared Error); for the second term, we choose between MSE and MAE based on performance on the validation sets, ultimately selecting MAE. in the training loss function is set according to parameter selection on the validation sets, searching from {0.1, 0.2, .... 1.0). Here, we aim to obtain a more general parameter setting for À on crystalline structures, and thus we determined À based on the overall performance on the validation sets of the six crystalline databases and the searched value is 0.3. To ensure the convergence of the TraceGrad method, we set the maximum training epochs to 5, 000. Other hyper-parameters and configurations are the same as DeepH-E3 (Gong et al. 2023): the initial learning rates for experiments on the MG, MM, BG, BB, BT, and BS databases are set to 0.003, 0.005, 0.003, 0.005, 0.004, and 0.005, respectively: the training batch size is set as 1; the optimizer is chosen as Adam; the scheduler is configured as a slippery slope scheduler.
F.2. IMPLEMENTATION DETAILS ON THE QH9 BENCHMARK SERIES
The software environment used is Pytorch 1.11.0 for experiments on the two molecular databases from the QH9 benchmark series. When combining the proposed TraceGrad method with the QHNet architecture, the implementation of QHNet is based on the project provided by Yu et al. (2023b), keeping the architecture and network contigurations consistent with their setup. In our framework, we use the same number of encoding modules as QHNet: 5 node feature encoding modules and 2 edge feature encoding modules. We opt to apply the Inontin (-) module proposed in Section and 5 to each -quiivaqun edge feature. We set the number of channels for (k) (1 < k < 2) as 1024. The neural network module Spontor (.) within Inontin is implemented as a llyees fully-connected module: the input size is set to 1024, consistent with u (R), the hidden layer size is also 1024, with SiLU as the non-linear activation function and LayerNorm as the normalization mechanism, and the output layer size (i.e., the dimensionality of z! (())) is set be equal the number of basic blocks for Hamiltonian matrix, which is 36 for QS and QD databases. The S0uu)-qquivariant decoder we adopt is the same as that of QHNet; the SO(3) decoder we adopt is a four-layers fully-connecter module: the input size is
https: wwt.o//wwwwww///ammmmco
18

---

TraceGrad
Table 3: Statistical information of the six benchmark databases, i.e., Monolayer Graphene (MG), Monolayer MoS2 (MM), Bilayer Graphene (BG), Bilayer Bismuthene (BB), Bilayer Bi2Te3 (BT). Bilayer Bi2Se3 (BS), from the DeepH benchmark series (Li et al. 2022; Gong et al., 2023). SOC: effects of Spin-Orbit Coupling. m: number of samples in the current dataset; amaz= maximum number of atoms from a unit cell in the current dataset. amin:: minimum number of atoms from unit cell in the current dataset. n/: non-twisted samples. twisted samples.
Statistic Types
MG
MM
BG
BB
BT
BS
Flements
C
Mo, S
C
Bi
Bi, Te
Bi, Se
SOC
weak
weak
weak
strong
strong
strong
m
270
300
I80
231
204
231
Training (nz)
Omaz
72
75
64
36
90
90
Omin
72
75
64
36
90
90
m
90
100
60
113
38
113
Validation (nl)
Omaz
72
75
64
90
Omin
72
75
64
36
90
m
90
100
60
113
113
Testing (nd)
Omaz
72
75
64
36
90
90
Omin
72
75
64
36
90
90
9
2
Testing (I)
Omaz
1084
244
130
190
Ommin
28
28
70
70
oPooo
Monolayer MoS2
Monolayer Graphene
Figure 3: Visualization of testing samples exhibiting non-rigid deformations due thermal motions
2 (K) times of the dimensionality of ( e.g., 72 for QS and QD; the hidden layers have 1024 neurons with SiLU as the non-linear activation function and LayerNorm as the normalization mechanism; the size the output layer is the number of basic blocks for an atomic pairwise Hamiltonian matrix. Regarding the error metric in the loss function Eq. (3). for the first term, we follow QHNet to use a combination of MSE and MAE; for the second term, we choose between MSE and MAE based on performance on the validation sets, ultimately selecting MAE. À in the training loss function is set according parameter selection on the validation sets, searching from {0.1, 0.2, ..., 1.0). Here, aim obtain more general parameter setting for À molecular structures, and thus we determined À based on the overall performance on the validation sets of the two molecular databases and the searched value is 0.2. Other byper-parameters and configurations are the same as QHNet: the maximum training steps are set as 300, 000 for QS and 260, 000 for QD, the initial learning rates for all experiments are set as x 10-4, the training batch size is set as 32, the optimizer set as Adam and a learning rate scheduler is implemented: the scheduler gradually increases the learning rate from 0 to a maximum value of 10 over the first 1, 000 warm-up steps. Subsequently, the scheduler linearly reduces the learning rate, ensuring reaches 17 by the final step.
G. Visualization of Block-level MAE Statisties
As shown in Fig. 2 each Hamiltonian matrix consists of numerous basic blocks, with each basic block representing the direct product of two degrees. Here, we follow (Yin et al.) (2244) to measure the MAE performance of deep models on each basic block, denoted M MAE block The values of M MAE Block for the two setups, i.e., DeepH-E3 and DeepH- Tr+raac on different blocks of the Hamiltonian matrix for six databases from the DeepH benchmark series are illustrated in Fig. 5 and
6 Fig. 5) presents the results for monolayer structures, while Fig. 6 focuses on bilayer structures. From these figures, it can be observed that our method, TraceGrad, brings significant accuracy improvements over the baseline method, DeepH-E3, across the vast majority of blocks of the Hamiltonian matrices, particularly on blocks where DeepH-E3 struggles with lower accuracy.
19

---

TraceGrad
Table 4: Statistical information of the two benchmark databases, QH9-stable (QS) and QH9-dynamic (QD), from the QH9 benchmark series (Yu et al. 2023a). The QS database is split using the 'ood' strategy, while the QD database is using the 'mol' strategy. m: number of samples the current dataset. 0maz= maximum number of atoms for sample in the current dataset. Omin minimum number of atoms for sample in the current dataset.
split
Statistic Types
QS
QD
Elements
C, H. O,
C,H.O.N,
104,001
79,900
Training
Omna
20
19
amin
3
10
m7
17,495
9,900
Validation
Oynax
22
19
amin
21
m
9,335
10,100
Testing
Omer
29
amin
23
10
Twisted Bilayer Bismuthene
Twisted Bilayer Bi2Se3
Twisted Bilayer Graphene
Figure 4: Visualization of testing samples with interlayer twists.
H. Ablation Study and Comparison between the Proposed Gradient-based Mechanism with the Gated Activation Mechanism
We conduct fine-grained ablation study on the six databases from DeepH benchmark series, comparing results from the four setups:
DeepH-B3 (Gong et al.
2023): the baseline model.
DeepH-E3+ Trace: this experimental setup, an ablation term, only implements half part of our method. Specifically, it extends the architecture of DeepH-E3 E3 by adding our
encoding and decoding branches and using the trace (H)) train them. As for ablation study, this setup does not include the
quantity T ur(H- Ht) tr tr gradient-based mechanism delivering non-linear expressiveness from SO(3)-invarial features to encode S0(3)-equivariant features; instead, it directly uses the SO(3)-equivariant features outputted by the SO(3)-equivariant encoders of DeepH- E3 for Hamiltonian regression. In this configuration, the SO(3)-invrrian branches only contribute indirectly during the training phase by backpropagating the supervision signals from the trace quantity to the earlier layers.
DeepH-E3+Gradd this setup is also an ablation term and implements the other half part of our method in contrast the
20

---

TraceGrad
previous ablation term. Specifically, it incorporates our SO(3)-invariant encoder branch well the gradient-induced operator deliver non-linear expressiveness from SO(3)-invariant features encode -Quii-qqivaai features. As for ablation study, this setup continues to use the single-task training pipeline of DeepH-E3, supervised only with the Hamil- tonian label without joint supervised training through the trace of Hamiltonian.
ReepH-E33TraaeGrrra this is a complete implementation of our framework extending beyond of the architecture and training pipeline of DeepH-E3, at the label level, we introduce the trace quantity to guide learning of S0(3)-invariant features; Meanwhile, at the representation level, we leverage the gradient operator yield equivvaann non-linear features for Hamiltonian prediction.
Table 5: MAE results (meV) of multiple experimental settings on the Monolayer Graphene (MG) and Monolayer MoS2 (MM) databases, + means lower values of the metrics correspond to better accuracy.
MG
MM
Methods
MAE (4)
MAE all
MAE
MAE cha
MAE
M MAE
MAE b
cho
DeepH-E3 (Baseline)
0.251
0.357
0.362
0.406
0.574
1.103
DecpH-E3TTrace
0.230
0.344
0.348
0.378
0.537
1.091
DeepH-E3+Gate
0.232
0.340
0.351
0.364
0.527
1.086
DeepH-E3+Gras
0.185
0.269
0.258
0.308
0.453
0.924
CeepH-EE3TraceGGaa
0.203
0.295
0.305
0.346
0.491
0.912
Traacerrrd
0.175
0.257
0.228
0.285
0.412
0.808
Furthermore, to compare our proposed gradient-basse based mechanism with the gated activation mechanism (Weiler et al.
2018) which is widely used in SOu3)-eqiivarian neural networks (Gong et al. 2023 Liao Smidi 2023). the following experimental setups are also included:
DeepH-E3+Gate: This variant modifies the experimental setup from DeepH- by replacing the gradient mech- anism, which constructs equivariant features features f. All other aspects remain the same.
of 0s with the gated activation mechanism, constructing equivariant
DeepH-F33+ CeopH---E-+TrrCeGaaee TraceGate: This variant modifies the experimental setup from DeepH-E3+ TraceeGra by replacing the gradi- ent mechanism (v
0z with the gated activation mechanism (v z -f). All other aspects remain the same.
Experimental results for all of the six setups are listed in Table and presenss the results for monolayer structures, while Table 6 focuses on bilayer structures. We have the following observations:
First, by comparing among the results of DeepH-EB, DeepH--E3Trraee, DeepHl-E3+Grad, and Deeptt-E3+TraceGrad
0.25
0.5
0.278
0
0.6
8.25
16
21
33
36
06
32
36
Block
bl6ck
Figure 5: Visualization of M A. Mock on each basic block of the Hamiltonian matrices for the Monolayer Graphene (MG) and Monolayer MoS2 (M M) databases.
21

---

TraceGrad
BC
0.544
Deepi--53
1,042
Deepl- E3
Despri ET-Tacarre
0.5
0LD
0:50
04
233
D.4
6.1
0.0
0.0
Blook
Black
1.4
Deepii-ES
1.-
1.0
0.6-
0.044
0.
0.7
0.6
30
22
33
36
43
46
35
23
20
35
42
46
Block
BCF
B6F
0.6
0.615
Desptte2
2,399
Deeihh-22
Dosgi 821 TraceGrod
DoGSH 63 TracoGrod
0.5
2.0
0.
1942 (103)
1.5
1,24
3.0
0.
0.5
0.0
D'1"
BS'
2.20
DeeptisE3 Deepli El-TaceGred
Dooph E3
0.5-
0.2-
0.0
32
46
B1ck
Figure 6: Visualization of MAE on each basic block of the Hamiltonian matrices for the non-twisted (marked with superscripts nl) and twisted (marked with superscripts () testing subsets of the Bilayer Graphene (BG), Bilayer Bismuthene (BB), Bilayer Bi2Te3 (BT). and Bilayer Bi2Se3 (BS) databases.
22

---

TraceGrad
Table 6: MAE results (meV) of multiple experimental settings on the Bilayer Graphene (BG), Bilayer Bismuthene (BB), Bilayer Bi2Te3 (BT),, and Bilayer Bi2Se3 (BS) databases. The superscripts nt and respectively denote the non-twisted and twisted subsets. + means lower values of the metrics correspond better accuracy.
BG
BG
Methods
MAE (4)
MAE
MAE cha_s
M MAEH
MAEE nll
MAE
MAEH
cho_5
cho_s
chabb
DeepH-E3 (Baseline)
0.389
0.453
0.644
0.264
0.429
0.609
DeepH-E3+Trase
0.362
0.417
0.593
0.251
0.401
0.480
DeepH-E3+Gaate
0,368
0,441
0.626
0.241
0,405
0,602
DeepH-E3+Grn
0.320
0.356
0.511
0.222
0.389
0.446
DeepH-EE+Trraceee
0.354
0.403
0.580
0.239
0.388
0.464
DeepH-E3+ Traccerrd
0.291
0.323
0.430
0.198
0.372
0.406
m
BB
Methods
MAEH
M MAEH
ch
MAEH
ch cho.b
M MAEH all
M MAE....
M MAEH
cha.b
DeepH-E3 (Baseline)
0.274
0.304
1.042
0.468
0.602
2.399
DeepH-E3+Trace
0.259
0.285
0.928
0.429
0.570
1.782
Geap
0.268
0.301
0.991
0.450
0.593
2.276
DeepH-E3+Grad
0.243
0.272
0.824
0.406
0.542
1.431
CeepH-E3+Traceee
0.252
0.279
0.908
0.417
0.561
1.740
Deepll-E3+TraceGrn
0.226
0.256
0.740
0.384
0.503
1.284
BT
BT*
Methods
MAEH all
MAEH
cha.s
MAEH
cho.b
MAEH
M MAEH
cha.
M MAEH
DeepH-E3 (Baseline)
0.447
0.480
1.387
0.831
0.850
4.572
DeepHH-E3+Trace
0.406
0.462
1.239
0.784
0.812
4.520
Geap
0.423
0.469
1.196
0.813
0.829
4.494
DeepHH-E3+Grad
0.342
0.365
0.750
0.742
0.786
4.463
DeppH-E3+Traceain
0.368
0.377
0.982
0.778
0.783
4,489
SrepH-E3+TraeCCii
0.295
0.312
0.718
0.735
0.755
4.418
BSSt
BSt
Methods
M MAEal
MAEH
MAER
MAEau
MAEH
MAEH
cha.s
cho
cho.s
cha.b
DeepH-E3 (Baseline)
0.397
0.424
0.867
0.370
0.390
0.875
DeepH-E3+Trace
0.382
0.397
0.843
0.351
0.367
0.838
DeepH-E33+Gate
0.385
0,401
0.852
0.369
0.373
0,820
DeepHH-E3+Grad
0.343
0.365
0.696
0.324
0.339
0.746
DeepH-E33 BeepH-E3TTraee TraaeeceGte
0.364
0.386
0.765
0.337
0.348
0.798
DeepH-E33 E3+TraceGrad
0.300
0.332
0.644
0.291
0.302
0.674
can conclude that the two core mechanisms of our method, i.e., the SO(3)-invariant trace supervision mechanism (Trace) the label level as well the gradient-based induction mechanism (Grad) at the representation level, can contribute the performance individually. Moreover, their combination provides even better performance. This is because, on one induction mechanism bridge, the non-linear expressiveness of SO(3)-invariant features
hand, with the gradient-basee learned from the trace label can be transformed into the SO())-qqivvariat representations during inference; the other hand, with trace label, the SO(3)-invariant network branch has a strong supervisory signal, enabling it to learn the intrinsic symmetry and complexity the regression targets, enhancing the quality of -nvaaiin features and ultimately ben- efits the encoding of -0euivaauaaarian features. The value of such complementarity has been fully demonstrated in the experimental results.
23

---

TraceGrad
Second, it is evident that both Geep and DetTac-E3 show improvement over the baseline DeepH- E3, these configurations introduce additional neural network module Saonlin() (+) in learning the feature 2, which increases model capacity. Moreover, Drepca-E3+TracGatt also incorporates our proposed trace supervision signal, which guides the learning of SO(3)-invariant features. However, their performance falls short of DeepH-E3+Grad and DeepH-E3+Traceerrrd, respectively, indicating that the gated activation mechanism may not fully capture the system's 's non-linearity patterns, whereas the proposed gradient-baser mechanism demonstrates stronger generalization performance on Hamiltonian prediction and may be better choice in terms of expressive capability.
I. Theoretical Analysis on Computational Complexity
The total number of non-zero Hamiltonian matrix elements to be calculated is proportional to the number of local atomic pairs in the system, with a complexity of D(NE), where is the total number of atoms and E is the average number of neighboring atoms within the cutoff radius per atom. Since the atomic orbital basis set has a finite range, the non-zero matrix elements vanish beyond a certain distance for given atom. In small systems, where all atoms lie within each other's cutoff radius, E scales with N, resulting in total number of non-zero elements proportional to N However, in sufficiently large systems, the finite range of the atomic orbitals ensures that E remains constant, independent of N. As a result, for large atomic systems, the total number of non-zero elements simplifies to ((N).
The baseline models we select, whether DeepH-E3 or QHNet, are S0(3)-eqiivarian graph neural network models with efficient information aggregation and pressage-passinge mechanisms. These models cleverly balance the locality of Hamil- tonian definitions with the long-range interactions present in the system. As a result, the computational complexity asymp- totically scales (N) as N increases, which is consistent with the growth of the scales of non-zero Hamiltonian matrix elements. The proposed TraceGrad method directly updates each SO(3) -equi--quivarian feature of the baseline models, and the computational amount is proportinal to the number of features of the baseline models. Therefore, combining TraceGrad, the computational complexity also scales as (NN)
Traditional DFT methods require iterations of diagonalizing N x N matrices, each with a time complexity of O(N*), because all occupied states are needed to compute the charge density. As N increases, this cubic complexity leads to significant computational overhead, making it challenging to simulate large atomic systems within a reasonable time frame. In contrast, our deep learning framework enables the efficient and accurate prediction of the Hamiltonian for large atomic systems with a linear time complexity of O(N), eliminating the need for self-consistent iterations. Moreover, since most physical properties, such as transport, optical, and topological properties, depend only on the energy bands near the Fermi level. it is unnecessary to solve for the eigenfunctions of all occupied states once the Hamiltonian is known. Since the Hamiltonian matrix is sparse and only a limited number of bands near the Fermi level are needed, these eigenstates can be efficiently computed using methods like the shift-invert approach available in the ARPACK package (Lehoueq et al., 1998), with a computational complexity of O(N).
A Joint Discussion on GPU Time Costs and Performance Gains
We here provide a joint comparison of the GPU time cost and corresponding accuracy of different models across four databases as representative: Monolayer Graphene (MG)), Monolayer MoS2 (M M). QH9-stable (QS), and QH9-dynamic D). This comparison includes the average inference time per sample for each model, with the test batch size set I and hardware environment set as Nvidia RTX A6000 GPU in single-task mode without computational sharing with other processes.
For MG and M we compare among DeepH-E3, DeepHH-E3+TraeeGree TraceGGrdd and DeepH-E3 x2, where DeepH-E33*2 refers to a model obtained by doubling the number of encoding blocks in DeepH-E3 and training it from scratch until conver- gence. For QS and QD, we compare among QHNet, @HNet+TraceGrad, and QHNet*2 where QHNet*2 refers to a model obtained by doubling the number of encoding blocks in QHNet and training it from scratch until convergence. The experimental results are documented in the Table 7
From this Table, we find that adding the TraceGrad module results in only a slight increase in inference time compared to the baseline models. Given the substantial accuracy improvements introduced by the TraceGrad method, we con- sider this minor increase in computational time acceptable for practical applications. In contrast, simply increasing the depth of DecpH-E3 or QHNet results in a significant rise in inference time but yields only limited accuracy improve- ments. In contrast, E3raaeEa+T demonstrates significantly better accuracy performance than DeepH-E3x2 achieves notably higher accuracy than Q2 Furthermore, the inference time of
and similarly,
24

---

TraceGrad
Table 7: Average inference time (denoted as Time, in seconds) and MAE performance (M AMAE in meV) for testing samples from the Monolayer Graphene (MG), Monolayer MoS2 (MM), Qt9-stable (QS), and QH9-dynamic (QD) databases. & indicates that lower values of the metrics correspond to better accuracy. All models are tested individually on a single Nvidia RTX A6000 single-task mode.
MG
MM
Methods
Time (1)
M MAE (1)
Time (4)
MA (1)
DeepH-E3 (Baseline)
0.247
0.251
0.256
0.406
DeepH-E3x2
0.483
0.244
0.510
0.387
DeepHl-E3+Traceraa
0.264
0.175
0.274
0.285
QS
QD
Methods
Time (4)
MAE (1)
Time (4)
M (I)
QHNet (Baseline)
0.233
1.962
0.174
4.733
QHNet*2
0.497
1.845
0.385
4.532
Qraet-TraceGrad
0.248
1.191
0.187
2.819
DeepHl-E3+Traceraad and CHNet+TraceGrad are both lower than their respective DeepH-E3*2 and QHNe* counter- parts. These results underscore the superiority of the TraceGrad method in enhancing expressive capability and improving accuracy performance, while maintaining time efficiency.
K. Acceleration Performance for the Convergence of Traditional DFT Algorithms
Despite the increasing ability of deep learning models independently handle more electronic-strucctre computation tasks, there still applications with extremely high numerical precision requirements and very low tolerance for error, where traditional DFT algorithms must perform the final calculations. In such cases, the predictions from deep models can be used as initial matrices to accelerate the convergence of traditional DFT algorithms. We evaluate the acceleration performance brought by the proposed method for the convergence classical DFT algorithms implemented by PySCF (Sun et al. 2018). Specifically, we adopt the two groups of metrics on acceleration performance:
The first group of metrics are defined by Yu et al. 2023a), follows:
Achieved ratio. This metric calculates the number DFT optimization steps taken when initializing with the Hamil- tonian matrices predicted by the deep model compared using initial guess methods like minao and 1e.
Error-level ratio. This metric measures the number of DFT optimization steps required, starting from random initial- ization, reach the same error level as the deep model's predictions, relative to the total number of steps in the DFT process.
Experimental results on these metrics are recorded in Table 8 where the results for the compared method QHNet are taken from Yu et al. (2023)). while the results of QHNet+TraceGrad, come from our experiments. In our experiments, the DFT calculation settings follow those in Section 4 of Yu et al. (2023a) (the DFT parameters and the 50 testing samples), except for the CPU environment, where we use single thread of an Intel(R) Xeon(R) Gold 6330 2.00GHz CPU single-task mode. worth noting that this does not affect the fairness of the comparison, as the achieved ratio and error-level ratio measure the ratio of iteration counts rather than runtime, and differences CPU computation times are negligible in these metrics. From Table X could observe that the proposed TraceGrad method brings significant improvements to the baseline model QHNet on the acceleration ratio of DFT calculation, notably reducing the achieved ratio and enhancing the levor-level ratio.
The second group of metrics measures the wall time savings that deep learning methods contribute to DFT calculations, specifically quantifying the incremental time savings brought by our proposed TraceGrad method in accelerating DFT caleulations. For fair comparison, we report the average wall time costs of testing samples (/s) across three metrics:
11: The wall time required for a DFT calculation initialized with a random guess initialization, such as le or mi nao.
25

---

TraceGrad
Table 8: The acceleration ratios of QHNet and TraeaaTra for DFT calculation. Both models are evaluated on a set of 50 molecules chosen by Yu et al. 2023a), with the mean and standard deviation of the metrics across these samples reported. means lower values correspond to better accuracy, while means higher values correspond to better performance.
Training databases
DFT initialization
Methods
Metric
Ratio
Achieved ratio +
0.400 + 0.030
le
Error-level ratio f 0.620 = 0.037
QS
Achieved ratio +
0.715 = 0.033
minao
Error-level ratio f
0.406 + 0.021
QHNet (Baseline)
Achieved ratio 4 0.512 + 0.138
le
Error-level ratio f
0.622 + 0.048
QD
Achieved ratio
0.882 = 0.217
minao
Error-level ratio f
0.406 + 0.066
Achieved ratio 4
0.345 + 0.038
le
Error-level ratio f
0.685 = 0.037
QS
Achieved ratio +
0.647 + 0.061
minao
Error-level ratio 1 0.466 = 0.035
QHNet+TraceGrad
Achieved ratio 1.
0,440 = 0.101
le
Error-level ratio f
0.645 + 0.046
QD
Achieved ratio +
0.761 + 0.167
minao
Error-level ratio f
0.435 + 0.052
t2: The wall time for inference using deep learning methods (i.e., QHNet or (HNet+TraceGrad),
13: The total wall time for the combined process, including both deep learning inference and the subsequent DFT calculation initialized with the deep model's outputs. t3 here provides more comprehensive evaluation of the actual time savings achieved when incorporating deep learning methods.
Experimental results on these metrics are recorded Table 9 Here all time-related measurements are conducted on single thread of Intel(R) Xeon(R) Gold 6330 @ 2.00GHz CPU, including the experiments for QHNet, which are reproduced under the same conditions ensure fairness. Unlike the time measurements in Appendix J, which performed GPUs, all deep learning models here are evaluated the CPU thread to maintain consistency in the comparison with DFT software. From the experimental results, we observe three key findings:
Comparing t2 and f1, deep models are significantly faster than DFT calculations, achieving speeds tens of times greater for the testing samples. worth noting that, given that the testing samples here are all small molecular systems, deep models have already demonstrated sigmificant time efficiency advantage compared DFT. Based the computational complexity analysis of deep learning methods compared to DFT methods in Appendix is reasonable to infer that for larger atomic systems, the disparity between /2 and (1 will expand rapidly.
Comparing t2 values between QHNet and THaeerTraeeGrad, the additional runtime introduced by TraceGrad is relatively minor on the CPU, consistent with the GPU-based results reported in Appendix J.
Comparing 13 and (1, the total runtime of using deep model to predict initial values and then running DFT cal- culations is significantly lower than performing DFT calculations from a random guess initialization. Particularly, comparing row and row 16, row 7 and row 19, row 10 and row 22, as well as row 13 and row 25 Table 9 can be observed that combing TraceGrad further reduces t3, demonstrating that the time saved by TraceGrad in DFT calculations far exceeds the minimal additional time required for its inference.
26

---

TraceGrad
Table 9: Average wall time costs per sample (/s) for three experimental settings: #1 represents the wall time required for DFT calculation with random initial guess like le of minao; t2 denotes the wall time for inference of deep learning methods (i.e., QHNet THNet+TraceGraad
is the total time of the process including deep learning inference and the DFT calculation that uses the outputs of deep learning methods initialization. Both models are evaluated molecules chosen by (Yu al., 2023a), with the mean and standard deviation of the metrics across these samples reported. indicates that lower values correspond to better performance. In order to ensure fairness in comparison, all experimental settings including DFT calculation and deep learning inference are measured on single thread of an Intel(R) Xcon(R) Gold 6330 2.00GHz CPU single-task mode.
set of 50
Training databases
DFT initialization
Methods
Metric
Time
tl
120.896 + 9.134
le
12
1.724 + 0.025
t3
48.604 + 7.741
QS
tl
63.193 + 5.335
minao
t2 4
1.724 + 0.025
t3
48.604 + 7.741
QHNet (Baseline)
tl
87.161 = 12.075
le
12
1.280 0.019
t3
44.146 + 9.342
QD
tl
51.396 5.870
minao
12 1
1.280 0.019
t3
44.146 + 9.342
t1
120.896 + 9.134
le
12
1.852 + 0.020
t3
41.941 + 6.783
QS
tl
63.193 + 5.335
minao
12 4
1.852 + 0.020
t3
41.941 + 6.783
Traerraa
tl
87.161 = 12.075
le
12
1.361 0.010
13
39.712 + 9.076
QD
tl
51.396 + 5.870
minao
12 4
1.361 0.010
t3 4
39.712 + 9.076
27

---

TraceGrad
MM
0.40
HarmoSE
0.406
HarmoSE + TraceGrad
0.35
0.296
MAE (meV)
0.25
0.20
0.15
0.00
11 16
21
26 31 36 41 46
Figure 7: Comparison the M A.F.
block metric for HarmoSE and
the MM M database.
Table 10: MAE results (meV) for DeepH-2, HarmoSE, and values of the metrics correspond to better accuracy. The results of the compared methods are taken from the corresponding literature (Wang et al. 2024b Yin et al. 2024), where the empty items are due to the data not being provided in the original paper.
the MM database. + means lower
MM
Methods
MAE (L)
MAEH
MAEH
MAEH
cha,s
cha.b
DeepH-2 (Wang et al. 2024b)
0.21
HarmoSE (Yin et al. 2024)
0.233
0.293
0.406
HarmoSSE+Tracerraa
0.178
0.228
0.296
L. Empirical Study on Combining Our Method with Approximately S0(3)-equivariance Framework
While non-strict -quivvaauu,, which may limit the depth of theoretical exploration, is not the main focus of this study aiming at bridging rigorous SO(()-eqqivariance with the non-linear expressive capabilities of neural networks, con- sidering that they remain of interest in a few numerical computation applications where precision is highlighted over strict equivariance, we also conduct empirical study combining our method with approximately S0(3)-equivariant techniques. Taking the Monolayer MoS2 (MMM) database as case study, we evaluate the performance combining our trace super- vision and gradient induction method (TraceGrad) with the an approximately equivariant approach HarmoSE (Yin al.] 2024). We here take HarmoSE the backbone encoder, and yields features by TraceGrad to enrich its representations. The experimental results in Table TO and Fig. 7 demonstrate that TraceGrrad significantly enhances the accuracy of HarmoSE, surpassing DeepH-2 (Wang et al. 2024b) and achieving SOTA results. Both DeepH-2 and HarmoSE sacrificed strict So(3)--quivariance to fully release the expressive capabilities of graph Transformers. aiming for the ultimate in prediction accuracy. Despite this, our method still manages to significantly exceed their accuracy, further confirming the superiority of our method in learning expressive representations of physical systems for Hamiltonian prediction.
M. Extending Our Method to Energy and Force Prediction Task
In this section, we extend the application of our proposed framework, TraceGrad, beyond electronic-structte Hamilto- nian prediction to address energy and force prediction task. This task crucial for variety of atomic system modeling
28

---

TraceGrad
Table 11: MAE results of Equiformer and Tquiformer+TraceGraad methods on energy and force prediction, evaluated on the MD17-Aspirin and MDI7-Malonaldehyde datasets. lmax corresponds to the maximum degree of features used. MAE of energy (F) and force (F) are in units of meV and meV/ respectively.
Method
Aspirin
Malonaldehyde
E
F
E
F
Equiformer Omax 2)
5.3
7.2
3.3
5.8
Aggufforme+-TracGiiaae (max 2)
5.06
5.65
3.21
4.68
applications across chemistry, materials science, and biology.
In predictions of energy and force, the regression target E (energy) is an SO(3)-invariant quantity (1 = 0), while F (force) an SO(3)-equivarian quantity = 1). The regression model typically learns S0(3)-quivariant features f, which then transformed into SO(3)-invariant features, from which E is regressed. Subsequently, the force field a given position DE where is the position vector of the i-th atom. This approach ensures energy conservation. Given the specificity of this task, we integrate the baseline model, Equiformer (Liao & Smidt (2023), with our proposed TraceGrad method, where the regression target shifts from dlectronic-structun Hamiltonian prediction to energy and force prediction, as outlined below:
is obtained by differentiating E with respect to the atomic coordinates: F; őr.
First, we use the -quuivaauivariiat features f encoded by the baseline model Equiformer as input, and construct SO(3)- invariant non-linear features according to our method (Sections 4 and of our paper). We then use the trace quantity T to supervise the learning of Given F as column vector with induce the SO(3)-equivariant features v with more non-linearity, which are then fed back into the baseline model for the subsequent encoding and decoding phases, where E' is regressed and finally F is constructed from the gradients of E.
here T simplifies T FT F. From 2,
We conduct experiments on two datasets as representative, i.e., MD17-Aspirin and MD17- MiDo-Malonaldehyde (Chmiela et al. 2017 Schütt et al. 2017, Chmiela et al. 2018). We use the same training, validation and testing sets as Liao Smidt (2023). We train Equiformer+
under the same training conditions as their open source codes []] The optimizer
used is AdamW, and a cosine learning rate scheduler with linear warmup is employed, where the warmup epochs are set as 10. The maximum learning rate is set to 5 10 4 and the batch size is 8. The total number of training epochs is 1,500. The weight decay parameter is set to 1 x 10-6 The weight for the energy loss is set to I, while the weight for the force loss is set to 80 for experiments on the Aspirin dataset and 100 for Malonaldehyde. The experimental results are listed in Table II where the results for Equiformer are taken from the original paper.
The experimental results show that our TraceGrad method improves the prediction accuracy for both energy and force. Particularly, the accuracy improvement for force, i.e., the SO(3)-equivariant regression target, is especially significant, as the new supervision signal (T F F) introduced is designed targeting force. Nevertheless, it is very interesting that even for the energy, i.e., the SO(3)-invariant quantity, TraceGrad also leads to accuracy improvement. This is because both energy and force predictions are based on the SO(3)) enhances the non-linear expressiveness of these SO33)-equivariant features, thus improving the accuracy of the induced physical quantities. These results show that our method can be effectively transferred to energy/force prediction, and its effectiveness is not limited to the prediction electronic-structur Hamiltonians and their downstream physical quantities.
features of the network, and our approach
former
29


Reviews:
[{"review_id": "PcwGmTWIdj", "reviewer": "Reviewer_1Tq8", "summary": "This paper proposes to enhance the Hamiltonian prediction networks with additional invariant supervision and an additional gradient branch. The authors observe that the trace of $H H^T$ is rotation invariant and can be used to supervise the learning of zero order features. Additionally, the gradient of a network using the zero order feature as input w.r.t. the zero and non zero order features is rotation equivariant, so it can be added back to the original equivariant feature. The authors combine these two techniques and test with DeepH-E3 and QHNet models and test on material datasets and small molecule datasets.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "1.**Clarification regarding whether the gradient mechanism can change the direction of features**: In Theorem 2 of our paper, for simplicity  and to highlight the core ideas, we selected $\\textbf{f}$ as a basic feature component of degree $l$. In this case, applying the gradient operation directly to $\\textbf{f}$ results in $\\textbf{v}$ that only changes its magnitude but not its direction. However, in the actual application of this theory, a series of feature components $\\textbf{f}_1, \\textbf{f}_2, \\cdots, \\textbf{f}_C$ with degree $l$ are applied with the gradient operation to obtain corresponding $\\textbf{v}_1, \\textbf{v}_2, \\cdots, \\textbf{v}_C$, a new series of feature components. \n\nThese components are then combined to form a new feature, $\\textbf{v} = \\sum_{1 \\leq c \\leq C} w_c \\textbf{v}_c$, where $w_c$ represents the combinational coefficients. In this case, both the direction and magnitude of the new feature $\\textbf{v}$ are different from those of the original features $\\textbf{f}_1$, $\\textbf{f}_2$, ..., $\\textbf{f}_C$. In fact, our method involves a more in-depth extension (please refer to lines 178-249 in section 5.1 of our paper for details). For different feature components, e.g., $\\textbf{f}^{(k)l_i}$ and $\\textbf{f}^{(k)l_j}$ in Eq. (2) of our paper, as long as they share the same degree $(l_i=l_j)$, we can construct SO(3)-invariant and SO(3)-equivariant features using them, even if their magnitudes and directions are different. Through such operations and  linear combinations, we can ultimately encode new features with flexible direction and magnitude, significantly enhancing the expressive power of the features. **Therefore, while the gradient operation on a single feature component itself does not change its direction, by combining multiple feature components and leveraging gradients and linear combinations, we can encode richer, more expressive feature representations with varying directions and magnitudes.**\n\nWe would like to thank the reviewer for their valuable feedback. We will clarify this analysis in the theoretical section of our paper to ensure a more accurate and comprehensive understanding of the gradient-based mechanism we proposed.\n\n2.We will move the ablation studies into the main text of the revised version of our paper. \n\n3.We have added ablation studies on the QH9-Stable (QS) database, where the notation for each experimental setup is consistent with the definitions provided in Appendix H. The results are summarized in the table below. As can be observed,  the key conclusion on this database is consistent with that in the DeepH-E3 benchmark series: each individual component (Trace and Grad) of our method contributes positively to the overall performance, and their combination leads to further improvements. Notably, the gradient-based mechanism (Grad) consistently achieves higher accuracy compared to the gated mechanism (Gate).\n\n**Table**: Experimental results measured by the $MAE^H_{all}$, $MAE^H_{diag}$, $MAE^H_{nondiag}$, $MAE^{\\epsilon}$, and $Sim(\\psi)$ metrics on the QH9-stable (QS)  database using 'ood'  split strategy. $\\downarrow$ means lower values correspond to better accuracy, while $\\uparrow$ means higher values correspond to better performance. The units of MAE metrics are meV, while $Sim(\\psi)$ is the cosine similarity which is dimensionless.\n\n| **Method**           | $MAE^H_{all}$ ↓ | $MAE^H_{diag}$ ↓ | $MAE^H_{nondiag}$ ↓ | $MAE^{\\epsilon}$ ↓ | $Sim(\\psi)$ ↑ |\n|----------------------|--------------------|----------------------|----------------------------|----------------|------------|\n| QHNet (Baseline)     | 1.962              | 3.040                | 1.902                      | 17.528         | 0.937      |\n| QHNet+Trace          | 1.874              | 2.936                | 1.815                      | 16.724         | 0.940      |\n| QHNet+Gate           | 1.796              | 2.845                | 1.741                      | 15.696         | 0.940      |\n| QHNet+Grad           | 1.604              | 2.587                | 1.558                      | 11.393         | 0.942      |\n| QHNet+TraceGate      | 1.516              | 2.426                | 1.465                      | 10.568         | 0.945      |\n| QHNet+TraceGrad      | **1.191**          | **2.125**            | **1.139**                  | **8.579**      | **0.948**  |", "claims_and_evidence": "- Using the trace of $H H^T$ to supervise zero order feature learning seems to be a reasonable design.\n\n- However, I have a doubt regarding the gradient part. Wouldn't taking the local gradient w.r.t. to the equivariant feature result in a linear scaling of the original equivariant feature? i.e., it won't change the direction of the equivariant feature. For example, for order 1 feature $v\\in \\mathbb{R}^3$, the zero-order feature from the CG decomposition of $v\\bigotimes v$ would be proportional to the dot product $v\\cdot v$. As a result, if we take the gradient w.r.t. $v$, the quadratic term in dot product will reduce to be linear w.r.t. $v$. This seems to hold even after applying the non-linear neural networks due to the chain rule. So the final equivariant feature from the gradient branch would be a scaling of the original equivariant feature.", "methods_and_evaluation_criteria": "- The benchmarked dataset includes both periodic and non-periodic systems, which is a strength.", "theoretical_claims": "- Most equations are descriptive rather than proving something.", "experimental_designs_or_analyses": "- The experimental design seems to be sound.", "relation_to_broader_scientific_literature": "- The related works are adequately discussed.", "essential_references_not_discussed": "- Not I am aware of.", "other_strengths_and_weaknesses": "- The experimental results are promising.\n\n- The theoretical analysis of the proposed gradient branch might be enhanced.", "other_comments_or_suggestions": "- I did not notice typo.\n\n- Some ablation studies like (+Trace, +Gate, +Grad)  might be presented in the main text.", "questions_for_authors": "- I noticed the ablation studies in the appendix (Table 5) about DeepH-E3 (+Trace, +Gate, +Grad), which I think are important results. Did the authors observe similar trends for QH9 datasets?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "bIHAMy05XC", "reviewer": "Reviewer_gNai", "summary": "This paper introduces TraceGrad, a framework that integrates strong non-linear expressiveness with strict SO(3)-equivariance for electronic structure Hamiltonian prediction. The approach first constructs theoretical SO(3)-invariant trace quantities derived from Hamiltonian targets, using them as supervisory signals to learn invariant features. A gradient-based mechanism is then employed to generate SO(3)-equivariant encodings of varying degrees from these learned invariant features. Empirical evaluations on eight benchmark datasets demonstrate improvements in predicting physical quantities and accelerating density functional theory (DFT) computations.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "1.**Response to the reviewer's question about the computational burden**: \nFirst, the branch decoding the trace quantity $\\mathbf{T}$ from the SO(3)-invariant features $z$ is only required during the training phase and does not need to be activated during inference; thus, the parameters associated with this branch are unnecessary at inference time. Second, the key difference between our proposed gradient-based mechanism for constructing non-linear SO(3)-equivariant features and the original gated mechanism is that the feature construction has changed from $\\mathbf{v} = z \\cdot \\mathbf{f}$ to $\\mathbf{v} = \\frac{\\partial z}{\\partial \\mathbf{f}}$. Given the network that learns the SO(3)-invariant features $z$, the gradient operation itself does not introduce additional parameters.\n\nThe reviewer's suggestion to analyze the computational cost is reasonable. Unfortunately, existing automated FLOPs counting tools like fvcore, torchprofile, and torchstat do not support precise quantification of the computational complexity for frameworks based on equivariant neural network packages such as E3NN. Therefore, in our paper, we measure the computational efficiency of our method by testing the average inference time under the same GPU/CPU hardware conditions.\n\nIn the paper’s Appendix J,  we provide a comprehensive comparison of GPU computational costs and corresponding accuracy for different models across four representative databases. From the experimental results in Table 7 of our paper, **we find that incorporating our TraceGrad method results in only a slight increase in inference time compared to the baseline models, i.e., DeepH-E3 or QHNet**. Given the substantial accuracy improvements introduced by the TraceGrad method, this minor increase in computational time is considered acceptable for practical applications. **In contrast, simply doubling the depth of DeepH-E3 or QHNet leads to a significant rise in inference time while providing only limited accuracy improvements.**\nIn contrast, DeepH-E3+TraceGrad and QHNet+TraceGrad exhibit significantly better accuracy performance compared to DeepH-E3$^{\\times 2}$ and QHNet$^{\\times 2}$, respectively. At the same time, the inference times of DeepH-E3+TraceGrad and QHNet+TraceGrad are considerably lower than those of DeepH-E3$^{\\times 2}$ and QHNet$^{\\times 2}$, respectively. **These findings highlight the superiority of the TraceGrad method in enhancing model expressiveness and improving accuracy performance while maintaining computational efficiency.**\n\nTo further address the reviewer's concern regarding the efficiency differences between the classical gated mechanism and our TraceGrad method, we introduce an additional experimental setup to evaluate inference efficiency. Specifically, we test QHNet+Gate, where the gradient-based mechanism for constructing the non-linear equivariant feature $ \\mathbf{v} $ is replaced with a classical gated mechanism, following the definition provided in Appendix H of the paper.\nExperimental results show that **QHNet+Gate** achieves inference times of 0.243s and 0.184s on the QS and QD datasets, with $ \\text{MAE}^H_{\\text{all}} $ values of 1.796 meV and 4.217 meV, respectively. In comparison, **QHNet+TraceGrad** takes 0.248s and 0.187s on the same datasets, with $ \\text{MAE}^H_{\\text{all}} $ values of 1.191 meV and 2.819 meV, respectively. These results demonstrate that the TraceGrad method introduces only a minor increase in inference time compared to the traditional gated mechanism, while achieving significant improvements in accuracy.\n\nIn addition to reporting the GPU inference time, **we also present the inference times of QHNet and QHNet+TraceGrad on a single CPU thread in the paper’s Appendix K**. Experimental results show that while combining TraceGrad introduces only a slight increase from the inference time of QHNet on the CPU, it delivers significant improvements in accelerating the convergence of DFT methods. **Notably, the time saved by TraceGrad for DFT calculations far exceeds the minimal additional time introduced by TraceGrad for the deep model's inference.**\n\n2.We agree with the reviewer’s view that TraceGrad’s contribution to equivariant graph neural networks is indeed general. In addition, **we have found that it demonstrates its effectiveness on another task, namely energy and force field prediction**. For more details, please refer to the 3rd item of our response to Reviewer AGnk .\n\n3.We have reduced the hidden size of $g_{\\text{nonlin}}(\\cdot)$ by half (from 1024 to 512) and conducted experiments on the QHNet-Stable (QS) database. The experimental results show that the $MAE^H_{\\text{all}}$ metric increase from 1.191 to 1.347, but still significantly outperforms the baseline method (1.962). This suggests that while the parameter size does impact the accuracy improvements introduced by our method, the effect is relatively limited.", "claims_and_evidence": "The paper makes two key claims:\n* Addressing the Challenge of Combining Non-Linear Expressiveness with SO(3)-Equivariance\n    * The authors propose a novel approach to this challenge by systematically bridging SO(3)-invariant and SO(3)-equivariant representations.\n    * The framework first supervises SO(3)-invariant features to ensure strong non-linear expressiveness and subsequently derives SO(3)-equivariant representations through a gradient-based mechanism.\n    * The mathematical derivations in Section 4 provide a solid theoretical foundation, and empirical results in Section 6 and Appendix H support the claim.\n* Significant Performance Gains in Hamiltonian Prediction on Eight Benchmark Datasets\n    * The authors report that TraceGrad outperforms state-of-the-art methods across eight datasets from the DeepH and QH9 benchmark series.\n    * This claim is supported by empirical results in Section 6 and Appendix H, which demonstrate improved prediction accuracy and DFT acceleration.", "methods_and_evaluation_criteria": "* Methods: The approach is well-grounded in equivariant neural networks for Hamiltonian prediction.\n* Evaluation Criteria: The chosen metrics align with standard benchmarks in quantum chemistry.", "theoretical_claims": "* The construction of SO(3)-invariant trace quantities and the gradient-based mechanism linking invariant and equivariant representations are mathematically sound.\n* The authors assert in Introduction and Remark 4.3 that the gradient mechanism induces expressive SO(3)-equivariant representations while maintaining physical interpretability, providing an advantage over gated mechanisms. Empirical ablation results (Appendix H) partially support this claim, but it remains unclear whether model variants using gating mechanisms have comparable parameter counts to those using TraceGrad. A direct comparison in terms of model complexity would strengthen this argument.", "experimental_designs_or_analyses": "* The experimental setup is well-structured and comprehensive.\n* However, there is no detailed analysis of parameter counts, which is crucial for assessing model expressiveness. The authors state that  $g_{\\text{nonlin}}(\\cdot)$  is implemented as a three-layer fully connected module with a large hidden size, injecting significantly more parameters into the architecture. Since larger models tend to have higher expressiveness, a fair comparison requires increasing the parameter counts of baseline models (e.g., by widening hidden layers or deepening architectures). Providing such comparisons would clarify TraceGrad’s advantages over other methods.", "relation_to_broader_scientific_literature": "TraceGrad’s contribution to equivariant graph neural networks is general and could inspire new model designs for other molecular property prediction tasks beyond Hamiltonian prediction.", "essential_references_not_discussed": "No essential references appear to be missing from the discussion, as the submission adequately contextualizes its contributions.", "other_strengths_and_weaknesses": "Strengths\n* The paper is well-organized and clearly written.\n* The proposed method is theoretically sound and provides an innovative solution to the equivariance-expressiveness tradeoff.", "other_comments_or_suggestions": "* A detailed parameter count analysis for various model variants would strengthen the paper and clarify the technical contribution of TraceGrad.", "questions_for_authors": "* Have the authors analyzed the impact of the hidden size of  $g_{\\text{nonlin}}(\\cdot)$ on model performance?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "QaLpmcsRNs", "reviewer": "Reviewer_7USG", "summary": "The author propose a new technique to enhance the non-linear expressiveness for equivariant architectures. The idea is simple and straightforward, first generate an invariant feature (like energy) and equivariant feature (like position). The gradient of the invariant feature (energy) with respect to the equivariant feature (position) is an equivariant feature (force). The non-linearity is enforced on the generation of the invariant features. The author validate their methods on two K-S Hamiltonian prediction benchmarks.\n\nPost rebuttal: The authors have addressed most of my concern.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "1.**We have found that our TraceGrad method also significantly improves energy/force prediction tasks** (please refer to the 3rd item of our response to Reviewer AGnk), demonstrating the generality of our method. We plan to follow the reviewer's suggestion and validate our method on more molecular property prediction tasks in the future work.\n\n2.**Although both our method and that of Li et al. construct SO(3)-invariant quantities for supervised learning, the motivations and mechanisms are very different**. Li et al. construct SO(3)-invariant quantities, which are the orbital energies after diagonalizing the Hamiltonian. Using this quantity to supervise the regression model aims to ensure that the orbital energies derived from the model’s output Hamiltonian are as close as possible to the true energies.\n\nIn contrast, our introduction of SO(3)-invariant quantities is intended to supervise high-quality SO(3)-invariant non-linear representations. The goal of these features is to effectively enhance non-linear expressiveness into SO(3)-equivariant features, enabling the fine regression of complex SO(3)-equivariant targets. **The work of Li et al. mentioned by the reviewer, was accepted and published online only about a week before the ICML 2025 submission deadline**, which left us with insufficient time to incorporate it into our work. However, we will certainly cite this paper and explore the possibility of combining it with our approach in future research.\n\nThe SO(3)-invariant quantities we construct, $\\mathbf{T} = \\text{tr}(\\mathbf{H} \\cdot \\mathbf{H}^\\dagger)$, quantify the total coupling strength encoded in the Hamiltonian matrix, representing the overall amplitude of electronic interactions in the system. It is a symmetry-invariant scalar that reflects the global energy scale of $\\mathbf{H}$ and serves as a physically meaningful regularization target.\nFrom a machine learning perspective, it serves as an excellent target to supervise SO(3)-invariant representations.\n**Moreover, this construction is easily extendable to other physical quantities**. For example, in force field prediction tasks, one could construct the trace quantity of the force, which reflects the strength of the force,\n as a supervision signal, train high-quality SO(3)-invariant feature representations, and then induce SO(3)-equivariant features through the gradient mechanism. \n\nIn some cases, using $tr( \\mathbf{C} \\mathbf{H} \\mathbf{H}^\\dagger \\mathbf{C}^\\dagger)$ can also be a reasonable choice. In particular, when $\\mathbf{C}^\\dagger \\mathbf{C} = \\mathbf{I}$, this expression reduces exactly to our definition, $tr( \\mathbf{H} \\mathbf{H}^\\dagger)$. However, its applicability is not as general as that of our proposed formulation, and it is also less straightforward to implement in practice. For instance, in the case of other tensorial physical quantities, such as the force $\\boldsymbol{F}$, which is a vector, or the electron-phonon coupling tensor, which is a rank-3 tensor, our definition remains valid, whereas $tr( \\mathbf{C} \\mathbf{H} \\mathbf{H}^\\dagger \\mathbf{C}^\\dagger)$ may no longer be applicable. Therefore, our trace-based construction offers a unified, symmetry-invariant, and easily implementable approach for supervising or regularizing such tensor quantities across a wide range of physical systems.\n\n3.**A theoretical analysis on the advantages of the TraceGrad method**.\nFirst, our method constructs SO(3)-invariant quantities, i.e., the trace quantity, to directly supervise SO(3)-invariant features, allowing for the effective learning of informative SO(3)-invariant features that capture the intrinsic symmetry properties of the mathematical structure of $\\mathbf{H}$. This helps the model learn informative invariant features and ultimately deliver them to the equivariant features to assist in Hamiltonian prediction. Second, the proposed gradient mechanism, i.e., $\\mathbf{v} = \\frac{\\partial z}{\\partial \\mathbf{f}}$, where $\\mathbf{v}$ and $\\mathbf{f}$ are used to regress $\\mathbf{H}$ and $z$ is used to regress $\\mathbf{T}$, reflects the partial derivative relationship between $\\mathbf{H}$ and $\\mathbf{T}$, i.e., $\\mathbf{H} = \\frac{\\partial \\mathbf{T}}{\\partial Conj(\\mathbf{H})}$, where $Conj(\\cdot)$ denotes the complex conjugate, imposing stronger physical constraints on the relationships between the components of the equivariant features. In contrast to the conventional gated activation mechanism, which can be expressed as $\\mathbf{v} = z \\cdot \\mathbf{f}$, this approach enables effective joint learning of $z$ and $\\mathbf{v}$, with supervision provided by $\\mathbf{T}$ and $\\mathbf{H}$.\n\n4.We will present the related works on Hamiltonian prediction in a dedicated section for clarity.\n\n5.We will revise the equation rendering to ensure standard formatting.\n\n6.Please refer to the 1st item of responses to Reviewer gNai for the discussion on computational cost.\n\n7.We use layer normalization techniques to stabilize the training process.", "claims_and_evidence": "To evaluate the effectiveness of tracegrad, more generalized tasks such as machine learning force field and quantum property regression are needed.", "methods_and_evaluation_criteria": "Yes, but does regressing against  $\\mathrm{tr}(H H^*)$ have any physical meaning? I think a more sensible or physical-inspired regression should be against the basis transformed version as in [1] i.e.  $\\mathrm{tr}(CH H^*C^\\*)$  .\n\n[1]: Li Y, Xia Z, Huang L, et al. Enhancing the Scalability and Applicability of Kohn-Sham Hamiltonians for Molecular Systems[C]//The Thirteenth International Conference on Learning Representations.", "theoretical_claims": "Why tracegrad is better than the paradigm of scalar-tensor interaction and enforcing the non-linearity on the scalar? Does the expressivity differ? Maybe a theoretical analysis could be beneficial here.", "experimental_designs_or_analyses": "Yes, but an efficiency comparison is needed. I am worried that tracegrad will be siginificantly more expensive than scalar tensor interaction paradigm.", "relation_to_broader_scientific_literature": "It is useful for general geometric deep learning.", "essential_references_not_discussed": "Yes, but I believe a special related work section on predicting the Hamiltonian matrix is needed.", "other_strengths_and_weaknesses": "The idea on using the gradient as a vector feature is novel, and could serve as an useful alternative to the scalar-tensor interaction paradigm or CG-tensor products. But the application on predicting the Hamiltonian is questionable. More extensive evaluations on other tasks such as machine learning force fields, or discussion on why the method only works for the Hamiltonian is necessary.", "other_comments_or_suggestions": "1. Many equations are not carefully written, and things like $loss$ should be wrapped with \\mathrm{} in latex.", "questions_for_authors": "1. Can you generalize your idea, for example, the Hessian matrix? I think the Hessian matrix also has equivariant properties and could be decomposed into a set of irreducible representations. A more generalized derivation to higher-order spherical tensor could be beneficial.\n\n2. Can you provide a detailed breakdown of your computational costs? \n\n3. How do you approach the instability (occurrence of NaNs) of the network induced by taking the gradient? Do you make any normalizations? If so, how?\n\n4. Can you apply your method to machine learning force field or molecular property prediction (e.g. dipole moment)?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "b0cS6JZP61", "reviewer": "Reviewer_AGnk", "summary": "This paper presents TraceGrad, a strategy to learn SO(3)-equivariant Hamiltonian with SO(3)-invariant trace as a guiding label, built upon their mathematical relations. It aims to overcome the tradeoff between SO(3)-equivariance constraints and NNs’ nonlinear expressiveness. TraceGrad brings improvement to baseline models in ablation studies.\n\n## update after rebuttal\nThe authors' answers to my questions are thorough and informative. I've maintained my score.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "1. \nWe will correct the quote errors and reduce redundancy in the discussion.\n\n\n2. \nIn this work, we explicitly introduce the trace quantity, i.e., the square of the Frobenius norm of the Hamiltonian $\\mathbf{H}$, as an SO(3)-invariant quantity in our neural network. In physics, invariants often reflect the fundamental mathematical structure underlying physical laws and can serve as the foundation for deriving other quantities with equivariant properties. For example, in special relativity, the Lorentz spacetime interval remains invariant under changes in reference frames, from which one can derive equivariant quantities such as velocity and momentum. Similarly, in molecular systems, energy is an invariant quantity, and its gradient with respect to atomic coordinates yields the force, which is equivariant.\nWe extend this principle—where invariant quantities induce equivariant ones—from specific physical examples to the context of neural network representation learning. This provides a structured inductive bias that guides the model in learning physically meaningful, symmetry-respecting representations.\n\n3. \nAs Reviewer gNai pointed out, the TraceGrad mechanism is highly generalizable. For other symmetry-equivariant tensorial physical quantities, the squares of their Frobenius norms are always symmetry-invariant quantities. These invariants can serve as supervision signals from which the corresponding equivariant quantities can be derived.\n\n We validate the effectiveness of the proposed TraceGrad method on the energy/force field prediction task. Due to the limited time during the discussion period, we conducted experiments on two datasets as representative, i.e., MD17-Aspirin and MD17-Malonaldehyde [1], as a representative example. We use the same setup of this dataset as Liao and Smidt [2]. ote that in this task, the regression target $ E $ (energy) is an SO(3)-invariant quantity ($ l=0 $), while $\\mathbf{F}$ (force) is an SO(3)-equivariant quantity ($ l=1 $). The model typically learns SO(3)-equivariant features $\\mathbf{f}$, which are then transformed into SO(3)-invariant features, from which $E$ is regressed.\nSubsequently, the force field at a given position is obtained by differentiating $ E $ with respect to the atomic coordinates: $\\mathbf{F}_i = -\\frac{\\partial E}{\\partial \\mathbf{r}_i}$, where $\\mathbf{r}_i$ is the position vector of the $ i $-th atom. This approach ensures energy conservation. Given the specificity of this task, we integrate the baseline model, namely Equiformer [2] with our proposed TraceGrad method as follows:\n\nFirst, we use the SO(3)-equivariant features $\\mathbf{f}$ encoded by the baseline model Equiformer as input, and construct SO(3)-invariant non-linear features $z$ according to our method (Sections 4 and 5 of our paper). We then use the trace quantity $\\mathbf{T}$ to supervise the learning of $z$. Given $\\mathbf{F}$ as a column vector with $l=1$, here $\\mathbf{T}$ simplifies to $\\mathbf{T} = \\mathbf{F}^T \\cdot \\mathbf{F}$. From $z$, we induce the SO(3)-equivariant features $\\mathbf{v}$ with more non-linearity, which are then fed back into the baseline model for the subsequent encoding and decoding phases, where $E$ is regressed and finally $\\mathbf{F}$ is constructed from the gradients of $E$. We train Equiformer+TraceGrad under the same experimental conditions as those used in the original Equiformer paper [2], with maximum feature degree ($l_{max}$) set as 2. The experimental results are as below: \n\nFor the MD17-Aspirin dataset, Equiformer achieves an energy MAE of 5.3 meV and a force MAE of 7.2 meV/Å. In comparison, Equiformer+TraceGrad  achieves lower MAE values, with 5.06 meV for energy and 5.65 meV/Å for force. For the MD17-Malonaldehyde dataset, Equiformer achieves  an energy MAE of 3.3 meV and a force MAE of 5.8 meV/Å，while Equiformer+TraceGrad, yields improved performance with an energy MAE of 3.21 meV and a force MAE of 4.68 meV/Å. **These experimental results show that our TraceGrad method improves the prediction accuracy for both energy and force, demonstrating that its effectiveness is not limited to the prediction of electronic-structure Hamiltonians and their downstream physical quantities, but has broader application potential**. Since our method is fundamentally general, we plan to extend it in the future to predict other physical properties, such as the force constant matrix, Born effective charges, and more. **We will add these experimental results and discussions in the revised paper**.\n\n[1] Chmiela et al. Towards exact molecular dynamics simulations with machine-learned force fields. Nature Communications, 2018.\n\n[2] Liao and Smidt. Equiformer: Equivariant graph attention transformer for 3d atomistic graphs. ICLR, 2023.\n\n4. \nTo apply the TraceGrad method, the backbone model only needs to be an end-to-end differentiable strictly SO(3)-equivariant neural network model, such as QHNet, DeepH-E3, or Equiformer.", "claims_and_evidence": "Claims about improvement in prediction accuracy are supported by experimental results.", "methods_and_evaluation_criteria": "The benchmark datasets are comprehensive, and evaluation criteria are reasonable.", "theoretical_claims": "I could not check the correctness of proofs.", "experimental_designs_or_analyses": "Experimental design and analyses look sound to me.", "relation_to_broader_scientific_literature": "Equivariant ML models are actively developed in the AI4Materials community. I’m not sure how widely TraceGrad is applicable to such models (see Questions).", "essential_references_not_discussed": "Related works to my knowledge are comprehensively discussed.", "other_strengths_and_weaknesses": "None.", "other_comments_or_suggestions": "* There’re some errors in using single quotes, e.g., Line 309.\n* The discussion of equivariant operators in Sec. 5.1 seems redundant with Related Works.", "questions_for_authors": "* The expressive is mostly added through nonlinear transformations on the invariant features. Could you comment on the inductive bias this approach assumes, and whether it’s reasonable in physics?\n* Is the TraceGrad strategy applicable to learning materials properties other than full Hamiltonian?\n* What’s the scope of applicability of TraceGrad? In other words, what requirements does the backbone model need to satisfy for TraceGrad to be applicable?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: The paper presents a mean-field electronic Hamiltonian prediction network that guarantees the equivariance of the high-order tensors. The equivariance is achieved by taking the gradient of an invariant feature, similar to the vector-equivariant case. The invariant feature can be processed by arbitrary nonlinearity to provide expressiveness, while also holds the interaction results between tensors of different orders through tensor product. In addition, the final resulting invariant feature is guided by the trace of the Hamiltonian, an invariant quantity. The design makes a good sense, and experiments on common benchmarks show improved accuracy than existing models.

Reviewers asked for more details in more aspects on assessing the work.
* Training and inference efficiency, as asked by Reviewers gNai and 7USG. This is a reasonable ask as gradient-based models typically require more computation and memory. The authors provided a few running time comparisons, showing that not significantly more time cost is induced. Nevertheless, to make it solid, the authors need to further append hardware details right after these results.
* Understanding the expressiveness of the model where tensorial feature is updated with gradient w.r.t a scalar (Reviewers 1Tq8, 7USG). The reply regarding whether the gradient expression can rotate the input tensor makes sense to me, while using $\\mathbf{H} = \\frac{\\partial\\mathbf{T}}{\\partial \\mathrm{Conj}(\\mathbf{H})}$ to explain the mutual benefit between $\\mathbf{H}$ and $\\mathbf{T}$ supervision seems a bit stretching, as the gradient in the model is not directly taken on $\\mathbf{T}$ and $\\mathbf{H}$ is not directly produced from the gradient. I hope the authors could provide further comment on the expressiveness/universality, following the consideration that the gradient w.r.t a vector can only model a conservative vector field, and if the combination of gradient tensors can circumvent a similar constraint. (This may be out of the scope, though.)
* Reviewer gNai asked about the model size (parameter count) in empirical comparisons. While the authors claimed no additional parameters are introduced when implementing the TraceGrad operation, I agree with the reviewer that there are still some more. I hope the authors could provide further information in the final paper.
* Energy/force prediction performance as asked by Reviewer 7USG. The provided results during rebuttal seem supportive.
* Literals in equations should be put upright.

Overall, the issues seem minor compared with the contribution.
Decision: Accept (poster)