Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Recurrent Equivariant Constraint Modulation: Learning Per-Layer Symmetry Relaxation from Data

Stefanos Pertigkiozoglou, Mircea Petrache, Shubhendu Trivedi, Kostas Daniilidis

TL;DR

This work tackles the challenge that strict equivariance can hinder learning by introducing Recurrent Equivariant Constraint Modulation (RECM), a per-layer mechanism that automatically learns the appropriate level of symmetry relaxation from data. Each layer combines an equivariant component with multiple unconstrained terms, whose relaxation weights are updated via a learned recurrence, with the limit state $h^*$ bounding the relaxation accordingly. The authors prove that when the input-target distribution is fully symmetric, the modulation converges to zero (fully equivariant), while non-symmetric distributions allow learned non-equivariant solutions; they demonstrate empirical gains across exact and approximate symmetry tasks, including large-scale conformer generation on GEOM-Drugs. RECM thus eliminates the need for task-specific relaxation schedules and achieves adaptive expressivity with only modest training overhead, improving performance relative to prior relaxation approaches. $|h^*|$ and $|\alpha_i^*|$ are bounded by symmetry-aware distances such as $W_1(p,p_G)$, ensuring principled behavior across task types.

Abstract

Equivariant neural networks exploit underlying task symmetries to improve generalization, but strict equivariance constraints can induce more complex optimization dynamics that can hinder learning. Prior work addresses these limitations by relaxing strict equivariance during training, but typically relies on prespecified, explicit, or implicit target levels of relaxation for each network layer, which are task-dependent and costly to tune. We propose Recurrent Equivariant Constraint Modulation (RECM), a layer-wise constraint modulation mechanism that learns appropriate relaxation levels solely from the training signal and the symmetry properties of each layer's input-target distribution, without requiring any prior knowledge about the task-dependent target relaxation level. We demonstrate that under the proposed RECM update, the relaxation level of each layer provably converges to a value upper-bounded by its symmetry gap, namely the degree to which its input-target distribution deviates from exact symmetry. Consequently, layers processing symmetric distributions recover full equivariance, while those with approximate symmetries retain sufficient flexibility to learn non-symmetric solutions when warranted by the data. Empirically, RECM outperforms prior methods across diverse exact and approximate equivariant tasks, including the challenging molecular conformer generation on the GEOM-Drugs dataset.

Recurrent Equivariant Constraint Modulation: Learning Per-Layer Symmetry Relaxation from Data

TL;DR

This work tackles the challenge that strict equivariance can hinder learning by introducing Recurrent Equivariant Constraint Modulation (RECM), a per-layer mechanism that automatically learns the appropriate level of symmetry relaxation from data. Each layer combines an equivariant component with multiple unconstrained terms, whose relaxation weights are updated via a learned recurrence, with the limit state bounding the relaxation accordingly. The authors prove that when the input-target distribution is fully symmetric, the modulation converges to zero (fully equivariant), while non-symmetric distributions allow learned non-equivariant solutions; they demonstrate empirical gains across exact and approximate symmetry tasks, including large-scale conformer generation on GEOM-Drugs. RECM thus eliminates the need for task-specific relaxation schedules and achieves adaptive expressivity with only modest training overhead, improving performance relative to prior relaxation approaches. and are bounded by symmetry-aware distances such as , ensuring principled behavior across task types.

Abstract

Equivariant neural networks exploit underlying task symmetries to improve generalization, but strict equivariance constraints can induce more complex optimization dynamics that can hinder learning. Prior work addresses these limitations by relaxing strict equivariance during training, but typically relies on prespecified, explicit, or implicit target levels of relaxation for each network layer, which are task-dependent and costly to tune. We propose Recurrent Equivariant Constraint Modulation (RECM), a layer-wise constraint modulation mechanism that learns appropriate relaxation levels solely from the training signal and the symmetry properties of each layer's input-target distribution, without requiring any prior knowledge about the task-dependent target relaxation level. We demonstrate that under the proposed RECM update, the relaxation level of each layer provably converges to a value upper-bounded by its symmetry gap, namely the degree to which its input-target distribution deviates from exact symmetry. Consequently, layers processing symmetric distributions recover full equivariance, while those with approximate symmetries retain sufficient flexibility to learn non-symmetric solutions when warranted by the data. Empirically, RECM outperforms prior methods across diverse exact and approximate equivariant tasks, including the challenging molecular conformer generation on the GEOM-Drugs dataset.
Paper Structure (17 sections, 4 theorems, 47 equations, 5 figures, 4 tables)

This paper contains 17 sections, 4 theorems, 47 equations, 5 figures, 4 tables.

Table of Contents

  1. Introduction
  2. Related Work
  3. Preliminaries
  4. Method
  5. Recurrent Equivariant Constraint Modulation
  6. Design of update function $\l_\theta$
  7. Experiments
  8. Ablation Study
  9. Comparisons On Tasks Requiring Different Level of Relaxation
  10. Conformer Generation
  11. Conclusion
  12. Proof of main results
  13. Extension to multivariable state $h_t$
  14. Modulation parameters evolution
  15. Implementation Details
  16. ...and 2 more sections

Key Result

Lemma 4.1

[Proof provided in Appendix proof:l1] Assume at time $t$ we sample $(z_t,y_t)$ from a distribution with density $p_t$, where $p_t$ converges in $1$-Wasserstein distance to distribution $p$. Also assume that $l_{\theta_t}$ is bounded, Lipschitz continuous and converges uniformly to $l^*$. Then we hav

Figures (5)

  • Figure 1: [Left] In our proposed Recurrent Equivariant Constraint Modulation (RECM) framework, the equivariant layer is relaxed by the addition of non-equivariant terms modulated by $\alpha_{i,t}$. These modulation parameters are controlled by an optimization state vector $h_t$, which is recurrently updated at each training step through the update rule described in Equation \ref{['eq:update']} (visualized as $\phi$ here).[Right] RECM guarantees that in cases of fully symmetric input-target distributions the non-equivariant modulation parameters $\alpha_i$ converge to zero, recovering fully equivariant layers, while in non-symmetric cases the network has flexibility to learn approximate equivariant solutions.
  • Figure 2: Evolution of relaxation modulation parameter $\alpha_w$ of RECM applied on ModelNet shapenet2015 classification of a full $\mathrm{SO}(3)$ symmetric dataset (input shapes with random poses) and a non-symmetric aligned dataset ( input shapes with pre-aligned poses)
  • Figure 3: Coverage Precision and Recall achieved on GEOM-DRUGS ($\delta=0.75\text{\AA}$) for different method along with there combined F1 score. The F1 score of each point of the two dimensional diagram is visualized by the underlying color, with green regions corresponding to points with higher F1-score values. Additionally, the diagonal lines correspond to points with the same F1 scores.
  • Figure 4: Evolution of relaxation modulation parameter $\alpha_b,\alpha_\text{noise}$ of RECM applied on ModelNet40 shapenet2015 classification of a fully $\mathrm{SO}(3)$ symmetric dataset (where the input shapes have random poses) and a non-symmetric aligned dataset (where the input shapes have pre-aligned poses)
  • Figure 5: Evolution of relaxation modulation parameter $\alpha_W$ for different sizes of generating set $C_G$, when applying RECM on the VN-PointNet model trained on the ModelNet40 classification with a full symmetric SO(3) dataset and an aligned dataset.

Theorems & Definitions (8)

  • Lemma 4.1: Convergence of $h_t$
  • Theorem 4.2
  • Lemma 4.3
  • proof
  • proof
  • proof
  • Lemma 1.1
  • proof