Search papers, labs, and topics across Lattice.
This paper introduces Flash EQ-Linear, an acceleration algorithm for equivariant linear layers that leverages the Fourier convolution theorem to significantly reduce computational complexity from \(\mathcal{O}(NDC)\) to \(\mathcal{O}(NDC/T)\). By implementing dedicated CUDA kernels, the authors achieve up to 2x speedup in forward operations and up to 1.7x end-to-end speedup in equivariant networks like Flash EQ-ViT and Flash EQ-Swin compared to both equivariant and non-equivariant baselines. This advancement allows equivariant networks to outperform non-equivariant architectures across accuracy, parameter efficiency, and inference speed for the first time.
Equivariant networks can now achieve superior performance across accuracy, efficiency, and speed, challenging the dominance of non-equivariant architectures.
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.
