Title:Training Dynamics of Learning 3D-Rotational Equivariance
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 step…
Title:Training Dynamics of Learning 3D-Rotational Equivariance
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.
| Comments: | Accepted to Transactions on Machine Learning Research (TMLR) |
| Subjects: | Machine Learning (cs.LG); Biomolecules (q-bio.BM) |
| Cite as: | arXiv:2512.02303 [cs.LG] |
| (or arXiv:2512.02303v1 [cs.LG] for this version) | |
| https://doi.org/10.48550/arXiv.2512.02303 arXiv-issued DOI via DataCite (pending registration) |
Submission history
From: Max Shen [view email] [v1] Tue, 2 Dec 2025 00:48:09 UTC (6,266 KB)