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

TL;DR

This paper introduces a diagrammatic framework based on category theory to significantly accelerate matrix multiplication in group equivariant neural networks, enhancing computational efficiency for multiple symmetry groups.

Contribution

It develops a novel diagrammatic approach that enables optimal factorization of weight matrices, reducing computational complexity in high-order tensor-based neural networks.

Findings

Achieves exponential improvement in time complexity over naive methods

Provides a unified framework for symmetric, orthogonal, special orthogonal, and symplectic groups

Facilitates practical implementation of high-order tensor neural networks

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\"{i}ve matrix multiplication.