2.1.1 · 哈希与聚类: 用内容建立稀疏注意力邻域

长序列注意力的难点可以写成一个近邻搜索问题. 对查询 $q_i\in\mathbb{R}^{d_h}$, 稠密注意力读取全部 $k_j\in\mathbb{R}^{d_h}$, 计算 $q_i^\top k_j$, 再由 softmax 决定哪些位置真正获得权重. 如果只有少数点积进入 softmax 后仍然重要, 一条自然路线是先找近邻, 再只对近邻执行精确注意力. 哈希和聚类都在生成这个候选集合: 哈希把向量映射为离散码, 聚类把向量分配给中心, 同码或同中心的 token 才互相计算.

Reformer 使用局部敏感哈希(Locality-Sensitive Hashing, LSH)近似角度相似性, 将相近向量以较高概率送进同一桶. Routing Transformer 则使用在线 k-means, 让 query 和 key 按与聚类中心的距离进入内容簇. 两者都把注意力图从完整的 $n\times n$ 矩阵改成若干小子矩阵, 但索引的性质不同: LSH 的桶由随机旋转和排序得到, k-means 的簇中心从数据中学习. 桶只负责候选召回, 桶内的 QK 点积、causal mask、softmax 和 Value 聚合仍然是精确注意力.

1. 从全量点积到内容桶

1.1. 输入、输出与邻接集合

设 batch 为 $B$, 注意力头数为 $H$, 序列长度为 $n$, 单头维度为 $d_h$. 投影后的张量为 $Q,K,V\in\mathbb{R}^{B\times H\times n\times d_h}$. 为简化记号, 先固定一个样本和一个 head, 写成 $Q,K,V\in\mathbb{R}^{n\times d_h}$. 路由器接收 $Q$ 与 $K$, 输出每个查询的候选索引 $I\in\mathbb{N}^{n\times m}$; $m$ 是每个查询最多读取的 key 数. 第 $i$ 行 $I_i$ 定义邻接集合 $\mathcal N(i)$. 稀疏注意力只在 $j\in\mathcal N(i)$ 上计算:

输出 $O\in\mathbb{R}^{n\times d_v}$ 与稠密 attention 的 shape 相同, 所以后续输出投影与残差路径无需改变. 变化发生在 softmax 的定义域. 一旦关键位置 $j^$j∗ 没进入 $\mathcal N(i)$, 式 (1) 的分母和加权和都看不见 $j^$. 增大桶、使用多轮哈希或同时保留局部窗口, 都是在降低这种漏选概率. 候选集合还必须服从任务掩码. 自回归模型要求 $j\le i$; padding 位置不得参与; 某些实现还禁止 query 与自身比较. 因此完整集合应写成:

其中 $\mathcal B(i)$ 是内容路由给出的桶, $\mathcal C(i)$ 是 causal 可见集合, $\mathcal P(i)$ 是有效 token 集合. 先分桶再加 causal mask 不会自动保持固定候选数: 排序后的一个桶可能大半来自未来位置. kernel 若按固定块分配空间, 需要把被屏蔽槽位保留为 padding, 或额外读取相邻桶补足历史候选.

1.2. 路由成本决定复杂度是否成立

若每个桶含 $m$ 个 token, 桶内 QK 与 AV 的计算约为 $O(nmd_h)$. 这个表达式没有包含建桶. 假如为了决定候选而先计算 $QK^\top$, 路由阶段已经支付 $O(n^2d_h)$, 后续稀疏化只减少 softmax 和 AV. 哈希的目的正是用 $O(nd_hR)$ 左右的随机投影与 $O(n\log n)$ 排序替代全量点积; 聚类则用 $O(nCd_h)$ 的中心打分替代 $n^2$ 配对, $C$ 是中心数. 以等大桶近似, $C=n/m$. 聚类打分和桶内 attention 合计为:

式 (3) 第一项来自每个 token 与 $C=n/m$ 个中心比较, 第二项来自桶内配对. 令两项同阶可得 $m\approx\sqrt n$, 因而总量为 $O(n^{3/2}d_h)$. Routing Transformer 报告的 $O(n^{1.5}d)$ 口径就来自这项平衡. 它低于平方, 却不是线性; 中心数太少会形成大桶, 中心数太多又让路由打分变贵. 稀疏率只描述式 (1) 保留多少边, 无法单独预测端到端速度. 排序、分桶、前缀和、索引重排、gather/scatter 与不均衡 padding 都会占时间. 小序列上, 一次规则的稠密矩阵乘通常比多次不规则 kernel 更快. 长序列上也要确认实现没有先生成稠密 mask, 再调用普通 attention.

这里的哈希投影成本还依赖桶方向数 $b$. 直接计算式 (4) 的稠密随机矩阵乘, $R$ 轮需要 $O(nd_h bR)$, 再加排序和 chunk 内 attention. 写成 $O(nd_hR)$ 需要把 $b$ 当作固定参数, 或明确使用了哪一种更便宜的投影实现. 当长度扩大时桶数也增加, 这项成本必须重新核算.

2. Reformer 的局部敏感哈希

Reformer v2 §2 共享 Q/K 的线性投影, 打分时令 $k_j=q_j/|q_j|$, query 保留原范数. 正比例归一化不改变式 (4) 的桶号, 却改变 logits. 下面二维例子使用原始点积, 只帮助理解分桶后的 softmax, 复现模型时应采用上述归一化.

2.1. 随机旋转怎样产生桶号

Reformer 的出发点是点积注意力与角度近邻之间的关系. 当 query 和 key 共享投影并做长度归一化时, 较大的内积对应较小的角度. 对向量 $x\in\mathbb{R}^{d_h}$, 取随机矩阵 $R\in\mathbb{R}^{d_h\times b/2}$, 计算旋转投影 $u=xR\in\mathbb{R}^{b/2}$. 再把正负方向拼接:

$b$ 是哈希桶方向数. 式 (4) 把球面划成 $b$ 个区域. 两个方向接近的向量更可能在同一随机旋转下取得相同最大坐标, 但结果不是确定的最近邻保证. 随机矩阵不同, 分区边界也不同. 因此 Reformer 使用多轮哈希: 第 $r$ 轮得到桶号 $h_r(x)$, 一个 query 可从任意一轮同桶 token 中获得候选. 原始 Reformer 的 self-attention 令 $Q=K$, 避免 query 与 key 采用不同表示后难以用一个哈希排序对齐. 这项约束会改变标准 Transformer 的参数化. 哈希不是可插拔的通用 ANN 服务: 它利用了共享 QK 和角度相似性. 如果模型仍分别学习 $W_Q$ 与 $W_K$, 应先证明哈希空间保持原 QK 排序, 或专门训练一个路由投影.

形状上, 一轮哈希的输入为 $X\in\mathbb{R}^{B\times H\times n\times d_h}$, 随机旋转可为每个 head 独立, 也可共享. 桶号张量为 $H_r\in\mathbb{N}^{B\times H\times n}$. 实现把桶号与位置组合为排序键, 对 token 做稳定排序, 得到置换 $\pi\in\mathbb{N}^{B\times H\times n}$. 排序后的 Q/K/V 在连续内存里按 chunk 切分, 每个 query chunk 访问自身 chunk 和前一个 chunk. 访问前一块补足排序后跨 chunk 的同桶成员; 合法边仍由桶号和原始位置决定.

2.2. 排序、分块与反置换

一轮 LSH attention 对共享投影按式 (4) 计算桶号, 注意力打分中的 K 再按长度归一化并把 (bucket, position) 编成排序键. 按键排序得到置换 $\pi$, Q、K、V 依照该置换重排并切成长度 $m$ 的 chunk. 每个chunk与自身及前一个chunk计算局部attention, 同时应用原始位置上的causal与padding mask. 计算完成后, 逆置换 $\pi^{-1}$ 把输出恢复到原序列顺序. 桶号相同不等于一个桶恰好有 $m$ 个 token. 某个方向可能吸收大量向量, 另一些方向为空. 排序后固定长度切块是在离散桶之外再施加工程分块. 同一哈希桶被切成多个 chunk 时, 只看当前与相邻 chunk 会漏掉同桶的远端成员; 不同桶在边界处也可能进入相邻 chunk. 因而算法里的「同桶注意力」与 kernel 实际计算的候选集合并不逐项相等.

多轮哈希可把每轮输出分别计算后合并. 设第 $r$ 轮对位置 $i$ 产生归一化常数 $Z_i^{(r)}$ 与输出 $o_i^{(r)}$. 不能简单平均每轮 softmax 输出, 因为同一 key 可能重复出现, 不同轮的分母也不同. 工程实现需要对重复配对做掩码, 并按各轮的 log-sum-exp 权重合并. 多轮提高召回, 代价近似随轮数 $R$ 线性增长, 索引与排序也要重复. LSH 的性能来源包含两层结构: 概率哈希缩小语义候选, 排序后的定长 chunk 把候选变成可批处理的矩形. 只有第一层会留下尺寸不定的 ragged 桶, 很难直接映射到高吞吐矩阵核; 只有第二层又会让分块边界主导召回.

一个二维手算。 取四个未归一化的二维向量, 仅示意一轮哈希与原始点积 attention, 就不写 batch、head 和位置编码:

取两个随机方向组成 $R=\begin{bmatrix}1&0\0&1\end{bmatrix}$, 式 (4) 拼接后的坐标依次对应 $(+e_1,+e_2,-e_1,-e_2)$. $x_1$ 与 $x_2$ 的最大坐标都是第一维, 所以 $h(x_1)=h(x_2)=1$; $x_3$ 进入桶 2; $x_4$ 的最大坐标是 $0.8=-(-0.8)$, 进入桶 3. 排序顺序可写成 $(1,2,3,4)$, 桶大小为 $(2,1,1)$. 令 query 为 $q=x_2$, 桶内只有 $k_1=x_1$ 与 $k_2=x_2$. 省略 $\sqrt{d_h}$ 缩放, 得分为 $0.9$ 与 $0.82$. 若 $v_1=(2,0)$, $v_2=(0,2)$, 则:

稠密 attention 还会计算 $q^\top x_3=0.1$ 与 $q^\top x_4=-0.70$. 四项指数和约为 $6.330$, 被删两项占约 $25.3$ 的概率质量. 输出误差还取决于未给出的 Value. 若把 $x_3$ 改成 $(0.71,0.70)$, 它进入桶 1; 把坐标交换成 $(0.70,0.71)$, 就进入桶 2; 若分区边界轻微变化, 一个与 $q$ 角度接近的向量也可能落入相邻桶. 多轮随机旋转就是为这种边界误差提供第二次相遇机会.

3. Routing Transformer 的在线 k-means

Routing Transformer v3 §4.1 与 Algorithm 1 将最近中心分配用于更新统计, 执行 attention 时则由每个中心选定额的高分 token, 允许一个 token 进入多个中心. 下文先用唯一归属模型推导邻域和手算, 再接到定额执行. 中心更新的分配与 attention 的候选是两个对象.

3.1. 中心、分配和稀疏矩阵

Routing Transformer 为每个 attention head 维护 $C$ 个中心 $\mu_c\in\mathbb{R}^{d_h}$, 组成矩阵 $M\in\mathbb{R}^{C\times d_h}$. 先用每个 token 唯一归属一个中心的简化模型理解内容邻域. 当中心等范数时, 最大内积与最近欧氏距离给出相同分配:

候选集合由相同中心连接:

实际方法还会限制每个簇保留的 query 和 key 数, 并与局部 attention head 组合. 式 (9) 给出内容拓扑: 属于同一语义区域的 token 即使位置相距很远, 也能在一层内直接交互. 与 LSH 相比, 中心由训练数据持续更新, 分区能贴合模型表示的分布; 相应地, 中心漂移、空簇与热点簇都进入训练稳定性问题. 路由打分张量为 $S_Q=QM^\top\in\mathbb{R}^{n\times C}$ 和 $S_K=KM^\top\in\mathbb{R}^{n\times C}$. argmax 后得到两个长度为 $n$ 的簇 ID 向量. 对 ID 排序便可把同簇 token 放到连续区间, 再为每个簇构造 QK 子矩阵. 若每簇最多容纳 $m$ 个 query 和 $m$ 个 key, 物理 attention 张量可整理为 $C\times m\times m$. 输出再根据 query 的逆置换散回 $n\times d_v$.

这种布局让复杂度容易核算. 当 $C\approx\sqrt n$, $m\approx\sqrt n$ 时, 中心打分为 $nC=O(n^{3/2})$, 簇内配对为 $Cm^2=O(n^{3/2})$. 如果数据分布均匀, 两项平衡. 若一个簇实际收到 $u\gg m$ 个 token, 实现必须截断、溢出到备用簇, 或把该簇拆成多个物理块. 截断损害召回, 溢出增加控制流, padding 则浪费算力.

定额执行令 $J_c^Q=\operatorname{Top}_m(S_Q[:,c])$, $J_c^K=\operatorname{Top}_m(S_K[:,c])$, gather 后得到 $C\times m\times d_h$, 计算 $C$ 个 $m\times m$ 子矩阵. token 可能重复归属, 输出需要按原 query 索引 scatter 并处理重叠, 无法只靠一个逆置换恢复.

3.2. 在线更新与训练阶段

中心不宜只靠主任务对 argmax 直接反向传播. argmax 对簇 ID 的导数几乎处处为零, 未被选择的中心也收不到正常 attention 路径的梯度. Routing Transformer 使用在线 k-means 风格更新, 根据分配到中心的向量移动质心. 对中心 $c$, 一个批次内收到的向量集合为 $A_c$, 批均值为:

用批均值 EMA 示意中心的状态更新, 动量为 $\beta$:

式 (11) 把中心视为训练状态. 分布式训练时, $|A_c|$ 与向量和需要跨设备聚合, 否则不同副本会形成不同路由空间. 空簇没有 $\bar x_c$, 需要保留旧中心或从数据重置. 更新过快会让簇边界抖动, 更新过慢又跟不上表示层随训练变化. 中心的学习率、归一化和同步频率都会影响最终拓扑. 局部 attention 与 routing attention 通常分配给不同 head. 局部 head 保留相邻语法与连续状态, routing head 提供远程内容边. 这种组合还有一个训练上的保护作用: 即便远程路由早期不稳定, 信息仍可沿局部窗口传播. 代价是动态路由并没有支配全部 head, 总体复杂度和质量取决于两类 head 的比例.

推理阶段中心停止按请求更新, 否则相同前缀可能因之前服务过的请求不同而产生不同路由. 固定中心后, 每次输入仍要计算 $QM^\top$ 与 $KM^\top$, 所以索引是动态的, 参数是静态的. Prefill 可一次处理整段并排序; 自回归 Decode 每步只有一个新 query, 但新 key 要加入对应簇的数据结构. 若实现每步重排全部历史 token, 算法收益会消失. 更合理的 Decode 路径是为每个簇维护追加式 KV 列表, query 只查对应列表, 同时处理容量和分页.

六 token 的聚类手算。 取 $d_h=2$, 两个单位中心 $\mu_1=(1,0)$, $\mu_2=(0,1)$. 六个 key 为:

按式 (8), 前三个 key 进入簇 1, 后三个进入簇 2. 取查询 $q=(0.7,0.3)$, 有 $q^\top\mu_1=0.7$, $q^\top\mu_2=0.3$, 所以只读取 $k_1,k_2,k_3$. 三个精确得分分别为 $0.7,0.62,0.54$. 为方便手算, 令对应 Value 是标量 $v_1=1,v_2=2,v_3=4$. softmax 分母为:

权重约为 $(0.360,0.333,0.307)$, 输出为:

现在把 $k_6$ 改为 $(0.49,0.51)$. 它仍被分到簇 2, 但与 $q$ 的点积为 $0.496$, 已接近簇 1 中的 $k_3$. 硬分配在边界处产生不连续漏选. 若允许 query 访问 top-2 中心, $k_6$ 会被召回, 但本例只有两个中心, 候选立刻退化为全量. 更大系统中可以取少量多中心路由, 其作用与 LSH 多轮相似: 用额外预算覆盖分区边界. 按式 (10) 更新簇 1, 均值为 $\bar x_1=(0.8,0.2)$. 若 $\beta=0.9$, 新中心为 $(0.98,0.02)$, 再归一化后方向略向簇内数据转动. 这一步显示在线聚类的状态来源: 中心由历史 mini-batch 的分配累计形成, 不是当前样本即时计算的均值.

4. 质量、性能与工程边界

4.1. 召回误差从哪里产生

哈希的主要误差是概率碰撞与分区边界. 相近向量可能落入不同桶, 不相近向量也可能碰撞; 多轮哈希提高前者的相遇概率, 同时扩大候选和重复项. 聚类的主要误差是中心表达能力与硬分配. 一个多义 token 只能进入少量簇, 长尾语义可能被大簇吞没. 两者都可能在平均 loss 变化很小时漏掉单个关键证据.

表示空间也会改变邻域含义. attention 真正排序的是带 RoPE 等位置变换后的 QK 分数, 哈希或聚类若作用在变换前向量上, 内容近邻未必等于最终 attention 近邻. 若在变换后路由, 每个位置的旋转又会让同一内容随距离改变方向. 论文或实现必须说明路由发生在投影与位置编码的哪一侧, 不能只写「按相似度分桶」. 评测候选质量时, 可用稠密 attention 的 top-$t$ 作为参照. 对 query $i$, 稠密重要集合为 $G_i$, 稀疏候选为 $\mathcal N(i)$, 召回率为:

还应计算被候选覆盖的稠密概率质量 $\sum_{j\in\mathcal N(i)}\alpha^{\mathrm{dense}}_{ij}$. top-$t$ 召回把第 1 名和第 $t$ 名视为同等, 概率质量更接近输出误差, 但两者都不能替代任务准确率. 多跳问答、唯一证据检索与代码变量引用对单次漏选尤其敏感.

4.2. kernel、负载和分布式执行

哈希和聚类都需要把逻辑分组转换为连续块. 典型流水线是生成整数 ID, radix sort 或 bucket sort, 计算每组 offset, gather Q/K/V, 执行分块 attention, 再 scatter 输出. 排序会读写索引与特征, 当 $d_h$ 较小或桶内计算很稀时, 数据搬运可能超过点积. 若每层都重新分桶, 这项成本按层重复; 若多层共享桶, 路由更便宜, 却使用了较陈旧的表示空间.

GPU 喜欢尺寸规则的矩形块. 大小不同的桶通常被补齐到上限, 或按多个 kernel 分批. 设真实桶大小为 $u_c$, 物理容量为 $m_c$, attention 有效率可粗写为 $\sum_cu_c^2/\sum_cm_c^2$. 少数热点桶会同时拉低效率与召回: 不截断时拖慢整批, 截断时丢 token. 因而负载均衡不是附属指标, 应报告桶大小分布、最大/平均比、截断率和 padding 比.

张量并行会让问题更复杂. 若每张设备持有一部分 head, 每个 head 的路由可本地完成; 若按序列切分, 同簇 KV 分散在设备间, 候选生成后需要 all-to-all 或远程 gather. 通信量取决于路由结果, 热门簇会造成目的端拥塞. 复制中心几乎没有压力, 复制全部热门 KV 则增加显存. 算法上的内容邻域只有与数据放置共同设计, 才能在多卡上保持稀疏收益.

训练与推理要分别核算。 Reformer 和 Routing Transformer 都改变了模型训练时的注意力结构, 质量结论来自按该结构训练的模型. 它们不能直接保证把一个已训练的稠密 LLM 换成分桶 attention 后仍保持行为. 稠密模型可能把某些 head 用作低幅度但分布广泛的聚合, 硬桶会改变 softmax 分母; 训练期使用的局部 head、共享 QK 或中心状态也不是现成模型必然具备的.

训练时要保存置换、桶 offset、softmax 统计量和路由状态供反向传播. 哈希随机性影响可复现性; 在线中心需要跨数据并行副本同步. 推理时不再保存反向激活, 但要处理变长请求和 KV Cache. Prefill 的排序能对长矩阵摊销, Decode 每步的索引维护与小 kernel 启动更突出. 所以同一方法在训练、Prefill 和 Decode 上可能有三套不同的瓶颈.

实际选择可以落到几个明确条件. 序列足够长、内容近邻稳定、桶较均匀并且有原生分桶 kernel 时, 哈希或聚类能减少关系数. 任务高度依赖唯一远程证据时, 应增加多轮、多中心或固定全局边. 现有稠密模型不能重训时, 更适合采用后续 MInference 一类按既有 attention 模式做近似的方法. 短序列、强局部任务或高度不均衡表示下, 规则窗口与优化良好的稠密 FlashAttention 往往更直接.

5. 索引维护与更严格的对照

5.1. 碰撞概率、容量与不对称路由

LSH 的价值不在哈希表查找本身, 而在碰撞概率随相似度单调变化. 设一轮哈希中, 目标近邻与 query 同桶的概率为 $p_1$, 一个无关点同桶的概率为 $p_2$, 且 $p_1>p_2$. 使用 $R$ 轮相互独立的哈希, 只要任意一轮同桶就进入候选, 目标近邻至少一次被召回的概率为:

若 $p_1=0.6$, 一轮召回概率为 $0.6$, 两轮为 $1-0.4^2=0.84$, 四轮为 $1-0.4^4=0.9744$. 同时, 无关点至少一次碰撞的概率也从 $p_2$ 增为 $1-(1-p_2)^R$. 取 $p_2=0.1$, 四轮后已是 $0.3439$. 多轮既提高召回也扩大候选, 并不会免费接近精确最近邻. 式 (16)还假设各轮独立. 若各轮随机旋转相关, 独立公式就不再直接适用. 向量位于低维子空间本身不等于各轮投影相关. 训练过程中共享随机种子有利于复现, 但固定分区可能让模型表示主动适应桶边界. 每步重采样提高分区多样性, 又会增加梯度噪声. Reformer 的实现选择必须与论文配置一起读, 不能把「LSH」当成唯一确定的算子.

候选预算还受桶大小影响. 设一轮平均桶大小为 $m$, $R$ 轮不去重的候选上限约为 $Rm$. 实际唯一候选数量为并集大小, 与各轮桶重合有关. 如果表示坍缩到少数方向, $m$ 会远大于设计值, 关系数重新接近平方. 所以监控哈希质量时, 除 recall@k 外还要记录每轮桶大小分位数、跨轮候选并集、重复率以及被截断的 token 数.

聚类容量不是普通超参数。 在线 k-means 常被简写成「分配到最近中心」, 但 attention kernel 需要有界容量. 设每个中心的容量为 $m$, 第 $c$ 个中心收到 $u_c$ 个 token. 当 $u_c>m$ 时, 至少 $u_c-m$ 个 token 无法留在该物理桶. 可以保留最接近中心的 token, 把溢出项送到第二近中心, 或把逻辑簇拆成多个物理 chunk. 三种处理分别改变召回、索引成本和同簇连通性.

保留最近 token 会偏向簇中心, 恰好位于语义边界的 token 更容易丢失. 送往第二中心需要保存 top-2 分配, 并可能把拥塞传递给另一簇. 拆 chunk 保留全部成员, 但同簇 query 若只访问一个 chunk, 仍看不见其他 chunk; 若访问全部 chunk, 热门簇成本为 $O(u_c^2)$. 因此「均衡簇」不是单纯为了 GPU 整齐, 它决定式 (9)是否仍代表实际邻接集合. 可以为中心增加容量正则, 让分配比例接近 $1/C$. 设批内中心使用率为 $p_c=u_c/n$, 均匀目标为 $1/C$, 一个简单惩罚为:

惩罚过弱时热点持续存在, 过强时语义天然集中的 token 被迫分散. 最小化式 (17)也不保证每个 mini-batch 都均衡, 更不保证分布式每张设备的通信均衡. 对训练系统而言, 应同时观察语义损失与物理负载, 不能用单一系数代替容量设计.

Q 与 K 不对称时发生什么。 标准 attention 分别学习 $W_Q$ 与 $W_K$, 相似性是双线性关系 $x_i^TW_QW_K^Tx_j$. Reformer 通过共享 QK 让一个向量能同时作为查询和索引项. Routing Transformer 可以分别把 Q 与 K 分配给同一组中心, 但同中心并不保证原始 QK 分数高. 如果 Q 分布和 K 分布尺度不同, 中心可能更贴近其中一侧. 一种修正是分别维护 query 中心与 key 中心, 再学习中心之间的连接. 设 $M_Q,M_K\in\mathbb{R}^{C\times d_h}$, query 进入 $c_q$, key 进入 $c_k$, 只允许中心对 $(c_q,c_k)$ 属于集合 $E_C$. 这样表达力更强, 但中心图本身需要训练或搜索, 候选不再由同 ID 直接得到. 另一种做法是先把 Q/K 投影到共享路由空间 $d_r<d_h$, 用便宜的路由内积选候选, 再回到原空间计算式 (1). 这正是后续许多 indexer 路线采用的结构.

路由空间的监督目标应与最终 attention 对齐. 仅让投影重构 Q/K, 未必保持 top-k 排序. 可以从稠密 attention 提取教师集合 $G_i$, 训练路由分数 $r_{ij}$ 让正例超过负例, 例如使用 margin 损失:

式 (18)直接优化相对排序, 但构造 $G_i$ 需要稠密教师, 训练期仍可能支付平方成本. 采样负例降低训练量, 会漏掉最难干扰项. 哈希与在线聚类绕开显式教师, 代价是只能通过结构假设近似原 QK 邻域.

5.2. Prefill、Decode 与 KV 状态

Prefill 可以全局排序, Decode 需要增量插入。 Prefill 一次拿到整段 Q/K/V. 哈希可批量生成桶号并排序 $n$ 个 token, 聚类可批量计算 $n\times C$ 的中心分数. 排序开销由整段 attention 摊销, 同桶矩阵也足够大. 反向传播时, 同一置换还能用于组织梯度. 对训练和长 prompt Prefill, 这是分桶路线最自然的执行形态. Decode 每次只新增一个 $k_t,v_t$, 历史 KV 已存于 cache. LSH 若每步用新随机旋转或因表示变化重算历史桶, 成本不可接受. 固定哈希函数后, 新 key 可在 $O(d_hR)$ 时间得到桶号并追加到桶列表; query 计算桶号后读取相应列表. 聚类同理, 新 key 与 $C$ 个中心比较后插入. 每个桶因此需要可增长的页表、长度和容量管理, 不能只保存 Prefill 的连续排序结果.

历史 KV 的物理布局通常按时间分页, 分桶索引只保存逻辑位置. query 命中一个桶后仍要从多个页面 gather K/V. 若为了连续访问再复制一份按桶排列的 KV, cache 容量近似翻倍, 更新也变复杂. 若直接重排唯一 KV, causal 时间顺序与其他 head 的访问布局会受影响. 因此 Decode 的理论 $O(m)$ 候选未必换来同等带宽下降.

KV 量化会改变路由排序。 KV Cache 量化通常将历史 K/V 存为低比特值并配套尺度. 如果路由器直接在量化 K 上算哈希或中心分数, 量化误差可能让边界附近 token 换桶. 若为路由保留一份高精度低维摘要, 索引更稳定, 但增加额外 cache. 设路由 key 为 $k_j^r=P k_j\in\mathbb{R}^{d_r}$, 可以高精度保存 $k_j^r$, 主 attention 命中后再反量化完整 $k_j,v_j$. 这种双表示的容量为每 token 的量化 KV 加 $d_r$ 个路由元素. 当 $d_r$ 很小, 开销可控; 若每层每 head 独立保存摘要, 总量仍明显. 摘要还要与 RoPE 位置保持一致. 先量化再旋转、先旋转再量化以及把位置项从路由分数中分离, 会得到不同桶. 内核融合也不同: 路由阶段只需 K 摘要, 精排阶段才读取和反量化完整 KV.

量化误差应通过路由 recall 与最终输出共同评估. 稀疏选择和量化是两次近似, 误差可能叠加. 若量化恰好改变边界分配, 关键 key 会在精确 attention 之前消失, 后续高精度计算无法恢复. 一个稳健做法是扩大路由候选后再用原 Q 与反量化 K 精排, 以更多带宽换回召回.

cache 淘汰与路由漏选不能混为一谈。 某一步没有选择 key $j$, 只表示 $j\notin\mathcal N(i)$; key 仍在 cache, 后续 query 可以再次命中. cache 淘汰则永久删除当前请求中的某段 KV, 除非从更早状态重算. 两者对错误恢复的含义不同. 哈希或聚类路由可以通过多轮、多中心扩大候选, 却无法找回已经淘汰的 KV. 如果系统同时使用动态路由与 KV 淘汰, 淘汰策略不应只看过去被选次数. 哈希边界可能让一个未来重要 key 长时间未命中, 频率统计会进一步把它删除. 应把路由置信度、局部保留和特殊 token 保护纳入驱逐策略, 并分别报告「存在但未选」与「已不存在」的错误. 这两类错误在输出上相似, 修复手段完全不同.

失效模式与评测协议。 边界、热点和路由抖动。 最常见的失效来自分区边界. 两个几乎相同的向量可能因极小扰动进入不同桶, 从而让邻接集合发生离散跳变. 训练时这种抖动增加梯度方差; Decode 时相邻 query 的候选变化过大, 会降低 K/V cache 层次中的复用. 多路哈希与多中心分配能平滑召回, 也让索引和候选增长.

热点是第二类失效. attention sink、高频标点或相似模板 token 可能集中进同一桶. 热点桶既增加计算又挤掉容量内的长尾 token. 单看平均桶大小会隐藏问题, 应报告最大值、P95/P99、Gini 系数和截断比例. 分布式场景还要按目标设备统计接收字节, 因为全局桶均衡不代表网络链路均衡. 第三类失效是路由空间与任务空间错位. 语言模型 loss 主要由局部可预测 token 构成, 路由器可能学会稳定处理常见局部结构, 却在罕见跨段引用上失误. 因此仅比较困惑度不足. 评测集应控制证据频率、距离、干扰项相似度和需要汇总的证据数量.

一个可复现的评测矩阵。 算法层先固定稠密教师与候选预算. 对每层每 head 计算 recall@$m$、稠密概率质量覆盖、候选重复率与桶容量溢出. recall@$m$ 的教师 top-k 数不能随方法改变. 若一个方法访问整块, 既报告逻辑命中的 token 数, 也报告整块实际读取的 token 数. kernel 层记录建桶、排序、gather、attention、scatter 各段时间, 以及 HBM 读写、有效 FLOPS 与 padding. 对照至少包含同 dtype 的优化稠密 kernel和规则窗口. 序列长度应跨过方法的收益门槛, batch 与 head 维度保持一致. 预热、编译和首次内存分配不能混入稳态数字.

任务层同时测试局部语言建模、远程单证据、多证据聚合、相似干扰和长度外推. 训练型方法还要区分训练长度内与更长长度, 并报告中心或哈希配置是否随长度改变. 性能与质量最终画在同一预算轴上: 横轴为实际读取的 KV 字节或物理关系数, 纵轴分别为任务质量和时延, 才能判断某个召回增益是否值得.

选择路线的判断。 哈希适合角度相似性明确、希望避免学习大型路由器并能接受概率召回的结构. 聚类适合表示分布可由有限中心描述、训练过程能够维护中心且愿意处理容量均衡的结构. 两者都更适合从训练阶段纳入架构. 对已有稠密模型, 另行训练路由投影或直接使用观察到的空间模式往往更稳妥.

无论选择哪条路线, 首要问题都是索引器读了什么. 只读低维摘要, 节省可能成立; 扫描全部高维 K, 可能仍受内存带宽限制; 计算完整 QK 后才选 top-k, 则没有减少主配对. 第二个问题是候选如何进入 kernel. 连续定长块能兑现吞吐, 任意索引会付出 gather 和负载代价. 第三个问题才是稀疏率. 内容路由的完整成本是「建索引 + 组织数据 + 精确 attention」, 完整质量是「候选召回 + 桶内计算 + 未命中补救」. 用这两个等式检查实现, 才能判断哈希或聚类是在减少工作, 还是把平方计算换成了另一组同样昂贵的索引操作.

从算法图到可运行实现。 索引张量必须携带哪些信息。 一个可运行的分桶 attention 不能只输出桶号. kernel 至少需要每个 query 块的起始偏移、有效长度、对应 key 块列表以及原位置, 才能恢复 causal 关系. 对 batch 和 head 独立分桶时, 常用压缩行格式: row_ptr 记录每个 query 组在索引数组中的起止, col_idx 记录 key 组编号, perm 与 inverse_perm 负责重排和恢复. 这些整数张量虽然小于 KV, 仍会被每层读取. 变长 batch 需要把样本边界编码进排序键, 否则相同桶号会把不同请求的 token 排在一起. 一个安全排序键可由 (batch_id, head_id, round_id, bucket_id, position) 组成. causal mask 使用的是原始 position, 不是排序后的槽位. 只按槽位比较先后会允许未来 token 泄漏, 训练 loss 可能因此异常变好.

反向传播还要保存或重建置换. 保存所有层的 perm 增加激活内存; 重建要求哈希随机数与中心状态完全一致. activation checkpointing 若在重算期间中心已被式 (11) 更新, 会产生与前向不同的桶. 因而中心更新应位于优化步骤的确定边界, 重算阶段读取同一版本.

重复候选和归一化。 多轮哈希、多中心分配以及局部窗口组合都会产生重复 key. 假设同一 $j$ 在两条候选路径出现, 若直接拼接后 softmax, $e^{s_{ij}}$ 会在分母出现两次, 等价于人为给该 key 增加 $\log2$ 的分数. 重复不是无害的多算, 会改变输出. 去重应发生在 softmax 前, 或由互斥的块划分保证. 显式排序唯一化会增加索引开销. 块级 kernel 可以规定优先级: 先计算局部块, 内容桶只计算局部范围外的元素; 多轮哈希则给每个配对指定最小轮号, 后续轮屏蔽已出现配对. 这种规则需要能在 kernel 内廉价判断, 否则去重 bitmap 本身可能接近 $n^2$.

不同候选分支也不能各自 softmax 后相加. 正确方法是共享行最大值与指数和, 或用 log-sum-exp 合并分支. 若局部分支有 64 个 key, 路由分支有 8 个 key, 分支级平均会让后者获得与前者相同总质量, 与单一候选集合上的 softmax 不同.

最小正确性测试。 小尺寸测试应先绕开性能. 取 $n\le16$, 明确列出每行候选, 用普通张量索引实现参考结果, 再与稀疏 kernel 比较. 测例要覆盖空桶、满桶、溢出、重复候选、序列首尾、padding、causal 边界、多轮哈希以及 GQA 共享. 只与稠密 attention 比较不够, 因为稀疏候选本来就允许输出不同; 应先验证 kernel 与同一候选集合的参考实现一致. 随后固定候选为全量, 稀疏实现应退化到稠密结果. 固定为局部窗口, 应与规则窗口实现一致. 这两个退化测试能发现 softmax 合并、逆置换和 causal mask 的错误. 退化测试通过后再测哈希或聚类生成的真实候选, 把算法近似误差与 kernel 数值误差分开.

性能测试则必须关闭调试排序与 Python 循环, 预热编译, 同步计时. 如果参考实现逐 head 循环而稠密基线批量执行, 得出的速度没有比较意义. 只有相同 batch、head、dtype、序列长度和硬件下的端到端层时间, 才能回答分桶是否值得. 还有一个必要基线: 在相同物理读取预算下比较随机桶、固定窗口和内容桶. 若内容路由只比稠密快, 却不比相同预算的固定窗口准, 在线建索引没有产生可识别收益. 反过来, 内容桶召回明显更高但端到端更慢, 说明算法信号成立而执行路径尚未兑现. 将质量增益与索引增量分别报告, 才能判断下一步应改路由目标还是改 kernel.

生产统计还应按请求隔离. 汇总后的均衡分布可能由一半请求集中在簇 1、另一半集中在簇 2 形成, 单请求依旧严重热点. 延迟尾部取决于单请求最大桶, 训练稳定性则更关心跨批中心使用率. 两个尺度不能互相替代.

6. 误差传播、状态维护与路线选择

6.1. LSH 召回与候选近似

先用随机超平面的一位哈希推导碰撞概率, 再看多 bit 与多轮的预算关系. 这一哈希家族与式 (4) 的有符号坐标 argmax 不同, 下面的角度概率不能直接代入 Reformer 的桶数.

LSH 的召回来自角度碰撞概率。 局部敏感哈希用整数桶号表达一种概率关系：向量越相似，落进同一桶的机会越高。Reformer 使用随机投影构造 angular LSH，再把同桶 token 排序到连续位置。碰撞概率把哈希轮数、桶数和候选预算连了起来。

随机超平面的基本概率.

取单位向量 $x,y$，夹角为 $\theta\in[0,\pi]$。随机采样方向 $r$，用一位哈希

随机超平面将两向量分到同侧的概率为

角度越小，碰撞概率越高；正交向量概率为 $1/2$，方向相反时为 0。这个公式不保证任何特定相似 key 一定碰撞，只给单轮概率。 若把 $b$ 个独立 bit 拼成桶号，两向量同桶概率为

增加 bit 会快速降低随机碰撞，桶更细；真正近邻也更容易被切开。假设夹角 $30^\circ$，单 bit 碰撞概率 $5/6$，$b=8$ 后同桶概率约 $0.233$。仅用一张哈希表会漏掉多数这样的 pair。

多轮哈希把漏召回降下来.

使用 $R$ 轮独立哈希，只要任一轮同桶就成为候选，命中概率为

仍以 $p_b=0.233$ 为例，$R=4$ 时命中约 $65.4$，$R=8$ 时约 $88.1$。召回提高的代价是候选并集、排序和去重成本增加。 轮次之间若共享投影或只有轻微扰动，独立公式会高估收益。训练中固定随机旋转，模型也可能适应桶边界；推理换种子会改变候选。论文配置中的轮数必须和随机生成方式一起解释。 多轮候选会重复。重复 pair 若直接进入 softmax，相当于提升其 logit；正确集合语义要去重。实际物理预算近似由各轮桶大小之和决定，逻辑唯一候选则取并集，两者之间的差值就是重复开销。

桶边界产生不连续路由.

向量轻微变化时，只要跨过哈希超平面，桶号就会跳变。Attention 分数对 Q/K 连续，候选集合却离散。落在边界附近的 pair 可能在相邻训练 step 反复出现和消失，形成路由噪声。

多轮哈希提供冗余：某轮跨界，其他轮仍可能保留。也可以检查相邻桶或采用 multi-probe，但候选预算增加。把路由视为固定 mask 时，梯度只通过被选边的 attention 分数传播，桶号本身通常没有可微梯度。 边界鲁棒性可以用扰动测试：固定输入，给 Q/K 加小幅各向同性噪声，统计候选 Jaccard、目标近邻召回和最终输出变化。平均召回相同的两套哈希，边界稳定性可能不同。

Angular LSH 与点积排序并非完全一致.

Angular LSH按方向相似度路由，attention logit 为 $q^\top k/\sqrt d$，同时受范数影响。LayerNorm 或投影后的范数分布较窄时，两者接近；范数变化大时，一个方向略差但范数很大的 key 可能拥有更高点积，却被哈希漏掉。

可以将向量归一化后路由、保留原向量计算 attention。路由目标变成 cosine 近邻，最终目标仍是点积 top-k。二者偏差要通过 oracle 对照测量：全量点积 top-k 中有多少进入哈希候选，而非只测 cosine 近邻召回。

Reformer 的排序分块是一种候选近似。 同桶内两两 attention 仍可能形成大块。Reformer 将哈希桶排序后切成固定 chunk，并让 query 关注本 chunk 和相邻 chunk，以控制最大工作量。桶语义与执行 chunk 不是同一层概念。

从桶到 chunk 会截断大桶.

某个桶含 $m$ 个 token，完整桶内 attention 需要 $m^2$ 配对。固定 chunk size 为 $c$，只读自身和前一个 chunk时，每个 query 候选约为 $2c$，总成本线性于 $m c$。相距多个 chunk 的同桶 token不会直接比较。 哈希碰撞只说明排到同一桶，不能保证进入同一执行邻域。评测路由召回时应分两层：目标 key 是否同桶，以及排序分块后是否真正进入 query 候选。只报桶碰撞会高估最终召回。 桶小于 chunk 时，排序数组可能用 padding 或拼接下一个桶。不同桶之间必须用 bucket mask 隔离，否则执行便利会引入错误边。大桶与小桶的 padding比例不同，负载分布直接影响物理工作。 8.2. 排序键必须保持因果信息.

排序会把原序列位置打乱。Causal attention 的合法性仍由原位置 $j\le i$ 判断，而非排序后索引。若只对排序矩阵施加下三角 mask，早出现但排序靠后的 key 会被错删，晚出现但排序靠前的未来 key会泄漏。 正确候选同时满足同桶/邻 chunk与原始因果条件。排序数组需要携带原 position，softmax 前逐元素比较。逆置换只负责把输出放回原 query 顺序，不能修复已经进入计算的未来边。 小例子可将原位置 $[0,1,2,3]$ 排成 $[2,0,3,1]$。排序下三角对 query 原位置 0 会允许读取排序在前的原位置 2，立即暴露错误。这个测试比随机 loss 更可靠。

哈希 Q 与 K 的一致性.

Self-attention 中 query 和 key 来自不同线性投影。若分别哈希，相关 q/k 可能因投影空间不同而不碰撞。Reformer 采用共享 QK 表示，使同一投影既作 query 又作 key，路由和打分处于共同空间。

共享 QK 降低模型自由度。标准 attention 可以让读取条件与被读取索引使用不同表示；强制共享将两者耦合。性能收益与架构变化同时存在，比较时不能把全部质量差异归因于 LSH。 若保留独立 Q/K，可训练额外路由投影 $r_q,r_k$ 对齐到共同空间。路由召回有监督或蒸馏目标，复杂度增加，也进入学习型索引范畴。

可逆层与 LSH 是两项独立设计.

Reformer 还采用 reversible residual layers 来减少激活内存。可逆结构允许反向从输出重建输入，LSH 减少 attention 配对。两者同时降低训练内存，却作用不同。 Attention matrix recomputation、可逆层和稀疏候选的成本应分别核算。重算哈希要求随机旋转、排序和桶配置与前向一致；随机状态不同会产生另一张图，梯度不再对应前向函数。

复现实验时，质量问题先固定哈希种子验证候选，内存问题再检查可逆与 checkpoint。把三项绑成一个开关，很难定位收益和错误。

K-means 路由是在线向量量化。 Routing Transformer 为 Q/K 分配聚类中心，只在同簇内计算 attention。它把近邻搜索改成中心分配，再用中心索引召回候选。近似误差包含量化误差、簇容量截断和中心漂移。

最近中心分配的几何边界.

给定中心 $\mu_1,\ldots,\mu_C$，向量 $x$ 分配到

两个中心之间的边界是超平面，空间被划成 Voronoi cells。真正相似的 q/k 若落在边界两侧，会被分到不同簇；距离中心很远的离群点，即使同簇，彼此也未必相似。 可以让每个向量分配多个最近中心，提高边界召回。Top-$m$ 分配把候选和中心打分成本放大约 $m$ 倍，还产生重复 key。另一方法检查相邻中心，但需要中心邻接表。 量化残差 $|x-\mu_{a(x)}|$ 是局部误差指标。残差大的 token 更可能错分，可给它们更多中心或 fallback；固定每个 token相同路由预算则更规则。

中心更新与非平稳表示.

训练期间 Q/K 投影不断变化，向量分布也移动。中心若更新太慢，会落后于表示；更新太快，又随 batch 噪声抖动。在线 EMA 可写为

这里沿用式 (11) 的批均值 EMA, $0\le\eta\le1$; $n_c$ 是本批分配的向量数. 空簇没有样本更新，可能长期失活；热点簇接收大量异质 token，容量溢出。 中心更新不一定通过反向梯度。若属于外部状态，checkpoint 必须保存它们，分布式 rank 需要同步统计。只恢复模型权重、不恢复中心，会让同一输入路由改变。

簇均衡与语义纯度冲突.

理想 K-means 最小化平方量化误差，不保证簇大小均衡。自然数据存在高频模式，热点中心可能吸收大量 token。固定 kernel希望每簇接近容量 $m=n/C$，否则最大桶决定尾延迟。 加入 capacity 后，超过上限的 token 必须丢弃、改投次近中心或进入 overflow。丢弃降低召回，次近分配提高路由成本，overflow 可能形成大桶。容量因子与 MoE 类似，但这里溢出会删除 attention 候选，直接影响信息路径。

强制完全均衡会把相似 token 拆开，增加量化误差。应同时报告簇大小分位数、中心残差和 oracle top-k 召回。只优化均衡度，可能得到执行漂亮却语义无效的分桶。

Query 与 key 的共同聚类目标.

若 Q/K 使用共享中心，目标是让高点积 pair 分到至少一个共同簇。普通 K-means分别最小化向量到中心距离，并未直接优化 pair recall。Q 与 K 分布不对称时，共同中心可能偏向样本更多或范数更大的那侧。 一种诊断是分别计算 query 残差和 key 残差，再按 attention top-k pair 统计共簇率。若单侧量化误差低、pair recall 仍差，中心目标与 attention 关系不匹配。 学习路由可以用 dense teacher top-k 作为监督，直接奖励候选覆盖；代价是训练期需要 dense 或近似 oracle。哈希/聚类的无监督简单性与学习索引的任务适应性形成明确交换。

6.2. 从候选召回到输出误差

从候选召回到 attention 输出误差。 候选 recall 是中间指标。漏掉一个 key 对输出的影响取决于其 dense attention 概率和 value。数量 recall 对所有 top-k 等权，无法描述概率质量。

被删概率质量.

Dense attention 权重为 $p_j$，候选集合为 $C$，被删总质量

Sparse softmax 在 $C$ 内重新归一化，权重为 $p_j/(1-\epsilon)$。若 value 范数有界 $|v_j|\le V$，可得到输出差异与 $\epsilon V$ 同阶的上界。$\epsilon$ 比候选数量 recall 更接近函数误差。 漏掉十个极小权重 key 可能几乎无影响，漏掉唯一高权重 key 会严重偏移。路由评测应报告 captured mass：$\sum_{j\in C}p_j=1-\epsilon$，并按 head、层和距离分桶。

Softmax 重归一化会放大剩余边.

即使被删 value 总和接近零，删除它们仍改变分母。剩余权重统一放大 $1/(1-\epsilon)$。Sink 或低 value norm token 可能主要贡献分母，数量/语义判断都会低估它们。 候选分支各自 softmax 后相加，会给每个分支固定总质量，与统一集合重归一化不同。哈希多轮、局部窗口与聚类候选合并时，应共享行最大值与指数和，或精确实现等价 log-sum-exp。 10.3. Oracle 路由分离索引与稀疏容量.

固定候选数 $k$，使用全量点积 top-k 作为 oracle mask，测量理想选择下的质量；再换成哈希或聚类候选。点积 top-k 也下降时, 还要检查是否丢失聚合所需的低分边; 点积最高的集合未必最有利于任务. 它保持质量、实际路由下降时, 候选差异值得进一步分析. 同一 $k$ 下加入固定窗口、随机候选作基线。内容路由若不优于随机，桶未捕获语义；不优于窗口，任务依赖主要局部或路由开销没有价值。两种候选的差距是诊断线索, 不构成任务质量的严格上限.

多层误差不会简单相加.

某层漏边改变 hidden state，后续 Q/K、桶号和候选都会变化。早期小误差可能跨越边界导致路由离散跳变，形成级联；残差也可能抑制影响。单层 captured mass 无法直接预测最终 loss。 可采用 teacher forcing 式诊断：只替换某一层为稀疏，其余层保持 dense，扫描层号；再逐步增加稀疏层。早层敏感说明表示与后续路由放大误差，晚层敏感可能直接影响输出。