Diagonal Symmetrization of Neural Network Solvers for the Many-Electron Schrödinger Equation
Authors: Kevin Han Huang
Ni Zhan
Elif Ertekin
Peter Orbanz
Ryan P Adams
Venue: ICML 2025 poster (2025)
Keywords: quantum Monte Carlo
ab initio methods
Schrodinger equation
neural network wavefunctions
many-body methods
invariance and symmetry
Primary area: applications->chemistry_physics_and_earth_sciences
TL;DR: We study ways of incorporating diagonal symmetry into neural network wave functions, and show that post hoc averaging outperforms conventional in-training symmetrization.

Abstract: Incorporating group symmetries into neural networks has been a cornerstone of success in many AI-for-science applications. Diagonal groups of isometries, which describe the invariance under a simultaneous movement of multiple objects, arise naturally in many-body quantum problems. Despite their importance, diagonal groups have received relatively little attention, as they lack a natural choice of invariant maps except in special cases. We study different ways of incorporating diagonal invariance in neural network ansatze trained via variational Monte Carlo methods, and consider specifically data augmentation, group averaging and canonicalization. We show that, contrary to standard ML setups, in-training symmetrization destabilizes training and can lead to worse performance. Our theoretical and numerical results indicate that this unexpected behavior may arise from a unique computational-statistical tradeoff not found in standard ML analyses of symmetrization. Meanwhile, we demonstrate that post hoc averaging is less sensitive to such tradeoffs and emerges as a simple, flexible and effective method for improving neural network solvers.

Diagonal Symmetrization of Neural Network Solvers for the Many-Electron Schrödinger Equation
Kevin Han Huang'
Ni Zhan
Elif Ertekin
Peter Orbanz' Ryan P. Adams?
Abstract
physical phenomena. Ab initio computational methods seek to solve the non- relativvviiistic electronic Schrödinger equation from first principles. There, the computation is performed directly from physical constraints and without relying on empirical approximations or training data, with the promise of producing
Incorporating group symmetries into neural net- works has been a cornerstone of success in many Al-for-science applications. Diagonal groups of isometries, which describe the invariance under a simultaneous movement of multiple objects, arise naturally in many-body quantum problems. De- spite their importance, diagonal groups have re- ceived relatively little attention, as they lack a natural choice of invariant maps except in special cases. We study different ways of incorporat- ing diagonal invariance in neural network ansätze trained via variational Monte Carlo methods, and consider specifically data augmentation, group averaging and canonicalization. We show that, contrary to standard ML setups, in-training sym- metrization destabilizes training and can lead worse performance. Our theoretical and numeri- cal results indicate that this unexpected behavior may arise from unique computational-statiie tradeoff not found in standard ML analyses of symmetrization. Meanwhile, we demonstrate that post hoc averaging is less sensitive to such trade- offs and emerges as simple, flexible and effec- tive method for improving neural network solvers.
electromic wavefunctions. How- ever, the strict requirement on physical constraints makes it challenging to incorporate neural networks into these meth- ods. Carleo & Troyer (2017), Hermann et al. (2020) and FermiNet (Pfau et al., 2020) are some of the first successful ab initio neural network methods, which learn the ground state wavefunctions atoms and molecules via a variational Monte Carlo (VMC) approach. Many neural network meth- ods have emerged since then (Li al., 2022; von Glehn et al., 2023; Cassella al., 2023), each seeking to improve how physical constraints are modelled in different systems. While these approaches have produced state-of-the sults on ground state energy and other physical properties, one notable drawback is their exorbitant training cost com- pared to classical VMC methods. This issue becomes dire for modelling large systems, as the Hilbert space of wave- functions grows exponentially with the number of electrons.
re-
In other Al-for-science approaches, symmetrization has proved to be successful both for modelling physical con- straints and for improving neural network performance (Batatia et al., 2022; Du et al., 2022; Batzner et al., 2022; Duval et al., 2023). One VMC example DeepSolid (Li et al., 2022), a FermiNet-type wavefunction that employs translationally invariant features model periodic solids. However, recent findings in protein structures (Abramson et al., 2024), atomic potential (Qu & Krishnapriyan, 2024) and machine learning theory (Huang et al., 2022; Balestriero et al., 2022) demonstrate that symmetries may be unnec- essary and could sometimes be harmful for performance. This raises the question how much the previously observed improvements can be attributed symmetries versus other factors such hyperparameter and architecture choices.
1. Introduction
We study the effect of symmetrizing neural network solu- tions to the Schrödinger equation. Solving the many-body Schrödinger equation is of fundamental importance in sci- ence, because it provides the key to understanding and pre- dicting the behavior of quantum systems and thereby many
'Gatsby Unit, University College London, London, UK 2 'Department of Computer Science, Princeton University, Prince- ton. NJ, USA 'Department of Mechanical Science and Engineering, University of Illinois at Urbana-Champaign, Champaign. IL. USA. Correspondence to: Kevin Huang <han.huang. frainn.20@uuc 20@ u...k..c..kk
Motivated by these questions, we examine the effectiveness of symmetrizing ab initio neural network solvers under the natural symmetry groups of a many-body problem, which are the diagonal groups of isometries. These groups, roughly speaking, describe the invariance of the system under
Proceedings of the 42nd International Conference Machine Learning, Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Diagonal Symmetrization of Neural Network Solvers for the Schrüdinger Equation
3
3
2
-2
2
displacement
andisplacement
diddiplaacemen
(a) log ( (OG)) original DeepSolid
(b) log 1
(PA:G) post hoc averaglig
(c) Improvement of log |e) (OG) 1a over log 1 (PAA
40
+3
V (Bohi)
Pel. line
-9
31
M
reA rene
10
91
000
-2.
-0.6
17 (Bohr)
(d) graphene x system,
(Bohr)
(Bohr)
6mm point group
(e) a synmettric configuration of 12 electrons
(i) a shifted version of (e)
Figure 1. Visualizations of the (partial) diagonal invariance of an unsymmetrized wavefunction versus a symmetrized wavefunction in a graphene system. (a) and (b) generated hy evaluating log (x) + f,. +)1² under simultaneous 2d translation of the configuration (71 x given by the 12 blue spheres in (e). The red overlay indicates the unit cell in (e) such that the ref. line exactly the origin when 0. (d) shows atoms (orange spheres) in an planar supercell, with G illustrated by the I objects. (c) and (f) are shifted copies of the same configuration, with their positions marked in (b). This method only visualizes the partial 3m1 symmetry; see Appendix C for a method that shows the full P 6mm symmetry of v/g
(PA Details setups are in Sec. 6, 7 and Appendix B.
simultaneous movement of multiple objects; see Sec. 2 for a detailed review. Common symmetrization approaches in machine learning (ML) roughly fall under three categories: Randomly transforming data by symmetry operations (data augmentation);
in an "infinite data" regime, and every symmetry operation comes at the cost of forgoing one new data point. Hold- ing the computational budget constant, symmetrization can destabilize training and lead to worse performance. This is demonstrated by theoretical and numerical results in Sec.
Averaging over group operations (group averaging);
Post hoc symmetrization helps. At inference time, VMC solvers are less sensitive computational costs, and allow for averaging over a moderate number of group operations (Sec. 5). We show that post hoc averaging (PA) leads to improved energy, variance and symmetry properties of the learned wavefunction (Fig. 1, Table 2). In one case, post hoc averaged DeepSolid achieves performance close to Deep- Solid trained with 10x more computational budget (Sec. 7).
Invariant maps (invariant features and canonicalization) We compare these approaches on DeepSolid (Li et al., 2022), the state-of-the-arts architecture for VMC on solids, for different solid systems. The emphasis is on providing an apples-to-apples comparison by fixing the architecture and hyperparameters, while varying the symmetry parameters. We find that, perhaps surprisingly, the effects of diagonal symmetrization for VMC problems mixed and nuanced. This arises due to the unique combination of challenges posed by VMC and diagonal invariance, as well as the joint consideration of computational cost, statistical behaviors and physical constraints. In particular, our analyses indicate: In-training symmetrization can hurt. VMC training operates
The remainder of the paper provides mathematical discus- sions and computational tools for understanding diagonal symmetries VMC. Sec. 2 reviews the concept of diagonal invariance and discusses why, except for simple cases e.g.. translations or E(3), finding a natural smooth invariant map is difficult. This further corroborated by mathematical and
2

---

Diagonal Symmetrization of Neural Network Solvers for the Many-Electron Schrüdinger Equation
terms. If the set of atom positions {ri} } is invariant under some G, VCoul is invariant under simultaneous transfor- mation of all x2 is by the same G. Note however that VCoul does not satisfy separate invariance, i.c. invariance does not hold æt and 22 are transformed by different group elements, since ((n(=))==2(2)))) * ||ay= in gen- eral. Mathematically, separate invariance can be modelled by the product group G acting on TR and diagonal invari- ance arises as a specific subgroup of G", i.e. Gaing C G". Fig. 2 illlusrraees the difference between Gain and G in 2-electron system: Under the given symmetry, potential function is left unchanged when both electrons reflected, but not when one is reflected and the other is rotated.
1
Figure 2. Different invariances for two electrons in graphene system. Left. Separate invariance under a reflection and rotation. Right. Diagonal invariance under simultancous reflection.
empirical results in the special case of a smoothed canoni- calization (Sec. 4.3 and Appendix E). Sec. 3 briefly reviews the VMC setup and how different computational costs arise. Sec. and respectively examine in in-training and post hoc symmetrization. Sec. 6 discusses our evaluation metrics and develops method for visualizing diagonal symmetries (Fig. 1). Sec. 7 discusses experiment details. Additional results and proofs are included in the appendix.
diig- invariant wavefunction. Throughout this paper, we use the shorthand g(x) = (9(#1)..... ((()) for g E G, and focus on -invarinn potentials V, i.e..
V(x) ((((x)) for all 9 G.
(3)
1.1. Many-body Schrödinger equation
Invariance of V does not imply the invariance of a!: For a translation-invariant V with one electron, Bloch (1929) proves that & is only invariant up to a phase factor, and non-invariant solutions can occur when the ground state degenerate (Tinkham, 2003). Nonetheless, invariant so- lution can always be constructed from a linear combination of these states. The next result confirms this for the general case of n 2 1 electrons and a diagonal group of isometries.
Throughout, we shall consider finding the re-lectron ground state wavefunction ab R 3m + C and the corresponding minimal energy E E R of the Schrödinger equation
H (x)x) = (x(x,, x = (ZI...., 2.)) E 3
(1)
The Hamiltonian H is given by (x)x) = ())x(x) + (x)) A is the Laplacian representing the kinetic en- ergy and V 33 + IR is the potential energy of the physi- cal system. We also denote neural network ansätze by the, parametrized by some network weights 0 RV, Note that in general, the wavefunction depends each electron via (z,, o.), where E {f, 4) is the spin, and db is required to be anti-symmetric with respect to permutations of (æ,, a.)) We focus on the case with fixed spins for simplicity, as is done in FermiNet and DeepSolid.
Fact 2.1. Suppose (v, E) solves (1) and V is invariant under some Gann induced by a group G of isometries. Then for any finite subset G Gf
1/ (x)
1 T E (g((x))
(4)
also solves (1) with respect to same energy E. In particular, if G a subgroup, v is invariant under G.
2. Diagonal invariance
Fact 2.1 motivates us to seek wavefunctions that respect the -idariivvri of the system. Fig. (()) also shows that an unsymmetrized, well-trained wavefunction already attempts to achieve some extent of approximate invariance.
Many physical systems of interest exhibit symmetry under some group G of isometries. ( consists of maps of the form
x
+ Ax + b,
2 E R
(2)
For more references on how Guing- invariinee arises and plays an important role in modelling vi, see Rajagopal et al. (1995); Zicovich- Wilson & Dovesi (1998).
for some orthogonal E R3x3 and some translation E RR3 We focus on groups that are countable. For systems with electrons, G typically causes the potential V in (1) to be invariant under diagonal group Gaiag acting on 3
Challenges in modelling -angrinvree Despite suc- cesses in modelling simple n-electron symmetries such as the Euclidean group E33 (Batzner al., 2022) and transla- tions (Whitehead et al., 2016), generalizing the approaches in those settings to general isometries can be challenging:
Gaiug (gg.....) 9) G)
To how Gaing may arise, consider the Coulomb potential
Vxoul(x) = ='e ||,
+ E,, [z, rzll
Translation groups possess simple and well-understood symmetries. Common symmetrizations include periodic Fourier bases (Rajagopal et al., 1995) and projection via taking a modulus (Dym et al., 2024). The former trades off representation power against computational cost as the
Each ry is the fixed, known position the I-th atom (un- der the Born-Ope Senn-Opennhemerr approximation), || is the Eu- clidean norm, and ... are the omitted electron-indepen independent
3

---

Diagonal Symmetrization of Neural Network Solvers for the Many-Electron Schrüdinger Equation
number of bases used varies, whereas the latter suffers from discontinuity and requires smoothing (Sec. 4.3). Neither has been extended to general Gding.
E(3). A non-exhaustive list of those works can be found in Mahajan & Sharma (2019); Lin et al. (2023); Luo et al. (2023); Zhang et al. (2025) and further comparisons are discussed in Section 4.3.
E(3) consists of all isometries in R², and also admits sim- ple invariant features (Batzner et al., 2022). However, the additional symmetries under F(3) are undesirable when G is only a subgroup of E(3). since the missing asymme- tries represent a loss of information and therefore limit the representative power of the feature. Indeed, building an invariant map without losing information necessitates the availability of maximal invariant (Lehmann et al., 1986), whereas building& a maximal invariant for Gaiag while respecting continuity constraints requires extend- ing the mathematical theory of orbifolds (Ratcliffe, 1994; Adams & Orbanz, 2023) to Gaing of isometries. Such extensions are not known to the best of our knowledge.
Symmetrization of many-body wavefunctions. We shall focus on generic symmetrization techniques from ML that can be applied directly to many-electron wavefunctions. This is to be distinguished with classical techniques, c.g. Zicovich- Miich--ilssn& & Dovesi (1998): There, symmetrization is performed by building invariance into single-electron features, which does not generalize well to offc--f-thhearr neural network solvers that are many-body by design.
3. Neural network VMC solver
We briefly review the VMC approach for training our ab initio neural network solvers. VMC seeks to solve the mini- mum eigenvalue problem of (1) via the optimization
In short, invariant maps are well-understood for groups with simple isometries and with many isometries, and the difficult cases are often found in the ones with restricted symmetries. This problem is particularly pronounced in the case of a diagonal group. Mathematically, the symmetries of Gaiag are described by fixed group G in IR3, but act on a larger and larger space 3m the number of electrons n grows. This means that Gdiag admits substantially less structure than that of e.g. the product group G
(we. Hvo) (wh, vo)
argmin ØER
argmin E E(Erocataooo (X), DERv
(5)
where X X~ Pve and Viocal:vo (x) == H (o)/ve(X))) is
the local energy. (f.g) == f(x))* g(x) dx is the com- plex inner product, and ps = [+((x2)] bility distribution obtained by normalizing a wavefunc- tion w. The optimization may be performed by first or second-order methods, which can be represented by the generic functions Exoe = F(vu(x),, (x)(x)) E IR9 and Qxpo = Q(vx(x)), Awo (x)) E R2x9 as
is the proba-
(W.0)
Space group Gap for a crystal system. A particular diffi- cult case of
in view of the above discus- sion, is the one induced by a space group Gsp in a crystal lattice. These groups are described by (2) with a finite num- ber of orthogonal matrices A and countable number of translations b, which tile the R33 space with a finite unit vol- ume called the fundamental region. Fig. (d) illustrates the P 66m point group in graphene system, where the space is tiled by a triangular fundamental region. The study of space group is a fundamental subject of solid state physics, and we refer interested readers Ashcroft & Mermin (1976).
0 + A - B[Fx; 0 - 0 - B[Qx
E Fx
Examples of Fx:ee and Qx.v can be found in Pfau et al. (2020). Notably, the expectation formulation above converts the expensive integral over the entire space into an expec- tation, which can then be approximated by Monte Carlo averages computed on finitely many samples from Paie
Since Gap is infinite, Fact 2.1 does not apply directly to G= Gsp. In practice however, to avoid computing an infi- nite crystal lattice, a common VMC practice is to restrict the system to a finite volume called the supercell (Esler et al., 2010; Kittel& & McEuen, 2018). This effectively reduces the set of translations b in (2) to a finite set and hence Gap to a finite group, which allows Fact 2.1 to apply.
Training (optimization) phase. Every training step con- sists of two sub-steps: (i) Sampling. Samples are obtained from running N independent MCMC chains with Pe, the target distribution; (ii) Gradient computation. Fx (or Qx:,g) ) is computed on the N samples. Since (i) typically requires computing only the derivative Ugwe, whereas (ii) involves at least do and &e the one-step computational costs of (i) and (ii) typically compare as Csamp < Cyrad- This is particularly true for neural network solvers, where short chains (20- 100 steps) are typically used due the expensive gradient evaluation and that any small increase in per-step cost is amplified by the large number of training steps. We verify this cost comparison in Table and discuss the alternative case with Csamp 2 Cgrad in Appendix I.
While our theoretical results apply to all finite groups of isometries, our numerical experiments focus on G.p. Each Gap is denoted by standard shorthands c.g., P 6mm, and an exhaustive list of Gup can be found in Brock et al. (2016). We emphasize that our focus on the diagonal action of Cap. a difficult class of restricted symmetries found in crystals, sets our work apart from existing VMC works that have inves- tigated translations, simple point symmetries in molecules, as well as the continuous symmetries in SO(2), (2) and
Inference phase. Having obtained a trained wavefunction bh parametrized by 0, we draw samples from long chains
4

---

Diagonal Symmetrization of Neural Network Solvers for the Many-Electrop Schrüdinger Equation
x10-4
that target Pay- These samples used to compute various physical properties of Va that can be expressed expec- tations of P see Sec. 6 for details. Note that the only computational cost occurred here is in terms of Caamp-
6
Ja |arr|660]]1
Symmetrization can be performed during training, infer- ence or both. We analyze symmetrization techniques in the two phases separately in Sec. 4 and Sec. as they involve different computational tradeoffs. Our theoretical analysis considers first-order methods for simplicity. Our numerical results focus on the symmetrization of DeepSolid (Li et al., 2022), a state-of-the-arts neural network solver for solid systems with Gap -Symmetry, trained with the second-order method KFAC (Martens & Grosse, 2015); see Appendix B.I.
OG, 1=1000
DA, 90,
GA, N= =-0 -0,k=11
GA. N 1000 k
1000,
50
100
150
200
250
GPU hours
Figure 3. Normalized variance of differently symmetrized gradient
updates against GPU hours. Experiment dctails in Sec. 7.
Method
N
Ceamp
Cgrad (8)
Total (s)
4. Symmetrization during training
OG
1000
0.16(3)
2.4(3)
2.5(3)
DA
90
12
0.041(5)
2.4(2)
2.5(2)
We first present theoretical analysis of the behavior of gradient updates under different in-training symmetrization techniques. Let .,.... XN be i.i.d. samples from Peo the distribution of m-th step MCMC chain with Pve
GA
90
12
0.16(2)
2.6(1)
2.7(1)
GA
1000
12
1.50(1)
24(1)
25(1)
(m)
GAs
1000
0.16(1)
2.4(1)
2.5(1)
Table I. Computational cost per training step. Details Sec. 7.
the target distribution. Our benchmark for comparison is
the update rule with the original (OG) unsymmetrized the.
mented N samples incur the same gradient evaluation cost
Cgrad, which dominates the overall computational cost.
0 + 0 5g((G)
(O()G)
N EigN Fx
(6)
While the sampling cost, Csamp/ kk enjoys a minor speed-up. the next result shows that this comes at the cost of increased instability of the gradient estimate. Below, we denote the
Sample size bottlenecked by computational cost. A key element of our analysis is that the VMC methods are in an infinite data regime. More precisely, unlike setups where sample size is constrained by the number of data points commonly found in many theoretical analyses of sym- metrization techniques (Chen et al., 2020; Lyle al., 2020: Huang et al., 2022)
(m)) distribution of X = g1, (X1) by :o and define distance between two distributions p, p' on R33 as
d=(p,r)) supfer [[E[/[X)]] E,, [XX]]
where F is class of 3n + Rp+p? test functions such that
theoretically allowed to draw infinitely many samples from the sampling step during train- ing. The practical limitation comes from Csamp and Cgrad, both of which are affected by the batch size N as well as the symmetrization techniques used. By taking the computa- tional effects into account, we shall see that symmetrization techniques exhibit substantially different statistical behay- iors from those observed in the existing literature.
{x + (Fx F
|@
R$} S F.
dy is called an integral probability metric (Müller, 1997). Proposition 4.1. Fix 0 E IRP. Then
[[(])(DA) - -[ & dz (p Po
(m)
(mm)
Var[sg - -Vr[[g((OG)] (A [Fx Xi]
4.1. Pitfalls of data augmentation (DA)
1+ 1111500000111 d (p(p)))) (P P (mm))
Since "training data" correspond samples drawn from (m)
(m)) dr dp Per (mm)
Pe. a k-fold data augmentation is performed as follows:
N
In particular, the distribution (m) is invariant under Criug.
Sample N/k ....... .... Lid.d (m) Polo
iii) Sample g..l...., BN/K, ii.d. from some distribu- tion on G;
i.e. g(X,) f all g E Grung we have
((80(DA)] = E[50(0G)] and Var[60(DA)] Var(80 where is the Loewner order of non-negative matrices.
(iii Compute the DA update as 0 f 50(DA), where
DA) =
1 N WESN/h Ejst R (x. (x):%
The error d=<p(m)) P describes how much an aug- mented sample from m-step chain deviates from an unaug-
(m)
(m²)
Notice that the sample size in (i) is reduced to N/k, since a set of k-times augmented N/k samples and set of unaug-
mented sample on average, as measured through the gradient
5

---

Diagonal Symmetrization of Neural Network Solvers for the Many-Eleetrer Schrüdinger Equation
Faniio and the squared gradient F For early training, we expect this error to have small contributions to the overall optimization compared to other sources of noise, e.g., the error incurred by running short chains instead of long chains and by stochastic gradients. At the final steps of training,
canon.
f(w) 0
smoothed canon...
J
Figure 4. Canonicalization functions for 1d unit translations.
expect this error to be small as P ( becomes approxi- mately invariant; Fig. (a) shows that this is the case for an unsymmetrized, well-trained neural network.
under KFAC, a second-order method used by DeepSolid, is also inflated under DA. Table also verifies that the compu- tational speedup from DA is negligible. Fig. 5 and Table 2 show that training with DA leads to worse performance.
Several messages follow from Proposition 4.1:
DA leads to similar gradients in expectation but a pos- sibly worse variance. This is in stark contrast to known analyses of DA in the ML literature, where DA for empirical averages is expected to improve the variance (Chen et al., 2020; Huang et al., 2022). This surprising difference arises because those analyses focus on statistical errors arising from augmenting a size-N real-life dataset to a size Nk dataset, whereas our analysis pays attention to both statis- tical errors and computational errors in a setup that com- pares a size N dataset versus a size N/k x l augmented dataset. Indeed, since the only computational saving of augmentation is the sampling cost Cramp < Coptim. every augmentation comes at the cost of one i.i.d. sample unused.
4.2. Group-averaging (GA)
Fix 9 with |g] = and let vig be the averaged network obtained from taking di to in (4). As with DA, since k-times more derivatives are computed, we need to use a batch size of N/k to maintain the same computational cost. Draw X
X
i..d. (m) The GA update rule is Pug
N/R
~
θ - 5B(GA)
50(GA) = N/k EISNI ExP:08
The mean and variance of 50(GA) is straightforward:
Lemma 4.2. Fix O E RP. Then
Var[o0(GA)]
Instability of DA is not specific to mean and variance. While Proposition 4. only controls the mean and the vari- ance, they do describe the distributions of 5g(DA) and 58(OG) well, even in the digeeniimesion regime where the number of parameters g is large compared to the batch size N. This is due to recent results on high- nigh-dimensional Central Limit Theorem (CLT): In Theorem D.I in the ap- pendix, we adapt results from Chernozhukov et al. (2017) to show that (D(DA) and (OG) are approximately normal
N/k
Remark 4.3. As with DA, we can mean and variance as proxies to understand the distribution of 50(GA) since a high-dimensional CLT does apply here; see Appendix D.
GA may also destabilize gradients. The decrease in sample size is again visible in the variance in Lemma 4.2, which suffers from Vi blowup. Notice that the reference mean and variance are stated in terms of the gradient DA parameter update is a general feature of the distribution, and not just specific to the mean and the variance.
in an appropriate sense. This shows that the instability of
Fx "xfref- which depends on the GA wavefunction both through the samples X and through the gradient evaluation at vy. Unlike the discussion in Proposition 4.1, we no longer expect that the mean and variance of Fxp:ug are close to the unsymmetrized analogue Fx since it does not suffice for X and X₃ to have similar distributions. In general, Var[50(GA)] increases if and only if the ratio
Applicability to multi-step updates. Proposition 4.1 con- cerns one-step gradients, but the analysis applies to multi- step updates: VMC methods draws fresh samples at every step conditionally on the parameter 0 from the previous step, so the same bounds hold with E[ and Var[ replaced by the conditional counterparts B[ O and Var (0).
VarlFxx iv
Var [
One limitation of the above analysis is that it is restricted to first-order updates. Extending similar analyses to second-order updates is a known challenge in the litera- ture: Those methods typically pre-multiply the empirical average N EieN/K EşSk (x)vve by the inverse of an empirical Fisher information matrix, which is also affected by augmentation, and there exist pathological examples where augmentation may increase decrease the variance depending on problem-specific parameters (see e.g., ridge regression analysis in Huang et al. (2022)). Nevertheless, we confirm numerically in Fig. 3 that the gradient variance
Empirically, we see that variance increase for (0(GA) is visible for KFAC in Fig. 3 compared to 50(OG) with similar computational costs (Table 1).
We also include two further comparisons:
GA with subsampling (GAs). One way to circumvent this computational hurdle is to average over size-k uniform subsample of G at every training step, and use the full g at inference time. We numerically investigate the effects of keeping N constant and uniformly sample l = 1 element: While GAs further destabilizes the gradient (Fig. 3), Table 2
6

---

Diagonal Symmetrization of Neural Network Solvers for the Many-Electron Schrüdinger Equation
shows that its performance improves from OG. Yet, it falls short of the performance obtained by post hoc averaging directly on the original wavefunction.
referenced above address this by building simpler early- layer features that reduce the amount of averaging required. However, this approach restricts the representation power of earlier layers, and is not directly compatible with the DeepSolid architecture. It remains unclear whether those ar- chitectures can be adapted to perform well in solid systems. We focus on DeepSolid for the purpose of our investigation.
GA with same N. Say the batch size N is kept the same. While the per-step cost of a k-fold averaging increases by a k times (Table 1), one may ask if the number of steps until convergence may be reduced under the symmetrized wavefunction, such that the overall training cost is constant. Fig. 3 shows that while such a reduction does appear, it is not enough to offset the increase in per-step cost.
5. Post hoc symmetrization
An alternative to in-training symmetrization is post hoc sym- metrization: We may first train 0 with unsymmetrized up- dates (e.g. (6)), and seek to symmetrize 33 during inference. In contrast to Sec. 4, post hoc symmetrization no longer in- curs the cost of Cgrad- While computing properties based on Monte Carlo estimates still incurs Caamples samples are typ- ically obtained in large batches from long chains only once, before being used for multiple downstream computations. This allows us to perform moderate amount of averaging without compromising on sample size. Meanwhile, a di- rect evaluation of the wavefunction, e.g. for producing the visualization in Fig. 1, also does not incur Csample-
4.3. Smoothed canonicalization (SC)
Another symmetrization method that gained traction in the theoretical ML community is canonicalization, defined as the projection to the fundamental region of a given group (Kaba al., 2023; Dym et al., 2024). For ld unit translations, an example of canonicalization is the map + z mod 1, illustrated in blue in Fig. 4. Canonicaliza- tion for space groups Gap is possible but suffers from non- smoothness at the boundary, as visible at 0 and in Fig. 4 and as we show in Appendix E. Jastrow factors (Whitehead et al., 2016) from VMC methods can be viewed as a way to smooth 1 canonicalization along each lattice vector, but the construction is specific to translations. Dym et al. (2024) proposes smoothed canonicalization (SC) for large permu- tation and rotation groups by taking weighted averages at the boundary. In Appendix E, we adapt their idea to develop an SC for diagonal invariance under Gap that also respects the anti-symmetry constraint of (1). The orange curve in Fig. 4 illustrates our method for d translations. However, we demonstrate in Appendix E that SC via weighted aver- aging requires averaging over |Ge] elements, where Ge is some carefully chosen subset of Gap and n is the number of electrons; the additional cost of arises from an anti- symmetry requirement. Therefore, SC suffers from similar computational bottlenecks as DA and GA and typically to a worse extent. This renders SC unsuitable for training. Since SC for Gain may be of independent interest, we provide theoretical guarantees and discussion of its weaknesses in Appendix E.
A wavefunction the and should not involve exogeneous randomness. As such, we do not consider DA for post- processro,ss,ngg. and discuss only group averaging and canonicalization in this setup.
physical object, is deterministic
Post hoc averaging (PA). Since the cost of averaging still scales linearly with [G], taking an average over the entire G can still be prohibitive when G is large. Fortunately, Fact 2.1 ensures the validity of averaging over any finite subset
G C G of our choice. One may choose G to be a subgroup of interest, or a generating set of G. Given a trained wavefunction vig. the corresponding PA wavefunction reads
Ws (PA.C))
(x)
= 1 Ege vg(g(x)). 1
Estimates of physical quantities are obtained by drawing N samples from P Paaaa incurring a computational cost of N [g] Czamp per MCMC step. For 9 of a moderate size, we can compare N samples from P Pagerwe) directly with N samples from Peg. without incurring a larger statistical error. (PA:9) offers two clear advantages over V/a (b B
SC is an architecture- agnossiic invariant map for Gaiag, and its shortcomings demonstrate our discussion in Section 2 that obtaining invariant maps for such restricted symme- tries can be difficult. An alternative approach (Han et al., 2019; Zepeda-Núñez et al., 2021: Gao & Günnemann, 2021: Gerard et al., 2022) is to introduce invariance implicitly, by building feature maps based on electron- -tttoon-attoooooo distances. One must then ensure the network is invariant under per- mutations of atomic indices. Introducing this permutation invariance naively, by averaging, requires an expensive sum- mation over a permutation group, and introduces additional complications with the boundary condition. The works
(i) More symmetry. By construction, us (T(A:) exhibits a higher degree of symmetry than v/g. the extent of which depends on the subset G chosen and the degree of symmetry already present in v/g- See Fig. 1 for the improved symmetry.
(ii) Robustness to outliers. The ground-truth wavefunc- tions are required to be non-smooth whenever an electron coincides with an atom or an electron. At those regions of x, the Hamiltonian H Hiw(s (x) diverges, but for ground- truth wavefunctions as well as carefully constrained clas-
7

---

Diagonal Symmetrization of Neural Network Solvers for the Many-Electron Schrödinger Equation
75.97
1000
x10 2
6. Evaluation and visualization methods
90. k-12 -2
-75.98
GA. N= N 12
GA, N= 1000. k=12
1.2
VMC wavefunctions are typically assessed via
GAS, --1000,
-75.99
PA. NA 1000,
[R] Flocalii (X) and Var(Fhocal:ei (X) X ~ Pve
(7)
-76.00
0.8
The energy our optimization objective (5). The variance is another measure of fit for (1): It admits lower bound 0 that is attained by any true solution v. to (1), since Floculeyt is everywhere constant. See Kent et al. (1999).
-76.01
4-
0.6
-76.02
0.4
-76.03
2
To show the amount of approximate symmetry already present in the OG wavefunction v/ we compare us ((OG) against PA wavefunctions averaged over different G's. This is reported as Var[PA/OG] in Table 2, which stands for
-76.01
0.2
-76.05
0.0
100
200
100
200
100 200
(a) Energy (Ha)
(b) Var (lacel (Ha*) (c) Symmetry
Performance of wavefunctions against GPU hours, ob- tained with different symmetrization methods. Metrics are defined in Sec. and experiment details in Sec. 7. (c) is computed via Var Har[AAOG in Sec. 6 but with 5% (OG) replaced by different thy 's.
Figure 5.
Var
(OG) ges % (0)) (g(())))))))))))
(OG) (x)
X
We also seek to visualize the for its symmetry under Gaing Visualizing diagonal symmetry can be challenging, as Gaing acts on a tigh-dimensional space R3, We propose vi- sualization method that exploits the fact that our Gaing is completely described by group G of isometries in 13:
sical wavefunction ansatz, the local energy Flocal::), (x) (x)) /// (x) in (5) stays finite due to the renormalization by the (x). The same is not necessarily true for neural net- work solvers, which are much more flexible by design: The Laplacian term in H Ho (x) may become large even when the magnitude of (x) stays put. This issue also manifests in the evaluation of other empirical quantities, e.g., the gra- dient terms computed in MCMC. The averaging in v/g (PA:))
ii Let G be group acting on I defined as TR 3x3 A(+) + E G for some b E IR3): (ii) Fix Xaymm E 3 configuration of electrons such that for every g E G, ((aymmm) = Xaymm
= {A E
(iii) Given a function f 3t + R to visualize, we plot the function f(t) == f f(aymmm + t) with t E R², i.c. all electrons are translated by simultaneously.
improves robustness in the sense that, if (g())) takes a large numerical value for one particular g E 9 and not for
The next result confirms the validity of this method:
the rest, the numerical outlier is rescaled by factor - The level of robustness also increases with the size of G.
Lemma 6.1. Let f (R + R be a function invariant under permutations of its arguments. Then for any g E G and R3, f(g(t)) f()) [[g(( ()) f(x 1).
Fig. 5 and Table 2 show that (PA:9) outperforms in-training symmetrization with the same computational costs in all metrics considered. Compare, for example, PA with 40k steps of training with N = 1000 versus GA with 10k steps of training, N = 1000 and h = 12: PA achieves a lower energy with lower variance as well as perfect symmetry, with only 1/4 of the training budget. Among methods with similar end-of-trainine energies, PA also attain lower ergy and variance with fewer training steps (Fig. 5). We also remark that GA and GAs by default implement PA at inference, so their only difference with PA is from training.
The permutation invariance assumption holds for //o and Eiocal:te since the is anti-symmetric. Fig. ((a), (b) and (f) are plotted with this method, with f(x) log (x)) Xaymi given in Fig. I(e) and G G = P3m1, which illus- trates the partial symmetries of the P 6mm group of graphene. To see the G P 6mm symmetry, different Xsymm is re- quired since G a G in this case; see Appendix C.
7. Experimental details
Issues with post hoc canonicalization (PC). One may also use smooth canonicalization post hoc to symmetrize trained wavefunction 4/-- For completeness. we record the performance of PC in Table 2. The results are signifi- cantly worse than other methods, despite being applied to well-trained wavefunction. The issue might arise from the fact that a weighted averaging near the boundary leads to a blowup in second derivatives, and we examine it in detail in Appendix E. This makes PC unsuitable specifically for our problem, since Frocal involves the Hamiltonian.
All code and data are available at github..ccm com/ PrincetonLIPS
DeepSolid. Experi- ments are performed with DeepSolid (Li al., 2022) on crystalline solids. Each network is evaluated by sampling from MCMC chains with 30k length, and the model from the last training step is used unless otherwise specified. Su- percell size is included in the first column of Table 2. Ap- pendix includes further specifications, experiments and a remark about why energy improvements in the decimals are considered substantial. few remarks about each system:
8

---

Diagonal Symmetrization of Neural Network Solvers for the Many-Electree Schrüdinger Equation
System
g
Method
NV
Steps
GPU hours
Energy (Ha)
Varj Eiocst (Ha")
VarlA/OGG
OG
1000
S0, 000
281
76.039(G]
2.02(3)
DA
12
80, DUU
277
76.039(3)
2.48(5)
Graphenc
GA
12
50, 000
304
76.048 76.049(3)
0.58(2)
GA
1000
12
10, 000
351
76.034(5)
1.2(2)
6mm
GAs
1000
1
80, 000
281
-6.0749 (4)04)(3)
0.48(2)
PA
1000
12
80, 000
281
76.050(3)
0.33
0.00180(4)
80,
-70.1(9)
8(1) 10
OG
-..38(()
0.06(1]
Lithium Hydride
PA
2
8.144(1)
0.0344(9)
(3)))3(4)
22/m
-8.148(1)
0.0197(6)
0.0249(5)
(1.HH)
F222
PA
4000
16
30, 000
571
8.1495(9)
(.0162(
0.0265
2 x
2m3m
48
8.1502(7)
0.0122(7)
0.0289 (8)
Fm3m
PA
192
8. (50007
0.0118(7)
0.0289
Metallic Lithium (hcc-1.i)
OG
15.01I(1)
0.059(2)
P4/mmm
16
15.021(2)
0.033(2)
0. 0.092(5)
Emmm
PA
3000
32
20, 000
462
15. (5.020(1)
0.036(2)
(.0101(4)
2 x 2 x
1m3m
15.022(3)
0.031(3)
0.0139(6)
Table Performance of symmetrization methods with similar computational budgets. Energy and variance are both reported at the per unit cell level. *See Appendix B.3 for specifications of the GPUs used for training.
Graphene. This is the setup considered in Fig. 1,3,5 and Table 1. PA outperforms other methods both in terms of the metrics in Table and speed of convergence in Fig. 5.
Regarding the implications of our findings for other settings, we note number of points:
Computationallstatistti tradcoff's for DA and GA. We expect the in-training tradeoffs observed in Section to be applicable to other ML setups where sampling performed in between gradient updates. A non-physics example is the contrastive divergence algorithm for training energy-based models (Hinton, 2002; Du et al., 2020).
LiH. We use nested subgroups of Fm3m and observe that the performance improves with k. For comparison, Li al. (2022) report the energy -8115096(1) for DeepSolid trained with 3e5 steps and batch 4096. PA attains comparable per- formance in similar systems' with 3e4 steps and batch 4000.
bec-Li. This is a known difficult case for ab initio meth- ods including DeepSolid (Yao et al., 1996; Li et al., 2022) and PA again helps. Fmmm and P4 /mmm are subgroups of Im3m containing different symmetries, and each offers similar improvements. For both LiH and bec-Li, we also observe a saturation effect: The improvement saturates once sufficiently many symmetries are incorporated. We do not know whether an optimal choice of subgroup exists that balances performance and computational cost.
Computational cost of SC. The computational cost dis- cussed Section 4.3 is specific to the canonicalization we have adopted from Dym al. (2024), which adopts an seerrage-nea-tth--buueeee
approach. It is an open ques- tion whether more efficient canonicalizations exist for Gap-
Applicability beyond DeepSolid and VMC. We conjec- ture that both our negative in-training findings and positive post-training findings are applicable beyond DeepSolid and VMC, especially in cases where difficult cases of restricted symmetries arise. As discussed in Section 2, if the system of interest possesses simpler symmetries such translations E33), there often exist more efficient symmetrizations than averaging or augmentation. These may avoid the in-training statistical-computtiinn tradeoffs we observe.
8. Discussion
For ML problems that exhibit geometric structure, conven- tional wisdom holds that incorporating the correct sym- metries should improve performance. This is called into question by the recent success of AlphaFold3 (Abramson et al., 2024), where substantial accuracy improvement is obtained without incorporating invariance into its architec- ture, well as by other works atomic potential and ML theory referenced in Section 1. Our work investigates such tradeoffs in the context of wavefunctions, and specifically of VMC. To this end, we compare different model-agnostic symmetrizations. To ensure comparability, we do so on a single architecture, and choose hyperparameters as in the original DeepSolid work (Li et al., 2022). We not claim have exhausted all possible ways in which symmetries may affect VMC.
Theoretical analyses. While our experiments focus on DeepSolid, our theoretical analyses describe more generally how DA, GA and SC interact with VMC.
'The energy by al. (2022) for lattice vector 4.0 Á. We followed the Materials Project (Jain al., 2013) to use 1.02 A.
9

---

Diagonal Symmetrization of Neural Network Solvers for the Many-Electron Schrödinger Equation
Acknowledgement
kristallgittem. Zeitschrifi für physik, 52(7):555-600, 1929.
This work was partially supported by NSF OAC 2820. KHH and PO are supported by the Gatsby Charitable Foun- dation (GAT3850). NZ acknowledges support from the Princeton AIS initiative. This work used Princeton ionic cluster and Delta GPU at the National Center for Supercom- puting Applications through allocation MAT220011 1 from the Advanced Cyberinfrastruccure Coordination Ecosys- tem: Services & Support (ACCESS) program, which is supported by National Science Foundation grants #2138259, #2138286, #2138307, #2137603, and #2138296. We also thank anonymous reviewer for pointing out regimes where per-step sampling cost could be larger than per-step optimisation cost.
Brock, C. P., Hahn, T., Wondratschek, H., Müller, U., Shmueli, U., Prince, E., Authier, A., Kopský, V., Litvin, D., Amold, E., al. International tables for crystallog- raphy volume A: gpace-group symmetry. Wiley Online Library, 2016.
Carleo, G. and Troyer, M. Solving the quantum many-body problem with artificial neural networks. Science, 355 (6325):602-606, 2017.
Cassella, G., Sutterud, H., Azadi, S., Drummond, N., Pfau, D., Spencer, J. S., and Foulkes, W. M. C. Discover- ing quantum phase transitions with fermionic neural net- works. Physical Review Letters, 130(3)003401, 2023.
Impact Statement
Chen, L. H., Goldstein, L., and Shao, Q.-M. Normal approx- imation by Stein's 's method, volume 2. Springer Science & Business Media, Berlin, 2011.
This paper presents work whose goal advance the field of Machine Learning. There are many potential societal consequences of our work, none which we feel must be specifically highlighted here.
Chen, S., Dobriban, E., and Lee, J. H. A group-theoretic framework for data augmentation. Journal of Machine Learning Research, 21(245):1-71, 2020
References
Chernozhukov, V., Chetverikov, D., and Kato, K. Central limit theorems and bootstrap in high dimensions. Annals of Probability, 45(4):2309 2352, 2017.
Abramson, J., Adler, J., Dunger, J., Evans, R., Green, T., Pritzel, A., Ronneberger, O., Willmore, L., Ballard, A. J., Bambrick, J., et al. Accurate structure prediction of biomolecular interactions with alphafold 3. Nature, 630 (8016):493-500 2024.
Du, W., Zhang, H., Du, Y., Meng, Q. Chen, W., Zheng, N., Shao, B., and Liu, T.-Y. SE(3) equivariant graph neural networks with complete local frames. In Interna- tional Conference on Machine Learning, pp. 5583-5608. PMLR, 2022.
Adams, R. P. and Orbanz, P. Representing and learning functions invariant under crystallographic groups. arXiv preprint ar wiXx:23366002661 2023.
Du, Y., Li, S., Tenenbaum, J., and Mordatch, Improved contrastive divergence training of energy based models. arXiv preprint arXiv:2012.0133: 2020.
Ashcroft, N. and Mermin, N. Solid State Physics. Cengage Learning, 1976.
Balestriero, R., Bottou, L.,, and LeCun, Y. The effects of regularization and data augmentation are class dependent. In Conference on Neural Information Processing Systems, 2022.
Duval, A. A., Schmidt, V., Sernández-Garcaa, A., Miret, S., Malliaros, F. D., Bengio, Y., and Rolnick, D. Faenet: Frame averaging equivariant gnn for materials modeling, In International Conference on Machine Learning, pp. 9013-9033. PMLR, 2023.
Batatia, I., Kovacs, D. P., Simm, G., Ortner, C., and Csányi, G. Mace: Higher order equivariant message passing neural networks for fast and accurate force fields. Ad- vances in Neural Information Processing Systems, 35: 11423-11436, 2022.
Dym, N., Lawrence, H., and Siegel, J. W. Equivariant frames and the impossibility of continuous canonicaliza- tion. arXiv preprint arXiv:2402. 16077. 2024.
Esler, K., Cohen, R., Militzer, B., Kim, J., Needs, R., and Towler, M. Fundamental high- gigh-pressure calibration from all-electron quantum monte carlo calculations. Physical Review Letters, 104(18):185700 2010.
Batzner, S., Musaelian, A., Sun, L., Geiger, M., Mailoa, J. P., Kornbluth, M., Molinari, N., Smidt, T. E., and Kozinsky, B. E(3)-equivariant graph neural networks for data-efficient and accurate interatomic potentials. Nature communications, 13(1):2453, 2022.
Gao, N. and Günnemann, S. Ab-initio potential energy sur- faces by pairing gnns with neural wave functions. arXiv preprint arXiv:2 10.05064, 2021.
Bloch, Über die quantenmechanik der elektronen in
10

---

Diagonal Symmetrization of Neural Network Solvers for the Many-Electron Schrüdinger Equation
Gerard, L.,, Scherbela, M., Marquetand, P., and Grohs, P. Gold-standard solutions to the schrödinger equation using deep learning: How much physies do we need? Advances in Neural Information Processing Systems, 35:10282- 10294, 2022.
Lin, J., Goldshlager, G., and Lin, L. Explicitly antisym- metrized neural network layers for variational monte carlo simulation. Journal of Computational Physics, 474: 111765, 2023.
Luo, D., Chen, Z.,, Hu, K., Zhao, Z., Hur, V. M., and Clark, B. K. Gauge-invariant and anyonic-smmmett autoregres- sive neural network for quantum lattice models. Physical Review Research, 5(1):013216, 2023.
Han, J., Zhang, L.., et al. Solving many-electron schrödinger equation using deep neural networks. Journal of Compu- tational Physics, 399: 198929,,,, 2019.
Hermann, J., Schätzle, Z., and Noé, F. Deep-neural-network solution of the electronic schrödinger equation. Nature Chemistry, 12(10):891-897, 2020.
Lyle, C., van der Wilk, M., Kwiatkowska, M., Gal, Y., and Bloem-Reddy, B. On the benefits of invariance in neural networks. arXiv preprint arXiv:2005.00172 2020.
Hinton, G. E. Training products of experts by minimizing contrastive divergence. Neural compulation, 14(8):1771- 1800, 2002.
Mahajan, A. and Sharma, S.
jastrow mean-field wave function in variational monte carlo. The Journal of Physical Chemistry A, 123(17)33111-3221 2019.
Huang, K. H., Orbanz, P., and Austern, M. Gaussian and non-gaussian universality of data augmentation. arXiv preprint arXiv:2202.0913:4 1, 2022.
Martens, J. and Grosse, R. Optimizing neural networks with Cronecker-factored approximate curvature. In Interna- tional conference on machine learning, pp. 2408-2417. PMLR, 2015.
Huang, K. H., Liu, X., Duncan, A., and Gandy, A. A high- dimensional convergence theorem for u-statistics with applications to kernel-based testing. In The Thirty Sixth Annual Conference on Learning Theory, pp. 3827-3918. PMLR, 2023.
Müller, A. Integral probability metrics and their generating classes of functions. Advances applied probability, 29 (2):429-443, 1997.
Jain, A., Ong. S. P., Hautier, G., Chen, W., Richards, W. D., Dacek, S., Cholia, S., Gunter, D., Skinner, D., Ceder, G., et al. Commentary: The materials project: A materials genome approach accelerating materials innovation. APL materials, 1(1), 2013.
Perdew, P., Burke, K., and Ernzerhof, M. Generalized gradient approximation made simple. Physical review letters, 77(18)33865, 1996.
Pfau, D., Spencer, S., Matthews, G., and Foulkes, W. M. C. Ab initio solution of the many-electron schrödinger equation with deep neural networks. Physical review research, 2(3):033429, 2020.
Kaba, S.-O., Mondal, A. K., Zhang, Y., Bengio, Y., and Ravanbakhsh, S. Equivariance with learned canonicaliza- tion functions. In International Conference on Machine Learning, pp. 15546-15566. PMLR, 2023.
Qu, E. and Krishnapriyan, A. The importance of being scalable: Improving the speed and accuracy of neural network interatomic potentials across chemical domains. Advances in Neural Information Processing Systems, 37: 139030-139053, 2024.
Kent, P. R., Needs, R., and Rajagopal, G. Monte carlo energy and auriance-mnimizzan techniques for opti- mizing many-body wave functions. Physical Review B, 59(19):12344) 1999.
Rajagopal, G., Needs, R., James, A., Kenny, S., and Foulkes, W. Variational and diffusion quantum monte carlo calcu- lations at nonzero wave vectors: Theory and application to seamond-structure germanium. Physical Review B, 51 (16): 10591, 1995.
Kittel, C. and McEuen, P. Introduction to solid state physics.
John Wiley & Sons, 2018.
Lehmann, E. L., Romano, J. P., and Casella, G. Testing statistical hypotheses, volume 3. Springer, 1986.
Rateliffe, J. Foundations of hyperbolic manifolds, Graduate Texts in Mathematics/ppringer Veriiiggge--eelgg... 149, 1994.
Li, R., Ye, H., Jiang, D., Wen, X., Wang, C., Li, Z., Li, X., He. D., Chen, J., Ren, W., et al. A computational frame- work for neural network-based variational monte carlo with forward laplacian. Nature Machine Intelligence, 6 (2):209-219, 2024.
Shevtsova, I. Optimization of the structure of the moment
bounds for accuracy of normal approximation for the distributions of sums of independent random variables. PhD thesis, Moscow State University, 2013.
Li, X., Li, Z., and Chen, J. Ab initio caleulation real solids via neural network ansatz. Nature Communications, 13 (1):7895, 2022.
Tinkham, M. Group theory and quantum mechanics. Courier Corporation, 2003.

---

Diagonal Symmetrization of Neural Network Solvers for the Many-Electron Schrüdinger Equation
von Glehn, I., Spencer, J. S., and Pfau, D. A self-attention ansatz for ab-initio quantum chemistry. In The Kleventh International Conference on Learning Representations, 2023.
Whitehead, T., Michael, M., and Conduit, G. Jastrow corre- lation factor for periodic systems. Physical Review B, 94 (3):035157, 2016.
Williams, K. T., Yao, Y., J., Chen, I.,, Shi, H.,, Motta, M., Niu, C., Ray, U., Guo, S., Anderson, R. J., al. Direct comparison of many-body methods for realistic electronic hamiltonians. Physical Review X, (O<1):011041, 2020.
Yao, G., Xu, J., and Wang, Pseudopotential variational quantum monte carlo approach to bec lithium. Physical Review B, 54(12):8393, 1996.
Zúñed,-Núñez, L.,, Chen, Y., Zhang, J., Jia, W., Zhang, L., and Lin, L. Deep density: circumventing the kohn- sham equations via symmetry preserving neural networks.
Journal of Computational Physics, 443:110523, 2021.
Zhang, Y., Jiang. B., and Guo, H. Schroödingernet: A universal neural network solver for the schrödinger equa- tion. Journal of Chemical Theory and Computation, 21 (2):670-677, 2025.
Dicovich-Wilson, C. and Dovesi, R. On the use of symmetry- adapted crystalline orbitals sef-lcao periodic calcu- lations. ii. implementation of the self-consistent-feet field scheme and examples. International journal of quan- tum chemistry, 67533311-33 1998.
12

---

Diagonal Symmetrization of Neural Network Solvers for the Many-Electron Schrödinger Equation
The appendix is organized follows:
Appendix collects the notation used throughout the paper.
Appendix B includes additional figures and details on the experiments.
Appendix C proves the validity of the diagonal visualization method (Lemma 6.1) and includes additional details on how to visualize the full symmetry beyond Fig. 1.
Appendix states the high-dimensional central limit theorem discussed Section 4.
Appendix E develops and discusses the smoothed canonicalization method discussed in Sections 4.3 and 5 for diagonal invariance under space group.
Appendices prove all mathematical results developed in this paper.
Appendix includes an additional discussion on the case where per-iter sampling cost is greater than per-iter gradient cost.
A. Notation
Tt: number of electrons
N: number of Markov chains number of Markov chain samples
length of Markov chain
c: number of symmetry operations
z: -aluued position of a single electron. ă = (x, 0) E R x { additionally specifies the spin of the electron.
- positions of electrons, also called an electron configuration. The diagonally transformed configuration .(.))) Similarly, x E (R3 {1,4))" is the n-electron analogue of ă.
is denoted as (xx
Space group acting on 3R3
Go: Point group subgroup of G, consisting of elements of the form x - A(x) where A is an orthogonal 3 x matrix
Gaig:: Diagonal group induced by the space group, acting electron configuration in 3
G: subset (not necessarily subgroup) of elements of Gaing
Pa: permutation group acting on an electron configuration in 3m
T sup: translation group acting on R that represents the supercell assumption
@: the tensor product
||: Fuclidean norm
); I-th coordinate of vector
absolute constant: number that does not depend on any other variables, and in particular not on n, l or M
equality in distribution
B. Experimental details and additional results
B.1. DeepSolid architecture
All our experiments use DeepSolid (Li et al., 2022) as tho the baseline unsymmetrized network. DeepSolid adapts FermiNet, molecular neural network ansatz (Pfau et al., 2020), to infinite periodic solids. We briefly review their architectures.
rrchtecttur-wwis FermiNet takes as inputs n electron positions and split them into two groups according to spins. For some distance function d, the electrons are encoded through feature maps of the form d(x, - .x)) and d(x, - TT), i.c electron- electron interactions and deectrnnnuuleen interactions, which are passed into fully-connected network with tanh activations.
13

---

Diagonal Symmetrization of Neural Network Solvers for the Many-Electron Schrödinger Equation
The outputs are orbital information the single-electron level that depends on electron-clectron interactions, which combined by weighted of Slater determinants form antisymmetric many-electron wavefunctions. For details, Fig. and Algorithm of (Pfau et al., 2020).
DeepSolid makes several adaptations of FermiNet that are crucial for modelling periodic solids. Of those, the most relevant one to our discussion are their choice of periodic features d(ay - mx) and d(zz - TK k). stated in (2) of (Pfau et al., 2020). Their periodic features introduce separate invariance with respect to supercell translations, such that the n electrons are effectively restricted within the supercell. The construction is based on Jastrow correlation factors (Whitehead al., 2016) and can be viewed version of smoothed canonicalization for 1d translations (Section 4.3).
B.2. Physical systems
Physical systems
Graphene 1
LiH 2 x
boc-Li x
Number electrons
12
Point group symmetry the system
P 6mm
Fm3m
Im 3m
Number of point group elements
12
192
96
Lattice vector length (conventional)
2.4612
4.02
3.4268 Á
Table List of physical systems considered in the experiments, where the number x indicates the size of the supercell for system and a x 5 x e is that for 3d system. The primitive cell (corresponding to x x supercell) of graphene is visualized in Fig. (d)) and those of LiH and bco-Li are in Fig. 6(a) and Fig. 6(c).
18 3
11
14
7.0
7.06
7.8g
5.7
5.7
3.5
1.9
0:0
1,9
3.8
5.7
1.9
3.8
5.7
1.9
3.8
5.7
(b) LiH x with atoms from
(a)
primitive
(c) symmetry of LiH
1
adjacent
cubic structure
12
7
6.4
4.9
3.8
3.2
5:2
1.6
0.6
5.7
7.6
3.2
4.9
4.9
(f) bec-Li 1 with atoms from
(d) P2/m symmetry LIH 1 x
(e) bec-Li 1 x primitive cell
adjacent cells show the cubic structure
Figure Visualization primitive cells LiH and bec-Li.
The experimental benchmark used for comparing symmetrization strategies is graphene with an supercell for
14

---

Diagonal Symmetrization of Neural Network Solvers for the Many-Electron Schrüdinger Equation
computational convenience. Effects of post-hoc averaging over different choices of g are illustrated through LiH and bcc-Li, which possess substantially more symmetries. Two types of simple symmetries in LiH, P1 and P2/m, are illustrated in Fig. 6(c) and (d); we refer readers to (Brock et al., 2016) for an exhaustive visualization of all 2d and 3d space groups.
B.3. Parameter settings
The symmetrization strategies consider, i.e. DA, GA, SC, PA and PC, are all implemented wrapper functions around the original DeepSolid model. Architecturally this can be viewed inserting layer each before and after the DeepSolid processing pipeline. Notably, this allows us to leave most of the training parameters from DeepSolid unchanged and retain the performance from the original DeepSolid.
To ensure a fair comparison, we keep the default network and training parameter settings from DeepSolid across all our experiments, except that we vary the training batch size according to the number of symmetry operations used. Graphene training uses NVIDIA GeForce RTX 2080 Ti (12GB) and LiH and bec-Li use NVIDIA A100 SXM4 (40GB).
At the inference stage, we collect samples from independent MCMC chains with length 30, 000. The numbers of samples collected are respectively 50, 000 for graphene and 20, 000 for LiH and bec-Li. The choice of smoothing in post-hoc canonicalization is &oo defined in Appendix E.1; see Appendix E for details.
B.4. Additional results
A remark about the scale of energy improvements. Notice that the energy improvements reported in Table 2 are in the third decimals in Hartree. Improvements at this scale are crucial in physics and chemistry. In the physical systems we consider, core electron binding energies are on the order of keV (1 keV = 37 Hartree), while valence electron binding energies are in the range of 10 eV (Rs 0.037 Hartree). To obtain wavefunetions with accurate ground state energy, therefore crucial obtain energy improvements on the order of 10-3 Hartree or smaller. In fact, the term "chemical accuracy" used in computational chemistry describe the level of precision needed for calculated energies provide meaningful predictions of chemical phenomena, is given by 1 kcal/ / mol (= 0.00159 Hartree). For more references, Williams et al. (2020); Perdew et al. (1996).
Varying number of subsamples in GAs. In Table 2 and Fig. 5, group-averaging with subsampling (GAs) achieves the closest performance to post hoc averaging on graphene 1. In Fig. below, we verify that the performance does not improve if one varies k; the number of subsampled group elements, and that a larger also leads to gradient destabilization.
x10-4
1.8
-76.020
GAS, 1000, k=
GAs, 50000,
GAs, =333,
1.6
-76.025
1.4
-76.030
-76.035
4
-76.040
OG, N 1000
3
GAs, N 500, 2
GAS, N =0000 1=
-76.045
0.6
GAs, N 333, k=
0.4
-76.050
50
100
150
200
250
100
200
300
100 200 300
(a) + [ar]]] normalized variance gradient updates
(b) Energy (Ha)
(c) Var Etocall (Ha?)
Figure Performance of GAs plotted against GPU hours and across different subsample size k.
Averaging diagonal translations. The results in Table 2 focus averaging over point groups, and may ask whether averaging additionally over translations helps. We perform a preliminary investigation in LiH and bec-Li, and find that incorporating translations does not lead to significant performance improvement. The results are indicated in Table 4 by xT, and we include the corresponding results for point groups without translations for comparison. Note that as averaging over translations incurs 8x the computational cost, we have computed the statistics only on 10, 000 samples.
Averaging over subsets of group elements. Fact 2.1 ensures the validity of averaging over any subset of group elements
15

---

Diagonal Symmetrization of Neural Network Solvers for the Many- Schrüdinger Equation
Sysiem
Method
NV
k
Steps
GPU hours
Energy (Ha)
Var Blorall (Ha
0.06(1)
LiH
Pm.mm Fm X
4000
48
30, 000
571
8.1502(7)
0.0122(7)
2 x
384
-8.1488(7)
0.012(1)
15.01T(1)
0.0599[2
boe-Li
94/mam
PA
3000
16
20, 000
462
15 021(2)
0.033(2)
2 x
15.017(6)
0.05(2)
Table Performance of post hoc averaging with translations.
and not just subgroups. The only caveat is that for general finite subset G, the average [9] EGEn (bg(2)) is no longer guaranteed to be invariant under g. Nevertheless, one may still ask averaging over subsets of a group G that are "sufficiently representative", c.g. the generators the group G, can helpful. We perform a preliminary investigation on graphene and LiH for PA with Gen(G), fixed set of generators of plus the identity element. The statistics are computed 40, 000 samples and reported in Table 5. The results are inconclusive: We find that for graphene, PA with Gen(G) improves energy but significantly inflates variance, whereas for LiH, PA with GG) has worse energy and variance. As sanity check, we also compute Var[PA G /vv] for each wavefunction to verify that the PA with Gen(G) is closer to PA computed on G compared to the original wavefunction.
System
Method OG
Steps
GPU hours
Energy (Ha)
Bauaal (Ha-)
Par//// 1.80(4) 10
-76.039(6)
2.02(3)
Graphene L 1
P.6m
PA
1000
12
80, 000
281
- 76 (50(3)
0.33(1)
0.0
Gen( Grm)
PA
768006((5)
1.04(2)
3.3(1) x 10 4
0G
-8. 138(2)
0.06(1)
2.65(5) 10
LaF
F-222 G22(
PA
16
30, 000
8,1495(9)
0.0162(7)
0.0
2 x
PA
5
8.1456(7)
0.0235(5)
6.6(1) x 10-3
Table 5. Performance of post hoc averaging with subsets of group elements.
C. Additional details on Lemma 6.1 and visualizing diagonal invariance
We first state the proof of Lemma 6.1, which is illustrative for understanding the necessity of using modified group G = {A E 3x3 A(+) ) + b e & for some b 3R
Proof of Lemma 6.1. By the definition of
f(g(t)) f(xyymm + A(t)) + b)
(a)
f(A(Xeymm) + A(f) b)
( f(A(Xyymm +t)+b) = f(g(x + t))
In (a), we have used that Syym is invariant under A E G. Using the definition of again to note that f(t) = f(x + +)) finishes the proof.
As mentioned in Section 6, & and the point group G, of G are in general two different groups. When G. is the P 6mmn point group, both G, and G consist of 12 elements. However, G, is generated by P 331 and the reflection R indicated in Fig. 8, whereas G is generated by P3m1 and the reflection It indicated in Fig. 8. Applying our method to C allows for visualizing the full 6mm symmetry of the wavefunction (Fig. 9).
R
R
4.0
3.0
(Bohr)
2.0
1.0
-0.0
-2.3
-0.6,
1.2
2.9
4.7
(Bohr)
Figure 8. A G-symmetric configuration of 12 electrons, where C is obtained from P 6mm
16

---

Diagonal Symmetrization of Neural Network Solvers for the Mlay-Eleetron Schrüdinger Equation
3
2
3
y-displacement
0-
2
displacement
2 displacerent
(a) log (1% 2 original DeepSolid
(b) log (
(PA:G) post hoc averaging
Figure 9. Visualization of the full 6mn-diagonal invariance of 1/ (OG) versus 1/ (PA,S) Same setup and wavefunctions as Fig. 1, except that the configuration to be translated, Xsymm, given by Fig. 8.
D. CLT results for DA and GA
We present results that characterize the distributions of the parameter update under data augmentation and group-averaging First focus the parameter update under DA:
(DA)
N EigNIK EIENIKESK (X.)(X
where X X,,..., XN/k
Lid. (m) ~ Par and gi....... BNyk, are i.i.d. samples from some distribution on C.
Notice that due to augmentations, (DA) involves correlated sum. Nevertheless, for each / p. we can re-express the I-th coordinate of the DA parameter update
50 50(DA)
(AA)
1 E N/k
EiSNjE FR
(A)
where Fdl (DA) 1 E/= (FR1,((0)))0)1
an empirical average of i.i.d. univariate random variables across 1 S < N/k, which allows for the application of a CLT. A coordinate-wise CLT, i.e. the normal approximation 50 for any fixed S p. is straightforward from classical CLT results. Since we are concerned about stability of the gradient estimate, it is more crucial to study the behavior of the coordinate of maximum deviation, i.e.
(80( EEj0
maxigp
and verify that its behavior is completely described by the mean and the variance. The next result complements Proposition 4.1 by providing CLT results for both individual coordinates and the coordinates of maximum deviation in the high-dimensional regime, where parameter dimension - the number of weights in our neural network - is allowed to be much larger than the batch size N/k - the number of Markov chains per training step. For convenience, we denote the standard deviation of F (DA)
(DA) = V Var (P(A)
Var E si.. (X1):w))) )c]
1)Cov((P (X,) bwoo (X,) (P (X)vvo)))
+
17

---

Diagonal Symmetrization of Neural Network Solvers for the Many-Electron Schrüdinger Equation
Theorem D.I. Fix 0 E IR9 Let Z -N(0, be a standard Gaussian vector and Z be its I-th coordinate. Then there exists some absolute constant C1 such that for every 1 p.
G GRIADAA
P( ((80(DA)
5 1) ( + (Varl66 Zi s s1)
sup ER
S
v N/R (of
Moreover, assume a mild tail condition that for every V' S p with oi ((DA) O, we have
(a(DA))-111) -
] 2 where F(DA)
(ff ( (1 1/99
E
exp
max ISp with (DA) >
max (E(3.4)
P(DA)
Then the coordinate of maximum deviation of 0 ((DA)
also satisfies a CLT: There is some absolute constant C2 > 0 such that
sup 1E3
IP
max 1Sp
(00
B[ s -P(
max
(Var(60!DA)))) Ze s s SO2 ISP
(F(DAI)2 (1gg(NN/))) 1/6
N/k
The bounds in Theorem D.1 each control difference in distribution function between 50(DA) and normal distribution, measured through either arbitrarily fixed coordinate or the coordinate of maximum deviation from its mean. In particular, they say that the distribution of (DA) is approximately normal and therefore completely characterized by the mean and the variance studied in Proposition 4.1. To interpret them in details:
The coordinate- wise bound follows from the classical Berry- Theorem (see e.g. Theorem 3.7 of (Chen et al., 201 1)), and the normal approximation error does not involve the dimension p.
For the coordinate of maximum deviation, the normal approximation error is small so long as p is much smaller than a constant multiple of exp((N/A))")), which in particular true even if the parameter dimension p larger than the batch size N/k. The moment assumption amounts to a light-tailed condition on the distribution of the updates, and they be relaxed at the cost of more complicated bounds- see Chernozhukov et al. (2017).
We also remark that Theorem D.1 is a special case of the universality result of Huang et al. (2022) for data augmentation, but with sharper bound and in Kolmogorov distance.
Analogous CLTs hold for both the unaugmented update and the group-averaging update. Recall that these updates are 1
(0!OG) =
1 N
EiSN Fx,.V
and
(g(GA) = N/K EieNyR v
where X ..... X N/R Pog transformation, and (G(GA) is the same as 80()G) with N replaced by N/k and Pwo replaced by Pig Therefore the CLT
iid. (m) In particular, (0()G) is the same as (DA) with h = 1 and gı,1 set to the identity (m)
results for (B(GG) and 50(GA) are direct consequences of Theorem D.1 by defining 50(0G) F F(GG), (.(GA) of ((GA) and F(GA) as the analogous quantities. We do not state these results for brevity.
" (OG) f(OG), 50
E. Smoothed canonicalization
This appendix expands on the discussion in Section 4.3 and Section 5. We restrict our attention to a diagonal group G induced by a space group G Gap-
The essential idea behind canonicalization that, to construct an invariant layer under Gaizg. one may perform a "projection" onto a "smallest invariant set" II c R called the fundamental region. Formally, II is a fixed set that contains exactly one point from each orbit of Gaing called the representative point of the orbit and canonicalization is an R - II map that brings an input z to its corresponding representative. In the simple case of unit translations and n 1, standard canonicalization is æ + x mod 1 with II |0. 1), visualized in Fig. 4. A naive canonicalization suffers from boundary discontinuity that violates physical constraints: In the case of mod, this arises because
lim (emod 1) 0 a 1 lim (odd 1).
(8)
18

---

Diagonal Symmetrization of Neural Network Solvers for the Many-Electron Schrödinger Equation
Resolving this requires introducing smoothing the boundary. As an example, consider the supercell assumption (Rajagopal al., 1995; Kittel & McEuen, 2018), which corresponds separate invariance under translations along each 1d lattice vector, Jastrow factors ((Whitehead et al., 2016); also see (Li et al., 2022) for its implementation in DeepSolid) exactly perform smoothed version of canonicalization along each 1d lattice vector: These factors take the form ha mod 1), where z E R is a suitably rescaled input along one lattice vector and h is a piecewise polynomial chosen such that æ > h(a mod 1) is twice continuously differentiable. However, the construction of is easy here only because the fundamental regions of Id translations have simple boundaries (two endpoints), and does not generalize well to more complicated symmetries.
For groups such as permutations and rotations, Dym, Lawrence, and Siegel (2024) proposes smoothing method by taking weighted averages at the boundary. We adapt this method for our Guiag induced by space group, and show how can be achieved by operating with the group G acting on single electrons. This requires two ingredients:
Fundamental region of Gaiag There are infinitely many choices of fundamental regions, and we shall fix one that is convenient for our canonicalization. Fix 1o C R a fundamental region of IR under G. Then Il == 1l x (I²))=11 forms a fundamental region under Gaig- as it is intersected exactly once by any orbit {(9(2))........ g(x))) ) g E G}. An advantage of this choice is that, since the boundary of TT is completely determined by the boundary of Np the space of the first electron, we may perform smoothing directly near the boundary of Пy. Also notice that we may equivalently choose the fundamental region [] (IR x Ho (R*)"=k-1 for any 1 < k S
Smoothing factor. Fix some small 0. For each x E R3, consider the set of group elements that moves r to be at most c-away from Ng:
(e(x) to G G d(o(z), II) s sc,
where is some distance preserved by the isometries in G:
d(g(z).g(Tn)) = d(x, IT)) for all E G and x E R
Adapting the weighted averaging idea from Dym al. (2024), we perform smoothing in the e-neighborhood of 1lg by assigning weights to the different group operations: For each g E G, we define
Ax(dl((a), Io)) 2 Eg'eal (d)
[0,1].
where A₂ R + [0, 1] is strictly decreasing function such that (w) for all w S 0 and >e(w) 0 for all 2 e. If d(g(2). II) 2 c, g moves "too far", and is assigned zero weights. If (l(g(z)) (o) = 0, g keeps within the closure of Io,, and is assigned weight of 1. In particular, if II is point of high symmetry and is small enough such that g(x) = for all g (x), the weight of every g (xx is |. Here, all operations that preserve are assigned equal weights. See Appendix E. for explicit choices d and X
With these tools at hand, we define the smoothed canonicalization (SC) of an unsymmetrized wavefunction fo
(SC) Vec
(x)
Σ
(SC:k) 8'
(x)
where
(SC.)) 8
(x)
Eyes.( wx(xx) vo(g(x))
(9)
Appendix E.2 illustrates (9) through one-electron example. The next result guarantees the validity of ve (x)
(SC)
(SC.k) (x) symmetrized wavefunction: Theorem E.) proves diagonal invariance, (ii) proves anti-symmetry and (iii) ensures smoothness under appropriate choices of X and d.
R k oo
(SC) Theorem E.1. Fix 0 E R2 and e > 0. The following properties hold for oo and (SC.) for all 1 s k S n:
t0
(i) te (SCM)) (g(x)) =
- (SC:A) (x) for all g E G and x E 3
(ii) Let the dpin-depen in tho be explicit, i.e. we view the as an (R3 x (t.4))" + function. If the is anti-symmetric with respect to permutations of (xi, σ,), then so is vo
(SC).
19

---

Diagonal Symmetrization of Neural Network Solvers for the Many-Electron Schrüdinger Equation
1.00
8 dsg
10
if 800
2
0.75
5
8152
0.50
1
0.25
900
-5
0.00
0
110
0.0
0.5
1.0
0.0
0.5
1.0
0.0
0.5
1.0
1.00
01
1000
0.75
500
0.50
-10
0.25
A1000
87,00
-500
8 Axpoo
2,
-20
aA,,
8 Nr12
0.00
10000
0.11
0.0
0.1
0.2
0.3
-0.1
0.0
0.1
0.2
0.3
-0.1
0.0
0.1
0.2
0.3
6
W
Figure 10. Illustration of different choices of step functions s, the corresponding A. 's and their first two derivatives. = 0.1 above.
(iii) Suppose de and d(., ((1o)) are p-times continuously differentiable for all G, and that tho is p-times continuously differentiable at g(x) E 3 for all g E G. Then V is also p-times continuously differentiable at x. Moreover
S <9s S p, the q-th derivative of Voc at x (315..., can be computed as
(SC)
(SC) V
(x1: ... Pn) = - k= EgEg(((w)
4 Vi (SCck)
8,ss (x)
where have defined, for 1 k g G and 21,1, ,,,...,, ,n R3.
(SC:k) Vo (y), (un,...,yu)) ,Bn) = de (d(yv.. ((To))) Egres de (d(yk. ((ll) Volo (in), g (yn))
Computational cost compared to DA and GA. As mentioned in Section 4.3, SC suffers from similar computational issue as DA and GA as it involves averaging. However, the amount of averaging SC requires is different. The additional averaging over n fundamental regions is crucial for ensuring anti-symmetry, which is made precise in Theorem E. and its proof in Appendix H. Averaging is also required over the set Gc, which necessarily includes all point group operations and unit translations of G (note that G. is not group). For a fair comparison, compare SC DA and GA that average over the group G. For small systems with 1 x x 1 supercell, G, is strictly larger than the point group G, and the amount of averaging required by SC, na G. is strictly larger. For large systems, the cost of SC scales differently from DA and GA: While the cost of DA and GA increase as a function of the increasing group size |G], for SC, G. is fixed and the cost increases function of the number of electrons n.
Post hoc smoothed canonicalization (PC), As discussed in Section 5, we may perform canonicalization only at evaluation. However, there is an inherent tradeoff in the choice of If is large, averaging happens in a large neighborhood of the boundary, which can cause the performance of t/ d' If e is small, there may be a blowup in the derivatives of the weights via the smoothing function:
(SC) to deviate significantly from the well-traines w/ (,(0)) (see Appendix E.2).
Lemma E.2. If X twice continuously differentiable, there are y1 y22 E [0,e] s.I. (v) = -1, 02 g2 (y2) = -2
Since the local energy calculation involves 02 Xe, this can lead to an 2 blowup of local energy in region of size O resulting in larger fluctuations of local energy. This inflation in the variance clearly visible in Table 2.
E. Choices of X and
To obtain a p-times continuously differentiable (/ Theorem E. (iii) requires a smoothing function X and a distance function d( d((,90 (To)) that are p-times continuous differentiable for all g E G.
20

---

Diagonal Symmetrization of Neural Network Solvers for the Many-Electron Schrüdinger Equation
Construction of Ag- Since de I [0, 1] required to be strictly decreasing with >x(w)) 1 for w and X(w) 0 for all w 2 e, the function
s(w) = (e(111 (1-(ww) w))
is p-times continuously differentiable approximation of step function: is strictly increasing, s(w) = 0 for all $ 0
and s(w) 1 for all w 2 1. Many choices are available for such a function: To have an infinitely differentiable Ae one may consider the smoothed step function
(w)
where <(w) = (exp(-w-))
w > 0
850 (w) = ((ar) + +0(1- - w)
if w S 0
and use the relationship de(w)) = s(1 ww/c to obtain
6(1 w//c)
Aepo (w)
(1 w/e) 6(w/e)
As we only need to evaluate the Hamiltonian, twice continuous differentiability typically suffices for our problem. Another choice of the step function s and the corresponding de are the degree-three polynomial splines
if w < 0
1
if w S 0,
if E (0, g] if E (g.3),
1
if w E (0, gl if w E (g 31, if E 2cc .]
2
2
(2
+
3 2
+1
(e2(w) =
9w
+
(3w-e) 2
+1
2c
911 w1-
+1
if E (2. wf > 1,
9(e ((--1))
2r3
if
Fig. 10 plots S $2, Aeços Ae and their derivatives. Note that the derivative plots for X 's verify Lemma E.2. To achieve general p-times continuous differentiability, we refer readers to the many constructions of such step functions available in the statistics literature, c.g. Theorem 3.3 of Chen et al. (2011) or Lemma 34 of Huang et al. (2023). In particular, our s2 is related to hmur; constructed in Lemma 34 of Huang et al. (2023) via $2(w) = h2:1;1 (w); one can similarly take their hg;1;1 to obtain a p-times continuously differentiable sp and obtain the corresponding Xe:( (w) == 1 sp(w/c).
Constructing d(- d(O,g( ((اا)) via the step function s. A p-times continuously differentiable step function s, discussed above, induces a p-times continuously differentiable approximation of max{ +, 0), given as
s(w) w swr)
In partieular sww) if and only if S 0, and ((w) for 1. To utilize ă construct d( ((( ".g.ggg(n (Ta)), notice that each g(11)) itself is fundamental region, and its closure is simplex completely characterized by the relation
(x- ( < 1 for 1 < m
x E ((lo)
|m 1
where co is the center of Tog and ny is the normal vector starting from co that is normal to the I-th face of Tl. This allows us to define
da. (To)) = Em (s((5- g(co))
1))²
||
This distance function satisfies that
d(g(x), g(Ln)) = d(x, Ho) for all g E G, which can be verified by noting that g( = =A(+) + b for an orthogonal matrix A E R3x3 and a translation vector b E TR3
d(z, I1) 0 if and only if z II, and the map + d(z, 11) is p-times continuously differentiable.
To visualize this distance function in the case of 1d system, see Fig. 14.
21

---

Diagonal Symmetrization of Neural Network Solvers for the Many-Electron Schrüdinger Equation
0.001
0.1
0.5
1,00
0.75
0.50
0.25
0.00
0.75
(2)
Roc
0.50
0.00
1.00
0.75
(a +11
0.50
0.25
0.00
0.0
1.5
-0.5
1.0
-0.5
0.0
4
Figure 11. Plots of (0 1), we (x) and (x 1) defined via either 8ec or 82 the smonthed step function used in A and and for different values of c.
E.2. An one-electron example of Vile
(SC)
For simplicity, we first illustrate vi;; group G is generated by translations length I, and the fundamental region is II, = [0, 1), as illustrated Fig. 13. As discussed at (8), standard canonicalization to take æ + x mod 1, which suffers from discontinuity at the boundary, Our proposed smoothed canonicalization this case becomes
(SC) for single Id-electron system with unit translation invariance. In this case, the
(SC)
(x) = ELEZsr d(a++.(0.1)) ErEZ (1,1)) Ar (d(e &1,
A, + t, (0, 1)))
de(r + +)
Pie
Choosing the distance function d( [0, 1)) as Appendix E.1, we have (0, = 3(2(z 1)) + 3(-2x)
where
((w) = ws(w)
and is a p-times differentiable approximation of step function. Fig. 14 plots d(z, [0, 1)) and its derivatives under s = 8.0 and $2 from Appendix 1, and verifies that d(x: [0, 1)) twice continuously differentiable under either case.
Notice by construction that (SC) is 1-periodic (also see Theorem E. for the proof in the general case), so it suffices to consider the value of (S(C) (z) within the interval æ [0, Clearly d(x,, [0, 1)) 0. By construction of the distance function, the only possible translations t with non-zero (unnormalized) weight, defined as
(x + t) = (e t, (0.11))
t +1. In other words, for [0, 1],
express
(SC) V0:c
(x)
6x vol 1) +
(x(x
vu(x)) + Ce( 1 Se(x)
the (x 1)
(c(c)
Sc(x)
22

---

Gradient blowup as decreases. Lemma E.2 and Fig. 10 both show that the derivatives of X blow up as e gets small. This Diagonal Symmetrization of Neural Network Solvers for the Many-Electron Schrüdinger Equation
0.001
0.1
0.5
100
8o
8rc
20-
0
0-
0
20
-10
1000
60000
3000
500
100
#2
2000
#2
#2
i (3r)
200
20000
1000
0
0-
11000
200000
400
0.5
0.0
0.5
1.0
1.5
-0.5
0.0
0.5
1.0
1.5
-0.5
0.0
0.5
1.0
1.5
a
x
a
Figure 12. Plots of derivatives of i. (x) for different values of e.
1
-1
1 --
21
2
3
To
Figure 13. Illustration of smoothed canonicalization (o..)
(SC)
for a single 1d electron
where we have denoted the normalizing term S,(x) u. - T. (x) + (x 1). Fig. 11 plots the values of the three unnormalized weights as varies, under different choices of c. Several observations follow:
For æ E [0, 1], we always have (x)) = 1, whereas we (ar - 1) is non-zero only for æ close to 1, and We w. (x non-zero only for close to O. In particular, vi (1(5)) (x) is exactly the original wavefunction the (zr when is far away from the edge of the interval, whereas an weighted average is taken with v + if a is close to 0, and with vo(x 1) if is close to 1.
+ 1) is
r = 0 and r = 1, (SC) (x) ) is exactly the arithmetic mean of (0) and (1). Indeed, smooth canonicalization D:
allows the wavefunction to smoothly interpolate to an average exactly at the boundary of the fundamental region.
(SC) Effect on the performance of vo as increases. Compare the exact canonicalization va mod 1) with (SC), In the V former case, the wavefunction value depends only on the values of the within the fundamental region II [0, 1). whereas in the latter case, vire depends on dhe evaluated on a slightly larger set {æ E R d(z,, [o, 1)) $ <), represented as the interval [x=e, +c] in both Fig. 13 and Fig. 11. Meanwhile, (x mod 1) = (ar) for all E (0, 1), whereas v/ ((SC) (zx = vo(#))
(SC)
only for slightly smaller set E [xr,, -e] labelled in Fig. 13 and Fig. 11. Notably for post hoc canonicalization, as increases, the performance of // (SC) relies on the to be a well-trained on larger and larger region, and t only recovers the trained vie on a smaller and smaller region, thus retaining fewer benefits from training.
(SC)
also leads to a blowup in derivatives of (.(x) and hence those of e (SC) near the boundary. Since the derivatives of w (.(x) involve products of derivatives of X₂ and derivatives
by the chain rule, the difference in the magnitude of blowup across
s = 82 and æ = 3o is more pronounced, as illustrated in Fig. 12: defining the weight function w. via 8 leads to a higher degree of smoothness, but at the cost of a larger blowup in the derivative.
23

---

Diagonal Symmetrization of Neural Network Solvers for the Many-Electron Schrödinger Equation
d(. [0, 1))
0,d(2, [0. 1))
Rdd(r, [0, 1))
62 -2.
-25
Figure 14. Plots of d(z,, [0, 1)) and its derivatives under different choices of the step function
F. Proof of Fact 2.1
Fix E G with its action on x E I represented by g(x) = Ar + b. Given that (v), E) solves (I), we seek to show that under the stated conditions, vi (x) = (x((x)) is also an snti-symmetric solution to (1) with respect to the same energy E. For E R and g E G, denote Xg = g(x). Then for every x E RR3"
(-1 V2 + V(x)) (x(x)) = (-2 v2 v22 + v(x)) 2s (AA1 + b....., Ar + b)
(A A) (x,,) + (xx))(xx,)
(o) (x))) + ((x)) @ Exb(xy) = Ex)y(x).
In (a), we have used the Gdin--nvvriiane of V; in (b), we used that (w. E) solve (1). Since the above holds for all x E R3"", we get that (zh., E) also solves the Schrödinger's ' equation. To verify the anti-symmetric requirement, we write the wavefunetion as (xx)= = (........ ),). where each é; (I; o.) now additionally depends on the spin o, (t:+4)) Since G only acts the spatial position in R and leaves the spins invariant, we can WLOG express the g-transformed version of i ( (g(x)), o,). Moreover, since E P. commutes with the diagonal action of g G on (R3 x {1)"
(x(())) = ((goo()) = v6 o gXx) sgn(a) (b((())) = sgn(o) (x))),
In (c), we have used the anti-symmetry of sh. This proves that (v)g, E) solves (1) with the correct anti-symmetric requirement. To prove the theorem statement, we see that & is a linear combination of finitely many eigenfunctions (Vy))eg with the same eigenvalue E and therefore yields an anti-symmetric solution with the same energy E.
G. Proofs for data augmentation and group-averaging
G.1. Proof of Proposition 4.1
We first compute the difference in expectation
[[()(DA)) |[
E N Eis iSN/K Ejst (x,s(X
-E Fxmeel)) I|l Eien
N
[E[ (X, (Xx)) hool B[FFxx;ollll
B [ M
(om)
[Fxcool E Ex (m)) [x
In the last line, we used linearity of expectation, the fact that gi,j (X:)'s (X.)( are identically distributed and that X, x'' are identically distributed. Recall that + (Fx:s. ) 0 E IRy) C F and dp(p,g) = supser [I Ey [x (fY)) This implies that
[Fxwo]] - E Ex [Fr 5 dr (p (mm)) -
X
24

---

Diagonal Symmetrization of Neural Network Solvers for the Many-Electree Schrödinger Equation
which proves the first bound. For the second bound, notice that
Var Var[6g(DA)]
Var
EiSNik Ejsh (x,
N/k L Eish (Xii(oe
Var
(6)
Varl 81, (X.))o)
gi, + 1- [o (X.): P (x))))))
Var[F 81,1 (X, (3))))))) k Var 8), 1x, (oo [1]
N
In (a), we have noted that the summands are i.i.d. across i < N/k; in (b), we have computed the variance of the sum explicitly by expanding the expectation of a double-sum; in (c), we have applied the law total covariance to obtain that
Cov[FRi.N ( F (x1.2(X1))0))
= Cov [
(
X1]
B[F R1.2 (X X1]] +E Cov [F
(X1)68 2x
X1
Var 21,1 (11 [X]]
Meanwhile, the same calculation with h = 1 and g1,1 replaced by identity gives
Var(5g(OG]] Var Fx
N
Taking difference and applying the triangle inequality twice, we have
Var (JA)) Varl50(OG))
Var
[
F.... (X.)
Xi]
N Varl (x. Var[Fx
E s2
E
P
|||
| (X) mnvo
182
R. [R
602
RI.1 (X, litto
N
S2
Ri.1 x Joobo
22
|||
+
N
$1,1 (:0k)
+ B Fx
|| (F [FF]))BB(Fxbbb [EXx)1)]]
(m)) d
Peg: Pora fm)
[2| Fx Ill dr (P )())))
(m)
(m) dy (P (mm)
N
||
(mm) + d Pig
x dy (p Peo:DA (m))
(10)
This proves the second bound. In the case when s(x,) X2 for all g Gaagg have (m)) bounds above evaluate to zero. In this case we have E[50(DA)] = E[j8((G)] and
(m)
Poo:DA -Po and therefore the
Var[5g(DA)) - Var [B0(0))]]
k N
Var 3-'VarB (F [x]
which is positive semi-definite.
G.2. Proof of Lemma 4.2
N/K EiSNjR Fx is a size-N/k empirical average of i.idd. quantities. The mean and variance formulas thus follows directly from standard computation:
By construction, (GA) 1
VarlF 10
E[80(GA)] =E,Fx Exp sg
and
Var 5a
N/k
G.3. Proof of Theorem D.I
To prove the coordinate-wise bound, fix < p. Note that if 01 (DA) = 0, then 50 1ES0 with probability 1, implying that the distribution difference zero and hence the bound satisfied. In the case of (DA) > 0, 50(DA)
(DA) is an average of i..d. univariate random variables with positive variance. By renormalizing and N/E Σ EigN/R
25

---

Diagonal Symmetrization of Neural Network Solvers for the Many-Electron Schrüdinger Equation
applying the Berry-Esséen theorem (see Theorem 3.7 of Chen al. (2011) Shevtsova (2013) for the version with tight
constant 1 0.469) applied to N7k Σ iSN/K Fut, we get that
2 <) P E [| + (Var(50 Za
s1)
sup rEB
(DA)
C1 BI N/k (o, (OA
sup tR N/K E NigN/K (1 - ((P(DAM)) s <1)-1(( -P VN/k
IP
Ze s1)
This proves the first set of bounds. To prove the second set of bounds, first denote the mean-zero variable Fur - F(DA)) + E[ and let (Z11.... Znp) be an R"r-valued Gaussian vector with the same mean and variance as (F (Fi1.... Pap). The difference distribution function can be re-expressed as
sup ER
IP
max PSP
80 ([
S t)
IP
max PSP
(DA())((A)))))) (DA) [Var 1) 1)112 Zi S st]
sup
P
max iSP
Σ
EigN/R Fa 1) IP
max ISP
Σ EigN/K
iSN/K Zu S 1)1
(I1)
IER
VNJK
NNjk
where we have used the definition of ((DA) and also replaced t by I/vNk. Note that Fat 's are i.i.d. mean-zero across for all - S p. the two random
1 1 i S N/k and Za's 's are i.i.d. mean-zero across i S N/k. As before, if oj (DDA) variables to be compared are both 0 with probability and the distributional difference above evaluates to zero. If there is least one such that of coordinates with zero variance. As such,
(DA) O, we can restrict both maxima above to over / S p such that (DA) 0 and ignore the
WLOG assume that o (DA) 0 for all S p. Now write
F (ai Fi...... ,......,07 o,
Z. = (0i Z,1 ................................ op (Zp)
and denote the eyperrrectangular set A(t) == [-o t, t, (1) 1- t, to [t] S Rd. We can now express the difference above further
(11) sup
P (N/R R
R REA - -P( (NNNT EisN/R 2 E ((t(())
VN/K
EE
This is a difference in distribution functions between a normalized empirical average of p-dimensional vectors with zero mean and identity covariance and standard Gaussian in IR? measured through a subset of hyperrectangles. In particular, this is the quantity controlled by Chernozhukov et al. (2017): Under the stated moment conditions and applying their Proposition 2.1, we have that for some absolute constant C2 > O,
(F(DA)} (log(pN/k))3 1/6
(11) I C2
N/k
H. Proofs for results on canonicalization
H.1. Proof of Theorem E.1
To prove (i), we fix any ğ E G, and WLOG let k = 1. Then by the definition of (SC:1)
Ve:
Vo. ()C:)(((x)).)))))))))))))))))))) (G(x:)....,gaa g((n)) =
Σ
Xe (d((((1), g(Ta)))
)=(d((()1). x we (z(1)) g-=".,(,,, =((xx)))
E d'EGsL
SEG ct.
d(g(x1)) (H ((
d(g(2)).9(1a))
Relahelling g by åg and of by åg' in the sums above, and noting that d(g(<1), g(To))((lo)) obtain that
V (SC:1) (f(21).....
Σ
A. (d(3(a))
d'eGu d( (o))s
(d(0(a1), (Ilo))
ge
d(0(#1)
x to 2 g0
-((x)))
Σ
{(d(2,,
volg v0 ... (xx))
Σ Mas (1a))s
d'eGs
g
gEGst.
d(
(SC:1) vec
(21,..,,,,)
26

---

Diagonal Symmetrization of Neural Network Solvers for the Many-Electron Schrüdinger Equation
This proves diagonal G-invariance.
To prove (ii), we shall make the spin-dependence explicit and write ã; = (xi, o.) and g(ā;) = (g(x)), oz). First consider ", a transposition that swaps the indices 1 and 2. Then
(SC) (Tx(1)) -..)...... Fx(x)) 1 k= Vo
((CCK) 03.,81,33....
By the anti-symmetry of the we see that for h: 2 3.
" 0: (i2, 31, .........a)) Fn) = E (xxx
(SCk))
(ix)) - -yge wl( (11), (t) (2)).... ,9 (in))
(SCAA) (SCC)) (fI
Vo:e
For the case 1,2, we can apply a similar calculation while noting that the weights remain unchanged, and obtain
Vo
(SC:1)
... (īz T1
[i2, 31 53 ......,,, Fn))
w/ (8),
(SC:2)
52, (,,, ................................ ..)
(
Vo. (. (5,, 52 5,)
This implies that
V0
(SC)
1 Σ k=1 Vo:
(SC:k)
02, ... ..)
sgn(T) (/ (5)) (3....,,3 5n) (
Since the choice of indices and 2 are arbitrary, the above in fact holds for all transpositions n, which implies
Dve (T+(1)...., (,()) = sgu(r) vo: (Z.....,ā5 in)
(SC)
(SC)
for all permutations 7 of the n electrons. This proves (ii).
To prove (iii), suffices to show that for every k S n, the function (SC:k) defined above p-times continuously Vee
differentiable at x, Lee V exists and is continuous at x. Again it suffices to show this for the case = 1. Let = (Z1...., F.) be vector in sufficiently small neighborhood of a fixed x = ( and recall that
(.)) = {gg E G ((ā). S sc}.
For x in sufficiently small neighborhood of x, Ge (ū)) takes value in [G where
Gi = (x1)) \ gं
and .f. Phs G is an enumeration of all group elements such that for 0 OS1S M.
d(m1 (Io))
Therefore ti
(SC:1) (x) takes values in (1hg, (x) Jos where de ( (
1h. (x: Eges, Eyrea de t ((
g ( ... g (ix))
Notice that at x = x, by the definition of of. we have
(.(((r) X th vo (x1), 9 (xn))
dig, (x) = Ey Egeea do (d(r), g' (To)))
Egea, (x1) Eyeanto de (d(z). ((o))((mo)))) (1o)))
X (d(z),
voolg (009-1((1)) .... g- (xx))
vi (x)
(SC:1)
27

---

Diagonal Symmetrization of Neural Network Solvers for the Many-Electron Schrödinger Equation
Since the above argument works with replaced by x, also have
V0:e
(SC:1(())
Eooa, Eyeo, A g(((( (0 ..... -x
de
(12)
for 0 I I M, where we have defined
G = g((1)
\ 0i .f.
and Øi. .i....., Fhr E G is an enumeration of all group elements such that d(i 0i (To)) = for O
Notice that for x in a sufficiently small neighborhood of x, we have (rdig = (9iheM, in which case (12) implies
(SC:)) V0::
(x) vic (x) for all 0 S M.
In other words. have shown that in sufficiently small neighborhood of x, F1 equals ,g for all 1 S / S M. Recall that A and d( g(To)) are p-times continuously differentiable for all G, and tho is p-times continuously differentiable at g(x) for all g E G by assumption. This implies that ,.g is also p-times continuously differentiable and v Moreover, for 0 S q S p, the derivative can be computed as
(SC:1)
4 Vo (o
(SC:1)
(x) ('))a(()) Σ (E))
V (SCcc) Vig.en (x).
where we recall that for k S n, n, G and x, yı.... Un R
kSC;k)
de (d(ua (((o)) (V1 U1...... Un) Egesset) de (d(v.. ((o)) voel (e(gg'((1) (v1), .... g (yn))
The same argument applies to all Vo 's with & I n and therefore for 0 9 P.
(SC
V 0:c (z), (ZI,...,,
(SC) xxn) 1 n 2==
(SC:)) (I (T,) = >
7 th D.egg. (x).
(SC:A)
which proves (iii).
H.2. Proof of Lemma E.2
Since (0) 1, (e) 0 and A is continuously differentiable, by the mean value theorem, there is y1 (0, such that
(v)) = (1 - 0)/c = c
Mcanwhile, since (w) for all < 0, (w) for all w and since aλ, is continuous, (A, By the twice continuous differentiability of de and the mean value theorem again, there exists of (0, C (0, such that
8 xr,,(1h) = (e 0)/v1 =
1 ev
1 2 z
Since (w) = 0 for all w < 0 and Xe is continuous, (2) = 0. As (2)c(vi) 2 2 by the intermediate value theorem, there exists some 1/2 E [0, 31] S [0, c] such that 82> -
I. Discussion on the case where per-iter sampling cost is greater than per-iter gradient cost
An anonymous reviewer has pointed out regimes where, unlike Section 3, Csamp may take larger values than Cgrad- We briefly discuss this point. There are two cases where Csamp > Cgrad may occur: (i) There exists significantly accelerated implementations of gradient computation (Li et al., 2024) such that it becomes more efficient than sampling; (ii) One typically needs to scale the number of sampling steps and neural network size both with the physical system size (von Glehn et al., 2023; Li et al., 2022), and the increase in per-step sampling cost could outweigh the increase in per-step gradient cost. In these regimes, our theoretical claims do hold, but come with some additional nuances:
28

---

Diagonal Symmetrization of Neural Network Solvers for the Many-Electron Schrüdinger Equation
DA instability. When Ceamp & Cgrad- our DA batch size in Section 4. chosen k x N' N, where N' == N/k is the number of i.i.d. drawn samples and k is the number of augmentations, The reduction in the number of i..dd. drawn samples was key to the destabilization effect. If instead Camp > Cgrad. may increase N,, the number of i.i.d. samples drawn, though the same conclusion holds as long as N' < N. Meanwhile if N = N, DA does not destabilize, but the k augmentations always increase the computational cost. The andistical-compuaatiinnn tradeoff now depends on how large N is and therefore has a more intricate dependence on the computational cost ratio /samp/ Cgrad.
GA instability.
The destabilization effect still occurs, as group-averaging always increases both sampling and gradient evaluation cost, which necessitates a reduction in the number of i.i.d. samples drawn.
PA benefits. PA is still computationally attractive compared to in-training symmetrizations, since the overall training cost still typically outweighs the inference cost.
Empirically, we do not observe Csamp >> Cgrad for the DeepSolid experiments we have performed.
29


Reviews:
[{"review_id": "u5A0aZdIAa", "reviewer": "Reviewer_mfuU", "summary": "This paper investigates methods for incorporating diagonal group symmetries into neural network wave function ansatze for solving the many-electron Schrödinger equation via Variational Monte Carlo (VMC). The authors compare three main approaches: data augmentation (DA), group averaging (GA) and canonicalization. The central claim is that, unlike typical machine learning scenarios, explicitly enforcing symmetry in-training can destabilize the optimization process and potentially lead to worse performance compared to post-hoc symmetrization. The methods are evaluated using the DeepSolid architecture on systems like H2, LiH, and graphene. The paper also introduces a method for visualizing diagonal symmetries in the high-dimensional wavefunction space.\n\n## Update after rebuttal\n\nThank you for the detailed rebuttal. Your clarifications on the broader context has convinced me about it's relevance to ICML. I am increasing my score accordingly and look forward to the revised manuscript.", "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": "We thank the reviewer for the positive and helpful feedback. We address the questions below.\n\n---\n1) **Relevance to other areas of physics simulation with similar symmetry constraints arise.** We agree that this would add value and is a very interesting avenue of future work. **In the revision, we will have a discussion and limitation section to discuss the relevance and comparisons with known cases of symmetries in other setups.** We do note that, as the other reviewers have also brought up comparisons with symmetrisation approaches in VMC for non-solid contexts (see response 1.2 to sZMY), we will prioritise those comparisons first and, if space permits, include a brief discussion on implications for other physics setups.\n---\n2) **Interest for the ICML community.**  We believe the general message of this paper, especially the part regarding the surprising effects of DA and GA in training, **is of broad interest to the ML community even beyond a physics setup**. While symmetry has become the staple in highly structured problems such as those in AI4Science, the observed benefits of symmetry in practice are often entangled with architectural changes, optimizer changes and hyperparamter choices. **Recent findings in protein structures [1] and atomic potential [2] showed that symmetries may actually be unnecessary for performance improvement**, calling into question how much symmetry helps with performance versus other factors e.g. hyperparameters. Our work can be viewed as **a first step towards this question** in the specific context of VMC for solids, as we perform **a strict apples-to-apples comparison on a fixed architecture with fixed optimizer and hyperparameters** and only vary the symmetrization methods. The computational-statistical tradeoffs we see for DA and GA are in fact applicable to any ML setups where sampling is performed in between gradient updates; one non-physics example in ML would be the contrastive divergence algorithm for training energy-based models (see e.g. [3] and [4]). **We will make sure this general applicability is mentioned in the revision.** \n---\n3) **Sensitivity of the results, e.g.~the observed instability, to the choice of architecture and hyperparameters.** We thank the reviewer for the interesting question. For varying batch sizes, we have included a preliminary ablation test in Appendix B.4. For the learning rate, we used the default setup from DeepSolid to ensure an apples-to-apples comparison across all setups. In view of our theoretical results, we believe that the instability findings and the computational-statistical tradeoffs are independent of these hyperparameters. **For architectures, we only used DeepSolid since it is the only architecture we know that performs well for VMC with infinite periodic solids, and because we want to perform an apples-to-apples comparison to an architecture that is known to do well in these problems.** We do conjecture that the effects of symmetrization may change in the context of other architectures. One example is the situation where, the system of interest possesses simpler symmetries and where more efficient symmetrizations exist (discussed in more detail in response 1.2 to reviewer szMY). Meanwhile, we agree that it is interesting to explore whether there exists a way to tweak DeepSolid in the ``most optimal way\" for DA or GA such that it outperforms PA. Although we feel that this is out of scope for the current paper, we do hope that our findings pave the way for these investigations: This is both by our discussions on the computational and statistical costs of DA and GA that one needs to be aware of when benchmarking the performances, and by offering tools such as the visualization method for understanding symmetry improvements.\n---\nWe hope these answer your questions. If our responses are helpful, we would be grateful if you would consider raising your score.\n\n---\n\n**References used for this response and other responses**\n\n[1] Abramson, Josh, et al. \"Accurate structure prediction of biomolecular interactions with AlphaFold 3.\" Nature 2024\n\n[2] Qu, E, and Aditi K. \"The importance of being scalable: Improving the speed and accuracy of neural network interatomic potentials across chemical domains.\" NeurIPS 2024\n\n[3] Hinton, G. E. \"Training products of experts by minimizing contrastive divergence.\" Neural computation 2002\n\n[4] Du, Y, et al. \"Improved contrastive divergence training of energy based models.\" ICML 2021\n\n[5] Chen, S., Edgar D., and Jane H. Lee. \"A group-theoretic framework for data augmentation.\" JMLR 2020\n\n[6] Lyle, C., et al. \"On the benefits of invariance in neural networks.\" arXiv:2005.00178\n\n[7] Benton, G., et al. \"Learning invariances in neural networks from training data.\" NeurIPS 2020\n\n[8] Yang, Jianke, et al. \"Generative adversarial symmetry discovery.\" ICML 2023\n\n[9] Balestriero, R, Leon B, and Yann L. \"The effects of regularization and data augmentation are class dependent.\" NeurIPS 2022", "claims_and_evidence": "The central claims are: in-training symmetrization (specifically group averaging) can destabilize VMC training and lead to poorer performance; post hoc averaging (PA) improves energy, variance, and symmetry effectively and efficiently. The claims are supported by both proofs and empirical results across different crystalline solids.", "methods_and_evaluation_criteria": "Yes, the methods and evaluation criteria are largely appropriate for the problem of developing symmetry-aware neural network solvers for quantum many-body problems.\n\n* The paper investigates symmetrization techniques: DA, GA, PA and PC\n* The authors use appropriate metrics (energy and variance in in local energy) and propose a new metric Var[PA/OG] to provide a quantitative measure for symmetry. They also use GPU hours as a measure of computational cost.\n* The models are benchmarked on 3 crystalline solids with DeepSolid as the baseline.", "theoretical_claims": "As a reviewer whose expertise lies in physics-informed ML for DFT and not specifically in VMC and associated ML methods, I have reviewed the statements of the theoretical results but have not rigorously checked the correctness of the proofs provided in the appendices. They seem plausible and the authors strengthen the claims with empirical results and have detailed proofs in appendices.", "experimental_designs_or_analyses": "As stated in methods and evaluation criteria, the experimental design and analyses are reasonable, with good baselines, datasets, metrics and ablation studies.", "relation_to_broader_scientific_literature": "The paper is well-written and establishes itself well in the context of NN solvers for ab-intio methods for the Schrödinger equation, symmetrization in machine learning and VMC methods.", "essential_references_not_discussed": "I cant definitively identify essential missing references within the core VMC or equivariant ML literature for diagonal groups. The paper appears to cite key works in neural network wavefunctions, general ML symmetrization and VMC.", "other_strengths_and_weaknesses": "The paper makes useful claims about properties of different symmetriziation approaches, including the counterintutive claim (from a standard ML viewpoint), that data augmentation can destablilize training. It provides a systematic comparison of these approaches applied during training and post-hoc. It is well-written and supports its main claims with both proofs and emperical results.\n\nWhile this is a well-written paper with detailed proofs and emperical results, I am concerned that it is too specialized for a venue like ICML and might be more suited to a physics journal.", "other_comments_or_suggestions": "1. The paper could benefit from a brief discussion of how the findings might apply to other areas of physics simulation where similar symmetry constraints arise.\n2. The visualization method in App C is a nice addition.", "questions_for_authors": "How sensitive are the results, particularly the ranking of the methods (OG, DA, GA, GAs, PA), to the choice of neural network architecture and hyperparameters like the learning rate and batch size? Could the instability observed with GA be tackled with careful hyperparameter tuning specific to that method?", "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": "GpMs1rKOJn", "reviewer": "Reviewer_szMY", "summary": "This paper investigated different methods for incorporating diagonal symmetries in neural network solvers for the many-electron Schrödinger equation, with a particular focus on variational Monte Carlo (VMC) methods. Specifically, the authors studied three main approaches to enforce diagonal invariance: data augmentation, group averaging and canonicalization (adding invariant features). Surprisingly, the authors found that in-training symmetrization often hurts performance, while posthoc averaging improves accuracy and symmetry of the learned wavefunctions. Both theoretical and empirical results are presented to justify the efficacy of posthoc averaging.\n\n## Update After Rebuttal\n\nThe reviewer is generally positive about the results presented in the paper and satisfied with the authors' responses to the questions raised here. Hence, the reviewer would like to remain the score.", "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": "We thank the reviewer for the constructive feedback. As all reviewers suggest additional references with some overlaps and similar questions on extendability, and due to the 5000 character limit per response, we address them together below.\n\n---\n1) **Essential references.** We thank all reviewers for mentioning interesting works. In the revision, we will **extend the introduction** and **include a new discussion and limitation section** to discuss these works. In detail,\n\n    **1.1) Essential pioneering works on neural network VMC solvers:** We will highlight [1] suggested by szMY (and also by wZs6) and other works by authors of [1]. \n\n    **1.2) Comparisons to symmetrizations that do not address diagonal space group symmetries or architectures that are not known to be effective in solids.** We first stress that we focus on **a class of difficult restricted symmetries**, i.e. diagonal space groups, which notably has only *discrete* point group symmetries. As observed by Reviewer wZs6, this is quite different from continuous symmetries like SO(2) and SU(2). **We realize this may not be sufficiently emphasized in our abstract and intro and will do so in the revision.** Also see our Section 2 for why these restricted symmetries are harder than e.g. the richer continuous symmetries in E(3), where the same discussion can be extended to SO(2) and SU(2). Architecture-wise, we stress that in the context of diagonal space group symmetries, **DeepSolid is the only architecture we know that performs well on infinite periodic solids**. Reviewer vG6b has pointed out interesting papers under Essential References and notes alternative approaches that focus on symmetries of observables rather than of wavefunctions. We agree that these are parallel and complementary approaches of interest and will cite these works. We clarify that for solids, our considered symmetries are typically applicable to ground state wavefunctions that are described by states with zero crystal momentums. We also note that, it is still an open question whether the architectures in those papers are effective in infinite period solids (since no solid benchmarks are reported) or for addressing the symmetries of the periodic solid. Making them work with solids would require additional tweaks, and **it will not be clear whether performance improvements or drops come from symmetrization or from additional architecture / hyperparameter tweaks**. This would invalidate our apples-to-apples comparison. However, we **wholeheartedly agree that it is important to survey more existing works on other symmetries**. We will cite [2,3,4,5,7] suggested by szMY and the symmetry papers suggested by wZs6 and vG6b, and highlight the differences. See also the response to wZs6 about why our findings seemingly contradict known benefits of symmetries.\n\n    **1.3) Extensions to broader contexts.** **We agree with reviewer szMY** that it is interesting to understand the effects of DA, GA and PA in other solvers like DMC. **We also agree with reviewer vG6b that** it is interesting whether our insights are applicable to open boundary problems and SO(3)-equivariant force fields. We do want to stress that **our findings are specific to complicated symmetries for which natural invariant maps are unknown**, which force us to adopt *model-agnostic approaches from conventional ML*. For simple symmetries or continuous symmetries like SO(3), where many *architectural symmetrization approaches* are available, we expect different results, e.g. no tradeoffs in in-training symmetrization. Nevertheless, **given the recent findings in protein structures and atomic potential ([1] and [2] cited in the response to Reviewer mfuU) that symmetries are unnecessary for training**, it is indeed of general interest how much symmetry plays a role in performance versus other factors e.g. hyperparameters. Our work can be viewed as **a first step towards this question in VMC**, by doing the first apples-to-apples comparison on an architecture before and after symmetrization. **The suggested extensions by the reviewers are very interesting future steps and we will mention them in the new discussion and limitation section.**\n---\n2) **Theory papers, szMY.** We thank the reviewer for suggesting these very interesting works on analysing high-dimensional Schrödinger equations (including [6] suggested by the reviewer in **essential references**). Indeed, our analyses are constrained by our desire to stay as close as possible to the DeepSolid VMC setup; a more careful analysis involves overcoming theoretical obstacles such as convergence guarantee of KFAC and short MCMC chains in high dimensions. We will discuss these suggestions as future directions to explore.\n---\nWe hope these answer your questions. If our responses are helpful, we would be grateful if you would consider raising your score.", "claims_and_evidence": "Here is the main claim made in this paper: posthoc averaging is a better choice to enforce diagonal invariance compared to in-training symmetrization. Sufficiently many experiments on Graphene, Lithium Hydride (LiH) and Metallic Lithium (bcc-Li) are included to compare postdoc averaging with other in-training symmetrization methods like data augmentation and group averaging, which supports the main claim in a clear and convincing way.", "methods_and_evaluation_criteria": "Yes. The experiments are mainly based on Graphene, Lithium Hydride (LiH) and Metallic Lithium (bcc-Li), which are standard examples used in the DeepSolid paper. The authors studies different symmetrization methods (data augmentation, group averaging and posthoc averaging) and compare their performance based on the induced wave functions and associated quantities like the ground state energy.", "theoretical_claims": "The only theoretical claim in the main text of the paper is Proposition 4.1. Its proof in Appendix G has been verified to be correct.", "experimental_designs_or_analyses": "Yes, please refer to the \"Methods And Evaluation Criteria\" section above.", "relation_to_broader_scientific_literature": "This work is mainly posited within the AI4Science literature on solving high-dimensional many-body Schrödinger equation based on wavefunctions parametrized via neural networks, which has applications in physics, chemistry and material sciences. In terms of methodology used in this paper, it mainly falls within the category of symmetry + machine learning.", "essential_references_not_discussed": "The reviewer finds that several key references are not discussed in this work. For instance, the authors didn't cite [1], which is, to the best of the reviewer's knowledge, one of the pioneering works that try to solve the many-body Schrödinger equation based on neural networks. Furthermore, in addition to the literature cited in the article, some other work like [2,3,4,5,6,7] have also studied the problem of incorporating symmetry within the machine learning based solvers of many-body Schrödinger equation. It might be meaningful for the authors to cite these articles and briefly discuss them as related work. \n\nReferences:\n\n[1] Carleo, Giuseppe, and Matthias Troyer. \"Solving the quantum many-body problem with artificial neural networks.\" Science 355, no. 6325 (2017): 602-606.\n\n[2] Mahajan, Ankit, and Sandeep Sharma. \"Symmetry-projected Jastrow mean-field wave function in variational Monte Carlo.\" The Journal of Physical Chemistry A 123, no. 17 (2019): 3911-3921.\n\n[3] Han, Jiequn, Linfeng Zhang, and E. Weinan. \"Solving many-electron Schrödinger equation using deep neural networks.\" Journal of Computational Physics 399 (2019): 108929.\n\n[4] Zepeda-Núñez, Leonardo, Yixiao Chen, Jiefu Zhang, Weile Jia, Linfeng Zhang, and Lin Lin. \"Deep Density: circumventing the Kohn-Sham equations via symmetry preserving neural networks.\" Journal of Computational Physics 443 (2021): 110523.\n\n[5] Lin, Jeffmin, Gil Goldshlager, and Lin Lin. \"Explicitly antisymmetrized neural network layers for variational Monte Carlo simulation.\" Journal of Computational Physics 474 (2023): 111765.\n\n[6] Abrahamsen, Nilin, Zhiyan Ding, Gil Goldshlager, and Lin Lin. \"Convergence of variational Monte Carlo simulation and scale-invariant pre-training.\" Journal of Computational Physics 513 (2024): 113140.\n\n[7] Zhang, Yaolong, Bin Jiang, and Hua Guo. \"SchrödingerNet: A Universal Neural Network Solver for the Schrödinger Equation.\" Journal of Chemical Theory and Computation 21, no. 2 (2025): 670-677.", "other_strengths_and_weaknesses": "This paper used a mixture of theory and practical experiments to justify the advantage of posthoc averaging compared to other approaches like data augmentation and group averaging, which is one of the first work that compares different ways of enforcing symmetry in the context of solving high-dimensional many-body Schrödinger equation. However, some potential drawbacks also exist. Firstly, the study is mainly limited to the Variational Monte Carlo (VMC) solver, so it might be meaningful to explore other solvers like Diffusion Monte Carlo (DMC) as well - see for instance [1]. Secondly, it seems that the claims and proofs in the paper can be made in a more mathematically rigorous way (which can be left as future work). For instance, below are some related literature that studies solving high-dimensional partial differential equations (PDEs) and symmetry from the perspective of statistical learning theory: [2,3,4,5,6,7,8,9,10].\n\nReferences:\n\n[1] Han, Jiequn, Jianfeng Lu, and Mo Zhou. \"Solving high-dimensional eigenvalue problems using deep neural networks: A diffusion Monte Carlo like approach.\" Journal of Computational Physics 423 (2020): 109792.\n\n[2] Jiao, Yuling, Yanming Lai, Dingwei Li, Xiliang Lu, Fengru Wang, Yang Wang, and Jerry Zhijian Yang. \"A rate of convergence of physics informed neural networks for the linear second order elliptic pdes.\" arXiv preprint arXiv:2109.01780 (2021).\n\n[3] Duan, Chenguang, Yuling Jiao, Yanming Lai, Xiliang Lu, and Zhijian Yang. \"Convergence rate analysis for deep ritz method.\" arXiv preprint arXiv:2103.13330 (2021).\n\n[4] Lu, Yiping, Haoxuan Chen, Jianfeng Lu, Lexing Ying, and Jose Blanchet. \"Machine learning for elliptic PDEs: Fast rate generalization bound, neural scaling law and minimax optimality.\" arXiv preprint arXiv:2110.06897 (2021).\n\n[5] Lu, Jianfeng, and Yulong Lu. \"A priori generalization error analysis of two-layer neural networks for solving high dimensional Schrödinger eigenvalue problems.\" Communications of the American Mathematical Society 2, no. 1 (2022): 1-21.\n\n[6] Lu, Yulong, Jianfeng Lu, and Min Wang. \"A priori generalization analysis of the deep Ritz method for solving high dimensional elliptic partial differential equations.\" In Conference on learning theory, pp. 3196-3241. PMLR, 2021.\n\n[5] Zweig, Aaron, and Joan Bruna. \"Symmetric single index learning.\" arXiv preprint arXiv:2310.02117 (2023).\n\n[6] Zweig, Aaron, and Joan Bruna. \"Exponential separations in symmetric neural networks.\" Advances in Neural Information Processing Systems 35 (2022): 33134-33145.\n\n[7] Zweig, Aaron, and Joan Bruna. \"Towards antisymmetric neural ansatz separation.\" arXiv preprint arXiv:2208.03264 (2022).\n\n[8] Zweig, Aaron, and Joan Bruna. \"A functional perspective on learning symmetric functions with neural networks.\" In International Conference on Machine Learning, pp. 13023-13032. PMLR, 2021.\n\n[9] Soleymani, Ashkan, Behrooz Tahmasebi, Stefanie Jegelka, and Patrick Jaillet. \"Learning with Exact Invariances in Polynomial Time.\" arXiv preprint arXiv:2502.19758 (2025).\n\n[10] Tahmasebi, Behrooz, and Stefanie Jegelka. \"The exact sample complexity gain from invariances for kernel regression.\" Advances in Neural Information Processing Systems 36 (2023).", "other_comments_or_suggestions": "N/A", "questions_for_authors": "N/A", "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": {"ethical_review_concerns": "N/A"}}, {"review_id": "0WQfSqMIxj", "reviewer": "Reviewer_wZs6", "summary": "The paper studies diagonal group symmetries in neural network solvers for many-electron Schrödinger equations, comparing different symmetrization approaches: data augmentation (DA), group averaging (GA), and post-hoc averaging (PA). The main claim is that in-training symmetrization can hurt performance while post-hoc averaging helps.", "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": "We thank the reviewer for the constructive feedback, which we address below.  We will add **a new discussion and limitation section** to reflect the discussions.\n\n---\n1) **PA's effectiveness on more architectures.** Indeed it's interesting whether PA is effective beyond DeepSolid. We stress that, as motivated in Sec. 2, we focus on a class of difficult restricted symmetries, i.e. diagonal space groups, for which only **model-agnostic symmetrizations are known**, which lead to our choice of DA, GA and PA. In the context of these symmetries, **DeepSolid is the only architecture we know that performs well on infinite periodic solids**, and is itself based on FermiNet. We do expect that for the *non-solid* systems handled by the original FermiNet and PauliNet, **architectural symmetrization** (e.g. the invariant maps discussed in our Sect. 2) may be possible and can be more effective than PA -- see response 1) to szMY.\n---\n2) **Conflicts with known benefits of symmetries.** The seemingly conflicting conclusions is related to point 1 in that we were forced to consider *model-agnostic symmetrizations*. For symmetries where *architectural symmetrization* is known, we don't expect them to suffer from the same cost tradeoffs or the same instability we saw for DA and GA. **We agree that symmetries may act via mechanisms we didn't address.** Here, we **try to minimize the influence of other mechanisms** in 2 ways: a) Focusing on the apples-to-apples comparison on a **fixed architecture with fixed optimizer and hyperparameters** and only vary the symmetrization; b) Citing a high-d CLT result that reduces the effect of DA to bias and variance (considered in a rich body of work on Gaussian universality and e.g. [5] cited in the response to Reviewer mfuU), and understand their effects empirically by stability and energy performance. **There could be other mechanisms that we didn't observe, and we'll mention this as a limitation.** We also stress that, in view of recent findings in AI4Science (see our response 1.3) to szMY), whether symmetry helps with performance is an open question in general, and **our apples-to-apples comparison consitutes a concrete attempt to minimize confounding effects** in the cases of DA, GA and PA.\n---\n3) **10x compute difference.** Apologies if it's unclear. GA with 12 symmetries requires 12x more evaluations of the network at the different inputs, which is expected to increase the cost significantly. This is why batch size was set to N/k to keep the same cost.\n---\n4) **KFAC vs alternatives.** We choose KFAC as it's the choice for DeepSolid and to ensure a fair comparison with the original model; see 2) above.\n---\n5) **Explanations of SC.** The limitation is specific to the ``average-near-the-boundary\" approach of building SC. Detailed explanations are in Appendix E, as the precise mathematical setup of SC is tedious (also the case in Dym et al. '24, who proposed the SC method we adapt). **Whether an efficient SC is possible for complicated groups is an open theoretical question:** Dym et al. only recently establish impossibility results for specific groups, and as mentioned in our Sec 2, building SC for diagonal space group is related to the problem of maximal invariants and orbifolds, which are unsolved math problems. **We'll ensure these are clarified in the SC section.**\n---\n6) **Batch-size ratios.** If we read the qn correctly, different ratios **are already investigated** in Fig. 7 of Appendix B.4.\n---\n7) **MCMC length.** We picked 30k as we empirically found that the sampled values had stabilised at this length. To address the review, we obtained additional runs for LiH PA with F222: The energy and var values are (-8.147(1), 0.018(1)) for 20k and (-8.1486(9), 0.0167(9)) for 40k, both within error range of the 30k result. Note that PA's benefits aren't from better quality of sampling but from the fact that it is sampling a different wavefunction.\n---\n8) **References; compare with other symmetrizations.** See above comments and response 1 to szMY.\n---\n9) **Theory for 2nd order methods.** This was discussed at the end of Sec 4.1, which highlights that the key barrier is analysing the high-dimensional Hessian matrix. As a further remark, a precise analysis rests on analysing the eigenvalues of a large random matrix; when the matrix does not have i.i.d. entries, this is a known hard problem in random matrix theory and no general results are yet known.\n---\n10) **Choice of groups.** The interesting observations in LiH and bcc-Li are indeed why we include the results. We don't have a concrete answer, but conjecture the effects to be a) system-specific and b) dependent on what approx. invariances the wavefunction was pretrained to possess. **We'll discuss these as interesting directions for future work.**\n---\nWe hope we've addressed as many comments as possible within the response limit. If they are helpful in addressing your concerns, we would be grateful if you would consider raising the score.", "claims_and_evidence": "The paper's theoretical analysis of gradient variance (Proposition 4.1) is well-derived and shows how DA introduces a variance inflation factor (k-1)/N. This manifests empirically in Figure 3, where DA with k=12 shows ~1.5x higher normalized gradient variance compared to baseline.\n\nThe post-hoc averaging results are compelling for the tested systems. e.g. improves energy from -8.138 to -8.1507 and reduces variance. However, the paper only tests three systems all using the DeepSolid architecture. Testing on more architectures (e.g., FermiNet, PauliNet) would strengthen the claims about PA's general effectiveness.\n\nThe claim that \"in-training symmetrization destabilizes training\" seems to conflict with the finding in (https://journals.aps.org/prresearch/abstract/10.1103/PhysRevResearch.5.013216), which demonstrates that properly incorporating symmetries into network architecture can actually improve training stability and efficiency in quantum lattice models. While the settings are different (lattice vs continuous systems), this contradiction deserves more discussion. The authors' theoretical analysis focuses on computational-statistical tradeoffs, but may be missing other mechanisms through which symmetries could aid optimization.", "methods_and_evaluation_criteria": "The evaluation uses standard VMC metrics: ground state energy and local energy variance. The energy improvements are significant given that chemical accuracy is ~0.00159. However, the paper lacks evaluation of other physical properties like electron density or correlation functions that could provide additional validation.\n\nThe computational cost analysis in Table 1 is detailed but raises questions. For graphene 1x1, GA with N=1000, k=12 takes 25s per step versus 2.5s for baseline, which is a 10x slowdown. The paper doesn't fully explain this large difference given that only k=12 group operations are being averaged. Additionally, the choice of KFAC optimizer is not justified - recent work has shown other second-order methods can be more effective for quantum many-body problems (e.g. https://www.nature.com/articles/s41567-024-02566-1).", "theoretical_claims": "The high-dimensional CLT result (Theorem D.1) provides rigorous bounds on the distribution of gradient updates. However, the key assumption that seems strong and its practical validity isn't verified empirically.\n\nThe smoothed canonicalization analysis shows an O(nk) cost scaling with number of electrons and group size. This explains why SC performs poorly in practice, but the paper could better explain if this limitation is fundamental or implementation-specific.", "experimental_designs_or_analyses": "The training setup uses KFAC optimization with batch sizes adjusted for computational budget. However, the sensitivity to these choices isn't explored - would different batch size ratios change the relative performance?\n\nThe MCMC chain length of 30000 steps for evaluation seems reasonable but isn't justified. Given that PA's benefits come from better sampling, analyzing how results vary with chain length would be valuable.", "relation_to_broader_scientific_literature": "The paper provides some coverage of recent neural network VMC methods but comparisons with other symmetry approaches are limited.", "essential_references_not_discussed": "https://www.nature.com/articles/s41567-024-02566-1 (second order optimization)\n\nhttps://journals.aps.org/prresearch/abstract/10.1103/PhysRevResearch.5.013216 (symmetry)\n\nAnd many other foundational works in VMC are also missing, including\nhttps://www.science.org/doi/10.1126/science.aag2302\nand many other works from the authors of the papers mentioned above", "other_strengths_and_weaknesses": "Already mentioned above.", "other_comments_or_suggestions": "No other comments.", "questions_for_authors": "1. For LiH, PA with F222 (k=16) achieves most of the improvement of Fm$\\bar{3}$m (k=192). Is there a systematic way to choose minimal symmetry groups that capture most benefits?\n\n2. The gradient variance analysis assumes first-order optimization, but experiments use KFAC. Can the theory be extended to second-order methods? Would this explain why GA performs better than DA despite similar variance bounds?\n\n3. For bcc-Li, different subgroups (P4/mmm, Fmmm) give similar improvements despite capturing different symmetries. What determines which symmetries are most important for improving the wavefunction?\n\n4. Why was KFAC chosen as the optimizer? Have you compared with other second-order methods like (https://www.nature.com/articles/s41567-024-02566-1)\n\n5. How does your post-hoc averaging approach compare with other symmetry preservation methods (e.g. https://journals.aps.org/prresearch/abstract/10.1103/PhysRevResearch.5.013216, where they show incorporating symmetries architecturally can improve training in quantum lattice models. \n\n6. How does your post-hoc averaging approach handle solutions allowing for different symmetry sectors?", "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": "UOMW8AKTOk", "reviewer": "Reviewer_vG6b", "summary": "This work investigates the impact of symmetrization in neural network wave functions for periodic solid systems. Specifically, the authors compare data augmentation, group averaging, and canonicalization. Contrary to other fields, the authors find that symmetrization may hurt performance, but post-hoc averaging generally yields better estimates.\n\n## After the rebuttal\n\nI increased my score in light of the authors’ rebuttal. However, I remain skeptical about the correctness of the evaluation. Biased estimates in VMC have frequently been a point in NN-VMC, and the reason why the number of MCMC steps typically increases linearly with the system size, as seen in [1]. Still, my main concern is the limited scope of this work. While it is certainly very valuable, I find the current scope too limited for a top-tier conference. I encourage the authors to broaden their evaluation, supporting wider applicable statements.\n\n[1] von Glehn et al. \"A Self-Attention Ansatz for Ab-initio Quantum Chemistry\"", "questions": "", "limitations": "", "rating": 2, "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": "We thank the reviewer for the constructive feedback and apologize for any confusion that arise from presentation issues or differences in understanding, which we now address.\n\n---\n**Correctness of gradient.** We thank the reviewer for the typo: $\\partial_\\theta \\partial_x^2$ should indeed be $\\partial_x^2$. The expectation is **otherwise correct**: F and Q are just generic update terms that\n- accommodate specific implementations of the gradient e.g. KFAC;\n- allow Prop 4.1, Lem 4.2, Thm D.1 to be *not tied to* a specific implementation. \n\nWe also stress that **the reviewer's formula is exactly the one that we compute numerically**; DA is run directly on the original DeepSolid except that gradient is computed on an augmented set of samples.\n\n---\n**Validity of DA,** which seems to be the main concern.\n- We stress that DA is *not* the method we champion. It's a popular ML method that one expects to help under approx. invariance; we follow its standard usage and discuss how it can fail.\n- As stated in Sec. 4.1, we follow standard usage to randomly sample DA **with replacement** (note the i.i.d. part). Unlike what the review suggests, the data is **not guaranteed** to include *either* the original *or* all the augmented data.\n- Indeed DA computes $E[F_{X,\\psi}]$ under a different distribution of X without changing $\\psi$. Yet this gives a **biased estimate, not an invalid estimate**. There are many other sources of bias e.g. stochastic gradient, KFAC approx etc., **some of which are known to help training**, and the gradient estimate is never unbiased in practice. The DA estimate is physically valid when $\\psi$ is exactly invariant -- where DA does nothing -- and gives a regularized estimate when $\\psi$ is approx. invariant. In ML, DA is routinely used under approx. invariance ([6,7,8] cited in the response to Reviewer mfuU) and **its bias is known to enforce regularization towards symmetry** ([5,9] cited in the response to Reviewer mfuU). For us, approx. invariance of $\\psi$ is enforced by pre-training and visible in Fig. 5(c). Moreover, DA gives **a physically valid energy** (Fig. 5(a)), just not an improved one. **We do agree this is important to clarify and will add this point.**\n---\n**Sampling vs gradient cost.** We thank the reviewer for highlighting the nuance that there may be better algorithms s.t. per-iter gradient compute is faster than per-iter sampling, esp. if the system requires more MCMC steps for good sampling. While we don't observe this in our setup, **we'll address it in the discussion**. We stress that the per-iter compute comparison is *not central* to our arguments, with some small additional nuances:\n- DA instability: Our DA batch size is k x N/k = N. If sampling dominates, one may exploit the sampling speedup by DA to get N' > N/k samples. Yet if N' < N, the use of fewer i.i.d. samples still inflates the variance; if N'=N, DA doesn't destabilize but the k DAs always increase compute cost. The tradeoff now depends on how large N' can be and thus how fast gradient compute is v.s. sampling.\n- GA: Unchanged; both sampling and gradient costs increase.\n- PA: Still computationally attractive, as overall training cost typically outweighs inference cost. The per-iter MCMC steps is multiplied by a large no. of training iters.\n---\n\n**Correctness of robustness.** We stress that the reviewer gives a practically useful but **mathematically equivalent** formula to $\\psi^{-1} H \\psi$, to which the *exact same discussion* is true, just more cumbersome. Note that differentiating $\\ln \\psi$ gives $\\psi^{-1}$.\n\n---\n**Over-emphasis on the negative results; more scalable symmetrization needed.** We note that our work **mainly seeks to offer perspectives rather than a universally good symmetrization**: We study a particularly hard case of symmetries and examine the strenghts and limitations of known ML symmetrizations. Given the success of posthoc averaging, one naturally expects in-training averaging / augmentations would help even more. **Our negative results are crucial in examining their potential failure points** and why that's not the case for PA. We also expect the negative results to be of **independent interest** to the applications of DA and GA **beyond a physics context**. We do agree that finding more scalable symmetrization for a large class of setups is an interesting open problem. Yet as the reviewer suggested, this requires finding the best combination of architecture, optimizer and symmetrization, and it's hard to tell whether the improvement is from symmetry or from other factors; see response 1.3) to szMY. We believe **part of the merits of our work is on an apples-to-apples comparison on a fixed architecture with v.s. without symmetry.** We will clarify these in the revision.\n\n---\n**References and questions**: See response 1) to szMY.\n\n---\nIf the above help to answer essential concerns behind the rejection, e.g. validity of DA, we'd be grateful if you'd consider raising the score.", "claims_and_evidence": "The authors make the following claims:\n1. symmetrization to diagonal groups destabilizes training and can lead to worse performance.\n2. post hoc averaging is effective in improving neural network solvers.\n\nThe paper supports these claims but to a somewhat limited degree due to the limited evaluation and specific problem statement. There appear to be theoretical misunderstandings about how to compute energy gradients in VMC. This most likely invalidates the data augmentation results. In general, contrary to the text statement, computing energy gradients does not involve computing $\\partial_\\theta\\partial^2_x\\psi_\\theta$ but is done by taking the derivative of $\\nabla_\\theta E_{\\psi^2}[\\psi^{-1}H\\psi]$ which corresponds to $E[(E_L - E[E_L])\\nabla_\\theta\\ln\\psi]$ where $E_L=\\psi^{-1}H\\psi$ as the Hamiltonian is hermitian, i.e., $E[\\nabla_\\theta E_L]=0$. Thus, the gradient we're computing is $\\int (\\nabla \\frac{\\psi(x)^2}{\\langle\\psi\\vert \\psi\\rangle}) E_L(x) dx$. Getting the expectations right is important. In the data averaging setup, the distribution of the expectations is changed yielding incorrect estimates and gradients. Further, contrary to standard data augmentation, a batch always includes the original sample and all augmentations, which is not the case in the standard case.\n\nAnother problem to support these claims is that the evaluation is limited to a single neural network wave function, limiting the impact of the statements on this neural network and the concrete symmetry groups. However, I acknowledge that the list of known symmetries is short for open boundary conditions.", "methods_and_evaluation_criteria": "The chosen periodic compounds are well chosen but the selection could be expanded to a cover a larger number of systems.", "theoretical_claims": "It is unclear to me what the message of Proposition 4.1 is. I generally found the notation unnecessarily complicated and partly undefined, e.g., the definitions of F and Q in l. 169/170.\n\nWhether sampling or gradient computations dominate the computations typically depends on the system size in VMC. For many electrons, sampling may take more time; e.g., in [1], the authors increase the number of sampling steps with the system size, and the gradient computation can be significantly accelerated [2].\n\n[1] von Glehn et al. \"A Self-Attention Ansatz for Ab-initio Quantum Chemistry\"\n[2] Li et al. \"A computational framework for neural network-based variational Monte Carlo with Forward Laplacian\"", "experimental_designs_or_analyses": "See above regarding the validity of the data augmentation scheme. Additional wave functions would be interesting, e.g., [1].\n\n[1] Gerard et al. \"Transferable Neural Wavefunctions for Solids\"", "relation_to_broader_scientific_literature": "Prior literature in neural network VMC focused primarily on accurate calculations by opening all degrees of freedom. This is mostly because most works focus on open boundary conditions where little is known about the symmetries of the wave function itself, only about the symmetries of its observables.", "essential_references_not_discussed": "The authors could discuss the difference to canonicalization schemes like [1,2] which focus on obtaining the right symmetries for observables rather than the wave function itself. \n\n[1] Gao et al. \"Ab-Initio Potential Energy Surfaces by Pairing GNNs with Neural Wave Functions\"\n[2] Gerard et al. \"Gold-standard solutions to the Schrödinger equation using deep learning: How much physics do we need?\"", "other_strengths_and_weaknesses": "The authors strongly focus on the negative results. Unfortunately, I don't see these supported very well due to the limited evaluation and partially flawed evaluation. The post-hoc averaging technique presents a valuable improvement for future solid VMC calculations but may be too limited in scope to justify an acceptance. Extending the scope to other structures and models or finding efficient symmetrization techniques that scale well would strengthen the position of the paper.", "other_comments_or_suggestions": "The justification for not including DA as a post hoc method is very obscure (l. 358), but given that it would yield incorrect estimates, it should anyway not be included.\nThe statement in robustnuss to outliers about the laplacian is incorrect; in practice, one does not compute $H\\psi$ and then divides by $\\psi$ but instead directly uses the formulation $\\psi^{-1} H \\psi=\\nabla^2\\ln\\psi + (\\nabla\\ln\\psi)^2$.", "questions_for_authors": "* Is there a way to translate these insights to open boundary conditions?\n* Can we incorporate the right symmetries directly in the architectures like SO(3)-equivariant force fields?", "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 authors compare different strategies for incorporating discrete symmetries of the sort that exist in materials into neural network ansatzes. They somewhat surprisingly find that, while explicitly symmetrizing during training can lead to instabilities, training an ansatz which does not respect discrete symmetries and then symmetrizing at inference time actually improves performance relative to the trained nonsymmetric ansatz. The reviews were generally positive. The paper could be improved with the inclusion of experiments on larger systems - only quite small systems were studied (LiH, graphene) and only a single unit cell, without twist averaging, so the results cannot be compared to experiment. Nevertheless, it's a surprising finding that could be useful for evaluating models trained on larger systems.
Decision: Accept (poster)