Fast computation of permutation equivariant layers with the partition algebraLinear 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.
arxiv.org