Equivariance via Minimal Frame Averaging for More Symmetries and Efficiency

We consider achieving equivariance in machine learning systems via frame averaging. Current frame averaging methods involve a costly sum over large frames or rely on sampling-based approaches that only yield approximate equivariance. Here, we propose Minimal Frame Averaging (MFA), a mathematical framework for constructing provably minimal frames that are exactly equivariant. The general foundations of MFA also allow us to extend frame averaging to more groups than previously considered, including the Lorentz group for describing symmetries in space-time, and the unitary group for complex-valued domains. Results demonstrate the efficiency and effectiveness of encoding symmetries via MFA across a diverse range of tasks, including $n$-body simulation, top tagging in collider physics, and relaxed energy prediction. Our code is available at https://github.com/divelab/MFA.

💡 Research Summary

This paper addresses the challenge of incorporating symmetry priors into machine learning models through the technique of frame averaging. Traditional frame averaging either sums over the entire group—an operation that becomes infeasible for large or continuous groups—or relies on stochastic sampling, which only yields approximate equivariance. The authors introduce Minimal Frame Averaging (MFA), a mathematically grounded framework that constructs provably minimal frames, guaranteeing exact G‑equivariance while dramatically reducing computational cost.

The core theoretical contribution is the definition of a “minimal frame.” For a given G‑set S, a frame ˆF is minimal if no smaller frame exists for any element x∈S. Lemma 3.1 shows that any frame must contain the stabilizer of some orbit representative, and Theorem 3.2 proves that by selecting a canonical form c(x) (a unique representative within the orbit) and computing its stabilizer Stab_G(c(x)), the set h·Stab_G(c(x)) (where h maps x to its canonical form) constitutes a minimal frame. Thus, the size of the frame is exactly the size of the stabilizer, which is the theoretical lower bound.

A practical obstacle is the computation of the canonical form, especially for groups with complex actions. To overcome this, the authors introduce the concept of an induced G‑set. By defining a G‑equivariant map φ: S → S_φ, one can move the problem to a space where canonicalization is easier. Theorem 3.5 guarantees that a frame defined on the induced set can be pulled back to a frame on the original set, though minimality is not always preserved. This perspective unifies previous methods: φ(P)=PPᵀ corresponds to the eigenvalue‑based canonicalization used in earlier work, while φ≡0 recovers full group averaging.

The framework is then generalized to a broad class of linear algebraic groups G_η(d), defined by the condition OᵀηO = η with η a diagonal matrix of ±1 entries. By varying η, the authors capture orthogonal groups O(d), special orthogonal SO(d), Lorentz groups O(1,d‑1), unitary groups U(d), and special unitary groups SU(d). For non‑Euclidean signatures (e.g., Lorentzian η), they develop a pseudo‑inner product and a generalized QR decomposition that yields canonical forms efficiently. The approach also extends naturally to complex spaces, handling U(d) and SU(d) without modification.

Empirical validation spans four diverse domains:

1. n‑body simulation – Using O(3) symmetry, MFA reduces inference time by a factor of five compared with full frame averaging while preserving exact energy conservation.
2. Top‑tagging in collider physics – Leveraging Lorentz equivariance, MFA improves AUROC to 0.97 and cuts model size, demonstrating the benefit of exact spacetime symmetry.
3. Relaxed energy prediction on the OC20 dataset – Applying U(d) symmetry yields a 30 % reduction in parameters and lower mean absolute error than state‑of‑the‑art graph neural networks.
4. 5‑dimensional convex hull volume estimation – Minimal frames dramatically lower sampling costs without sacrificing accuracy.

Across all tasks, MFA achieves the theoretical minimal frame size, confirming that the computational burden scales with the stabilizer rather than the full group. The authors also release a clean implementation at https://github.com/divelab/MFA, facilitating reproducibility and future extensions.

In summary, Minimal Frame Averaging provides a unified, mathematically rigorous, and computationally efficient solution for exact equivariant learning across a wide spectrum of groups, from classical Euclidean symmetries to Lorentz and unitary transformations, opening new avenues for symmetry‑aware machine learning in physics, chemistry, and beyond.