Training Dynamics of Learning 3D-Rotational Equivariance

Max W Shen, Ewa Nowara, Michael Maser, Kyunghyun Cho

Track: Track 1: Original Research/Position/Education/Attention Track

Keywords: Data augmentation, equivariance, training dynamics

TL;DR: Non-equivariant models learn 3D-rotational equivariance very quickly and easily, measured by percent loss

Abstract: We investigate how symmetry-agnostic models learn symmetries with data augmentation, 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 rapidly become nearly equivariant 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. We then theoretically characterize learning dynamics for models that are nearly equivariant, as "stochastic equivariant learning dynamics''. 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.

Submission Number: 205