Flash EQ-Linear: Accelerating Equivariant Linear Layers via Group-wise Discrete Fourier Transform

Equivariant neural networks are highly valued for their ability to embed geometric symmetries (such as rotation or reflection) directly into their architecture. By sharing weights across group dimensions, these models achieve remarkable parameter efficiency and robustness. However, this efficiency has historically been theoretical. In practice, existing implementations treat the structured, shared weights as large, dense matrices, which are then processed by standard, non-optimized kernels. This leads to significant operational overhead, often making equivariant layers slower than their standard, non-equivariant counterparts despite having fewer parameters.

The authors identify that the equivariant linear (EQ-Linear) layer is fundamentally a circular convolution along the group dimension, followed by a linear transformation along the channel dimension. By leveraging the Fourier convolution theorem, the authors transform this group-wise convolution into pointwise multiplication in the frequency domain.

To maximize efficiency, the algorithm further exploits the conjugate symmetry of the real Discrete Fourier Transform (DFT). Because the input features and weights are real-valued, the frequency components are redundant; the algorithm computes only the non-redundant components, effectively halving the required operations. This reduces the computational complexity from O(NDC) to O(NDC/T).

The authors provide dedicated CUDA kernels for Flash EQ-Linear that support both forward and backward passes in FP32 and FP16 precision. These kernels are optimized for memory access patterns, parallelism, and frequency-domain layout. Empirical results show that Flash EQ-Linear achieves up to a 2x speedup at the operator level compared to standard linear layers. At the network level, models like Flash EQ-ViT and Flash EQ-Swin demonstrate up to a 1.7x end-to-end speedup, marking the first time equivariant networks have simultaneously outperformed standard baselines in accuracy, parameter count, and inference speed.