Categorification of Group Equivariant Neural Networks
We present a novel application of category theory for deep learning. We show
how category theory can be used to understand and work with the linear layer
functions of group equivariant neural networks whose layers are some tensor
power space of $\mathbb{R}^{n}$ for the groups $S_n$, $O(n)$, $Sp(n)$, and
$SO(n)$. By using category theoretic constructions, we build a richer structure
that is not seen in the original formulation of these neural networks, leading
to new insights. In particular, we outline the development of an algorithm for
quickly computing the result of a vector that is passed through an equivariant,
linear layer for each group in question. The success of our approach suggests
that category theory could be beneficial for other areas of deep learning.
Edward Pearce-Crump
# Categorification of Group Equivariant Neural Networks

###### Abstract

We present a novel application of category theory for deep learning. We show how category theory can be used to understand and work with the linear layer functions of group equivariant neural networks whose layers are some tensor power space of \(\mathbb{R}^{n}\) for the groups \(S_{n}\), \(O(n)\), \(Sp(n)\), and \(SO(n)\). By using category theoretic constructions, we build a richer structure that is not seen in the original formulation of these neural networks, leading to new insights. In particular, we outline the development of an algorithm for quickly computing the result of a vector that is passed through an equivariant, linear layer for each group in question. The success of our approach suggests that category theory could be beneficial for other areas of deep learning.

## 1 Introduction

Despite the numerous advances that have been made in many areas of deep learning, it is well known that the field is lacking a rigourous theoretical framework to support the applications that are being developed. Practitioners typically spend a significant amount of time and effort searching for a neural network architecture that works well for the problem that they wish to solve. The architectures are often designed using heuristics that have been shown to work well in practice, despite them being poorly understood in theory. Notably, small modifications to the architecture can often result in a significant reduction in performance.

It has only become apparent very recently to researchers in the deep learning community that there is potential for category theory to provide a new set of tools for developing the theory of deep learning. The hope is that category theory will provide the rigourous theoretical framework in which all existing and future results can be placed. Category theory is based on the core concept of _compositionality_; that complex systems can be built out of smaller parts, and that the entire system can be understood by studying these smaller parts. Category theory was first used in pure mathematics in the 1940s as a way of establishing a higher-level structure for understanding a number of algebraic objects (sets, vector spaces, topological spaces etc.) and their maps (functions, linear maps, continuous maps etc.) that shared similar characteristics. It has since been applied successfully to many other disciplines of science, such as physics, chemistry and computer science. Given that many deep learning architectures share similar characteristics, in that they are typically built out of layers and maps between these layers, it is no surprise that deep learning researchers are looking to category theory to achieve a similar outcome for their own field.

In this paper, we present a novel application of category theory for deep learning. We show that a number of group equivariant neural networks - for the groups \(S_{n}\), \(O(n)\), \(Sp(n)\) and \(SO(n)\) - whose layers are some tensor power of \(\mathbb{R}^{n}\) that have recently appeared in the literature (Maron et al. (2019); Pearce-Crump (2022a,b); Godfrey et al. (2023)) can be understood in category theoretic
terms. We call this the "Categorification of Group Equivariant Neural Networks" because, in proving this result, we replace a number of set-theoretic constructions with category theoretic notions that results in a deeper structure for understanding and working with the layer functions of the neural networks themselves. We wish to emphasise that the outcome of this process is not simply a case of rewriting the existing results in a different language, but that, crucially, we obtain new insights into these neural networks from the richer structure that is established. One particularly important consequence that we show is that any of the weight matrices that appear in the neural networks in question can be understood solely by using a certain type of combinatorial diagram that has a string-like quality to it. By pulling on the strings or dragging their ends to different locations, we can use category theory to obtain new results for these group equivariant neural networks. We describe a very powerful example of this idea, where the properties of the categorification lead to a recovery of the algorithm - using a very different method - proposed by Godfrey et al. (2023) for computing the result of a vector that is passed through a symmetric group equivariant linear layer. We suggest that our approach can be adapted to obtain an algorithm for computing the same procedure for the other groups mentioned in this paper; this result will appear in another paper by the same authors.

## 2 Preliminaries

We choose our field of scalars to be \(\mathbb{R}\) throughout. Tensor products are also taken over \(\mathbb{R}\), unless otherwise stated. Also, we let \([n]\) represent the set \(\{1,\ldots,n\}\).

Recall that a representation of a group \(G\) is a choice of vector space \(V\) over \(\mathbb{R}\) and a group homomorphism

\[\rho_{V}:G\to GL(V) \tag{1}\]

Furthermore, recall that a map \(\phi:V\to W\) between two representations of \(G\) is said to be \(G\)-equivariant if, for all \(g\in G\) and \(v\in V\),

\[\phi(\rho_{V}(g)[v])=\rho_{W}(g)[\phi(v)] \tag{2}\]

We denote the set of all _linear_\(G\)-equivariant maps between \(V\) and \(W\) by \(\operatorname{Hom}_{G}(V,W)\). It can be shown that \(\operatorname{Hom}_{G}(V,W)\) is a vector space over \(\mathbb{R}\). See Segal (2014) for more details.

### Tensor Power Spaces as Group Representations

The groups of interest, namely, \(S_{n}\), \(O(n)\), \(Sp(n)\), and \(SO(n)\), can all be viewed as subgroups of \(GL(n)\). We use the symbol \(G\) to refer to any of these groups in the following. Recall that \(\mathbb{R}^{n}\) has a standard basis that is given by \(\{e_{i}\mid i\in[n]\}\), where \(e_{i}\) has a \(1\) in the \(i^{\text{th}}\) position and is \(0\) otherwise.

(Note that if \(G=Sp(n)\), then \(n=2m\), and we label the indices by \(1,1^{\prime},\ldots,m,m^{\prime}\) and call the standard basis of \(\mathbb{R}^{n}\) the symplectic basis.)

There exists a (left) action of \(G\) on \(\mathbb{R}^{n}\) that is given by left multiplication on the standard basis, which can be extended linearly to obtain a representation \(G\to GL(\mathbb{R}^{n})\).

Moreover, since the elements

\[e_{I}\coloneqq e_{i_{1}}\otimes e_{i_{2}}\otimes\cdots\otimes e_{i_{k}} \tag{3}\]

for all \(I\coloneqq(i_{1},i_{2},\ldots,i_{k})\in[n]^{k}\) form a basis of \((\mathbb{R}^{n})^{\otimes k}\), the \(k\)-tensor power space of \(\mathbb{R}^{n}\), there also exists a (left) action of \(G\) on \((\mathbb{R}^{n})^{\otimes k}\) that is given by

\[g\cdot e_{I}\coloneqq ge_{i_{1}}\otimes ge_{i_{2}}\otimes\cdots\otimes ge_{i_{ k}} \tag{4}\]

Again, this action can be extended linearly to obtain a representation \(\rho_{k}:G\to GL((\mathbb{R}^{n})^{\otimes k})\).

Figure 1: Examples of \((6,4)\)–partition diagrams. b) is also a \((6,4)\)–Brauer diagram, and c) is also a \(10\backslash 6\)–diagram.
We are interested in the space of \(G\)-equivariant linear maps between any two tensor power spaces of \(\mathbb{R}^{n}\), \(\operatorname{Hom}_{G}((\mathbb{R}^{n})^{\otimes k},(\mathbb{R}^{n})^{\otimes l})\), since these maps are the linear layer functions in the group equivariant neural networks of interest.

### Set Partition Diagrams

Pearce-Crump (2022a,b) showed that, for the groups \(G\) in question, \(\operatorname{Hom}_{G}((\mathbb{R}^{n})^{\otimes k},(\mathbb{R}^{n})^{\otimes l})\) can be constructed from certain set partitions of \([l+k]\), and in particular, from their corresponding set partition diagrams. We review these constructions below.

For \(l,k\in\mathbb{Z}_{\geq 0}\), consider the set \([l+k]\coloneqq\{1,\ldots,l+k\}\) having \(l+k\) elements. We can create a set partition of \([l+k]\) by partitioning it into a number of subsets. We call the subsets of a set partition _blocks_. Let \(\Pi_{l+k}\) be the set of all set partitions of \([l+k]\). Then, for each set partition \(\pi\) in \(\Pi_{l+k}\), we can associate to it a diagram \(d_{\pi}\), called a \((k,l)\)-partition diagram, consisting of two rows of vertices and edges between vertices such that there are

* \(l\) vertices on the top row, labelled left to right by \(1,\ldots,l\)
* \(k\) vertices on the bottom row, labelled left to right by \(l+1,\ldots,l+k\), and
* the edges between the vertices correspond to the connected components of \(\pi\).

As a result, \(d_{\pi}\) represents the equivalence class of all diagrams with connected components equal to the blocks of \(\pi\).

There are special types of \((k,l)\)-partition diagrams that we are interested in, namely:

* A \((k,l)\)-Brauer diagram \(d_{\beta}\) is a \((k,l)\)-partition diagram where the size of every block in \(\beta\) is exactly two.
* Given \(k\) and \(l\), an \((l+k)\backslash n\)-diagram \(d_{\alpha}\) is a \((k,l)\)-partition diagram where exactly \(n\) blocks in \(\alpha\) have size one, with the rest having exactly size two. The vertices corresponding to the blocks of size one are called free vertices.

We give examples of these diagrams in Figure 1.

We can form a number of vector spaces as the \(\mathbb{R}\)-linear span of certain subsets of \((k,l)\)-partition diagrams, as follows:

* The partition space \(P^{l}_{k}(n)\) is defined to be the \(\mathbb{R}\)-linear span of the set of all \((k,l)\)-partition diagrams.
* The Brauer space \(B^{l}_{k}(n)\) is defined to be the \(\mathbb{R}\)-linear span of the set of all \((k,l)\)-Brauer diagrams.
* The Brauer-Good space \(D^{l}_{k}(n)\) is defined to be the \(\mathbb{R}\)-linear span of the set of all \((k,l)\)-Brauer diagrams together with the set of all \((l+k)\backslash n\)-diagrams.

Furthermore, we can define two \(\mathbb{R}\)-bilinear operations on \((k,l)\)-partition diagrams

composition: \[\bullet:P^{m}_{l}(n)\times P^{l}_{k}(n)\to P^{m}_{k}(n)\] (5) tensor product: \[\otimes:P^{l}_{k}(n)\times P^{m}_{q}(n)\to P^{l+m}_{k+q}(n)\] (6)

as follows:

Composition: Let \(d_{\pi_{1}}\in P^{l}_{k}(n)\) and \(d_{\pi_{2}}\in P^{m}_{l}(n)\). First, we concatenate the diagrams, written \(d_{\pi_{2}}\circ d_{\pi_{1}}\), by putting \(d_{\pi_{1}}\) below \(d_{\pi_{2}}\), concatenating the edges in the middle row of vertices, and then removing all connected components that lie entirely in the middle row of the concatenated diagrams. Let \(c(d_{\pi_{2}},d_{\pi_{1}})\) be the number of connected components that are removed from the middle row in \(d_{\pi_{2}}\circ d_{\pi_{1}}\). Then the composition is defined, using infix notation, as

\[d_{\pi_{2}}\bullet d_{\pi_{1}}\coloneqq n^{c(d_{\pi_{2}},d_{\pi_{1}})}(d_{\pi _{2}}\circ d_{\pi_{1}}) \tag{7}\]

Tensor Product: Let \(d_{\pi_{1}}\in P^{l}_{k}(n)\) and \(d_{\pi_{2}}\in P^{m}_{q}(n)\). Then \(d_{\pi_{1}}\otimes d_{\pi_{2}}\) is defined to be the \((k+q,l+m)\)-partition diagram obtained by horizontally placing \(d_{\pi_{1}}\) to the left of \(d_{\pi_{2}}\) without any overlapping of vertices.
It is clear that both of these operations are associative.

The composition and tensor product operations for \(B^{l}_{k}(n)\) are inherited from the composition and tensor product operations for \(P^{l}_{k}(n)\), defined in (5) and (6) respectively. However, the composition and tensor product operations for \(D^{l}_{k}(n)\) are rather more involved; full details of their formulation can be found in the Technical Appendix.

### Group Equivariant Linear Layers

From the vector spaces defined in Section 2.2, it is possible to obtain either a spanning set or a basis for \(\operatorname{Hom}_{G}((\mathbb{R}^{n})^{\otimes k},(\mathbb{R}^{n})^{\otimes l})\). We give the form of the the spanning sets/bases, expressed in the basis of matrix units for \(\operatorname{Hom}((\mathbb{R}^{n})^{\otimes k},(\mathbb{R}^{n})^{\otimes l})\), in the Technical Appendix. Here, we reproduce a number of results which describe the existence of a surjective map from each of the vector spaces defined in Section 2.2 onto its corresponding vector space of \(G\)-equivariant linear maps, which arises from the spanning sets/bases.

**Theorem 2.1** (Diagram Basis when \(G=S_{n}\)).: _(Godfrey et al., 2023, Theorem 5.4) For any \(k,l\in\mathbb{Z}_{\geq 0}\) and any \(n\in\mathbb{Z}_{\geq 1}\), there is a surjection of vector spaces_

\[\Theta^{l}_{k,n}:P^{l}_{k}(n)\to\operatorname{Hom}_{S_{n}}((\mathbb{R}^{n})^{ \otimes k},(\mathbb{R}^{n})^{\otimes l}) \tag{8}\]

_that is given by_

\[d_{\pi}\mapsto E_{\pi} \tag{9}\]

_for all \((k,l)\)-partition diagrams \(d_{\pi}\), where \(E_{\pi}\) is given in the Technical Appendix._

_In particular, the set_

\[\{E_{\pi}\mid d_{\pi}\text{ is a }(k,l)\text{-partition diagram having at most }n\text{ blocks}\} \tag{10}\]

_is a basis for \(\operatorname{Hom}_{S_{n}}((\mathbb{R}^{n})^{\otimes k},(\mathbb{R}^{n})^{ \otimes l})\) in the standard basis of \(\mathbb{R}^{n}\), of size \(\operatorname{B}(l+k,n)\coloneqq\sum_{t=1}^{n}\left\{\begin{smallmatrix}l+k\\ t\end{smallmatrix}\right\}\), where \(\left\{\begin{smallmatrix}l+k\\ t\end{smallmatrix}\right\}\) is the Stirling number of the second kind._

**Theorem 2.2** (Spanning set when \(G=O(n)\)).: _(Pearce-Crump, 2022b, Theorem 6.5) For any \(k,l\in\mathbb{Z}_{\geq 0}\) and any \(n\in\mathbb{Z}_{\geq 1}\), there is a surjection of vector spaces_

\[\Phi^{l}_{k,n}:B^{l}_{k}(n)\to\operatorname{Hom}_{O(n)}((\mathbb{R}^{n})^{ \otimes k},(\mathbb{R}^{n})^{\otimes l}) \tag{11}\]

_that is given by_

\[d_{\beta}\mapsto E_{\beta} \tag{12}\]

_for all \((k,l)\)-Brauer diagrams \(d_{\beta}\), where \(E_{\beta}\) is given in the Technical Appendix._

_In particular, the set_

\[\{E_{\beta}\mid d_{\beta}\text{ is a }(k,l)\text{-Brauer diagram}\} \tag{13}\]

_is a spanning set for \(\operatorname{Hom}_{O(n)}((\mathbb{R}^{n})^{\otimes k},(\mathbb{R}^{n})^{ \otimes l})\) in the standard basis of \(\mathbb{R}^{n}\), of size \(0\) when \(l+k\) is odd, and of size \((l+k-1)!!\) when \(l+k\) is even._

**Theorem 2.3** (Spanning set when \(G=Sp(n),n=2m\)).: _(Pearce-Crump, 2022b, Theorem 6.6) For any \(k,l\in\mathbb{Z}_{\geq 0}\) and any \(n\in\mathbb{Z}_{\geq 2}\) such that \(n=2m\), there is a surjection of vector spaces_

\[X^{l}_{k,n}:B^{l}_{k}(n)\to\operatorname{Hom}_{Sp(n)}((\mathbb{R}^{n})^{ \otimes k},(\mathbb{R}^{n})^{\otimes l}) \tag{14}\]

_that is given by_

\[d_{\beta}\mapsto F_{\beta} \tag{15}\]

_for all \((k,l)\)-Brauer diagrams \(d_{\beta}\), where \(F_{\beta}\) is given in the Technical Appendix._

_In particular, the set_

\[\{F_{\beta}\mid d_{\beta}\text{ is a }(k,l)\text{-Brauer diagram}\} \tag{16}\]

_is a spanning set for \(\operatorname{Hom}_{Sp(n)}((\mathbb{R}^{n})^{\otimes k},(\mathbb{R}^{n})^{ \otimes l})\), for \(n=2m\), in the symplectic basis of \(\mathbb{R}^{n}\), of size \(0\) when \(l+k\) is odd, and of size \((l+k-1)!!\) when \(l+k\) is even._
**Theorem 2.4** (Spanning set when \(G=SO(n)\)).: _(Pearce-Crump, 2022b, Theorem 6.7)_

_For any \(k,l\in\mathbb{Z}_{\geq 0}\) and any \(n\in\mathbb{Z}_{\geq 1}\), there is a surjection of vector spaces_

\[\Psi_{k,n}^{l}:D_{k}^{l}(n)\to\operatorname{Hom}_{SO(n)}((\mathbb{R}^{n})^{ \otimes k},(\mathbb{R}^{n})^{\otimes l}) \tag{17}\]

_that is given by_

\[d_{\beta}\mapsto E_{\beta} \tag{18}\]

_if \(d_{\beta}\) is a \((k,l)\)-Brauer diagram, where \(E_{\beta}\) is given in the Technical Appendix, and by_

\[d_{\alpha}\mapsto H_{\alpha} \tag{19}\]

_if \(d_{\alpha}\) is a \((k+l)\backslash n\)-diagram, where \(H_{\alpha}\) is is also given in the Technical Appendix_

_In particular, the set_

\[\{E_{\beta}\}_{\beta}\cup\{H_{\alpha}\}_{\alpha} \tag{20}\]

_where \(d_{\beta}\) is a \((k,l)\)-Brauer diagram, and \(d_{\alpha}\) is a \((l+k)\backslash n\)-diagram, is a spanning set for \(\operatorname{Hom}_{SO(n)}((\mathbb{R}^{n})^{\otimes k},(\mathbb{R}^{n})^{ \otimes l})\) in the standard basis of \(\mathbb{R}^{n}\)._

## 3 Strict \(\mathbb{R}\)-Linear Monoidal Categories and String Diagrams

We appreciate that the language of category theory is not commonplace in the machine learning literature. To aid the reader, we have provided some foundational material in the Technical Appendix. Other good references are Mac Lane (1998); Kock (2003); Turaev & Virelizier (2017).

At a basic level, category theory is concerned with objects and the relationships between objects. These relationships are called morphisms. A collection of objects and a collection of morphisms between objects (satisfying some additional conditions) form a category. We can perform operations with morphisms, such as (vertically) composing them, to form new morphisms between objects. Category theory makes it possible to abstract away specific details of structures, to focus instead on the relationships between them. We are interested not only in the relationships within a category but also in how relationships are preserved across different categories. These are described by functors.

In this paper, we are interested in categories that have a specific property called _monoidal_, and the (monoidal) functors between these categories. We will see in Section 5 that it is this property that has important implications for the group equivariant neural networks that we look at in this paper. The monoidal property gives additional structure to the way in which objects and morphisms can be related. In particular, monoidal categories have an additional operation, known as a tensor product, that allows objects and morphisms to be composed in a different way, which we call horizontal. Monoidal functors preserve the tensor product across monoidal categories.

We assume throughout that all categories are _locally small_; that is, that the collection of morphisms between any two objects is a set. In fact, all of the categories that we consider in this paper have morphism sets that are vector spaces: in particular, the morphisms between objects become linear maps. We follow the presentation given in Hu (2019) and Savage (2021) below.

### Strict \(\mathbb{R}\)-Linear Monoidal Categories

_Definition 3.1_.: A category \(\mathcal{C}\) is said to be _strict monoidal_ if it comes with a bifunctor \(\otimes:\mathcal{C}\times\mathcal{C}\to\mathcal{C}\), called the tensor product, and a unit object \(\mathds{1}\), such that, for all objects \(X,Y,Z\) in \(\mathcal{C}\), we have that

\[(X\otimes Y)\otimes Z=X\otimes(Y\otimes Z) \tag{21}\]

\[(\mathds{1}\otimes X)=X=(X\otimes\mathds{1}) \tag{22}\]

and, for all morphisms \(f,g,h\) in \(\mathcal{C}\), we have that

\[(f\otimes g)\otimes h=f\otimes(g\otimes h) \tag{23}\]

\[(1_{\mathds{1}}\otimes f)=f=(f\otimes 1_{\mathds{1}}) \tag{24}\]

where \(1_{\mathds{1}}\) is the identity morphism \(\mathds{1}\to\mathds{1}\).
_Remark 3.2_.: We can assume that all monoidal categories are strict (nonstrict monoidal categories would have isomorphisms in the place of the equalities given in Definition 3.1) owing to a technical result known as Mac Lane's Coherence Theorem. See Mac Lane (1998) for more details.

_Definition 3.3_.: A category \(\mathcal{C}\) is said to be \(\mathbb{R}\)-_linear_ if, for any two objects \(X,Y\) in \(\mathcal{C}\), the morphism space \(\operatorname{Hom}_{\mathcal{C}}(X,Y)\) is a vector space over \(\mathbb{R}\), and the composition of morphisms is \(\mathbb{R}\)-bilinear.

Combining Definitions 3.1 and 3.3, we get

_Definition 3.4_.: A category \(\mathcal{C}\) is said to be _strict \(\mathbb{R}\)-linear monoidal_ if it is a category that is both strict monoidal and \(\mathbb{R}\)-linear, such that the bifunctor \(\otimes\) is \(\mathbb{R}\)-bilinear.

Analogous to how there exists maps between sets, there exists "maps" between categories, known as functors. In particular, we are interested in the following type of functors:

_Definition 3.5_.: Suppose that \((\mathcal{C},\otimes_{\mathcal{C}},\mathds{1}_{\mathcal{C}})\) and \((\mathcal{D},\otimes_{\mathcal{D}},\mathds{1}_{\mathcal{D}})\) are two strict \(\mathbb{R}\)-linear monoidal categories.

A _strict \(\mathbb{R}\)-linear monoidal functor_ from \(\mathcal{C}\) to \(\mathcal{D}\) is a functor \(\mathcal{F}:\mathcal{C}\to\mathcal{D}\) such that

1. for all objects \(X,Y\) in \(\mathcal{C}\), \(\mathcal{F}(X\otimes_{\mathcal{C}}Y)=\mathcal{F}(X)\otimes_{\mathcal{D}} \mathcal{F}(Y)\)
2. for all morphisms \(f,g\) in \(\mathcal{C}\), \(\mathcal{F}(f\otimes_{\mathcal{C}}g)=\mathcal{F}(f)\otimes_{\mathcal{D}} \mathcal{F}(g)\)
3. \(\mathcal{F}(\mathds{1}_{\mathcal{C}})=\mathds{1}_{\mathcal{D}}\), and
4. for all objects \(X,Y\) in \(\mathcal{C}\), the map \[\operatorname{Hom}_{\mathcal{C}}(X,Y)\to\operatorname{Hom}_{\mathcal{D}}( \mathcal{F}(X),\mathcal{F}(Y))\] (25) given by \(f\mapsto\mathcal{F}(f)\) is \(\mathbb{R}\)-linear.

The following definition will also prove to be very important in what follows.

_Definition 3.6_.: Given any two (locally small) categories \(\mathcal{C}\) and \(\mathcal{D}\) (not necessarily strict \(\mathbb{R}\)-linear monoidal), a functor \(\mathcal{F}:\mathcal{C}\to\mathcal{D}\) is said to be _full_ if the map (25) is surjective for all objects \(X,Y\) in \(\mathcal{C}\).

### String Diagrams

Strict monoidal categories are particularly interesting because they can be represented by a very useful diagrammatic language known as string diagrams. As this language is, in some sense, geometric in nature, we will see that it is much easier to work with these diagrams than with their equivalent algebraic form.

_Definition 3.7_ (String Diagrams).: Suppose that \(\mathcal{C}\) is a strict monoidal category. Let \(W,X,Y\) and \(Z\) be objects in \(\mathcal{C}\), and let \(f:X\to Y\), \(g:Y\to Z\), and \(h:W\to Z\) be morphisms in \(\mathcal{C}\). Then we can represent the morphisms \(1_{X}:X\to X\), \(f:X\to Y\), \(g\circ f:X\to Z\) and \(f\otimes h:X\otimes W\to Y\otimes Z\) as diagrams in the following way:

(26)

In particular, the vertical composition of morphisms \(g\circ f\) is obtained by placing \(g\) above \(f\), and the horizontal composition of morphisms \(f\otimes h\) is obtained by horizontally placing \(f\) to the left of \(h\).

We will often omit the labelling of the objects when they are clear or when they are not important.

As an example of how useful string diagrams are when working with strict monoidal categories, the associativity of the bifunctor given in (23) becomes immediately apparent. Another, more involved, example is given by the interchange law that exists for any strict monoidal category. It can be expressed algebraically as

\[(\mathds{1}\otimes g)\circ(f\otimes\mathds{1})=f\otimes g=(f\otimes\mathds{1 })\circ(\mathds{1}\otimes g) \tag{27}\]
Without string diagrams, it is somewhat tedious to prove this result - see (Savage, 2021, Section 2.2) - but with them, the result is intuitively obvious, if we allow ourselves to deform the diagrams by pulling on the strings:

(28)

## 4 Categorification

At this point, we have defined a vector space for each \(k,l\in\mathbb{Z}_{\geq 0}\) that is the \(\mathbb{R}\)-linear span of a certain subset of \((k,l)\)-partition diagrams. However, it should be apparent that, for all values of \(k\) and \(l\), these vector spaces are all similar in nature, in that the set partition diagrams only differ by the number of vertices that appear in each row and by the connections that are made between vertices. Moreover, the astute reader may have noticed that set partition diagrams look like string diagrams. Given that string diagrams represent strict monoidal categories, and that we have a collection of vector spaces for certain subsets of set partition diagrams, this implies that we should have a number of strict \(\mathbb{R}\)-linear monoidal categories! Indeed we do; we formalise this intuition below.

### Category Definitions

We assume throughout that \(n\in\mathbb{Z}_{>0}\).

_Definition 4.1_.: We define the partition category \(\mathcal{P}(n)\) to be the category whose objects are the non-negative integers \(\mathbb{Z}_{\geq 0}=\{0,1,2,\dots\}\), and, for any pair of objects \(k\) and \(l\), the morphism space \(\operatorname{Hom}_{\mathcal{P}(n)}(k,l)\) is \(P_{k}^{l}(n)\).

The vertical composition of morphisms is given by the composition of partition diagrams defined in (5); the bifunctor (the horizontal composition of morphisms) is given by the tensor product of partition diagrams defined in (6); and the unit object is 0.

_Definition 4.2_.: We define the Brauer category \(\mathcal{B}(n)\) to be the category whose objects are the same as those of \(\mathcal{P}(n)\) and, for any pair of objects \(k\) and \(l\), the morphism space \(\operatorname{Hom}_{\mathcal{B}(n)}(k,l)\) is \(B_{k}^{l}(n)\).

The vertical composition of morphisms, the horizontal composition of morphisms and the unit object are the same as those of \(\mathcal{P}(n)\).

_Definition 4.3_.: We define the Brauer-Grood category \(\mathcal{SG}(n)\) to be the category whose objects are the same as those of \(\mathcal{P}(n)\) and, for any pair of objects \(k\) and \(l\), the morphism space \(\operatorname{Hom}_{\mathcal{SG}(n)}(k,l)\) is \(D_{k}^{l}(n)\).

The vertical composition of morphisms and the horizontal composition of morphisms are the same as those defined for \(D_{k}^{l}(n)\), which can be found in the Technical Appendix. The unit object is 0.

It is easy to show that \(\mathcal{P}(n)\), \(\mathcal{B}(n)\) and \(\mathcal{SG}(n)\) are strict \(\mathbb{R}\)-linear monoidal categories.

Also, for the four groups of interest, we can define the following category.

_Definition 4.4_.: If \(G\) is any of the groups \(S_{n},O(n),Sp(n)\) or \(SO(n)\), then we define \(\mathcal{C}(G)\) to be the category whose objects are pairs \(\{((\mathbb{R}^{n})^{\otimes k},\rho_{k})\}_{k\in\mathbb{Z}_{\geq 0}}\), where \(\rho_{k}:G\to GL((\mathbb{R}^{n})^{\otimes k})\) is the representation of \(G\) given in Section 2.1, and, for any pair of objects \(((\mathbb{R}^{n})^{\otimes k},\rho_{k})\) and \(((\mathbb{R}^{n})^{\otimes l},\rho_{l})\), the morphism space, \(\operatorname{Hom}_{\mathcal{C}(G)}(((\mathbb{R}^{n})^{\otimes k},\rho_{k}), ((\mathbb{R}^{n})^{\otimes l},\rho_{l}))\) is precisely \(\operatorname{Hom}_{G}((\mathbb{R}^{n})^{\otimes k},(\mathbb{R}^{n})^{ \otimes l})\).

The vertical composition of morphisms is given by the usual composition of linear maps, the horizontal composition of morphisms is given by the usual tensor product of linear maps, both of which are associative operations, and the unit object is given by \((\mathbb{R},1_{\mathbb{R}})\), where \(1_{\mathbb{R}}\) is the one-dimensional trivial representation of \(G\).

It can be shown that \(\mathcal{C}(G)\) is a subcategory of the category of representations of \(G\), \(\operatorname{Rep}(G)\). See the Technical Appendix for more details. In particular, it is also a strict \(\mathbb{R}\)-linear monoidal category.
### Full, Strict \(\mathbb{R}\)-Linear Monoidal Functors

Given that we have a number of categories from the vector spaces of the different types of partition diagrams, that we have a category from the group equivariant linear maps between tensor power spaces, and that we have a number of maps between these vector spaces - as seen in Section 2.3 - we should have a number of functors between the newly defined categories. Indeed, we have that

**Theorem 4.5**.: _There exists a full, strict \(\mathbb{R}\)-linear monoidal functor_

\[\Theta:\mathcal{P}(n)\to\mathcal{C}(S_{n}) \tag{29}\]

_that is defined on the objects of \(\mathcal{P}(n)\) by \(\Theta(k)\coloneqq((\mathbb{R}^{n})^{\otimes k},\rho_{k})\) and, for any objects \(k,l\) of \(\mathcal{P}(n)\), the map_

\[\operatorname{Hom}_{\mathcal{P}(n)}(k,l)\to\operatorname{Hom}_{\mathcal{C}(S_ {n})}(\Theta(k),\Theta(l)) \tag{30}\]

_is precisely the map_

\[\Theta^{l}_{k,n}:P^{l}_{k}(n)\to\operatorname{Hom}_{S_{n}}((\mathbb{R}^{n})^{ \otimes k},(\mathbb{R}^{n})^{\otimes l}) \tag{31}\]

_given in Theorem 2.1._

**Theorem 4.6**.: _There exists a full, strict \(\mathbb{R}\)-linear monoidal functor_

\[\Phi:\mathcal{B}(n)\to\mathcal{C}(O(n)) \tag{32}\]

_that is defined on the objects of \(\mathcal{B}(n)\) by \(\Phi(k)\coloneqq((\mathbb{R}^{n})^{\otimes k},\rho_{k})\) and, for any objects \(k,l\) of \(\mathcal{B}(n)\), the map_

\[\operatorname{Hom}_{\mathcal{B}(n)}(k,l)\to\operatorname{Hom}_{\mathcal{C}(O (n))}(\Phi(k),\Phi(l)) \tag{33}\]

_is the map_

\[\Phi^{l}_{k,n}:B^{l}_{k}(n)\to\operatorname{Hom}_{O(n)}((\mathbb{R}^{n})^{ \otimes k},(\mathbb{R}^{n})^{\otimes l}) \tag{34}\]

_given in Theorem 2.2._

**Theorem 4.7**.: _There exists a full, strict \(\mathbb{R}\)-linear monoidal functor_

\[X:\mathcal{B}(n)\to\mathcal{C}(Sp(n)) \tag{35}\]

_that is defined on the objects of \(\mathcal{B}(n)\) by \(X(k)\coloneqq((\mathbb{R}^{n})^{\otimes k},\rho_{k})\) and, for any objects \(k,l\) of \(\mathcal{B}(n)\), the map_

\[\operatorname{Hom}_{\mathcal{B}(n)}(k,l)\to\operatorname{Hom}_{\mathcal{C}( Sp(n))}(\Phi(k),\Phi(l)) \tag{36}\]

_is the map_

\[X^{l}_{k,n}:B^{l}_{k}(n)\to\operatorname{Hom}_{Sp(n)}((\mathbb{R}^{n})^{ \otimes k},(\mathbb{R}^{n})^{\otimes l}) \tag{37}\]

_given in Theorem 2.3._

**Theorem 4.8**.: _There exists a full, strict \(\mathbb{R}\)-linear monoidal functor_

\[\Psi:\mathcal{B}\mathcal{G}(n)\to\mathcal{C}(SO(n)) \tag{38}\]

_that is defined on the objects of \(\mathcal{B}\mathcal{G}(n)\) by \(\Psi(k)\coloneqq((\mathbb{R}^{n})^{\otimes k},\rho_{k})\) and, for any objects \(k,l\) of \(\mathcal{B}(n)\), the map_

\[\operatorname{Hom}_{\mathcal{B}\mathcal{G}(n)}(k,l)\to\operatorname{Hom}_{ \mathcal{C}(SO(n))}(\Phi(k),\Phi(l)) \tag{39}\]

_is the map_

\[\Psi^{l}_{k,n}:D^{l}_{k}(n)\to\operatorname{Hom}_{SO(n)}((\mathbb{R}^{n})^{ \otimes k},(\mathbb{R}^{n})^{\otimes l}) \tag{40}\]

_given in Theorem 2.4._

Proofs of these results are given in the Technical Appendix.

Figure 2: We can use the string-like aspect of \((k,l)\)–partition diagrams to factor them as a composition of a permutation in \(S_{k}\), a _planar_\((k,l)\)–partition diagram, and a permutation in \(S_{l}\).
## 5 Implications for Group Equivariant Neural Networks

The fullness of the functors given in Section 4.2 is especially important. This condition immediately implies that, to understand any \(G\)-equivariant linear map in \(\operatorname{Hom}_{G}((\mathbb{R}^{n})^{\otimes k},(\mathbb{R}^{n})^{\otimes l})\), it is enough to work with the subset of \((k,l)\)-partition diagrams that correspond to \(G\), since we can apply the appropriate functor to obtain the equivariant maps themselves.

Furthermore, as the \((k,l)\)-partition diagrams have a string-like aspect to them - because they are morphisms in a strict \(\mathbb{R}\)-linear monoidal category - we are able to drag and bend the strings and/or move the vertices to obtain new partition diagrams, and, in the process, new \(G\)-equivariant linear maps via the appropriate functor!

One very powerful use of this idea can be seen in Figure 2. On the left hand side, we have an arbitrary \((5,3)\)-partition diagram. Suppose that we wish to multiply a vector \(v\in(\mathbb{R}^{n})^{\otimes 5}\), expressed in the standard basis, by the matrix that the diagram corresponds to under the functor \(\Theta\) given in Theorem 4.5. We assume here that \(n\geq 4\), since the number of blocks in the set partition corresponding to the \((5,3)\)-partition diagram is \(4\). One option would be to multiply the vector by the matrix as given. However, we can vastly improve the speed of the computation by performing a number of deformations to the diagram as shown in the figure.

At each stage, we drag and bend the strings representing the connected components of the set partition to obtain a factoring of the original \((5,3)\)-partition diagram in terms of a composition of three other diagrams: a \((5,5)\)-partition diagram that is not only Brauer but also a diagram that represents a permutation in the symmetric group \(S_{5}\); another \((5,3)\)-partition diagram that is _planar_ - that is, none of the connected components in the diagram intersect each other - and, finally, a diagram representing another permutation, this time in the symmetric group \(S_{3}\).

Since the middle diagram is planar, this means that it can be decomposed as a tensor product of a number of simpler diagrams, using the strict monoidal property of \(\mathcal{P}(n)\). The tensor product decomposition is shown in Figure 3. The key point is that, under the functor \(\Theta\), the image of the planar partition diagram will be the Kronecker product of the images of these simpler diagrams, since the functor is strict \(\mathbb{R}\)-linear monoidal. This will make the multiplication of any vector by this matrix significantly quicker to perform.

Hence, we have outlined a procedure which shows that, to multiply \(v\) by the linear map that is the image of the left hand diagram under \(\Theta\)_quickly_, we can first apply a permutation to the indices of the standard basis elements appearing in the input vector \(v\), then apply the Kronecker product matrix that is the image under \(\Theta\) of the planar \((5,3)\)-partition diagram, and then finally apply another permutation to the indices of the standard basis elements in \((\mathbb{R}^{n})^{\otimes 3}\) appearing in the resulting vector.

By generalising this example to any \((k,l)\)-partition diagram, we will recover - with one key distinction - the algorithm of Godfrey et al. (2023) for applying symmetric group equivariant layer functions on tensor power spaces of \(\mathbb{R}^{n}\) to input vectors; however, we have used a very different approach to obtain it. The key distinction between the two versions comes from making the middle diagram in the composition planar. Moreover, it is not hard to see that this idea will generalise to give an algorithm for applying group equivariant linear maps to input vectors for the other groups presented in this paper. This result will appear in another paper by the same authors.

## 6 Related Work

Maron et al. (2019) studied the classification of linear permutation equivariant and invariant neural network layers. They characterised the learnable, linear, permutation equivariant layer functions in \(\operatorname{Hom}_{S_{n}}((\mathbb{R}^{n})^{\otimes k},(\mathbb{R}^{n})^{ \otimes l})\) for \(n\geq k+l\), using the orbit basis. However, Godfrey et al. (2023) discovered that using the diagram basis, first constructed by Jones (1994) in the case \(k=l\), is beneficial for permutation equivariant neural network computations. Pearce-Crump (2022) estab

Figure 3: The decomposition of the planar \((5,3)\)–partition diagram into a tensor product of simpler partition diagrams.
lished the connection between permutation equivariant linear layers and the partition algebra, using Schur-Weyl duality. They fully characterised these layer functions for all values of \(n\) and tensor power space orders, revealing that the dimension of the layer function space is not independent of \(n\). The partition algebra was introduced by Martin (1990, 1994, 1996) and expanded upon by Jones (1994). Recent papers by Benkart & Halverson (2019, 2019) and Benkart et al. (2017) show how the partition algebra can be used to construct the invariant theory of the symmetric group. Comes (2020) discovered the partition category and expressed it in terms of generators and relations. Pearce-Crump (2022b) characterised all learnable, linear, equivariant layer functions in \(\mathrm{Hom}_{G}((\mathbb{R}^{n})^{\otimes k},(\mathbb{R}^{n})^{\otimes l})\) for \(G=O(n)\), \(Sp(n)\), and \(SO(n)\), using various sets of set partition diagrams. This characterisation was adapted from the combinatorial representation theory of the Brauer algebra, first developed by Brauer (1937). Grood (1999) studied the representation theory of the Brauer-Grood algebra, while the Brauer category first appeared in Lehrer & Zhang (2012). The same authors investigated the theory behind what we have termed the Brauer-Grood category in Lehrer & Zhang (2018).

## 7 Conclusion

In this paper, we showed how category theory can be applied to the linear layer functions of group equivariant neural networks for the groups \(S_{n}\), \(O(n)\), \(Sp(n)\), and \(SO(n)\), resulting in a richer structure and a deeper understanding of these layer functions. In particular, we outlined the development of an algorithm for computing the result of a vector that is passed through an equivariant, linear layer for each group. The success of our approach suggests that category theory could be beneficial for other areas of deep learning, leading to new insights and approaches.

## 8 Acknowledgments

The author would like to thank his PhD supervisor Professor William J. Knottenbelt for being generous with his time throughout the author's period of research prior to the publication of this paper.

This work was funded by the Doctoral Scholarship for Applied Research which was awarded to the author under Imperial College London's Department of Computing Applied Research scheme. This work will form part of the author's PhD thesis at Imperial College London.
