# Fast computation of permutation equivariant layers with the partition algebra

arXiv: https://arxiv.org/abs/2303.06208 · HTML: https://mlsanity.com/paper/2303.06208/

## Key facts

- Cited by 7 papers in the map (data through 14 Oct 2026).
- Builds on 12 of its 25 references that are in the map.
- First posted on arXiv 10 Mar 2023.

## Builds on (most cited first, 12 of 12)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [Geometric Deep Learning: Grids, Groups, Graphs, Geodesics, and Gauges](https://mlsanity.com/paper/2104.13478/) (2021)
- [Invariant and Equivariant Graph Networks](https://mlsanity.com/paper/1812.09902/) (2018)
- [Sign and Basis Invariant Networks for Spectral Graph Representation Learning](https://mlsanity.com/paper/2202.13013/) (2022)
- [Scalars are universal: Equivariant machine learning, structured like classical physics](https://mlsanity.com/paper/2106.06610/) (2021)
- [Deep Models of Interactions Across Sets](https://mlsanity.com/paper/1803.02879/) (2018)
- [Universal Approximation of Functions on Sets](https://mlsanity.com/paper/2107.01959/) (2021)
- [Transformers Generalize DeepSets and Can be Extended to Graphs and Hypergraphs](https://mlsanity.com/paper/2110.14416/) (2021)
- [The general theory of permutation equivarant neural networks and higher order graph variational encoders](https://mlsanity.com/paper/2004.03990/) (2020)
- [Brauer's Group Equivariant Neural Networks](https://mlsanity.com/paper/2212.08630/) (2022)
- [Connecting Permutation Equivariant Neural Networks and Partition Diagrams](https://mlsanity.com/paper/2212.08648/) (2022)

## Built on by (most cited first, 7 of 7)

- [Graph Automorphism Group Equivariant Neural Networks](https://mlsanity.com/paper/2307.07810/) (2023)
- [Equivariant neural networks and piecewise linear representation theory](https://mlsanity.com/paper/2408.00949/) (2024)
- [Categorification of Group Equivariant Neural Networks](https://mlsanity.com/paper/2304.14144/) (2023)
- [Current Symmetry Group Equivariant Convolution Frameworks for Representation Learning](https://mlsanity.com/paper/2409.07327/) (2024)
- [An Algorithm for Computing with Brauer's Group Equivariant Neural Network Layers](https://mlsanity.com/paper/2304.14165/) (2023)
- [Compact Matrix Quantum Group Equivariant Neural Networks](https://mlsanity.com/paper/2311.06358/) (2023)
- [Permutation Equivariant Neural Networks for Symmetric Tensors](https://mlsanity.com/paper/2503.11276/) (2025)

## Abstract (arXiv)

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.