arXiv:2607.21271cs.CV2026-07

提出Flash EQ-Linear,让等变层推理速度提升2倍,同时保持参数高效和精度优势。

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

论文配图:Flash EQ-Linear: Accelerating Equivariant Linear Layers via Group-wise Discrete Fourier Transform
图 1 · 摘自论文原文
  • 利用群上离散傅里叶变换与共轭对称性,将等变线性层复杂度从O(NDC)降至O(NDC/T)
  • 在FP32/FP16下实现高达2倍的前向加速,网络级速度提升1.7倍
  • 首个在精度、参数量、速度三方面全面超越非等变模型的等变架构

等变网络通过权重共享嵌入几何对称性,在视觉任务中表现出卓越的参数效率。然而,这种参数效率并未转化为计算效率:现有实现将结构化权重展开为稠密矩阵并调用通用稠密核,导致等变层的浮点运算量与非等变层相当。本文观察到,等变线性(EQ-Linear)层——现代等变架构中最基础且常用模块——本质上是沿群维度的循环卷积与沿通道维度的线性变换的组合。基于此,我们提出Flash EQ-Linear,一种精确加速算法,通过在群维度结合傅里叶卷积定理与实数DFT的共轭对称性,将复杂度从$ \mathcal{O}(NDC)$降至$ \mathcal{O}(NDC/T)$。我们进一步实现了针对Flash EQ-Linear的专用CUDA内核,覆盖前向与反向传播及FP32、FP16精度。在算子层面,Flash EQ-Linear相较PyTorch的F.linear实现最高2倍前向加速;在网络层面,Flash EQ-ViT与Flash EQ-Swin相较等变与非等变基线实现最高1.7倍端到端加速。据我们所知,这是首次等变网络在精度、参数效率与推理速度三个维度上均严格优于非等变模型。代码已开源于https://github.com/zhongchenzhao/FlashEQLinear。

原文摘要 · Abstract (English)

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.

等变网络加速算法傅里叶变换CUDA优化

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。