Energy Loss Functions for Physical Systems

AI summary

Overview

Research area: Machine learning for the physical sciences — specifically, physics-informed objectives for regression, generative modeling, and classification on molecular and spin systems.

Technical level: Intermediate. Readers should be comfortable with machine learning training objectives (MSE, cross-entropy, KL divergence, diffusion noise prediction) and basic statistical-physics concepts (Boltzmann distributions, potential energy, partition functions).

Scope: The paper proposes a general framework for deriving "energy loss functions" from a reverse KL divergence against a Boltzmann distribution centered on the data, instantiates it for atomistic and spin systems, and evaluates it on shape generation, molecule generation, and spin ground-state prediction.

What This Paper Is About

Most prior work injects physical knowledge into machine learning models through architecture design, for example by building in equivariance. This paper instead asks whether the loss function itself can be made physically meaningful. The authors show that if you assume each data sample sits in thermal equilibrium on an approximate energy landscape, the reverse KL divergence against a Boltzmann distribution yields a loss equal to an energy difference between the data and the model prediction — a physically grounded alternative to generic losses like MSE or cross-entropy.

Key Contributions

1. A general recipe for physically grounded losses. A methodology for deriving loss functions by minimizing the reverse KL divergence between the model prediction and a Boltzmann distribution centered on each data point, giving a loss of the form E(ŷ, y)/T + log Z(y, T), where the log-partition term does not depend on the model parameters.

2. Atomistic instantiation with invariance analysis. A distance-based quadratic pair potential for atomic positions, plus a formal analysis of its symmetries (invariance to E(d) and to permutation automorphisms of the ground-truth distance matrix) and a characterization of its family of global minimizers.

3. Extension to diffusion models. A derivation showing that when the energy loss uses constant coefficients, its minimizers approximate the true score up to a translation in the direction of rigid motions, that the minimum-norm minimizer matches the true score, and that the resulting estimator has approximately zero bias and variance no greater than the MSE-based estimator.

4. Instantiation for discrete systems. A variational free-energy loss for classification and a local-field energy for spin systems, plus empirical evaluation showing consistent improvements over baselines across shape prediction, molecule generation, and spin ground-state prediction.

Main Findings

• MSE is an energy loss with a physically meaningless energy. The Gaussian conditional distribution used implicitly by MSE corresponds to T = 2σ² and an isotropic harmonic potential E(ŷ, y) = ||ŷ − y||², which models an external force pulling particles back to the data rather than realistic interactions.

• Distance-based energy losses are invariant to symmetry. With invariant coefficients, the atomistic energy loss is invariant to E(d) × (Aut(k(y)) ∩ Aut(Δy)), so no loss is incurred for predictions related to the data by rigid motion or by symmetries of the distance matrix. Unlike Kabsch-style aligned losses, this requires no expensive alignment or minimization procedure.

• Sparse rigid graphs give linear scaling without changing global optima. The full pair energy has a quadratic number of terms in the number of particles; using constructions from rigidity theory reduces this to O(N) operations while preserving global optima. In the regular-shape experiment, the sparse version remained equally performant as the complete-edge energy loss as the number of vertices scaled.

• Regular shape prediction: As the rotation augmentation angle θ_aug increases, MSE performance deteriorates while invariant losses (energy loss and Kabsch-aligned MSE) remain performant. Models trained with an invariant loss automatically learned to produce canonical orientations of shapes.

• Molecule generation, GEOM-Drugs: GDM-aug with energy loss reached 24.6% molecule stability, 96.0% atom stability, 89.7% validity, and 100% uniqueness, compared with 0.8%, 85.6%, 94.8%, and 100% for MSE.

• Molecule generation, QM9: With GDM-aug, energy loss reached 89.8% ± 2.8 molecule stability, 99.3% ± 0.3 atom stability, 97.7% ± 1.4 validity, and 99.9% ± 0.002 uniqueness, versus MSE at 83.7% ± 2.3, 98.3% ± 0.004, 93.6% ± 1.7, and 100.0% ± 0.0, and Kabsch alignment at 82.3% ± 0.5, 97.8% ± 0.004, 90.8% ± 2.0, and 100.0% ± 0.0. Sparse energy loss landed between the two (89.1% ± 0.9 molecule stability).

• Energy loss beats equivariance at low cost. Results on QM9 show that energy loss with a non-equivariant architecture produced more improvement than using an equivariant architecture (EDM with MSE reached 82.4% ± 3.4 molecule stability and 93.0% ± 2.5 validity), at negligible computational cost.

• Data efficiency: Energy loss allowed training of capable molecular generative models producing over 75% stable molecules using only 50% of the training set (50K samples), and gave +10% molecule stability over MSE in the low-data regime.

• JODO results: Adding energy loss to the JODO model improved nearly all 3D metrics and all alignment metrics relative to the default Kabsch-aligned loss. JODO + Energy (Exp.) reached 99.6% atom stability, 96.6% molecule stability, 98.4% validity, 98.4% completion, FCD 1.495, and bond/angle/dihedral MMDs of 0.0928, 0.0142, and 4.97e-3; JODO + Energy (Inv.) reached 99.4%, 94.3%, 97.1%, 97.0%, FCD 0.892, and 0.1125, 0.0046, 4.95e-4.

• Coefficient ablation (QM9, GDM-aug): Exponential distance decay performed best (89.8% ± 2.8 molecule stability), followed by inverse squared distance (84.6% ± 1.8), inverse distance (84.5% ± 2.1), and constant coefficients (83.6% ± 1.5). The authors note that inverse distance coefficients worked well for JODO, and a less stark decay worked better for the sparse loss on large molecules, concluding that coefficients must be ablated on new tasks.

• Score estimation is still correct. For the constant-coefficient case, the energy loss's minimizer set is the true score up to a translation along rigid motions, and the minimum-norm minimizer equals the true score.

• Spin systems: The local-field energy is convex because of its linear dependence on the prediction, and admits the data point as its unique minimum if h⁰ > 4 (or h⁰ > 0 if the data is a ground state of the true energy), making it a proper classification objective unlike the true energy. The specific numerical results of this experiment are not included in the provided paper content, which is truncated mid-sentence in Section 6.3.

Methodology in Plain English

The starting point is a standard way of turning a model into a loss function: instead of asking how likely the data is under the model (forward KL, which gives MSE under a Gaussian assumption), the authors ask how likely the prediction is under a distribution placed around the data (reverse KL). This is convenient because the normalizing constant — the partition function — then depends only on the data, not on the model parameters, and because the authors only need to define sensible distributions around real data samples rather than around potentially poor predictions.

For physical systems in thermal equilibrium, the natural choice for that distribution is a Boltzmann distribution. This turns the loss into an energy difference: the model is penalized by how much the approximate energy increases when you move from the data configuration to the predicted one. For atoms, the energy chosen is a quadratic penalty on pairwise distance mismatches, which is a second-order Taylor approximation of an interaction potential; the spring coefficients can be constant or decay with distance (inverse, inverse-squared, or exponential). Because the loss only depends on distances, it is automatically invariant to rotations and translations. For discrete systems like spins, the authors use the variational free energy of the predicted categorical distribution, with a local-field energy that measures the energy cost of flipping spins relative to the data configuration.

The empirical work spans three settings: a synthetic regular-polygon generation task with a two-hidden-layer MLP, molecular diffusion models (EDM, GDM, GDM-aug, JODO) on QM9 and GEOM-Drugs, and a CNN predicting ground states of spin glasses on 16×16 grids from Ising couplings sampled uniformly in [−1, 1], with targets computed by integer linear programming. All molecular comparisons were compute-matched, with exhaustive sweeps over learning rate and the weighting between position and atom-type losses.

Why This Matters

Impact on research: The paper reframes a common assumption — that MSE is the neutral default for continuous prediction — as a specific, and physically dubious, energy choice. It connects loss design to statistical mechanics and to rigidity theory, and it shows that improvements from better losses are complementary to, not subsumed by, architectural equivariance. This opens a distinct axis for physics-informed machine learning research that does not require changing the model architecture.

Real-world applications:

• Drug discovery: Molecular generative models that produce higher rates of valid, stable molecules with fewer training samples could reduce the cost of exploring chemical space. The paper explicitly evaluates atom stability, molecule stability, validity, and uniqueness because they matter to the drug discovery pipeline.
• Materials and crystal structure prediction: The atomistic loss applies wherever coordinates of particles must be predicted from scarce, expensive data.
• Combinatorial optimization: Spin systems with unstructured couplings are noted as being of interest in combinatorial optimization as well as physics, which connects the discrete energy loss to optimization problems.
• Large-scale structure modeling: The sparse, O(N)-scaling variant targets macromolecules and crystals with large unit cells, where the quadratic cost of pairwise terms would otherwise be a bottleneck.

Industry relevance: Any pipeline that trains generative or regression models on physical data — molecular design, protein structure, materials discovery — could swap in an energy loss without redesigning architectures. The paper reports that the gains come at negligible computational cost, and the provided code lowers the barrier to adoption.

Future Directions

• Coefficient selection remains task-specific. The ablation shows the best functional form for the spring coefficients differs between EDM/GDM, JODO, and the sparse large-molecule loss. The authors state it is necessary to ablate these coefficients on new tasks — a general rule for choosing them is an open question.

• Reconciling score bias with better samples. For more physically motivated coefficients, the theoretical score-recovery results no longer strictly apply, yet these variants performed better empirically. The authors hypothesize this is because evaluation measures physical plausibility rather than agreement between learned and data distributions, leaving a gap between theory and practice to explain.

• Scaling the sparse construction. Rigidity-theoretic sparse graphs maintained performance in the regular-shape experiment. Whether the same construction holds up for macromolecules and large-unit-cell crystals at scale is not demonstrated in the reported results.

• Scope boundaries. The framework is deliberately not aimed at sampling problems where a ground-truth energy is given (for example Boltzmann generators); extending it to those settings, and beyond atomistic and spin systems to other physical system types, is left open.

Target Audience

Machine learning researchers working on physics-informed models, generative modeling for scientific data, and equivariant or invariant architectures will get the most from this paper, particularly those interested in how loss design interacts with architectural symmetry. Practitioners in computational chemistry, materials science, and molecular design who train diffusion models on structural data will find the empirical results and released code directly applicable. Readers without prior exposure to Boltzmann distributions or KL divergences will need some background, as the derivations build on those concepts throughout.

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