Linear neural network layers that are either equivariant or invariant to permutations of their inputs form core building blocks of modern deep learning architectures. Examples include the layers of DeepSets, as well as linear layers occurring in attention blocks of transformers and some graph neural networks. The space of permutation equivariant linear layers can be identified as the invariant subspace of a certain symmetric group representation, and recent work parameterized this space by exhibiting a basis whose vectors are sums over orbits of standard basis elements with respect to the symmetric group action. A parameterization opens up the possibility of learning the weights of permutation equivariant linear layers via gradient descent. The space of permutation equivariant linear layers is a generalization of the partition algebra, an object first discovered in statistical physics with deep connections to the representation theory of the symmetric group, and the basis described above generalizes the so-called orbit basis of the partition algebra. We exhibit an alternative basis, generalizing the diagram basis of the partition algebra, with computational benefits stemming from the fact that the tensors making up the basis are low rank in the sense that they naturally factorize into Kronecker products. Just as multiplication by a rank one matrix is far less expensive than multiplication by an arbitrary matrix, multiplication with these low rank tensors is far less expensive than multiplication with elements of the orbit basis. Finally, we describe an algorithm implementing multiplication with these basis elements.

翻译：对输入置换具有等变或不变性的线性神经网络层是现代深度学习架构的核心构建模块。例如DeepSets中的层、Transformer注意力模块以及某些图神经网络中的线性层。置换等变线性层的空间可被识别为特定对称群表示的不变子空间，近期研究通过构造一组基来参数化该空间，这些基向量是关于对称群作用的标准基元素轨道的和。这种参数化使得通过梯度下降学习置换等变线性层的权重成为可能。置换等变线性层的空间是划分代数（partition algebra）的推广——该对象最初发现于统计物理，与对称群表示理论有深刻联系，上述基是对划分代数所谓轨道基的推广。我们提出另一种替代基，它推广了划分代数图基（diagram basis），其计算优势源于构成该基的张量在自然分解为克罗内克积时具有低秩特性。正如与秩一矩阵相乘远优于与任意矩阵相乘，与这些低秩张量相乘的计算代价远低于与轨道基元素相乘。最后，我们描述了一种实现与这些基元素相乘的算法。