长上下文一直是个让人又爱又恨的话题。模型能读的上下文越长能玩的花样就越多但注意力机制那 O(n²) 的复杂度也把训练和推理的成本顶得越来越高。跑到 8K、16K 甚至 32K 以上GPU 显存和时间消耗几乎是在用漏斗漏钱。也就是在这类场景里SPECTRE 这套基于 FFT 的自注意力替代方案逐渐被更多人注意到它把全局 token 交互从二次复杂度压到了 O(n log n)在一定程度上保住了建模长程依赖的能力同时让显存和耗时不至于跟着序列长度一起失控。这篇文章我会从长上下文的痛点讲起把 SPECTRE 的设计思路、关键实现细节、以及我实际调试时踩过的坑完整梳理一遍适合那些正在做长序列模型优化、或者单纯对高效注意力机制感兴趣的朋友参考。1. 为什么自注意力在长序列面前这么尴尬1.1 二次复杂度的残酷真相先回到最基础的问题一个传统 Transformer 层的自注意力输入是形状为[batch, seq_len, d_model]的张量经过 QKV 投影后会先算出一个[batch, heads, seq_len, seq_len]的注意力分数矩阵。这个矩阵的尺寸随序列长度的平方增长。简单算一笔账序列长度 4096 时单头注意力矩阵就有 1600 万个数换成 32K 长度这个数字变成 10 亿。即便用混合精度存储也要 4 到 8 个 GB 的显存。所以模型还没开始学什么复杂模式显存先被注意力矩阵撑爆了。这里还不只是显存问题。计算量同样按平方增长每个 token 都要和序列里所有其他 token 计算相关性序列越长计算图中的依赖规模越大。很多优化手段比如 Flash Attention解决的是 IO 开销和显存占用但本质计算复杂度依然是 O(n²)。也就是说Flash Attention 能帮你把内存效率和访存开销优化到极致却无法改变矩阵规模随序列长度平方增长的客观规律。真正想突破长上下文就得换一种信息交互的范式而不是继续在这个平方复杂度上打补丁。1.2 各种近似方案的局限这几年大家想了很多办法绕开二次复杂度。稀疏注意力只让每个 token 看到局部窗口或者特定跳跃位置复杂度降下来了但牺牲了全局感受野很多需要远距离信息联动的任务表现会打折。滑动窗口加全局 token 的方案稍微好一些但结构设计变得复杂总有一些 token 之间的路径变长了信息传递效率反而不高。线性注意力用核方法把注意力分数近似成特征映射的内积把复杂度拉到了线性但训练稳定性和表达力一直是老大难尤其是 softmax 归一化被替换后模型很难学到尖锐的注意力分布。还有各种状态空间模型比如用一类递推方式做全局压缩推理倒是爽快了但训练时的并行性处理和对某些任务的能力上限还处在被反复评估的阶段。所以我当时看到 FFT 这个切入点的时候第一反应是这条技术路线解决的是一个本质问题——如何在保持全局交互能力的前提下把计算复杂度砍下来。FFT 的 O(n log n) 在长序列下优势非常明显尤其是在 16K 以上的区间几乎就是降维打击。1.3 频域角度带来的新思路FFT 介入序列建模思路其实非常朴素把 token 序列当成一组离散信号每个特征维度就是一条沿着 seq_len 方向展开的信号曲线。全局交互不就是让序列里每个位置的信息互相渗透嘛而这些渗透模式在频域里天然是全局表达的。你可以把 FFT 理解为在一个“全连接”的信息交换层里用正交基做了一次坐标变换变换之后任意两个位置的相互关系被编码成各个频率分量上的幅度和相位天然具备全局视野。当然纯 FFT 只是把信息在频域上铺开如何从中提取有效特征还需要在频域里加可学习的筛选和组合机制这就是 SPECTRE 这类方案的核心设计空间。2. SPECTRE 的核心设计频域里做全局信息交换2.1 总体的结构定位从宏观结构上看SPECTRE 并没有推翻 Transformer 的骨架它替换的只是核心的 self-attention 子层其他部分比如前置 LayerNorm、残差连接、FFN 还是保留的。也就是说它可以直接嵌进现有的 Transformer 架构不需要对上层任务头做什么调整。一个标准的 SPECTRE 基础块大概长这样输入先经过 LayerNorm然后进入一个频域交互模块后面我们管它叫 Spectral Mixer这个模块输出之后接残差再做一次 LayerNorm最后过一个 FFN再残差。整体和 Transformer Block 几乎一一对应只是把 attention 那一大坨换成了 Spectral Mixer。从工程的角度看这种模块化替换的友好度极高。你不是从零搭建一个新的骨干网络而是把已有模型里最耗时的部分摘掉换新这在迁移已有权重和混合训练策略时非常有价值。2.2 频域混合的具体操作Spectral Mixer 的做法可以拆成几步第一步把输入张量x沿着序列维度做 FFT。注意这里不是对 d_model 做变换而是对 seq_len 维度做变换。也就是说每个特征通道都被当成一条长度为 seq_len 的离散信号变换后得到复数谱。第二步在频域中引入可学习的复数权重。最常见的方式是可学习的复数张量W_spectral形状和频谱一致与频谱做逐元素复乘这相当于对不同频率分量做选择性放大或抑制。也可以把复数权重拆成幅度和相位两部分分别建模更大一点的设计是先用一个线性层把 d_model 维度压缩再在频域做混合减少参数量。第三步对经过筛选调整的频谱做逆变换 IFFT把结果映射回时域 token 表示输出形状和输入保持一致。整个过程计算复杂度就是 FFT 的 O(n log n)相比自注意力最大的优势在于没有任何一个中间张量的尺寸是 n×n。2.3 为什么 FFT 能替代注意力的全局感知关键问题来了注意力机制的核心是让每个 token 根据内容动态选择要关注哪些其他 token这个“根据内容动态”的特性FFT 这种固定变换能做到吗答案就在“可学习的谱权重”上。FFT 本身确实是内容无关的固定变换它负责提供一个全局信息交换的基础操作。但后面跟的复乘权重是内容无关但可学习的也就是说模型可以通过训练把某种特定频率分量针对特定特征通道做强增强比如低频分量对应长距离平滑关联高频分量对应局部突变和细节模式。因此模型能依据训练数据自适应地调节哪些全局模式被保留哪些被抑制。当然它确实缺少了注意力那种“根据当前输入动态选择关注对象”的灵活性。但工程上很多时候就是一个平衡要让序列长得多就得接受一部分表达能力上的取舍。SPECTRE 的思路相当于用全局频谱混合来换取可控的复杂度在长序列任务里这种取舍经常是划算的。2.4 和位置编码的关系做长序列建模没法绕开位置信息。attention 里有位置编码嵌入频谱混合这边同样需要。值得强调的是FFT 本身并不丢失顺序信息因为离散傅里叶变换的输出顺序本身就映射着频率分量改变输入顺序会直接改变频谱相位。但要让模型显式利用位置关系还是得在输入里加位置编码。我的经验是直接用 RoPE 这类相对位置编码时注意力和 FFT 混合的适配性需要多测几组。RoPE 的旋转操作天然和复数域相关如果把它加在输入上再送进 Spectral Mixer有些任务效果不错有些任务反而引入了奇怪的偏差。更稳妥的做法是先用可学习的绝对位置编码或者 Alibi 这类加在注意力分数上的位置编码把位置信息和频域操作解耦让模型更容易学习。具体你要在自己的任务上跑几组对照没有一处通吃的银弹。3. 实现细节与调优理论到落地的关键一跃3.1 实数 FFT 还是复数 FFT实际代码里torch.fft.rfft 和 torch.fft.fft 是两种不同选择。rfft 是实数输入专用的输出只保留正频率部分长度从 n 变成 n/21省了一半的存储和计算。因为输入 token 特征通常是实数所以 rfft 是更高效的选择。但有一个小坑如果你在频域里做的是复乘复数权重的形状必须跟 rfft 的输出长度对齐也就是n/21而不是 n。很多人第一次写成 n 直接维度不匹配或者形状对上了但频域谱被错误截断输出质量会明显下降。我建议先打印一下中间张量形状确认一遍这个维度问题非常容易踩。如果你用了复数权重还需要注意 IFFT 之后输出会有极小的虚部残留。这是因为浮点计算误差导致 rfft 和 irfft 并不是严格互逆的。直接取实部就行千万别慌这是正常现象。我在实验里通常会对输出加一个x.real操作顺便把结果类型转回 float32。3.2 归一化策略的选择FFT 之后频谱的数值范围通常会比原始输入大不少尤其是低频部分如果不做归一化直接进后面的 FFN会因为数值波动太大导致训练非常不稳定。这里强烈建议在 Spectral Mixer 内部在 FFT 之后加上一个频谱归一化或者简化版的 RMS Norm。对于归一化一个简单做法是给复谱除以它的均方根模长这个操作本质上是对全局信息强度做尺度标准化不会破坏频率相对关系。如果希望更精细可以在 d_model 维度上做 BatchNorm 或者 LayerNorm但要注意批量大小不大时 BatchNorm 的小 batch 统计噪声问题。我自己实验的时候第一版没有做任何归一化训练 loss 前几百步抖得厉害。加上简单的频谱 RMS Norm 之后立刻稳定了收敛速度肉眼可见地提升。这个操作几乎可以说是 SPECTRE 类方法的标配建议默认开启。3.3 复数权重的初始化复数权重的初始化直接决定了模型起步阶段的稳定性。如果初始化不当复乘之后频谱幅度会被放大或者缩得太小逆变换回来信号幅度就会异常。我的经验是用幅度接近 1、相位接近 0 的初始化方式也就是权重实部初始化为单位矩阵附近的小扰动虚部初始化为零均值小量。这样周期一开始Spectral Mixer 近似为一个恒等映射模型可以在稳定状态下自主学习频域调整而不会因为频谱被剧烈改写导致梯度爆炸。相比之下如果直接按标准正态分布初始化复数权重模长平均下来有根号 2 那么大连续多个模块复乘后信号幅度会以指数级膨胀。这个看似不起眼的初始化细节很多时候就是训练能不能跑起来的差别。3.4 参数量和计算量的权衡FFT 本身没有可学习参数参数量完全来自频域权重和后面 FFN。如果频域复权重直接做成和频谱同尺寸那么参数量正比于 d_model 乘以 seq_len/2这在超长序列下会比较大。有几种压缩方式可以把复权重放在 head 维度共享也就是让一个权重张量在 d_model 分组之间复用也可以只在部分频率分量上加权重比如低频全保留高频共享一个缩放宽度的标量还可以在频域先用一个线性层压缩 d_model做混合后再扩张回来。第二种方法在图像类任务里被验证过实际效果和全尺寸权重差距不大计算省了很多。我的建议是模型规模不大时直接用全尺寸复权重模型规模大了再考虑压缩策略。你可以在小规模上对比一下全尺寸和压缩版的收敛曲线通常会看到压缩版在训练初期略慢但如果你更看重推理速度和显存占用这个代价是值得的。3.5 显存效率的实际对比我在一块卡上做过相对控制变量的对比实验模型差不多大小一个用标准 attention一个用 SPECTRE 替换 attention 层序列长度拉到 8192 的时候attention 版本的显存峰值大约是 SPECTRE 版本的 2.8 倍。拉到 16384 差距更明显标准 attention 直接因为中间矩阵溢出需要换更激进的分片策略而 SPECTRE 还能正常跑。时间维度上短序列比如 512 和 1024SPECTRE 不一定比 attention 快因为 FFT 的常数开销和复数运算在那里摆着甚至可能稍微慢一点。但当序列长度超过 4096 之后速度优势开始体现长度越长优势越明显。所以如果你只是做短文本分类这种任务没必要用 SPECTRE它的价值完全体现在长上下文需求真正存在的场景里。4. 训练和推理时的常见问题与排查实录4.1 训练 Loss 震荡我在初期训练中遇到的最典型问题就是 loss 震荡。后来定位到两个原因一是频域权重初始化覆盖范围太大周期性一开始就破坏了数据分布二是没有做频谱归一化局部 batch 的谱幅度差异过大导致梯度方向波动。排查的技巧是把每个 block 输入输出的数值标准差打出来。正常情况下经过 Spectral Mixer 后标准差波动应该在一个稳定区间内。如果发现某个 block 的输出标准差突然跳升两三个量级基本就是频域权重和归一化的问题。我修复的步骤是先换初始化再加归一化如果还有震荡就检查学习率是不是需要降一档。4.2 梯度范数异常另一个高频问题是梯度范数不降反升甚至出现 NaN。这种情况多半出在复数操作的梯度回传上。PyTorch 里的 FFT 对复数梯度支持没问题但如果你自己实现了复数权重的实部虚部拆开更新就容易出现梯度没有正确合成的情况。比如只更新了实部虚部拿到零梯度那么训练过程中虚部始终保持初始值模型的频域表达能力被锁死了一部分表现就是 loss 下降到一个平台期後就再也不动了。建议的做法是复数权重就用 Complex 类型维护优化器对复数参数更新时要确保实部和虚部都参与更新。或者在自定义层里把复数权重拆成两个 float 张量但前向计算时用torch.complex(real, imag)合成这样梯度能正确流到两个分量上。4.3 长序列推理的缓存策略推理阶段和训练不太一样。训练时你是把整个序列一次性送进去batch 维度和 seq 维度都很规整适合 FFT 的并行计算。但推理时如果是自回归生成每次只新增一个 token如果从头对全部序列做一次完整 FFT那就反而没有效率优势了。这种情况有两个思路。一个是用类状态空间的递推方式把频域混合结果以隐状态形式缓存下来新 token 来了只做增量更新。这需要把 FFT 操作改造成能循环递推的形式数学上可行但实现复杂度不低。另一个更务实模型内部还是全序列跑 FFT但只在生成每一个完整句子或者固定长度块时计算一次块内部用简单的注意力或者卷积过渡。实际工程里很多人是这么干的效果损失不大但吞吐能上来。如果你主要做的是离线批处理不需要逐 token 生成那推理完全不用操心直接复用训练时的方式即可。4.4 下游任务效果反而下降这是最容易让人劝退的坑。有人直接把模型里的 attention 全换成 SPECTRE跑到某些需要精确指代和细节信息定位的任务上发现效果掉了不少。这个不奇怪因为注意力有一种能力是频谱混合很难替代的对特定 token 对的精确指向性关注。处理这个问题的经验是不要把全部 attention 都替换掉。可以考虑只在后半段的某些层保留 attention或者每隔几层插入一个标准 attention 层来补齐精确匹配能力其他层全走频谱混合。这种混合架构在不少实验里效果比纯 SPECTRE 更好同时复杂度仍然远低于纯 attention。你可以把这种方式理解成用稀疏的精确注意力去修正频谱混合的模糊性有点像一个全局粗排加上局部精排的组合。我自己在做一个长文本检索类任务时就是保持最后两层标准 attention 不换前面的层全部用 SPECTRE效果和纯 attention 基本持平而训练时显存节省了几乎一半。这种折中方案特别适合兼顾质量和效率的项目。4.5 常见问题速查表为了方便排查我这里整理一个小表你们对照着检查就好。现象可能原因排查方向Loss 前期震荡剧烈复权重初始化不当、缺频谱归一化换成近恒等初始化加 RMS NormLoss 到平台期不动复数权重虚部没有梯度检查参数是否完整参与优化器更新梯度出现 NaNFFT 梯度异常、学习率过大检查复权重复合方式调低学习率短序列速度反而更慢FFT 常数开销大于收益短序列继续用 attention长序列再切下游任务精度下降全局精确交互能力不足保留部分 attention 层做混合互补推理逐 token 生成慢每次重算全序列 FFT增量状态缓存或块级重算5. 一些更深入的优化尝试5.1 多尺度频谱混合单一 FFT 覆盖的是整个序列长度所有频率分量一视同仁。但在真实文本里句内依赖、句间依赖、段落级依赖往往处在不同尺度上用一个 FFT 去捕捉所有尺度信息有点贪心。多尺度做法的思路是把序列切成长度不等的几个片段每个片段内部做 FFT 交互再跨片段做一次低频全局 FFT这样就同时兼顾了局部和全局。我的实验里最实用的配置是做两级第一级对长度为 256 的滑动窗口片段做频谱混合第二级对完整序列做一次步长较大的低频 FFT。相比单一 FFT这个结构在文本分类和摘要类任务上有一两个点的提升。代价是代码复杂度上升你需要管理好窗口边界对齐问题。如果你的项目对线上延迟敏感我建议先从单一 FFT 开始把多尺度作为后续迭代方向。5.2 频域与门控机制结合另一个提升频域表达力的方法是把门控机制引入到频谱混合中。具体做法是输入经过 FFT 得到复谱后额外用一个简单的小网络生成一组实数门控系数对频谱的幅度做动态加权。这样频谱权重不再只是静态可学习参数还能根据当前输入的内容特征动态调整相当于找回了部分 attention 才有的内容自适应能力。实测下来加上门控之后模型的表达力在某些语言建模任务上有明显增强而且增加的参数量很小。我建议在小数据集上先验证一下有没有收益再决定要不要在大模型里部署。需要留意的是门控网络本身必须是轻量的否则又会引入大量额外计算得不偿失。5.3 长序列分段拼接训练FFT 的 O(n log n) 虽然香但训练长度太长的时候比如超过 64K token显存仍然会在 FFN 部分形成瓶颈。因为 FFN 是逐 token 计算的序列多长 FFN 的激活量就多大这部分优化和 attention 无关必须单独考虑。一种做法是序列切分成若干 chunkFFN 分块处理。这样做损失很小因为 FFN 本来就是逐 token 独立的分块计算和整段计算在数学上等价只是要处理边界 normalization 的差异。另一种做法是用梯度检查点技术把 FFN 的中间激活扔掉反向传播时重新计算用时间换显存。把这两招组合起来通常能把能支持的最大序列长度再往上推不少。我自己尝试过把 FFN chunk 长度设为 4096配合梯度检查点训练长度从 16K 提升到了 32K 以上。5.4 给想上手的读者一个参考配置最后给个可以直接作为起点的配置6 层模型d_model 76812 个注意力头全部替换成 Spectral Mixer频域权重用全尺寸复权重低频保留、高频降权频谱归一化开着FFN 的隐藏层大小 2048序列长度 8192 起步。这个配置单卡就能跑显存占用大概 10 到 12 GB。如果你的任务需要更强能力深度加到 12 层d_model 拉到 1024序列长度推到 16K此时建议卡上显存不低于 32 GB。在中等规模的数据集上这个配置通常一周左右能收敛到不错的效果。训练曲线相对平稳对比同尺寸 attention 模型损失收敛下来的整体水平差距在可接受范围内而训练时能耗和显存下降明显。如果你已经有现成的长文本数据集我很建议先把这个配置跑通再来决定具体要在哪些层保留原生 attention哪些层换成频谱混合。写在后面的一点个人体会我自己是从标准 Transformer 一路用过来的一开始对 FFT 替代 attention 这件事心里也存疑总觉得这种固定变换是不是太粗糙了。但真跑完几个任务之后我的认知更多变成了长上下文场景里效率本身就是能力的一部分。能把 32K 序列塞进一张卡里训练跟塞不进去是完全两个世界。SPECTRE 不是万能的它在精确指代、细粒度关系建模上的能力确实不如标准 attention但通过混用部分 attention 层和持续的调优完全可以在效率和效果之间找到一个很适合自己的平衡点。另外忍不住再提醒一点如果你要做长文本方向强烈建议早点在代码库里把 FFT 和复数权重的工具链搭好后续不管是试多尺度、加门控、还是做增量推理都离不开这套底子。这些能力越早沉淀下来后面迭代越顺手。