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

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.

Equivariant neural networks constrain the structure of their internal learnable weight matrices to guarantee that a geometric transformation of the input results in a corresponding predefined transformation of its output. These constraints allow for better generalization to transformed inputs, but can also make the learning process more difficult. Additionally, if these constraints are misspecified, they can cause a drop in overall network performance. In this work, we propose to avoid the above limitations by relaxing the equivariant constraints during training. Specifically, we control the level of relaxation via a per-layer modulation parameter, which is updated at each learning step by a learnable update rule. By design, our update rule guarantees that our constraint modulation will converge to an equivariant layer if the layer's input-target output distribution is symmetric, and to a possibly non-equivariant layer otherwise. Here, we refer to an input-target output distribution as symmetric if the input-target output pairs transform according to the predefined transformation. Additionally, we show that the level of relaxation to which the constraint modulation converges is analogous to how close the distribution is to being symmetric. We evaluate our method on both symmetric and non-symmetric tasks, showing improvements over previous state-of-the-art methods.