How unconstrained machine-learning models learn physical symmetries
Authors: Michelangelo Domina, Joseph William Abbott, Paolo Pegolo, Filippo Bigi, Michele Ceriotti
Organizations: Laboratory of Computational Science and Modeling, Institut des Matériaux, École Polytechnique Fédérale de Lausanne, 1015 Lausanne, Switzerland
The requirement of generating predictions that exactly fulfill the fundamental symmetry of the corresponding physical quantities has profoundly shaped the development of machine-learning models for physical simulations. In many cases, models are built using constrained mathematical forms that ensure that symmetries are enforced exactly. However, unconstrained models that do not obey rotational symmetries are often found to have competitive performance, and to be able to learn to a high level of accuracy an approximate equivariant behavior with a simple data augmentation strategy. In this paper, we introduce rigorous metrics to measure the symmetry content of the learned representations in such models, and assess the accuracy by which the outputs fulfill the equivariant condition. We apply these metrics to two unconstrained, transformer-based models operating on decorated point clouds (a graph neural network for atomistic simulations and a PointNet-style architecture for particle physics) to investigate how symmetry information is processed across architectural layers and is learned during training. Based on these insights, we establish a rigorous framework for diagnosing spectral failure modes in ML models. Enabled by this analysis, we demonstrate that one can achieve superior stability and accuracy by strategically injecting the minimum required inductive biases, preserving the high expressivity and scalability of unconstrained architectures while guaranteeing physical fidelity.
Figures & tables
Figure 1: Overview of the structure of a symmetry-aware ML model, the conditions of group equivariance, and the symmetry diagnostic metrics introduced in this work. a) the ML model is represented by a generic smooth function, f , that predicts the physical properties (tensors of different rank), y^ , of an input, e.g. a decorated point cloud, x . f can be a symmetry-preserving (i.e. equivariant) or unconstrained model. b) Group equivariance is preserved if and only if the model predictions transform like the inputs under the action of the appropriate group symmetry operations. c) The metrics Aα and Bα introduced in this work quantify the equivariance error of model predictions and the group symmetry content of internal features. For a set of inputs given by Haar integration of x over the group, the equivariance error , Aα , is given by the variance of back-transformed model predictions, while the character projections , Bα , gives the group symmetry decomposition of model features from arbitrary layers.
Figure 2: Equivariance diagnostics for a PET MLIP. Top: distributions of the absolute error (AE) and equivariance error, Aα , for energy E , non-conservative forces fNC , and non-conservative stress SNC . The arrows on the x-axis indicate the distribution medians. Bottom: normalized character projections, Bα , for the corresponding quantities as a function of the probed angular momentum channel λ . Solid lines and markers are averaged over 150 randomly sampled test structures, while faint lines show the individual structure projections.
Figure 3: (a) Character projection heatmaps report on the magnitude of the character projection Bα for a quantity as a function of the character α of the relevant group and along successive epochs of a training run. In this case the characters are the (λ,σ) irreps of the O(3) group. (b) Training curves and character heatmaps for energy, E , non-conservative forces, fNC , and the two irreducible spherical components of the non-conservative stress, SNC .
Figure 4: Overview of a PET MLIP architecture and dynamical evolution of the internal features. (a) Local atomic environments of atoms in a molecule/material are represented as decorated point clouds. The atomic species of the central atom ( ai ) and its neighbors ( aij ), together with the edge vectors ( rij) and their magnitudes ( rij ), form the inputs to PET. (b) Backbone modules: embedding modules map the inputs to the latent space via Center, Neighbor, and Geometry Embedders. These feed into one Edge Embedder per GNN layer, which then enters the Transformer Layers. (c) Full architecture: the complete model pipeline leading to predictions for energy E , non-conservative forces fNC , and non-conservative stress SNC . The following architecture hyperparameters were chosen: number of GNN layers nGNN=2 , number of transformer layers nTL=2 , cutoff rcut=4.5 Å, edge feature size dPET=128 , node feature size dnode=256 , LLF size dhead=128 . The surrounding heatmaps, labeled with red letters (A–L), correspond to specific points in the architecture. They display the intensity of the normalized character projections, Bα/⟨∥t∥22⟩G , as a function of the training epoch ( x -axis, log-scale) and the probed λ channel ( y -axis). Within each heatmap group, the top and bottom panels represent the σ channels, respectively (as detailed in heatmap A). The thin, isolated column on the far left of each heatmap represents the untrained model, denoted by “U” on the x -axis.
Figure 5: Model error and equivariance error for the energy E , non-conservative forces fNC and λ=0,2 components of the non-conservative stress SNC for a universal PET model (trained and tested on MAD-1.5) whose energy and non-conservative force readout layers have been retrained with a loss L=Lμ+γLσ combining model error and an equivariance error penalty. The marker colors correspond to the weighting of the equivariance penalty, γ .
Figure 6: Equivariance diagnostic for the geometric pseudoscalar training. Top: Character projection heatmaps for a PET model across architectural layers over the course of training. Bottom: Test set RMSE and average equivariance error. The region between the red and the green dashed lines marks the onset of rapid learning. This phase is driven by a strong activation of the pseudoscalar channel across all capable layers, accompanied by a broader, weaker activation of higher-order tensorial ( σ=+1 ) and pseudo-tensorial ( σ=−1 ) channels (indicated on the right).
Figure 7: Impact of higher-angular order inductive biases on learning high-angular order targets. Standard PET (first column) and modified variants—which replace standard embeddings with solid spherical harmonics up to λmaxemb∈{2,4,8} —are trained on the λtarget=8 channel of the electron density projections ( w ). Top: proper-tensorial ( σ=+1 ) heatmaps tracking the angular order ( λ∈{0,...,10} ) evolution of features across key network layers during training. Colors denote the normalized intensity of the character projection. Bottom: validation metrics per epoch, including prediction RMSE (black), equivariance error (cyan), and the predicted-to-target norm ratio (orange).
Figure 8: Schematic representation of the PoLAr-MAE architecture, with the character decomposition of the internal tokens and the segmentation outputs plotted for 1000 test events. For representation purposes we show the decomposition after aggregation over the points of each event. Equivariance errors for the model class probabilities are shown for a representative event (the same as in Figure 9 ), showing how values of A correlate either with branching points or trajectory segments for which the classification is unstable with respect to rigid rotations of the reference frame.
Figure 9: A representative input and output of the PoLAr-MAE model. For a given input event (left panel), e.g., clusters of particle tracks (point clouds) in a liquid argon detector decorated with particle energies, the model makes point-wise classifications (right panel) of the discrete event class. Points surrounded by a red circle change classification when the structure is rigidly rotated. Axis labels are given in terms of voxel indexes.