Equivariant non-linear maps for neural networks on homogeneous spaces (2504.20974v1)

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.

• The paper introduces a framework that extends linear G-CNNs to non-linear equivariant neural network layers via integral transforms with feature-dependent kernels.
• It proves that any G-equivariant operator can be expressed as an integral transform satisfying a generalized steerability constraint.
• This unifying framework connects diverse architectures, including self-attention and LieTransformer, offering insights for designing expressive equivariant networks.

This paper introduces a novel mathematical framework for constructing and understanding non-linear equivariant neural network layers on homogeneous spaces. Building upon the established theory of linear equivariant convolutions (G-CNNs) on these spaces, the authors generalize the concept to arbitrary non-linear operators using integral transforms with input-dependent kernels.

The core idea is to represent an operator $\Phi$ mapping input feature maps $f$ (sections of a vector bundle over a homogeneous space $G/H$, equivalent to functions in an induced representation $\mathcal{I}_\rho$) to output feature maps (in $\mathcal{I}_\sigma$) as an integral:

$[\Phi f](g) = \int_{G'} \omega(f, g, g') dg'$

where $g, g' \in G$ (or possibly different groups $G, G'$) and $$\omega: \mathcal{I}_\rho \times G \times G' \to V_\sigma$$ is a map that takes the entire input feature field $f$ as an argument, in addition to points $g$ and $g'$. The space of such maps $\omega$ must satisfy a compatibility constraint $\omega(f, gh, g') = \sigma(h^{-1})\omega(f, g, g')$ for $h \in H$ to ensure the output is in the target induced representation $\mathcal{I}_\sigma$. The authors prove that any map $\Phi: \mathcal{I}_\rho \to \mathcal{I}_\sigma$ can be expressed in this integral form by choosing $\omega(f, g, g') = \delta(g')\lambda[f](g)$, where $\lambda = \Phi$ and $\delta$ is the Dirac delta distribution, interpreted in a distributional sense.

The paper then derives the crucial constraint on $\omega$ for the operator $\Phi$ to be $G$-equivariant, meaning $\Phi(kf) = k(\Phi f)$ for all group elements $k \in G$. This condition is $\omega(f, g, g') = \omega(kf, kg, g')$, which means $\omega$ must be constant on the $G$-orbits of the combined space $\mathcal{I}_\rho \times G \times G'$. This equivariance allows the map $\omega$ to be reduced to a two-argument map $\hat{\omega}(g^{-1}f, g')$, where $g^{-1}f$ represents the input feature field centered at $g$. The resulting equivariant non-linear operator takes the form:

$$[\Phi f](g) = \int_{G'} \hat{\omega}(g^{-1}f, g') dg'$$

The map $$\hat{\omega}: \mathcal{I}_\rho \times G' \to V_\sigma$$ must satisfy a generalized steerability constraint: $$\hat{\omega}(hf, g'h') = \sigma(h)\hat{\omega}(f, g')$$ for $h \in H, h' \in H'$, where the action on $f \in \mathcal{I}_\rho$ is defined by the induced representation structure. The authors show that any $G$-equivariant operator $\lambda: \mathcal{I}_\rho \to \mathcal{I}_\sigma$ can be represented by this form with $$\hat{\omega}(g^{-1}f, g') = \delta(g')\lambda[g^{-1}f](e)$$.

The paper highlights the distinction between this non-linear framework and the linear G-CNNs. In the linear case, the kernel $\hat{\kappa}(g')$ is independent of the feature map $f$, and the integrand is a simple product $\hat{\kappa}(g^{-1}g')f(g')$. The domain of the kernel can be further reduced from $G$ to the double coset space $H \setminus G / H'$. In the non-linear setting, the integrand $\hat{\omega}(g^{-1}f, g')$ depends arbitrarily on the feature map $f$, making the operation non-linear. The domain of $\hat{\omega}$ is $\mathcal{I}_\rho \times G'$, and the dependence on $f$ prevents the same kind of domain reduction seen in the linear case.

A significant contribution is demonstrating how various established equivariant neural network architectures are special cases of this framework, characterized by specific forms of the $\hat{\omega}$ map:

• G-CNNs (linear convolutions): $$\hat{\omega}(g^{-1}f, g') = \hat{\kappa}(g')[g^{-1}f](g')$$, where $\hat{\kappa}$ is a steerable kernel independent of $f$. [cohen-theory-equivariant-hom]
• Implicit Steerable CNNs [implicit-kernels]: For the Euclidean group on $\mathbb{R}^d$, this corresponds to a kernel $\kappa$ dependent on position and features: $$\hat{\omega}(g^{-1} f, g') = \kappa( g', [g^{-1}f](e), [g^{-1}f](g') ) [g^{-1} f](g')$$. This generalizes the steerability condition to include feature dependence.
• Standard Self-Attention [vaswaniAttentionAllYou2017]: For a finite set and permutation group $S_n$, with trivial $H$ representation, the Softmax attention mechanism takes the form of the general integral operator with $\hat{\omega}$ dependent on query, key, and value projections of the features, including the Softmax normalization factor.
• Relative Position Bias Self-Attention [shaw-etal-2018-self]: For the translation group on integers, this corresponds to a similar Softmax form, but with the attention score incorporating a positional bias term $\psi(g')$ that depends only on the relative position $g'$.
• LieTransformer [hutchinsonLieTransformerEquivariantSelfattention2021]: This continuous attention mechanism on homogeneous spaces, originally formulated for trivial feature representations, is shown to fit the framework with $$\hat{\omega}(g^{-1}f,g') = \frac{1}{\mathcal{Z}_\alpha(g^{-1}f)} \alpha( [g^{-1}f](e), [g^{-1}f](g'), g' ) W_V[g^{-1}f](g')$$, where $\alpha$ is the attention kernel function and $\mathcal{Z}_\alpha$ is the normalization.

The authors also discuss the mathematical foundations in a distributional sense in an appendix, providing a more rigorous backing for the universality claims.

In conclusion, the paper provides a unifying theoretical framework for non-linear equivariant operators on homogeneous spaces, demonstrating that both generalized convolutions and attention mechanisms can be seen as special cases of an integral transform with a feature-dependent kernel ($\hat{\omega}$). This framework offers insights into the structural requirements for designing novel equivariant layers and connects disparate architectures within a common mathematical language based on group theory and functional analysis. The practical implications include guiding the design of more expressive and efficient equivariant neural networks, potentially leveraging techniques like implicit kernels for implementation.