Flash EQ-Linear: Accelerating Equivariant Linear Layers via Group-wise Discrete Fourier Transform
Equivariant networks embed geometric symmetries as structural priors through weight sharing, achieving remarkable parameter efficiency across vision tasks. However, this parameter efficiency does not translate into compute efficiency: existing implementations unroll the structured weights into dense matrices and dispatch them to generic dense kernels, so the FLOPs of an equivariant layer are no smaller than those of a non-equivariant counterpart. In this paper, we observe that the equivariant linear (EQ-Linear) layer---the most fundamental and frequently used module in modern equivariant architectures---is essentially a circular convolution along the group dimension composed with a linear transform along the channel dimension. Building on this observation, we propose Flash EQ-Linear, an exact acceleration algorithm that reduces the complexity from $\mathcal{O}(NDC)$ to $\mathcal{O}(NDC/T)$ by combining the Fourier convolution theorem along the group dimension with the conjugate symmetry of the real DFT. We further provide dedicated CUDA kernels for Flash EQ-Linear, covering both forward and backward passes and both FP32 and FP16 precision. At the operator level, Flash EQ-Linear achieves up to ${2\times}$ forward speedup over PyTorch's F.linear; at the network level, Flash EQ-ViT and Flash EQ-Swin achieve up to ${1.7\times}$ end-to-end speedup over both equivariant and non-equivariant baselines. To our knowledge, this is the first time equivariant networks strictly dominate their non-equivariant counterparts along all three axes simultaneously: accuracy, parameter efficiency, and inference speed.Code is available at https://github.com/zhongchenzhao/FlashEQLinear.
Code (0)
등록된 구현이 없습니다.
Similar Papers 제목 키워드 기반
In Search of Projectively Equivariant Networks
Equivariance of linear neural network layers is well studied. In this work, we relax the equivariance condition to only be true in a projective sense. We propose a way to construct a projectively equivariant neural netwo…
Equivariant non-linear maps for neural networks on homogeneous spaces
This paper presents a novel framework for non-linear equivariant neural network layers on homogeneous spaces. The seminal work of Cohen et al. on equivariant $G$-CNNs on homogeneous spaces characterized the representatio…
Equivariant neural networks and piecewise linear representation theory
Equivariant neural networks are neural networks with symmetry. Motivated by the theory of group representations, we decompose the layers of an equivariant neural network into simple representations. The nonlinear activat…
Invariant and Equivariant Graph Networks
Invariant and equivariant networks have been successfully used for learning images, sets, point clouds, and graphs. A basic challenge in developing such networks is finding the maximal collection of invariant and equivar…
2kFast computation of permutation equivariant layers with the partition algebra
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 l…