A Diagrammatic Approach to Improve Computational Efficiency in Group Equivariant Neural Networks
This work addresses a critical bottleneck for researchers and practitioners using group equivariant neural networks, enabling more efficient computations in applications with known symmetries, though it is incremental as it builds on existing frameworks.
The authors tackled the computational inefficiency of implementing group equivariant neural networks by developing a fast matrix multiplication algorithm for equivariant weight matrices in tensor power layers, achieving an exponential improvement in Big-O time complexity compared to naive methods.
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.