Flash EQ-Linear:让等变网络比标准线性层更快

快速导读:提出Flash EQ-Linear算法,利用群维度上的DFT卷积定理和实数DFT共轭对称性,将等变线性层复杂度从O(NDC)降至O(NDC/T),实现比标准线性层更快的推理速度。

等变网络通过嵌入几何对称性先验,在参数效率上具有显著优势。然而,现有实现往往因未充分利用群结构的循环特性,导致计算效率低下,推理速度慢于非等变对应模型。

本文提出Flash EQ-Linear算法,核心思想是将等变线性层视为群维度上的循环卷积,并利用离散傅里叶变换(DFT)的卷积定理将其转化为频域逐点乘法。结合实数DFT的共轭对称性,算法消除了冗余频率分量的计算,从而大幅降低复杂度。

英文题目:Flash EQ-Linear: Accelerating Equivariant Linear Layers via Group-wise Discrete Fourier Transform

论文出处:arXiv 每日论文精选 · arXiv:2607.21271

原始论文:PDF / 论文页面

对应视频标题:Flash EQ-Linear:让等变网络比标准线性层更快|推荐指数:★★★★★

这篇论文解决什么问题?

在EQ-ViT等架构中,EQ-Linear层占据了约50%的FLOPs和延迟。现有的Naive实现通常将权重展开为稠密矩阵,并使用通用内核进行计算。这种方法忽略了群维度上的循环结构,导致大量冗余计算。

尽管G-CNN等早期工作引入了群等变卷积,但主要关注理论框架。FFT-based CNN acceleration等工作利用FFT加速空间卷积,但未将其思想迁移至群维度的线性层。因此,如何在保持等变性的同时提升计算效率,是一个亟待解决的研究空白。

此外,现有加速方法往往需要重新训练或微调模型,这限制了其在实际部署中的应用。因此,一种无需重新训练、能直接替换现有模块的精确加速方法显得尤为重要。

核心创新

  • 算法创新:提出Flash EQ-Linear,将群卷积重构为频域逐点乘法,利用共轭对称性将理论复杂度从NDC MACs降至2NDC/T MACs。
  • 精确性保证:基于DFT可逆性和卷积定理,算法输出与Naive EQ-Linear完全一致,无需重新训练或微调。
  • 系统实现:提供覆盖前向/反向传播及FP32/FP16精度的专用CUDA算子,通过算子融合和内存访问优化实现实际加速。
  • 全面优势:首次使等变网络在准确率、参数效率和推理速度三个维度上同时优于非等变对应模型。

方法概览

将EQ-Linear视为群维度上的循环卷积,结合DFT卷积定理将其转化为频域逐点乘法,并利用实数DFT的共轭对称性消除冗余频率分量计算。提供专用CUDA内核,融合DFT、复数乘法和IDFT步骤。

  • 算法核心:将EQ-Linear视为群维度上的循环卷积。根据DFT卷积定理,时域卷积等价于频域逐点乘法。
  • 共轭对称性利用:对于实数输入,DFT结果具有共轭对称性。算法利用这一性质,仅计算一半的频率分量,从而将理论复杂度从NDC MACs降至2NDC/T MACs。
  • 专用CUDA内核:提供覆盖前向和反向传播的专用CUDA算子。通过融合DFT、复数乘法和IDFT步骤,减少内存访问开销。
  • 精度支持:支持FP32和FP16精度,确保在不同硬件平台上的适用性。
  • 精确性保证:基于DFT的可逆性和卷积定理,算法输出与Naive EQ-Linear完全一致,无需重新训练或微调。

逐图理解论文

现有瓶颈

现有瓶颈
Figure 1: Illustration of the EQ-Linear. The weight matrix W is constructed by cyclically shifting and replicating a set of learnable parameters ˜ W = [ ˜ W0, ˜ W1, · · · , ˜ WT−1] along the group dimension.

图1展示了EQ-Linear的权重矩阵W的构造方式。W由一组可学习参数沿群维度循环移位和复制而成。这种结构揭示了群维度上的循环卷积特性,为后续利用DFT加速提供了理论基础。

方法创新

方法创新
Figure 2: Illustration of the Flash EQ-Linear Algorithm for the p4 Rotation Group.

图2详细说明了Flash EQ-Linear算法在p4旋转群下的工作流程。算法首先对输入进行DFT,然后在频域进行逐点乘法,最后通过IDFT恢复时域输出。利用共轭对称性,仅计算一半的频率分量,从而减少计算量。

实验验证

实验验证
Table 3: Inference latency comparison of Flash EQ-ViT/Swin against standard non-EQ and naive EQ baselines. Speedup ratios (green) are measured against standard ViT/Swin baselines.

表3展示了Flash EQ-ViT/Swin与标准非等变和Naive EQ基线的端到端推理延迟比较。在FP32下,Flash EQ-ViT/Swin的速度提升1.4x-1.7x,证明了算法在网络级的高效性。

结论与局限

结论与局限
Table 1: Forward latency comparison of Flash EQ-Linear against Standard Linear and Naive EQ-Linear baselines under different channel numbers (64 →2048). Speedup ratios (green) are measured against the standard Linear (PyTorch’s F.linear) baseline.

表1比较了Flash EQ-Linear与Standard Linear和Naive EQ-Linear在不同通道数下的前向延迟。结果显示,随着通道数增加,Flash EQ-Linear的加速比显著提升,最高可达2.1x。

实验如何设计?

  • 硬件平台:NVIDIA RTX 4090 GPU。
  • 基准模型:EQ-ViT和EQ-Swin Transformer。
  • 对比基线:Standard Linear (PyTorch F.linear) 和 Naive EQ-Linear。
  • 评估指标:前向/反向延迟、端到端推理速度、相对L2误差、等变性NMSE。
  • 数据集:ImageNet-100。

关键结果与论文证据

  • 算子级加速:在RTX 4090上,Flash EQ-Linear前向延迟较PyTorch F.linear快1.3x-2.1x,FLOPs减少61.2% (Page 9)。
  • 网络级加速:在EQ-ViT和EQ-Swin上,FP32下端到端速度提升1.4x-1.7x,FP16下提升1.2x-1.3x,保持ImageNet-100 Top-1准确率不变 (Page 11)。
  • 数值验证:Flash EQ-Linear与Naive EQ-Linear的相对L2误差在FP32下为10^-7量级,FP16下为10^-4量级,符合浮点精度预期 (Page 12)。
  • 等变性验证:在p4旋转群下,Flash EQ-Linear的等变性误差(NMSE)为4.0e-8,与Naive实现相当,远优于标准线性层 (Page 12)。

阅读时需要注意

  • 当前CUDA实现仅针对p4旋转群(90度旋转),未泛化至其他变换群如反射群。
  • 反向传播在FP16小通道数下加速有限,因缺乏针对PyTorch生产级F.linear反向传播的深度内存访问优化。
  • 主要加速线性层,其他等变算子如等变卷积尚未应用此原理。

关联工作

  • G-CNN [Cohen and Welling, 2016a] 引入群等变卷积,通过群维度权重共享实现等变性。
  • EQ-ViT [He et al., 2021; Hutchinson et al., 2021; Fu et al., 2026] 将等变性引入Vision Transformer架构,本文基于其进行加速。
  • FFT-based CNN acceleration [Mathieu et al., 2013; Vasilache et al., 2014] 利用FFT加速空间卷积,本文将其思想迁移至群维度线性层。

展开:论文全文中文翻译

以下译文用于快速探索和学习,技术术语按需要保留英文;正式引用和精确表述请以原论文为准。

第 1 页

Flash EQ-Linear:通过群离散傅里叶变换加速等变线性层

赵忠辰1,王吉鑫1,齐骥1,*,林辉1,张磊2,3,孟德宇1,徐宗本1 1西安交通大学 2香港理工大学 3OPPO研究院 zhongchenzhao@stu.xjtu.edu.cn, xie.qi@mail.xjtu.edu.cn

摘要

等变网络通过权重共享将几何对称性作为结构先验嵌入网络,在视觉任务中实现了显著的参数效率。然而,这种参数效率并未转化为计算效率:现有实现将结构化权重展开为密集矩阵,并分派给通用的密集内核,因此等变层的浮点运算次数(FLOPs)并不小于其非等变对应层。在本文中,我们观察到等变线性层(EQ-Linear)——现代等变架构中最基本且频繁使用的模块——本质上是由沿群维度的循环卷积和沿通道维度的线性变换组成的。基于这一观察,我们提出了 Flash EQ-Linear,这是一种精确的加速算法,通过结合群维度上的傅里叶卷积定理和实数离散傅里叶变换(DFT)的共轭对称性,将复杂度从 $O(NDC)$ 降低至 $O(NDC/T)$。我们进一步为 Flash EQ-Linear 提供了专用的 CUDA 内核,涵盖前向和反向传播,以及 FP32 和 FP16 精度。在算子层面,Flash EQ-Linear 相比 PyTorch 的 F.linear 实现了高达 2 倍的前向加速;在网络层面,Flash EQ-ViT 和 Flash EQ-Swin 相比等变和非等变基线实现了高达 1.7 倍的端到端加速。据我们所知,这是首次等变网络在准确性、参数效率和推理速度这三个维度上同时严格优于其非等变对应网络 [cs.CV]。代码可在 https://github.com/zhongchenzhao/FlashEQLinear 获取。

1 引言

等变网络提供了一种原则性的方法,将几何对称性先验(例如平移、旋转、反射对称性先验)显式嵌入到网络架构中,不仅带来了性能提升,还增强了视觉模型的根本鲁棒性和泛化能力。通过在网络模块中通过权重共享将几何对称性作为显式结构先验嵌入 [Ravanbakhsh et al., 2017],等变架构消除了学习这些对称性所需的高昂数据增强成本,显著提高了参数效率,并提供了理论保证:输入上的几何变换会在输出上诱导可预测的变换。基于这一原理,等变卷积神经网络(CNNs)[Cohen and Welling, 2016a,b, Weiler et al., 2018, Kondor and Trivedi, 2018, Weiler and Cesa, 2019, Shen et al., 2020, 2021, Xie et al., 2022, 2025]、等变视觉 Transformer(EQ-ViTs)[He et al., 2021, Hutchinson et al., 2021, Fu et al., 2026] 以及最近的等变 Visual Mamba [Zhao et al., 2026] 可以作为其标准对应物的即插即用替代品,在图像分类、超分辨率、去噪和重建任务中,始终以更少的可学习参数实现更强的性能。

*通讯作者。

1

第 2 页

离散傅里叶变换 循环移位与共享 共轭 对称性

逆离散傅里叶变换 共轭

去除冗余 形状:卷积定理

实数乘法 实数乘法 复数乘法 实数乘法

等变参数量 # 高效等变参数量 # 高效等变参数量 # 高效 实用性 (a) 标准线性层 (b) 朴素 EQ-Linear 层 (c) Flash EQ-Linear 层 图 1:EQ-Linear 的示意图。权重矩阵 W 是通过沿群维度循环移位并复制一组可学习参数 $\tilde{W} = [\tilde{W}_0, \tilde{W}_1, \dots, \tilde{W}_{T-1}]$ 构建而成的。

尽管 equivariant 网络在参数效率和性能方面表现卓越,但在实际应用中却面临一个关键瓶颈:这种参数效率并未转化为计算效率。根本原因在于群维度参数共享机制本身。按其原生实现方式,共享权重并非直接就地计算;相反,它们被复制、置换并沿群维度平铺,以组装成一个大型稠密权重张量,随后传递给标准的 PyTorch 稠密内核。这种朴素的复制与稠密乘法策略引入了显著的操作开销,导致延迟远超理论 FLOPs 所暗示的水平。因此,即使 FLOPs 计数相同,equivariant 层通常也比其非 equivariant 对应层更慢 [Gerken et al., 2022]。加剧这一问题的还有,equivariant 网络目前仍主要用于学术研究。它们缺乏标准网络层所受益的高度优化的 CUDA 内核 [Chetlur et al., 2014],这使得它们在计算可信度上更低,在实际部署中更难。

在现代 equivariant 架构的组成部分中,提高计算效率的关键在于 equivariant 线性 (EQ-Linear) 层 [Finzi et al., 2021, Xie et al., 2025],因为它是最基础且使用最频繁的模块。目前,EQ-Linear 在接入最先进架构时已展现出令人羡慕的强大下游性能,并成为计算占比最高的 equivariant 模块之一。例如,在 EQ-ViT 中,EQ-Linear 层出现在每个注意力投影和每个 FFN/MLP 块中,约占 FLOPs 和推理延迟的 50%。在现有实现中,EQ-Linear 操作首先将共享权重矩阵复制并循环移位 T 次(其中 T 为群大小),以获得一个 $C \times C$(其中 $CT$ 为通道大小$^1$)的稠密矩阵,然后将其与 $C$ 维输入特征相乘(如图 1 (b) 所示)。这种方法在正向和反向传播中都存在重复的数据移动和相同的稠密矩阵乘法 FLOPs,拖慢了训练和推理的速度。一种常见的变通方法是在推理时预组装权重矩阵。这样 EQ-Linear 的运行速度与普通线性层完全一样。然而,这仅仅是隐藏了冗余而非利用它,参数效率仍未转化为计算效率。

事实上,将 EQ-Linear 权重矩阵视为普通稠密矩阵进行乘法运算,忽视了可以通过利用其高度规则的内部结构来获得计算增益。如图 1(b) 所示,EQ-Linear 的参数形成了一种群循环结构,自然地编码了数据中的循环对称性。这种结构规律引发了一个关键问题:我们能否利用这种规律性设计一种加速算法,使 EQ-Linear 的计算速度快于标准线性层,从而逆转 equivariant 网络的速度劣势,并首次实现

$^1$为了保持总通道数匹配,equivariant 网络将每组的通道宽度设置为非 equivariant 对应层的 $1/T$。

2

第 3 页

准确率、参数量以及推理速度? 在本文中,我们以肯定的方式回答了这个问题。我们的关键洞察在于,EQ-Linear 的计算并非通用的密集矩阵乘法。它实际上是在群维度上进行的循环卷积,并与沿通道维度的标准线性变换相结合。如图 1 所示,EQ-Linear 权重矩阵具有群循环结构:共享权重沿群轴循环移位并复制 T 次。这恰好构成了一种循环卷积。

这一观察结果提示了一种更快的 EQ-Linear 计算方法。(i) 首先,根据离散傅里叶变换 (DFT) 的卷积定理,原始域中的群卷积可以精确地计算为频域中的逐点乘法 [Oppenheim and Schafer, 2009]。这将主要计算量从 $C^2$ 次实数乘法减少到 $C^2T$ 次复数乘法。(ii) 此外,我们进一步利用了实数 DFT 的共轭对称性:由于输入和权重均为实数,约有一半的频率分量是另一半的复共轭,无需独立计算。这使得总成本降低至大约 $2C^2/T$ 次实数乘法。

基于上述分析,我们提出了 Flash EQ-Linear,这是一种精确的加速算法,在计算受限场景下为等变线性层提供约 2 倍的加速。据我们所知,这是首次使等变线性在墙钟时间(wall-clock time)上严格快于其非等变对应物(PyTorch 的 F.linear)。其意义深远:等变网络首次同时在三个维度上超越其非等变对应物——准确率性能、参数效率以及推理速度。本文的主要贡献总结如下:

• 算法方面。我们提出了 Flash EQ-Linear:一种用于等变线性的精确加速算法。它将群卷积重构为频域中的逐点乘法,并进一步利用实数 DFT 的共轭对称性去除冗余的频域操作,将理论复杂度从 $NDC$ 次 MACs 降低至 $2NDC/T$ 次 MACs。

• 精确性。Flash EQ-Linear 具有理论保证的精确性:它仅依赖于 DFT 的可逆性和卷积定理的严格等价性,其输出与朴素 EQ-Linear 的输出完全相同。此外,该算法无需训练且即插即用,可以直接替换 EQ-ViT 等预训练等变网络中的 EQ-Linear 层,无需重新训练、微调或更改架构。

• 实现。我们为 Flash EQ-Linear 提供了专用的 CUDA 内核,涵盖前向和后向传播,以及 FP32 和 FP16 精度。这些内核对频域布局、复数乘法、合并内存访问和并行粒度进行了系统级优化,并作为开箱即用的库向社区开源。

• 实证验证。在算子层面,Flash EQ-Linear 相比 PyTorch 的 F.linear 实现了高达 2 倍的前向加速。在网络层面,Flash EQ-ViT 和 Flash EQ-Swin 相比标准 ViT/Swin 及相应的等变基线实现了高达 1.7 倍的端到端加速。这些结果与我们的理论分析高度吻合。

2 相关工作

2.1 群等变神经网络

群等变神经网络将几何对称性嵌入到架构中,确保输入的几何变换会引发可预测且相应的输出变换。这一思想由 G-CNN [Cohen and Welling, 2016a] 引入,它通过在群维度上共享权重,将标准卷积推广为群卷积。后续工作扩展了这一

第 4 页

沿三个方向的原则:(i) 扩大支持的群,例如从 90◦ 旋转扩展到更精细的旋转、反射以及更一般的几何群 [Cohen and Welling, 2016a, Weiler and Cesa, 2019];(ii) 改进群表示的参数化,例如从多项式基扩展到傅里叶基和双三次基 [Weiler and Cesa, 2019, Xie et al., 2022, 2025];以及 (iii) 将等变性引入现代架构家族,例如从等变 CNN 到等变 ViT [He et al., 2021, Hutchinson et al., 2021, Fu et al., 2026] 和等变 VMamba [Zhao et al., 2026]。这些工作表明,几何对称性是一种强大的结构先验,使等变网络能够以显著更少的参数匹配或超越其非等变对应模型。

然而,现有研究侧重于表达能力、泛化能力和参数效率,而忽视了等变算子的运行时效率。据我们所知,这是第一项通过专用 CUDA 内核在系统层面加速等变算子的研究。

2.2 神经网络加速

通用神经网络加速方法包括剪枝、量化、知识蒸馏、稀疏化和低秩分解 [Cheng et al., 2018, Han et al., 2016, Hinton et al., 2015, Denton et al., 2014, Frankle and Carbin, 2019]。这些方法通常通过移除冗余参数、降低数值精度、强制稀疏计算、用低秩结构近似权重矩阵或训练更小的学生网络来提高效率 [Han et al., 2016, Hinton et al., 2015, Denton et al., 2014]。然而,它们通常需要重新训练、微调或专门的硬件支持,并且往往会产生显著的性能-速度权衡 [Dong et al., 2020, Liu et al., 2019]。此外,FLOPs 或参数的理论减少并不总是转化为成比例的墙钟时间加速 [Ma et al., 2018, Williams et al., 2009],这使得同时满足准确性、速度和可用性这三个目标变得困难。

另一条互补路线追求基于傅里叶卷积定理的精确无损加速。Mathieu 等人 [2013] 首次将 Cooley–Tukey FFT [Cooley and Tukey, 1965] 应用于 CNN 空间卷积,将其重写为频域中的逐点乘法;Vasilache 等人 [2014] 随后将其工程化为 cuDNN-FFT 后端。这些工作作用于图像卷积的空间维度。相比之下,我们在等变线性层的群维度上应用傅里叶变换,并结合实数 DFT 的共轭对称性与卷积定理,获得一种精确、免训练、即插即用的加速算法。

3 Flash EQ-Linear:原理、算法与实现

在本节中,我们开发了 Flash EQ-Linear,这是一种利用离散傅里叶变换 (DFT) 将 EQ-Linear 复杂度从 O(NDC) 降低到 O(NDC/T) 的精确算法。我们首先介绍底层加速原理(第 3.1 节),然后推导加速算法(第 3.2 节),并分析其复杂度(第 3.3 节)和特性(第 3.4 节)。最后,我们描述 CUDA 内核实现(第 3.5 节)。

3.1 加速原理

EQ-Linear 的公式化。等变线性层在群结构化特征空间之间进行映射,同时保持群作用下的等变性。如图 1 所示,其权重矩阵 W 具有块循环结构,由沿群维度循环移位一组可学习参数块 W˜ = [W0,˜ W1,˜ . . . , WT−1]˜ 形成。给定定义在变换群上的群结构化输入 X ∈R CT ×T

第 5 页

算法 1 Flash EQ-Linear ▷总乘加运算次数 (MACs) ≈ NCT + DC + 2NDC/T + NDT 输入:实值输入特征 $X \in \mathbb{R}^{N \times C_T \times T}$,参数 $\tilde{W} \in \mathbb{R}^{D_T \times C_T \times T}$,以及偏置 $\tilde{b} \in \mathbb{R}^{D_T}$。 1: 计算 $X$ 和 $\tilde{W}$ 的群离散傅里叶变换 (DFT) 的非冗余频率分量 $0, \dots, \lfloor T/2 \rfloor$: $\hat{X}^{G_k} = F_G(X)^{G_k}, \hat{W}^{G_k} = F_G(\tilde{W})^{G_k}, k=0, \dots, \lfloor T/2 \rfloor$. ▷ MACs=$2(N + D_T)C(\lfloor T/2 \rfloor+ 1)$ 2: 在每个非冗余频率处独立执行复数矩阵乘法: $\hat{Y}^{G_k} = \hat{X}^{G_k} \cdot (\hat{W}^{G_k})^\top, k=0, \dots, \lfloor T/2 \rfloor$. ▷ MACs= $\frac{4NDCT}{2} (\lfloor T/2 \rfloor+ 1)$ 3: 利用共轭对称性恢复剩余频率分量: $\hat{Y}^{G_{T-k}} = \overline{\hat{Y}^{G_k}}, k=1, \dots, \lfloor (T-1)/2 \rfloor$. ▷ MACs=0 4: 应用群逆离散傅里叶变换并加上偏置:$Y = F_G^{-1}(\hat{Y}) + \tilde{b}$. ▷ MACs=NDT 输出:$Y \in \mathbb{R}^{N \times D_T \times T}$

$T \times T$ 的群 $G = \{G_t \mid t = 0, 1, \dots, T-1\}$,EQ-Linear 将其映射为输出 $Y \in \mathbb{R}^{D \times C \times T}$

$\begin{bmatrix} Y^{G_0} \\ Y^{G_1} \\ \vdots \\ Y^{G_{T-1}} \end{bmatrix} = \begin{bmatrix} \tilde{W}_0 & \tilde{W}_1 & \cdots & \tilde{W}_{T-1} \\ \tilde{W}_{T-1} & \tilde{W}_0 & \cdots & \tilde{W}_{T-2} \\ \vdots & \vdots & \ddots & \vdots \\ \tilde{W}_1 & \tilde{W}_2 & \cdots & \tilde{W}_0 \end{bmatrix} \begin{bmatrix} X^{G_0} \\ X^{G_1} \\ \vdots \\ X^{G_{T-1}} \end{bmatrix} + \begin{bmatrix} \tilde{b} \\ \tilde{b} \\ \vdots \\ \tilde{b} \end{bmatrix}, \quad (1)$

其中,$X^{G_t} \in \mathbb{R}^{C_T}$ 和 $Y^{G_t} \in \mathbb{R}^{D_T}$ 分别表示 $X$ 和 $Y$ 的第 $t$ 个群分量$^2$。权重矩阵 $W \in \mathbb{R}^{D \times C}$ 和偏置 $b \in \mathbb{R}^D$ 由可学习参数 $\tilde{W} \in \mathbb{R}^{D_T \times C_T}$ 和 $\tilde{b} \in \mathbb{R}^{D_T}$ 平铺而成,其中 $C_T$ 和 $D_T$ 是每个群元素的输入和输出通道数。 朴素 EQ-Linear 的 MACs。如公式 (1) 所示,朴素实现忽略了 $W$ 的块循环结构,将 $\tilde{W}$ 显式展开为稠密矩阵 $W$,并使用通用矩阵乘法内核计算 $Y$。因此,它需要 $NDC$ 次乘加运算 (MACs),与具有相同输入和输出维度的非等变线性层的成本相匹配。 EQ-Linear 作为群循环卷积。我们的关键观察是,EQ-Linear 本质上是在群维度上对输入 $X$ 和参数 $\tilde{W}$ 执行循环卷积。具体而言,公式 (1) 可以等价地写为

$Y = X \circledast_G \tilde{W} + \tilde{b}, \quad \text{其中 } Y^{G_t} = \sum_{s=0}^{T-1} X^{G_s} \cdot \tilde{W}^{(s-t) \pmod T} + \tilde{b}, \forall G_t \in G, \quad (2)$

这正是在群维度上的循环卷积,并沿通道维度组合了矩阵乘法。 加速原理。通过在群维度上应用 DFT 卷积定理$^3$,公式 (2) 中的群循环卷积在原域中需要 $NDC$ 次 MACs——可以等价地转换为频域中的逐元素乘法,将主导复杂度降低至 $NDC/T$ 次 MACs:

$F_G(X \circledast_G \tilde{W}) = F_G(X) \odot_G F_G(\tilde{W}) \implies X \circledast_G \tilde{W} = F_G^{-1}(F_G(X) \odot_G F_G(\tilde{W})) . \quad (3)$

此外,由于输入特征 $X$ 和参数 $\tilde{W}$ 均为实值,其傅里叶系数

$^2 G_t \in G$ 既表示群变换(例如旋转矩阵),也表示沿群维度的索引。 $^3$ DFT 卷积定理指出,原域中的循环卷积变为频域中的逐点乘法:$F(X \circledast W) = F(X) \odot F(W)$,其中 $F(\cdot)$、$\circledast$ 和 $\odot$ 分别表示离散傅里叶变换、循环卷积和逐点乘法。

5

第 6 页

表现出共轭对称性。4 因此,我们只需计算非冗余的频率分量,从而将频域计算量几乎减半。

3.2 加速算法

3.2.1 适用于一般有限循环群的 Flash EQ-Linear

基于上述原理,我们开发了 Flash EQ-Linear,如算法 1 所述。该算法包含四个主要步骤:

步骤 1:群离散傅里叶变换(Group-wise DFT)。给定具有群结构输入特征 $X \in \mathbb{R}^{N \times C_T \times T}$,参数 $\tilde{W} \in \mathbb{R}^{D_T \times C_T \times T}$ 和偏置 $\tilde{b} \in \mathbb{R}^{D_T}$,我们首先沿群维度对 $X$ 和 $\tilde{W}$ 应用 DFT:

$$ \hat{X}^G_k = F_G(X)^G_k = \sum_{t=0}^{T-1} X^G_t e^{-i2\pi kt/T}, \quad \hat{W}^G_k = F_G(\tilde{W})^G_k, \quad k = 0, \dots, \lfloor T/2 \rfloor. \quad (4) $$

其中 $\hat{X} \in \mathbb{C}^{N \times C_T \times T}$ 和 $\hat{W} \in \mathbb{C}^{D_T \times C_T \times T}$ 为复数值。由于 $X$ 和 $\tilde{W}$ 是实数值,仅需显式计算 $\hat{X}$ 和 $\hat{W}$ 的前 $\lfloor T/2 \rfloor + 1$ 个非冗余频率分量,其余分量由共轭对称性确定。

步骤 2:逐频率矩阵乘法。如公式 3 所示,群循环卷积分解为跨频率的独立复数矩阵乘法:

$$ \hat{Y}^G_k = \hat{X}^G_k \cdot \hat{W}^{G\top}_k, \quad k = 0, \dots, \lfloor T/2 \rfloor. \quad (5) $$

每次乘法将固定频率处的逐点交互与跨通道维度的投影相结合。关键在于,不同频率之间相互独立,可以并行处理。

步骤 3:共轭对称恢复。利用实数 DFT 的共轭对称性和复共轭的乘积性质(即对于任意 $x, y \in \mathbb{C}$,有 $\overline{x \cdot y} = \bar{x} \cdot \bar{y}$),$\hat{Y}$ 的其余频率分量无需额外计算即可恢复:

$$ \hat{Y}^G_{T-k} = \hat{X}^G_{T-k} \cdot \hat{W}^{G\top}_{T-k} = \overline{\hat{X}^G_k} \cdot \overline{\hat{W}^{G\top}_k} = \overline{\hat{X}^G_k \cdot \hat{W}^{G\top}_k} = \overline{\hat{Y}^G_k}, \quad k=1, \dots, \lfloor (T-1)/2 \rfloor. \quad (6) $$

步骤 4:群逆离散傅里叶变换(Group-wise IDFT)。最后,我们沿群维度对 $\hat{Y}$ 应用逆 DFT 并加上群共享偏置:

$$ Y = F^{-1}_G(\hat{Y}) + \tilde{b}, \quad Y \in \mathbb{R}^{N \times D_T}, \quad (7) $$

其中 $\tilde{b}$ 在样本和群维度上进行广播。

3.2.2 适用于 p4 旋转群的 Flash EQ-Linear

如图 2 所示,对于广泛使用的 p4 旋转群 [Cohen and Welling, 2016a],将算法 1 专门化应用于 $T=4$ 可得到 Flash EQ-Linear 的大幅简化形式。

步骤 1:群离散傅里叶变换。对于 p4 旋转群($T=4$),$X$ 的群离散傅里叶变换特化为以下显式形式:

$$ \begin{bmatrix} \hat{X}^G_0 \\ \hat{X}^G_1 \\ \hat{X}^G_2 \\ \hat{X}^G_3 \end{bmatrix} = \begin{bmatrix} 1 & 1 & 1 & 1 \\ 1 & -i & -1 & i \\ 1 & -1 & 1 & -1 \\ 1 & i & -1 & -i \end{bmatrix} \begin{bmatrix} X^G_0 \\ X^G_1 \\ X^G_2 \\ X^G_3 \end{bmatrix} \implies \begin{bmatrix} \hat{X}^G_0 \\ \hat{X}^G_1 \\ \hat{X}^G_2 \\ \hat{X}^G_3 \end{bmatrix} = \begin{bmatrix} 1 & 1 & 1 & 1 \\ 1 & 0 & -1 & 0 \\ 1 & -1 & 1 & -1 \\ 0 & -1 & 0 & 1 \end{bmatrix} \begin{bmatrix} X^G_0 \\ \text{Re}(\hat{X}^G_1) \\ X^G_2 \\ \text{Im}(\hat{X}^G_1) \end{bmatrix}. \quad (8) $$

4 对于实值序列 $X \in \mathbb{R}^T$,其 DFT 满足 $\overline{F(X)_{T-k}} = F(X)_k$,其中 $\overline{\cdot}$ 表示复共轭。

6

第 7 页

图2:Flash EQ-Linear算法在p4旋转群下的示意图。

由于 $X$ 为实值,自共轭分量 $\hat{X}_{G0}$ 和 $\hat{X}_{G2}$ 为实值,而剩余一对满足 $\hat{X}_{G3} = \overline{\hat{X}_{G1}}$。因此,该变换完全由四个实值分量确定:$\hat{X}_{G0}$、$\text{Re}(\hat{X}_{G1})$、$\hat{X}_{G2}$ 和 $\text{Im}(\hat{X}_{G1})$。此外,式(8)中的简化变换矩阵仅包含0和 $\pm 1$。因此,群离散傅里叶变换完全简化为加法和减法,无需非平凡乘法。相同的简化也适用于实值参数 $W$ 的群离散傅里叶变换。

步骤2:逐频率傅里叶域矩阵乘法。利用步骤1中导出的四个实值傅里叶分量,式(5)中的逐频率乘法特化为:

$$ \begin{aligned} \hat{Y}_{G0} &= \hat{X}_{G0} \cdot \hat{W}_{G0}^\top, \quad \hat{Y}_{G2} = \hat{X}_{G2} \cdot \hat{W}_{G2}^\top, \\ \text{Re}(\hat{Y}_{G1}) &= \text{Re}(\hat{X}_{G1}) \cdot \text{Re}(\hat{W}_{G1})^\top – \text{Im}(\hat{X}_{G1}) \cdot \text{Im}(\hat{W}_{G1})^\top, \\ \text{Im}(\hat{Y}_{G1}) &= \text{Re}(\hat{X}_{G1}) \cdot \text{Im}(\hat{W}_{G1})^\top + \text{Im}(\hat{X}_{G1}) \cdot \text{Re}(\hat{W}_{G1})^\top. \end{aligned} $$

因此,步骤2简化为六个实值矩阵乘法,可以并行执行。由于每个乘法将 $C$ 个输入通道映射到 $D$ 个输出通道,总成本为 $16NDC/6$ 次乘加运算(MACs)。

步骤3:通过共轭对称性恢复。由于输入和权重均为实值,傅里叶域输出满足 $\hat{Y}_{G3} = \overline{\hat{Y}_{G1}}$。因此,我们避免显式生成冗余的 $G3$ 分量,并在逆变换期间隐式重建其贡献。此步骤既不需要额外的乘加运算,也不需要为 $\hat{Y}_{G3}$ 分配存储空间。

步骤4:群逆离散傅里叶变换。利用步骤3中的共轭对称性,p4旋转群的群逆离散傅里叶变换可以直接使用四个实值傅里叶分量表示为:

$$ \begin{bmatrix} Y_{G0} \\ Y_{G1} \\ Y_{G2} \\ Y_{G3} \end{bmatrix} = \begin{bmatrix} 1/4 & 1/4 & 1/4 & 1/4 \\ 1/4 & i/4 & -1/4 & -i/4 \\ 1/4 & -1/4 & 1/4 & -1/4 \\ 1/4 & -i/4 & -1/4 & i/4 \end{bmatrix} \begin{bmatrix} \hat{Y}_{G0} \\ \hat{Y}_{G1} \\ \hat{Y}_{G2} \\ \hat{Y}_{G3} \end{bmatrix} = \begin{bmatrix} 1/4 & 1/4 & 1/4 & 1/4 \\ 1/4 & i/4 & -1/4 & -i/4 \\ 1/4 & -1/4 & 1/4 & -1/4 \\ 1/4 & -i/4 & -1/4 & i/4 \end{bmatrix} \begin{bmatrix} \hat{Y}_{G0} \\ \hat{Y}_{G1} \\ \hat{Y}_{G2} \\ \overline{\hat{Y}_{G1}} \end{bmatrix} = \begin{bmatrix} 1/4 & 1/2 & 1/4 & 0 \\ 1/4 & 0 & -1/4 & -1/2 \\ 1/4 & -1/2 & 1/4 & 0 \\ 1/4 & 0 & -1/4 & 1/2 \end{bmatrix} \begin{bmatrix} \hat{Y}_{G0} \\ \text{Re}(\hat{Y}_{G1}) \\ \hat{Y}_{G2} \\ \text{Im}(\hat{Y}_{G1}) \end{bmatrix}, $$

第二个等式代入了 $\hat{Y}_{G3} = \overline{\hat{Y}_{G1}}$,而最终等式将 $\hat{Y}_{G1}$ 展开为其实部和虚部。所得的实值变换仅包含二进系数 $\{0, \pm 1/4, \pm 1/2\}$,因此简化为加法、减法和2的幂次缩放。因此,在最终实现中,步骤4不产生额外的乘加运算。

3.3 复杂度分析

我们通过分析算法1中每一步的乘加运算次数来量化 Flash EQ-Linear 的计算效率。所得复杂度总结如下,完整推导见附录6.1。

7

第 8 页

定理 1(Flash EQ-Linear 的复杂度)给定输入 $X \in \mathbb{R}^{N \times C \times T}$,输出 $Y \in \mathbb{R}^{N \times D \times T}$,所提出的 Flash EQ-Linear 算法产生的乘加运算次数(MACs)约为

$$ \frac{4NDC}{T} (\lfloor T/2 \rfloor + 1) \approx \frac{2NDC}{T} \quad (11) $$

对于 p4 旋转群($T=4$),复杂度可进一步降低至 $16NDC/6$ 次 MACs。

加速分析。相比之下,朴素 EQ-Linear 以及具有相同总通道维度的标准非等变线性层在计算稠密 $C \times D$ 投影时需要 $NDC$ 次 MACs。因此,Flash EQ-Linear 实现了理论加速比

$$ \text{Speedup} = \frac{\text{MACs}_{\text{Naive}}}{\text{MACs}_{\text{Flash}}} \approx \frac{NDC}{2NDC/T} = \frac{T}{2}. \quad (12) $$

对于 $T=4$,针对 p4 的特化版本产生的理论加速比为 $16/6 \approx 2.67\times$。

3.4 Flash EQ-Linear 的性质

精确性与等变性保持。Flash EQ-Linear 在构造上是精确的:DFT 的可逆性和卷积定理使其在代数上等价于朴素 EQ-Linear。因此,它保留了原始算子的旋转等变性。如表 4 所示,两者 FP32 输出之间的相对 $L_2$ 差异在 $10^{-7}$ 数量级,与数值精度一致。

加速比随群大小线性扩展。如公式 (12) 所述,Flash EQ-Linear 的理论加速比随群大小 $T$ 线性扩展。因此,随着等变网络采用更精细的群(更大的 $T$)以捕捉更丰富的几何结构,Flash EQ-Linear 的优势日益显著:更强的等变性带来更大的加速比。

免训练且即插即用。Flash EQ-Linear 可以即插即用地直接替换预训练等变网络中的朴素 EQ-Linear,无需重新训练、微调或更改架构。因此,现有模型如 EQ-ViT 和 EQ-Swin 可以在不修改其训练或推理流程的情况下采用 Flash EQ-Linear。

3.5 CUDA 内核实现

Alg. 1 的直观 PyTorch 实现将计算分解为单独的 DFT、复数乘法和 IDFT 算子。由此产生的内核启动、布局转换和全局内存流量可能主导运行时间并削弱算法收益 [Williams et al., 2009, Dao et al., 2022]。为了将理论加速比转化为实际加速比,我们为 p4 旋转群开发了 Flash EQ-Linear 的专用 CUDA 内核,支持 FP32 和 FP16 的前向和反向传播。

我们的实现将 Alg. 1 的所有阶段融合为类似 GEMM 的数据流,消除了昂贵的往返全局内存的中间过程。对于 $T=4$ 的 p4,DFT 和 IDFT 被完全展开为加法、减法和固定重缩放,中间傅里叶域片段保留在寄存器或共享内存中。权重以频域平面存储,而输入激活值保持其原始布局,避免了辅助打包张量和额外的 GPU 内存分配。我们进一步提供了形状特化的内核优化,针对不同 $(B, N, C, D)$ 情况调整分块大小、共享内存暂存、持久性、Tensor Core 利用率和向量化内存访问。更多实现细节见附录 6.2。

8

第 9 页

表 1:Flash EQ-Linear 与标准线性(Standard Linear)及朴素等变线性(Naive EQ-Linear)基线在不同通道数(64 → 2048)下的前向延迟对比。加速比(绿色)是相对于标准线性(PyTorch 的 F.linear)基线测量的。

| 方法 | #Param. (M)↓ | FLOPs (G)↓ | 前向延迟 (ms)↓ | | | | | | | :— | :— | :— | :— | :— | :— | :— | :— | :— | | | | | 64 | 128 | 256 | 512 | 1024 | 2048 |

单精度浮点 (FP32)

标准线性 | 0.066 | 0.134 | 0.11 | 0.43 | 1.52 | 6.05 | 24.48 | 104.96 朴素等变线性 | 0.016 | 0.134 | 0.11 | 0.44 | 1.54 | 6.25 | 25.16 | 105.74 Flash EQ-Linear | 0.016 | 0.052 | 0.07 | 0.25 | 0.79 | 3.03 | 11.92 | 49.14 加速比/减少量 ↑ | 75.8% | 61.2% | 1.6× | 1.7× | 1.9× | 2.0× | 2.1× | 2.1×

半精度浮点 (FP16)

标准线性 | 0.066 | 0.134 | 0.05 | 0.11 | 0.44 | 1.70 | 6.79 | 30.08 朴素等变线性 | 0.016 | 0.134 | 0.05 | 0.12 | 0.46 | 1.73 | 6.90 | 30.59 Flash EQ-Linear | 0.016 | 0.052 | 0.04 | 0.09 | 0.31 | 1.04 | 3.23 | 14.07 加速比/减少量 ↑ | 75.8% | 61.2% | 1.3× | 1.2× | 1.5× | 1.6× | 2.1× | 2.1×

4 实验

4.1 算子级加速

实验设置。我们将 Flash EQ-Linear 与两个基线进行基准测试:标准线性(PyTorch 的 F.linear [Paszke et al., 2019])和显式展开等变权重的朴素等变线性(Naive EQ-Linear)。所有方法接收相同的输入 $X \in \mathbb{R}^{B \times N \times C_T \times T}$,其中批量大小 $B = 32$,序列长度 $N = 1024$,组大小 $T = 4$ 对应于 p4 旋转群。我们将每组通道宽度 $c = C_T$ 从 64 变化到 2048,并设置输出宽度 $D$ 等于 $C$。为了公平比较,标准线性应用密集 $C \rightarrow C$ 投影。参数和 FLOPs 在 $c = 64$ 时报告。我们在单个 NVIDIA RTX 4090 GPU 上测量 FP32 和 FP16 下的前向和后向延迟,平均 200 次运行。所有加速比均相对于标准线性测量。

实验结果。表 1 和表 2 分别报告了不同通道宽度下的前向和后向延迟。(i)前向延迟。随着通道宽度 $C$ 从 64 增加到 2048,Flash EQ-Linear 在 FP32 下将前向传播加速了 1.6×–2.1×,在 FP16 下加速了 1.3×–2.1×。在较小的通道宽度下,性能主要受内存流量、布局转换和内核启动开销的主导。随着负载变为计算密集型,加速比趋向于理论上限 2.67×。Flash EQ-Linear 还将计算成本从 0.134G 降低到 0.052G FLOPs,减少了 2.6 倍,这与我们理论复杂度分析一致。(ii)后向延迟。在 FP32 下,Flash EQ-Linear 一致地加速了相对于标准线性的后向传播,尽管加速幅度小于前向传播,因为梯度计算引入了额外的算术运算和内存流量。在 FP16 下,在较小的通道宽度下增益有限,在某些设置中甚至略有减速。这是由于工程优化不足,而非算法限制:我们的后向内核缺乏 PyTorch 生产级 F.linear 后向所具备的重度调优内存访问优化。总体而言。这些结果证实,计算的理论减少转化为实质性的墙钟时间加速。在计算密集型模式下,Flash EQ-Linear 相对于高度优化的非等变基线实现了高达 2.1× 的前向加速,同时完全保持等变性。

第 10 页

表2:Flash EQ-Linear 与 Standard Linear 和 Naive EQ-Linear 基线在不同通道数(16→2048)下的反向延迟对比。

反向延迟 (ms)↓ 方法 16 32 64 128 256 512 1024 2048

单精度浮点 (FP32)

Standard Linear 0.17 0.14 0.16 0.46 1.67 5.42 20.80 83.11 Naive EQ-Linear 0.34 0.33 0.33 0.53 1.80 5.68 22.35 89.34 Flash EQ-Linear 0.14 0.13 0.13 0.39 1.29 4.70 17.79 70.91 加速比 (vs. Standard) ↑ 1.2× 1.1× 1.3× 1.2× 1.3× 1.2× 1.2× 1.2× 加速比 (vs. Naive) ↑ 2.5× 2.6× 2.6× 1.4× 1.4× 1.2× 1.3× 1.3×

半精度浮点 (FP16)

Standard Linear 0.18 0.14 0.14 0.30 1.06 3.90 14.67 59.48 Naive EQ-Linear 0.35 0.33 0.33 0.57 1.21 4.28 16.00 65.07 Flash EQ-Linear 0.12 0.13 0.23 0.33 1.28 3.90 12.84 45.02 加速比 (vs. Standard) ↑ 1.5× 1.1× 0.6× 0.9× 0.8× 1.0× 1.1× 1.3× 加速比 (vs. Naive) ↑ 2.9× 2.5× 1.4× 1.7× 0.9× 1.1× 1.2× 1.4×

4.2 网络级加速

实验设置。我们将 Flash EQ-Linear 集成到 EQ-ViT 和 EQ-Swin [Fu et al., 2026] 中,通过替换所有 Naive EQ-Linear 层,同时保持其余架构不变。我们将 Flash EQ-ViT/Swin 与标准的非等变 ViT/Swin [Dosovitskiy et al., 2021, Liu et al., 2021] 及其朴素等变对应模型在多个模型规模(Tiny/Small/Base/Large/Huge)下进行对比。在每个模型规模内,所有方法均使用相同的总通道宽度、批量大小 128、输入分辨率 224 × 224,以及对应于 p4 旋转群的组大小 T = 4。我们报告 ImageNet-100 上的 Top-1 准确率,这是 ImageNet1K [Deng et al., 2009] 的一个 100 类子集,同时报告参数量、FLOPs、累积线性层延迟、线性层到网络延迟比率、端到端延迟和吞吐量 (imgs/s)。延迟是在单块 NVIDIA RTX 4090 GPU 上,在 FP32 和 FP16 下测量并取 200 次运行的平均值。完整的实验结果见附录 6.3 中的表 7 和表 8。

实验结果。表 3 报告了 Flash EQ-ViT/Swin 在不同模型规模下的推理性能。(i) 线性层延迟。用 Flash EQ-Linear 替换 Naive EQ-Linear 使所有规模的累积线性层延迟降低了约 2 倍,这与表 1 中计算密集型算子级别的实验结果一致。这些结果表明,算子级别的增益直接转移到了全网络执行中。(ii) 端到端加速。在网络层面,Flash EQ-Linear 在 FP32 下实现了 1.4×–1.7× 的端到端加速,在 FP16 下实现了 1.2×–1.3× 的加速。随着模型规模的增大,增益也随之增加,因为较大的模型在线性层上花费了更大比例的运行时间——从 ViT-S 的 58.4% 到 ViT-H 的 77.7%——这正是 Flash EQ-Linear 加速的组件。此外,Flash EQ-ViT/Swin 保持了相应朴素等变模型在 ImageNet-100 上的 Top-1 准确率,证实了其无损、免训练且即插即用的加速特性。总体而言,Flash EQ-ViT/Swin 保留了等变网络的参数效率(对于 T = 4,参数量比非等变基线少约 4 倍),同时在推理速度上优于朴素等变模型和高度优化的非等变基线。这些结果表明,Flash EQ-Linear 的理论效率在实际中转化为显著的端到端加速。

4.3 Flash EQ-Linear 的精度验证

我们进一步验证了第 3.4 节中确立的两个属性:与 Naive EQ-Linear 的数值等价性以及旋转等变性的保持。

精确性验证。我们在随机输入上将 Flash EQ-Linear 与 Naive EQ-Linear 进行对比:

10

第 11 页

表3:Flash EQ-ViT/Swin 与标准非等变(non-EQ)及朴素等变(naive EQ)基线的推理延迟对比。加速比(绿色)是相对于标准 ViT/Swin 基线测得的。

所有线性层 | 整个网络 —|— 方法 | 参数量 (M)↓ | Top-1 (%)↑ | FLOPs (G)↓ | 延迟 (ms)↓ | 加速比 | FLOPs (G)↓ | 延迟 (ms)↓ | 吞吐量 (imgs/s)↑

单精度浮点 (FP32)

ViT-S | 33.8 | 76.5 | 13.1 | 0.352 | 58.4% | 14.1 | 0.603 | 1658 Naive EQ-ViT-S | 8.4 | 78.2 | 13.1 | 0.355 | 58.7% | 14.1 | 0.604 | 1656 Flash EQ-ViT-S | 8.4 | 78.2 | 5.0 | 0.183 | 42.6% | 6.0 | 0.429 | 2332 (1.4×) ViT-B | 85.9 | 77.1 | 33.5 | 0.761 | 67.1% | 35.1 | 1.135 | 881 Naive EQ-ViT-B | 21.4 | 80.1 | 33.5 | 0.757 | 67.1% | 35.1 | 1.128 | 886 Flash EQ-ViT-B | 21.4 | 80.1 | 12.7 | 0.400 | 52.3% | 14.3 | 0.764 | 1309 (1.5×) ViT-L | 303.4 | 78.4 | 119.0 | 2.672 | 71.9% | 123.1 | 3.715 | 269 Naive EQ-ViT-L | 75.8 | 80.6 | 119.0 | 2.672 | 71.9% | 123.1 | 3.715 | 269 Flash EQ-ViT-L | 75.8 | 80.6 | 44.8 | 1.336 | 56.7% | 49.0 | 2.356 | 424 (1.6×) ViT-H | 631.1 | 81.3 | 248.0 | 5.586 | 77.7% | 254.7 | 7.192 | 139 Naive EQ-ViT-H | 157.7 | 81.5 | 248.0 | 5.586 | 77.7% | 254.7 | 7.192 | 139 Flash EQ-ViT-H | 157.7 | 81.5 | 93.3 | 2.719 | 63.2% | 100.1 | 4.304 | 232 (1.7×) Swin-H | 655.1 | 87.2 | 228.6 | 5.232 | 70.4% | 230.3 | 7.430 | 135 Naive EQ-Swin-H | 164.0 | 88.3 | 228.6 | 5.232 | 70.4% | 230.3 | 7.430 | 135 Flash EQ-Swin-H | 164.0 | 88.3 | 84.3 | 2.474 | 52.8% | 90.8 | 4.687 | 213 (1.6×)

半精度浮点 (FP16)

ViT-H | 631.1 | 81.3 | 248.0 | 1.541 | 64.8% | 254.7 | 2.379 | 420 Naive EQ-ViT-H | 157.7 | 81.5 | 248.0 | 1.557 | 65.1% | 254.7 | 2.392 | 418 Flash EQ-ViT-H | 157.7 | 81.5 | 93.3 | 0.976 | 54.2% | 100.1 | 1.801 | 555 (1.3×) Swin-H | 655.1 | 87.2 | 228.6 | 1.428 | 59.2% | 230.3 | 2.412 | 415 Naive EQ-Swin-H | 164.0 | 88.3 | 228.6 | 1.406 | 58.4% | 230.3 | 2.408 | 415 Flash EQ-Swin-H | 164.0 | 88.3 | 84.3 | 0.958 | 49.0% | 90.8 | 1.955 | 512 (1.2×)

使用形状为 (32, 1024, 64, 4) 的张量,通过四个指标进行评估:相对 L2 误差(主要指标)、最大和平均绝对误差,以及第 99 百分位相对误差。我们在一个完整的训练步骤中评估所有张量:前向输出 $Y$ 以及梯度 $\nabla X$、$\nabla \tilde{W}$ 和 $\nabla \tilde{b}$。如表 4 所示,它们在 FP32 下的相对 L2 误差为 $10^{-7}$–$10^{-6}$,在 FP16 下约为 $10^{-4}$,这与相应精度下的浮点舍入一致 [Goldberg, 1991]。所有张量均满足匹配准则,证实 Flash EQ-Linear 除了有限精度效应外,没有引入数值误差。

等变性验证。在 p4 旋转群下,旋转输入会在输出的群维度上引起循环移位。因此,等变线性层 $f(\cdot)$ 应满足 $f(G(X))=G(f(X)), \forall G \in G_{p4}$。遵循 Xie 等人 [2025] 的方法,我们使用 $f(G(X))$ 和 $G(f(X))$ 之间的归一化均方误差 (NMSE) 来量化 FP32 下的等变误差。如表 5 所示,标准线性层产生了较大的误差 (1.4),表明无约束的线性层不能保持群作用。Flash EQ-Linear 将误差降低至 $4.0 \times 10^{-8}$,与朴素等变线性层 (Naive EQ-Linear) 的 $1.3 \times 10^{-7}$ 相当,证实了它在浮点舍入范围内保持了等变性。

11

第 12 页

表 4:Flash EQ-Linear 与朴素 EQ-Linear 数值等价性的验证。相对 L2 误差(Rel L2)是主要指标。

| Tensor Dtype | Shape | Rel L2 ↓ | Max Abs ↓ | Mean Abs ↓ | p99 Rel ↓ | Match | | :— | :— | :— | :— | :— | :— | :— | | Forward output Y | FP32 | B×N×DT×T | 2.6×10⁻⁷ | 1.2×10⁻⁶ | 1.1×10⁻⁷ | 1.7×10⁻⁶ | ✓ | | Backward input grad ∇X | FP32 | B×N×CT×T | 2.2×10⁻⁷ | 8.3×10⁻⁷ | 9.4×10⁻⁸ | 1.1×10⁻⁵ | ✓ | | Backward weight grad ∇W | FP32 | CT×CT×T | 5.9×10⁻⁷ | 2.7×10⁻⁴ | 2.9×10⁻⁵ | 3.2×10⁻⁵ | ✓ | | Backward bias grad ∇b | FP32 | DT | 1.5×10⁻⁷ | 2.3×10⁻⁵ | 7.2×10⁻⁶ | 4.5×10⁻⁶ | ✓ | | Forward output Y | FP16 | B×N×DT×T | 5.1×10⁻⁴ | 2.0×10⁻³ | 1.9×10⁻⁴ | 2.8×10⁻³ | ✓ | | Backward input grad ∇X | FP16 | B×N×CT×T | 4.1×10⁻⁴ | 2.0×10⁻³ | 1.4×10⁻⁴ | 1.8×10⁻² | ✓ | | Backward weight grad ∇W | FP16 | DT×CT×T | 6.3×10⁻⁴ | 2.5×10⁻¹ | 2.8×10⁻² | 3.0×10⁻² | ✓ | | Backward bias grad ∇b | FP16 | DT | 2.1×10⁻⁴ | 6.3×10⁻² | 6.3×10⁻³ | 9.2×10⁻³ | ✓ |

表 5:Flash EQ-Linear 在 p4 旋转群上的等变性验证。

| Method | Rel L2 ↓ | NMSE ↓ | Max Abs ↓ | Mean Abs ↓ | p99 Rel ↓ | | :— | :— | :— | :— | :— | :— | | Standard Linear | 1.4×10⁰ | 1.4×10⁰ | 4.1×10⁰ | 6.5×10⁻¹ | 6.4×10¹ | | Naive EQ-Linear | 1.4×10⁻⁷ | 1.3×10⁻⁷ | 7.3×10⁻⁷ | 6.0×10⁻⁸ | 7.5×10⁻⁶ | | Flash EQ-Linear | 5.4×10⁻⁸ | 4.0×10⁻⁸ | 3.2×10⁻⁷ | 1.9×10⁻⁸ | 2.4×10⁻⁶ |

5 结论

在本文中,我们指出了等变线性中长期存在的一个差距:参数效率并未转化为计算效率,因为现有的实现将结构化权重展开为密集矩阵,并将其分派给通用内核。为了弥补这一差距,我们提出了 Flash EQ-Linear,这是一种精确的加速算法,它利用傅里叶卷积定理和实数 DFT 的共轭对称性,将复杂度从 $NDC$ 次乘加运算(MACs)降低到 $2NDC/T$ 次乘加运算,从而实现了 $T/2$ 的理论加速比。在算子层面,Flash EQ-Linear 相比 PyTorch 的 F.linear 实现了高达 2 倍的推理加速;在网络层面,Flash EQ-ViT 相比等变和非等变基线实现了高达 1.7 倍的推理加速。此外,我们的工作揭示了一个更广泛的见解:等变性可以提供无损的计算加速,而不仅仅是参数效率。

局限性与未来工作。我们当前的 CUDA 实现专注于广泛使用的 p4 旋转群(即 90 度旋转群)。一个自然的扩展是将该算法推广到其他变换群,例如反射群。此外,虽然这项工作针对基本的等变线性层,但相同的算法原理可能扩展到其他等变算子,例如等变卷积,从而实现等变网络更全面的加速。

参考文献

Yu Cheng, Duo Wang, Pan Zhou, and Tao Zhang. A survey of model compression and acceleration for deep neural networks. IEEE Signal Processing Magazine, 35(1):126–136, 2018.

Sharan Chetlur, Cliff Woolley, Philippe Vandermersch, Jonathan Cohen, John Tran, Bryan Catanzaro, and Evan Shelhamer. cudnn: Efficient primitives for deep learning. arXiv preprint arXiv:1410.0759, 2014.

Taco S. Cohen and Max Welling. Group equivariant convolutional networks. In International Conference on Machine Learning, pages 2990–2999, 2016a.

Taco S Cohen and Max Welling. Steerable cnns. arXiv preprint arXiv:1612.08498, 2016b.

12

第 13 页

James W Cooley 和 John W Tukey。用于计算机计算复傅里叶级数的算法。计算数学,19(90):297–301,1965。

Tri Dao, Daniel Y Fu, Stefano Ermon, Atri Rudra, 和 Christopher R´e。Flashattention:具有 I/O 感知能力的快速且内存高效的精确注意力机制。在神经信息处理系统进展,第 35 卷,第 16344–16359 页,2022。

Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, 和 Li Fei-Fei。ImageNet:大规模分层图像数据库。在 IEEE 计算机视觉与模式识别会议,第 248–255 页,2009。

Emily L. Denton, Wojciech Zaremba, Joan Bruna, Yann LeCun, 和 Rob Fergus。利用卷积网络内的线性结构进行高效评估。在神经信息处理系统进展,2014。

Zhen Dong, Zhewei Yao, Yaohui Cai, Daiyaan Arfeen, Amir Gholami, Michael W. Mahoney, 和 Kurt Keutzer。Hawq-v2:基于海森矩阵感知迹加权的神经网络量化。在神经信息处理系统进展,2020。

Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Un-terthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkoreit, 和 Neil Houlsby。一张图片抵得上 16×16 个单词:用于大规模图像识别的 Transformer。在国际学习表征会议,2021。

Marc Finzi, Max Welling, 和 Andrew Gordon Wilson。为任意矩阵群构建等变多层感知器的实用方法。在第 38 届国际机器学习会议论文集,机器学习研究论文集第 139 卷,第 3318–3328 页。PMLR,2021。

Jonathan Frankle 和 Michael Carbin。彩票假设:寻找稀疏、可训练的神经网络。在国际学习表征会议,2019。

Jiahong Fu, Qi Xie, Deyu Meng, 和 Zongben Xu。Vanilla 群等变视觉 Transformer:简单且有效。arXiv 预印本 arXiv:2602.08047,2026。

Jan Gerken, Oscar Carlsson, Hampus Linander, Fredrik Ohlsson, Christoffer Petersson, 和 Daniel Persson。球面图像的等变性与数据增强。在第 39 届国际机器学习会议论文集,机器学习研究论文集第 162 卷,第 7404–7421 页。PMLR,2022。

David Goldberg。计算机科学家应知的浮点算术知识。ACM 计算调查,23(1):5–48,1991。doi: 10.1145/103162.103163。

Song Han, Huizi Mao, 和 William J. Dally。深度压缩:通过剪枝、训练后量化和霍夫曼编码压缩深度神经网络。在国际学习表征会议,2016。

Lingshen He, Yuxuan Chen, Yiming Dong, Yisen Wang, Zhouchen Lin, 等人。高效等变网络。神经信息处理系统进展,34:5290–5302,2021。

Geoffrey Hinton, Oriol Vinyals, 和 Jeff Dean。蒸馏神经网络中的知识。arXiv 预印本 arXiv:1503.02531,2015。

13

第 14 页

Michael J Hutchinson, Charline Le Lan, Sheheryar Zaidi, Emilien Dupont, Yee Whye Teh, 和 Hyunjik Kim。 Lietransformer:李群上的等变自注意力。在 International Conference on Machine Learning 中, 页码 4533–4543。PMLR,2021。

Risi Kondor 和 Shubhendu Trivedi。论等变性和神经网络中卷积到紧群作用的推广。在 International Conference on Machine Learning 中,页码 2747–2755。PMLR,2018。

Ze Liu, Yutong Lin, Yue Cao, Han Hu, Yixuan Wei, Zheng Zhang, Stephen Lin, 和 Baining Guo。Swin Transformer:使用移位窗口的分层视觉 Transformer。在 IEEE/CVF International Conference on Computer Vision 论文集,页码 10012–10022,2021。

Zhuang Liu, Mingjie Sun, Tinghui Zhou, Gao Huang, 和 Trevor Darrell。重新思考网络剪枝的价值。在 International Conference on Learning Representations,2019。

Ningning Ma, Xiangyu Zhang, Hai-Tao Zheng, 和 Jian Sun。Shufflenet v2:高效 CNN 架构设计的实用指南。在 European Conference on Computer Vision (ECCV) 论文集,页码 116–131,2018。

Michael Mathieu, Mikael Henaff, 和 Yann LeCun。通过 FFT 快速训练卷积网络。 arXiv 预印本 arXiv:1312.5851,2013。

Alan V Oppenheim 和 Ronald W Schafer。离散时间信号处理。Pearson,第 3 版,2009。

Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Kopf, Edward Yang, Zachary DeVito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, 和 Soumith Chintala。PyTorch:一种命令式风格的高性能深度学习库。在 Advances in Neural Information Processing Systems 中,第 32 卷,页码 8024–8035,2019。

Siamak Ravanbakhsh, Jeff Schneider, 和 Barnab´as P´oczos。通过参数共享实现等变性。在 International Conference on Machine Learning 中,页码 2892–2901。PMLR,2017。

Zhengyang Shen, Lingshen He, Zhouchen Lin, 和 Jinwen Ma。Pdo-econvs:基于偏微分算子的等变卷积。在 International Conference on Machine Learning 中,页码 8697–8706。 PMLR,2020。

Zhengyang Shen, Tiancheng Shen, Zhouchen Lin, 和 Jinwen Ma。Pdo-es2cnns:基于偏微分算子的等变球面 CNN。在 AAAI 人工智能会议论文集, 第 35 卷,页码 9585–9593,2021。

Nicolas Vasilache, Jeff Johnson, Michael Mathieu, Soumith Chintala, Serkan Piantino, 和 Yann LeCun。使用 fbfft 的快速卷积网络:GPU 性能评估。arXiv 预印本 arXiv:1412.7580,2014。

Maurice Weiler 和 Gabriele Cesa。通用 E(2)-等变可旋转 CNN。在 Advances in Neural Information Processing Systems 中,2019。

Maurice Weiler, Fred A Hamprecht, 和 Martin Storath。学习用于旋转等变 CNN 的可旋转滤波器。在 IEEE Conference on Computer Vision and Pattern Recognition 论文集,页码 849–858, 2018。

Samuel Williams, Andrew Waterman, 和 David Patterson。Roofline:一种针对多核架构的直观视觉性能模型。ACM 通讯,52(4):65–76,2009。

14

第 15 页

Qi Xie, Qian Zhao, Zongben Xu, 和 Deyu Meng. 基于傅里叶级数展开的滤波器参数化用于等变卷积。IEEE 模式分析与机器智能汇刊, 45(4): 4537–4551, 2022.

Qi Xie, Jiahong Fu, Zongben Xu, 和 Deyu Meng. 旋转等变任意尺度图像超分辨率。IEEE 模式分析与机器智能汇刊, 2025.

Zhongchen Zhao, Qi Xie, Keyu Huang, Lei Zhang, Deyu Meng, 和 Zongben Xu. 用于视觉任务的旋转等变 Mamba。arXiv 预印本 arXiv:2603.09138, 2026.

15

第 16 页

6 附录

6.1 Flash EQ-Linear 的详细复杂度分析

步骤 1 的 MACs:群离散傅里叶变换。根据定义,群离散傅里叶变换是沿群维度的复矩阵乘法,可写为 $\hat{X} = XF$,其中 $F \in \mathbb{C}^{T \times T}$ 为 DFT 矩阵。因此,计算前 $\lfloor T/2 \rfloor + 1$ 个非冗余频率分量相当于将 $X \in \mathbb{R}^{N \times CT \times T}$ 乘以 $F$ 的相应列:

$$ \hat{X}_k = X \cdot F_{:,k}, \quad k = 0, \dots, \lfloor T/2 \rfloor, \quad (13) $$

这需要 $2NC(\lfloor T/2 \rfloor + 1)$ 次实数 MACs。类似地,变换 $W \in \mathbb{R}^{DT \times CT \times T}$ 需要 $2DC_T(\lfloor T/2 \rfloor + 1)$ 次 MACs。因此,步骤 1(公式 4)的总成本为 $2(N + D_T)C(\lfloor T/2 \rfloor + 1)$。

步骤 2 的 MACs:逐频率矩阵乘法(主导项)。在 $\lfloor T/2 \rfloor + 1$ 个非冗余频率中的每一个上,公式 5 中的逐通道复矩阵乘法

$$ \hat{Y}^G_k = \hat{X}^G_k \cdot \hat{W}^{G\top}_k \quad (14) $$

涉及 $NDC$ 次复数乘法。由于复数乘法 $(a + bi)(c + di) = (ac – bd) + (ad + bc)i$ 需要四次实数乘法,因此每个频率的成本为 $4NDC$ 次 MACs。对所有非冗余频率求和,$\lfloor T/2 \rfloor + 1 \approx T/2$,得到 $2NDC T$,这主导了 Flash EQ-Linear 的整体成本。

步骤 3 的 MACs:共轭对称恢复。$\hat{Y}$ 的剩余频率分量根据公式 6 通过共轭复制恢复。因此,此步骤不产生任何 MACs。

步骤 4 的 MACs:群离散逆傅里叶变换。与步骤 1 类似,公式 7 中的群离散逆傅里叶变换可表示为 $Y = \hat{Y} F^{-1}$,其中 $F^{-1} \in \mathbb{C}^{T \times T}$ 为逆 DFT 矩阵。步骤 3 中建立的共轭对称性 $\hat{Y}^{G}_{T-k} = \hat{Y}^{G}_k$ 保证了 $Y$ 是实值的$^5$。利用这一对称性,每对共轭项可简化为实数表达式:

$$ \hat{Y}^G_k F_{t,k-1} + \hat{Y}^G_{T-k} F_{t,T-k-1} = \hat{Y}^G_k F_{t,k-1} + \hat{Y}^{G*}_k F^{*}_{t,k-1} = 2 \text{Re}(\hat{Y}^G_k)\text{Re}(F_{t,k-1}) – 2 \text{Im}(\hat{Y}^G_k)\text{Im}(F_{t,k-1}), \quad (15) $$

该计算仅需两次实数乘法而非八次,相比朴素复数运算实现了 4 倍的缩减。因此,在一般情况下,步骤 4 产生 $NDT$ 次 MACs。

总 MACs。结合四个步骤,Flash EQ-Linear 的总计算成本为

$$ \text{MACs}_{\text{Flash}} = \underbrace{2(N + D_T)C(\lfloor T/2 \rfloor + 1)}_{\text{群 DFT (步骤 1)}} + \underbrace{2NDC \frac{T}{2}(\lfloor T/2 \rfloor + 1)}_{\text{复矩阵乘法 (步骤 2)}} + \underbrace{NDT}_{\text{群 IDFT (步骤 4)}}. \quad (16) $$

在典型的等变网络中,当 $N \gg T$ 且 $C, D \gg T$ 时,主导项为

$$ \text{MACs}_{\text{Flash}} \approx 2NDC \frac{T}{2} \cdot \frac{T}{2} \approx 2NDC, \quad (17) $$

而群 DFT 和 IDFT 仅贡献低阶开销。

对于常见的设置 $T = 2$ 和 $T = 4$,DFT 和 IDFT 简化为固定的加法、减法和二进缩放。这些缩放操作可以折叠到内核的后处理阶段,并根据我们的计数约定不产生额外的 MACs。因此,在这些专用实现中,步骤 1 和 4 不产生乘法成本,使步骤 2 成为唯一的 MACs 来源。

$^5$ 根据实 DFT 对偶性,信号为实值当且仅当其 DFT 为共轭对称。

16

第 17 页

加速分析。相比之下,朴素 EQ-Linear 以及具有相同总通道维度的标准非等变线性层,在计算稠密 $C \times D$ 投影时均需要 $NDC$ 次乘加运算(MACs)。因此,Flash EQ-Linear 实现了理论加速比:

$$ \text{Speedup} = \frac{\text{MACs}_{\text{Naive}}}{\text{MACs}_{\text{Flash}}} \approx \frac{NDC T}{2NDC/T} = \frac{T^2}{2}. \quad (18) $$

对于 $T=4$,针对 p4 群的精确特化产生了 $16/6 \approx 2.67\times$ 的理论加速比。

注:为何共轭对称恢复至关重要。若无第 3 步,仅第 2 步将在所有 $T$ 个频率上产生 $4NDCT$ 次实数乘法,这等于 $T=4$ 时朴素方法 $NDCT^2$ 的成本,且不提供任何加速。因此,共轭对称性缩减是必要的:它使得 Flash EQ-Linear 对于常见的 p4 旋转群($T=4$)严格快于朴素基线。

6.2 详细的 CUDA 内核实现

算法 1 的直接 PyTorch 实现会将该方法表示为独立的群离散傅里叶变换(DFT)、频域乘法和群逆离散傅里叶变换(IDFT)算子。尽管这在数学上是忠实的,但它会实例化中间频率张量,并引入额外的布局转换、内存流量和内核启动开销。这种开销在 p4 群中尤为明显,因为其群维度较小,且变换本身仅包含少量固定加法。因此,我们的 CUDA 实现将第 3.2 节中的 FFT 公式视为原始等变线性映射的向量化规则,而非一系列独立库调用。沿群维度的固定变换、逐通道矩阵乘法和逆变换被融合到相同的执行数据流中,使得频率域变量仅在图块(tile)生命周期内作为寄存器/共享内存片段存在。

向量化群 DFT 实现。对于实验中使用的 p4 群,群大小固定为 $T=4$。在这种情况下,公式 4 中的群 DFT 矩阵简化为一个常数 $4 \times 4$ 变换,其条目仅为符号和虚数单位的幂。在利用实信号的共轭对称性后,非冗余坐标可由四个实标量表示:平凡表示分量、自共轭二阶分量,以及非平凡共轭表示对的实部和虚部。因此,不需要通用 FFT 内核、动态旋转因子表或显式复数张量构造。给定输入特征在一个批次-令牌-通道位置处的四个群分量 $(x_0, x_1, x_2, x_3)$,内核计算向量化群 DFT 坐标为:

$$ f_0 = x_0 + x_1 + x_2 + x_3, \quad f_1 = x_0 – x_2, \quad (19) $$ $$ f_2 = x_0 – x_1 + x_2 – x_3, \quad f_3 = x_3 – x_1. $$

这四个表达式是算法 1 中群 DFT 步骤的向量化 CUDA 实现。它们在加载时为每个输入图块进行评估。因此,该实现通过四个固定的算术表达式实现了基于 FFT 的公式,这些表达式可以被编译器完全展开。频域特征因此是片上量,保存在寄存器中或暂存在共享内存中,而不是写入并重新从全局内存读取的完整张量。这是我们的 CUDA 算子与数学算法的字面实现之间的主要系统差异。

内核数学形式。令 $n$ 索引展平的批次-令牌维度,$c$ 为输入通道,$d$ 为输出通道。对于每个输出图块,前向内核沿 $c$ 流式传输并直接累加

17

第 18 页

结构化频域乘积:

a0,n,d = Σc f0,n,c w0,d,c, a1,n,d = Σc f1,n,c w1,d,c − f3,n,c w3,d,c , (20) a2,n,d = Σc f2,n,c w2,d,c, a3,n,d = Σc f1,n,c w3,d,c + f3,n,c w1,d,c .

其中 w0 对应 C4 傅里叶基的平凡表示(k=0)分量,w2 对应自共轭的二阶分量(k = 2),而 (w1, w3) 参数化了非平凡共轭对(k =1, 3)的实部和虚部。因此,第二和第四个方程正是公式 5 中的复数乘法,写成了实值融合操作。重要的是,该内核从不构建等效的稠密 C × D 矩阵,也从不将六个实数 GEMM 打包表示作为持久中间变量构建。相反,每个 CTA 读取原始群结构化输入的图块,即时将其转换为 (f0, f1, f2, f3),并直接累加 (a0, a1, a2, a3)。在 c 上的规约之后,内核在后处理阶段应用相应的群离散傅里叶逆变换(IDFT):

y0,n,d = 4a0,n,d + 12a1,n,d + 14a2,n,d, y1,n,d = 14a0,n,d − 14a2,n,d − 12a3,n,d, (21) y2,n,d = 14a0,n,d − 12a1,n,d + 14a2,n,d, y3,n,d = 14a0,n,d − 14a2,n,d + 12a3,n,d.

后处理阶段也是固定的向量化变换。它仅使用常数 {1/4, 1/2} 以及加减法,因此可以与最终的全局存储融合。因此,CUDA 内核实现了与方法部分完全相同的方程,但避免了显式构建展开的稠密 EQ-Linear 矩阵、打包输入张量和打包频域输出张量。计算路径仍然是基于 FFT 的 Flash EQ-Linear 路径;区别在于傅里叶变换被实现为围绕类 GEMM 通道规约的局部向量化基变换。

内存布局与向量化。外部张量布局保持为 [B, N, CT , T],群维度连续。这种布局使得一个行-通道对的四个群分量在内存中相邻,因此内核可以将群元组作为紧凑向量加载,立即应用四个群离散傅里叶变换公式,并随后连续存储四个输出分量。批次和令牌/空间维度被展平为行维度,通道规约以类 GEMM 的方式在行维度和输出通道维度上进行分块。权重存储在相应的频域平面中,因此内核可以在不改变向用户暴露的数学表示的情况下发出公式 5 中的结构化乘积。在每个图块内,输入片段和权重片段被暂存以匹配合并的全局访问和高效的矩阵乘法消耗。对于 FP16 执行,输入和权重图块以半精度存储,并由面向 Tensor Core 的主循环消耗,当需要额外的数值稳定性时使用混合累加变体。这种设计保持了高吞吐量的 GEMM 核心,同时消除了原本会出现在群离散傅里叶变换、乘法和群离散傅里叶逆变换之间的全局内存往返。

形状专用加速。不同的 (B, N, CT , DT ) 情况暴露出不同的瓶颈,从小通道宽度时的启动开销和图块占用率不足,到大通道宽度时的算术吞吐量。因此,我们提供了形状专用的 CUDA 变体,并将每个基准测试形状路由到最快的已验证实现。这些变体保留了上述相同的数学内核,

18

第 19 页

表 6:实验中使用的 ViT 和 Swin Transformer 变体的配置细节、参数量及总 FLOPs

| Model | Patch Size | Embed Dim. | Depth/Depths | Num. Heads | #Param. (M)↓ | FLOPs (G)↓ | | :— | :— | :— | :— | :— | :— | :— | | ViT | | | | | | | | Tiny | 16 | 384 | 12 | 3 | 21.7 | 9.2 | | Small | 16 | 480 | 12 | 6 | 33.8 | 14.1 | | Base | 16 | 768 | 12 | 12 | 85.9 | 35.1 | | Large | 16 | 1024 | 24 | 16 | 303.4 | 123.1 | | Huge | 16 | 1280 | 32 | 16 | 631.1 | 254.7 | | Swin | | | | | | | | Tiny | 4 | 96 | [2, 2, 6, 2] | [3, 6, 12, 24] | 27.6 | 8.99 | | Small | 4 | 96 | [2, 2, 18, 2] | [3, 6, 12, 24] | 48.9 | 17.5 | | Base | 4 | 128 | [2, 2, 18, 2] | [4, 8, 16, 32] | 86.9 | 30.9 | | Large | 4 | 192 | [2, 2, 18, 2] | [6, 12, 24, 48] | 195.2 | 69.0 | | Huge | 4 | 352 | [2, 2, 18, 2] | [11, 22, 44, 88] | 655.1 | 230.3 |

通过调整图块大小、持久化、共享内存暂存、向量化内存访问以及 Tensor Core 的使用,以加速不同形状下的计算。换言之,形状专业化是一种实现层面的优化;它不改变基于 FFT 的方程或 EQ-Linear 的参数化。

前向和反向传播。前向传播遵循上述的融合公式。反向传播以相反顺序使用相同的频域分解,并保持相对于原始输入张量和原始 EQ-Linear 参数的梯度定义,而非任何后端特定的打包布局。具体而言,梯度首先表示为相同的四坐标频率基,乘以相应的转置结构化权重,然后通过逆固定变换映射回原始空间。这保留了该层的即插即用语义:优化的 CUDA 布局是内部执行细节,而数学算子及其梯度仍保持为 Flash EQ-Linear 的定义。

6.3 更多实验结果

ViT 和 Swin 的配置细节。表 6 列出了我们在实验中使用的 ViT 和 Swin Transformer 变体的架构配置(第 4.2 节)。对于每个模型系列,我们遵循原始论文 [Dosovitskiy et al., 2021, Liu et al., 2021] 的标准缩放惯例,实例化了五种规模(Tiny / Small / Base / Large / Huge)。ViT 变体通过缩放嵌入维度和注意力头数来增加容量,参数量从 ViT-T 的 21.7M 到 ViT-H 的 631.1M 不等。Swin 变体采用四阶段分层深度模式,缩放嵌入维度和每阶段头数,参数量从 Swin-T 的 27.6M 到 Swin-H 的 655.1M 不等。

Naive EQ-ViT/Swin 和 Flash EQ-ViT/Swin 变体与其非等变对应物共享相同的架构配置,但在 p4 旋转群(T=4)上运行,每组通道宽度设置为非等变基线的 1/T,以保持总通道数匹配。

Flash EQ-ViT 和 Flash EQ-Swin。在正文(第 4.2 节,表 3)中,由于篇幅限制,我们仅在代表性模型规模下报告了推理效率。在此,我们提供了所有评估模型规模(ViT 和 Swin 的 Tiny/Small/Base/Large/Huge)下的完整结果,分别在 FP32(表 7)和 FP16(表 8)精度下。

这些扩展结果加强了正文中的两个观察结果。(i) 在所有模型规模和两种精度下,Flash EQ-Linear 始终将线性层延迟降低约 2 倍,从而在网络级别带来 1.1×–1.7× 的端到端加速。(ii) 端到端加速

19

第 20 页

表7:Flash EQ-ViT/Swin 与标准非等变(non-EQ)及朴素等变(naive EQ)基线在 FP32 下的推理效率对比。加速比(绿色)是相对于标准基线测量的。

所有线性层 | 总网络 方法 | 参数量 (M)↓ | Top-1 (%)↑ | FLOPs (G)↓ | 延迟 (ms)↓ | 加速比 | FLOPs (G)↓ | 延迟 (ms)↓ | 吞吐量 (imgs/s)↑

单精度浮点 (FP32)

ViT-T | 21.7 | 76.5 | 8.4 | 0.218 | 60.8% | 9.2 | 0.358 | 2796 Naive EQ-ViT-T | 5.4 | 77.3 | 8.4 | 0.218 | 60.8% | 9.2 | 0.358 | 2796 Flash EQ-ViT-T | 5.4 | 77.3 | 3.2 | 0.115 | 46.2% | 4.0 | 0.249 | 4012 (1.4×) ViT-S | 33.8 | 76.5 | 13.1 | 0.352 | 58.4% | 14.1 | 0.603 | 1658 Naive EQ-ViT-S | 8.4 | 78.2 | 13.1 | 0.355 | 58.7% | 14.1 | 0.604 | 1656 Flash EQ-ViT-S | 8.4 | 78.2 | 5.0 | 0.183 | 42.6% | 6.0 | 0.429 | 2332 (1.4×) ViT-B | 85.9 | 77.1 | 33.5 | 0.761 | 67.1% | 35.1 | 1.135 | 881 Naive EQ-ViT-B | 21.4 | 80.1 | 33.5 | 0.757 | 67.1% | 35.1 | 1.128 | 886 Flash EQ-ViT-B | 21.4 | 80.1 | 12.7 | 0.400 | 52.3% | 14.3 | 0.764 | 1309 (1.5×) ViT-L | 303.4 | 78.4 | 119.0 | 2.672 | 71.9% | 123.1 | 3.715 | 269 Naive EQ-ViT-L | 75.8 | 80.6 | 119.0 | 2.672 | 71.9% | 123.1 | 3.715 | 269 Flash EQ-ViT-L | 75.8 | 80.6 | 44.8 | 1.336 | 56.7% | 49.0 | 2.356 | 424 (1.6×) ViT-H | 631.1 | 81.3 | 248.0 | 5.586 | 77.7% | 254.7 | 7.192 | 139 Naive EQ-ViT-H | 157.7 | 81.5 | 248.0 | 5.586 | 77.7% | 254.7 | 7.192 | 139 Flash EQ-ViT-H | 157.7 | 81.5 | 93.3 | 2.719 | 63.2% | 100.1 | 4.304 | 232 (1.7×)

Swin-T | 27.6 | 87.1 | 8.68 | 0.288 | 43.9% | 8.99 | 0.657 | 1521 Naive EQ-Swin-T | 6.9 | 87.3 | 8.68 | 0.285 | 43.3% | 8.99 | 0.657 | 1522 Flash EQ-Swin-T | 6.9 | 87.3 | 3.20 | 0.142 | 27.5% | 3.85 | 0.515 | 1942 (1.3×) Swin-S | 48.9 | 87.6 | 17.0 | 0.514 | 47.2% | 17.5 | 1.088 | 919 Naive EQ-Swin-S | 12.3 | 87.7 | 17.0 | 0.505 | 46.4% | 17.5 | 1.088 | 920 Flash EQ-Swin-S | 12.3 | 87.7 | 6.70 | 0.256 | 30.7% | 7.20 | 0.834 | 1200 (1.3×) Swin-B | 86.9 | 87.8 | 30.2 | 0.792 | 50.4% | 30.9 | 1.570 | 637 Naive EQ-Swin-B | 21.8 | 88.6 | 30.2 | 0.772 | 49.4% | 30.9 | 1.563 | 640 Flash EQ-Swin-B | 21.8 | 88.6 | 11.8 | 0.415 | 34.6% | 12.5 | 1.198 | 835 (1.3×) Swin-L | 195.2 | 87.5 | 68.0 | 1.630 | 56.3% | 69.0 | 2.894 | 346 Naive EQ-Swin-L | 48.9 | 88.2 | 68.0 | 1.630 | 56.3% | 69.0 | 2.894 | 346 Flash EQ-Swin-L | 48.9 | 88.2 | 25.2 | 0.812 | 40.3% | 27.6 | 2.015 | 496 (1.4×) Swin-H | 655.1 | 87.2 | 228.6 | 5.232 | 70.4% | 230.3 | 7.430 | 135 Naive EQ-Swin-H | 164.0 | 88.3 | 228.6 | 5.232 | 70.4% | 230.3 | 7.430 | 135 Flash EQ-Swin-H | 164.0 | 88.3 | 84.3 | 2.474 | 52.8% | 90.8 | 4.687 | 213 (1.6×)

加速比随模型规模单调增长:较小模型(如 ViT-T)获得的收益较小,因为线性层占总延迟预算的比例较小;而较大模型(如 ViT-H)则接近算子级别的 ~2× 加速,因为线性层主导了运行时间。重要的是,Flash EQ-ViT/Swin 在所有规模下都保留了 Naive EQ-ViT/Swin 对应模型的 ImageNet-100 Top-1 准确率,这证实了该加速在构造上是无损的。

20

第 21 页

表 8:在 FP16 下,Flash EQ-ViT/Swin 与标准非等变(non-EQ)和朴素等变(naive EQ)基线的推理效率对比。加速比(绿色)是相对于标准基线测量的。

所有线性层 | 整个网络 —|— 方法 | 参数量 (M)↓ | Top-1 (%)↑ | FLOPs (G)↓ | 延迟 (ms)↓ | 加速比 | FLOPs (G)↓ | 延迟 (ms)↓ | 吞吐量 (imgs/s)↑

半精度浮点数 (FP16)

ViT-T | 21.7 | 76.5 | 8.4 | 0.074 | 57.4% | 9.2 | 0.128 | 7794 Naive EQ-ViT-T | 5.4 | 77.3 | 8.4 | 0.074 | 57.4% | 9.2 | 0.128 | 7794 Flash EQ-ViT-T | 5.4 | 77.3 | 3.2 | 0.058 | 51.3% | 4.0 | 0.113 | 8848 (1.1×) ViT-S | 33.8 | 76.5 | 13.1 | 0.105 | 52.6% | 14.1 | 0.199 | 5030 Naive EQ-ViT-S | 8.4 | 78.2 | 13.1 | 0.109 | 54.5% | 14.1 | 0.200 | 5010 Flash EQ-ViT-S | 8.4 | 78.2 | 5.0 | 0.084 | 48.4% | 6.0 | 0.174 | 5750 (1.1×) ViT-B | 85.9 | 77.1 | 33.5 | 0.215 | 57.6% | 35.1 | 0.372 | 2688 Naive EQ-ViT-B | 21.4 | 80.1 | 33.5 | 0.218 | 58.3% | 35.1 | 0.374 | 2676 Flash EQ-ViT-B | 21.4 | 80.1 | 12.7 | 0.147 | 49.1% | 14.3 | 0.300 | 3330 (1.2×) ViT-L | 303.4 | 78.4 | 119.0 | 0.742 | 59.0% | 123.1 | 1.258 | 795 Naive EQ-ViT-L | 75.8 | 80.6 | 119.0 | 0.752 | 59.6% | 123.1 | 1.262 | 792 Flash EQ-ViT-L | 75.8 | 80.6 | 44.8 | 0.497 | 49.5% | 49.0 | 1.004 | 996 (1.3×) ViT-H | 631.1 | 81.3 | 248.0 | 1.541 | 64.8% | 254.7 | 2.379 | 420 Naive EQ-ViT-H | 157.7 | 81.5 | 248.0 | 1.557 | 65.1% | 254.7 | 2.392 | 418 Flash EQ-ViT-H | 157.7 | 81.5 | 93.3 | 0.976 | 54.2% | 100.1 | 1.801 | 555 (1.3×)

Swin-T | 27.6 | 87.1 | 8.68 | 0.091 | 39.1% | 8.99 | 0.232 | 4305 Naive EQ-Swin-T | 6.9 | 87.3 | 8.68 | 0.089 | 38.8% | 8.99 | 0.230 | 4350 Flash EQ-Swin-T | 6.9 | 87.3 | 3.20 | 0.085 | 38.5% | 3.85 | 0.220 | 4554 (1.1×) Swin-S | 48.9 | 87.6 | 17.0 | 0.173 | 47.8% | 17.5 | 0.362 | 2766 Naive EQ-Swin-S | 12.3 | 87.7 | 17.0 | 0.174 | 48.1% | 17.5 | 0.361 | 2770 Flash EQ-Swin-S | 12.3 | 87.7 | 6.70 | 0.150 | 45.0% | 7.20 | 0.332 | 3009 (1.1×) Swin-B | 86.9 | 87.8 | 30.2 | 0.240 | 45.5% | 30.9 | 0.527 | 1899 Naive EQ-Swin-B | 21.8 | 88.6 | 30.2 | 0.238 | 45.2% | 30.9 | 0.526 | 1902 Flash EQ-Swin-B | 21.8 | 88.6 | 11.8 | 0.185 | 39.4% | 12.5 | 0.470 | 2129 (1.1×) Swin-L | 195.2 | 87.5 | 68.0 | 0.454 | 47.7% | 69.0 | 0.952 | 1051 Naive EQ-Swin-L | 48.9 | 88.2 | 68.0 | 0.448 | 47.2% | 69.0 | 0.949 | 1053 Flash EQ-Swin-L | 48.9 | 88.2 | 25.2 | 0.317 | 38.8% | 27.6 | 0.817 | 1224 (1.2×) Swin-H | 655.1 | 87.2 | 228.6 | 1.428 | 59.2% | 230.3 | 2.412 | 415 Naive EQ-Swin-H | 164.0 | 88.3 | 228.6 | 1.406 | 58.4% | 230.3 | 2.408 | 415 Flash EQ-Swin-H | 164.0 | 88.3 | 84.3 | 0.958 | 49.0% | 90.8 | 1.955 | 512 (1.2×)

21

发表评论