长序列模型做久了,大概率都会遇到同一个坎:序列长度一过 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 为 2,8 个注意力头,序列长度 8192,fp32 精度下,光是一个注意力分数矩阵就是 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 三个线性投影,得到三组矩阵,每个都按注意力头拆分。然后把每个头里的序列维度做实数 FFT(rfft,因为输入是实数,rfft 输出只保留一半多的频点)。在频域里,Q 和 K 的谱被拼接起来,通过一个轻量级 MLP(多层感知机)生成一组复数滤波系数。这个系数和 V 的频谱逐元素相乘,最后通过逆 FFT 回到时间域,经过输出投影得到这一层的输出。
整个过程不再出现 [B,H,N,N] 的注意力矩阵。取而代之的是 [B,H,N/2+1,2E] 规模的频域中间量,N/2+1 远小于 N,而且在当前深度学习框架里都有高效实现。
可能有人会问:FFT 是线性操作,表达能力会不会不够?答案是够用的,但需要正确理解它的作用位置。SPECTRE 的滤波系数由 q 和 k 的频谱通过非线性 MLP 生成,因此滤波本身是非线性的、数据依赖的。频域里的逐元素乘法,等价于时域里一个高动态范围的动态卷积核在做全局混合。模型的非线性主要由逐层的 MLP 和激活函数提供,这点和注意力机制里 softmax 的作用不完全一样,但实践中并不构成瓶颈。
2.3 复杂度与理论表达能力分析
标准稠密注意力的时间和空间复杂度都是 O(n²d)。SPECTRE 的时间复杂度主要来自三段:三个 FFT 变换(Q、K、V 各一次,O(n log n)、频域 MLP(O(n d²))、频域逐元素乘法与逆 FFT(O(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/2+1 个频点,这个动态核可以捕捉非常复杂的时域交互。而由于滤波系数由全部 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_len=4096, dropout=0.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, bias=False) self.k_proj = nn.Linear(d_model, d_model, bias=False) self.v_proj = nn.Linear(d_model, d_model, bias=False) 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//2+1, 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//2+1, 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, n=N, 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//2+1,因为实数序列的频谱满足共轭对称性,不需要存储完整的复频谱。但这也意味着spectral_scale初始化的最大频点数必须覆盖你训练时遇到的最长序列。如果序列长度变化,要按当前 N 切取前 N//2+1 个频点,不要直接索引到越界位置。
第二,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_model=768 的模型上测试过: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 的误差积累。短序列(N<1024)时误差几乎不可见,但长序列(N>8192)时,连续好几层 SPECTRE 的频域乘加会把累计误差放大,表现为 loss 在训练到一半时突然小幅反弹。解决办法不是全链路 fp32,而是每隔几层在频域分支入口做一次 RMSNorm。这个额外开销很小,但能把长序列数值稳定性拉回 fp32 的水平。
第二个教训是频谱坐标的尺度不均衡。低频分量(特别靠近 0 频)的能量通常比高频分量高很多,如果不做任何处理,学习出来的滤波系数主要优化低频,高频细节被忽略。我在每个频点上都安排了可学习的 per-frequency scale,让模型自由调节不同频率段的增益,相当于给每个频点一组"音量旋钮"。这个设计让训练收敛均衡了不少。
第三个教训是 irfft 后微小的虚部噪声。虽然 irfft 理论上返回实数,但在 fp16 下偶尔会出现 1e-7 量级的残差虚部,某些算子会报警告甚至直接抛出。稳妥做法是out = out.real显式取实部,不要依赖隐式转换。这个小地方不值得踩。
5.3 掩码与边界效应
开头提到因果性,我再来细讲边界效应。
直接拿双向版本的输出做截断来冒充因果版本,我测试过,效果是灾难性的。因为 FFT 是周期性的,序列两端会被"卷绕"到一起,导致未来信息从序列尾部回流到头部。正确做法是零填充到 2N 再 FFT,这样才能把线性卷积和循环卷积区分开,尾部信息不会污染因果输出。如果你用的框架里 rfft 支持传入n参数,记得在调用时显式指定n=2*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 到 8192,12 层 Transformer,隐藏维度 768,8 头。基线是标准稠密注意力,对照组是稀疏局部注意力,实验组是 SPECTRE 双向 + 浅层局部窗口的混合配置。
训练到相同步数时,标准注意力的显存峰值约 14GB,SPECTRE 混合配置约 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 做注意力替代,上面这些实现细节和坑位应该能帮你少走不少夜路。至少对我来说,这套方案的训练速度和最终效果,比我在一开始设想的预期值要好得多。