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对每个特征维度单独做 FFTimport numpy as np def fft_attention(Q, K, V): # Q, K, V shape: (batch, heads, L, d_head) # L65536, 实测约 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, 980189SPECTRE (FFT)41.56, 35058128K 序列长度下标准 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切成若干块块内用标准 AttentionO(b²) 在 b4096 时可接受块间用 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-32KSPECTRE 开始追平并小幅反超显存优势约为 30%。超长上下文≥64KSPECTRE 拉开明显差距且序列越长优势越大。所以这个架构的合适应用区间是需要处理 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 我用了一个简单的决策流程给新项目做架构选择时我通常按下面这个流程判断省了非常多纠结的时间序列长度是否经常超过 8K如果不是直接用标准 Attention。别给自己找额外复杂度。8K 以上但任务里局部信息占比极高比如密集抽取式 QA选窗口稀疏注意力保留局部精确度。8K 以上且任务依赖全局信息整合摘要、长程推理、检索定位选 SPECTRE 式 FFT。显存是你的第一瓶颈任务刚好也是长文本为主不用犹豫SPECTRE 是最优解。有条件做分层混合强烈建议前 1/3 层用 FFT 或频域注意力后 2/3 层用标准 Attention这是我在多组实验里效果和效率最平衡的方案。6.3 给团队项目落地时的心得如果你打算在团队项目里引入 FFT 注意力我觉得提前想清楚两件事会非常有帮助。第一FFT 计算图和标准自注意力的计算图差异较大现有的许多工程基建需要额外适配比如参数初始化方式、学习率策略。我踩过的坑是直接把 Transformer 的标准初始化用到 SPECTRE 结构上前几轮训练 loss 下降异常缓慢后来发现频域复数滤波器的权重尺度需要更小的初始化。调整后训练表现恢复正常。第二这个架构更擅长统计全局模式而非精确定位局部值这是频谱压缩的天然代价。不要期待它在每个任务上都超过标准 Attention它的收益在长线上——省出的显存和算力可以让模型跑更大的 batch、更多的数据、更长的序列。我最后那版长文本模型正是在同样硬件预算下靠 SPECTRE 省出的计算量把训练数据从 300 亿 token 扩到了 500 亿 token整体效果反而迎来了质的提升。对长文本这个赛道来说效率优势可以反哺规模而规模带来的收益往往比架构本身更大。最后提一个值得关注的扩展方向把 FFT 和多尺度融合——低频段用较长窗口、高频段用较短窗口模拟出类似 Wavelet 的自适应分辨率效果。我做了初步尝试在保持效率的同时把摘要任务的分数据回升到了接近标准注意力的水平说明频域注意力的潜力还远没有挖完。