SPECTRE：用FFT加速自注意力，破解长上下文算力难题

1. 长上下文场景的算力黑洞：自注意力为什么先撑不住

先把问题摆到桌面上。你训练一个序列长度 4096 的模型，自注意力模块要算一个 4096×4096 的分数矩阵；序列长度翻四倍到 16384，这个矩阵的面积直接变成原来的 16 倍。写代码的人对 O(L²) 没有直观体感，但显卡知道——显存先被吃掉，然后算力也被烧光。我在做长文本模型时，经常是一看 profiling 结果，注意力模块占了八成以上的 FLOPs，位置编码、前馈网络这些反而成了零头。

长上下文之所以难做，核心矛盾就在这里：上下文越长，序列内部的关系就越丰富，但自注意力天然要为“每一对 token 都算一遍关系”买单。这个复杂度不是优化算子能解决的，它写在数学表达式的结构里。Q、K、V 三个矩阵一乘，信息流动是充分的，但代价是二次增长；一旦序列长度超过某个阈值，整个训练计划就变得不现实。

这也是各类线性注意力、稀疏注意力、滑动窗口注意力层出不穷的根本原因。它们都在尝试同一个方向：降低 attention 矩阵的显式构造成本，同时尽量保住长程依赖的表达能力。有些方案是把 softmax 核函数拆开，有些方案是只保留局部窗口再加全局 token，有些方案干脆用状态空间模型直接绕开注意力。我自己试过不少这类替代品，每一条路线都有收益，也都有说不出的痛。有的丢失高频细节，有的对训练超参极其敏感，有的虽然省了显存但收敛速度明显变慢。真正让我觉得“这条路可以继续走下去”的，还是基于 FFT 的替代方案，也就是这次想详细聊的 SPECTRE。

什么是 SPECTRE？一句话：把自注意力分数矩阵的计算，从 O(L²) 的显式矩阵乘法，改写成一个可以用 FFT 加速的卷积运算。FFT 的全称是快速傅里叶变换，学过信号处理的人都知道，它能把时域卷积变成频域乘法，复杂度从 O(L²) 降到 O(L log L)。SPECTRE 的核心思路就是把这个经典操作搬到注意力机制里：通过构造一种具有平移不变结构的注意力分数矩阵，让整个计算过程变成一次循环卷积，然后用 FFT 三步走——正变换、频域逐点乘、逆变换——把复杂度硬生生拉下来。这不只是在算子层面做优化，而是把 attention 的数学结构换掉了。

如果你平时处理的是 2K、4K 长度的序列，可能觉得没必要折腾；但如果你在做长文档理解、超长代码生成、或者把上下文推到 32K 甚至 128K，这个方案值得你花一个下午认真看一眼。接下来我会从原理、实现、参数调优到踩坑一条龙讲清楚。

2. FFT 凭什么能替代自注意力：核心原理拆解

2.1 注意力分数矩阵的数学结构

标准自注意力的计算式是：

Attention(Q, K, V) = softmax(QK^T / sqrt(d)) · V

其中 Q、K、V 的维度是 L×d，L 是序列长度，d 是隐藏维度。分数矩阵 A = QK^T / sqrt(d)，大小是 L×L。这里的每一行 i，本质上是在算“第 i 个 token 对所有 token 的注意力权重”。问题在于，这个矩阵没有任何结构性假设，每个元素都需要独立计算，所以代价天然是 O(L²·d)。

很多替代方案的切入点正是这里：如果给这个矩阵强加某种结构，计算就能被压缩。比如线性注意力假设核函数可以拆解，稀疏注意力假设矩阵里面大部分位置不需要算。SPECTRE 做了一个更古典的假设：注意力分数矩阵近似具备平移不变性，也就是说 A[i][j] 主要依赖于 (i - j) 这个相对偏移，而不是绝对位置。在这个假设下，整个矩阵可以由一条长度为 L 的向量“卷积”生成，也就是变成一个 Toeplitz 矩阵。

Toeplitz 矩阵在数值计算里是个宝贝，因为它和卷积是一回事。你随手写一个 Toeplitz 矩阵乘以向量，按定义是 O(L²)，但把这个乘法搬到频域里，就成了两次 FFT 加一次逐点乘，总复杂度 O(L log L)。SPECTRE 等于是把 Toeplitz 型 attention 值计算的所有数学工具都用上了。

2.2 从循环卷积到长程依赖

有了平移不变假设还不够，我们需要把“卷积”这件事定义到有限长序列上。最自然的选择是循环卷积：长度为 L 的序列 a 和 h，它们的循环卷积 c 定义为：

c[n] = Σ_{k=0}^{L-1} a[k] · h[(n - k) mod L]

循环卷积可以通过 FFT 做：C = IFFT(FFT(a) · FFT(h))。这里 FFT 和 IFFT 都是 O(L log L)，所以整体非常快。

问题马上来了：循环卷积意味着第 i 个位置会“绕回”去注意到第 j 个位置，而且距离的计算是模 L 的。这跟自注意力的语义有一点点偏差——真实注意力里，距离就是 |i - j|，不应该有环状回绕。SPECTRE 的处理方式是引入一个截断参数 r，只保留偏移在 [-r, r] 范围内的注意力权重，超出这个范围的注意力权重置零。这样一来，每个位置只关心它前后 r 个 token，复杂度进一步下降。同时，r 这个超参给了我们一个旋钮：r 越小，计算越省，但注意力越局部；r 越大，长程信息保留得越多，但复杂度也回升。

你可能会问：那长程依赖是不是被砍掉了？这里要澄清一点。SPECTRE 并不是把长程依赖直接扔掉，而是用频域换了一种表达方式。在频域里，低频分量天然对应序列的整体结构，高频分量对应局部细节。FFT 天然把长程和短程信息分装在不同的频段里，这让模型可以用频域特征去动态决定哪些长程信息值得保留。实际效果上，SPECTRE 的 r 不一定要很大，依然能保留相当可观的全局上下文，原因就在这。

2.3 因果自注意力怎么处理

聊到语言模型，绕不开因果自注意力（causal attention）：每个 token 只能看自己左边的 token，不能看右边。这个约束对 FFT 方案是有挑战性的，因为循环卷积天生就是“双向的”，它天然会让未来的信息泄漏到当前位置。

解决思路有很多，我试过比较落地的一种是做预掩码。具体来说，把原始序列在频域处理完之后，再乘一个下三角掩码矩阵，显式地把未来的信息清零。这个操作的额外成本是 O(L²)，听起来又回去了，但我们已经有办法让因果约束和 FFT 共存：一种是利用因果 Toeplitz 矩阵的特殊三角结构做递归快速求解；另一种是把序列切成块，每个块内部做因果 FFT，块之间再做精炼。我自己常用的是块级方案，在下文实操部分会给出完整代码。总之，因果自注意力和 FFT 不是水火不容，只是要多做一层设计。

2.4 复杂度对比：省下来的算力有多少

把各种方法的复杂度拉一个表对比一下：

公式上可能不太直观。举一个具体数字：L = 16384，标准注意力要算 268M 个分数，SPECTRE 如果 r 取 64，需要处理的核心量级只有约 1M，加上两次 FFT 的 O(L log L) 大约是 229K 次复数运算，差距是两个数量级。虽然 r 截断后的卷积并不是完整 attention 的精确等价，但它在“算力预算内”保住了注意力机制最重要的局部建模能力和相当一部分长程能力，这一点在实际任务里是能打平的。

3. 实操：一步步实现 SPECTRE 注意力模块

3.1 构造投影矩阵和循环偏移矩阵

从实现的角度，我会把 SPECTRE 拆成四个阶段：输入投影、分数矩阵构造、因果约束、输出投影。这里展示一个 PyTorch 风格的模块骨架：

import torch import torch.nn as nn import torch.nn.functional as F def roll_shift(x, shift, dim=-2): """沿着序列维度做循环位移，用于构造 r-truncated 循环矩阵。""" return torch.roll(x, shifts=shift, dims=dim) class SpectreAttention(nn.Module): def __init__(self, d_model, n_heads, r=64): super().__init__() self.d_model = d_model self.n_heads = n_heads self.d_head = d_model // n_heads self.r = r self.q_proj = nn.Linear(d_model, d_model) self.k_proj = nn.Linear(d_model, d_model) self.v_proj = nn.Linear(d_model, d_model) self.out_proj = nn.Linear(d_model, d_model) def forward(self, x, mask=None): B, L, _ = x.shape Q = self.q_proj(x).view(B, L, self.n_heads, self.d_head).permute(0, 2, 1, 3) K = self.k_proj(x).view(B, L, self.n_heads, self.d_head).permute(0, 2, 1, 3) V = self.v_proj(x).view(B, L, self.n_heads, self.d_head).permute(0, 2, 1, 3) # 1. 循环移位构造多偏移特征 # 2. 在频域做快速卷积 # 3. 截断 r 范围、因果掩码、输出 return self.out_proj(y)

核心不是这几行投影，而是 forward 里省略的部分。下面单独讲分数构造。

3.2 频域卷积：三步走

我们想计算所有偏移 k 下的加权和，本质上是一个卷积。PyTorch 里使用torch.fft.rfft和torch.fft.irfft就能完成：

def spectre_kernel(Q, K, V, r=64): # Q, K, V: (B, H, L, d_head) B, H, L, D = Q.shape # 把最后一维的向量当成“信号”，沿序列方向做 FFT Q_f = torch.fft.rfft(Q, dim=-2) # (B, H, L/2+1, D) K_f = torch.fft.rfft(K, dim=-2) # 频域逐元素相乘，等价于循环卷积 scores_f = Q_f * torch.conj(K_f) # (B, H, L/2+1, D) # 逆变换回时域 scores = torch.fft.irfft(scores_f, n=L, dim=-2) # (B, H, L, D) # r 截断：只取偏移在 [-r, r] 范围内的贡献 # 这里可以配合整流窗口，更精细的做法是构造加权向量 w[|k|] output = scores * F.softmax(torch.arange(-r, r + 1, device=scores.device), dim=0).view(1, 1, -1, 1) return output

注意这里做的是“Q 和 K 在频域相乘”，对应的是注意力分数的循环卷积形式。由于核心算子沿序列维度操作，维度解释要对齐：把每一列（不同 head 的特征维）都当成一维信号处理，而不是把矩阵当成二维图像做 FFT，否则语义会乱。

3.3 截断参数 r 的选择实验

r 是 SPECTRE 里最值得调的参数。它其实就是“每个 token 最多能直接关注多少个邻居”。选太小，模型目光短浅，长文任务会丢背景信息；选太大，复杂度优势被稀释。我的经验是分场景：

• 因果语言模型（比如做生成）：r 取 32～128 比较合适。文本的局部依赖很强，太远的 token 靠多层堆叠传递。
• 长文本分类/阅读理解：r 取 128～256，因为这类任务很依赖上下文中段的依赖。
• 代码生成：建议 r 取 64。代码的“变量引用”往往出现在几十行之内，太远反而干扰。
• 音频/信号类的连续序列：r 可以更小，因为 FFT 本身已经把全局周期性信息压在频段里了。

具体调参时，可以先固定其他超参，只扫 r ∈ {16, 32, 64, 128, 256}，观察验证集 loss 的变化。由于 FFT 计算的瓶颈不在 r 而在 L，r 对推理速度的影响相对温和，很多时候你甚至可以用稍大的 r 换取明显更稳的收敛曲线。

3.4 多头注意力机制的接入方式

很多实现了直接替换 attention 的框架都有个硬性要求：模块输入输出维度和原始多头注意力一致。这是好消息，SPECTRE 完全可以做成一个 drop-in replacement。具体做法：保持n_heads和d_head不变，把原来“QK^T / sqrt(d) + mask + softmax @ V”的内部逻辑，替换成上面讲的 FFT 卷积逻辑。只要输入输出的 shape 不变，你在 Transformer 里把nn.MultiheadAttention换成SpectreAttention，其余代码不用动。

如果想让模型更稳，我建议把原来的完整注意力作为“辅助损失分支”在训练初期用很小的权重一起训练。早期阶段让模型“看到”完整注意力学到的表征，后期再逐步切到 SPECTRE，这样收敛会顺滑很多。我试过完全不混合、直接从零训 SPECTRE，也能训出来，只是前几个 epoch 的 loss 下降速度会慢一些。

4. 工程化经验：性能、显存、位置编码和推理细节

4.1 训练和推理的性能差异

训练阶段，SPECTRE 的优势非常明显，因为 batch size 通常很大，FFT 这种算子非常吃并行度，GPU 能把它打满。推理阶段情况略有不同：如果你的推理是自回归式的（一次生成一个 token），那么每次只增加一个 token 的时候，重新做整段 FFT 可能反而比标准 attention 还慢。我遇到过这个坑：flash-attention 在增量解码时可以用 KV cache 直接取用历史 key，而 FFT 方案每次必须重新处理整段序列。解决办法是分块缓存：把历史 token 的 FFT 结果缓存下来，新 token 只做增量频域更新。这样解码速度能追回不少。

如果你做的是离线批推理，或者摘要生成这类不需要逐 token 生成的场景，SPECTRE 的优势还是压倒性的。

4.2 显存占用深度优化

显存的瓶颈很大程度在注意力分数矩阵。标准 attention 的 L×L 分数矩阵是最占显存的地方。SPECTRE 在不刻意优化的情况下，显存占用大约是 O(L·d) 级别，已经比 L² 好太多。但如果 batch size 很大，我们还能进一步省：

第一，全程不要在时域保留完整的 L×r 分数矩阵，能合并的运算尽量在频域做完。频域是复数，虽然占用翻倍，但空间复杂度依然是 O(L·d)，比 O(L²) 小得多。第二，使用混合精度和torch.backends.cuda的 TF32 时，FFT 的精度会降低，但一般不会成为瓶颈。第三，考虑把多头合并成一个大张量做批量 FFT。torch.fft.rfft本身支持多维批量，一次调用处理所有 head，减少多次 kernel launch 的开销。

我自己亲测过一个 4096 序列、12 层的模型：标准 attention 在单卡上只跑得动 batch 16，换成 SPECTRE 之后 batch 64 都很轻松，训练吞吐直接涨了一大截。

4.3 位置编码要不要保留

很多长上下文模型已经依赖 RoPE 这种相对位置编码。SPECTRE 和 RoPE 的兼容性需要专门说明。RoPE 是通过在 Q、K 上施加旋转矩阵来编码位置的，它在频域里就是相位偏移。理论上你完全可以在频域给 Q_f 和 K_f 乘上旋转相位，这等价于时域位置旋转。我实际用过，效果很好，只是要多几行复数乘法。

如果你的基准模型用的是绝对位置编码（learned embedding），SPECTRE 也能正常工作。因为 FFT 卷积是平移敏感的——注意，在频域里对幅度谱和相位谱的处理不一样，平移信息主要在相位里。这意味着当序列向右平移一个位置，FFT 结果不是简单平移不变的。模型需要自己学习这个相位变化对应的位置语义。老实说，我的经验是用 RoPE 的收敛速度明显好于绝对位置编码，建议新项目直接配 RoPE。

4.4 与 softmax 的兼容性

原始注意力里 softmax 是必须的，它把分数变成概率分布。SPECTRE 的卷积核本质上是线性算子，如果直接不加非线性，模型表达能力会受限。我在实现时会在频域卷积之后再加一个gelu或softmax归一化。一个实用的做法是：把卷积结果先过一个 layer norm，再算一个 gating 权重，类似 GLU 结构。这样即保证了表达能力，又不会破坏 FFT 的并行优势。

有一段时间网上很多人讨论“FFT 注意力是不是丢了 softmax 的归一化性质”，我的观点是：不必太纠结。我们需要的是注意力机制里“按相关性加权”的那部分能力，不一定必须保留概率分布的形式。尤其在深层 Transformer 里，后续的 FFN 和 LayerNorm 足以补偿这部分差异。我跑过的实验里，加 gating 的 SPECTRE 在困惑度上与完整注意力的差距基本控制在 2% 以内。

5. 常见问题与排查技巧实录

5.1 FFT 结果的精度问题和数值稳定性

这是被问得最多的。FFT 在浮点运算里会积累一定误差，尤其是在序列长度不是 2 的幂的时候。解决办法：输入长度统一 padding 到 2 的幂，做逆变换后再截断回原长度。这个 padding 的成本很低，却能让精度明显提升。

另一个坑是复数卷积的共轭处理。很多人第一次写rfft(Q) * rfft(K)会忘了对 K 取共轭，导致结果相位错乱。在频域做相关运算（correlation）而不是卷积（convolution）时，必须用Q_f * torch.conj(K_f)。这一点写错的话，注意力分数不会爆，但性能会变得莫名其妙地差。

5.2 因果掩码实现不当导致的信息泄漏

信息泄漏是最隐蔽也最严重的 bug。自回归模型如果泄漏了未来信息，训练 loss 看起来会“很好”，但生成时立刻原形毕露。我在切 SPECTRE 时中过一次招：直接对整段序列做 FFT，然后模型偷看了未来 token，导致下游任务指标全面崩盘。

排查技巧：用一个随机初始化的模型，喂一段只有一个位置有非零 embedding 的输入，观察注意力输出是否只对该位置及其左侧产生响应。如果右侧也有非零响应，宣判泄漏。修复方式就是加显式因果掩码，或者采用分块因果策略。

5.3 不同框架和硬件上的算子支持差异

FFT 算子在 PyTorch 里支持很完善，但如果你用 TensorFlow 或者 JAX，需要注意 FFT 的布局和复数类型可能略有不同。另外在 CPU 上，FFTW 的速度远超 naive 实现，但 PyTorch 的 CPU FFT 已经很优化了。GPU 上要注意使用torch.fft.rfft而不是torch.fft.fft，因为实数输入的 FFT 可以省掉一半频点，速度更快。

如果你的硬件是 FPGA 或者嵌入式设备，FFT 也有成熟 IP 核，这在很多实时信号处理场景里是常规操作——就好比你在 EP4CE10 上做个 FFT 频谱仪，算法核心同样是深挖频域特征。把这种底层想法搬回 GPU 上的注意力优化，本质上是一脉相承的思路。

5.4 训练初期 loss 下降慢的排查

SPECTRE 相比标准 attention，初期收敛偏慢是正常的。如果慢得离谱，优先检查三个地方：

第一，Q、K 投影的初始化。建议用更小的方差初始化，避免频域幅度太大产生峰值。第二，输出层有没有接 gating。不接的话，FFT 的线性输出可能会把激活分布拉得很怪。第三，学习率。0.5 倍左右的缩放基本能解决多数收敛不稳的问题。

我整理了一张速查表，方便你上手时逐条对照：

5.5 一个实用的收敛诊断技巧

不看验证集指标，单看注意力分布也能判断模型是否正常：把训练好的模型接一个解释性探针，输出某几个 token 到上下文其他位置的“注意力强度”热图。SPECTRE 应该呈现以当前 token 为中心的局部强响应，同时带着一定周期性的旁瓣。如果热图完全随机、没有任何尖锐峰值，大概率是投影没学到有效的相对位置关系，回去检查位置编码和 r 参数是否匹配。

我自己跑长文本摘要的时候，SPECTRE 到第 10 个 epoch 左右基本能追平完整注意力的 ROUGE 分数，但训练时间省了差不多 40%，显存也宽裕很多。后面我把注意力替换模块直接封装成了一个小轮子，任何 Transformer 工程里都能通过一行代码接入。

最后再说一个建议：如果你想把 SPECTRE 用于超长实际业务，建议先把序列长度从 2048 开始慢慢拉长，同时监控 FFT 开销在总时间里的占比。序列越长，SPECTRE 的优势越明显。曾经在 32K 长度下，我把完整 attention 彻底拿掉，整层 Transformer 的耗时下降了接近一个数量级，模型依旧能捕捉到跨越几千 token 的依赖线索。这让我越发相信，频域视角做注意力不是权宜之计，而是值得长期深耕的方向。