FreeSolo：线性时间稀疏注意力实现零样本长度外推

分离记忆定位与证据读取的稀疏注意力，4k 训练零样本外推到 256k

前置知识

RoPE 旋转位置编码

RoPE 在计算注意力分数前，分别用位置 $i$ 和 $j$ 的旋转算子 $R(i),R(j)$ 旋转 query 和 key，使得 $q_i^\top k_j$ 只依赖相对偏移 $j-i$。模型训练时只见过有限范围内的相对偏移，超出该范围的偏移属于分布外输入，往往导致注意力分数失真、检索崩溃。

FreeSolo 的关键位置设计是 RoPE 只在固定 512 宽的局部窗口内使用，长程分支完全不旋转；理解 RoPE 的相对偏移敏感性问题，才能明白这套「局部 RoPE + 长程 NoPE」设计为何能零样本外推。

原生稀疏注意力（NSA / MoBA）

指在预训练阶段就引入稀疏性的注意力架构，每个 query 只对选中的少数块做注意力。MoBA 用各块的均值 key 与 query 打分路由选出 top 块再读 token；NSA 学习压缩的 key/value，用压缩注意力导出选择分数并把压缩值混入读出。两者在固定块因子 $b$ 下复杂度为 $O(N^2/b)$。

FreeSolo 延续了 MoBA 的均值 key 路由思想，但换成了范数先验、纯 token 读取和局部 RoPE；对照这两个基线才能看清本文的增量与本质区别。

循环状态线性模型（Mamba / GDN）

线性时间序列模型把历史压缩进固定大小的循环状态，每步 $O(1)$ 更新、整序列 $O(Nd)$。推理状态恒定是优点，但状态容量有限，远距离精确检索（如大海捞针）信息会被压缩丢失。GDN（Gated DeltaNet）是其中的代表性门控delta规则模型。

循环模型是本文线性时间组的主要对照：GDN 在 800M/16k 的 RULER 均值仅 10.4、NAtS-L 仅 2.8，说明单纯压缩历史不足以支撑长程检索，这正是 FreeSolo 要解决的核心矛盾。

RULER 长上下文基准

NVIDIA 提出的长上下文能力评测套件，含单针/多针大海捞针（S-NIAH/MK-NIAH/MV-NIAH）、变量追踪 VT、多跳问答 QA 等任务，论文设置下每任务每长度 100 个样本、贪心解码、官方 0-100 分。论文强调用检索分数而非困惑度，因为 token 平均会掩盖依赖远距离上下文的错误。

论文全部核心结论（零样本外推、256k 单针检索、任务强弱差异）都以 RULER 分数量化，读懂主表 Table 1 的结构和数字必须先熟悉该基准。

持续预训练（CPT）

在已有预训练 checkpoint 上继续训练一个短阶段，通常重置优化器与学习率调度。长上下文研究中 CPT 常用于在更长序列上适配位置编码或稀疏模式，但长序列训练代价高昂。本文的特殊用法是：CPT 仍在 4096 长度上进行，只用于让模型表示适应剪枝等稀疏约束，而非延长训练上下文。

FS-L 和 FS-B 都由基线 FS 经极短 CPT（约预训练 token 量的 5%，0.42B-1.57B token）得到，两阶段训练配方是理解方法与复现结果的关键。

NoPE（无位置编码注意力）

注意力计算中完全不加位置信息的方案。因果掩码本身已隐含顺序信息，且任何序列长度下打分函数形式完全相同、不存在分布外的位置偏移，因此 NoPE 天然具备一定的长度外推能力。erf 距离先验等偏置可以只作用于检索打分而不进入最终 softmax。

FreeSolo 的长程分支就是 NoPE token 注意力：位置偏置只影响 cell 选择，readout 完全无位置编码，这是长距离检索可以零样本外推的机制基础。

研究动机

长上下文语言建模需要在超出训练长度后依然保持高效注意力与精确检索，但现有方案各有硬伤。全注意力（TF++）代价为 $O(N^2 d)$，且零样本外推彻底崩溃：800M 的 TF++ 在 4k 训练后，S-NIAH-1 从 4k 的 95 跌到 8k 的 0，16k 的 RULER 均值为 0.0。以 Mamba、GDN 为代表的线性模型把历史压进固定大小循环状态，容量受限：800M 的 GDN 在 16k 的 RULER 均值只有 10.4，S-NIAH-3 只有 5；混入保留 token 注意力的 NAtS-L 均值仅 2.8、S-NIAH-3 为 0。MoBA 与 NSA 等原生稀疏方案保留 token 级访问，但要对随上下文线性增长的块集合持续打分（$O(N^2/b)$），实验显示它们同样无法在 8k 之外维持召回：NSA 在 800M/16k 均值仅 0.1，MoBA 在 340M/8k 均值已归零。业界的常规补救是继续预训练（CPT）在更长序列上扩上下文，如位置插值与 YaRN，但这带来额外的高昂训练成本，且超出训练长度后的泛化依然不稳定。

本文的目标是本文的目标是设计一种原生稀疏注意力机制，在全程只用 4096 token 训练长度的前提下实现零样本长度外推与长程检索，同时把注意力代价压到线性或严格亚二次。具体拆为三个可检验的子目标：其一，长程检索不能依赖对不断增长的前缀逐块穷举打分，每 query 的读预算必须随历史增长保持有界（每头最多读 $Ks=128$ 个 token）；其二，机制必须免疫位置分布偏移，长距离分支不接触训练中未见过的 RoPE 相对偏移；其三，短上下文通用能力不能因稀疏化明显损失，800M 规模下游 8 任务平均要接近全注意力 TF++ 的 56.8。围绕这些目标，作者提出基线机制 FreeSolo（FS）及两个变体：线性时间的 FS-L（$O(N)$、状态有界）和平衡树搜索的 FS-B（$O(N\log^2 N)$ 含 radix sort），并要求两者都只需极短的 CPT 即可获得外推能力。

与已有工作不同的是，本文最独特的切入角度是把「记忆定位」与「证据读取」彻底分离：路由层只需用极粗的免学习统计摘要判断「证据可能在哪」，真正的注意力读取始终发生在原始 token key/value 上。此前工作要么像 NSA 学习压缩 key/value 并混入读出，要么像 Landmark/InfLLM 额外学习 landmark 或代表 token，都引入新的可学习模块；FreeSolo 只复用每层已有的 $W_k$ 投影，用均值 key $\bar{k}_c$ 与最大 key 范数 $a_c$ 两个统计量完成路由与剪枝。第二，位置编码的头内混合是以往没有的：RoPE/NoPE 混合（Kimi Linear 等）都在层级别实现，FreeSolo 在每个头内部同时耦合局部 RoPE 窗口与稀疏 NoPE 长程分支，4k 训练全程不产生任何未见过的旋转偏移。第三，FS-L 把 H2O/SnapKV 那类依赖累积注意力分数的事后 KV 剪枝，改造成「固定规则 + CPT 适配」：保留规则固定为高范数优先，让表示在短 CPT 中主动适应被限制的访问模式。

核心方法

FreeSolo 的直觉是：在不断增长的记忆里找证据，不必对每条记忆精细打分，只需足够好的「目录页」。每个注意力头有两个分支：局部滑窗分支用宽度 $w=512$ 的窗口做标准 RoPE 注意力；长程分支把前缀切成 $s=32$ token 的 cell，用均值 key $\bar{k}_c$ 做内容路由、用 cell 内最大 key 范数 $a_c$ 做查询无关的重要性先验，为每个 query 打分选出 $K=4$ 个 cell，再只在这 128 个原始 token 上做无位置编码的 softmax 读出。两分支共享投影，输出按每头可学习系数合成；cell 打分由内容匹配项加范数加权的 erf 相对位置偏置构成。骨架上有三个变体：FS 对整个前缀打分，二次代价；FS-L 先按范数剪枝，每头只留 $P=64$ 个高范数 cell 加 $W=8$ 个最近 cell，线性代价、状态有界；FS-B 在完整档案上建两棵平衡树做 beam search，代价 $O(N\log^2 N)$。查询按 $B=128$ 分块，块内共享前缀可用集，因果掩码允许并行训练。

核心创新有三点。第一，「摘要只为定位、读取只用原始 token」：NSA 把学习的压缩 KV 混入读出，MoBA 虽读选中 token 但无系统外推设计；FreeSolo 的 cell 摘要只进路由打分，softmax 读出完全在原始 token KV 上，每 query 每头读预算固定 $Ks=128$，不随历史增长。第二，查询无关的范数剪枝：FS-L 每头只留 $P=64$ 个最高范数 cell 与 $W=8$ 个最近 cell；与 H2O/SnapKV 事后依赖注意力分数不同，该规则训练时就固定，短 CPT 让表示主动适应约束——340M 的 S-NIAH-1 从直接剪枝的 62 恢复到 76-90。第三，头内 RoPE/NoPE 耦合：局部窗口用 RoPE，长程路由与读出完全不旋转，长距离检索永不遇到未见过的 rotary offset。消融显示去掉长程 erf 项仍保留强 S-NIAH-1 召回，外推应归于 NoPE 设计而非距离先验；但局部 NoPE 对照下游准确率从 51.9 掉到 38.3，说明局部 RoPE 对短上下文质量必不可少。

方法步骤详情

流程分六步。①共享投影：隐状态经 RMSNorm 与 $W_q,W_k,W_v$ 得 $q,k,v$，两分支共用。②局部分支：在 $w=512$ 窗口内做 RoPE 注意力。③cell 化：前缀按 $s=32$ 切 cell，算均值 key 与最大 key 范数；查询按 $B=128$ 分块，块内共享可用集。④可用集收缩：FS 保留全前缀；FS-L 每头只留 $P=64$ 个最高范数 cell 加 $W=8$ 个最近 cell。⑤打分选择：按「内容匹配加范数加权 erf 位置偏置」打分，每头选 top-$K=4$ 个 cell 的 128 个 token；FS-B 改为在两棵平衡森林（全前缀/高范数）上做每层 16 节点、对数层深的 beam search，再用 token key 重排。⑥读出合成：选中 token 上做 NoPE softmax 再 RMSNorm，与局部分支相加。训练两阶段：先训基线 FS，再开剪枝做 CPT（FS-L 用 0.73B/1.57B token，FS-B 用 0.42B token 加检索加权损失）；硬 top-P/top-K 索引无梯度。

技术新颖性

新颖性有四层。架构上：此前没有工作在每个注意力头内同时实现局部 RoPE 窗口与稀疏全局 NoPE 检索——已有 RoPE/NoPE 混合是层级混合；FreeSolo 把两者耦合进同一个头，训练与推理位置分布一致。记忆上：与 GSA/Raven 的槽内存、Mamba 系循环状态不同，FreeSolo 保留原始 token KV，摘要只是索引；与 QUEST 的 key 边界页选择相比无需额外学习；与 DRT/RAMba 不同，FreeSolo 硬选择后只做一次 softmax。训练上：两阶段配方（先 FS 后开剪枝 CPT）是新发现——从头带剪枝训练的 800M 模型检索很差，先建表示再施加约束只需约 5% 预训练 token 量；FS-B 的检索加权损失对重复 span 与切换位置加权，训练用全前缀选择、树只在推理使用。算法上：FS-B 双森林保证「任何 cell 不被永久丢弃」与「高范数证据被强调」，节点界由均值 key、半径与最大范数给出可靠上界，检索从 $O(N^2/s)$ 降到 $O(N\log^2 N)$（含 radix sort），严格亚二次且档案空间 $O(Nd)$。

**Figure 2**: FreeSolo attention. FS scores the full prefix; FS-L prunes per head, independently of the query. Both select K = 4 cells per query and head. We use P = 64 high-norm cells, W = 8 recent cells, cell size s = 32, query blocks of B = 128 tokens, and SR radius w = 512.

**Figure 3**: FS-B long-range selection, per head. The query searches two causal forests over the prefix: one over every cell and one over cells with $a_c \geq \tau_\ell^h$. Each beam keeps 16 nodes per level down to the leaves.

实验结果

（1）零样本外推：所有模型 4k 训练，8k 起为外推。800M 的 FS-B 在 16k 的 S-NIAH-1/2/3 达 100/88/84，RULER 均值 30.4；FS-L 得 81/40（Raven 26/0），均值 22.3 领先 Raven/GDN/NAtS-L 的 12.8/10.4/2.8。（2）超长单针：340M 的 FS 到 256k（64 倍训练长度）仍保持 98-100，FS-L 得 84（Raven 131k 已跌至 20）；800M 的 FS/FS-L 在 256k 得 96/85。（3）短上下文：800M 下游 8 任务平均 FS 56.5、FS-L 55.2，接近 TF++ 的 56.8，340M 的 FS 51.9 还略高于 TF++ 51.8。（4）效率：FS-L 注意力缓冲从 4k 到 256k 恒定（766/1149 MiB）、解码 89/75 token/s，但训练吞吐 98.4K 低于全注意力 303K。匹配 CPT 对照下两者 16k 的 S-NIAH-2/3 为 84/38 对 80/42：短 CPT 几乎不损检索地换成线性代价。

**Table 1**: RULER accuracy (0–100). All models train at 4k; results at 8k and above test zero-shot extrapolation. Means cover available tasks.

**Table 2**: Zero-shot downstream accuracy (0–100). Rows are ordered by average within each scale; bold and underline mark the best and second-best scores.

**Table 4**: Extended RULER task scores including multi-key / multi-value NIAH at 800M (referenced for FS-L vs linear baselines).

**Table 6**: Efficiency comparison (masked implementation vs bounded-cache streaming implementation).

**Table 12**: FS vs FS-L RULER comparison at 800M (Appendix): the CPT-ed FS-L has a lower RULER mean than base FS at every evaluated length.

查看结构化数据

任务	指标	本文	基线	提升
S-NIAH-1 超长单针检索（340M，256k = 64 倍训练长度）	准确率 0-100（每任务每长度 100 样本）	FS 全程 98-100，FS-L 84	Raven 在 131k 跌至 20；NAtS-L 800M 在 256k 为 0	FS 零退化完成 64 倍外推，FS-L 比同距离 Raven（41@64k）高出 40+ 个点
S-NIAH-2/3 困难单针（800M，16k 外推）	准确率 0-100	FS-L 81/40；FS-B 88/84	Raven 26/0，NAtS-L 3/0，GDN 13/5，NSA 0/0	较最强线性基线 Raven 提升 55/40 个点
RULER 套件均值（800M，16k）	12-13 任务平均分 0-100	FS-B 30.4（超线性组最高），FS 23.1，FS-L 22.3（线性组最高）	Raven-full 23.8，Raven 12.8，GDN 10.4，NAtS-L 2.8，TF++ 0.0	FS-B 领先最强基线 Raven-full 6.6 分；FS-L 领先线性组 Raven 9.5 分
MK-NIAH-1 多针检索（800M，16k）	准确率 0-100	FS-L 23	Raven 3，NAtS-L 3	较两个线性基线均提升 20 个点
MV-NIAH 多值检索（800M，16k）	准确率 0-100	FS-L 17	Raven 9，NAtS-L 0.5	较 Raven 提升 8 个点、较 NAtS-L 提升 16.5 个点
零样本下游 8 任务平均（HellaSwag/PIQA/ARC 等，800M）	平均准确率 0-100	FS 56.5，FS-B 56.2，FS-L 55.2	TF++ 56.8，SWA 56.2，Raven 54.6，MoBA 45.3，NSA 44.5	与全注意力 TF++ 差距仅 0.3 分，且远高于 MoBA/NSA 约 11 分
推理资源（4k 到 256k 上下文，A100-40GB）	注意力缓冲大小 / 解码速度	缓冲恒定 766 MiB（340M）/1149 MiB（800M），解码 89/75 token/s，prefill 2.1-2.2K token/s	全注意力 KV 缓存随长度线性增长；masked 参考实现 prefill 34.6-44.2K token/s 但缓存为 O(Nd)	存储 O(1)，总处理时间随输入长度近似线性增长

局限与改进

作者承认三点：实验止于 800M 规模；FS-L 的固定保留预算可能在 query 到来前丢弃证据，且 800M CPT 后的 RULER 均值在每个评测长度都低于不剪枝的基线 FS（Table 12），「既省内存又不掉点的 CPT 目标」仍开放；FS-B 在 800M 的多针与多值检索上仍落后全注意力的 Raven-full（Table 4）。我的观察：其一，erf 位置先验名不副实——可用 cell 满足 $i-p_c>0$，作者证明 $g_h\ge\mathrm{erf}(2)\approx0.9953$，该项实为范数 bonus，去掉后 S-NIAH-3 反而更好，距离建模并未真正起作用；其二，路由标量 $\sigma_h=4$、$\lambda_h=1$ 完全无梯度、手工固定，敏感性未研究；其三，$P=16$ 时精度随预算非单调，超参鲁棒性存疑；其四，所有模型（含 FS 系）在变量追踪 VT 等需要状态累积的任务上训练长度内也失败；其五，训练比全注意力慢约 3 倍，理论线性未在工程上兑现；其六，FS-B 流式缓存有重复缓冲拷贝（附录 E.2 自述），实测效率可能低于上界。

独立分析的弱点

四个弱点。第一，FS-L 的精度-效率权衡未真正解决：800M 的 RULER 均值（16k 22.3）全面低于 FS（23.1），论文没有给出「什么任务分布下剪枝净收益为正」的判据；固定保留 72 个 cell 对证据分散的语料（如每个键值对等概率被查询的 MV-NIAH）会系统性失败（附录 D.5 自证）。改进方向：把 FS-B 已验证的检索加权 CPT 损失迁移到 FS-L，或让 $P$ 随输入信息熵自适应。第二，FS-B 的树搜索是纯推理期固定算法，训练仍用全前缀打分，训练与推理选择分布不一致且策略不可学习；可在 CPT 中以树检索 top-K 为目标，用直通估计或 Gumbel 松弛端到端训练。第三，范数先验是单 cell 级统计量，表达不了「高范数 cell 簇」等结构信息，节点上界用最大范数偏松，beam 宽 16 可能偏大，更紧的界可缩小搜索量。第四，评测局限：RULER 以人工检索为主，缺长文摘要、代码库理解等真实任务；部分基线在长长度缺席（Raven 800M 64k 后无 S-NIAH-3），线性组参数量 335M-380M 不等，容量差异可能影响组内对比公平性。

未来方向

作者提出三个方向：找到能在 FS-L 内存预算内提升 RULER 均值的 CPT 目标（FS-B 的检索加权损失是已验证起点）；让 FS-B 学习搜索策略而非固定推理算法；扩展到 800M 以上规模。基于成果还可延伸：其一，把「局部 RoPE + 长程 NoPE」头内耦合作为通用组件移植到 Llama、Qwen 等开源模型做稀疏化改造，短 CPT 即可检验外推收益，工程路径清晰；其二，范数保留规则与 KV 量化、分页结合——高范数 cell 本身就是天然的冷热分层依据，可构建流式长上下文服务；其三，erf 先验退化为范数 bonus 的事实提示应显式学习真正距离敏感的打分函数，或验证更小 cell 粒度下距离项是否重新有用；其四，用可微排序/选择方法处理 top-P/top-K 硬索引不可导问题，让路由端到端训练，替代当前「路由标量无梯度」的妥协；其五，FS-B 依赖 radix sort 的树构建可改为 GPU 友好的增量维护，支持真正的流式服务；其六，多针多值任务上与 Raven-full 的差距说明固定 $K=4$ 读预算可能不足，按任务复杂度动态分配读预算值得研究。

复现评估

作者承诺开源代码，细节披露程度高：数据用公开 FineWeb-Edu，32k LLaMA tokenizer，训练上下文固定 4096，bf16 加 AdamW；基线共享语料与训练设置，参数量逐一列出（340M 档 335M-380M，800M 档 694M-1128M）。预训练预算 340M/15B、800M/32B token，CPT 仅 0.42B-1.57B token（约 5%）；按同类工作估算，340M/15B 约需 8 卡 A100 数天，800M/32B 更多；但只复现核心贡献「短 CPT 适应剪枝/树搜索」可从公开 checkpoint 起步，成本低得多。评测协议完整：RULER 每任务每长度 100 样本、贪心解码、官方评分，附录 B-E 给全部任务分数、效率测量、消融与复杂度推导。风险点：流式缓存含重复缓冲拷贝（附录 E.2 自述），效率数字可能因实现而异；from-scratch 与 CPT 结果有跨 seed 波动（附录 D.4），建议多 seed 验证；8 卡训练吞吐 98.4K token/s 低于全注意力的 303K，预训练时间成本不可忽视。

FreeSolo：线性时间稀疏注意力实现零样本长度外推 FreeSolo: Linear-time Sparse Attention with Zero-Shot Length Extrapolation