Training Dynamics of Learning 3D-Rotational Equivariance

2025-12-02

AI Summary

Existing studies lack quantitative analysis of how deep learning models learn 3D rotational equivariance during training for high-dimensional molecular tasks. Method: We propose a differentiable, scale-invariant equivariance error metric based on convex loss functions, and systematically evaluate it across flow matching, force-field prediction, and voxel denoising. Contribution/Results: We find that equivariance learning converges significantly faster than the primary task—equivariance error drops to ≤2% of the retained loss within 1k–10k optimization steps. Its optimization trajectory is more stable, with better-conditioned loss landscapes, and these findings hold robustly across varying model sizes and dataset scales. This work provides the first empirical evidence that rotational equivariance can be efficiently and implicitly acquired as an inductive bias during training, revealing a fundamental mechanism by which symmetry emerges from data-driven optimization. Our results establish a new paradigm for both theoretical understanding and practical design of symmetry-aware molecular models.

Technology Category

• Machine Learning: Learning with Manifolds
• Search and Optimization: Learning to Search
• Computer Vision: Learning & Optimization for CV

Application Category

• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs

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.

Problem

• Measure equivariance error in symmetry-agnostic models
• Analyze 3D-rotation equivariance learning in molecular tasks
• Compare performance of equivariant and non-equivariant models

Innovation

• Derived principled measure of equivariance error for convex losses
• Models achieve ≤2% equivariance error quickly in 3D-rotation tasks
• Non-equivariant models may outperform unless efficiency gap is narrowed