Training Dynamics of Learning 3D-Rotational Equivariance

Fuente: arXiv
Salvato in:
Dettagli Bibliografici
Autori principali: Shen, Max W., Nowara, Ewa, Maser, Michael, Cho, Kyunghyun
Natura: Preprint
Pubblicazione: 2025
Soggetti:
Accesso online:
Tags: Aggiungi Tag
Nessun Tag, puoi essere il primo ad aggiungerne!!
_version_ 1866915647985811456
author Shen, Max W.
Nowara, Ewa
Maser, Michael
Cho, Kyunghyun
author_facet Shen, Max W.
Nowara, Ewa
Maser, Michael
Cho, Kyunghyun
contents 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.
format Preprint
id arxiv_https___arxiv_org_abs_2512_02303
institution arXiv
publishDate 2025
record_format arxiv
spellingShingle Training Dynamics of Learning 3D-Rotational Equivariance
Shen, Max W.
Nowara, Ewa
Maser, Michael
Cho, Kyunghyun
Machine Learning
Biomolecules
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.
title Training Dynamics of Learning 3D-Rotational Equivariance
topic Machine Learning
Biomolecules
url https://arxiv.org/abs/2512.02303