Any-dimensional equivariant neural networks
Traditional supervised learning aims to learn an unknown mapping by fitting a
function to a set of input-output pairs with a fixed dimension. The fitted
function is then defined on inputs of the same dimension. However, in many
settings, the unknown mapping takes inputs in any dimension; examples include
graph parameters defined on graphs of any size and physics quantities defined
on an arbitrary number of particles. We leverage a newly-discovered phenomenon
in algebraic topology, called representation stability, to define equivariant
neural networks that can be trained with data in a fixed dimension and then
extended to accept inputs in any dimension. Our approach is user-friendly,
requiring only the network architecture and the groups for equivariance, and
can be combined with any training procedure. We provide a simple open-source
implementation of our methods and offer preliminary numerical experiments.
Eitan Levin, Mateo Díaz
# Any-dimensional equivariant neural networks

###### Abstract

Traditional supervised learning aims to learn an unknown mapping by fitting a function to a set of input-output pairs with a _fixed dimension_. The fitted function is then defined on inputs of the same dimension. However, in many settings, the unknown mapping takes inputs in _any dimension_; examples include graph parameters defined on graphs of any size and physics quantities defined on an arbitrary number of particles. We leverage a newly-discovered phenomenon in algebraic topology, called representation stability, to define equivariant neural networks that can be trained with data in a fixed dimension and then extended to accept inputs in any dimension. Our approach is user-friendly, requiring only the network architecture and the groups for equivariance, and can be combined with any training procedure. We provide a simple open-source implementation of our methods and offer preliminary numerical experiments.

## 1 Introduction

Researchers are often interested in learning mappings that are defined for inputs of different sizes. We call such objects "_free_" as they are dimension-agnostic. Free objects are pervasive in a range of scientific and engineering fields. For example, in physics, quantities like electromagnetic fields are defined for any number of particles. In signal processing, problems such as deconvolution are well-defined for signals of any length, while image and text classification tasks are meaningful regardless of the input size. In mathematics, norms of vectors and matrices are defined for any dimension; sorting algorithms handle arrays of any length, and graph parameters, such as the max-cut value, are defined for graphs of any size.

In contrast, most existing supervised learning methods aim to learn mappings by fitting a parametrized function \(\widehat{f}\colon\mathbb{R}^{d}\to\mathbb{R}^{k}\) to a set of input-output pairs \(\mathcal{S}=\{(X_{1},y_{1}),\ldots,(X_{n},y_{n})\}\subseteq\mathbb{R}^{d} \times\mathbb{R}^{k}\) with fixed dimensions \(d\) and \(k\). Naturally, the function \(\widehat{f}\) is only defined on inputs of dimension \(d\), and _a priori_ cannot be applied to inputs in other dimensions. Practitioners use ad-hoc heuristics, such as downsampling, to resize the dimension of new inputs and match that of the learned map [15]. The lack of systematic methodology for applying learned functions to inputs of different sizes motivates the main question of this work:

_How can we learn mappings that accept inputs in any dimension?_

We answer this question for mappings that are invariant or equivariant under the action of groups. Such mappings are ubiquitous since many application domains exhibit symmetry, e.g., physical quantities are equivariant under translations and rotations because they only depend on
the relative position of particles [36], while graph parameters are equivariant under permutations because they only depend on the underlying topology rather than on its labelling [3]. We propose to learn such mappings using free equivariant neural networks. Formally, we consider a sequence of neural networks and a sequence of groups, such that a network at any level in the sequence is equivariant with respect to the group at the same level; see Figure 1. We train an equivariant network with data \(\mathcal{S}\) in a fixed dimension and leverage symmetry to extend it to other dimensions.

**Core contributions.** The question posed above presents three key challenges. First, how do we encode infinite sequences of neural networks, one in each dimension, using only finitely-many parameters? Doing so is necessary if we want to learn from a finite amount of data. Second, how do we ensure that the networks we learn from data in a fixed dimension generalize well to higher dimensions? Third, is there a user-friendly procedure for learning these networks? We proceed to tackle these challenges.

_Free equivariant neural networks._ Equivariance is the key ingredient in tackling the first challenge. The authors of [25] showed that the dimension of permutation-equivariant linear layers between tensors is independent of the size of the tensors, and gave free bases for such linear layers. For example, the space of permutation-equivariant linear layers for graph adjacency matrices is 15-dimensional for any \(n\), and its basis contains the identity, taking a transpose and extracting the diagonal, among others. The authors of that paper use this fact to derive finite parameterizations of linear layers

\[W=\alpha_{1}A_{1}^{(n)}+\cdots+\alpha_{15}A_{15}^{(n)} \tag{1}\]

where \(\left\{A_{1}^{(n)},\ldots,A_{15}^{(n)}\right\}\) is a basis of free equivariant maps for graphs on \(n\) nodes. This parameterization allows them to instantiate linear layers in any dimension using only the fifteen parameters \(\alpha_{i}\), which can be learned from data in a fixed dimension. Our **first contribution** is showing that this is not a coincidence but rather stems from a general, recently-identified phenomenon known as _representation stability_. We utilize this phenomenon to show that the dimension of equivariant linear layers stabilizes for a number of group actions, including those induced by (signed) permutations, the orthogonal groups \(\mathrm{O}(n)\), rotations \(\mathrm{SO}(n)\), and the Lorentz groups \(\mathrm{O}(1,n)\). We utilize this observation to derive finite parametrizations of infinite sequences of equivariant neural networks.

_Generalization across dimensions._ The authors of [25] observed that neural networks obtained using (1) do not always generalize well to dimensions different from the one used for training. Intuitively, this is due to overparametrization, as there are many ways to represent the same function in a fixed dimension, but not all of them extend correctly to other dimensions. Our **second contribution** is to introduce a compatibility condition relating the mappings in different dimensions, which has a regularization effect that often leads to better generalization. Formally, it entails to the commutativity of the diagrams in Figure 1. We further explain how to impose this condition on the network architecture.

_Computational recipe._ The authors of [25]_manually_ found the free basis elements in (1). However, for more complicated groups and spaces of linear layers, manually finding a free basis becomes prohibitive. In a fixed dimension, one can circumvent this issue using an algorithm proposed in [11] to compute a basis. For our **third contribution**, we extend this basis to any other dimension by solving a sparse linear system, enabling us to obtain free bases computationally.

**Notation.** Given two finite-dimensional vector spaces \(U\) and \(V\) we use \(\mathcal{L}(V,U)\) to denote the set of linear mappings between them. For a group \(G\) acting on both \(U\) and \(V\), we denote invariant elements as \(U^{G}=\left\{u\in U\mid g\cdot u=u\right\}\), and equivariant linear maps as \(\mathcal{L}(V,U)^{G}=\left\{A\in\mathcal{L}(V,U)\mid gAg^{-1}=A\right\}.\) The map \(\mathcal{P}_{W}\colon U\to U\) denotes the projection onto a subspace \(W\subseteq U\). The symbols \(\mathbb{1}_{n}\) and \(I_{n}\) represent the all-ones vector and the identity in dimension \(n\), respectively. The symbol \(\mathbb{R}[G]\) denotes
the set of finite linear combinations of elements in \(G\). To avoid cluttering notation, we use bold font symbols whenever possible to denote the \(n\)th element in a sequence, e.g., the symbol \(\mathbf{G}\) denotes \(G^{(n)}\). We add a "+" superscript to denote the \((n+1)\)th element, for instance, \(G^{(n+1)}\) is denoted by \(\mathbf{G}^{+}\).

**Outline.** The remainder of this section focuses on related work. Section 2 formally defines free neural networks. Section 3 introduces a compatibility condition that ensures good generalization across different dimensions. In Section 4, we describe our computational recipe to learn free neural networks from data in a fixed dimension and extend to other dimensions. Section 5 provides numerical experiments. We close the paper in Section 6 with conclusions, limitations, and future work. We defer several technical details to the appendix.

### Related work

**Equivariant learning.** The benefits of equivariant architectures first became apparent with the success of convolutional neural networks (CNNs) in computer vision [22]. Since then, equivariance has been applied to a range of applications. Examples include DeepSets [41] and graph neural networks [25, 40] using permutation equivariance to process sets and graphs; AlphaFold 2 [16] and ARES [33] using SE(3)-equivariance for protein and RNA structure prediction and assessment; steerable [8] and spherical [20] CNNs using rotation-equivariance to classify images; and physics-informed neural networks are equivariant under the symmetries of the corresponding physical systems [17]. Many of these architectures have been shown to implement a generalized notion of convolution over the groups at hand [7, 19]. Under additional assumptions, equivariant architectures have been derived using invariant theory to explicitly parametrize polynomial equivariant maps [36]. We refer the reader to [2, 24] for an introduction to equivariant deep learning.

**Dimension-free learning.** Certain architectures processing data of a particular type are defined for inputs of different sizes. Many architectures processing graph-based data update the features at each vertex by applying the same function of the features of the vertex's neighbors, hence can be applied to graphs of any size [13, 35, 40]. Convolutional neural networks (CNNs) processing signals and images convolve their inputs with filters of constant size, hence can also be applied to inputs of different sizes. Nevertheless, it has been observed that naively applying CNNs to inputs of size different from that used during training leads to artifacts, and several downsampling techniques have been proposed to resize inputs to a CNN [15]. Networks processing natural language embed words as vectors of the same length but are defined for arbitrarily long sequences of such vectors. Recurrent neural networks process such sequences one-by-one, using cycles in their architectures to sequentially combine a new input vector with a function of the previous inputs [26]. Recursive neural networks apply the same weights recursively to pairs of input vectors to combine them into one, until the entire sequence is reduced to one vector [32]. Notably, all these architectures must process the input in particular ways motivated by the specific application at hand, whereas we only need to assume that the network processes its input equivariantly.

**Representation stability.** Representation stability considers nested sequences of groups and their representation, and implies that such sequences of representations often stabilize. Specifically, there is a labelling of the irreducibles of these groups such that the decomposition of the representations in the sequence contain the same irreducibles with the same multiplicities. This phenomenon has been formalized in [5] and further studied in [6, 12, 28, 29, 31, 39]. It has been applied to study free convex sets and algebraic varieties [1, 4, 9, 23, 34], though to our knowledge, this paper is the first to apply it to equivariant deep learning. We refer the reader to [10, 30, 38] for introductions to this area.
## 2 Free neural networks

In this section, we introduce the concept of free neural networks, i.e., networks that can be instantiated in every dimension. To set the stage, let us recall the classical notion of a neural network (NN). A NN is a mapping \(f=f_{L}\circ\ldots\circ f_{1}\) where \(f_{i}\colon V_{i}\to V_{i+1}\) is a composition \(f_{i}(x)=\sigma_{i}(W_{i}x+b_{i})\) of an affine map \(x\mapsto W_{i}x+b_{i}\) and an activation map \(\sigma_{i}\colon U_{i}\to V_{i+1}\). This yields a family of mappings parametrized by the weights \(W_{i}\in\mathcal{L}(V_{i},U_{i})\) and biases \(b_{i}\in U_{i}\). In turn, these parameters are chosen to minimize a prescribed loss function.

In several settings, we want \(f\) to respect symmetries in its inputs. Formally, we require \(f\) to be equivariant with respect to the action of a group \(G\), i.e., \(f\circ g=g\circ f\) for all \(g\in G\). A natural way to obtain equivariance is by ensuring that each building block of \(f\) is equivariant, i.e., \(W_{i}\in\mathcal{L}(V_{i},U_{i})^{G}\), \(b_{i}\in U_{i}^{G}\), and \(\sigma_{i}\) is \(G\)-equivariant. NNs satisfying these properties are called _equivariant_. Equivariant NNs can be trained by computing a basis for the space of weights and biases and optimizing over the coefficients in this basis [11].

In this work, we seek equivariant NNs that extend to inputs in _any dimension_. To this end, we consider sequences \(\{\mathbf{U}_{i}\}\) and \(\{\mathbf{V}_{i}\}\) of nested1 vector spaces \(\mathbf{U}_{i}\subseteq\mathbf{U}_{i}^{+}\) and \(\mathbf{V}_{i}\subseteq\mathbf{V}_{i}^{+}\), with actions of a nested sequence of groups \(\{\mathbf{G}\}\), see Figure 1. For example, we have inclusions \(\mathbb{R}^{n}\subseteq\mathbb{R}^{n+1}\) by zero-padding on which the group of permutations on \(n\) letters \(\mathbf{G}=\mathrm{S}_{n}\) acts by permuting coordinates. These define a sequence of NNs whose weights \(\mathbf{W}_{i}\colon\mathbf{V}_{i}\to\mathbf{U}_{i}\) and biases \(\mathbf{b}_{i}\in\mathbf{U}_{i}\) increase in size. The activation functions \(\mathbf{\sigma}\) are often defined for every dimension, as the next section shows. Thus, we focus here on extending the linear layers.

Footnote 1: Formally, we have embeddings \(\mathbf{V}\hookrightarrow\mathbf{V}^{+}\) and \(\mathbf{U}\hookrightarrow\mathbf{U}^{+}\).

**Free equivariant networks.** Because the dimensions of the space of linear layers \(\mathcal{L}(\mathbf{U}_{i},\mathbf{V}_{i})\) and the vector spaces \(\mathbf{U}_{i}\) usually grow, sequences of _general_ NNs are not finitely-parameterizable and cannot be learned from a finite amount of data. The situation radically simplifies when considering equivariant networks since the dimensions of \(\mathcal{L}(\mathbf{U}_{i},\mathbf{V}_{i})^{\mathbf{G}}\) and \(\mathbf{U}_{i}^{\mathbf{G}}\) often stabilize. This was previously observed in the context of simple graph NNs by [25], and follows from a general phenomenon known as representation stability [5]. To illustrate it with a concrete example, consider the nested sequence \(\mathbf{U}=\mathbf{V}=\mathbb{R}^{n}\) with the action of \(\mathbf{G}=\mathrm{S}_{n}\) as above. Then \(\mathbf{V}^{\mathbf{G}}\) is one-dimensional and spanned by \(\mathbbm{1}_{n}\) while \(\mathcal{L}(\mathbf{U},\mathbf{V})^{\mathbf{G}}\) is two-dimensional and spanned by \(I_{n},\mathbbm{1}_{n}\mathbbm{1}_{n}^{\top}\), which are basis elements defined for every \(n\). Similarly, [25] obtained explicit free bases for spaces of invariant tensors.

Interestingly, all the above basis elements project onto each other, e.g., the orthgonal projection

Figure 1: Equivariant free neural networks. We use bold font to denote the \(n\)th element in a sequence; see Notation in Section 1 for details.
of \(\mathbbm{1}_{n}\) onto \(\mathbb{R}^{n-1}\) is \(\mathbbm{1}_{n-1}\), and similarly \(I_{n},\mathbbm{1}_{n}\mathbbm{1}_{n}^{\top}\) project onto \(I_{n-1},\mathbbm{1}_{n-1}\mathbbm{1}_{n-1}^{\top}\). Motivated by this observation, we say that a sequence of equivariant NNs is _free_ if the weights \(\mathbf{W}_{i}\colon\mathbf{V}_{i}\to\mathbf{U}_{i}\) and biases \(\mathbf{b}_{i}\in\mathbf{U}_{i}\) satisfy

\[\mathbf{W}_{i}=\mathcal{P}_{\mathbf{U}_{i}}\left.\mathbf{W}_{i}^{+}\right|_{\mathbf{V}_{i}}, \qquad\text{and}\qquad\mathbf{b}_{i}=\mathcal{P}_{\mathbf{U}_{i}}\mathbf{b}_{i}^{+}.\] (FreeNN)

For many sequences \(\{\mathbf{V}_{i}\},\{\mathbf{U}_{i}\}\), equation (FreeNN) uniquely determines \(\mathbf{W}_{i}^{+},\mathbf{b}_{i}^{+}\) from \(\mathbf{W}_{i},\mathbf{b}_{i}\), allowing us to uniquely extend a network to accept larger inputs.

**Representation stability.** To understand for which sequences (FreeNN) allows us to extend, we introduce some key definitions from the representation stability literature [5, 6].

**Definition 2.1** (Consistent sequences).: Fix a sequence of compact groups \(\mathcal{G}=\{\mathbf{G}\}_{n\in\mathbb{N}}\) such that \(\mathbf{G}\subseteq\mathbf{G}^{+}\). The family \(\mathcal{V}=\{(\mathbf{V},\mathbf{\varphi})\}_{n\in\mathbb{N}}\) is a _consistent sequence_ of \(\mathcal{G}\)-representations if the following hold true for all \(n\):

1. (**Representations**) The set \(\mathbf{V}\) is an orthogonal \(\mathbf{G}\)-representation;
2. (**Equivariant isometries**) The map \(\mathbf{\varphi}\colon\mathbf{V}\hookrightarrow\mathbf{V}^{+}\) is a \(\mathbf{G}\)-equivariant isometry.

We will identify \(\mathbf{V}\) with its image \(\mathbf{\varphi}(\mathbf{V})\subseteq\mathbf{V}^{+}\), and omit the inclusions \(\mathbf{\varphi}\) unless needed. Importantly, consistent sequences can be combined to form more complex sequences. In particular, if \(\mathcal{V}=\{(\mathbf{V},\mathbf{\varphi})\}\) and \(\mathcal{U}=\{(\mathbf{U},\mathbf{\psi})\}\) are consistent sequences of \(\{\mathbf{G}\}\)-representations, then so are their sum and tensor product:

\[\mathcal{V}\oplus\mathcal{U}=\{(\mathbf{V}\oplus\mathbf{U},\mathbf{\varphi}\oplus\mathbf{ \psi})\}\quad\text{ and }\quad\mathcal{V}\otimes\mathcal{U}=\{(\mathbf{V}\otimes\mathbf{U},\mathbf{\varphi} \otimes\mathbf{\psi})\}.\]

This follows directly from Definition 2.1. We use \(\mathcal{V}^{\otimes k}\) (resp. \(\mathcal{V}^{\otimes k}\)) to denote the direct sum (resp., tensor product) of \(\mathcal{V}\) taken \(k\) times. As a direct by-product of this observation, we obtain that the sequence of linear maps \(\mathcal{L}(\mathcal{V},\mathcal{U})=\{(\mathcal{L}(\mathbf{V},\mathbf{U}),\mathbf{\varphi }\otimes\mathbf{\psi})\}\) is consistent since it is isomorphic to \(\mathcal{V}\otimes\mathcal{U}\). The following parameter controls the complexity of consistent sequences and ensures that their spaces of invariants are eventually isomorphic.

**Definition 2.2** (Generation degree).: A consistent sequence \(\{\mathbf{V}\}\) is _generated in degree \(d\)_ if \(\mathbb{R}[\mathbf{G}]V^{(d)}=\mathbf{V}\) for all \(n\geq d\). The generation degree is the smallest such \(d\in\mathbb{N}\). The sequence is _finitely-generated_ if the generation degree is finite.

In words, a consistent sequence \(\{\mathbf{V}\}\) is generated in degree \(d\) if for all \(n\geq d\), \(\mathbf{V}\) is equal to the span of linear combinations of elements of \(\mathbf{G}\) applied to \(V^{(d)}\). For instance, if \(\mathbf{V}=\mathbb{R}^{d}\) with the action of permutations \(\mathbf{G}=\mathrm{S}_{n}\), then its generation degree is one, as any \(v\in\mathbf{V}\) can be obtained from \(e_{1}=(1,0,\ldots,0)\in V^{(1)}\subseteq\mathbf{V}\) via \(v=\sum v_{i}\,g_{i}\cdot e_{1}\) where \(g_{i}\) swaps the the first and \(i\)th components. The next proposition gives an isomorphism between spaces of invariants in a finitely-generated consistent sequence. This result was first established in the representation stability literature, albeit in a different form. We include proof for completion.

**Proposition 2.3** (Isomorphism of invariants).: _If \(\{\mathbf{V}\}\) is generated in degree \(d\), then the orthogonal projections \(\mathcal{P}_{\mathcal{V}}\colon(\mathbf{V}^{+})^{\mathbf{G}^{+}}\to\mathbf{V}^{\mathbf{G}}\) are injective for all \(n\geq d\), and isomorphisms for all large \(n\)._

Proof.: First, the projection \(\mathcal{P}_{\mathbf{V}}\) is \(\mathbf{G}\)-equivariant because \(\mathbf{G}\) acts orthogonally. Second, it maps \(\mathbf{G}^{+}\)-invariants in \(\mathbf{V}^{+}\) to \(\mathbf{G}\)-invariants in \(\mathbf{V}\) because \(\mathbf{G}\subseteq\mathbf{G}^{+}\). Third, we prove injectivity for \(n\geq d\). Suppose \(\mathcal{P}_{\mathbf{V}}v=0\) for some \(v\in(\mathbf{V}^{+})^{\mathbf{G}^{+}}\). For any \(u\in\mathbf{V}^{+}\), write \(u=\sum_{i}g_{i}u_{i}\) where \(u_{i}\in\mathbf{V}\) and \(g_{i}\in\mathbf{G}^{+}\). We then have \(\langle v,u\rangle=\langle v,\sum_{i}u_{i}\rangle=0\) because \(v\) is invariant and \(\sum_{i}u_{i}\in\mathbf{V}\). Since \(u\in\mathbf{V}^{+}\) was arbitrary, we conclude that \(v=0\). The injectivity of \(\mathcal{P}_{\mathbf{V}}\) shows that \(\dim(\mathbf{V})^{\mathbf{G}}\geq\dim(\mathbf{V}^{+})^{\mathbf{G}^{+}}\) for all \(n\geq d\), hence the sequence of dimensions \(\dim(\mathbf{V})^{\mathbf{G}}\) eventually stabilizes, at which point the projections become isomorphisms.
Note that (FreeNN) precisely requires \(\mathbf{W}^{+}\) and \(\mathbf{b}^{+}\) to project to \(\mathbf{W},\mathbf{b}\) inside the consistent sequences \(\mathcal{L}(\mathcal{V}_{i},\mathcal{U}_{i}),\mathcal{U}_{i}\), respectively. Thus, Proposition 2.3 implies that for all large \(n\), we can parameterize infinite sequences of equivariant linear layers with a finite set of parameters \(\alpha_{1},\ldots,\alpha_{\ell}\), just as in (1). This result enables us to train free NNs as we describe in Section 4.

**Examples.** We collect a few crucial examples and include additional ones in the appendix.

_Scalar sequence._ The sequence \(\mathcal{S}=\{\mathbb{R}\}\) together with \(\mathbf{\varphi}(x)=x\) yields a consistent sequence for any sequence of groups acting trivially, i.e., \(g\cdot x=x\). It is generated in degree one. By setting the last layer \(\mathbf{V}_{L+1}=\mathbb{R}\) in the free NN architecture of Figure 1, we obtain a sequence of invariant functions \(\mathbf{f}\).

_Permutation sequences._ Let \(\mathcal{R}=\{\mathbb{R}^{n}\}\) be the consistent sequence with zero-padding embeddings \(\mathbf{\varphi}(x)=(x^{\top},0)^{\top}\) and the action of the symmetric group \(\mathbf{G}=\mathrm{S}_{n}\) by permuting coordinates. Consider \(\mathcal{V}_{k}=\mathcal{U}_{k-1}=\mathcal{R}^{\otimes k}\), which is generated in degree \(k\). The dimensions of the space of weights and biases for low-order tensors are eventually

\[\dim\mathcal{L}(\mathbf{V}_{1},\mathbf{U}_{1})^{\mathbf{G}}=4,\ \dim\mathbf{U}_{1}^{\mathbf{G}}=2,\ \ \dim\mathcal{L}(\mathbf{V}_{2},\mathbf{U}_{2})^{\mathbf{G}}=52,\ \ \text{and}\ \ \dim\mathbf{U}_{2}^{\mathbf{G}}=5.\]

**When does stability kick in?** Understanding the exact level at which the projections become isomorphisms is important, yet, an exact characterization is rather technical. Stabilization occurs at the so-called **presentation degree** of the sequence, which might be larger than the generation degree in general. However, they agree for many relevant examples, such as in the sequence \(\mathcal{R}^{\otimes k}\) above. We include a formal definition and a discussion in Appendix B.

## 3 Generalization across dimensions and compatibility

Free NNs often do not generalize correctly to other dimensions. In Section 5, we provide examples where a free NN yields excellent test error in its trained dimension while exhibiting errors \(10^{10}\) times worst in other dimensions. This discrepancy arises due to overparametrization, as there are numerous ways to encode the same function in a fixed dimension, and yet not every encoding extends seamlessly. In this section, we propose a regularization strategy that often leads to better generalization across dimensions.

Specifically, inspired by a similar condition for free convex sets [23], we introduce the following compatibility condition.

**Definition 3.1** (Compatible networks).: A sequence of maps \(\{\mathbf{f}\colon\mathbf{V}\to\mathbf{U}\}\) is (intersection-) _compatible_ if \(\mathbf{f}^{+}|_{\mathbf{V}}=\mathbf{f}\). A sequence of equivariant networks (Figure 1) is compatible if

\[\mathbf{W}_{i}^{+}|_{\mathbf{V}_{i}}=\mathbf{W}_{i},\qquad\mathbf{b}_{i}^{+}=\mathbf{b}_{i},\qquad \text{and}\qquad\mathbf{\sigma}_{i}^{+}|_{\mathbf{U}_{i}}=\mathbf{\sigma}_{i}.\] (CompNN)

Intuitively, this condition ensures that applying a function in a fixed dimension and then extending to a higher dimension is equivalent to first extending and then applying the function. Layer-wise compatibility (CompNN) guarantees that the sequence of NN \(\{\mathbf{f}\colon\mathbf{V}_{1}\to\mathbf{V}_{L+1}\}\), given by Figure 1, is compatible, or equivalently, that the diagram in that figure commutes. The condition (CompNN) is strictly stronger than (FreeNN), as compatible NNs are free, but free NNs are not compatible in general.

**Compatibility in the wild.** Compatibility is a natural condition satisfied by sequences of equivariant maps \(\{\mathbf{f}\}\) arising in many problems of interest. Here we list some examples.

_Graph parameters._ Graph invariants are \(\mathrm{S}_{n}\)-equivariant because they only depend on the underlying graph structure rather than on its labeling. Several graph invariants are also compatible with
zero-padding, which is equivalent to adding an isolated vertex. For example, the max-cut value, the number of triangles, and cycles do not change if we add an isolated vertex, hence they are compatible.

_Matrix mappings._ Linear algebra operations are often compatible. For instance, the experiments in [25] aimed to learn the following matrix mappings: \(X\mapsto\frac{1}{2}(X+X^{\top})\), \(X\mapsto\operatorname{diag}(\operatorname{diag}(X))\), \(X\mapsto\operatorname{tr}(X)\), and \(X\mapsto\operatorname{argmax}_{\|v\|=1}\|Xv\|\). All of these are compatible sequences with the zero-padding embedding described in Section 2, and \(\operatorname{S}_{n}\)-equivariant.

_Orthogonal invariance._ The papers [11, 36] consider the task of learning the function

\[\mathbf{f}(x_{1},x_{2})=\sin(\|x_{1}\|)-\frac{\|x_{2}\|^{3}}{2}+\frac{x_{1}^{\top}x _{2}}{\|x_{1}\|\|x_{2}\|}, \tag{2}\]

defined on \((\mathbb{R}^{n})^{\oplus 2}\) with \(n=5\) from evaluation data. This function is \(\operatorname{O}(n)\)-invariant if we let rotations \(g\) act by \((g\cdot x_{1},g\cdot x_{2})\). Embedding \((\mathbb{R}^{n})^{\oplus 2}\) into \((\mathbb{R}^{n+1})^{\oplus 2}\) by zero-padding each vector, we see that \(\{\mathbf{f}\}\) is a compatible sequence of invariant functions.

_Orthogonal equivariance._ The \(\operatorname{O}(3)\)-equivariant task in [11, 36] consists of taking as input \(n=5\) particles in space, given by their masses and position vectors \((m_{i},x_{i})_{i=1}^{n}\in(\mathbb{R}\oplus\mathbb{R}^{3})^{\oplus n}\), and outputting their moment of inertia matrix \(\mathbf{f}(m_{i},x_{i})=\sum_{i=1}^{n}m_{i}(x_{i}^{\top}x_{i}I_{3}-x_{i}x_{i}^{ \top}).\) Embedding \(\mathbf{V}=(\mathbb{R}\oplus\mathbb{R}^{3})^{\oplus n}\) into \(\mathbf{V}^{+}\) by zero-padding and letting \(\mathbf{U}=\mathbb{S}^{3}\), we see that \(\{\mathbf{f}\}\) is a compatible sequence of maps. Furthermore, let \(\mathbf{G}=\operatorname{S}_{n}\times\operatorname{O}(3)\) act on \(\mathbf{V}\) by \((\pi,g)\cdot(m_{i},x_{i})_{i=1}^{n}=(m_{\pi^{-1}(i)},gx_{\pi^{-1}(i)})_{i=1}^{n}\), and on \(\mathbf{U}\) by \((\pi,g)\cdot X=gXg^{-1}\) for \((\pi,g)\in\mathbf{G}\) (so permutations act trivially). Then each \(\mathbf{f}\) is \(\mathbf{G}\)-equivariant.

There are many more examples of compatible mappings in the literature [23]. Therefore, we proceed to study compatible NNs. We derive conditions ensuring the compatibility of the linear layers and show that several standard activation functions are compatible.

**Compatible linear layers.** Since we only have data in some finite level \(n_{0}\), we ask when do fixed-dimensional weights \(W_{i}^{(n_{0})}\) extend to a sequence satisfying (CompNN). The set of such \(W_{i}^{(n_{0})}\) forms a linear space, and the next theorem characterizes a subspace of it and enables us to find a basis for that subspace in Section 4. We defer its proof to Appendix C.

**Assumption 1**.: The sequences \(\mathcal{V}=\{\mathbf{V}\},\ \mathcal{U}=\{\mathbf{U}\}\) are obtained from direct sums and tensor products of the same sequence \(\mathcal{V}_{0}\). The sequence \(\mathcal{V}\) is generated in degree \(d_{\mathrm{g}}\) and presented in degree \(d_{\mathrm{p}}\). The mapping \(W^{(n_{0})}\in\mathcal{L}(V^{(n_{0})},U^{(n_{0})})^{G^{(n_{0})}}\) at level \(n_{0}\geq d_{\mathrm{p}}\) satisfies

\[W^{(n_{0})}(V^{(m)})\subseteq U^{(m)}\text{ for }m\leq d_{\mathrm{g}}.\]

**Theorem 3.2**.: _Suppose that the sequences \(\mathcal{V}=\{\mathbf{V}\},\ \mathcal{U}=\{\mathbf{U}\}\) and linear map \(W^{(n_{0})}\in\mathcal{L}(V^{(n_{0})},U^{(n_{0})})^{G^{(n_{0})}}\) satisfy Assumption 1. Then there is a unique extension \(\{\mathbf{W}\}\) of \(W^{(n_{0})}\) to a sequence of equivariant linear maps satisfying (CompNN)._

Recall that we define the presentation degree in the appendix. In words, Theorem 3.2 says that if \(W^{(n_{0})}\) is equivariant and its restrictions to lower-dimensional subspaces satisfy (CompNN), then it uniquely extends to higher-dimensional weights satisfying (CompNN). In Section 4, we use this result to find a basis for such \(W^{(n_{0})}\), and use it to train compatible equivariant NNs.

**Compatible activation functions.** The majority of equivariant nonlinearities proposed in the literature are compatible, such as the following few examples.
Entrywise activations and permutations.Let \(\sigma\colon\mathbb{R}\to\mathbb{R}\) be any nonlinear function. Define \(\mathbf{\sigma}\colon(\mathbb{R}^{n})^{\otimes k}\to(\mathbb{R}^{n})^{\otimes k}\) by applying \(\sigma\) to each entry of a tensor. Then each \(\mathbf{\sigma}\) is \(\mathbf{G}=\mathrm{S}_{n}\)-equivariant and \(\{\mathbf{\sigma}\}\) is a compatible sequence.

_Bilinear layers._ The authors of [11] map \((\mathbf{V}\otimes\mathbf{U})\oplus\mathbf{V}\to\mathbf{U}\) by sending \((v\otimes u,v^{\prime})\mapsto\langle v,v^{\prime}\rangle u\), which is equivariant for _any_ group action on \(\mathbf{V},\mathbf{U}\). This sequence of maps is also compatible.

_Gated nonlinearities._ The authors of [37] map \(\mathbf{V}\oplus\mathbb{R}\to\mathbf{V}\) by \((v,\alpha)\mapsto v\sigma(\alpha)\) where \(\sigma\colon\mathbb{R}\to\mathbb{R}\) is a nonlinearity. These are equivariant maps for any group acting trivially on the \(\mathbb{R}\) component, and they form a compatible sequence.

## 4 A computational recipe for learning free neural networks

In this section, we describe an algorithm to train free and compatible NNs. We do so in two stages. First, we train our NN at a large enough level \(n_{0}\) by finding bases for the weights and biases at that level and optimizing over the coefficients in this basis. Second, for any higher level \(n>n_{0}\), we extend the trained NN at level \(n_{0}\) by solving linear systems for the higher-dimensional weights and biases; for any lower level \(n<n_{0}\) we project the NN using (FreeNN). Throughout, we fix consistent sequences \(\mathcal{V}_{i},\mathcal{U}_{i}\) and sequences of compatible equivariant nonlinearities \(\{\mathbf{\sigma}_{i}\}\). We summarize our procedure in Algorithm 1. In Appendix D, we show that our procedure provably generates free and compatible NN.

**Finding bases for free NN.** To train free networks (FreeNN), we fix \(n_{0}\)_exceeding the presentation degrees_ of \(\mathcal{V}_{i}\otimes\mathcal{U}_{i}\) and \(\mathcal{U}_{i}\) for all \(i=1,\dots,L\), and find a basis for equivariant weights and invariant biases at level \(n_{0}\) using the algorithm of [11]. Specifically, Theorem 1 in [11] states that \(b\in(U_{i}^{(n_{0})})^{G^{(n_{0})}}\) if and only if \(b\) satisfies that

\[B\,b=0\quad\left(\text{for all }B\in\mathcal{B}_{i}^{(n_{0})}\right)\qquad\text{ and}\qquad(D-I)\,b=0\quad\left(\text{for all }D\in\mathcal{D}_{i}^{(n_{0})}\right) \tag{3}\]

where \(\mathcal{B}_{i}^{(n_{0})}\) is a basis for the Lie algebra of \(G^{(n_{0})}\) and \(\mathcal{D}_{i}^{(n_{0})}\) are discrete generators for \(G^{(n_{0})}\), which are finite sets. Here \(B\) and \(D\) are represented as matrices acting on \(U_{i}^{(n_{0})}\). Since equivariant linear maps \(W\colon V_{i}^{(n_{0})}\to U_{i}^{(n_{0})}\) are just the invariants in \(\mathcal{L}(V_{i}^{(n_{0})},U_{i}^{(n_{0})})\), they constitute the kernel of the analogously-defined set of equations. Thus, finding a basis for equivariant weights and invariant biases reduces to finding a basis for the kernel of a matrix (3), which is typically large and sparse. This can be tackled either using a Krylov-subspace method [11] or a sparse LU decomposition [14, 21].

**Finding bases for Compatible NN.** If we rather want to find basis for layers satisfying (CompNN), we fix \(n_{0}\)_exceeding the presentation degrees_ of \(\mathcal{V}_{i}\) and let \(d_{i}\) be its generation degree. For most representations, (CompNN) implies zero bias, so we focus on the weights. We find a basis for weights \(W_{i}^{(n_{0})}\) satisfying the hypotheses of Theorem 3.2 by noting that Assumption 1 holds if and only if \(W_{i}^{(n_{0})}\) is \(G^{(n_{0})}\)-equivariant and satisfies

\[\left[\mathcal{P}_{V_{i}^{(m)}}\otimes\left(I-\mathcal{P}_{U_{i}^{(m)}}\right) \right]\,\text{vec}\left(W_{i}^{(n_{0})}\right)=0\quad(\text{for all }m\leq d_{i}),\]

where vec denotes the vectorization of a matrix by stacking its columns. This characterization allows us to find a basis for weights by finding a basis for the kernel of a sparse matrix.

**Extending to arbitrary dimensions.** For any \(n>n_{0}\), we extend our trained network at level \(n_{0}\) by finding the unique equivariant weights and biases at level \(n\) projecting onto the trained ones as in (FreeNN). Formally, we find the unique \(W_{i}^{(n)}\) and \(b_{i}^{(n)}\) satisfying (3) with \(n_{0}\) replaced by \(n\)
and

\[\left[\mathcal{P}_{V_{i}^{(n_{0})}}\otimes\mathcal{P}_{U_{i}^{(n_{0})}}\right]\, \operatorname{vec}\left(W_{i}^{(n)}\right)=\operatorname{vec}\left(W_{i}^{(n_{0 })}\right),\quad\text{and}\qquad\mathcal{P}_{U_{i}^{(n_{0})}}\,b_{i}^{(n)}=b_{i }^{(n_{0})}.\] (ExtSyst)

This amounts to solving a linear system, which is again typically sparse. It can be solved via iterative methods, e.g., stochastic gradient descent or LSQR [27]. For any \(n<n_{0}\), we set \(W_{i}^{(n)}\) and \(b_{i}(n)\) via (FreeNN), which is equivalent to (ExtSyst) with \(n\) and \(n_{0}\) swapped.

## 5 Numerical Experiments

In this section, we use Algorithm 1 to learn some of the compatible mappings from Section 3. Our implementation is based on that of [11], and is available at

[https://github.com/mateodd25/free-nets](https://github.com/mateodd25/free-nets).

For each experiment, we aim to learn a sequence of mappings \(\{\boldsymbol{f}\colon\boldsymbol{V}_{1}\to\boldsymbol{V}_{L+1}\}\). To do so, we fix a level \(n_{0}\) and randomly generate data \(\{(X_{i},f^{(n_{0})}(X_{i}))\}\subseteq V_{1}^{(n_{0})}\times V_{L+1}^{(n_{0 })}\). To fit the data, we consider the architectures described in the Appendix E and compute bases for weights and biases as we describe in Section 4. We then optimize over the coefficients in those bases using ADAM [18] starting from random initialization. Finally, we extend the trained network at level \(n_{0}\) to several levels \(n\). We evaluate our extended network on random test data generated at each level from the ground-truth sequence of maps.

We consider five examples from Section 3: the trace \(f(X)=\operatorname{tr}(X)\), diagonal extraction \(f(x)=\operatorname{diag}(\operatorname{diag}(X))\), symmetric projection \(f(X)=(X+X^{\top})/2\), the top right-singular vector \(f(X)=\operatorname{argmax}_{\|v\|=1}\|Xv\|\), and the function defined in (2). The first and last examples are invariant functions; the rest are equivariant mappings. For the first four examples, we set \(\boldsymbol{G}=\operatorname{S}_{n}\), and for the last one, we set \(\boldsymbol{G}=\operatorname{O}(n).\) We train with 3000 data points generated randomly and evaluate the test error at each dimension on 1000 fresh samples.

We use the mean squared error (MSE) as our loss function for all the experiments except for learning the top singular vector. For that experiment, we use the squared sine loss proposed by [25], i.e., \(\ell(\widehat{y},y)=1-\langle\widehat{y},y\rangle^{2}/\|\widehat{y}\|^{2}\|y \|^{2}\). The resulting errors in each dimension for both free and compatible NNs are shown in Figure 2. The test errors at the trained dimension are competitive with those obtained in [11, 25, 36] for learning the same mappings.

Due to memory limitations, we set the training level to \(n_{0}=3\) for the orthogonal invariance task, whereas Algorithm 1 would require \(n_{0}=6\) to uniquely extend free networks2 since the presentation degree of the linear maps between the hidden layers equal to six -- see Theorem D.1 in the appendix. This highlights an advantage of imposing compatibility -- it allows us to uniquely extend a trained network from lower-dimensional data. Moreover, we see that imposing our compatibility condition yields substantially lower errors across dimensions. Remarkably, the test error for free NNs increases
by many orders of magnitude for simple functions, a phenomenon that was previously noted by [25]. This further underscores the importance of compatibility when generalizing to other dimensions.

## 6 Conclusions, limitations, and future work

We leveraged representation stability to prove that a broad family of equivariant NNs extends to higher dimensions, which enables us to train an infinite sequence of NNs using a finite amount of data. Extending networks to higher dimensions often results in substantial test errors. To improve generalization, we introduced a compatibility condition relating networks across dimensions. We characterized compatibility and developed an algorithm to train free and compatible NNs. Finally, we applied our method to several numerical examples from the literature. In these examples, free NNs trained without imposing compatibility generalize poorly to higher dimensions, even when learning simple linear functions. In contrast, compatible NNs generalize significantly better.

We hope to address several limitations of our approach in future work. First, while our compatibility condition improves generalization, it can be restrictive in certain settings. For example, it often yields zero bias. It would be interesting to learn the relation between maps in different dimensions directly from the data rather than assume it is known in advance. Second, for examples like the leading eigenvector, the error increases substantially when extending to higher dimensions, even if we impose compatibility. It would be interesting to improve generalization to higher dimensions in these examples by imposing additional compatibility conditions. Third, extending a trained network to higher dimensions in our examples involved solving large and sparse systems, which we currently do not fully exploit. Incorporating sparsity in the toolbox of [11] will enable training larger networks.

Figure 2: Test errors across dimensions for free and compatible networks. Each experiment is run three times; the lighter bands show the max and min runs, while the bold line shows the average. For diagonal extraction and symmetric projection, we measure the average MSE per entry. Vertical gray lines mark the dimension used for learning.
## Acknowledgments

We thank Venkat Chandrasekaran for his constant support and encouragement during the development of this work. Both authors were funded by Venkat Chandrasekaran through AFOSR grant FA9550-20-1-0320. MD was partially funded by Joel Tropp through ONR BRC Award N00014-18-1-2363, and NSF FRG Award 1952777.
