Approximate Equivariance via Projection-Based Regularisation

Equivariance is a powerful inductive bias in neural networks, improving generalisation and physical consistency. Recently, however, non-equivariant models have regained attention, due to their better runtime performance and imperfect symmetries that might arise in real-world applications. This has motivated the development of approximately equivariant models that strike a middle ground between respecting symmetries and fitting the data distribution. Existing approaches in this field either rely on sampling from a group, incurring a high sample complexity, or explicitly parameterise a model as a sum of an equivariant and non-equivariant network. This work instead approaches approximate equivariance via a projection-based regulariser which leverages a layer-wise orthogonal decomposition of a network's layers into equivariant and non-equivariant components. In contrast to existing methods, this penalises non-equivariance at an operator level across the full group orbit, rather than point-wise as in sample-based approaches. We present a mathematical framework for computing the non-equivariance penalty exactly and efficiently in both the spatial and spectral domains. In our experiments, our method consistently outperforms prior approximate equivariance approaches in both model performance and efficiency, achieving substantial runtime gains over sample-based regularisers.

When a molecule is rotated in 3D space, forces acting on each atom rotate equivalently while the total potential energy remains constant. Hence, when modelling these quantities using machine learning techniques, our models must obey similar transformation rules. However, enforcing these constraints comes with drawbacks, mainly computational cost. In this work, we propose a framework for building models which are flexible and cheap but still carry some physical inductive bias.