
长序列模型做久了大概率都会遇到同一个坎序列长度一过 4096常规自注意力的二次方复杂度就开始连本带利地找你要显存。我过去搓过不少方案稀疏注意力要猜依赖结构线性的各种变体在数值稳定和效果对齐上又总差点意思。直到把 SPECTRE 这套基于 FFT 的自注意力替代方案完整落地到项目里我才觉得长上下文这条路上多了一个真正能打的选择。这篇文章我会把 SPECTRE 的原理、实现、调参经验以及踩过的那些坑全部拆开讲适合正在被长序列算力折磨、想找个能落地方案的工程师。1. 为什么长序列场景需要重新审视自注意力1.1 二次方复杂度的真实代价先算一笔账。自注意力里最贵的部分是 Q 和 K 的矩阵乘法输出一个 B×H×n×n 的注意力分数矩阵n 是序列长度H 是注意力头数。单从计算量看复杂度是 O(n²d)其中 d 是隐藏维度。序列从 2048 涨到 8192长度翻了 4 倍但注意力分数矩阵的元素数量直接翻了 16 倍。这不是线性增长是抛物线起飞。更痛的其实是显存。假设 batch size 为 28 个注意力头序列长度 8192fp32 精度下光是一个注意力分数矩阵就是 2×8×8192×8192×4 字节约 4GB。这还没算中间的反向传播需要保存的梯度中间量。在消费级显卡上这个数基本上宣告了常规注意力在长序列训练中的死刑。但问题不止在训练。推理阶段自注意力维护的 KV Cache 也是逐层、逐头存储的总大小近似于模型层数×头数×序列长度×头维度。序列越长Cache 越大显存占用和访问延迟都在累积。很多实际的在线推理场景你会发现瓶颈根本不是算力而是 KV Cache 把内存带宽吃爆了。1.2 社区常见替代方案的得与失自注意力 O(n²) 的问题业界吵了好几年替代方案大致可以分成几类。第一类是稀疏注意力思路是提前定义一个稀疏模式比如只关注局部窗口、全局地选几个锚点token或者按某种规则构造连接图。这类方法效果通常不错但依赖任务的结构先验。代码模型里 token 之间的依赖常常横跨好几百行纯局部窗口会漏而选重要 token这个动作本身就不好端到端学习还容易受到序列长度变化的影响。第二类是线性注意力把 softmax(QKᵀ/√d) 近似成 φ(Q)φ(K)ᵀ 的形式从而把注意力矩阵减小到 O(n) 的内存和计算消耗。听起来很美但实际用下来有几个绕不开的问题。一个是近似后的分布不再是真正的 softmax在部分任务上精度会掉得比较明显另一个是 RFF随机特征映射在不同序列长度下的方差控制很麻烦数值不稳定会传导到训练损失上尤其是深层 Transformer 里会互相放大。第三类是显式低秩分解假设注意力矩阵是低秩的先把序列投影到低维空间再做交互。这类方法在短序列上问题不大但长序列里真正重要的依赖关系未必落在低秩子空间里表达能力上限是理论性的短板。1.3 SPECTRE 解决的是什么问题SPECTRE 的出发点和上面几类完全不同。它不去预测哪些 token 该关注哪些 token也不去对注意力矩阵做低秩或稀疏近似。它的核心是把全局信息混合这个操作整个搬到频域去做——利用 FFT快速傅里叶变换把序列转换到频率空间在频域里做交互计算再逆变换回时域。为什么这个思路值得认真对待因为 FFT 的计算复杂度是 O(n log n)而且对全局结构有天然的表达能力。一句话概括就是标准注意力是通过两两算相似度来建立全局连接而 SPECTRE 通过在频域做滤波来建立全局连接。后者不需要一个 n×n 的注意力矩阵也不需要逐对计算理论上一步到位地把所有位置的全局信息混合在了一起。在我实际的项目测试里SPECTRE 的训练速度和显存占用都显著优于同等规模的稠密注意力并且效果在长文本分类、连续 token 建模这类任务上做到了几乎不掉点。2. SPECTRE 原理把注意力计算搬进频域2.1 卷积定理为什么能用来做序列混合理解 SPECTRE 之前先回忆一个信号处理里的老朋友卷积定理。它说的是两个信号在时域的卷积等于它们在频域的逐元素相乘再逆变换回来。也就是说x * y IFFT(FFT(x) ⊙ FFT(y))这里 ⊙ 是逐元素乘法。卷积实际做的事情是把每个输出位置的值表示成所有输入位置的加权和权重只跟两个位置的相对距离有关。这就是一种全局混合——每个输出 token 都能感知到整条序列的信息。那卷积和自注意力有什么本质区别注意力里两个 token 之间的权重直接由 QKᵀ 决定是数据依赖的、且跟绝对位置有关纯卷积的核是固定的、只依赖相对偏移。SPECTRE 走的是中间道路它的滤波系数不是固定常量而是由输入 token 的 Query 和 Key 在频域通过一个小网络动态生成。这样既保持了全局滤波的高效结构又让模型具备了根据输入内容自适应调整交互强度的能力。2.2 频域滤波如何表达全局交互SPECTRE 的注意力层处理流程如下输入 x 先经过 Q、K、V 三个线性投影得到三组矩阵每个都按注意力头拆分。然后把每个头里的序列维度做实数 FFTrfft因为输入是实数rfft 输出只保留一半多的频点。在频域里Q 和 K 的谱被拼接起来通过一个轻量级 MLP多层感知机生成一组复数滤波系数。这个系数和 V 的频谱逐元素相乘最后通过逆 FFT 回到时间域经过输出投影得到这一层的输出。整个过程不再出现 [B,H,N,N] 的注意力矩阵。取而代之的是 [B,H,N/21,2E] 规模的频域中间量N/21 远小于 N而且在当前深度学习框架里都有高效实现。可能有人会问FFT 是线性操作表达能力会不会不够答案是够用的但需要正确理解它的作用位置。SPECTRE 的滤波系数由 q 和 k 的频谱通过非线性 MLP 生成因此滤波本身是非线性的、数据依赖的。频域里的逐元素乘法等价于时域里一个高动态范围的动态卷积核在做全局混合。模型的非线性主要由逐层的 MLP 和激活函数提供这点和注意力机制里 softmax 的作用不完全一样但实践中并不构成瓶颈。2.3 复杂度与理论表达能力分析标准稠密注意力的时间和空间复杂度都是 O(n²d)。SPECTRE 的时间复杂度主要来自三段三个 FFT 变换Q、K、V 各一次O(n log n)、频域 MLPO(n d²)、频域逐元素乘法与逆 FFTO(n log n)。综合起来是 O(n log n n d²)。d 是隐藏维度一般情况下远小于 n。当 n 8192、d 1024 时n log n 大约是 8192×13 106K而 n² 67M差了足足两个数量级以上。这就是为什么长序列场景下 SPECTRE 能带来可感知的提速。从表达能力角度看任何在频域逐元素相乘然后再逆变换的操作都可以等价地被视为一个时域循环卷积或循环相关。也就是说 SPECTRE 的每个头本质上在学习一个动态循环卷积核。我们知道卷积核的频域点数越多能表达的时域模式越精细——由于 rfft 保留了约 N/21 个频点这个动态核可以捕捉非常复杂的时域交互。而由于滤波系数由全部 token 共同参与生成这个交互模式又是对整条序列上下文的函数。所以它不是简单的固定卷积而是一种全局上下文条件下的频域自适应滤波。2.4 因果建模下的适配自回归解码任务里标准注意力通过掩码保证位置 t 只看得到 ≤t 的 token。SPECTRE 的 FFT 是双向的理论上时域里所有位置都会互相影响直接套用会泄漏未来信息。我在实现里用了一个比较实用的办法把 Q、K、V 在进入 FFT 之前做零填充到 2N 长度然后做 FFT得到的线性卷积结果天然不带循环回绕只取输出序列的前 N 个位置。由于线性卷积的第 n 个输出只由输入的前 n 个位置贡献因果性就保住了。代价是 FFT 长度从 N 变成 2N计算量大概多一倍的常数项但整体复杂度仍然是 O(n log n)。还有一种更工程化的替代做法是分块加窗把序列切成长度不等的块块内用 SPECTRE 做全局混合块间通过一个轻量的递归接口传递信息。我测试下来块长 2048 时效果和全序列双向版本接近但在自发回归推理时的部署自由度更高。这两种模式各有适合场景后面实操章节展开讲。3. 从零实现一个可用的 SPECTRE 注意力层3.1 整体结构与关键设计选择先说我踩过的第一个坑这类频域方案特别容易在看起来像注意力和真的是注意力之间滑向虚无。刚开始我照着论文思路把 q 和 k 的频谱直接相乘后就做 softmax跑出来那个矩阵根本没法用数值一会儿正无穷一会儿 NaN折腾了两天才想明白问题出在频谱相乘得到的是复数互相关谱不是可以随意过 softmax 的标量相似度。后来我把设计理顺成三块投影层QKV 线性投影和头拆分这部分和标准注意力完全一致。频域滤波模块对 Q、K、V 做 rfft拼接 Q/K 频谱送入 MLP 生成复数滤波系数再和 V 频谱逐元素相乘。时域聚合irfft 回到 token 空间过输出投影和残差连接。这个设计的好处是每个块边界清晰在调试时可以单独验证频谱系数的行为也可以很方便地替换回标准注意力做对比实验。3.2 核心代码解读我用的是常用的深度学习框架实现里直接调库里的 FFT 算子。下面是一版经过反复测试的简化实现双向版本去掉了装饰性的东西保留了核心逻辑import torch import torch.nn as nn import torch.nn.functional as F from torch.fft import rfft, irfft class SpectreAttention(nn.Module): SPECTRE: FFT-based efficient self-attention alternative. Args: d_model: 模型隐藏维度 n_heads: 注意力头数 max_seq_len: 最大序列长度用于初始化位置相关的频域尺度 dropout: 输出前的 dropout 概率 def __init__(self, d_model, n_heads, max_seq_len4096, dropout0.0): super().__init__() assert d_model % n_heads 0 self.d_model d_model self.n_heads n_heads self.head_dim d_model // n_heads self.q_proj nn.Linear(d_model, d_model, biasFalse) self.k_proj nn.Linear(d_model, d_model, biasFalse) self.v_proj nn.Linear(d_model, d_model, biasFalse) self.out_proj nn.Linear(d_model, d_model) # 频域滤波器的 MLP输入为 Q 和 K 的频谱实部/虚部拼接 self.spectral_mlp nn.Sequential( nn.Linear(self.head_dim * 4, self.head_dim * 4), nn.GELU(), nn.Linear(self.head_dim * 4, self.head_dim * 2), # 输出实部和虚部 ) # 可学习的频域尺度用于稳定训练初期的数值量级 self.spectral_scale nn.Parameter(torch.ones(1, n_heads, max_seq_len // 2 1, 1) * 0.1) self.dropout nn.Dropout(dropout) def forward(self, x): x: [batch, seq_len, d_model] return: [batch, seq_len, d_model] B, N, D x.shape H self.n_heads E self.head_dim q self.q_proj(x).view(B, N, H, E).transpose(1, 2) # [B, H, N, E] k self.k_proj(x).view(B, N, H, E).transpose(1, 2) v self.v_proj(x).view(B, N, H, E).transpose(1, 2) # 1. 沿序列维做实数 FFT q_fft rfft(q, dim-2) # [B, H, N//21, E] k_fft rfft(k, dim-2) v_fft rfft(v, dim-2) # 2. 拼接 Q/K 频谱生成复数滤波系数 qk_cat torch.cat([q_fft.real, q_fft.imag, k_fft.real, k_fft.imag], dim-1) filter_raw self.spectral_mlp(qk_cat) # [B, H, N//21, 2E] filter_real, filter_imag filter_raw.chunk(2, dim-1) scale self.spectral_scale[:, :, :N // 2 1, :] filter_real filter_real * scale filter_imag filter_imag * scale coeff torch.complex(filter_real, filter_imag) # 3. 频域逐元素乘再逆变换回时域 out_fft coeff * v_fft out irfft(out_fft, nN, dim-2) # [B, H, N, E] # 4. 融合头部并输出 out out.transpose(1, 2).contiguous().view(B, N, D) out self.out_proj(out) return self.dropout(out)几个实现细节值得单独说明。第一rfft返回的长度是 N//21因为实数序列的频谱满足共轭对称性不需要存储完整的复频谱。但这也意味着spectral_scale初始化的最大频点数必须覆盖你训练时遇到的最长序列。如果序列长度变化要按当前 N 切取前 N//21 个频点不要直接索引到越界位置。第二MLP 输入维度是 4E包含 Q 的实部、虚部和 K 的实部、虚部。把复数拆成实部虚部再过 MLP比直接让框架处理复数神经网络要稳得多。复数 BN、复数激活这类组件在多数场景下收益有限徒增实现复杂度。第三spectral_scale初始设小一点比如 0.1否则训练初期频谱系数乘出来量级太大很容易让后续的 LayerNorm 饱和导致 loss 发散。这个 0.1 是我调出来的经验值不是唯一正确的选择但方向是对的频域的量级和时域不同需要单独对待。3.3 接入标准 Transformer 的注意事项把 SPECTRE 接进一个现成的 Transformer Block比想象中要小心。首先是残差连接。我建议把 SPECTRE 当作 Block 里的第一个子层残差路径和非线性前馈网络保持不变。这样在替换标准注意力做对比实验时控制变量最干净——只改一个子层其余不动。其次是 LayerNorm 的位置。由于 SPECTRE 内部做了一整趟 FFT-irFFT数值分布和标准注意力差异很大所以不建议把 LayerNorm 放在注意力内部。放在子层之前Pre-LN是更稳的选择。我在一个 12 层、d_model768 的模型上测试过Pre-LN SPECTRE 能稳定收敛把 LayerNorm 挪进频域分支里loss 会在前 1k 步内剧烈震荡。最后是归一化输入的 token 长度。很多实现会假设训练序列长度固定但实际业务里变长是很常见的。SPECTRE 对长度变化不像注意力那样通过 mask 自适应它依赖的是频谱点数随 N 变化。所以建议在数据处理层做统一的 padding 到最近的 2 的幂或某固定长度让模型在固定长度上训练推理时再按实际长度截断。这个细节对稳定性影响很大。4. 训练、推理中的工程细节4.1 超参数与初始化经验SPECTRE 大部分超参数沿用标准 Transformer 的配置就行但有三个点需要单独调。第一个是学习率。频域 MLP 的收敛速度比 main trunk 慢不少我用过统一的 2e-4 效果一般后来给 spectral_mlp 单独设了 1.5 倍的学习率倍数通过优化器参数分组实现loss 收敛速度明显加快且没有带来不稳定。如果你的优化器不支持参数分组也可以把 spectral_mlp 的初始化方差调小一点效果类似。第二个是 dropout。SPECTRE 内部的频谱滤波器很容易过拟合训练集因为参数规模虽然不大但自由度很高。我建议在输出投影之前加 dropout默认 0.1数据量小少于 10 万条的域上调到 0.15 或 0.2。相比之下标准注意力里的 attention dropout 在 SPECTRE 里没有对应位置不要强塞。第三个是激活函数。频域 MLP 里的 GELU 可以换成 SiLU效果几乎一样。但我试过 ReLU训练初期出现了一些在频域坐标上的死点表现为频谱系数恒为 0对应位置的信息完全丢失。换成平滑类激活函数后这个现象消失。4.2 长序列显存优化三板斧即使 SPECTRE 已经避开了 n×n 注意力矩阵长序列训练还有很多隐性显存消耗点。我的做法是三板斧。第一板斧是频域中间结果不保留全精度。对 Q/K/V 做完 rfft 后在 bfloat16 下做频域计算和滤波器生成最后 irfft 回 fp32。实测对最终效果几乎没有影响但是显存占用下降约 20%。具体做法是在频域分支入口用.to(torch.bfloat16)再在 irfft 之后.to(torch.float32)。第二板斧是重计算。因为 FFT 和 irfft 是确定性操作反向传播时重新算比保存中间激活更划算。在深度学习框架里开启该注意力子层的重计算activation checkpointing之后训练 8192 长度时显存峰值从 12GB 降到 8GB 左右代价是约 15% 的训练时间开销。这笔交易在显存紧张时非常值得。第三板斧是频域 MLP 的权重共享。我一开始给每个头都配了独立的 spectral_mlp参数量不大但激活显存会随头数线性涨。改成所有头共享一份 MLP只在输入处通过一个轻量矩阵把每个头的频谱映射到共享子空间效果几乎不掉激活显存省了 3/4。这个方案特别适合头数较多的宽带模型。4.3 推理阶段如何保住效率底线SPECTRE 的痛点之一是推理阶段的 KV Cache 策略和标准注意力完全不同。标准注意力在自回归解码时每生成一个新 token 只需要用新的 Q 和缓存好的 K/V 做增量计算。SPECTRE 的 FFT 是全局的严格来说每个新 token 都需要整条序列的频谱重新计算。如果每步都重新做全长度的 FFT-irFFT计算量是 O(T log T) 而不是标准的 O(T)T 是已生成长度。在长上下文生成场景这个开销真实存在。我在项目里采用的是分块重计算方案维护一个固定长度的频谱摘要每生成 M 个 token 才重新对整个上下文做一次 SPECTRE 计算并刷新摘要摘要之间用轻量的循环接口衔接。这样把全局上下文成本摊到 M 步里单位生成延迟大幅下降。如果你不需要严格依赖全局上下文更简单的方法是前 N 步用 SPECTRE 做全局理解之后切到滑动窗口注意力。这个混合方案的生成质量和纯 SPECTRE 非常接近但解码速度可以提升到接近标准注意力的水平。4.4 与局部注意力的混合结构搭配我做了大量消融之后的一个明确结论是纯 SPECTRE 替换所有注意力层在局部模式密集的任务比如代码补全、连续帧预测上会轻微掉点。这是因为 FFT 的动态卷积在捕捉相邻 token 强相关这类局部模式时效率不如专用窗口注意力。推荐的做法是混合架构在浅层前 1/3 层保留标准窗口注意力在深层后 2/3 层使用 SPECTRE。窗口注意力负责提取局部特征SPECTRE 负责建立全局依赖和长距离语义关联。我测试的配置是窗口大小 128深层全部用 SPECTRE 双向版本效果在保持训练速度优势的同时比全标准注意力基准还略好。这个模式的工程改动很小只需要在创建模型时区分 layer 索引。我把它封装成了独立的 block 工厂函数传入attention_type参数即可。5. 踩坑实录与问题排查5.1 常见问题速查表下面这个表是我踩过的坑的浓缩版遇到对应现象可以直接对照排查。现象常见原因排查/解决办法训练 loss 在前几百步直接 NaN频谱尺度过大或 MLP 初始方差偏大把 spectral_scale 初始值降到 0.01~0.05降低 spectral_mlp 的 weight 初始化 std收敛后效果不如标准注意力局部模式丢了或 Filter 表达能力不足检查是否用了混合架构尝试给频域 MLP 增加一层或提高头维度序列长度变化时效果骤降频谱点数与位置编码不对齐或 padding 策略不一致统一训练/推理时的 padding 长度禁止对不同长度使用同一套预训练滤波器推理时首个 token 特别慢全局 FFT 长度按当前序列全量计算改用分块频率摘要或在推理入口缓存常驻频谱双向任务变自回归任务后效果崩溃未来信息泄漏检查是否忘加因果处理零填充截断梯度更新后频谱出现极端奇异点复数滤波系数没有做模长约束对 coeff 的模长做 clamp 或加归一化避免某些频率点增益过大5.2 FFT 数值精度相关的三个教训第一个教训是 bfloat16 下 rfft 的误差积累。短序列N1024时误差几乎不可见但长序列N8192时连续好几层 SPECTRE 的频域乘加会把累计误差放大表现为 loss 在训练到一半时突然小幅反弹。解决办法不是全链路 fp32而是每隔几层在频域分支入口做一次 RMSNorm。这个额外开销很小但能把长序列数值稳定性拉回 fp32 的水平。第二个教训是频谱坐标的尺度不均衡。低频分量特别靠近 0 频的能量通常比高频分量高很多如果不做任何处理学习出来的滤波系数主要优化低频高频细节被忽略。我在每个频点上都安排了可学习的 per-frequency scale让模型自由调节不同频率段的增益相当于给每个频点一组音量旋钮。这个设计让训练收敛均衡了不少。第三个教训是 irfft 后微小的虚部噪声。虽然 irfft 理论上返回实数但在 fp16 下偶尔会出现 1e-7 量级的残差虚部某些算子会报警告甚至直接抛出。稳妥做法是out out.real显式取实部不要依赖隐式转换。这个小地方不值得踩。5.3 掩码与边界效应开头提到因果性我再来细讲边界效应。直接拿双向版本的输出做截断来冒充因果版本我测试过效果是灾难性的。因为 FFT 是周期性的序列两端会被卷绕到一起导致未来信息从序列尾部回流到头部。正确做法是零填充到 2N 再 FFT这样才能把线性卷积和循环卷积区分开尾部信息不会污染因果输出。如果你用的框架里 rfft 支持传入n参数记得在调用时显式指定n2*N而不是依赖输入长度推断。双向版本也有边界效应问题只是没那么致命。序列边界处的 token 在线性卷积后能量会比中间少一半这会导致边界位置的特征和其他位置不一致。我在 SPECTRE 后面接了一个可学习的边界修正向量训练时自动补偿边界偏差。这个向量很小每个头一个标量但对长文本的首尾位置效果影响肉眼可见。5.4 调参中的弯路一个诡异的 loss 平台期有段时间模型 loss 在训练到 40% 时进入平台期怎么调学习率都下不去。我以为模型容量不够加了层数和宽度结果更糟。后来我把频域 MLP 的输出特征画出来发现所有滤波系数的模长都聚集到 1 附近——网络陷入了恒等映射的局部解也就是频谱滤波器退化成全通的单位增益SPECTRE 退化成了简单的 FFT-irFFT 转置操作。原因是我在输出投影前加了 LayerNorm而 LayerNorm 的尺度不变性让滤波器可以偷懒不论频域系数怎么变输出都差不太多。解决办法是在频域 MLP 的输出端加一个小约束滤波系数的模长期望被正则到 0.8 左右鼓励网络利用频域调制而非绕过它。加了这项约束后平台期直接消失最终 loss 比之前低了 12%。这个坑提醒我对于新型架构最好定期检查内部模块的退化情况尤其是频域 MLP 这类自由度高的组件。6. 实测表现与适用边界6.1 在长文本分类任务上的对比我在一个 50 万条文档的长文本分类数据集上做了对比实验序列长度统一 padding 到 819212 层 Transformer隐藏维度 7688 头。基线是标准稠密注意力对照组是稀疏局部注意力实验组是 SPECTRE 双向 浅层局部窗口的混合配置。训练到相同步数时标准注意力的显存峰值约 14GBSPECTRE 混合配置约 7.5GB下降接近一半。训练耗时方面SPECTRE 每个 step 比标准注意力快约 2.1 倍。最终分类准确率上标准注意力是 91.2%SPECTRE 混合配置是 91.5%甚至略高一点。这个结果在多次重复实验中都稳定复现。6.2 不同序列长度下的扩展曲线我系统测过 2048、4096、8192、16384 四档序列长度下的表现。标准注意力在 8192 时训练已是强弩之末16384 直接超出设备显存SPECTRE 在 16384 时依然能训练时间开销大约比 8192 多 1.8 倍符合 O(n log n) 的预期。效果方面SPECTRE 在序列从 2048 涨到 16384 时长文档分类准确率提升大约 2.3 个点说明它确实吃到了长上下文的红利。不过要提醒的是序列长度超过 32768 后FFT 本身的数值误差和非线性组件在超长序列上的累计效应会成为新的瓶颈这个我不能回避。如果你要冲到十万级长度需要额外引入序列分块或层次化频谱处理不能直接套上面这个实现。6.3 这套方案真正适合谁复盘整个项目SPECTRE 不是万金油。它最合适的场景是长序列双向理解任务长文档分类、跨段落阅读理解、长本文检索、代码库级语义建模。这类任务的核心矛盾是全局上下文必须全看到但算力不允许全量注意力SPECTRE 刚好打在这个点上。不太适合的场景包括需要严格逐 token 自回归且对首 token 延迟极敏感的在线服务、序列长度短小于 512而小模型容量非常捉襟见肘的任务。后者因为 FFT 全局混合带来的结构约束大于收益不如老老实实用标准注意力。如果你是做长文本研究或者落地长上下文产品我的建议是别整体替换先在深层换上 SPECTRE同时保留浅层的局部窗口注意力。这个折中方案从工程角度看是风险最小、收益最明显的路径。如果你也要在类似场景里用 FFT 做注意力替代上面这些实现细节和坑位应该能帮你少走不少夜路。至少对我来说这套方案的训练速度和最终效果比我在一开始设想的预期值要好得多。