Diagonal Symmetrization of Neural Network Solvers for the Many-Electron Schrödinger Equation

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.

---

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

• accommodate specific implementations of the gradient e.g. KFAC;
• allow Prop 4.1, Lem 4.2, Thm D.1 to be not tied to a specific implementation.

We 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.

---

Validity of DA, which seems to be the main concern.

• 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.
• 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.
• 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.

---

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:

• 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.
• GA: Unchanged; both sampling and gradient costs increase.
• 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.

---

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}$.

---

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.

---

References and questions: See response 1) to szMY.

---

If 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.

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.

---

1. 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.

---

2. 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.

---

3. 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.

---

4. 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.

---

5. 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.

---

6. Batch-size ratios. If we read the qn correctly, different ratios are already investigated in Fig. 7 of Appendix B.4.

---

7. 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.

---

8. References; compare with other symmetrizations. See above comments and response 1 to szMY.

---

9. 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.

---

10. 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.

---

We 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.

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.

---

1. 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,

   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].

   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.

   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.

---

2. 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.

---

We hope these answer your questions. If our responses are helpful, we would be grateful if you would consider raising your score.

We thank the reviewer for the positive and helpful feedback. We address the questions below.

---

1. 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.

---

2. 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.

---

3. 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.

---

We hope these answer your questions. If our responses are helpful, we would be grateful if you would consider raising your score.

---

References used for this response and other responses

[1] Abramson, Josh, et al. "Accurate structure prediction of biomolecular interactions with AlphaFold 3." Nature 2024

[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

[3] Hinton, G. E. "Training products of experts by minimizing contrastive divergence." Neural computation 2002

[4] Du, Y, et al. "Improved contrastive divergence training of energy based models." ICML 2021

[5] Chen, S., Edgar D., and Jane H. Lee. "A group-theoretic framework for data augmentation." JMLR 2020

[6] Lyle, C., et al. "On the benefits of invariance in neural networks." arXiv:2005.00178

[7] Benton, G., et al. "Learning invariances in neural networks from training data." NeurIPS 2020

[8] Yang, Jianke, et al. "Generative adversarial symmetry discovery." ICML 2023

[9] Balestriero, R, Leon B, and Yann L. "The effects of regularization and data augmentation are class dependent." NeurIPS 2022