Title: Energy Loss Functions for Physical Systems
Venue: NeurIPS 2025
Accepted: yes
Keywords: deep learning

AI for science

physics

physics-informed

diffusion

molecule

generative
Abstract: Effectively leveraging prior knowledge of a system’s physics is crucial for applications of machine learning to scientific domains. Previous approaches mostly focused on incorporating physical insights at the architectural level. In this paper, we propose a framework to leverage physical information directly into the loss function for prediction and generative modeling tasks on systems like molecules and spins. We derive energy loss functions assuming that each data sample is in thermal equilibrium with respect to an approximate energy landscape. By using the reverse KL divergence with a Boltzmann distribution around the data, we obtain the loss as an energy difference between the data and the model predictions. This perspective also recasts traditional objectives like MSE as energy-based, but with a physically meaningless energy. In contrast, our formulation yields physically grounded loss functions with gradients that better align with valid configurations, while being architecture-agnostic and computationally efficient. The energy loss functions also inherently respect physical symmetries. We demonstrate our approach on molecular generation and spin ground-state prediction and report significant improvements over baselines.


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe authors propose to replace generic standard training objectives like MSE or cross-entropy with energy-based loss functions that embed known physics of a system into the training signal: each datapoint is treated as an equilibrium state and the loss is simply the (approximate) energy difference between the model's prediction and that state, evaluated under a surrogate potential that respects the system's symmetries. The authors construct losses of this type for atomistic coordinates, diffusion-model training, and discrete spin systems and show that they are architecture-agnostic and symmetry-invariant without costly alignments. They demonstrate across shape-generation, molecular-generation, and spin-glass tasks that show that this physically grounded loss yields higher-quality samples and better data efficiency using standard datasets and models.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths\n\n- The idea is creative and a refreshing rethink of foundational concepts of ML\n- Through examples and theory, the authors show that their energy-based loss functions mirror real energy landscapes and can respect rigid motions and permutations. This is a strong property, since e.g. default DDPM optimizer gradients can pull predictions away from valid symmetry-equivalent structures.\n- The versatility of the recipe for constructing such loss functions is a strength; it works for several machine learning tasks without requiring architectural changes. \n- Losses lead to better gradients and avoids pathologies of e.g. MSE\n- On the benchmarks shown, the energy-based loss functions lead to tangible improvements. \n\nWeaknesses\n- The authors derive the loss function from a Boltzmann distribution; the derivation only makes sense here if training samples are near local minima of the true energy landscape. The authors note this assumption, but it is unclear to me how crucial it is to the framework. For example, if our data comes from MD simulations and contains e.g., non-equillibrium and transition states, can this result in loss gradients pointing the wrong way in situations where MSE stays agnostic?\n- The driving example of the paper uses a quadratic pair potential which results in a loss that scales with the number of particles N squared. While the authors mitigate this issue by selecting an O(N) subset of spring edges from a globally rigid graph, those missing long-range edges no longer penalise collective shears or electrostatic couplings, so it is unclear whether the sparse loss loses some of its informative gradient signal, even if its optima remain unchanged in theory. Table 2 shows that results are comparable to the non-sparse energy for QM9, but these are all small systems. I wonder what the tradeoff will be for systems that are larger and thus where the sparsification will be most significant for runtime.\n- Generally, I would have appreciated wall-clock speed comparing the default losses vs the dense and sparse training losses, preferably versus molecule size.\n- While the paper's derivations are mostly convincing, I find the experimental section thin where it matters most for adoption. From an application point of view, the molecule generation task with GDM and EDM is the most interesting, and the paper lacks a comparison that clearly indicates whether it can push molecule generation tasks beyond state of the art. For example, it would have been useful to see numbers from other publications reported in the papers' tables, in particular if the same architecture was used. For example, the JODO paper (Huang et al.) reports molecular stability numbers for QM9 of 93.4 % which is higher than the 89.8 reported in the present paper. It is plausible that the authors' energy-based loss objective could outperform if a similar newer architecture was used, but it is unclear whether the loss function improvement is orthogonal to the improvements from newer architectures with more embedded physics.  I also believe this would be very interesting in terms of the molecular stability on the GEOM-Drugs dataset, since this number seems to be either remarkably bad or remarkably good (I believe the latter is the case but see next point). It is, however, simply unclear where these numbers place the authors compared to prior work, and the experiments section thus does not lift this paper from theoretically interesting to broadly applicable, which I very well believe it could be. \n- Some background is missing. It would be useful to understand exactly how the stability was calculated, since the molecular stability number can be calculated in several ways (See e.g. Table 2 in the JODO paper, Huang et al., 2023). It would also be useful to discuss some of the points made in the discussion of stability calculations in GEOM-Drugs presented in Nikitin et al. (GEOM-Drugs Revisited:.. 2025), although this paper is very new and the authors haven't had a chance to include these results/discussions in their paper.  Generally, performing and discussing experiments with newer architectures and clearly reporting numbers that are fully comparable to previous publications could contribute to a better understanding of whether a modern architecture paired with an energy-based loss could achieve state-of-the-art molecular generation. If affirmative, this would very significantly impact my assessment of this paper.\n\nQUESTIONS:\nYour derivation relies on the data being close to local minima of the (unknown) true energy. How sensitive is the loss to violations of that assumption? Can you discuss or demonstrate this?\n\nCould you provide numbers on the wall-clock time and efficiency/accuracy tradeoffs of the different loss types on the more realistic benchmarks? The small reduction in accuracy for the sparse objective on QM9 would look stronger if you could demonstrate how much compute is saved\n\nCould you please clarify exactly how molecular stability was calculated and where the numbers place your results compared to prior publications? Preferably you could include numbers from other publications in the tables you already have produced.\n\nIn my opinion, the biggest improvement would be to show that using one of the more modern architectures (e.g. JODO or EQGAT) trained with your energy-based loss similarly leads to better results on the metrics typically reported. This would very likely push results beyond current state-of-the-art if the improvements are comparable. Even one small-scale experiment (same hyper-parameters, swapped-in loss) would show whether the gains you see are orthogonal to architecture advances, and would strengthen the claim that the loss is generally useful.\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 3, "CLARITY": 2, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-06-16"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposes a method to incorporate knowledge of physical systems into the loss function, whereas previous approaches have primarily focused on model architecture. This method is realized through a reverse KL-type loss function, which allows for an analytical expression of the loss. The authors demonstrate that this approach can be applied to diffusion models and enables the incorporation of invariances, with theoretical justification. Finally, they show the effectiveness of the method through several experiments, including molecule generation and spin ground-state prediction.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths\n\n・Overall, the paper is clearly written and easy to follow.\n\n・The idea of circumventing the intractability of the normalizing term by using the reverse KL divergence is simple yet interesting.\n\n・The proposed approach is broadly applicable to various machine learning problems, including generative tasks with diffusion models.\n\nWeaknesses\n\n・The difference between incorporating physical knowledge into model architectures and into loss functions should be more clearly articulated. At present, it remains unclear how important it is to include physical information in the loss.\n\nQUESTIONS:\n1. This energy-based physical knowledge can also be incorporated into model architectures. Could you discuss the strengths and weaknesses of incorporating it into the loss function, compared to incorporating it into the architecture?\n\n2. In Section 6.2, the results show that the energy-based loss yields stability. Could you provide a more detailed explanation of why the proposed approach leads to such stability?\n\n3. Can this approach be applied to other generative models, such as flow matching? Additionally, do the theoretical results discussed in Section 4.2 scale to those cases?\n\nLIMITATIONS:\nYes.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 2, "CLARITY": 4, "IMPACT": 2, "ORIGINALITY": 3, "DATE": "2025-06-25"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper introduces novel energy based loss functions for physical systems, like molecules and lattice systems. Although the authors mention that the energy of the target system could be utilized directly if available, they propose a different strategy and explain why using the true energy function can be infeasible. For many particle systems like molecules, they propose a loss based on distances between particles. In this case, the energy is given by the difference between the distance matrix of the prediction and the target. They use this to derive an alternative loss for diffusion models for many particle systems evaluate on small molecules, showing improvements over traditional diffusion models. \nMoreover, they also investigate a spin lattice system, where they use a different energy function based on the difference of the predicted spin and the target spin and its surroundings.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths:\n- The paper is well written and clearly motivated.\n- The idea of using a surrogate energy function in the loss in the proposed way is to the best of my knowledge novel. However, there are some related methods that utilize the actual target energy in the loss, which should be discussed in the paper (see suggestions)\n- They evaluate their method on very different systems, i.e., molecules and a lattice spin system, showing in both cases the advantages of their proposed loss. \n- As there are many applications of diffusion models for molecular systems, the proposed energy based loss instead of the usual diffusion loss might be useful for many applications.\n\nWeaknesses:\n- See suggestions and questions\n\nQUESTIONS:\n- Suggestions:\n    - Many flow modes, especially Boltzmann Generators, use the target energy function during training. Although it is different to what is discussed in this paper, I think it should be mentioned in the related work sections. And potentially also the abstract and introduction should be changed accordingly, as the phrase \"previous approaches mostly focused on incorporating physical insights at the architectural level. In this paper, we propose a framework to leverage physical information directly into the loss function ...\" is not correct.\n   - 44: Concurrent work of Klein et al., 2023 should be cited as they introduce the same training objective, i.e. \"Equivariant Flow Matching with Hybrid Probability Transport\" by Yuxuan Song et al.\n\n\n- Questions:\n  - Could the energy diffusion loss also be directly applied to flow matching? \n  - Why is the number of valid molecules in table 1 lower for the energy based loss? This seems counterintuitive, but maybe there is a good explanation. Moreover, are there errors for the results reported in this table?\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-01"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis work establishes a new method for physics-guided Machine Learning, where embedding physical principles into loss functions significantly improves performance in molecular generation and spin prediction. Its core contribution is shifting loss design from statistical error minimization to physics-driven optimization, ensuring the model contains physical constrains such as symmetry during predictions. The new model optimizes prediction errors, efficiency and reliability.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths:\n\nThis paper introduces energy loss functions, a novel framework that integrates thermal equilibrium properties of physical systems into Machine Learning loss design. Unlike traditional losses such as MSE, which optimize statistical prediction errors, this work leverages the Boltzmann distribution and reverse KL divergence to derive physics-based loss functions. This ensures model predictions adhere to physical constraints such as symmetry and energy minimization.\n\nOriginality: First systematic integration of physics into loss functions, beyond Machine Learning model architectural stage.\nQuality: A comprehensive theory with formal proofs in Appendix, extensive experiments (molecule generation and spin prediction) and code release commitment.\nClarity: Showing comparison between MSE and energy loss functions with clear physics figures such as Fig. 1. Comprehensive mathematical derivations in Appendix.\n\nWeaknesses:\n1. The selection of energy coefficients lacks discussion. The choice of Morse and L-J potential seems arbitrary, adding a few sentences to prove the completeness of the choice may help. In addition, it may help if the authors give a few examples of energy coefficients in different physical systems, such as protein and crystalline systems.\n2. Lack of discussions of how to build the sparse version of the energy loss. The main text mentions the sparse model reduce computation by O(N) operations, but without clear details.\n3. Temperature remains unexplored. Temperature is treated as a hyperparameter rather than a physical variable. The authors may need to add a few sentences to discuss what will change in energy loss building under different temperature.\n4. Incomplete comparison with equivariant architectures. May need to add studies to demonstrate the complementary of benefits of energy loss functions.\n\nQUESTIONS:\n1. Energy losses assume thermal equilibrium. How can the framework adapt to far-from-equilibrium systems (e.g., chemical reaction pathways)? Should dynamic energy landscapes be incorporated?\n2. The main focus is loss-level innovation. Could coupling with equivariant architectures enhance performance? For example: losses enforcing energy constraints + architectures handling symmetries.\n3. Will energy coefficients, the choice of potential energy equations and noise change under different temperature?\n4. How to modulate the energy loss functions when it comes to the physical systems that don’t obey Boltzmann distribution?\n\nLIMITATIONS:\n1. Atomic system coefficients (exponential decay) rely on experience without much theoretical guidance. A complex system (e.g., proteins) may need much more complicated energy approximations at prohibitive computational cost.\n2. Theoretical guarantees for score estimation valid only at low noise, high noise requires unaddressed corrections.\n3. Sparse rigid graphs reduce complexity but lack reported speedups, Fig. 3c only qualitatively claims O(N) scaling.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 3, "CLARITY": 3, "IMPACT": 2, "ORIGINALITY": 2, "DATE": "2025-07-02"}