Training Dynamics of Learning 3D-Rotational Equivariance
Authors: Max W Shen
Ewa Nowara
Michael Maser
Kyunghyun Cho
Venue: Accepted by TMLR ()

Abstract: While data augmentation is widely used to train symmetry-agnostic models, it remains unclear how quickly and effectively they learn to respect symmetries. We investigate this by deriving a principled measure of equivariance error that, for convex losses, calculates the percent of total loss attributable to imperfections in learned symmetry. We focus our empirical investigation to 3D-rotation equivariance on high-dimensional molecular tasks (flow matching, force field prediction, denoising voxels) and find that models reduce equivariance error quickly to $\leq$2\% held-out loss within 1k-10k training steps, a result robust to model and dataset size. This happens because learning 3D-rotational equivariance is an easier learning task, with a smoother and better-conditioned loss landscape, than the main prediction task. For 3D rotations, the loss penalty for non-equivariant models is small throughout training, so they may achieve lower test loss than equivariant models per GPU-hour unless the equivariant ``efficiency gap'' is narrowed. We also experimentally and theoretically investigate the relationships between relative equivariance error, learning gradients, and model parameters.

Published in Transactions on Machine Learning Research (12/2025)
Training Dynamics of Learning 3D-Rotational Equivariance
Max W. Shen'
sherm hww.wmn9@@en....m com
Genentech Computational Sciences
Ewa M. Nowara'
Genentech Computational Sciences
Michael Maser
Genentech Computational Sciences
Kyunghyun Cho
Genentech Computational Sciences & New York University
Reviewed on OpenReview: https: openreview. net/ forum? id=DLOIAW1883
Abstract
While data augmentation is widely used to train symmetry-agnoste models, it remains un- clear how quickly and effectively they learn to respect symmetries. We investigate this by deriving a principled measure of equivariance error that, for convex losses, calculates the percent of total loss attributable to imperfections in learned symmetry. We focus our em- pirical investigation to 3D-rotation equivariance on digh-dimensional molecular tasks (flow matching, force field prediction, denoising voxels) and find that models reduce equivariance error quickly <2% held-out loss within 1k- 10k training steps, a result robust model and dataset size. This happens because learning 3D-rotational equivariance is easier learning task, with a smoother and better-conditioned loss landscape, than the main prediction task. For 3D rotations, the loss penalty for non-equivariant models is small throughout training, so they may achieve lower test loss than equivariant models per GPU-hour unless the equiv- ariant "efficiency gap" is narrowed. We also experimentally and thcorctically investigate the relationships between relative equivariance error, learning gradients, and model parameters.
1 Introduction
Machine learning modeling of molecules generative modeling, properly prediction, simulating dynamics, etc. holds great potential for advancing scientific discovery and human health via therapeutics. Molecules are three-dimennsooall physical entities whose biochemical properties are invariant or equivariant to 3D rotations To model these symmetries, two approaches are common: 1) use symmetrr--rspeccing neural architectures, 2) training aygneicy-agnostic models with data augmentation, wherein training samples randomly transformed by the symmetry group. This choice is made at the start of any molecular modeling project and can have a significant impact engineering, training, and model performance, yet there has been lack of clarity on when to prefer which approach.
3D-rotational oquivariant architecturos use sophisticated tensor operations to maintain equivariance (Luo et al., 2024), achieve loss scaling curves similar to non-equivariant models (Brehmer et al., 2025; Mahan al., 2024), and are more parameter efficient than non-equivariant models on spherical image tasks (Gerken al,, 2022). Yet they can be much slower (10x-100x) than non-oquivariant models? (Gerken et al., 2022; Elhag et al., 2025; Brehmer et al., 2025), and they can be harder to optimize based on findings that breaking exact
*Co-first author.
Molecules also have symmetries translation, which commonly handled by centering molecule positions.
2 Slowness is also partially from less optimized code and GPU kernels.

---

Published in Transactions on Machine Learning Research (12/2025)
equivariance improves learning (Qu & Krishnapriyan, 2024; Portigkiozoglon et al., 2024; Canez et al., 2024). We call this the efficiency gop, arising both from optimization speed (training steps per second) and ease (loss reduction per training step). Meanwhile, recent work achieve strong performance on molecular machine learning tasks using non-equivariant architectures with data augmentation (Wang et al., 2024; Abramson et al., 2024; Qu & Krishnapriyan, 2024; Gellner et al., 2025).
d ERUUPP (Fors Beid prddction Gropp minarormer (14 cbt (h))
Proteto (lowren AFD the toneforner (130, (33V)
Vedded piono densking)
over 1- Ui fuereii chown with cotp
30 U-Me (1ri0)
MMeee Shoups Shogeanergge
(SNS Rocc
Encocer
parffcyyy apaiastant
Edge
Paa
Tuikte ttandonn
T
Vorraso blot aoe con-sagoosta
(Drarh aie (AOOS)
kalura
Iitipriinn bimond arielis BESK pi stang ahe aitention)
Loss desiepatiionn
a (Ai(uee)), 22 (17 7] On
vados
System
Peratom
seangy
vector
Hander bark
Farier LMAk increed ruir by
Yon org Aa sarce C003
MSE
Raaloarraroe BYTCK
MSF
VexMee 119M
Poras fald predictio Jeduwe aV, seniit))
Row naiching
Prones 9bb)
Figure 1: Overview of the paper. (a) Schematic of twisting and twirling, which underpin a principled measure of equivariance error. (b) Loss decomposition by Taylor expansion around the twirled prediction. (c) Loss landscapes for each loss component at early model checkpoints (step=500). (d) Archilectures of three non-equivariant models studied here. (e) For MSE loss, the loss decomposition holds exactly, enabling computing the percent validation loss from equivariance error, which is plotted by training step in three settings.
To answer "are symmetry-respecting architectures worth it?", one powerful principle is: the model that achieves better held-out loss. In fixed amount of GPU-hours, non-equivariant models could incur "unnecessary" equivariance error leading to higher loss, but equivariant models may achieve worse test loss due to the efficiency gap. In fact, we suggest the loss penalty vs. efficiency gap tradeoff is a general explanatory framework. This work focuses on equivariance, because on notation-invarian tasks like property prediction, symmetry-respecting architectures are relatively uncontroversial (Shoghi et al., 2024; Gasteiger et al., 2022; Nowara et al., 2025): they have a minimal efficiency gap to symmetry-agnostic architectures as rolation-invarianl features are informative and fast to compute, and standard deep learning operations easily preserve rotation invariance. In contrast, consider set permutation invariance where the symmetry- respecting architecture is the norm. This can be explained by observing that set transformers have minimal efficiency gap to symmetry--gnostic transformers, as set transformers simply ignore positional embeddings.
While is possible to directly compare efficiency gaps to loss penalties from imperfect symmetry, this is casily confounded by implementation details. To provide more fundamental insight, we instcad isolate and quantify a key source of potential underperformance in symmetry- to investigate: what is the percent of symmelry-ggnostic model's loss that comes only from its failure
models. We develop tools
2

---

Published in Transactions on Machine Learning Research (12/2025)
to be perfectly equivariant? ($2, Fig. 1A-B) In idcalized setting (ignoring efficiency differences), this characterizes the counterfactual error reduction if we had trained symmetry--especting model instead. Tn light of efficiency gaps for 3D-rotational equivariance, this metric quantifies how small the efficiency gap must become for equivariant models to outperform non-oquivariant models.
In this work, we focus our empirical investigations three high-dimensional (TR + R3N) molecular learning tasks satisfying 3D-rotational equivariance flow matching, molecular dynamics force field prediction, and denoising voxelized atomic densitics ($3, Fig. 1D). For any non-negative convex loss, decompose the total loss with data augmentation [()) mmean(0) Lequiv (0), where Lequiv captures all information about deviance from exact equivariance. In particular, exactly equivariant models have L(O)
Luean (0), i.c.,
Lequiv- = 0.
We find that equivariance error shrinks rapidly within 1k-10k training steps (minutes) to under 2% of the total loss (Fig. 1E). This occurs because Lequiv is a significantly casier learning task than Luean' the loss landscape for Lequiv is significantly smoother and better conditioned (Fig. 1C). Strikingly, this is robust to model size, training set size, batch size, and optimizer: we find it with standard batch sizes well as batch size I, on training sets of 1M molecules to small as 500 molecules, and model sizes from 1M to 400M.
Lastly, in $4, we conduct theoretical and experimental investigations to better understand the relationship between relative equivariance error, learning gradients, and parameters. We prove that under certain condi- tions, and experimentally, smaller equivariance error can increase the similarity between [(0) & Lmean (0). We also prove a quadratic relationship between Lequiv and the parameter deviation from the subspace of exactly equivariant functions for the modern graph transformer ASe.IP.
2 Measuring Equivariance & Loss Decompositions
Let f: RD be a learnable function and let G a compact group, for instance 3D rotations. We consider T as the matrix representation of the action of G on R.D A function f is G-equivariant if it commutes with all transformations T G, such that for any input æ E R", we have f(T(=)) T()())), also written (ʃ ( (x) (To (x). Rearranging, we observe that a perfectly equivariant function satisfies, for all x, T:
(T o foT)(()) = f(z)
(1)
We call (f=) the twisted prediction for x, from the twisted function oooT. To produce a twisted prediction" on molecules, we sample a random rotation, use to rotate the input molecule, pass this through the function, and un-rotate the output. The un-rotation step re-aligns the output to the "original frame" of the input molecule, which provides canonical frame to compare the impact of different transformations on the output.
In contrast to a perfectly equivariant function, a non-equivrriant function must have some distinct transfor- mations Ti, T2 where the twisted prediction is different: (7 f (z) * (T. (2). This property motivates analyzing the distribution of twisted predictions over a uniform distribution on the group, which is the usual choice for data augmentation. For a given x:
(T) e (T-1 (c)) T Uniform(G))
(2)
Its first central moment μ(x) is the group-averaged, twirled prediction.
((x) 2 of
(3)
3 The name reflects physical intuition introducing a twist in the middle rope with fixed endpoints: approaching the middle, the rope twists, and after the middle, it untwists.
3

---

Published in Transactions on Machine Learning Research (12/2025)
By the twirling formula, pz(z) is perfectly G-equivariant (Fulton & Harris, 1999). The second central moment of the twisted random variable is the covariance: (ovr((n(()) Er (7) - (2x)) - ==x))) J] The total variance the trace of the covariance matrix is natural measure of equivariance error:
D BErT |
(4)
This quantity measures the variance of the twisted predictions around their equivariant mean. An important property is that it is zero if and only if the function is perfectly equivariant. Higher moments of Zz()) likewise capture multivariate generalizations of skewness and kurtosis of the equivariance error.
2.1 Loss decomposition
Twisting and twirling provide machinery to understand a function's behavior around group actions. We can extend this machinery analyze losses used to train models under random data augmentation, where each training point is randomly rotatod. Let the data distribution p(z,y) and loss function i RDD RRD + R be invariant to G. That is, the joint data distribution p(z, for any transformation T satisfies: p(a,y) p(T(z), T())) and for any predictions and targets y, and for all T E G: I(T(z), 1(2, y) These conditions imply that the loss-optimal model is equivariant, and that: f o v)). The total loss over all data and transformations is:
T())) = 1((7-1
C(f) A ExaT [[([T]] ())),
(5)
We perform Taylor expansion of the total loss around the twirled prodiction yu(r)) (averaged over T), and obtain terms involving central moments of the twisted random variable:
C(f) = Ex Exp (e(()) v)]
+ 2 En. [tr (Hi(u(z), 8) Covr fo ())))) +0((88)1))))
twirled prediction
equivariance error
where 8 ==
u(z), Hi(, y) is the D x D Hessian matrix of the loss with respect to its first argument, and Covt is a D x D covariance matrix over the distribution of transformations T', Proposition 1. If (f y) bll== yl12 mean-squared error, then the total loss decomposes as:
C(f) (x,,)(((()) v EE,r [I[(T- Tx)) m(=)]1]]-
For MSE loss, our Taylor expansion reduces to a version of bias-variance decomposition. The equivariance error is identical to equation because MSE loss places equal weight on all dimensions. These two terms are central objects of study, so we name them:
Emean A En (pl))) v)]
(6)
1 En. [ D
Lequiv
(7)
Percent of loss from equivariance error. Denoting model parameters as 0, under MSE loss, we can express the total loss exactly L(0) Lmean (0) Lequiv (0). As all three terms are strictly non-negative, this implies:
% MSE loss from equivariance error = Lequiv
(8)
L())
Generalization to convex losses. We can further define a generalized mcasure of the percent of loss from oquivariance error for any convex loss function with non-negative outputs, such as KL divergence or cross- entropy. By Jensen's inequality, have Lueaa (0) L(O) and both terms non-negative. Furthermore,
4

---

Published in Transactions on Machine Learning Research (12/2025)
the two terms are equal if and only if the model is exactly equivariant. For convex losses, this motivates defining:
Lequiv _ [(0) meann (0)
(for convect losses)
(9)
as the non-negative difference, which is compatible with: % loss from equivariance error = Lequiv (()
Finite sample estimators. Denoting twisting as (x(x) (T f (T) the naive Monte Carlo esti- mates with N group samples are A(x) - + (xxx), and for MSE loss Emean (2) - (l((()) - yll (x)
ND Ei -(x))) However, Lequiv variance term, so it biased relative to the true population statistic, unless adjusted by the Bessel correction N/(N (([V-)), 1).. In $D.6, derive unbiased finite-sample esti- mators for Liman: Lequiv, and a bias-corrected finite-sample estimator for the percent loss from equivariance error: On 3D molecular tasks, we find that neural networks are typically smooth enough that only five to ten rotation samples are necessary to achieve a stable estimate of the twirled prediction and each loss component ($B.1).
Our derivations provide a principled framework for measuring and understanding degrocs of learnod equiv- ariance. An important property is that. for exactly equivariant architectures, Lequiv(0) 0, so that (()) (f). We remark that sometimes, non-equivariant models may be trained without data aug- mentation, so that this decompasition may not apply on the training loss. We stress that as long as we aim for these models to behave in an equivariant manner, then the loss decomposition is valid for held-out loss.
3 Experiments
To gain insight into the empirical learning behavior of non-equivariant models, we apply our loss decom- position framework three high-dimensional learning problems on 3D molecules, each with a distinct task and a modern non-equivariant model architecture. For each task, we follow the standard training procedure described in its original publication. The tasks span predictive regression tasks, autoencoding, generative modeling. Notably, all tasks use a mean-squared error loss, so our framework provides an exact decomposition of L(f) into Lmean and Lequiv. We report both of these metrics, as well as the percentage of the total loss altributable to the model's lack of equivariance, on a held-out dataset during training. We provide complete details on methods in $D.
well
Neural Interatomic Potential (NNIP): We consider force prediction with EScAIP (Qu & Krish- napriyan, 2024), a graph transformer architecture. The model predicts 3D force vector for each atom based density functional theory, mapping input molecule with N atoms to output in R3N This task is physically equivariant to the special orthogonal group SO(3) acting on atom coordinates in IR3.
Probabilistic Flow Matching: We study generative modeling lask with Proteína (GelTner et al., 2025), transformer- architecture with similarities to AlphaFold3. The model learns approximate the velocity field of probability flow that transforms random noise into structured protein backbones. For molecule with N alpha carbon atoms, the network maps noised atom coordinates and a time t E [0, 1] to a velocity vector in R3N, The learning task is made rotationally equivariant through data augmentation, aligning it with 53(3) acting on atom coordinates in R
Denoising Voxelized Atomic Densitics: We analyze denoising autoencoder task with VoxMol (Pin- heiro et. al., 2023; Nowara el. al., 2025), non-equivariant 3D convolutional neural network. Molecules are represented as densities in a cubic voxel grid. For a grid length g and a atom types, the input and output are tensors of shape [g. 9.9. a]. This lcarning task made rotationally equivariant through data augmentation using 16 axis-preserving 90-degree rotations of a cube, which do not introduce discretization artifacts due to aliasing. These rotations subset of the full octohedral group O.
5

---

Published in Transactions on Machine Learning Research (12/2025)
3.1 Force field prediction with EScAIP
We trained EScAIP 6M on subset of SPICE with 950k training examples used by Qu & Krishnapriyan (2024) for 30 epochs with batch size 64. SPICE is a dataset with of small molecule 3D conformers with energies and forces computed by quantum-mechanict density functional theory (Eastman et al., 2024). We varied model size from 1M, 4M and 6M, varied training set size from 950k, 50k, 5k, and 500 (with batch size 1); and varied the optimizer or learning rate. In this task, an equivariant model would output the same yet rotated force prediction, when the imput molecule rotates. We observe the following:
Equivariance is learned early and quickly, in a manner robust to training set size, model size, and optimizer and learning rate. The percent validation loss from equivariance error rapidly plummets in the first stage of training 0.1% within 1k-10k training steps (Fig. 2A-B). Notably, this speed is independent of epoch or training set size with a 950k training set, this occurs 25% through the first epoch. Training with 500 datapoints with batch size 1, this occurs at the fourth epoch. The dip is least affected by changing model size (Fig. 2E), and most affected by the optimizer and learning rate (Fig. 2F).
Equivariance is learned quickly because it is an easier learning task than the main prediction task. The loss landscape (Fig. 1C) for the equivariance error is much smoother and better conditioned, with a 1,000x lower condition number, than the loss landscape for the twirled prediction error.
After univerllv dip, percent loss from equivariance error can increase mildly. In the default setting, the percent incrcases from 0.1% to 0.3%. This is explained by a plateau in the equivariance error while the twirled prediction error continues to decrease (Fig. 2C).
Typical models converge to being nearly equivariant, with percent validation loss from equivariance error under 0.1%. The exception is training on 500 or 5k examples only: equivariance error continues to increase as training progresses, whereas equivariance error decreases in the long-term for larger training set sizes (Fig. 2D, Supp. Fig. 6).
3.2
Flow matching with Proteína
The conditional flow matching objective at each time optimizes a mean-squared error loss, which enables to apply loss decomposition to study the percent loss from equivariance crror. In this task, equivariant model would output the same yet rolated velocity, when the input noised molecule rotates Such equivariance is a common desired property for molecular generative models Geffner et al. (2025); Abramson et al. (2024). While many users may care more about final genorative quality metrics, the MSE loss plays a critical role in training, monitoring, and in defining the target loss-optimal velocity field. In particular, any non-equivariance in the model's learned generative distribution is caused by a non-equivariant. velocity field. This motivates tracking and understanding the percent loss from equivariance error of flow matching models at various which may enable refining training strategics to improve equivariance.
We trained Proteína. at 60M without triangular attention and 400M with triangular attention on the full Protein databank (PDB) dataset with 225k training examples. We also trained models on 1% of the PDB with 2k examples and 0.1% with 200 examples. Flow matching trains a model jointly over t, flow matching time, ranging from f = 0 for noise and t 1 for data. We measure metrics at / 0, 0.2, 0.4, 0.6, 0.8, 0.9, 0.95, and 0.99, and use red colors for high t close to the data, and blue-purple colors for low near noise Figure 3.
We observe that the equivariance learning dip occurs early for all t, in manner robust to training set size and model size. Following the dip at 1k-10k training steps (Fig. 3A), low t (closer to noise) are more equivariant, while high l (closer to data) are less equivariant, which spike following the dip. This holds for the 400M model (Fig. 3E), and 60M model trained on 1% and 0.1% of the PDB (Fig. 3F-H). Interestingly, the dip occurs at the same number of training steps despite different training set sizes, so it occurs 4% through one epoch when trained on the full PDB, but occurs around epoch 53 when trained on 0.1% of the PDB.
6

---

Published in Transactions on Machine Learning Research (12/2025)
Force field prediction with EScAIP
Eoptrartance emer dating traiiin
Equissrance encr durig hrining
decompootiion
Rapiczee
Tea
14:90
1225
100
Regatoo wur pan
10000
30030
3300
+8000
50003
18'
Travergeto
d
Vary Irain
mode
Modelicc
Opter
10.%%
1e-5
6.99
508 each ube 1ob, bertt Mir 04
0.1%
0.1%
Co 1006
10'
12'
12'
12
13'
12'
12' Thoking
Tooking
Figure 2: Training dynamics of learning equivariance in EScAIP (Force field prediction). (a-c) Validation losses and percent validation loss from equivariance error during training, early in training (a), with log-log axes (b), and decomposed into separate terms (c). (d-f) Impact of varying training set size (d), model size (e), and optimizer or learning rale (I).
After 1M training steps, the model's equivariance error is low across t, with maximum value around ~6% = 0.90, though it is lower at the extremes t = 0.99 at 3% and == 0 at 0.04% (Fig. 3B). This indicates that the time-conditional velocity field learned by the model is approximately equivariant. Our finding that at 1M training steps, t = 0.90 is the least equivariant timestep is relevant for designing training augmentation or test-lime strategies for improving equivariance. Notably, this finding highlights the advantage of our loss decomposition framework. Vonessen et al. (2025) measure equivariance error in Proteina by the variance of the normalized twisted prediction, but this not interpretable as percent of loss, and thus conflates task difficulty, which gets easier as - 1, with equivariance error. We correct for this issue, and find that t 0.9 is the most problematic time for non-equivariance, whereas they find 0.5 instead.
3.3 Denoising voxelized atomic densities with VoxMol
We trained VoxMol 111M on GEOM-drugs, a dataset of 3D structures of drug- molecules with 1.1M training examples. We also trained models on 1% (11k), 10% (110k). 25% (275k), and 50% (550k) examples, and models of varying size: full (111 M parameters), small (28 M), and tiny (7 M). In this autoencoding task, equivariant model would output rotated predicted reconstruction when the input molecule rotates; this is a commonly desired property when using the decoder a generative model Pinheiro et al. (2023). We observe that the equivariance learning dip occurs carly for all f, in a manner robust to training set and model size. Across the training set sizes, all models rapidly reduce their percent validation loss from equivariance error from an initial 80% to ~2% within 1k-10k training steps (Fig. 4A-B). At 50k training steps, models have around 5-10% validation loss from equivariance crror. Beyond 50k training steps, the twirled prediction error continues to decrease while the equivariance error plateaus, or decreases more slowly, below 1c-5 (Fig. 4C).
7

---

Published in Transactions on Machine Learning Research (12/2025)
matching Proteina
Tuddow aIPER OA
Noddoo Pusidae 502
Figure 3:
Training dynamics of learning equivariance in Proteína (Flow matching). Colors indicate flow matching
time, with noise at t 0 and data at = 1. (a) Percent validation loss from equivariance error during training. (b) Bar plot of the percent validation loss from equivariance error, by flow matching time, at final checkpoint after 1M training steps. (c-d) Validation losses by training step. (e-l) Impact varying model size (e), training set size (f-h..
Denoising voxelized atomic densities with VoxMol
Tertel
hatee dkleeriee
Tyr
2000
Figure 4: Training dynamics of learning equivariance in VoxMol (Denoising voxelized atomic densities). (a) Percent validation loss from equivariance error during training. (b-c) Validation losses by training step.
3.4 Loss Landscape Analysis
To better understand the initial dip, we studied loss landscapes for Liean and Lequiv at early checkpoints (500 steps). We computed the Hessian of each loss on a training batch for subset of 33k parameters including non-linear layers for EScAIP, 1.5k parameters in Proteína's 's linear head, and 6.9k parameters in final layer of VoxMol. For EScAIP, measured condition numbers around 1e9 for Lmesn and lef for Lequiv (~1,000x smaller). For Protcína, we measured 2c10 for Luean and 1c8 for Lequiv (~110x smaller). For VoxMol, we measured 5e9 and 6e8 respectively (~10x smaller). We calculate condition numbers for Lmean and Lequiv using the largest positive and smallest positive eigenvalues for each loss, and plot them over 20 minibatches in figures 12, 13. For loss landscape plotting, we chose two axes for plotting using the largest positive and smallest positive eigenvector on the total loss, and used the same step size and grid for Lmean and Leguiv- In both models, we find that Lequiv has substantially smoother loss landscape than Luean (Fig. 1C).
8

---

Published in Transactions on Machine Learning Research (12/2025)
3.5 Do Latent Representations Learn to Respect Equivariance?
The three model archilectures studied here have substantial differences (Fig. ID), yet they display some similarities in their training dynamics of learning equivariance. A natural question is whether their latent representations learn to respect equivariance during training. We provide an analysis here, with further details in $D.
EScAIP's latent representation is rotation-invariant. EScAIP uses rolation-invarianl features, and all intermediate layers maintain Intationinvariance Thus, the final representation is exactly rotation- invariant by design (Fig. ID). Note that the whole architecture is not invariant or equivariant due to the final prediction head (see eq. 17).
Proteina's latent representation is approximately equivariant. The architecture acts directly on 3D C-alpha coordinates, and does not use rolation-invariant or equivariant features. Its final latent sequence representation is mapped by a linear head into the model output (Fig. 1D). Thus, when the model empirically approximately equivariant, the final latent is also approximately equivariant.
VoxMol's latent representations are not equivariant nor invariant. Unlike EScAlP and Proteína where first-rriccipess reasoning suffices, we had to study VoxMol empirically. To evaluate if latents were equivariant, for an imput molecule and a given rotation, we measured the cosine similarity between the rotated latent, and the latent of the rotated molecule. We found median of 0.6, comparable to the cosine similarity between latents of different molecules, indicating a lack of equivariance (Fig. 11). To evaluate if latents were invariant, the same molecule as 0.64, which is statistically significantly higher but with small effect size than the cosine similarity between latents of different molecules at 0.58.
measured the cosine similarity between the latents of different rotations of
4 Impact of Lequiv < Emenn on Gradients and Parameters
Our empirical results showed that early in training, equivariance error rapidly diminishes, such that Lequiv is small relative to Emean- In this section, we investigate empirically and theoretically the connections between equivariance error, gradients and parameters. For non-negative convex losses, the loss is connected to the gradients by an analogous docomposition (eq. 9):
VC(0)= = V Vmmaan (0) + V (equiv
(10)
VC(0)= = V Cmeau (
(for exactly equivariant models)
(11)
When (lequiv vanishes, the model is said to follow "equivariant learning dynamics". While gradient norms in general have complex relationship with loss norms, in $4.1 we show that in certain regions of parameter space, as Lequiv shrinks, VL(() can become more similar to VL (0). Experimentally, we find moderate-to-stroog correlations between the percent. loss from equivariance error, and the similarity of (L( 2 Wmmenn (0), throughout training.
In $4.2, we consider parameter decomposition framework from Nordenfors et al. (2025), which shows that when certain assumptions hold, model parameters can be othogonally decomposed:
0 0g + 13g
(12)
where f3r is the model's parameter deviation from the subspace of perfectly equivariant functions E. These assumptions do not always hold, but they do hold for EScAIP's force prediction head, enabling us to prove a quadratic relationship between Lequiv and parameter deviation from the subspace of exactly equivariant functions for the modern graph transformer. Experimentally, we find the two are closely linked over training with Spearman correlation 0.99. Finally, in general settings where the parameter decomposition does hold, loss decomposition to derive some additional relationships between |0e_ to Lequiv for MSE loss.

---

Published in Transactions on Machine Learning Research (12/2025)
4.1 Equivariance error and gradients
We shall consider the relative loss ratio from equivariance error for non-negative convex losses:
c(0) Lequiv (0)
(13)
muean (0)
In proposition 2, we consider a mild assumption, commonly satisfied in practice, that the deep neural network is analytic, i.e,, it is composed from analytic activation functions, and that the losses are analytic. By standard real analysis results, analyticity implies smoot on compact subsets of parameter space (such as parameters explored in training; Lemma 7). A function f is My((U-smmoth z'| for all z, x' E X. Smoothness in turn can be usod derive a bound on the similarity of (()) z (meen (0) in terms of e(0). In practice, the bound can be weak when MU) can be large, and the bound becomes vacuous near saddle points where || (f)) - 0. Nevertheless, this result sheds light on the structure of the relationship between <(0) and the gradients.
a compact subset
if [[V - (( I [z)
Proposition 2. Let the model So be analylic (i.e., deep neural network constructed from analylic ac- tivation functions), and suppose the analylic loss satisfies [(0) meman (0) + Lequir(0) (i..., for conver losses via 9). Then, for any compact subset U in the parameter space of fo where <(0) is well-defined (the denominator is non-zero) and that includes an optima 6 with c(f²) 0, there exists finite constant M2(U)= = supecu [[)]l such that the approsimation VL(0) Rs VL (0) holds for all A E U with relative error bounded by:
VC((0) - C mean (0)1
Lmean (0) [ (0)) 2M,
< <(()) +
(14)
|V Leean [
Proof. Provided in C.3
Near the global optima, we can derive different bound in proposition 3 using the same assumptions, but employing theorem known as the Kurdyka-Eojasiewics (KE) inequality (Kurdyka, 1998; Dereich & Kassing, 2021). The Kurdyka-Pojasiewiez (KE.) inequality is generalization of the Polyak=Tojjssieiez condition, itself a generalization of convexity, which has been used to study convergence rates of stochastic gradient descent conditions that are more realistic to deep neural networks (Scaman et al., 2022). The Kurdyka- Eojasiewicz (KI.) inequality holds locally under mild conditions, only requiring analyticity, and states that there exists a compact local neighborhood U around any critical point 0" and constants 0, E [1,2) such that, for all 0 E U:
|V - 1(0')
(15)
This is mathematically an incquality in the opposite dircction as smoothness. Whereas smoothness ensures the function does not change too quickly, the KT. inequality says the function does not change too slowly, which is important for gradient descent convergence rates. For our purposes, smoothness gives an upper bound on the numerator, while the KL. inequality provides a lower bound on the denominator. Combining the two gives the ratio bound.
Proposition 3. Let the model fa be analytic (i.e., deep neural network constructed from analytic activation functions), and suppose L(0) = Lmean(f)) + (o) (i.e., for convex losses via 9) is analytic and has a minimum at 0. Then, there exists a compact neighborhood U around the global minimum 0* and finite constants e > 0, ( E [1,2) such that for all A E U, the approcimation VL(0) z [m holds with relative error bounded by:
(VL - me
2Mc (0)((
C Lequiv (f)
(16)
|| VV (!))
C Emean (0)a
c
10

---

Published in Transactions on Machine Learning Research (12/2025)
Proof. Provided in $C.4.
Experimental investigation. Our propositions 2, 3 show that in certain conditions, function of the loss ratio upper bounds the gradient norm ratio. We empirically investigated this by measuring the loss ratio and the gradient norm ratio during training, which we plot in $B.3. We find strong log-log correlations between the loss ratio and the gradient norm ratio of Pearson = 0.75 over training significant
EScAIP, and statistically 0.23 0.95 for Proteina across all times. Interestingly, we observe stronger correlations at smaller "lag" times, suggestive of stable "cocfficients" within basins, with sparse number of transitory windows where the models seem to "move between basins". Collectively, these experimental results provide complementary insight into the relationship between relative equivariance error and model gradients.
4.2 Equivariance error and deviation from equivariant parameter subspace
In the proceeding analysis, we adopt Nordenfors et al. (2025)'s mathematical framework for analyzing neural network parameters in terms of equivariant and non-equivariant parameter subspaces; under certain condi- tions, model parameters can be decomposed into orthogonal components: 0 = 0g fea. This framework relies three assumptions: (i) the symmetry group is compact and acts fiiiteddimensional hidden spaces; (ii) the neural net linarricssss are equivariant; and (iii) the loss is invariant. Under these conditions, the total parameter space is an inner product space, and the subspace of perfectly equivariant functions & is a linear subspace, which together enable the orthogonal decomposition 0 = 0g + f.g. This framework can be applied to a broad class of modern neural network operations and architectures, including fully connected layers with non-linearities, convolutions, residual connections, and attention layers. It also includes a broad class of symmetry groups including 5O(3) and all groups studied in this work. We provide more detail in $A.1 and refer the interested reader to Nordenfors et al. (2025).
Our scope and contributions here are as follows. Assumptions (i-iii) and the parameter decomposition = 0g + +51 do not always hold, but they do hold for EScAlP's force prediction head, enabling us to prove a quadratic relationship between Lequiv and parameter deviation from the subspace of exactly equivariant functions for the modern graph transformer. Furthermore, in general settings where the parameter decom- position does hold, Nordenfors et al. (2025) studies its implications, such as models following approximately equivariant learning dynamics when |0 i is small. In this general setting, we apply our loss decomposition to derive some relationships between ||0=_|| || to Lequiv.
What is the relationship between Lequiv and ||0,T]| when 0= 0g In general for neural networks, Lequiv is a complex, highly non-linear function of 0. However, we know that Lequiv is non-negative, continuous, and equal to zero iff 8gg = By these properties, we know that if ||0s]]] is small, then Lequiv is small. More formally, for any e > 0, there exists 8 > 0 such that if the parameter deviation small ([(0eill < 8), then the equivariance error is also small (Cequiv E).
The EScAIP architecture is modern graph transformer architecture that achieved strong results on NNIP energy and force prediction tasks, and satisfics assumptions (i-iii). The EScAIP architecture uses rotation- invariant features derived from an input molecular graph. Its hidden representations for atoms and edges, denoted h, are rolation-invariant. throughout. the network. Force prediction outputs a 3D force vector at each atom in a molecule. For a single atom with a set of 3D edge vectors E (the vectors pointing from one atom another atom) in a molecule æ, EScAIP predicts force voctors as:
Oz Oy = 0,
Te Wx Th(e, x) Cy wy (( x)
(17)
EEE
ez W. x((
where e E R33 is a 3D edge vector, hn,,, x) E A is the last hidden representation of the edge e in molecule æ, and W- = [wx, Wy. Wz]: where each w E R², are the parameters for a linear head with no bias. The 3D odge vectors e are rotation-oquivariant with respect to the input molccule, while the hidden represcntation h(e) is rotation-invariant to the input molecule, but composing these to form the output prediction generally breaks both invariance and equivariance.
11

---

Published in Transactions on Machine Learning Research (12/2025)
In particular, force predictions are equivariant if and only if the scalar projections of the hidden features are independent of the coordinate axis, i.e., WxTh(e, x) The x) w, W,,,(e, x), for all inputs. Under the mild assumption of a non-degenerate learned embedding function h(e,x), such that the set of all possible hidden vectors spans the feature space, this condition holds if and only if the parameter vectors themselves are identical: Wx = Wy = Wz- This condition defines the subspace & for the EScalP architecture. Using this, we decompose W We+ WET with an equivariant part We = [u,, w, wi where u 3 (Wxx Wy + and a non-equivariant part WEL [d,, dyy dz} E T3 where dN Wx w, and same for y:
With this setup, relationship with the magnitude of the parameter deviation from &, the space of perfectly equivariant func- tions.
establish that the equivariance error of the EScalP architecture has a quadratic
Theorem 4. For the EScAIP archilecture bruined with mean- squared ervor loss on non-degenerale dalasel, for any fized set of upstream parameters 9\\ W, there exist positive constants 0 < Aman S Amox (which depend on the model architecture, data distribution, and other parameters W. such that:
Amin WEL 1/ S Lequio(() S Xmax Weal
(18)
Proof. Provided in C.5.
Experimental investigation. In $B.4, we empirically plot the force prediction head's deviation from the mean, vs. equivariance error, and find a Pearson correlation of 0.94 and Spearman correlation of 0.99 over training, which supports our conclusion.
The preceding analysis can be generalized to a broader class of neural networks. Applying a Taylor expansion to Lequiv(0) for the neural net f on an input æ, we have: f(x; 0g + 071) f(x;0z) + Jes f(;; 00g + O(1(e111)) (13) where Joe. fx; 0g) is the Jacobian of the network output with respect to parameter components BEL, evaluated at 0g. The key structure, analogous to the FSScIPP argument, is the decomposition of the neural net output into purcly equivariant term, and a term linear in BEL, as well as remainder term in this setting. With this setup, for broad class of neural network architectures, we can relate locally near & that Cequiv is quadratic in |0s1 (Thm. 5), and its grad norm is linear in ||0=]]| (Thm. C.7).
Theorem 5. For any neural network whose perameters can be expressed as 0 = 0g + tgl with Gg E E and AET E EL, and for equivariance error Lequio defined by the voriance of the output with respect to transformations, there exist posilive constants Amin S Amaa such that for non-degene datasel, using || 1-1 denote La-norm:
Amiul 101 12 0(100111)) (+0(1=2')) s (equii(() S Ammaelleell 0(1(e111))
(19)
Proof. Provided in C.6.
Theorem 6. Under the same conditions as Thm. 5, the norm of the gradient of the equivariance loss with respect to the non-equieare con-equivarian paramelers bounded by the devialion ilself. Specifically, there crists a constont C such that:
Ve Ve Ceguri (0)l S C111l
Proof. Provided in C.7.
5 Related Work
Prior work have measured learned equivariance with a wide variety of approaches (Kvinge el. al., 2022; Karras et al., 2021; Geffner et al., 2025; Qu & Krishnapriyan, 2024; Gruver et al., 2023; Nowara et al., 2025; Fuchs et al., 2020), but our knowlodge, this work the first derive measure of oquivariance error that is interpreted as percent of loss. Notably, many prior measures effectively estimate equivariance error pairwise deviation using only two samples per datapoint, whereas estimate variance around mean
12

---

Published in Transactions on Machine Learning Research (12/2025)
using enough samples of the twisted prediction as necessary to obtain stable estimates. Vonessen et al. (2025) use the variance of the normalized twisted prediction, but this is not interpretable as percent of loss. They study flow matching, but their metric conflates task difficulty, which gets easier as t + 1, with equivariance error. We correct for this issue, and find that t = 0.9 is the most problematic time for non- equivariance, whereas they find 0.5 instead. Canes al. (2024) find that relaxing architectures from exact equivariance improves loss landscape conditioning and achieves better loss than perfectly equivariant architectures on image super-resolution and fluid dynamics modeling.
Twirling serves simple yet powerful postprocessing operation to transform any learned function into equivariant one at test time. Ideas like this have been explored in Pozdnyakov & Ceriotti (2023); Nordenfors et al. (2025). Preprocessing inputs to a canonical frame is another simple yet powerful postprocessing operation to convert any function into an equivariant one (Mondal et al., 2023; Gandikota et al., 2021).
6 Discussion
In this work, we found that 3D-rotational equivariance learned casily and quickly. We described two-phase learning dynamic: initially, models rapidly learn equivariance. This occurs because learning equivariance is easier task, with smoother and better-conditioned loss landscape, than the main prediction task. After training, the final percent loss from equivariance error is small for all models, but it is notably smaller for EScAIP at 0.006% than for Proteina and VaxMol (< 5%). While all of these loss penalties are small, and easily remedied by test-time postprocessing techniques like twirling or input frame canonicalization, this observation may also motivate research on architecture design to narrow this gap.
Intriguingly, equivariance is learned rapidly despite significant differences model architectures. EScAIP is "nearly equivariant", it becomes exactly equivariant with only a small change to its final linear head, yet its initial dip occurs just quickly as Proteína and VoxMol, which are distant from being architecturally equivariant. It is also interesting that each model's latents learn (or fail to learn) to respect symmetries in different ways.
Our work establishes a principled and unified framework for quantifying equivariance error for non-negative convex losses. We focused our empirical study on 3D rotations, as this is a physically important symmetry group for biomolecules, but other symmetry groups may be easier or harder learn. Looking forward, our framework could be used to study the learning dynamics of equivariance other symmetry groups.
Acknowledgements
We thank Pan Kessel and Saced Saremi for helpful discussions.
Code Availability
leaarriannne_eaannnn... learning. Our code simply adds callbacks to compute equivariance metries during training on top of the original EScATP, Proteína, and VoxMol codebases.
We provide code at https://githubb.cm wwatthub.com/ggneee
References
Josh Abramson, Jonas Adler, Jack Dunger, Richard Evans, Tim Green, Alexander Pritzel, Olaf Ronneberger, Lindsay Willmore, Andrew J Ballard, Joshua Bambrick, Sebastian W Bodenstein, David A Evans, Chia- Chun Hung, Michael O'Ncill, David Reiman, Kathryn Tunyasuvunakool, Zachary Wu, Akvilé Zemgulyté, Eirini Arvaniti, Charles Beattie, Ottavia Bertolli, Alex Bridgland, Alexey Cherepanov, Miles Congreve, Alexander I Rivers, Andrew Cowie, Michael Figurnov, Fabian B Fuchs, Hannah Gladman, Rishub Jain, Yousuf A Khan, Caroline M Low, Kuba Perlin, Anna Potapenko, Pascal Savy, Sukhdcop Singh, Adrian Stecula, Ashok Thillasuudaram, Catherine Tong, Sergei Yakneen, Ellen Zhong, Michal Zielinski, Augustin Zidek, Victor Bapst, Pushmeet Kohli, Max Jaderberg, Demis Hassabis, and John M Jumper.
13

---

Published in Transactions on Machine Learning Research (12/2025)
Accurate structure prediction of biomolecular interactions with AlphaFold 3. Nature, (8066)))) June 2024.
Johann Brehmer, Sönke Behrends, Pim De Haan, and Taco Cohen. Does equivariance matter at scale?, 2025. URL
Diego Canez, Nesta Midavaine, Thijs Stessen, Jiapeng Fan, Sebastian Arias, and Alejandro Garcia. Effect of equivariance on training dynamics, July 2024.
Steffen Doreich and Sebastian Kassing. Convergence stochastic gradient descent schemes for lojasiowicz- 09385.
landscapes. CoRR, abs/ abs/2100.09988, 2021. URL
Peter Eastman, Benjamin P. Pritchard, John D. Chodera, and Thomas E. Markland. Nutmeg and spice: Models and data for biomolecular machine learning. Journal of Chemical Theory and Computation, 20(19): 8583 8593, 2024. doi: /a002l/acs.jtt..col0 URL https://don fttps://doi.org/0.0/ atgg10..02//aee.... w0w.110221cc..aa wttaaacs.jjc...... Jctc.4c00794 PMID: 39318326.
Ahmed A. Elhag, T. Konstantin Rusch, Francesco Di Giovanni, and Michael Bronstein. Relaxed equivariance via multitask learning, 2025. URL https://arxiv wwteervvvvvorrvveeeeeee orgo httg//bb/2411..11.0
Fabian Fuchs, Daniel Worrall, Volker Fischer, and Max Welling. Be(3)-transfornn) 3d noto-translation equivariant attention networks. In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin (eds.), Advances in Newral Information Processing Systems, volume 33, pp. 1970-1981. Cur- ran Associates, Inc., 2020. URI. https:/ httpa::///ooped///eeeeee www.rip.cccpppprrrara..... 1522177cccbb7881332265666988 Paf
William Fulton and Joe Harris. Representation theory. Graduate texts in mathematics. Springer, New York, NY, edition, July 1999.
Kanchana Vaishnavi Gandikota, Jonas Geiping, Zorah Lähner, Adam Czapliískii, and Michael Moeller. Training or architecture? how to incorporate invariance in neural networks, 2021. URL https://arxix: org/abs/2100 10044.
Johannes Gasteiger, Muhammed Shuaibi, Amuroop Sriram, Stephan Günnemann, Zachary Ward Ulissi, C. Lawrence Zitnick, and Abhishek Das. Gemnet-OC: Developing graph ncural networks for large and diverse molecular simulation datasets. Trunsactions on Machine Learning Research, 2022. ISSN 2835-8856.
URL htttps://oppnreviiee wwt./eeruumeea.e ddu.r.sSsuu.....
Tomas Geffner, Kieran Didi, Zuobai Zhang, Danny Reidenbach, Zhonglin Cao, Jason Yim, Mario Geiger, Christian Dallago, Emine Kucukbenli, Arash Vahdat, and Karsten Krcis. Proteina: Scaling flow-based protein structure generative models. In International Conference on Learning Represenlations (ICLR), 2025.
Jan Gerken, Oscar Carlsson, Hampus Linander, Fredrik Ohlsson, Christoffer Petersson, and Daniel Persson. Equivariance versus augmentation for spherical images. In Kamalika Chaudhuri, Stefanie Jegelka, Le Song, Csaba Szepesvari, Gang Niu, and Sivan Sabato (eds.), Proceedings of the 39th International Conference on Machine Learning, volume 162 of Proceedings of Machine Learning Research, pp. 7404-7421. PMLR, 17-23 Jul 2022. URL ptttps:/prrceeediie
mlr wwwps/vveisiie gerre/2geree222a. html.
Nate Gruver, Mare Anton Finzi, Micah Goldblum, and Andrew Gordon Wilson. The lie derivative for measuring learned equivariance. In The Eleventh International Conference on Learning Representations, 2023. URL https://openrevier: http://wweetttorrriiiidd///a////
Tero Karras, Miika Ailtala, Samuli Laine, Erik Härkönen, Janne Hellsten, Jaakko Lehtinen, and Timo Aila. Alias-free generative adversarial networks. In M. Ranzato, A. Beygelzimer, Y. Dauphin, P.S. Liang, and
J. Wortman Vanghan (ods.), Advances in Neural Information Processing Systems, volume 34, pp. 852- 863. Curran Associates, Inc., 2021. URL /ttp wwt..orps.c.pppppetiiii..... (fopprr/ilee//pape//22// htl.://ccccaaabbbeafaa77ee222202 http:///wwa.....ooe..ppopppppppp
14

---

Published in Transactions on Machine Learning Research (12/2025)
Krzysztof Kurdyka. On gradients of functions definable in o-minimal structures. Annales de l'institut Fourier, 48(3):769-785 1998. URT. http::///eennnoooddoooo///oeee// orgg/oc/75332.
Henry Kvinge, Tegan Emerson, Grayson Jorgenson, Scott Vasquez, Timothy Doster, and Jesse Lew. In what ways are deep neural networks invariant and how should we measure this? In Alice H. Oh, Alckh Agarwal, Danielle Belgrave, and Kyunghyun Cho (eds.), Advances in Neural Information Processing Systems, 2022.
URL https: //perr/veee/eeeeeeeeeeeee
Shengjie Luo, Tianlang Chen, and Aditi S. Krishnapriyan. Enabling efficient equivariant operations in the fourier basis via gaunt tensor products. In The Twelfth International Conference on Learning Represen- tations, 2024. URL https://openrevvex net/ www.orrddmmmyllooo
Scott Mahan, Davis Brown, Timothy Doster, and Henry Kvinge. What makes machine learning task good candidate for equivariant network? In JCML 2024 Workshop on Geomelryy-rrunnde Represenlation Learning and Generative Modeling, 2024. URL http::///ppparreeaaaaaee//eotttp
Arnab Kumar Mondal, Siba Smarak Panigrahi, Sékou-Oumar Kaba, Sai Rajeswar, and Siamak Ravanbakhsh. Equivariant adaptation of large pretrained models. In Thirty-seventh Conference on Neural Information Processing Systems, 2023. URL https://opg http:////ppeeeeee////// wtwpearrwwieeeeeeeoooorriieeeeee wtt.:eeeooouuri.rrr..e..2.222
Oskar Nordenfors, Fredrik Ohlsson, and Axel Flinth. Optimization dynamics of equivariant and augmented neural networks. Transactions on Machine Learning Research, 2025. ISSN 2835-8856. URL https:
//openreview wtwpeevvvvvveee/fforeeeeeee het//forum/id--PT330225
Ewa Nowara, Pedro O Pinheiro, Sai Pooja Mahajan, Omar Mahmood, Andrew Martin Watkins, Saeed Saremi, and Michael Maser. Nebula: Neural empirical bayes under latent reprosentations for efficient and controllable design of molecular libraries. In JCML 2024 AT for Science Workshop, 2024.
Ewa M. Nowara, Joshua Rackers, Patricia Suriana, Pan Kessel, Max Shen, Andrew Martin Watkins, and Michael Maser. Do we need equivariant models for molecule generation?, 2025. URL https: org/abs/2507.005
Stefanos Pertigkiozoglou, Evangelos Chatzipantazis, Shubhendu Trivedi, and Kostas Daniilidis. Improving equivariant model training via constraint relaxation. In The Thirty-eighth Annual Conference on Neural Information Processing Systems, 2024. URL. https://openrevieee www.:/oormm.dd-iiiiiieee.e..
Pedro O. Pinheiro, Joshua Rackers, joseph Kleinhenz, Michael Maser, Omar Mahmood, Andrew Martin Watkins, Stephen Ra, Vishmu Sresht, and Saeed Saremi. 3d molecule generation by denoising voxel grids. In Thirty-seventh Conference Neural Information Processing Systems, 2023. URL https:// openreview met/forumiid-Zyyaeee
Sergey Pozdnyakov and Michele Ceriotti. Smooth, exact rotational symmetrization for deep learning on point clouds. In seirth Conference on Neural Information Processing Systems, 2023. URL https: //openreview.o htt.://ooorrriieaaaFsFff.....
Eric Qu and Aditi S. Krishnapriyan. The importance of being scalable: Improving the speed and accuracy of neural network interatomic potentials across chemical domains. In The Thirty-eightth Anmual Conference on Newral Information Processing Systems, 2024. URL
met/ffoumm? forum?id=Y4mBBazuury
Saced Saremi and Aapo Hyvärinen. Neural empirical bayes. Journal of Machine Learning Research, 20, 2019. ISSN 1532-4435.
Kevin Scaman, Cedric Malherbe, and Ludovic Dos Santos. Convergence rates of non-convex stochastic gradient descent under a generic lojasiewicz condition and local smoothness. In Kamalika Chaudhuri, Stefanie Jegelka, Le Song, Csaba Szepesvari, Gang Niu, and Sivan Sabato (eds.), Proceedings of the 39th International Conference on Machine Learning, volume 162 of Proceedings of Machine Learning Research, pp. 19810-19327. PMLR, 17-23 Jul 2022. URL https://proceeding mlr, html.
15

---

Published in Transactions on Machine Learning Research (12/2025)
Nima Shoghi, Adeesh Kolluru, John R. Kitchin, Zachary Ward Ulissi, C. Lawrence Zitnick, and Brandon M Wood. From molecules to materials: Pre-training large generalizable models for atomic property prediction. In The Twelfth International Conference on Learning Representations, 2024. URL net/forum het/forrmmidd-PPPnddaa
Carlos Vonessen, Charles Harris, Miruna Cretu, and Pictro Liò. Tabasco: A fast, simplified model for molecular generation with improved physical quality, 2025. URL
httpa/aivv..orgbaaa//a// httpa/2500009999 00899..
Yuyang Wang, Ahmed A. Elhag, Navdeep Jaitly, Joshua M. Susskind, and Miguel Angel Bautista. Swal- lowing the bitter pill: Simplified scalable conformer generation. In Porty-first International Conference Machine Learning, 2024.
16

---

Published in Transactions on Machine Learning Research (12/2025)
A Appendix
A.1 Parameter space decomposition
Here, we describe in greater detail Nordenfors et al. (2025)'s mathematical framework for analyzing the geometry of neural network parameters in terms of equivariant and non-equivariant parameter subspaces. This framework relies on three assumptions: (i) the symmetry group compact and acts on fimito-dimensional hidden spaces; (ii) the neural net non-linearities are equivariant; and (iii) the loss is invariant.
In practice, conditions (i) and (iii) are satisfied by many settings. Notably, the group SO(3) of 3D rotations is compact, but SE(3) which includes translations is not compact; but this is readily handled by restricting modeling to positionally centered data. Fimitte-iimensioaall hidden spaces is easily satisfied by neural net- works. Belief that a task is equivariant or invariant implics that the correct loss to use must be invariant. Condition (ii) is more setting-specinn this is satisfied in the VoxMol setting with voxel atomic densities with rotations, and non-linearities that act pixel-wise, but this is not generally satisfied by element-wiss non-linearities on linear transformations of 3D coordinates.
In the scope of this manuscript, use Nordenfors et al. (2025)'s framework to study the linear head of the ASP architecture, which does satisfy the conditions. We also use the framework to extend Proposition 3.14 in Nordenfors et al. (2025), which describes how models follow approximately equivariant learning dynamiies when 8.r is small. In our propositions 5, 6, we extend this with our loss decomposition framework, and derive a relationship between 0gg and Lequiv-
The foundation of this framework is the representation of network's parameters in all of its linear layers as a point in high-dimensional vector space, denoted H. This captures the dominant set of learnable parameters when non-linearities are fixed. The space is formally constructed the direct sum of the parameter spaces for cach individual layer: H = + Hom (o, Specific network architectures are assumed to have parameters in an affine subspace L C H, referred to as the space of "admissible layers". This setup is shown by construction to be expressive and capable of describing many modern neural network architectures and operations, including fully connected layers, convolutions, residual connections, and attention layers.
To define equivariance for multi-layer network, the framework supposes that the symmetry group G acts through series of representations, Ps. With this
all imput, hidden, and output spaces (X,, setup, the set of all parameter configurations where each linear layer is individually equivariant forms a linear subspace of H. denoted Ng. This set is a linear subspace because the group actions pi(g) is a linear operator, which means any linear combination of equivariant linear maps remains equivariant. For instance in the setting of rotations on 3D molecules, consider a linear layer with matrix with a rotation matrix R if it is equivariant, we have ARr RAr. If and are both equivariant to R, then C c + ca is also equivariant R: RCx R(c1A + c² (c1A + cx CRr.. Ho is thus a linear subspace that is closed under addition and scalar multiplication.
Algebraic manipulations show that TC = C,Tx,, using:
TC, = Ai
= CAATT + B,Tx
-c + 02 C2
=C,Tx
This subspace's linearity follows from the group's actions being linear transformations.
The parameters that are both architecturally admissible and perfectly equivariant then lie in the intersection of these spaces, E = L NG. It further follows that if non-linearities are equivariant (condition ii), then the entire neural network function is equivariant when its parameters are in E.
This geometric structure guarantees that any admissible paramcters 0 in L can be uniquely decomposed via orthogonal projection into two components: 0 == + This is possible because H being an inner product space allows for unique projection onto the tangent space of the subspace E. The component 0g is the
17

---

Published in Transactions on Machine Learning Research (12/2025)
projection of the parameters onto the subspace of equivariant functions (E), while 03p is the component in the orthogonal complement of this subspace, representing deviation from perfect equivariance.
B Supplementary Analyses & Figures
B.1 Sensitivity of percent loss from equivariance error to number of rotations used for estimation
Here, we plot bootstrapped estimates of standard error of percent loss from equivariance error, for model with 4 percent. loss from equivariance error, with varying number of rotations used to estimate the statistic. The model is a 60M Proteina model early in training at 500 steps, evaluated at t = 0.5. We compute loss metrics across 100 randomly sampled 3D rotations, and use this metrics set to derive subsampled bootstrap statistics for lower number of rotations. We find that the Proteina model is smooth enough, even early in training, over the group of 3D rotations. Our practical choice of estimating percent loss from equivariance error with 10 rotations has low standard error around 0.0015 compared to the measured percent loss from equivariance error of 0.04.
0.04
Percent loss Equivariance
0.03
0.02
0.01
0.00
Number of Rotations
Figure Bootstrapped estimate of standard error of percent loss from equivariance error, for model with percent loss from equivariance error.
B.2 Validation loss curves for EScAIP, by training set size
Here, we plot validation loss curves for EScAIP, split into total MSE loss, group-averaged loss, and equiv- ariance error, for models trained on different training set sizes. These results highlight that across a variety of training set sizes, group-averagod loss dominates the total loss. At smaller training set sizes, cquivariance error rises towards the end of training, suggesting some type of overfitting effect, while larger training set sizes, equivariance error continues to shrink.
18

---

Published in Transactions on Machine Learning Research (12/2025)
tc0tuiinng Loarphs
14 Foliig counglos
On Craaeunge Devitelles
Crup eyoee Conteflio
19
18
15"
Tario
Thaligg 10
Figure 6: EScAIP; Validation loss curves over training, varied by training set size.
B.3 Gradient norm ratios vs. loss ratios
Here, plot the loss ratio (percent loss from equivariance error) to the gradient norm ratio: the ratio of the gradient norm of the equivariance error, to the gradient norm of the group-averaged loss, over training. In general, gradient norms have a complex relationship with loss norms. In propositions 2 and 3, we show that under cortain local smoothness assumptions, the gradient norm ratio can be controlled by the loss ratio; however, this smoothness assumption may not apply in practice. With these empirical plots, we find generally strong correlation between loss ratio and gradient norm ratio.
Percent validation loss from equivariance ertrr
10.0%
350000
300000
250000
1.0%
200000
150000
0.1%
100000
50000
Speamnnn p: 0
10
Ratio of gradient eeguvvariance group-averaged ("twirlad")
Figure 7: ESSAIPP Percent validation loss from equivariance error vs. grad norm ratio, over training. Colored line indicates smoothed exponential moving average, colored by training step.
19

---

Published in Transactions on Machine Learning Research (12/2025)
1-0.2
-00%
Phrcatt maliom loa TE0t App.EEEE
Pannart www.eeeeeeeeeeeee
1,0%
Theasro slec
1.2%
Peanse Spaammar 2.0.3
0.64 Spearnat et 2.055
Rado of ga nom
Rado gradent nom equivarkane group
1-0.4
10.0%
100%
Parsent valcaaiin Kae NOTT MUNNEAEEE
1.0%
Spusmma st B.1
Spearman e: 0.5.5
10"
10
10"
Porks of guackent nor grappnaeaged
Porto gracient norm groppeweaaee
1-0.8
Penaase Ver Oaaa Soad from MQUANEMEEE
Parcart losa wew.haaaanncaaaaac
12.0%
1.0%
Spaarman
10"
10
10"
10'
Ratlo gradent
Rario grgncctnnt nor
Lequhorriance groppaneeeggel ("Iitled")
equuivriance grosp-awensgg ('wirkef")
30.0%
lose wqt.araaaaaaaoooooooooo
1.0%
1.0%
Pearson 20 10'
Pesrson
0.55
10"
12"
of droqukortance mor to group averaged ("wirled") loss Rorlo grrdlett
d'scukattncc eroo group averagel ("wilfed")
Figure 8: Proteina: Percent validation loss from equivariance error vs. grad norm ratio, over training, by flow matching time. Colored line indicates smoothed exponential moving average, colored by training step.
20

---

Published in Transactions on Machine Learning Research (12/2025)
B.4 Variance of EScAIP's linear force head controls equivariance error
Here, provide plots of the variance, or deviation from the mean, of EScAIP's linear force head, over training time, and compared to percent loss from equivariance error. Our proposition 4 states that equivariance error is controlled quadratically by the lincar head's deviation from its mcan. Empirically, we find strong agroement between the two metrics, with Pearson correlation = 0.94, and Spearman correlation 0.99 over training.
0.011
ESGAMP Inear Incar tace hoad weight vatance
0.007
0.006
100000
200000
300006
Stup
Figure 9: EScAlP: Variance of linear force head weights, by training time.
1.00%
400000
ERUMPP weight verianoe
200000
Triiinn
200000
150000
100000
50000
Pearson R: 0.94
Spearman 0.54
10
10"
10"
10
10"
10
Validation
Figure 10: EScAIP: Variance of linear force head weights, vs. percent validation loss from equivariance error, over training.
C Proofs
C.1 Proof of Proposition
Proposition. Ifi(z,y) = b112 yll mean-squared error, then the total loss decomposes as:
C(f) = (x)()(((()) y)l
D Er.y
Σ Va o
(20)
prediction
equévariance
21

---

Published in Transactions on Machine Learning Research (12/2025)
Proof. For mean-squared error, the Hessian is constant: Hi(z, v) D where I is the D D identity matrix. Furthermore, higher-order derivatives are zero, so the decomposition has no additional terms. The equivariance error simplifies as:
(6'1 Coryl...l)
2 Ex.y tr
= "D Bz.y [tr (Covrl....
(21)
C.2 Lemma: Analytic functions are smooth on compact sets
Lemma 7. Let f f:A. : - IR be a real-analytic function. Then, on any compact subset X C X, there erists a finite constant M such that f is -smooth on X; that is,
(22)
2
for some minimum x² E X.
Proof. Since f is analytic, all derivatives of f exist and are continuous. In particular, its Hessian f2)()) is continuous X. By the extreme value theorem, any continuous function on compact set attains a finite supremum. Therefore,
2 sup || [<>)l2 <0. zex
(23)
This boundedness of the Hessian implics that Vf is Lipschitz contimuou with constant, M on X, which is the M-smoothness condition. The inequality V M]]f(z) - f(''ll follows from integrating the gradient-Lipschitz property along the line segment between æ and x² or by standard smoothness results. 7
C.3 Proof of Proposition
Proposition. Let the model So be analytic (i.e., a deep neural network constructed from analylic acti- vation functions), and suppose the analytic loss satisfies L(0) (f) + Lequii(() (i.e., for conver losses via 9). Then, for any compact subset U in the porameter space of fo where e(0) is well-defined (the denominator is non-zero) and that includes optima 0 with e(0r) = 0, there erists finite constant MT) supeeu || [()(( such that the approrimation VL(0) 2 V. (0) holds for all O U with relative error bounded by:
VL(()) V memnn(()) (8) S 6()) + mmean (8) V (U)e(())(
(24)
| mean [
V Mmman (0)l
Proof. Using ((0) Emean (8) + Lequiv (0), the total loss is [(0) c(8)) (6). Differentiate with the product rule:
VC(0) = v[(1 + c(0)
(25)
Ve(() (0) (1 E Cmman (0)
(26)
VL(0) V Eieau (0) VL (0) Lieau ( Ve)
(27)
Now, we bound the norm of this difference using the triangle inequality:
VL() V Limonn (0) ( (0) Linaa Lm) Ve())
(28)
22

---

Published in Transactions on Machine Learning Research (12/2025)
The model and losses analytic, by lemma there exists finite MU) such that c(0) is M. over the compact parameter region U. Specifically, Mc(U) suPeEu [0)]- Applying the smoothness
property ||Ve(O)]I (2M((U)e(() relative to the optima 0* with E(0*) 0, we obtain the final result:
[VV VL (0)) S c(0) + [ (6) (M)(U)
mmman (0)
(9
[mean
C.4 Proof of Proposition
Proposition. Let the model Jo be analylic (i.e., a deep neural network constructed from analylic aclivation functions), and suppose C(0) [meon(f) + Leq (i.e., for convez losses via 9) is analytic and has a minimum Then, there erists a compact neighborhood constonts c 0, (0 (1,2) such that for all 0 E U, the approcimation VL(0) V.mean holds with relotive error bounded by:
around the global minimum Ha and finite
[V - mean
2M (on
C Lequiol (0) L Linean (f)a
(30)
V IV moaa (
Proof. By Kurdyka (1998); Dereich & Kassing (2021); any real-analytic function satisfies the Kurdyka- Lojasiewicz inequality, which states that there exists a compact local neighborhood U around any critical point 8" and constants > 0, E [1,2) such that, for all 0 U:
|V 2 cl((0) 1(0*)1"
(31)
When applied to the real-analytic function Lmnan at the global minimum 0" with Lmman = 0, the Hurdyka-Lojasiewin inequality states that there exists a neighborhood U around 6 and constants c > 0 and a E [1, 2) such that:
L (0)1 2 c. Lmeon (0)
(32)
Finally, by analyticity of the network and loss functions and applying lemma 7 relative the optima 0" with (equiv((*) O, there exists finite M Cequi (U) such that Lequiv is M Moqui (0) (0)( (U))))ooo(U))mooooo (Um-mmothh the compact neighborhood U:
Le (0)]1 < 2M (orgur (U) Lequiv
(33)
The final result follows algebraically by combining inequalities 32 and 33.
C.5 Proof of Proposition
Theorem. For the EScATP architecture trained with mean-squared error loss on a non-degenerate dataset, for any fixed set of upstream porameters W, there exist positive constants Amin S Aman such that:
Amain WELl < (fq < Amax |eal
(34)
Remarks. The constants Amin and Amax depend on the model architecture, data distribution, and other parameters 0 W.
23

---

Published in Transactions on Machine Learning Research (12/2025)
Proof. For molecule x, the k-th component of the predicted force vector decomposes into sum of contri- butions from We and WEL:
oR (z; W) = Σ ex (w"h(e))
Σ ex h())) CEE
(35)
cE
(oa We)
(xxL);WeL)
where the final hidden representation h depends on 0 W, the set of upstream parameters. Recall the equivariance error from Proposition 1, and observe that the variance of оμ Oeg + Ao₃ depends only on Aon, as Oeg is equivariant by construction. Thus, the equivariance error of the entire model, for a fixed set of upstream parameters and expressed as a function of the force prediction head parameters, is:
Lequiv (0) Bz,T [IAo W
By [To W
Now, let us denote: g(T, x. We1) = Ao We1). Observe that this function g is linear in our deviation parameters W11 By vectorizing the h xx3 parameter matrix We into a 34x1 column vector we can express this linear relationship as a matrix-vector product, for some matrix MT. with shape 3 x 3h: g(T, z, We1) = Mr Similarly, the atvario-averaged prediction g(x; WEL) = Er (T", æ, W&L) is also a lincar function, so we associate it with the matrix M,. The equivariance error term with these lincar matrix forms is:
vec( (WW.)).
Ea,Tlll( x, WEL) F(x W = ppT Op
(36)
where the matrix = E-Err((MM M, -M,) Finally, observe that positive definite, as the as equivariance error is strictly positive on a non-degenerate dataset whenever WE * O. By the properties of positive definite matrix, the quadratic form p Qp is lower-bounded by the smallest eigenvalue of Q, denoted Amin((), which is positive. It is also upper bounded by the largest cigenvalue Amaz(Q). This establishes the quadratic relationship on the equivariance loss stated in the theorem.
C.6 Proof of Proposition
Theorem. For ony neural nelwork whose parometers can be expressed as 0 OF + 0e with Of E E and l3l E EL, and for equivariance error Lequio defined by the voriance of the output with respect to trunsformations, there exist positive constants 0 < Ammin S Amaz such that for non-degenerate dataset, using || to denote Lz-norm:
Amin 101 12 0(1 S (0) DeLl1 0(10e1
(37)
Proof. Applying a Taylor expansion (oguiv(() for the neural net f on an input x around equivariant parameters 0g, have:
f(x 0g + 0el) = f(x; + Joc f(; + O(1 ll?)
(38)
where Jozn fx; 0g) is the Jacobian the network output with respect to parameter components fTg evalu- ated at 0g. As before, the term (a; 0g) is equivariant by construction, and thus drops out of the cquivariance error term. The term Jos, fax; 0g 0e is linear in dE7, which creates quadratic dependence on 0er in the variance term in Leguiv-
24

---

Published in Transactions on Machine Learning Research (12/2025)
The deviation from the twirled mean is the difference between the canonicalized prediction and its average over transformations. Let's expand this difference:
(T o f (2; 0) p<=; 0) - f (f)
(39)
Substituting the Taylor series and using the equivariance of f(z;
(J(x; 0E) + J0g ((T()):=))) + 0110= (2))
0z)] + 0 [*] - ET [f(2; 02) + (,0 Jo - (() (x): 0z1 + 0(0021 :)
(40)
(41)
Let AJ&T - T Jogu (TT))) [T J0p ((T The expression becomes A 881 + 0(100111
Lequiv (0) = Eg, |AJr 031 + 0(1001 11
(42)
D
1
E ET A 071 2(AJa.T 0._) -s_)) P) |0((0s) 1)]]]
(43)
The orders of the terms are:
||Ar, ll is 0(10s1113)
The cross-term is O111011 O 011101111 = 0110e1111)
The final term (O(1)0e |3)) = 01 11*).
We will study the leading term, which is quadratic in 0E1, and subsume the remainder into O(1101 I1*). As AJr,7 is linear function, we can define matrix Q that represents the averaged outer product of the Jacobian deviations: Q A 1 Ex,T (AJa,,) (AJ,,))) The equivariance error now be expressed conciscly:
Cequiv(0: + (11) ~ 05_ 00E1
(44)
The matrix @ is positive definite for non-degenerate dataset when 131 * 0. Using the Rayleigh-Ritz thcorem, this quadratic form thus bounded by the smallest and largest cigenvalues:
Amin 1 150 00g1 S Amax 10_l1
Reincorporating the remainder term in our Taylor expression, we arrive at:
Amin 10112 0(05111) x Lequiv (f) S Amax 00 101 0(10a1))
O
(45)
C.7 Proof of Proposition 6
Thcorem. Under the same conditions as the Taylor expansion theorem aboue, the norm of the gradient of the equivariance loss with respect to the non-equivariant parameters is bounded by the deviation itself. Specifically, there erists a constant C. such that:
| Coquio (0)1 S -
Proof. From our proof C.6, we know Lequiv (0) p Op, where p = voc(0()) The gradiont of a quadratic form is linear: VpCeuui = 2Qp. Taking norms, we get | Vp = ||Qpll < @2|[[]|]p) Setting C 22 2| 2112 gives the result.
25

---

Published in Transactions on Machine Learning Research (12/2025)
D Methods & Experimental Details
D.1 EScAIP
We trained EScAIP 6M on subset of SPICE with 950k training examples used by Qu & Krishnapriyan (2024) for 30 epochs with batch size 64. SPICE is dataset with of small molecule 3D conformers with energies and forces computed by quantum-mechaniee density functional theory (Eastman et al., 2024). We varied model size from 1M, 4M and 6M, varied training set size from 950k, 50k, 5k, and 500 (with batch size 1), and varied the optimizer or learning rate. The model predicts a 3D force vector for each atom based on density functional theory, mapping an input molecule with N atoms to an output in R&N, This task is physically equivariant to the special orthogonal group 50(3) acting on atom coordinates in R
We follow the same training recipe as the original repository, which does not data augmentation. We suspect that data augmentation is not important for EScAIP because it operates rotation-invariant features.
For further details and configuration files, please refer to our code repository.
D.2 Proteina
We trained Proteina at 60M without triangular attention and 400M with triangular attention on the full Protein databank (PDB) dataset with 225k training examples. We also trained models on 1% of the PDB with 2k examples and 0.1% with 200 examples. Flow matching trains model jointly over t, flow matching time, ranging from t = 0 for noise and t = 1 for data. We measure metrics at t 0.2, 0.4, 0.6, 0.8, 0.9, 0.95, and 0.99, and use red colors for high i close to the data, and blue-purple colors for low 1 near noise in Figure 3. The model learns to approximate the velocity field of a probability flow that transforms random noise into structured protein backbones. For a molccule with N alpha carbon atoms, the network maps noised alom coordinates and time / E [0, 1] to a velocity vector in RAN, The learning task is made rotationally equivariant through data augmentation, aligning with SO(3) acting on atom coordinates in R
For further details and configuration files, please refer to our code repository.
D.3 VoxMol
Following Pinheiro et al. (2023); we represent each molecule using 3D voxel grid by placing a continuous Gaussian density at each atom's 's position. Each atom type is assigned a distinct imput channel, producing a 4D tensor of shape [c x 7 x 1 x (, where c denotes the number of atom types and is the edge length of the voxel grid. The voxel values are normalized between and 1.
The denoising task arises from the use of walk-jump sampling for generating molecules (Saremi & Hyvärinen, 2019). This uses a two-step score-based sampling method. The "walk" phase involves running & steps of Langevin Markov chain Monte Carlo trajectory along a manifold. The "jump" phase applics a denoising autocncoder (DAE) to clean up the noisy sample using a forward pass of the trained model at step k. The DAE is trained on voxelized molecules corrupted with isotropic Gaussian noise, with a mean squared error (MSE) loss between prediction and ground truth. WJS provides a fast alternative to diffusion models by requiring only a single noise and denoise step (Pinheiro al., 2023; Nowara et al., 2024).
randomly initialized noisy voxel grid, simulating a stochastic
Architecture The VoxMol architecture is based on a 3D U-Net with convolutional layers spanning four resolution scales, and includes self-allention modules at the two coarsest levels Pinheiro al. (2023). During training, data augmentation is performed by applying random rotations and translations to each sample. For further architcctural and training details, refer to Pinheiro et al. (2023).
Mcasuring whether latent representations learn to respect cquívariance To evaluate whether VoxMol learns equivariant latent features, we analyze cosine similarity between latent embeddings under two scenarios.
26

---

Published in Transactions on Machine Learning Research (12/2025)
First, we examine representations of the same molecule under rotation. Let be molecule and R discrete rotation operator (e.g,, 90" around an axis). Using the encoder Φ(.) E BBCDDxH<W, with C = 512 and spatial dimensions 8 x 8 x 8, we define the spatially pooled latent vector:
x(x)
Σ d(x)li,d, h, d.h.a0
DHW
We then compute:
simeame cos (o(Rx(x)), (x(x(x)))
This measures whether encoding a rotated molecule is equivalent to rotating the latent vector of the original imput-a key signature of learned equivariance.
Second, to obtain bascline, we compute cosine similarities between embeddings of randomly selected dif- ferent molecules:
simait cos (Ø(x:); o(x;)), with x; Xx;
We compute these metrics across 1000 molecules for various rotation angles along all three axes. Cosine similarities are calculated over the 512-dimensional latent vectors and visualized using violin plots to capture the distributional differences Figure 11.
Findings. Cosine similarity between rotated versions of the same molecule tends to decrease rotation angle increases, reflecting imperfect latent equivariance. While same-moleccle embeddings remain more similar to each other than to embeddings of different molecules, the overlap between their distributions grows with rotation. This suggests that although the encoder partially preserves geometric structure, the latent space does not fully achieve rotation equivariance, indicating potential for improved regularization or architectural design.
Cosine Similarity of Latent Embeddinng
Cosine Similanty
0.-
0.4
Sarne Mol
Sare Mol
Same Mol
Sanie Mtl 1.3
Sanse Moi
Sanse Mol
Sane Mol yl
Same Mer
Same Mol
Same Mal
Sune Mai
Savie Moi
Sanse Med
Difieret Mois
0.5
Same
0.5
1.5
2.0
215
same
2.2
Figure 11: VoxMol: Cosine similarity of molecule latent representations with different rotations. y, indicate rotation axes, and mumbers 0, 0.5, 1, 1.5, 2 correspond to 0, 90, 180, 270, 360 degrees of rotation. The last column depicts cosine similarity between different molecules.
D.4 Metrics
To compute equivariance error, twirled prediction, error, percent MSE loss from equivariance error, and gradient norms, 10 rotations per sample used in EScAIP and Proteína. 4 rotations per sample were
27

---

Published in Transactions on Machine Learning Research (12/2025)
used for VoxMol. These numbers were found to be sufficient to provide stable signal for metrics which robust to randomness and resampling ($B.1). Fach measurement point uses different set of random rotations, so each metric's stability over time also reflects the stability of our measurements. For EScAIP; these metrics were computed on the first four (fixed) validation batches with batch size 16, for a total of 64 samples. For Proteína, these metrics were computed on the first. eight. (fixed) validation batches with batch size of 3, for total of 24 samples. The total MSE loss on these subsets was indicative of the total validation MSE loss, indicating these sample sizes were sufficient to provide stable and representative signal for these metrics.
D.5 Hessian Analysis and Condition Numbers
To plot the loss landscape, we selected subset of parameters in each architecture. For ESeATP, we used the final FFN (with non-linearity) and the final linear head, for combined total of 33k parameters. For Proteína, we used the final linear head with 1.5k parameters. We computed the Hessian of this parame- ter subset for the total MSE loss using one fixed training batch with ten rotations. We then performed eigendecomposition of the total MSE loss Hessian to find the eigenvectors for the largest positive eigenvalue, and minimum positive eigenvalue, which formed the two axes for plotting the loss landscape. We selected a step size approximately 2-3x the training slep size at that checkpoint, which estimated by multiplying the training learning rate with the total parameter gradient norm at that checkpoint. We then create a 2D grid of perturbations to the parameter subset, and compute Lmean and Lequiv at cach point the grid. Importantly, the axes and the step size are the same for both Lmean and Lequiv.
To compute the condition numbers, we computed the Hessian of the same parameter subsets for Lmaan and Lequiv separately, and performed cigendocomposition on them separately. We reported the condition number as the ratio between the largest positive eigenvalue and the minimum positive eigenvalue.
Here, we provide figures of the condition numbers across 20 minibatches.
EScAIP: Condition Number (33k params) after 500 training steps, over random minibatches
1011
1030
Condition Number
10a
10a
10a-
10a
Lmean
Lequiv
Figure 12: EScAIP: Condition numbers across 20 minibatches.
28

---

Published in Transactions on Machine Learning Research (12/2025)
Proteina: Condition Numbers (1. 11.5k params)
after 500 training steps, over random minibatches
1011
Conditio Number
1010
10"
10a
L_mean
Lequuv
Figure 13: Proteina: Condition numbers 20 minihatches.
aiter 502 traning st65s over rondom manbbrcches
Figure 14: VoxMol: Condition mimbers across 20 minibatches.
D.6 Bias-corrected finite sample estimators
In this section, we discuss unbiased finite sample estimators for Luean and Loquiv, which can be used to build a bias-corrected finite sample estimator for the percent loss from equivariance error. Suppose we have N > finite group samples. Denote the the exact group-averaged mean as μ (i.e., the exact expectation, which could be calculated with infinite group samples), and our finite sample estimate of the mean as p.
u(x) = Er [[T o T)(x)]
(46)
N M(x) N
[(T Tx)(
(47)
Similarly, denote o² as the exact variance over the group. Denote 8² as our biased finite sample estimate of the variance (biased because it divides by N, not -1): 1):
1 N
2²
(T
fo T.)xx) -(xx))
(48)
N
N
|
(49)
N
29

---

Published in Transactions on Machine Learning Research (12/2025)
The MSE loss of our finite sample estimate j is biased, relative to μ.. For exposition, denote μ fv
Br y)
(1 y)
(50)
N
squared
This can be derived as follows:
(i - y)² (ft + μ μr y)²
(51)
(ft p)) 2(1 y) 11) (a y)
(52)
Now, take an expectation over group samples
Eri(A y) (A+ μ 1 - y)?
(53)
= =El(AA - 11) + 2(μ - E. -1]]---]]]] + ExI(H - -y
(54)
-Erl(i- - 2 +(u y)²
(55)
=02/N
Thus, unbiased finite sample estimators are:
ET
(ft - y)²
y)² N
82
= (p- -
(56)
N
ET
62
= =²
(57)
N
Proposition 8. The finite somple estimator for the percent loss from equivariance error
02 N
(58)
( - y) +
has numerator that is unbiased to o², and has denominator is unbiased to (n y)² + 0².
Proof. For the numcrator, see equation 48. For the denominator: Erl(A y)² + 02] (μ y
y) + 02.
N -
Er((2 -) Ex (μ -y) + 0²
+
N o²
(59)
N
(μ y) + 0
(60)
A similar argument can be used motivate the N/(N 1) correction for general convex losses, where the correction is second-order accurate.
30


Reviews:
[{"review_id": "q6n1HVQIcL", "reviewer": "Reviewer_Rsqg", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "No", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This work investigated the training dynamics of learning 3D equivariance by decomposing the loss into a twisted prediction error and equivariance error. A few interesting observations were found in terms of fast learning of equivariance. The authors also tried to provide a theoretical justification.", "claims_and_evidence_yn": "No", "claims_explanation": "The authors provide a comprehensive numerical studies of fast convergence of equivariance error. However, I found most of their efforts in providing theoretical justifications questionable. See more details in my comments below.", "audience_yn": "Yes", "audience_explanation": "Achieving equivariance is an important task for machine learning, especially for non-equivariant models with data-augmentation.", "requested_changes": "Major issues:\n\nThe loss decomposition in 2.1 is formulated for a prediction test. It is not clear how it generalizes for training objectives of generative tasks like the conditional flow-matching loss for Proteina. The author should also discuss the interpreation of the loss for non-predictions tasks in section 3.2 and 3.3.\n\nSection 3.2, it is not clear how to interpret the traning dynamics for various flow time t, and why the loss with respect to t is interesting to the audience. In particular, the training of flow models is simultaneous for all t as a whole, and given that the sampling also requires using the velocity field at all times, why is the point of investigating the training dynamics of at a particular flow time t? The author should add discussion on the interpretation of these quanties that are numerically tested. \nUnfortunately, the author reports observations of the percentage validation loss for various but no conclusion/interpretation is given. \n\nFigure 3, it seems that the percentage equivariance error is decaying fast for noise regime (t close to zero). Is this because a Gaussian noise was used as it is rotation invariant automatically?\n\nIt is claimed in Section 4.1 that, as \\epsilon shrinks, \\nabla L_mean will dominate \\nabla_L. This is certainly not true in general because small loss does not imply small gradient. The author tried to prove this argument via the M-smooth assumption so that the gradient of \\epsilon is bounded by itself. However, this is not a mild condition and it is not true in general. The author should either provide insights why such condition can be met for certain models or at least numerically verify it. There is a paragraph discussing numerical validation but it is not clear why M-smoothness holds.\n\nThe proof to Proposition 3 is based on several previous results and mathematical notions that are not defined in the paper, making it challenging to read and understand the proof. The author should at least define important concepts (like KL inequality), and prove all claims such as analyticity implies M-smoothn ess. I suspect these claims are not true in general, for example x^4 is analytic but not M-smooth on the real line.\n\nSection 4.2, the decomposition here only makes sense when the equivariant models form a linear subspace.  The author tries to follow Nordenfors et al in Appendix A.1. However, this claim only holds for linear models. The author claims that ``element-wise applied nonlinearities are equivariant”, but this is incorrect for most activation functions. For example, ReLU is not rotation equivariant.\n\nMinor issues:\n\nPage 4: the author should (briefly) explain that the first term in the decomposition of the loss is averaged over T. \n\nFonts in some of figures (Figure 3) are too small to read.\n\nThe M_\\epsilon-smooth condition is never defined in the main text. There also seems to be a typo in its definition in the Proof of Proposition 2.\n\nTypos:\n\nPage 1, 2nd paragraph missing period.\n\nPage 3, extra period in the first sentence in section 3.3.\n\nPage 20, proposition 3 is labeled as proposition 7.\n\n\nNotations:\n\nThe ordering in the legend of Figure 2f should be fixed.\n\nFigure 4c, text should be added as in 4b to distinguish the prediction error and equivariance error.", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "PdKpGPI3IQ", "reviewer": "Reviewer_4MxY", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "• Introduces a “twist & twirl” framework to measure learned equivariance by comparing twisted vs. twirled predictions.  \n• Derives a principled loss decomposition under data augmentation: total loss = prediction error + equivariance error, exact for MSE. Provides an interpretable metric: % of loss due to non-equivariance.  \n• Empirical study on 3D rotations (force-field prediction, flow matching, voxel denoising) shows near-equivariance emerges early and robustly, with smoother optimization.  \n• Theoretical results connect equivariance loss to gradient alignment, parameter distance, and scaling laws.  \n• Practical insights: residual non-equivariance often has small penalties; non-equivariant models can be more efficient, while simple test-time twirling reduces loss.\n\nKey Strengths\n• Clear, general framework with exact decomposition under MSE.  \n• Strong, consistent empirical evidence across tasks and scales.  \n• Theory explains the fast emergence of near-equivariance and provides actionable test-time methods.  \n\nKey Weaknesses / Limitations\n• Exactness limited to convex/MSE losses; analysis approximate for others.  \n• Test-time twirling adds overhead without full trade-off analysis.", "claims_and_evidence_yn": "Yes", "claims_explanation": "• A formal loss decomposition (exact for MSE) is derived, so the “equivariance error” metric is well-defined and testable.  \n• Theory is tied to parameters (equivariant subspace vs. orthogonal part), with a quadratic relationship proved for EScAIP—making falsifiable predictions.  \n• Experiments across three distinct 3D molecular tasks show equivariance error plunges early (≈1k–10k steps) and stays small, robust to model/data/optimizer choices.  \n• Landscape probes indicate Lequiv is smoother/better-conditioned than Lmean, explaining the rapid convergence.  \n• Gradient/loss correlations support the “stochastic equivariant learning dynamics” explanation.", "audience_yn": "Yes", "audience_explanation": "• ML theory (loss decomposition, gradient dynamics). \n• Representation learning/symmetry (learned equivariance metrics). \n• Molecular ML practitioners (3D tasks, practical trade-offs, and twirling).", "requested_changes": "I believe the paper makes a solid contribution, and I would like the authors to perform the following to demonstrate the advantages of their work. \n1. Uncertainty quantification: For “% loss from non-equivariance” curves, add CIs (bootstrap over molecules & rotations) and a sensitivity study vs. # of rotations per sample.\n2. Fair baselines: Add equal-compute comparisons vs. strong equivariant models; report accuracy per GPU-hour and wall-clock; quantify test-time twirling overhead.\n3. Landscape analysis: Repeat conditioning/Hessian probes across multiple batches/checkpoints/parameter blocks; report dispersion (median/IQR).\n4. Scope of theory: Make MSE exactness explicit; either validate on another convex loss (e.g., MAE) or narrow claims accordingly.\n5. Data rigor: Specify precise train/val/test splits; check leakage (near-duplicates); add held-out test results for all settings.\n\nIn the following, I will list the changes that would strengthen the work even further, but are not necessary:\n• Parameter tracking: Track ∥θ_E⊥∥ over training and correlate with L_equiv; for EScAIP, show head weights (w_x,w_y,w_z) merging to a shared subspace.\n• Latent diagnostics: Add probes for learned (non)equivariance (e.g., predict rotation/canonical frame from latents) and report AUROC over training.\n• Augmentation ablations: None vs. single-axis vs. full SO(3) to isolate drivers of early equivariance.\n• Broader groups: Small demo on a second symmetry (e.g., permutations) or an appendix plan with expected dynamics.\n• Statistical reporting: Error bars/CIs on all time-series; include r and p-values for gradient-purity correlations.\n• Practical guidance: Add a 1-page table—when to prefer equivariant vs. symmetry-agnostic (+ aug/twirling) by data scale/compute/latency constraints.\n• Twirling recipe: Pseudocode + estimator variance notes; sensitivity to rotation count.\n• Writing/figures: Define “twist/twirl” once with a notation table; replace “dip” with precise wording; unify EScAIP naming; enlarge axis labels; move essential metric details from appendix to main text.", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "yoNUe5y2gw", "reviewer": "Reviewer_S3vC", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This work considers the setting of supervised learning in problems that feature equivariances. The typical choice in these problems is either to design an architecture that is inherently equivariant (but may be substantially more expensive for training and inference), or perform data augmentation with a standard architecture but penalise violations of equivariance. The work proposes to decompose the overall loss into two parts: one for the equivariance violations, and another for the main prediction task itself. The authors perform an empirical investigation on several 3D molecular prediction tasks, finding several interesting phenomena that are mostly consistent across tasks: a \"dip\" followed by a rise in the validation loss component, and a loss landscape for learning equivariance that is much smoother than the one for the main task. The work then derives a series of theoretical results on the learning dynamics, studying the magnitudes of the losses, their gradients, and the deviation in parameters from a reference perfectly equivariant function. The paper makes both theoretical and empirical contributions.", "claims_and_evidence_yn": "Yes", "claims_explanation": "Yes, both empirical and theoretical claims are appropriately supported by evidence. I lack the necessary background to check the correctness of the proofs, however.", "audience_yn": "Yes", "audience_explanation": "Yes, this paper is highly relevant for the TMLR audience and addresses an important problem with many applications in the sciences.", "requested_changes": "The paper is very well thought out,  well-executed, and well-written. My opinion of it is positive overall. I have a few suggestions that, in my opinion, would strengthen the paper. First the more important points:\n\nC1. The theoretical results are mostly disjoint from the empirical experiments, with the exception of those in 4.1, which do not hold consistently. Is there a way to also validate those in 4.2 and 4.3? If the answer is no, why is this challenging or impractical?\n\nC2. Can you make any statements about how difficult it would be to adapt the decomposition for classification tasks instead of regression with the MSE objectives? Would this still allow you to study the problem empirically even if not theoretically? \n\nC3. Do you intend to release the code publicly in the case of acceptance? Code was not submitted as supplementary material, which makes it harder to reproduce the work.\n\nTypos and small suggestions:\n\n- \"models rapidly become nearly equivariant within 1k-10k training steps\" -> seems this can vary substantially so \"rapid\" is sufficient for the abstract?\n- Footnote on author roles irrelevant for double-blind submission.\n- Bold part of point (ii) in Section 1 is an incomplete sentence.\n- 3.1, first bullet point \"its\" -> \"it is\"\n- 3.3, double period towards end of first paragraph\n- 3.5 first paragraph: \"we provide *an* analysis here\"\n- Section 4: can you provide a definition of gradient purity? The term is non-standard IMO.\n- Check references: Nowara et al. is duplicated; format for the conferences (e.g. NeurIPS) is inconsistent", "broader_impact_concerns": "None", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Decision: Accept (poster)