A Diagrammatic Approach to Improve Computational Efficiency in Group Equivariant Neural Networks

TL;DR

The paper tackles the high computational cost of applying equivariant weight matrices between tensor-power layers in group-equivariant neural networks. It develops a diagrammatic, category-theoretic framework that realises equivariant maps as images of monoidal functors from partition/Brauer categories to group-representation categories, enabling a fast forward pass. By introducing algorithmically planar set partition diagrams and a three-stage multiplication procedure (Factor, Permute, PlanarMult), the authors obtain substantial Big-$O$ improvements across $S_n$, $O(n)$, $Sp(n)$, and $SO(n)$, with tailored group-specific implementations and complexities. This approach promises practical acceleration and broader adoption of high-order tensor-power equivarient networks in domains with symmetry, while preserving exact equivariance via diagrammatic representations.

Abstract

Group equivariant neural networks are growing in importance owing to their ability to generalise well in applications where the data has known underlying symmetries. Recent characterisations of a class of these networks that use high-order tensor power spaces as their layers suggest that they have significant potential; however, their implementation remains challenging owing to the prohibitively expensive nature of the computations that are involved. In this work, we present a fast matrix multiplication algorithm for any equivariant weight matrix that maps between tensor power layer spaces in these networks for four groups: the symmetric, orthogonal, special orthogonal, and symplectic groups. We obtain this algorithm by developing a diagrammatic framework based on category theory that enables us to not only express each weight matrix as a linear combination of diagrams but also makes it possible for us to use these diagrams to factor the original computation into a series of steps that are optimal. We show that this algorithm improves the Big-$O$ time complexity exponentially in comparison to a naïve matrix multiplication.

Figures (9)

• Figure 1: We use the string-like property of $(k,l)$--partition diagrams to Factor them as a composition of a permutation in $S_k$, an algorithmically planar $(k,l)$--partition diagram, and a permutation in $S_l$. Here, $k=5$ and $l=4$.
• Figure 2: The decomposition of the algorithmically planar $(5,4)$--partition diagram that appears in Figure \ref{['symmfactoringSn']} into a tensor product of smaller partition diagrams. These diagrams correspond, from right-to-left, to tensor contraction, transfer, and copying operations under the functor $\Theta$. This tensor product decomposition is used in PlanarMult for the symmetric group $S_n$.
• Figure 3: We show how matrix multiplication is implemented in PlanarMult for $S_n$ using the tensor product decomposition of the algorithmically planar $(5,4)$--partition diagram given in Figure \ref{['tensorproddecompSn']} as an example. We perform the matrix multiplication as follows: first, we deform the entire tensor product decomposition diagram by pulling each individual diagram up one level higher than the previous one, going from right-to-left, and then we apply the functor $\Theta$ at each level. Finally, we perform matrix multiplication at each level to obtain the final output vector.
• Figure 4: We use the string-like aspect of $(k,l)$--Brauer diagrams to Factor them as a composition of a permutation in $S_k$, an algorithmically planar $(k,l)$--Brauer diagram, and a permutation in $S_l$. Here $k = l = 5$.
• Figure 5: The tensor product decomposition of the planar $(5,5)$--Brauer diagram that appears in Figure \ref{['symmfactoringOn']}.

Theorems & Definitions (50)

• Definition 1
• Example 1
• Definition 2
• Example 2
• Definition 3
• Example 3
• Definition 4
• Theorem 5: Diagram Basis when $G(n) = S_n$
• Corollary 6: Permutation Equivariant Weight Matrices
• Theorem 7: Spanning set when $G(n) = O(n)$