Browse this library

Chinese source article

2 · 动态路由:从内容中选择注意力连接

比较检索、聚类、哈希与学习路由,分析选择成本和召回边界

View original Markdown ↗

静态窗口提前规定每个 query 能看见哪些位置,执行规则,却无法知道远处哪段内容会在当前问题下重新重要。动态路由把邻接集合写成输入的函数 N(i;x)N(i;x):先用较便宜的信号缩小候选集,再对候选 key 执行精确注意力。选择器读取多少状态以后才能得到 top-k, 决定了「选 top-k」究竟省了多少工作。若路由阶段已经扫描全部 key 并计算完整点积,后面的稀疏 softmax 只节省物化与聚合,没有消除主要配对。

动态路由可沿四个量拆解:候选生成的计算和访存、关键 token 的召回率、离散选择怎样训练、索引怎样映射到块级 kernel。内容相关路由可能使用块摘要、低维投影、聚类中心、近似最近邻或可学习门控。它们对局部连续性、远距离检索和批内负载的假设不同,不能只用统一的稀疏率比较。

1. 路由器到底在计算什么

设输入序列的 key 被切成 mm 个块,每块包含 bb 个 token,路由器先为第 rr 个块构造摘要 srs_r。当前 query qiq_i 不再直接与全部 nn 个 key 计算高维点积,而是先计算 qiq_i 与 m=n/bm=n/b 个摘要的相关分数,从中选出 kk 个块。第二阶段才读取这些块内的原始 key 和 value。若摘要维度为 drd_r,粗略计算量从稠密注意力的 O(ndh)O(nd_h) 变成 O(mdr+kbdh)O(md_r+kbd_h);dhd_h 是 attention head 维度。这个式子同时暴露两个容易被忽略的条件:块数 mm 仍随序列长度增长,而且被选中的块通常要整块读取。

摘要怎样产生,会改变路由器能看到的信息。均值池化保留块的整体方向,却可能稀释只出现一次的关键实体;最大池化偏向局部强响应;额外训练的代表向量能适配任务,但会增加参数和训练目标。聚类方法把相近 key 放进同一桶,哈希方法通过离散码缩小候选,低维投影则用较便宜的内积近似原始注意力排序。几条路线都在估计同一件事:如果执行完整注意力,哪些位置会获得较大权重。估计错误发生在精确计算之前,后续 softmax 无法补回没有进入候选集的 key。

有些动态方法把路由放在 attention 层外,用前一层状态或独立索引预测候选;另一些方法复用当前层已经产生的局部统计。前者可以提前准备索引和数据搬运,后者与当前 query 更一致,却更难隐藏路由时延。训练时若索引随参数每步变化,还要考虑重建成本和旧索引带来的偏差。推理时中心或路由参数可以固定, 候选索引仍随请求和新 token 更新,批内每个请求选择的块数和位置不同,又会形成不规则 gather。算法图里的一次 top-k,落到执行层通常包含打分、选择、去重、排序、地址生成和数据搬运。

2. 从静态拓扑到内容相关连接

基础章中的滑动窗口把位置距离作为唯一选择信号。动态路由没有必要取消局部窗口。更常见的结构是保留一组确定的局部边,再增加少量内容相关的远程边。局部边保证相邻 token 的语法与短程状态稳定传播,动态边负责召回远处的实体、定义或约束。两部分预算应该分别报告,否则同样写成每个 query 访问 KK 个 token,可能一个方法将大部分预算留给局部窗口,另一个方法几乎全部用于远程检索。

块级选择与 token 级选择也不能直接比较。块级路由更容易生成规则 kernel,地址连续,代价是块内许多 token 只因邻近关键位置而被一并计算。token 级 top-k 更接近注意力权重的细粒度排序,但随机地址读取、去重与负载均衡更困难。两者之间还有分层选择:先选粗块,再在块内选 token。分层选择能控制索引规模,却会叠加两次召回误差。评测时需要分别记录粗选召回和细选召回,避免最终质量下降时只看到一个总路由准确率。

MInference 离线为每个 head 固定模式类型, 在线从当前输入估计该模式的索引,MoBA 将长序列划块并按 query 对块进行选择,NSA 则把压缩、选择与局部窗口组合进原生可训练的结构。它们都带有动态成分,但动态对象不同:有的方法在线选择列与位移, 有的方法选择块,有的方法同时保留多条注意力分支。比较时应回到 N(i)N(i) 的产生过程,写清每条分支提供哪些边、各自占多少预算,以及不同分支的输出怎样合并。方法名不能代替邻接关系。

2.1. 训练信号与失败方式

top-k 是离散操作。只有入选候选参与精确注意力时,未入选块通常收不到来自主任务的直接梯度。训练可以使用稠密 attention 产生教师排序,也可以为路由器增加对比损失、召回损失或负载均衡项;还可以在训练早期扩大候选集,随后逐步收紧预算。每种做法都改变训练成本。若教师分支需要计算完整注意力,训练仍可能保持平方复杂度,只在部署阶段获得稀疏收益。

路由器还会遇到「热门块」问题。许多 query 同时选择少数块时,模型行为可能没有立即恶化,执行系统却会在持有这些块的设备上形成热点。负载均衡损失可以鼓励候选分散,但分散本身不保证语义相关;惩罚太强会迫使 query 选择次优块。分布式部署还要决定被选块是远程读取、提前复制,还是按路由结果做 all-to-all。三种策略分别增加时延、容量和通信,对应的最优选择随 batch、块大小和拓扑变化。

关键证据被稳定漏掉, 比平均召回率的小幅下降更危险。长文中一条约束可能只出现一次,它在块摘要里权重很小,却决定最终答案。局部窗口可以在证据附近保存表示,但如果经过有限层数仍无法传播到 query,动态边的漏选就会直接表现为推理错误。因此质量评测要加入稀有实体、相似干扰块、多跳证据和不同相对距离;仅用语言模型 loss 会被大量容易预测的局部 token 稀释。

阅读时先固定预算:每个 query 最多访问多少 token 或块,路由器额外读多少字节,训练和推理是否使用同一选择规则。随后检查失败路径:关键块未入候选时能否由局部边或全局层补救,路由抖动是否破坏缓存复用,热点块是否让不同设备负载失衡。后续将以这套口径比较哈希、聚类、两阶段检索和学习式路由,并把算法召回与端到端时延放在同一张表中。

动态路由与下一章的 KV 选择共享索引,却改变不同对象。路由决定当前 query 读取哪些仍然存在的状态,KV 淘汰决定未来还能从哪些状态中选择。一个块本步没有被选中,不妨碍下一步重新进入候选;一个块从 cache 中删除后,任何路由器都无法恢复它。训练章会继续处理离散选择的梯度与稠密到稀疏迁移,内核章则检查候选索引能否在读取完整 KV 之前生成。

3. 导航

动态路由先从粗粒度选择进入. 这条路线把 token 整理成桶、簇或二维块, 用块内多算换取连续访存和规则 kernel. 哈希与聚类解释训练型内容分组, MInference解释已有模型如何离线选择 head 模式并在线构造 Prefill 索引. 后续的细粒度索引与跨层共享将沿同一套候选召回、索引成本和物理 tile 口径展开.

细粒度索引把最终候选压到 token级. DSA 与索引器从 Lightning Indexer 的目标、训练和 Sparse MLA 数据路径进入, 再比较 HISA、MISA 与 LISA 分别缩小 token 搜索域、active heads 或增加线性全局通道. 这条路线尤其需要同时核算低维全扫、top-k 和主 KV 的随机 gather.

跨层共享进一步减少逐层重复的 KV、索引与 Prefill 计算. CSA、HCA 与 CSA2解释 Full、Reindex、Reuse 层怎样共享主 KV、索引器 K 和候选位置; HySparse 与 HySparse2解释 full attention oracle、block/token 选择、KV Reuse 与两段 decoder 的 KV Bridging. 两条路线都必须按组尾误差和真实读取字节验收.

训练原生稀疏把稀疏结构贯穿前向、反向、Prefill 与 Decode. NSA以全局压缩、重要连续块和局部滑窗三分支为核心, 复用压缩分数做 block 选择, 并让 GQA 组共享候选以对齐实际 KV 读取. 这条路线必须用原生训练结果验证, 不能从稠密模型的后处理表现直接外推.

完整的动态 Sparse Attention 数据路径包括:摘要由哪些状态构造,query 与多少摘要计算,候选以 token 还是块为单位,top-k 结果如何变成地址,精确 attention 读取多少 KV,漏选后有没有补救通道。这条路径把论文复杂度、代码索引张量和 profiler 访存对应到同一个方法。

3.1. 选择器的统计目标

把稠密注意力当教师时,最直接的标签是每个 query 的真实 top-kk 集合 TiT_i。选择器输出候选集合 CiC_i,可以计算集合召回

Ri=∣Ti∩Ci∣∣Ti∣.R_i=\frac{|T_i\cap C_i|}{|T_i|}.

这个量容易理解,却把 top-kk 内的所有 key 看成同等重要。更贴近 attention 的统计量是教师概率质量召回:

Mi=∑j∈Cipij,pij=exp⁡(qi⊤kj/dh)∑t≤iexp⁡(qi⊤kt/dh).M_i=\sum_{j\in C_i}p_{ij}, \qquad p_{ij}=\frac{\exp(q_i^\top k_j/\sqrt{d_h})} {\sum_{t\le i}\exp(q_i^\top k_t/\sqrt{d_h})}.

Mi=0.95M_i=0.95 表示候选覆盖了稠密 softmax 的 95% 概率质量。候选内重新归一化后,总变差距离正好等于 1−Mi1-M_i。不过概率质量仍没有观察 value;若漏掉的 5% 全指向一个独特方向,输出误差可能大于漏掉许多冗余 value 的情形。因此还应计算教师输出 oio_i 与稀疏输出 o^i\hat o_i 的范数差,或比较它们经过输出投影后的变化。

训练数据的平均 MiM_i 很高,也可能掩盖长尾失败。可以按 query 熵、最大 logit 间隔、证据距离和任务类型分桶。低熵 query 的概率集中,选对少数 key 就足够;高熵 query 需要更大候选集。边界附近有许多相似分数时,top-kk 集合对微小扰动很敏感,集合召回会剧烈变化,输出却未必同样不稳定。路由评测同时观察集合、概率质量和输出三层,才不会把排序抖动误判成语义崩溃。

4. 预算应该随 query 改变

固定 kk 便于批处理,却不是唯一选择。设路由分数降序为 r(1)≥⋯≥r(m)r_{(1)}\ge\cdots\ge r_{(m)},可以选择最小的 kik_i,使近似概率累计达到阈值 1−ε1-\varepsilon:

ki=min⁡{k:∑t=1kexp⁡r(t)∑t=1mexp⁡r(t)≥1−ε}.k_i=\min\left\{k:\frac{\sum_{t=1}^{k}\exp r_{(t)}} {\sum_{t=1}^{m}\exp r_{(t)}}\ge1-\varepsilon\right\}.

候选分布尖锐时 kik_i 较小,分布平坦时自动扩大。问题在于求完整分母仍要遍历所有摘要,而且不同 query 的 kik_i 会造成不规则 shape。实际方法可以先设置块级上限,再用分数间隔或尾部上界提前停止。此时阈值既控制近似误差,也控制负载方差;只约束平均 kk 会允许少数 query 占用极大预算。

另一种做法是把路由写成带成本约束的优化:

min⁡θ  E[Ltask(Cθ(x))]+λ E[∣Cθ(x)∣]+γ Var(∣Cθ(x)∣).\min_\theta\; \mathbb E[\mathcal L_{\text{task}}(C_\theta(x))] +\lambda\,\mathbb E[|C_\theta(x)|] +\gamma\,\mathrm{Var}(|C_\theta(x)|).

第一项保护任务质量,第二项压低平均边数,第三项限制请求间负载抖动。固定 top-kk 相当于直接把集合大小写成硬约束;阈值路由则需要显式面对方差。若分布式系统还限制每个 KV 分片的接收量,可以对分片负载再加约束,但它会与语义最优候选发生冲突。

4.1. 离散路由怎样得到梯度

硬 top-kk 的集合在排序边界处跳变。常见训练方式有三类。第一类用稠密教师给索引器做监督,将路由学习变成排序或对比学习;主模型梯度不必穿过离散索引。第二类在训练时使用连续门值 gij∈[0,1]g_{ij}\in[0,1],推理时再截断;温度降低会让门值接近离散,但梯度也会集中在边界附近。第三类使用 straight-through 估计:前向执行硬选择,反向把离散算子近似为恒等映射或连续代理。

三种方法各有偏差。教师监督复制的是某个训练阶段的稠密注意力,学生参数变化后标签可能过时;连续门在训练时让大量弱边参与,部署时突然消失会产生间隙;straight-through 的反向梯度并非硬选择的真实导数。训练原生路线通常用额外的局部或压缩分支提供稳定通道,使路由器早期犯错时主任务仍能学习。随着选择器改善,再逐步提高稀疏分支的责任。

路由器的负样本也会改变学到的边界。随机负样本大多很容易,无法教会索引器区分同名实体、相邻代码块或共享大量词汇的段落。困难负样本可从教师排名边界、同一文档的相似块或上一版索引器的误召回中取得。若所有负样本都来自局部邻域,模型又可能把距离当成捷径,所以需要混合不同距离和内容相似度。

5. 跨头与跨层共享的判定条件

同一层不同 query head 共享 KV,并不意味着它们应共享候选。设第 hh 个头的教师分布为 p(h)p^{(h)},共享集合 CC 的最坏质量损失受

max⁡h(1−∑j∈Cpj(h))\max_h\left(1-\sum_{j\in C}p_j^{(h)}\right)

控制。按各头平均分数选集合,可能牺牲一个专门追踪远距实体的低占比头。更稳妥的合并方式是为每个头先取小集合再求并集,或按 GQA 组共享并给每个头保留少量独立槽位。并集会增大实际 KV 读取,所以论文中的「每头 top-kk」不能直接当作物理读取量;共享 KV 页允许复用, 是否只从 HBM 读取一次还取决于缓存命中和执行调度,互不重叠的候选却会扩大并集。

跨层复用还要测量排序漂移。令相邻层候选为 Ci(ℓ)C_i^{(\ell)} 和 Ci(ℓ+1)C_i^{(\ell+1)},Jaccard 重合率

Ji(ℓ)=∣Ci(ℓ)∩Ci(ℓ+1)∣∣Ci(ℓ)∪Ci(ℓ+1)∣J_i^{(\ell)}=\frac{|C_i^{(\ell)}\cap C_i^{(\ell+1)}|} {|C_i^{(\ell)}\cup C_i^{(\ell+1)}|}

只能反映集合变化;还需观察复用旧集合时的概率质量下降。高重合率支持共享,高质量下降则说明那一小部分变化恰好很重要。CSA、HCA、CSA2、HySparse2、LISA、MISA-2 与 ResMiSA 都可以落在这个坐标中:共享什么状态、每隔多少层重新索引、组内哪些头保留独立残差候选,以及组尾怎样吸收累积误差。

5.1. 可恢复性决定错误会积累多久

一次漏选的后果由下一次决策决定。逐层重索引时,位置 jj 在第 ℓ\ell 层漏选,仍可能在第 ℓ+1\ell+1 层重新进入;若一个层组复用同一集合,同一 token 的直接漏选可能持续到下一次重新索引, 信息仍可能经局部或压缩分支传入;若 KV 已被淘汰,原状态对以后所有 query 都不可见。可把恢复时间写成随机变量

Tij=min⁡{t>0:j∈Ci(ℓ+t)},T_{ij}=\min\{t>0:j\in C_i^{(\ell+t)}\},

并分别统计中位数、尾部分位数与永不恢复比例。平均召回相同的两个选择器,如果一个错误独立、另一个会连续漏掉相同实体,长任务表现会很不一样。

局部窗口和压缩分支提供替代路径,也改变恢复定义。原 token 没有重新入选,但它的信息可能已经写入邻近状态或压缩槽。此时只比较 token ID 会低估功能恢复,应加入干预实验:屏蔽候选路径、只保留替代通路,再观察目标输出变化。动态路由最终要解释的是信息是否送达,而不只是索引是否重复教师的编号。

6. 参考资料