稀疏注意力:每个 query 只挑一部分历史去读

摘要是把块里 8 个 key 拼起来再过线性层,它混进了这 8 个位置的全部信息。 如果 query 在块中间就去看这个摘要,等于偷看了自己后面的 token,训练时就泄漏了答案。 所以代码用 BLOCK_DONE 做掩码:块的最后一个位置 ≤ t,这块的摘要才对 t 可见。 query 所在的那块没有摘要可看,它的内容交给挑选支路(当前块必选)和滑窗支路。

全部被掩掉时 softmax 的分母是 0,会算出 NaN。代码加了一个永远可见的"空槽" sink_kv(可学习参数,初值为 0), 拼在所有摘要的最前面。于是 t = 0…6 的 query 至少能看见这一格,读数公式里压缩支路那一项也因此多了 +1。

NSA 论文(§3.3.1 式 (7))的压缩函数 φ 是一个带块内位置编码的 MLP,块长 l=32、步长 d=16, 相邻块有一半重叠,论文的解释是减轻信息被切碎。本关用一个 Linear 加可学习的 blk_pos 代替 MLP, 步长等于块长(不重叠),这样挑选块和压缩块能共用同一套切分,打分可以直接复用(下一步)。 另外,摘要 key 没有加 RoPE,块内先后顺序由 blk_pos 和拼接的顺序表达。

分数 = softmax(q·摘要key / √head_size),和 NSA 论文 §3.3.2 式 (8) 一致。 同一组 K/V 的两个 query 头把分数相加后一起挑(论文式 (10)),这样一组只需要加载一份被选中的 K/V。 当前块的分数被直接设成 1e4,保证入选。

整个挑选过程在 torch.no_grad() 里,挑哪几块这件事本身不产生梯度。 摘要 key 通过压缩支路的输出拿到梯度,越训越能反映块的内容,挑选复用它的分数,也就跟着变准。 挑中之后,块里 token 的注意力计算是正常可导的。

脚本只保存了最后一层、第 0 组 K/V 头的挑选结果(chosen[0, t])。 每层有 2 组 K/V 头,第 1 组挑了什么没有存;前三层挑了什么也没有存。 答案要从值 token 那里取,取的动作可以发生在更早的层,之后再由残差流一路带到最后一层, 所以最后一层这一组不一定还需要回头看钥匙。 门控数据也和这一点对得上:找钥匙模型第 2、3 层挑选门的均值是 0.609 与 0.565,最后一层是 0.458(STEP 4 可切换查看)。 这张图能说明的是:挑选确实会越过 64 个噪声位置跳回前面的钥匙区,但不能说明检索具体发生在哪一层。

当前块还没写满,压缩支路里没有它的摘要,而紧挨着 query 的几个 token 对预测下一个字最重要。 NSA 论文的做法类似:16 个挑选块里固定包含开头 1 块和本地 2 块(§4.1)。

NSA 论文 §3.2 写的就是 sigmoid:g ∈ [0,1],由输入特征经 MLP 和 sigmoid 得到。 softmax 会强迫三个门加起来等于 1,一条支路变重,另外两条必须变轻;sigmoid 允许三条同时开大或同时关小, 等于顺带学了一个整体缩放。图里语言建模 NSA 每层三个门之和从 0.65 到 1.03,找钥匙 NSA 在 1.08 到 1.26 之间,用的就是这个自由度。

两个数取自不同的前向:gate_mean_last_layer 只喂了 1 条验证样本(语言建模 NSA:cmp 0.145、sel 0.232、win 0.652), 图里的 gate_mean_per_layer 喂的是整批(语言建模 32 条,找钥匙 128 条),最后一层是 cmp 0.14、sel 0.255、win 0.637。 样本不同,平均值略有差别,两者给出的排序一致。

NSA 论文的理由是防止一条支路的梯度通过共享投影去影响另一条:滑窗支路学局部模式最快, 共用投影时它会主导 K/V 的学习,压缩和挑选支路就学不起来。代价是参数变多: 本关每层每条支路一个 Linear(128 → 2×2×32),NSA 那组总参数 0.947M,全注意力 0.743M。STEP 5 比 loss 时要算上这一点。

不能直接这么说。三个原因: ① 参数不一样多:NSA 0.947M,全注意力 0.743M,多 27.5%,多出来的主要是三套 K/V 投影和压缩器; ② 单种子:每组只跑了一次; ③ 评测本身有抖动:全注意力最后三次评测是 1.4946 → 1.4998 → 1.5038,相邻两次差到 0.005,NSA 是 1.4820 → 1.4970 → 1.4825,差到 0.015。

按这个抖动量级看:只滑窗 1.5052 与全注意力 1.5038 差 0.0014,是噪声;压缩+挑选(不带滑窗)1.5321,高出 0.028,更可能是真的, 说明去掉滑窗之后局部信息不够。NSA 低 0.021,方向和论文一致(论文 Figure 4:27B 模型上 NSA 的预训练 loss 低于全注意力),但这组数据不足以给出幅度。

tiny shakespeare 是字符级文本,预测下一个字符主要靠最近几十个字符,32 的窗口基本够用。 找钥匙任务是专门反过来设计的:128 个位置里,钥匙和值放在 1–32,中间 33–96 是噪声,提问在 97–127。 滑窗在提问区往回最远看到位置 66,碰不到任何一把钥匙。值有 32 种,乱猜的准确率是 1/32 = 3.125%,只滑窗那组 3.11% 就是乱猜。 全注意力 100%,NSA 99.85%:挑选支路把前面的块找了回来。

教学版是在稠密的 256×256 注意力矩阵上加掩码来"模拟"稀疏:F.scaled_dot_product_attention 照样把整张矩阵算完, 被掩掉的位置只是不参与 softmax。三条支路就是三次完整的注意力,再加压缩和打分,所以 MPS 上 NSA 那组训练 2756 秒,全注意力 801 秒。 算出来的数值和真稀疏实现一致,速度不一致。

真正变快要靠只取被选中块的 kernel。NSA 论文用 Triton 实现,在 8×A100 上 64K 长度训练前向快 9.0 倍、反向快 6.0 倍; 解码每步的访存量从 65,536 个 token 降到 5,632 个,按访存量估计的加速是 11.6 倍(表 4)。

稀疏注意力省"读":K/V 全部保存,每个 query 只读其中一部分,所以精确检索能力保留(本关找钥匙 99.85%),但 KV cache 照样随上下文增长。 线性注意力省"存":历史压进固定大小的状态,KV cache 不增长,代价是很难从状态里精确取回很久以前的某个 token。

两者可以叠加。Qwen3.8-Flash-Next 是 12 × (3 × Gated DeltaNet + 1 × QSA):四层里三层用线性注意力省存储,剩下一层用稀疏注意力省读取。 DeepSeek-V4 在稀疏之外还做了 KV 压缩(CSA 每 4 个 token 压成一条,HCA 每 128 个压成一条), 论文报告 V4-Pro 的 KV cache 是 V3.2 的 10%。