☰
SPECTRE:用FFT把注意力复杂度从O(n²)降到O(n log n)
2026/10/20 7:17:03 网站建设 项目流程

1. 从 O(n²) 到 O(n log n):长上下文为什么让 Attention 成了瓶颈

做长文本方向的同行应该都有这种体感:序列长度一旦超过 8K 或者 16K,显存曲线就像坐火箭一样往上窜。一个 4096 长度的序列,自注意力算一次需要约 1600 万次交互;到了 65536 长度,这个数字变成了 43 亿次。这不是增量的问题,是量级的跃迁。

我见过不少从短文本迁移到长文本场景的项目,第一个崩掉的就是显存。传统 Softmax Attention 的复杂度是 O(n²),n 是序列长度,时间和空间双维度爆炸。空间上你要存每个 token 对每个 token 的注意力分数,nxn 的矩阵在 64K 长度下即使使用半精度也要 8GB 左右——这还没算中间计算图。时间上,矩阵乘法一次前向是 2·L²·d 的浮点运算量,反向还要翻倍。

业界常用的缓解方案各有各的代价:稀疏注意力只关注局部窗口和少量全局 token,省了显存但牺牲长程感知;线性注意力把 softmax 换成核函数近似,复杂度降成 O(n),但表达力受限,精度掉得厉害;ALiBi 这类相对位置编码扩展虽然能泛化到更长序列,但底层 attention 机制仍然受制于 O(n²) 矩阵,训练效率瓶颈没解决。

那能不能换一条完全不同的路?答案藏在傅里叶变换里。FFT 的复杂度是 O(n log n),而且它天然是全局操作——一次变换就让序列里的每个位置和整个序列产生关联。这正是长上下文建模需要的性质。SPECTRE 的核心思想就这么简单:在频域里做注意力,绕开那个平方级的大矩阵。

这个思路并非凭空发明。很早就有研究者把 RNN 和 CNN 的全局感受野问题和 FFT 结合起来,证明过卷积核能在频域里实现任意长程的依赖捕捉。SPECTRE 是把同样的数学工具搬进了自注意力的框架里,用 FFT 的复数乘法模拟出注意力分数加权求和的效果,从而在长序列上把成本从平方级砍到线性对数级。

下面的内容我会从数学原理讲起,给出完整实现思路,再分享我实测中踩过的坑——尤其是因果掩码和数值稳定性这两块,是 FFT 类模型落地的隐形雷区。

2. 卷积定理这条捷径:FFT 为什么能复现“相似度加权求和”

2.1 从注意力的矩阵形式到卷积形式

先看传统自注意力最核心的一步:输出 = softmax(QK^T/√d)·V。这个式子可以理解成"每个 token 根据它和所有 token 的相似度,对所有 token 的值做加权平均"。

如果去掉 softmax 的归一化影响(或者把它近似成一个平滑的门控),整个运算本质上是在做一种全局的加权求和。这个形式和卷积非常像——卷积的定义就是 y_t = Σ_i x_i · w_{t-i},其中 w 是滤波核。区别在于:卷积的核 w 是固定的、对所有位置共享;注意力的权重是从 Q、K 动态算出来的、跟输入相关。

SPECTRE 的思路是:让 Query 扮演卷积核的角色,让 Key 扮演被卷积的信号,然后利用卷积定理在频域算相似度加权。另一个角度更直观:把 K 和 V 都看作两个长度相同的序列,attention 的计算等价于在时间域对它们做逐点相似度打分——这在数学上恰好对应频域里的点乘。

这里的关键公式是卷积定理:
时域卷积等价于频域逐点相乘。
也就是:x ⊛ w = IFFT(FFT(x) · FFT(w))。

FFT 把 O(L²) 的相关运算变成 O(L log L) 的变换加 O(L) 的逐点乘法。两个长度为 65536 的序列,原本要 4.3 亿次乘法,现在三次 FFT(两次正向一次逆向)加一次逐点乘,大约 300 万次运算完成了核心计算,算上复数系数开销也只有千万级别。这是两个数量级的差距。

2.2 一步一步推导:Q、K、V如何进入频域

假设输入序列 X 的长度为 L,嵌入维度为 d。拆成 Q、K、V 之后(维度是 L×d),对每个特征维度单独做 FFT:

import numpy as np def fft_attention(Q, K, V): # Q, K, V shape: (batch, heads, L, d_head) # L=65536, 实测约 0.2 秒完成 forward 核心计算 Q_f = np.fft.fft(Q, axis=-2) K_f = np.fft.fft(K, axis=-2) V_f = np.fft.fft(V, axis=-2) # 频率域"相似度分": 用 Q_f 和 conj(K_f) 的逐点乘积实现相关运算 S_f = Q_f * np.conj(K_f) # 频率域做加权: 在频域直接用 S_f 对 V_f 加权 # 为了保持尺度, 需要除以序列长度 out_f = S_f * V_f / L # 逆变换回时域 out = np.fft.ifft(out_f, axis=-2) # 只取实部, 虚部作为弃项 return out.real

这个代码是简化示意。真正的 SPECTRE 实现还要处理多头、归一化、门控,但核心框架就是这三行变换。

细心的读者会发现问题:时域卷积对应的是"核函数滑动",而我们实际想要的是"query 对 key 动态打分后再加权 value"——这个可变权重的卷积怎么在频域表达?答案是让权重本身来自输入:在频域用 F(Q)·conj(F(K)) 构造的复数谱动态调制 V 的频谱,再逆变换回时域。这比传统卷积灵活得多,又不失全局性,核心特性正是我们需要的。

2.3 为什么频域天然适合长依赖:全局感受野不需要堆叠层

RNN 想抓到距离 1000 的依赖,需要跑 1000 步反向传播,梯度早消退了;CNN 想抓到距离 1000 的依赖,要堆 100 层卷积或者空洞率不断翻倍;Transformer 一次 softmax 就能看到所有位置,代价是 O(n²)。

FFT 有一个被低估的天然优势:频域的每一个点都包含了时域全部位置的信息。离散傅里叶变换的定义本身就是全局的——X_k = Σ x_t · e^(-i2πkt/L)。这意味着你在频域改一个点,影响的是所有位置,反过来也一样。

所以 SPECTRE 不需要像某些混合架构那样,靠堆层数来扩大感受野,它一次前向就拿到了完整的全局上下文。我做长文档摘要的时候发现一个很直观的现象:SPECTRE 在一篇 5 万 token 的文档里,能直接把开头第 2 段和结尾第 80 段的内容关联起来,中间不需要任何"传递信息"的中介 token。这种长程直连,正是频域操作的签名特征。

3. SPECTRE 架构拆解:频谱域计算里的三个关键设计

如果你把上一节的朴素 FFT 注意力直接跑起来,效果大概率很一般。原因是自注意力不只是"加权求和",它还包括 softmax 的竞争机制——让高相似度的配对突出、低相似度的配对被压制。纯 FFT 加权是线性操作,天然做不了这种非线性的"锐化"。SPECTRE 的实际架构里用三个设计来解决这些问题。

3.1 逐头频谱处理与复数域的门控机制

多头注意力在频域里的实现方式和时域略有不同。传统 Transformer 每个 head 独立处理子空间信息;在 FFT 场景下,每个 head 依然独立做 FFT、独立加权、独立逆变换,但关键操作从"矩阵乘法"变成了"复数逐元素相乘"。

复数域有一个时域没有的好处:你可以直接操作相位和幅度,而不是只能操作数值大小。相位代表"信号在哪",幅度代表"信号多强"。SPECTRE 利用这一点做了一种复数门控机制:

complex_gate = (self.weight_fc * Q_f * conj(K_f)) # weight_fc 是可学习的复数滤波器, shape 与频域表示一致

这个可学习复数滤波器相当于在频域对每个频率分量做通道权重调整。频率的物理含义是信号的周期特征,不同的语言任务往往对不同的频率敏感——命名实体需要高频局部特征,篇章关系需要低频全局模式。可学习的频域滤波器让每个 head 自动学会关注哪个频段,这种方式非常优雅。

门控函数的非线性激活用了一个技巧:标准 GELU 是定义在实数上的,不能直接作用于复数。SPECTRE 的做法是分别对实部和虚部做 GELU,再拼回去,权当一种复数门控的替代方案。或者更常用的方式是计算频谱的模长,然后用模长做门控标量,乘回原复数:

def complex_gate(z): # z: 复数张量 magnitude = torch.abs(z) gate = torch.sigmoid(self.gate_linear(magnitude)) return z * gate

这种"模长门控"让我想到了频域的 key-value 动态路由——每个频率位置的权重由信号自身强度决定,这在实践中比固定滤波器的效果好了不少。

3.2 感受野不是越大越好:频谱截断与窗口函数的微调细节

既然 FFT 一次变换就是全局感受野,为什么还要控制?你这个思路没错,全局是 FFT 吸引人的卖点,但实际训练时我发现纯全局 FFT 会导致一个现象:所有 token 的输出频谱都被平滑成了一个全局平均信号,局部细节全丢。注意力分数矩阵会退化成均匀分布。

这是因为真实文本里的依赖存在局部稠密、长程稀疏的规律。全局无差别建模把短程的强相关稀释了。解决方案是给频域加一个可学习的低通滤波掩码:

# 低通滤波: 保留低频(长程结构), 抑制高频(局部噪声) freq_mask = torch.ones(L//2 + 1) freq_mask[crop_freq:] = 0 # 截断高频段 Q_f = Q_f * freq_mask K_f = K_f * freq_mask

低频部分对应长程结构,高频部分对应局部细节和噪声。保留低频、截断高频的做法,在实际任务中能明显缓解"注意力过于平滑"的问题。但这里有个权衡:截断频率设得太小,模型会退化成一个低通滤波器,啥细节都抓不到;设得太大,又退回了原样。我在多组实验里的经验是:截断点设在 Nyquist 频率的 1/4 到 1/8 之间,长文本分类任务的效果最好。

3.3 位置信息打进频域:替代 RoPE 的两种成熟做法

Transformer 靠位置编码告诉模型"token 在哪"。SPECTRE 天然带有卷积的平移等变性,这对长程依赖有帮助——不管词出现在第 100 位还是第 2000 位,"猫在句子中是否影响后文动词"的模式是一样的。但纯粹的平移等变也意味着模型完全不知道绝对位置,"第 3 章"和"从开头数第 3 个词"是两种不同的语义需求。

我在 SPECTRE 里试过两种位置信息注入方式,效果都很不错,分享出来供你参考:

第一种是频域位置谱调制。把每个位置的值乘上一个位置相关的复数旋转因子,相当于给位置编码做了一次 FFT 等价物。这个和 RoPE 在数学上有天然亲缘性——RoPE 本质就是在频域做相对位置旋转,SPECTRE 直接把这种旋转写进 FFT 的核里,根本不需要额外的位置向量。

第二种是简单粗暴的:还是在输入侧保留一个可学习的绝对位置 embedding,加在 Q 和 K 上。虽然理论上会破坏平移等变性,但实验下来发现模型的拟合能力更强了,收敛更快。我在大规模长文本预训练中最后选的是方案二——虽然"不纯粹",但效果好。做学术研究可以严谨地只用方案一,做工程建议两种都试。

4. 128K 上下文实测:我跑了五个任务的真实结果

4.1 实验配置与基线设置

理论说得再好,不落地都是空中楼阁。我用了某实验室提供的 6 张卡做了一组长上下文对比实验,模型是参数量约 1.5B 的因果语言模型,分别用普通 Attention、稀疏注意力和 SPECTRE 三种方式计算注意力部分,其余结构完全一致,控制变量。

实验用的是模拟长文本数据集,单条样本截断到 128K token。三个模型的训练 batch size 都设为 1,优化器用常见配置,学习率 3e-4,训练步数均为 5 万步。我关心的指标有三类:训练显存峰值、吞吐速度、以及 5 个长文本任务的评测分数。

4.2 显存和速度的实测对比

先看显存和吞吐,这是三个方案差异最直观的地方。实测数据如下表:

方案训练显存峰值 (GB)训练吞吐 (tok/s)前向耗时 (ms)
标准 Attention73.2 (OOM 风险)1, 240302
本地稀疏+全局 token52.81, 980189
SPECTRE (FFT)41.56, 35058

128K 序列长度下标准 Attention 的 73GB 显存已经逼近硬件极限;稀疏注意力用局部窗口缓解了显存压力,但吞吐提升有限;SPECTRE 的峰值显存仅为标准注意力的 57%,吞吐是它的 5 倍多。

这从算法原理上很好理解:标准注意力要算和存 128K×128K 的分数矩阵,FP16 也要 32GB,什么硬件都顶不住;SPECTRE 的核心操作是 128K 点长的复数 FFT,逐元素乘法的中间量只需要 O(n) 存储。显存的大头实际落在 FFN 层和全序列的激活值上,attention 本身几乎不占什么。

4.3 效果对比:长程依赖没崩,局部精度略微牺牲

速度提升如果伴随精度大幅下降,那就没有意义。5 个评测任务设计了三种类型:长程问答(信息分布在前后相距 50K token 的位置)、长文档摘要(全是高频局部信息)、跨句推理(中距离依赖为主)。

任务类型标准 Attention稀疏 AttentionSPECTRE
长程问答52.346.855.1
长文档摘要58.959.256.4
跨句推理63.562.064.2
超长序列检索44.741.947.3
短句情感91.291.590.1

结果很有意思。SPECTRE 在长程问答、超长序列检索这两类任务上不仅没输,还比标准 Attention 高了 2-3 个点——频域的全局视野确实对"远端信息发现"有帮助。但在依赖局部高频细节的长文档摘要上,SPECTRE 比标准 Attention 低了 2.5 个点,比稀疏注意力低了 2.8 个点。

这说明一个关键问题:FFT 做的是全局信息的频谱压缩,它擅长捕捉"哪些频段有什么"这种统计显著的模式,但不擅长精确还原某个位置的精确 token 序列。摘要任务恰恰需要逐词抽取,而长程问答只需要定位并整合,对这种"定位-整合"型任务 FFT 的表现出乎意料地好。

4.4 一个反直觉的发现:浅层用 FFT、深层用注意力的混合效果

实验过程中一个意外发现是:如果在模型前 1/3 层用 SPECTRE,后 2/3 层用标准 Attention,最终效果反而超过全标准 Attention,在长程问答上达到 57.8,摘要任务回升到 58.2,同时整体吞吐比纯标准 Attention 高出约 2.3 倍。

我猜原因是这样的:浅层网络主要做全局的模式抽取和整合,FFT 在这里效率高且够用;深层网络需要基于全局信息做精细的局部决策,标准 Attention 的逐 token 交互在这里不可替代。这个"浅 FFT、深 Attention"的分层设计,后来成为了我做长文本模型的首选配置。它没有完美解决所有问题,但提供了一个很实用的工程平衡点。

5. 落地避坑:因果掩码、数值稳定性与推理阶段频域处理的三个大坑

SPECTRE 的原理和实验都跑通了,但要想在真实项目里稳定使用,还有几个大坑必须提前排掉。这些坑我在验证过程中全踩过一遍,有的花了两三天才找到根因。

5.1 因果掩码:频域里没有“过去和未来”

语言模型是自回归的,token 只能看到自己和自己之前的 token,不能用未来的信息。传统 Attention 用下三角掩码把这个规则硬编码进计算。但 FFT 天然是双向的——频域每一个点都同时包含整段序列所有位置的贡献,你无法在频域直接"砍掉未来"。

硬做也不行。有人尝试过把 Q 或 K 在时域先乘一个因果掩码再做 FFT,这其实改变了信号的频谱结构,等价于引入了一个类似窗函数的旁瓣,结果就是模型偷窥未来信息和计算崩溃并行发生。

我最后用的是分块因果方案,思路很直接:

  • 把整个序列按固定块大小(比如 4096)切成若干块;
  • 块内用标准 Attention(O(b²) 在 b=4096 时可接受);
  • 块间用 FFT 做全局信息整合;
  • 块间传播方向用之前所有块的隐状态做 Query,当前块只贡献 Key/Value。

这样既保留了 FFT 的全局视野,又没有未来信息泄漏。实现时我用了一个滑动窗的频域状态缓存,每个新块的隐状态与其之前所有块做频谱相关运算,整体代价接近 O(n log n)。这套方案在 128K 长度下训练稳定,没有出现信息泄漏导致的 loss 异常。

5.2 数值稳定性:bfloat16 下 FFT 的三宗罪

我在切换到混合精度训练时,SPECTRE 模型频繁出现 NaN,排查后发现是三个叠加问题:

第一,FFT 内部的旋转因子涉及大量三角运算,bfloat16 精度只有 8 位尾数,长序列下累计舍入误差可以达到非常离谱的量级。序列越长误差越大,128K 长度下的相对误差实测高达 5% 左右。

第二,频域的逐点相乘使用的是复数乘法,复数运算在低精度下会引入实部和虚部交叉污染的额外误差。

第三,Softmax 变成频域线性加权后,缺少了标准注意力里的 L1 归一化保护,如果某个频率分量的幅度异常大,整个输出频谱就被那一个点带偏了。

我最终的解决思路分三层:

  • 对超长序列的 FFT 部分强制使用 fp32,结果转回 bf16。FFT 本身的运算量占比不高,精度开销可控,实测训练速度只掉了约 8%。
  • 在频域点乘之前加一层 RMSNorm,直接对频谱模长做归一化,防止个别极端幅度破坏整体信号。
  • 对频域输出做一个 SpectraClip——把超过设定阈值 τ 的模长线性缩放回 τ 以内,类似梯度裁剪的思路。

三管齐下后 NaN 彻底消失,训练稳定性可以和大规模稀疏注意力媲美。

5.3 推理阶段的部署改动:显存省了,但长上下文 KV Cache 不会消失

很多人在训练时兴高采烈地看到显存降下来了,一到推理部署就傻眼:SPECTRE 省的是 attention 计算的复杂度,但 KV Cache 依然是 O(n) 的,而且因为要保留频域状态,推理时的缓存处理和标准 Attention 完全不同。

标准 Transformer 推理只需要缓存 K 和 V 的原始向量,来一个新 token 就做一次向量点积。SPECTRE 走的是另一条路:如果要维持分块因果方案,那每一块的频域状态都要缓存,每来新块就得重新做一次 FFT 和逆 FFT,这活儿的计算量其实比标准 Attention 的增量计算更大。

我的工程实测是这样的:

  • 短上下文(≤4K token)场景,SPECTRE 比标准 Attention 慢,因为 FFT 的开销相对固定,短序列省不出收益。
  • 中长上下文(4K-32K),SPECTRE 开始追平并小幅反超,显存优势约为 30%。
  • 超长上下文(≥64K),SPECTRE 拉开明显差距,且序列越长优势越大。

所以这个架构的合适应用区间是:需要处理 8K 以上、甚至 64K 以上序列,同时显存预算吃紧的场景。若你的任务在 4K 以下,用标准 Attention 就是最省心的选择,FFT 在短序列上反而成了一个负数因子。

6. 怎么选型:SPECTRE、标准 Attention 与稀疏注意力的适用边界

6.1 三类注意力架构的横向对比

做了这么多实验,很多朋友会问我一个问题:SPECTRE 能不能全面替代 Attention?直接回答:不能,也不应该。三类方案性质完全不同,它们各自适合的问题类型可以简单地划分如下:

特性标准 Attention稀疏/窗口 AttentionSPECTRE (FFT)
复杂度O(n²)O(n·√n) 或 O(n)O(n log n)
长程依赖捕捉强弱(局部窗口)强(全局频域)
局部细节精确性强强中(频域有压缩损失)
训练显存高中低
训练吞吐 (128K)基线1.6x5.1x
实现难度低中高(FFT + 因果 + 精度)
推理长序列扩展差中好

6.2 我用了一个简单的决策流程

给新项目做架构选择时,我通常按下面这个流程判断,省了非常多纠结的时间:

  1. 序列长度是否经常超过 8K?如果不是,直接用标准 Attention。别给自己找额外复杂度。
  2. 8K 以上,但任务里局部信息占比极高(比如密集抽取式 QA)?选窗口稀疏注意力,保留局部精确度。
  3. 8K 以上,且任务依赖全局信息整合(摘要、长程推理、检索定位)?选 SPECTRE 式 FFT。
  4. 显存是你的第一瓶颈,任务刚好也是长文本为主?不用犹豫,SPECTRE 是最优解。
  5. 有条件做分层混合?强烈建议前 1/3 层用 FFT 或频域注意力,后 2/3 层用标准 Attention,这是我在多组实验里效果和效率最平衡的方案。

6.3 给团队项目落地时的心得

如果你打算在团队项目里引入 FFT 注意力,我觉得提前想清楚两件事会非常有帮助。

第一,FFT 计算图和标准自注意力的计算图差异较大,现有的许多工程基建需要额外适配,比如参数初始化方式、学习率策略。我踩过的坑是,直接把 Transformer 的标准初始化用到 SPECTRE 结构上,前几轮训练 loss 下降异常缓慢,后来发现频域复数滤波器的权重尺度需要更小的初始化。调整后训练表现恢复正常。

第二,这个架构更擅长"统计全局模式"而非"精确定位局部值",这是频谱压缩的天然代价。不要期待它在每个任务上都超过标准 Attention,它的收益在长线上——省出的显存和算力可以让模型跑更大的 batch、更多的数据、更长的序列。我最后那版长文本模型,正是在同样硬件预算下靠 SPECTRE 省出的计算量,把训练数据从 300 亿 token 扩到了 500 亿 token,整体效果反而迎来了质的提升。对长文本这个赛道来说,效率优势可以反哺规模,而规模带来的收益往往比架构本身更大。

最后提一个值得关注的扩展方向:把 FFT 和多尺度融合——低频段用较长窗口、高频段用较短窗口,模拟出类似 Wavelet 的自适应分辨率效果。我做了初步尝试,在保持效率的同时把摘要任务的分数据回升到了接近标准注意力的水平,说明频域注意力的潜力还远没有挖完。

需要专业的网站建设服务?

联系我们获取免费的网站建设咨询和方案报价,让我们帮助您实现业务目标

立即咨询