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

TL;DR

This paper introduces RECM, a novel layer-wise mechanism that learns to relax equivariance constraints in neural networks based on data, improving generalization and flexibility without prior tuning.

Contribution

RECM automatically learns appropriate symmetry relaxation levels from data, eliminating the need for task-specific tuning of equivariance constraints.

Findings

RECM converges to layer-specific relaxation levels bounded by symmetry gaps.

Layers with symmetric distributions recover full equivariance.

RECM outperforms prior methods on diverse equivariant tasks, including molecular conformer generation.

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.