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.
Training Transformer-based architectures with finite data augmentation has become an increasingly popular approach in geometric machine learning. Despite its empirical success, the interplay between the Transformer architecture, invariance to different symmetries, and augmentation budgets remains underexplored. In this paper, we study the ability of a vanilla Transformer to learn various symmetries through finite data augmentation for point cloud datasets. We identify an ordering of increasing learnability across the following symmetry groups: (i) non-angle-preserving symmetries, (ii) angle-preserving symmetries, and (iii) base angle-preserving subgroups, such as translation, rotation, and scale. For the base angle-preserving groups, we further investigate the Transformer's extrapolation behavior and conduct a structural analysis of the trained models, allowing us to identify interpretable mechanisms that induce invariance. Finally, we extend our analysis to equivariant functions and show that the detected mechanisms for approximate invariance can also provide a key building block for learned equivariance. Our project page is available at https://transformers-learn-symmetries.github.io/
Eduardo Santos-Escriche, Valerie Engelmayer, Ya-Wei Eileen Lin +1
Technical University of Munich Munich Center for Machine Learning · Massachusetts Institute of Technology
Rotational symmetry is one of the most important structural principles in machine learning on 3D data. In applications ranging from physics and materials science to 3D computer vision, predictions should not depend on an arbitrary choice of coordinate frame. Rotational equivariance captures this requirement mathematically by enforcing that a rotation of the input induces a corresponding transformation of the model output. This tutorial provides a comprehensive introduction to rotational equivariance, starting from the physical and geometric intuition behind coordinate independence and building up the necessary machinery from geometric deep learning, group theory, and representation theory. We introduce message passing on Euclidean graphs, group actions and representations, spherical harmonics, Wigner matrices, tensor products, and Clebsch-Gordan decomposition, and explain how these ingredients give rise to modern equivariant architectures. We then survey the principal strategies for incorporating rotational equivariance in deep learning, including group convolutions, internal tensorial representations, and canonicalization-based methods, and discuss their practical strengths and limitations. The tutorial aims to lower the barrier to the subject by connecting the underlying mathematics to practical model design, by unifying ideas that are often expressed in different formal languages, and by helping practitioners choose among competing approaches through a clear discussion of their trade-offs.
Peter Lippmann, Fred A. Hamprecht
Interdisciplinary Center for Scientific Computing (IWR), Heidelberg University, 69120 Heidelberg, Germany
Learning symmetries through data augmentation offers an alternative to architectural equivariance, but does not guarantee that a network respects those symmetries. Measuring learned equivariance is especially difficult for intermediate representations, whose group actions are not specified. In a finite discrete setting, we show that equivariance requires representations of original and transformed inputs to determine one another exactly, yielding a test based on normalized mutual information (NMI). For high-dimensional activations, we use centered kernel alignment (CKA) as a scalable test, with an exact guarantee for orthogonal feature transformations. Across architectures and tasks, augmentation usually increases transformation consistency, showing that learned equivariance can be probed directly from representations without recovering the hidden transformation rules that relate them.
Claire Schlesinger, Wolfgang Gatterbauer, Robin Walters
Khoury College of Computer Sciences Northeastern University Boston, MA 02115, USA