Equivariant non-linear maps for neural networks on homogeneous spaces

TL;DR

This work develops a universal framework for nonlinear equivariant neural network layers on homogeneous spaces by introducing nonlinear integral maps that respect group symmetries. It extends the linear theory of $G$-CNNs to arbitrary nonlinear operators via a steerable, kernel-based integral representation, and proves a universal construction with kernel-equivalence classes arising from redundancy. The framework subsumes and explains a spectrum of architectures, including $G$-CNNs, implicit steerable kernels, standard and relative self-attention, and LieTransformer, as concrete instantiations. This unification enables principled design of new equivariant nonlinear layers and provides a bridge between symmetry theory and practical nonlinear architectures with potential for broad applicability and efficient implementations.

Abstract

This paper presents a novel framework for non-linear equivariant neural network layers on homogeneous spaces. The seminal work of Cohen et al. on equivariant $G$-CNNs on homogeneous spaces characterized the representation theory of such layers in the linear setting, finding that they are given by convolutions with kernels satisfying so-called steerability constraints. Motivated by the empirical success of non-linear layers, such as self-attention or input dependent kernels, we set out to generalize these insights to the non-linear setting. We derive generalized steerability constraints that any such layer needs to satisfy and prove the universality of our construction. The insights gained into the symmetry-constrained functional dependence of equivariant operators on feature maps and group elements informs the design of future equivariant neural network layers. We demonstrate how several common equivariant network architectures - $G$-CNNs, implicit steerable kernel networks, conventional and relative position embedded attention based transformers, and LieTransformers - may be derived from our framework.

Figures (10)

• Figure 1: Architectures and models as special cases of our proposed framework.
• Figure 2: Specialisations of the general framework based on the integrand $\hat{\omega}:\mathcal{I}_\rho \times G\to V_\sigma$ to well-established architectures for different choices of the map $\hat{\omega}$. The $G$-CNN case is obtained when $\hat{\omega}$ factorises in a kernel $\hat{\kappa}: G \to \mathrm{Hom}(V_\rho,V_\sigma)$, which is independent of the input features, and the input feature $g^{-1}f:G \to V_\rho$. The other cases also appears through different factorisations as some type of kernel contracted with the feature $g^{-1}f$. However, in all cases the kernel is also dependent on the input feature $g^{-1}f$, which separates them from the $G$-CNN case and makes them inherently non-linear in $f$.
• Figure 3: A visualisation of the map $\mathrm{h}:G/H \times G\to H$ encoding the twist in the fibre.
• Figure 4: The Cartan mixing diagram for the principal $H$-bundle $G \to G/H$ and the two equivalent representations of data: either as an element $f_G$ of the induced representation $\mathcal{I}_\rho:\, G \,\to\, V_\rho$, or as a section $f_X \,\in\, \Gamma(G \times _\rho V_\rho)$ of the associated vector bundle.
• Figure 5: The Cartan mixing diagram perspective on equivalent formulations of the integral kernels of \ref{['thm:non-equivariant-kernels']}. Note that all solid arrows commute. The kernel can be viewed either as an element of the sections of an associated vector bundle $\kappa_X\in\Gamma((G\times G)\times_{\sigma\otimes \rho^*}\mathrm{Hom}(V_\rho,V_\sigma))$ or, equivalently, as elements of the induced representation $\kappa\in\mathcal{I}_{\sigma \otimes \rho^*}$.

Theorems & Definitions (109)

• Definition 1: Actions, representations
• Remark 1
• Remark 2
• Definition 2: Orbits, quotients
• Remark 3
• Example 1
• Definition 3: Equivariant, invariant maps, intertwiners
• Remark 4
• Definition 4: Homogeneous space
• Proposition 1