Kimi K3 · KDA 深度剖析 · 第二篇

Linear vs Sparse
压缩即智能

从 Full Attention 的写入 / 存储 / 读取三段论出发, 沿统一压缩粒度标尺拆解 Linear 与 Sparse 之争: 压缩下界、无法消除的 O(T)、Grassmann 流形上的 hidden-dim 压缩, 以及 Harness 时间轴上的 Compaction。

K3 TECHNICAL ESSAY·约 30 MIN READ·KDA / SPARSE ATTN / MEMORY

TL;DR

前面一篇文章《谈谈Kimi K3 的KDA(1): KDA 如何与 Gated MLA 以及 AttnRes 协同》详细分析了 K3 的混合注意力机制如何解决 《详细分析一下Attention, Sparse or Linear?》提出的 9 个 Linear Attention 的缺陷.

但是还有一个问题, KDA本身在处理过程中状态矩阵是固定的, 若某份信息在KDA层被写入状态矩阵时已经因为碰撞或者Delta覆盖而丢失, 后续的过程中是无法恢复的. 这是第一篇文章为什么要讨论KDA的代数结构是一个仿射收缩半群, 简单的6个字压缩了大量的信息:

但是第一篇文章还缺少一个详细的和Sparse Attention方案的对比. 其实这里学术界和工业界还是有很多争议的, 我一直的观点是对于未收敛的问题或许都是一个非常好的研究机会, 既然是争议自然各自都有优劣势, 那么就逐渐剥开来详细探讨, 因此本文从几个视角来详细分析:

  1. 将Attn的数据路径拆分为 KV写入段, KV存储段以及最后的 Query 读取段 , 再以数据路径不同的压缩机制和寻址机制来展开分析.
  2. 从体系结构的视角, Linear Attn 的 State 更像是体系结构中的Register File, 而传统的Full Attn / Sparse Attn 的 KV 类似于传统的 Cache 结构, 那么实质上模型架构和芯片/系统架构 co-design 的视角来看, 本质上是在设计 Memory subsytem, 而这里 Memory 一语双关, 既代表 LLM 中的记忆, 又代表真实芯片中的内存存储结构.
  3. 可能需要考虑的一个问题是 Agent Harness 和 RSI 场景下的 Memory , 所以整体上还需要覆盖 Test-Time Training 和 Harness 场景下的一些时空可组合性的问题.

这些是需求侧的分析, 然后在供给侧芯片的Memory wall的存在, 实质上算法/模型结构的竞争本质也变成了同等智力情况下的内存访问的差异. 这其实才是 Sparse 和 Linear 争论的焦点... 当然从体系结构的视角来看待 LLM memory的算法或许后面会单独写一篇.

本文目录如下:

1. Linear vs Sparse 的争议是什么?
1.1 从标准Full Attention 谈起
1.2 Linear的实质是结合律换掉的查询感知
1.3 Sparse Attention压缩是学习出来的池化
1.4 统一压缩粒度标尺

2. 压缩的难题
2.1 KV压缩的下界在说什么?
2.2 无法消除的 O(T)?
2.3 再看 Linear 和 Sparse 的争议
2.4 Linear Attn的视角
2.4.1 缝缝补补的过程
2.4.2 KDA不可避免的几个墙
2.5 Sparse Attn的视角

3. 压缩即智能
3.1 支持 O(T) 的Linear Attn ?
3.1.1 block size 的约束
3.1.2 基于 Prefix Cache 快照
3.1.3 block overlap 和 state 快慢更新
3.1.4 Query-Aware block selection
3.2 基于 Grassmann Manifold 的 hidden-dim 压缩 ?
3.2.1 从 JL 引理到 Grassmann Manifold
3.2.2 从全局 latent 到 block-local 子空间
3.2.3 在局部坐标内完成 Sparse Attention
3.2.4 到底能够省多少 ?
3.2.5 只让适合压缩的 block 进入低维空间
3.3 Harness? 时间轴上的压缩?
3.3.1 时间轴其实是三个时钟
3.3.2 Compaction! 时间轴的压缩
3.3.3 代数上到底谁可逆, 谁可交换?

4. 小结


1. Linear vs Sparse 的争议是什么?

1.1 从标准Full Attention 谈起

从计算机体系结构的视角来看, 标准的 Full Attention实质上包含一个完整的写入 / 存储 / 读取的过程. 一个 token 从进来到被用掉, 经过三段:

基于 Softmax 的 Full Attention 基线在这三段上都不压缩, 于是存储是 , 读取是 , 但是注意到一个不对称的地方:

从这个视角来看, 本质上 Linear Attn 和 Sparse Attn 的争议, 实质上是在阐述关于 Memory 的处理: "把压缩插在哪一段" 能够获得更高的效率并且在输出上尽量不打折扣. 唐杰老师最近有一篇《Memory for Large Language Models》的survey 写的挺好的, 后面我们将展开从这个视角进行分析.

实质上的选择是压缩的插入位置以及压缩沿哪条轴进行, 一份 KV 有两条轴可压, token 轴与通道轴, 而这两条轴的压缩后果截然不同, 因此构成了 4 种方案:

方案 A, 压缩插在写入段, 沿 token 轴. token 进来的瞬间就被折叠进一个定长对象, 写完即压完, 条历史被压成 1 个状态. 存储降到 , 这是最彻底的省法. 但是 token 轴被压到只剩一个对象, 于是读取端基于 Query 连"挑哪一个"这个动作都没有对象可挑. Linear Attention 与 KDA 走这条路.

方案 A', 压缩插在写入段, 但只沿通道轴. 通过对 KV 下投影的方式在写入端压缩, 以 MLA 为代表, 每个 token 只执行一次, 与任何 query 无关, 压缩比 倍. 但它一个 token 都没有合并: 条目数仍然是 , 每个 token 保留自己独立的一条. 读取时把 576 维上投影回完整的 K 与 V, 于是候选个数与基线完全相同, query 的选择权一分未失, 唯一的损失是每个候选变糊了.

两者的区别, 实质上是一个压缩是否剥夺读取端的选择权的取舍.压 token 轴动的是"有几个东西可挑", 压通道轴动的是"每个东西看得多清".

方案 B, 压缩插在读取段. 每来一个 query, 就针对这个 query 把历史压一次. query 感知拉满, 因为压缩函数可以自由地依赖 query. 但存储一分不省 . 你必须留着完整历史才有东西可压; 而且压缩成本从"每 token 一次"变成"每 query 一次", 计算反而更贵. 因此这条路工程上没人走.

方案 C, 压缩插在存储段, 两条轴一起压, 选择插在读取段. 先把 KV 的存储形态改一改 , token 轴上把 条池化成 条, 通道轴上每条槽是 512 维 , 这一步仍然与 query 无关; 但压完之后留下的东西仍然是一堆可寻址的 KV 条目, 只是条目少了. 于是读取段可以照旧按 query 挑一个子集来看. DSA 走这条路.

把四个方案与基线并排, 演进的逻辑就清楚了:

方案 压缩插在 沿哪条轴 存储 每查询候选数 压缩函数可依赖 query 读取可按 query 重组 代表
基线 不压 都不压 不适用 能 Full Softmax
A 写入段 token 轴, 1 个定长对象 不能, 因果律禁止 不能 Linear, KDA
A′ 写入段 通道轴, 不能, 因果律禁止 能 Gated MLA
B 读取段 任意 能 能 不省资源, 无人实现
C 存储段加读取段 两条轴, 且降到 512 不能 能 DSA, CSA

实质上就是压缩是有损的, 在压缩损失代价上需要和算力/存储做一个合理的平衡. 当 K3 以方案 A 把 token 轴压到底换 状态时, 不得不在中间插入一系列 方案 A'保住 token 轴换回可寻址性.

方案 C 是 A 与 B 之间的那个中间点, 而它之所以存在, 靠的是与方案 A′ 同一个不显眼的性质: 压缩之后的对象仍然可寻址. A′ 靠一个 token 都不合并来保住这一点, C 靠合并之后仍留下 条独立条目来保住这一点; 只有 A 把这一点丢了.

1.2 Linear的实质是结合律换掉的查询感知

KDA 的写入段是

而 与 的全部参数 , , , 都只由当前 token 的输入 生成. 展开成通解:

这个式子里没有 . 一个 token 在历史中占多大份额, 由它自己写入时的 与之后所有 token 的 连乘决定, 而这些量在任何 query 到来之前就已经全部确定. 写入段做完之后, 谁重要谁不重要这件事就已经被写死在 里了.

那么 query 到底还能做什么? 把读出展开成隐式核的形式:

实质上 query 可以调方向为 与 的夹角; 但 query 调不了包络, 即由 与算子范数的相容性,

右端一个 都没有. 于是有一个干脆的判断: query 在 KDA 里只能在一个已经画好的天花板底下挑方向, 天花板本身它碰不到. 而这个天花板沿距离单调递减 , 越老的 token 上限越低, 所以一条被写入时判断为"不重要"因而 很小的信息, 或者一条被后续 token 反复衰减的信息, 无论后来来了多么切题的 query, 都拿不回来.

这就把问题变成了一个下注问题. 写入段要在当下决定两件事: 决定这条信息记多久, 决定它写多重. 而这两件事的最优值取决于将来会来什么 query. 对两段前缀完全相同, 但后续 query 不同的序列, KDA 必然给出完全相同的 .

一个更直接的对比如下:

左式之所以能做到 每步, 唯一的原因是求和被移到了查询之前. 求和一旦完成, 历史就成了一个与查询无关的对象, 之后来多少查询都只花一次读取. 它是效率的取舍:你把"读之前先看一眼要读什么"这个权利卖掉了, 换回常数状态. 使得查询变成了在状态矩阵中选方向, 而不是像Sparse Attention那样可以选择对象.

的另一面是指固定的状态矩阵, 以当前 K3 为例, 实质上每一层能够容纳的上限为 元素 , 设Context 长度为 , 那么每 token 在 1M context 下保留的元素大概只有 1.5 个, 但是为了应对 Linear Attention 自身的一些缺陷 (例如前文列举的9个缺陷) 必须引入Hybrid Attention.

1.3 Sparse Attention压缩是学习出来的池化

以DeepSeek-V4为例, 压缩这一步值得细看, 因为它和 KDA 的"衰减覆盖"是完全不同的丢信息方式. HCA 的形式最简单: 令 , , 每 个条目按

压成一个, 其中 是可学习的窗口内位置偏置. 通过这个压缩操作, HCA将序列长度压缩为倍. 注意 是哈达玛积, 所以权重是逐通道的, 也就是说同一个窗口内不同通道可以选择不同的 token 作为主要来源. 这一点常被忽略, 它使压缩的表达力远高于标量加权平均.

CSA 的压缩率只有 4, 但多了一层重叠机制. 它算两套投影 与两套权重 , 每个压缩槽由 个条目产生:

用于 的 索引与用于 的 索引重叠, 所以序列长度仍压到 但相邻槽共享边界 token. 注意, 重叠不是平滑技巧那么简单, 它有更多深层次的用途. 拿到压缩槽之后, CSA 还可以基于Query再做一次 top-k indexer进行选择, 而 HCA 层由于压缩比 它对全部压缩槽做 dense attention.

1.4 统一压缩粒度标尺

令 为被合并成一个不可再分单元的 token 数, 则 就是查询能分辨的历史片段数. 四类层放在同一把尺子上:

机制 压缩粒度 1M 下可寻址单元数 每单元维数 每 token 每层保留元素 查询能做的事
KDA, 69 层 1 , 1M 下为 1.5 只能选方向
Gated MLA, 24 层 1 1,048,576 576 576 对全部单元连续加权
DSv4 HCA, 31 层 128 8,192 512 4 对全部单元连续加权
DSv4 CSA, 30 层 4, 含 2 倍重叠 262,144 512 128 先选 1024 个再加权

从压缩比来看, 1M 下 KDA 每 token 每层只留 1.5 个元素, 而 DSv4 的 HCA 层留 4 个, KDA 压得更狠, 但两者同量级. 其实这里我有一个疑问, 可能最终的消融实验来看, HCA估计这么高的压缩比下也没什么卵用, 是否可以算出某个压缩比的下界出来? 但是加权到全模型 K3 是 , DSv4 是 , K3 每层每 token 平均保留的元素反而是 DSv4 的 2.30 倍, 反问 Linear Attn 到底省了么?

另一个展现的差异是可寻址单元数: KDA 是 1, DSv4 CSA 是 262,144. 把它折成比特就是"一次读取里查询能做出多少选择决定": CSA 从 262,144 个单元里选 1024 个, 决定量为

而 KDA 的单元数为 1, 组合选择量恰好 bit. 于是"查询感知"有了一个可算的刻度:

另外对于KDA 和 DSA 两种压缩, 还有一个分析压缩干扰的视角, 对比如下:

什么情况下两条事实互相干扰 判据的性质 可预测性 有无缓解手段
KDA 两者 key 方向相近, 与位置无关 语义的 不可预测, 取决于学到的 key 分布 无, 写入前无法预知谁会撞车
DSA 两者落在同一压缩窗内, 与内容无关 位置的 可预测, 由 与位置直接算出 有, CSA 的两套投影重叠正为此

最后, 总结一下两者的差异:

KDA 的压缩 DSA 的压缩
压缩后的对象 一个不可再分的定长状态 一个仍可寻址的槽序列
写入是否覆盖 是, 衰减与 沿 擦除 否, 槽只追加不互相覆盖
可逆性 不可逆, 幺半群非群 缓存层面无损, 只有槽内池化有损
失效的形态 内容没了, 永久 这次没看到, 可重试
补救方向 无法补救, 只能由 MLA 从另一份副本取回 改进 indexer 即可, 下一层下一次都能重选

两条路线"把不确定性消解在写入时还是读取时"上分道: K3 在 69 层里下注, 用 24 层 MLA 与 8 个 AttnRes 归档做对冲; DSv4 全部推迟到读取段.

2. 压缩的难题

这一章的目的是从 KV 压缩的下界计算谈起, 然后展开分别谈论一下 Sparse 和 Linear 各自的优势和缺陷.

2.1 KV压缩的下界在说什么?

KDA 虽然可以说把复杂度从 降到 , 这个说法仅在 KDA 层级上, 对整个模型而言: K3 保留了 24 层完整 的 MLA 用于解决纯 Linear Attn的一系列缺陷. 因此, 我们可以得到一个猜想: KV的压缩的下界从模型层面必定以 的形式存在而无法消除, 实际上的区别只是 的常数项如何降低的取舍. K3 靠减少 层的数量, 从 93 层降到 24 层; DSv4 靠降低每层 项的单价, 压缩 4 倍(CSA)或 128 倍(HCA).

那么是否这个下界真的存在? 正如前一章所讲, 实质上要考虑在Query-blind的情况下的数据压缩下界的问题.

2.2 无法消除的 O(T) ?

对于这个猜想, 后来搜了一些论文, 发现在论文《Compression Barriers in Autoregressive Transformers》[1]中已经有一些详细的阐述了, 其中:

据说太多的数学内容被批“不拿读者当人看”, 那么大白话说说大概的证明过程, 大致就是把这个问题规约成一个通信索引问题. 也就是说我们把 Attn 本身当作一个访问内存的问题.

该问题的单向随机通信复杂度为.

类比 Attn 算法, 此处 Alice 持有的输入是 0/1 矩阵 , 共 个比特; Bob 需要查询其第 位. 其中 为 token 序列长度, 为hidden-dim. 然后将注意力用作按地址寻址的读取操作. 我们关注 Alice 这一侧针对 KV 的编码:

1. 以 value 承载待编码信息 :最直接的一步: 将矩阵 的第 行取作第 个 value 向量:

2.以 key 作为相互独立的地址: 构造需要 个互不干扰的地址向量. 在 中这是平凡的: 标准正交基 两两内积为 . 但我们期望将 key 的空间压缩, 例如 key 位于 维空间且 , 无法容纳 个严格正交的向量.

这里就引入了Johnson–Lindenstrauss（JL）随机投影, 它解决的是一个很漂亮的问题：能不能把大量高维向量压缩到低维, 同时近似保留它们两两之间的欧氏距离? 答案是可以.

Alice 输入完 个三元组后, 将 比特内存状态发送给 Bob. Bob 只需再向算法输入一个三元组 , 其中

即取第 个地址向量的方向, 并按因子 缩放 (该缩放技巧出自 Keles et al., 2022). 于是打分向量 满足

这正是论文所称 softmax 的可塑性 (malleability): 增大 可使 与其余 的间距任意增大, 而指数函数进一步将该间距放大, 使 softmax 由平滑的加权平均退化为近似的 argmax, 即在离散支持集上逼近一个 Dirac 的尖峰.

令 为 softmax 权重. 输出向量的第 个分量是

(Bob 输入的那一行满足 , 不贡献分子; 仅使分母增加一项 , 不改变量级.)

以下分两种情况. 记主项 , 干扰项之和的上界 :

情况 1: . 分子中仅剩干扰项, 且各 , 故

情况 2: . 分子至少包含主项本身, 故

协议成立的条件是 Bob 能够区分这两个区间, 即要求 . 而

任何 均可. 论文取 , , 此时

即两种情况下的输出相差 倍. 至此 的构造已经完成, 且 的取值范围已确定. 但 只保证了 Bob 能够答对, 距离 必须很大 还隔着整条归约的收尾. 最关键的是协议的正确率达标, 此时存在两个失败来源: JL 投影不满足内积条件 (概率 ), 以及算法自身出错 (概率 ). 取并集界, Bob 的成功概率至少为

满足索引问题对成功率的要求. 至此我们持有一个完整合法的索引问题协议:

JL定理断言: 任何满足上述条件的协议, 其通信量为 .因此我们可以夹出下界

论文还做了以下几点探讨:

  1. 低维空间下的边界: 当嵌入维度 较小 ()时, 论文证明了空间复杂度的下界为 , 并指出 Zandieh 等人提出的 SubGen 算法在这种特定场景下达到了这个理论下界. 但是通常现在的LLM
  2. 结构性假设的作用: 论文明确指出, 非结构化的稀疏性本身并不足以打破线性的空间壁垒. 也就是说, 仅仅知道注意力矩阵是稀疏的, 但不知道稀疏的位置, 依然无法实现亚线性空间. 这为许多依赖特定稀疏模式(如滑动窗口)的实用算法提供了理论依据.
  3. 滑动窗口注意力的扩展: 论文分析了一种更通用的滑动窗口注意力机制, 其中窗口外的 Value 向量并未被完全忽略. 针对此场景, 作者提出了一个新颖的, 基于蓄水池抽样 (Reservoir Sampling) 的亚线性空间算法, 并证明了其空间复杂度几乎是最优的.
  4. 时间复杂度下界: 论文还研究了 token 生成的时间复杂度, 证明了任何非自适应 (non-adaptive, 即提前确定数据访问模式) 的流式算法, 在最坏情况下处理最后一个 token 的时间复杂度至少是 .

2.3 再看 Linear 和 Sparse 的争议

其实这里就有一个问题了, 通常的论断会出来:

  1. Linear attention 状态固定为 , 与 无关, 必然丢信息, 无法复现 softmax attention 的输出;
  2. 所以它是一个没有任何保证的有损近似;
  3. 所以它没有意义, 反正 躲不掉, 不如老实存全量 KV, 然后根据 JL 引理去压缩 KV 的维度? 就像 MLA/DSA/CSA 那样?

固定的状态是Linear Attn最大的软肋, linear attention 的状态比 低了 倍 (取 , , 约 2000 倍), 它完完全全在下界之外, 这一条我觉得支持 Linear Attn 的同学也没什么好否认的, 单纯的 Linear Attn 前面一篇文章也讨论 9 种缺陷...

当然这样的结论是有些片面的, 因为Linear Attention的路径上已经发展出来Hybrid Attention了, 例如Kimi K3 或许应该把 3层 KDA + 1层Gated MLA作为一个整体来看待?

例如归约用的那个困难问题是索引通信问题, 即 Alice 藏一个比特, Bob 事后指定位置去取. 这本质上就是 needle-in-a-haystack: 从无结构的长上下文里精确检索单点信息. 所以定理不只是说"亚线性会有损", 它精确指出了有损会发生在哪种任务上. 而这正是经验上观察到的现象: Kimi K3 在长上下文精确召回类任务上可能会掉点, 在语言建模, 摘要, 常识推理这类信息高度冗余的任务上只有很少的损失? 因此整个 K3 block( 3x KDA + 1x MLA)是否也是一种不错的Attention结构? 因此出现一种争锋相对的观点: "我习惯的那类近似可以有损, 你那类不行", 而且我也不需要"必须有保证无损" , 模糊一点后面加一个 Full Attn 又不是不行?

另一方面论文也阐述了一下对数据不感知的情况下, 仅仅知道注意力矩阵是稀疏的, 但不知道稀疏的位置, 依然无法实现亚线性空间, 这也封死了一系列没有带 indexer 的 Sparse Attn 算法. 那么对于 Hybrid Attn , 可能支持 Linear Attn 的同学也会说状态矩阵是基于数据驱动的压缩, 并配合交替的 MLA 恢复了长程记忆, 然后还相对于 Sparse Attn 多了很多层的归纳偏置能力, 并且可以提高泛化能力?

因此 Linear Attn 和Sparse Attn 对比和争议已经改为 Hybrid Attn 和 DSA/CSA 这类带有 indexer 的 Sparse Attn 方案的对比.

2.4 Linear Attn的视角

2.4.1 缝缝补补的过程

前面一篇文章《谈谈Kimi K3 的KDA(1): KDA 如何与 Gated MLA 以及 AttnRes 协同》也详细谈到过, Linear Attention的一些问题导致 必须要加一个 conv 来修复一下, 然后 加了, 也得加, 加完还得整个非线性的SiLU, 这不想省的指数运算又被加回来了一些, 即便如此, 状态仍会累积和碰撞.

然后Linear Attention只能一直往状态矩阵里面输入, 没法让它遗忘一些东西. 于是在 Mamba-2 里面引入 , 这是一个输入相关的记忆寿命, 但遗忘过于粗粒度, 写入仍是相加. 于是出现了 DeltaNet, 即 , 可以沿当前 key 定向改写, 但是缺少快速清空全局旧状态的机制. 于是我们又搞出了Gated DeltaNet , 即 , 此时增加了全局遗忘和定向改写结合的能力, 但是每个 head 只有一个 decay. 接着就出现了KDA, , 引入了每个 key channel 独立 decay, 但是固定状态仍不等于完整 token-addressable memory.

然后我们继续考虑 transition . 它有 个特征值为 , 沿 方向的特征值为 如果不约束 , 即使 , 也可能出现 , 导致该方向在递推中被放大或反向过冲. 接着, 继续打补丁 Q/K 要加 L2Norm. 然后output gate也需要加一个Full-rank projection, 同时也需要 MLA 也同时加上, 整体激活参数又多了 8B.

这是在 KDA 算子层面, 更进一步从它的代数结构上来看, 它是一个仿射收缩半群,收缩的性质决定了它对一些长程检索的记忆能力是有局限性的, 因此必须采用Hybrid Attn, 通过 Full Attn 的 MLA 来做补充. 这是整个 Linear Attn 逐渐演进到 KDA 缝缝补补的过程.

2.4.2 KDA不可避免的几个墙

如果我们考虑未来 Linear Attn 能够避免 KV 的 MLA 一类的Full Attn 获得亚线性的收益, 其实等价的来说我们把 token 之间的 attn 连成线构成一个拓扑, 那么需要回答的一个问题是: 两个相距 的 token, 它们之间那条"连线"在 KDA 状态里能不能活到被读出来., KDA 的权重沿时间按 衰减, 所以这条连线的强度随距离指数掉. 设定一个"还算活着"的门槛 (衰减到 1% 以下就算死了), 得到

关键在于 变大只有一个办法, 把 往 1 顶. 而 会同时触发三件坏事:

第一重:需求本身就极端(衰减墙) : 意味着 每涨 10 倍, 就得小 10 倍. 只能撑 458 个 token, 要撑 1M 得 . 得精确到小数点后第 7 位.这是要求把一个数钉在 1 的邻域里一个极窄的位置上.

第二重:那个位置的梯度几乎是零(可训练墙): ,求导得 .要 就得 , sigmoid 的极左端, 趋零, 梯度小了四个数量级,等于告诉优化器"这个方向别走了".

第三重:就算学到了,也算不准(数值墙): 在前一篇文章§2.3.8 也谈到过这个问题, FP32 在 只剩 7 位有效位,也就是 1M 距离上的长程 Attn 拓扑环路在 FP32 下已经不可靠.

三重的关键在于"互相独立、同向变紧"

三者卡在同一个旋钮上,而且都朝同一方向恶化, 没有任何一档 能让三者同时满意. 而另一方面:

但是每步 delta 只能擦一个方向(Rank=1). 所以:选择性清除一个 维概念子空间至少需要 步. 衰减虽可快清但通道内无选择性. "又快又有选择性"的遗忘在单层单步内不可达, 只能靠多个 token 或多层接力.是不是这里也推导出需要连续几个 KDA 层配合, 以及整个模型需要更深的网络才能完成? 实际上深度轴的效率也受到了影响. 这产生一个辩证矛盾:

接下来还有一个效率的问题, 或者说我换个说法提问 为什么FlashKDA的chunksize = 16? 公式里 的除法是 chunkwise 形式的阿喀琉斯之踵: 是 保留因子的连乘积, 其倒数 随 chunk 内距离无界增长. 在负 Softplus 参数化下, 16 步累积衰减可以任意接近 , 有限精度下必然溢出.

Kimi Linear 的对策是把 chunk 细分为 16 token 的次级 tile: 非对角 tile 在 log 空间算相对衰减后可用 Tensor Core, 但对角 tile 仍须逐位置对显式计算, 成为 chunk 内计算的主要瓶颈. Kimi K3 的解法是从源头消除数值问题: 封底之后,

于是倒数重标定因子全程可表示, 对角 tile 与非对角 tile 统一为稠密 Tensor Core GEMM, 慢速路径被彻底删除. 这也解释了 FlashKDA 选择 的三重理由:

  1. 数值范围: 刚好落在 bf16 动态范围内, 无需任何 chunk 内 rescaling 技巧;
  2. 求逆便宜: 的逆用 Neumann 级数直接展开即可, 无需进一步分解;
  3. 指令映射: 全部 数学干净地映射到 SM80 MMA (m16n8k16) 指令, 内核简单且可移植.

好的说法是刚好配合 GPGPU 常见的 TensorCore 16x8x16的维度, 但是思考另一个问题对于 TPU 这些 DSA 架构的加速器通常 TensorCore 为128x128 / 256x256 应该怎么办呢? 大致的问题如下图所示:

看上去只要把 CHUNK 改成 128 就能解决一切. 但 K3 的数值约束会反噬: 单级重标定要求 (bf16/fp32 指数上限均为 ), 所以不能靠提高 买 MXU 友好度 , 那是改模型, 不是改 kernel. 与 CHUNK = 16 是联合标定的一对.

更深层的障碍在于 channel-wise 衰减: 危险因子出现在内积内部

是逐通道向量, 无法从点积里提出来. 若像 GDN / Mamba-2 那样用标量衰减, 是纯标量, 可以在 log 空间自由按块重锚定, 放大 chunk 几乎无痛. KDA 为表达力选了 channel-wise, 代价就是重锚定窗口被数值范围钉死在 16 token.

当然我们也可以考虑做一个 2 级循环分块的策略, 外层 CHUNK = 128 把不受 影响的 GEMM 按照 128x128 计算, 内层再采用拼接矩阵的方式提高 TensorCore 的利用率. 但总体来看 MFU 也并不会太好.

2.5 Sparse Attn的视角

其实 Sparse Attn 的视角是非常 straightforward 的, 首先在 EOT 的视角下SDPA是最优解, 因此 softmax 再慢也值得保留. 那么剩下几件事情: 从通道轴来看, 承认 KV 的稀疏性, 然后利用 JL 引理在 hidden-dim 上压缩, 这就是 MLA 做的事情. 从token序列轴的维度, 对内容一定层度的感知, 通过一个更低维度更低精度(FP4)的评估器(indexer) 来做选择, 这就是 DSA. 另一方面是序列长了 Softmax 峰值变小的问题, 一方面可以引入 attn-sink 的处理, 另一方面 topk 的选择也增加了长上下文的峰度.

然后进一步进行分块构造 的压缩, 并对于 top-k 可能存在的 block 边界跳跃的问题, 通过 KV 上的 overlap 来处理, 这就是 CSA. 可以看到 DeepSeek 整个演进的过程.

当然 Sparse 直接取某个常数的 topK 可能也是有一些问题的, 如何保证这种硬截断选出的 block 能够近似的拟合到 Full Attn, 简单推导一下:

对于 个 key, 排序后分数 , softmax 概率 , 其中 . Top-k 截断保留前 个并重新归一化: (), . 记尾部质量 .

头部每个 (分母变小了), 尾部每个 . 头部"多出来的"总量 , 尾部"丢掉的"总量也是 . 于是

对 , 似然比 是一个与 无关的常数. 常数的对数取期望还是它自己:

是 Top-k 重归一化独有的结构. 我们再来分析 score 的分布, 头部每项 , 共 项; 尾部每项 , 共 项. 代入 :

也就是说 每大 1, 误差上界约缩小 倍. 例如 , 时, 即保证 . 这个界只需两个数.

那么要保证 , 由 直接得到唯一一条要求:

因此可以做一个简单的分桶算法:

  1. 投桶: 算分数 时顺手算桶号 , 把整数计数 加一. 每个 key 只做一次减法, 一次取整, 一次计数.
  2. 查表: 桶 里的键最多值 , 这张表与数据无关, 编译期就能算好. "桶里有几个"乘上"每个最多值多少"就是该桶的质量上界.
  3. 累加: 从最右边 (分数最低) 的桶往左累加, 一超过预算就停. 停在哪个桶, 边界就是阈值 ; 保留 的键.

还有一个潜在的问题和 HCA 实现相关, 不管Sparse 怎么选择topK, 如何进一步提高单个 block 的压缩比也是一个值得去考虑的问题. HCA 的池化映射是

我们来构造一个最简模型: 门控近似均匀即 ; 窗口内恰有一个 token 携带查询所需的信号 , 尺度为 ; 其余 个 token 相互独立, 零均值, 同尺度. 池化后的槽为

信号幅度为 , 干扰幅度按独立叠加为 . 两者的 相消, 于是信噪比只剩一个极简的形式:
机制 信号 同窗干扰 SNR 分贝
MLA 1 0 --
CSA 4 0.577 dB
HCA 128 0.089 dB

另一方面把注意力分布的损失做 KL 分解:

HCA 的支撑集是全部压缩槽, 所以第一项对它恒为零, 虽然它的读取端是无损的, 全部损失集中在池化那一步. 实际是缺陷的来源: 因为损失全在写入侧, 而写入侧的决定与查询无关, 所以换一个查询什么也换不回来.

3. 压缩即智能

正如前文所讲, Sublinear的压缩可能要数据感知和数据驱动. 本文开头针对KDA的探讨就是一个很好的例子, 仿射收缩半群短短六个字就概括了 KDA 的很多性质, 但似乎现在的模型都还没有智能到这种层度.... 那么对于我们来说, 现阶段应该考虑哪几个方面的压缩呢?

我觉得至少第一点, 我们得承认过分追求 的 Linear Attn 带来的各种缺陷, 接受当前 的现状再来思考以下几个方向:

实质上的一切都是关于下式的内存访问优化

3.1 支持 O(T) 的Linear Attn ?

2.2 节的结论是: 如果要求对任意未来 query 都保持通用正确性, 固定精度存储的最坏情形下界仍然是 bit. 传统 Linear Attention 选择把整段历史压进一个 state, 于是所有 token 都在同一个固定容量矩阵里竞争; Sparse Attention 则保留 个可寻址条目. 两者之间还有一个值得研究的点, CSA/HCA 介于两者之间为 . 不过似乎 HCA 这样的 128 tokens block 的压缩并不是一个很好的方案.

一个很自然的想法是我们能否沿用 Linear Attn 的更新 State 的方法, 它的表达能力和顺序依赖能力其实是强于简单的池化方案的. 因此, 我们是否能构造一个新的 Block KDA 算法, 不缓存每个 token, 而是每 个 token 固化一个 recurrent state, 让长期记忆从一个全局状态增长为 个可选择状态. , 简单来说 Blockwise KDA 是把一次全局不可逆的下注, 拆成 次彼此隔离的局部下注.

另外还有一种形态, 例如《Log-Linear Attention》[2]采用一个随序列长度对数增长的 Fenwick 树层级取代单个固定尺寸状态, 以更细分辨率保留近期信息, 以较粗摘要保存远期历史. 另外还有一些《 Hybrid Associative Memories》[3]的方法构建混合存储, 引入新型 HAM 层, 对 RNN 来说意外的 token 使用逐个token的 KV 进行存储, 以互补方式组合 RNN 与注意力.

3.1.1 block size 的约束

但是此时还需要考虑一些数值精度的问题以及多个 head 和 state 矩阵的内存开销问题, 令 为真正启用 blockwise KDA 层数, 为冻结 state 的每元素字节数, 每个 block 另存 个 维路由 anchor, 每元素字节数 ; 先假设不同 cache 层共享这套 anchor. 不做 overlap 时, cache 约为

是容量规划上界, 包含已封存归档与至多一个热状态分配; 在线路由只扫已封存的 个 block. 于是给定 state cache 预算 , 忽略取整与很小的描述符项, block size 至少要满足

K3 的 69 个 KDA state matrix 以 FP32 保存合计 414 MiB; 若把 ShortConv 等活动态也算进一份完整检查点, 上界约 428.55 MiB. 对 的单条序列, 不计描述符与分配器元数据, 不同 的容量如下:

block 数 69 层 FP32 state
4,096 256 103.50 GiB
8,192 128 51.75 GiB
16,384 64 25.88 GiB
32,768 32 12.94 GiB
65,536 16 6.47 GiB
131,072 8 3.23 GiB

另一个视角是, 根据文章结构的评估, 按照常见分词器估算, 可以采用下面的 Token 容量层级.

内容单位 常见范围 推荐中值 典型组成
一个段落 Token Token 个句子
一个小节 Token Token 个段落
一个章节 Token Token 个小节
一本普通书 Token Token 个章节

不同类型差异较大:

因此一个比较合理的 block size 可能会在 32K~64K 附近.

3.1.2 基于 Prefix Cache 快照

其实有一个比较取巧的办法, 通常 KDA 一类的 Linear Attention 在构建 Prefix Cache 时会选择在一个请求的结尾, 详细的内容可以参考《当 Prefix Cache 遇见 KDA：Mooncake 如何 Day-0 支持 Kimi K3》[4], 那么实际上我们是有一个递归状态按照逐个请求的快照. 那么此时是否可以构建某种轻量级 indexer 的方式, 帮助模型选到某个快照, 或者在 topk 个快照中按照某种和 query 的相似性进行加权? 此时可能需要构造一个远小于 state 的路由描述符 用于indexer的计算或者某种加权权重和 query 进行运算.

3.1.3 block overlap 和 state 快慢更新

block 之间需要一定程度的 overlap 用于避免语义上的块间硬截断. 例如 blocksize = 64K 时, 我们可以按照 32K 进行 overlap. 例如 Seq = 33K 时, 实际上它需要更新前一个 64K block 的 state 以及下一个窗口 . 实际上我们就构成了一个基于 的短窗口快速更新(当前只有1K tokens), 和一个基于 的长窗口慢速更新(当前累积 33K tokens), 相当于有了一个多时间尺度的状态矩阵, 并且等效的增加了 head 的数量, 那么是否对应的 head 数也可以相对减少一点?

或者另一个维度来看, 由于单层实质上是有了两个不同状态 的更新, 是否可以把 Hybrid 3:1 的比例进一步下降到 2:1 这样可以在模型深度轴上赚回一些TPOT和避免多次 FFN/MoE 的开销?

另外对于计算代价而言, KDA 是一个仿射幺半群, 这也是 chunkwise 并行算法能够运行的原因, 因此对于 Overlap 的 State 在仿射结构下运算量并不大.

3.1.4 Query-Aware block selection

在 query 计算阶段, 设 blockwise state 序列为 , 是否存在一些类似于Sparse Attention 的TopK 选块机制 ? 例如通过 query 在一个更小的维度 indexer 计算出一些 block 的相关性分数, 再取TopK ? 正好搜了一下, 发现《Memory Caching: RNNs with Growing Memory》[5]也讨论过类似的问题.

比较恰当的办法是, 为每个 block 额外保存 个小 anchor . 为了让多个启用归档的层共享一次路由, 又不产生前后层循环依赖, 设 为第一个启用归档读取的层, 在进入该层前计算

并把同一组选中 block ID 广播给后续 个 cache 层, 或者每隔几层做一次 block 选择, 这样从芯片的角度来看, 后续的层如果共享前面一层选中的 block ID 也有更好的从 CPU DRAM 或者外部存储对 KVCache 进行 Prefetch 的能力. 共享 TopK 路由方案下:

其中 是强制保留的最近 block 集合. 选出 block 后, 可以把热状态与归档状态的输出做查询相关的混合. 每个 state 本来就是按层按头的, 默认让混合门控也保留这个粒度, 但共享候选集合 :

其中, 可以复用共享路由分数, 再加一个很小的层/头校准项.

3.2 基于 Grassmann Manifold 的 hidden-dim 压缩 ?

前一节考虑的是 token 序列轴: 将历史组织成仍然可寻址的 block, 再由 query 选择访问哪些 block. 另一个可以与它正交叠加的方向, 是在不合并 token的前提下压缩 hidden-dim 通道轴. 先用 表示一个 block 所含 token 数, 前面分析 HCA 时的问题恰恰在于

它在 query 到来之前就把 个 token 混成了一个槽. 一旦目标信号与同窗噪声在池化中相互污染, 读取端面对的只剩一个不可再分的对象. 那么一个自然的问题是: 能否保留这 个 token 的独立地址, 只让每个 token 从 维变窄到 维, 并且让不同语义 block 使用不同的低维空间? 这正是 Grassmann Manifold 视角比较适合处理的部分. 以前介绍过一篇《[大模型中的数学-103] Grassmann流形视角下的Attention算法》

3.2.1 从 JL 引理到 Grassmann Manifold

通道轴压缩有一个现成的经典起点. 注意力真正需要的只是内积 , 而 Johnson-Lindenstrauss 引理恰好就在说内积可以被低维投影保住: 对 中任意 个点, 存在一个映射 (而且随机高斯矩阵以高概率就是), 使得

这个结论有两处让人心动: 只依赖 与精度 , 与原始维度 无关; 而且 与数据无关, 不用看过 key 就能选定, 因此对所有未来 query 一视同仁.

但把它代进注意力的尺度就会发现这条路在这里走不通. 取一个 block 的 个 key, 即使只要求 这种很粗的精度, 也已经是几百的量级, 而我们的 本来就只有 128; 严格地说, 要覆盖任意未来 query 还得对 query 集合做 union bound 或 net 论证, 只会把 推得更大. JL 引理是为 极大, 的场景设计的; 在 head 维度只有一百多的地方, 随机投影要的预算比原始维度还宽, 压缩比直接小于一.

原因不难理解: JL 的保证是 distribution-free 的最坏情形保证, 它必须对任意点集成立, 所以只能把预算均匀洒在所有方向上. 而一段自然语言的 key 通常并不是最坏情形, 它们往往集中在少数几个方向附近. 想拿到真正的压缩, 就得放弃"与数据无关"这个优点, 换成数据自适应: 把随机的 换成这一段 key 自己的主子空间 . 误差随之从 JL 的相对形式 变成实际残差 (本节后面会用到这个界): 低秩假设成立时它可以远小于 JL 的界, 不成立时它也会诚实地变大. 附带的一个后果是: JL 保的是内积, 而 Value 侧需要保的是加权和的重构, 两者的最优子空间并不相同 , 这也是后面 K/V 各学一套基的原因.

一旦投影矩阵改成"由数据选, 且每个 block 各选一个", 表示上就多出一层冗余需要处理, 这才是 Grassmann Manifold 进场的地方. 决定 logit 与输出的其实只有投影算子 , 也就是 张成的那个子空间; 同一个子空间有无穷多组正交基, 彼此相差一个 . 把这层冗余商掉, 剩下的就是全体 维子空间构成的集合

即 Grassmann Manifold, 其中 是 Stiefel 流形. 它给这条路线三样具体的东西:

  1. 纪律, 只有 的函数才是合法的规范不变量, 可以放心用作打分信号或缓存元数据;
  2. 尺子, 两个 block 子空间之间有 principal angle 定义的距离, 可以判断它们能不能共享一组基, 或把相近的 block 聚成一档;
  3. 正确的优化几何, 学习 时梯度应该落在 Grassmann 的切空间上, 而不是把 个数当成自由参数硬更新.

按这个视角回看, 几种通道轴压缩其实是同一件事的不同取法, 区别只在"这个 维子空间是谁选的, 一共要付几份":

做法 子空间来自 付几份基 依赖的假设
JL 随机投影 随机, 与数据无关 1 份, 甚至不必存 无假设, 但预算 在 下不划算
MLA 式全局 latent 训练学出, 全局共享 1 份, 在权重里 整个数据分布近似低秩
本节 Grassmann codec 每个 block 由自己的数据选 份, 随 cache 存 每个语义 block 近似低秩, 假设更弱

三行都是在 上取点, 只是从"随机一个点"走到"学一个点", 再走到"每段各取一个点". 越往下假设越弱, 越可能吃到局部结构, 代价是多存基与多一次构建. 也需要提前说清: Grassmann 本身不产生任何压缩, 它只把"选子空间"这件事的自由度与不变量算清楚; 能不能兑现, 完全取决于 block-local 低秩这个假设在真实上下文上成立到什么程度.

3.2.2 从全局 latent 到 block-local 子空间

这里先把两级 block 区分开: 3.1 节 32K~64K 的 block 是粗粒度 state/routing block; 本节用 表示其中按段落或小节划分的较小 subspace block 长度. 同时, 下文的 表示当前被压缩的单个 attention head 的 K/V 通道维度, 而不是整个模型的 ; 如果压缩发生在 之前, 同一套推导也可以作用在完整 hidden-dim 上.

设第 个 subspace block 的表示为 , 为它选择一组正交基

并保存每个 token 在这组基下的坐标

这里真正有意义的对象不是某一组具体的基 , 而是它张成的子空间. 对任意正交矩阵 , 与 表示同一个子空间, 因此一个 block 的压缩对象应写成

这和 HCA 的差异非常直接: HCA 将 行压成 1 行; Grassmann block 仍然保留 行, 只是每行从 维变成 维. Grassmann 点 只是一份 block-local 字典, 才保存每个 token 的内容, 所以不能只存 然后声称整个 block 已经被表示了.

也可以把它看成 MLA 的 context-adaptive 版本. MLA 让所有 token 使用模型训练出的同一套 latent codec; 这里则允许小说段落、代码块、公式段落分别选择自己的局部坐标系. 赚到的前提不再是“整个数据分布共享一个低维空间”, 而是更弱的“每个局部语义块近似低秩”. 相应地, 每个 block 也要多付一份 的存储与构建成本.

机制 一个 -token block 压缩后的形态 仍可寻址的 token 数 主要丢失的信息
HCA 1 个 维槽 1 token 轴上的块内细节
MLA-like 全局 latent 个 维坐标 全局 codec 舍弃的通道方向
Grassmann block 1 组 基 + 个 维坐标 当前 block 子空间外的正交残差

3.2.3 在局部坐标内完成 Sparse Attention

Grassmann Manifold 本身不会凭空省下任何 FLOPs; 真正的收益来自不把每个 token 上投影回 维, 而是直接在局部 维坐标内完成选中 block 的注意力. 设一个 attention head 的 Key 与 Value 维度同为 . 更稳妥的默认方案是分别学习两个子空间

并对已经固化的 block 缓存

这里 Key 的子空间负责保留 query-key logit, Value 的子空间负责保留加权输出, 两者的优化目标并不相同; 因而严格地说, 每个 block 对应的是 上的一对子空间;

首先可以用 query 在 Key 子空间上的投影能量作为 block 路由信号:

它有两个讨人喜欢的性质:

  1. 规范不变: 换基 时投影矩阵 不变, 分数不依赖基的具体选择, 是良定义在 Grassmann 点 上的量.
  2. 计算便宜: 算它只要一次 的投影, 完全不用碰 的坐标矩阵 .

但便宜是有原因的: 它没读坐标, 也就看不见内容. 把 block 内某个 token 的真实 logit 沿子空间拆开 (记 ):

投影能量只给出第一项的许可上限: 由 Cauchy-Schwarz, . 能量大, 只说明这个 block 的字典方向与 query 对齐, 允许出现大 logit; block 内是否真有 token 在这些方向上放了内容, 由坐标 决定, 而 根本没读它. 一个与 query 同主题的 block 完全可能字典对齐得很好, 所有 token 的坐标却都落在与 近乎正交的方向上: 分数很高, 相关 token 一个也没有. 所以它回答的只是 "query 是否与这个子空间对齐", 单独拿它给 block 定序会产生系统性的假阳性; 它的正确位置是打分器的特征之一, 与坐标统计量, 内容描述符一起进训练.

反过来做排除倒是可以严格化, 但要把两项一起控制住. 上面的分解给出逐 block 的 logit 上界

其中 query 侧的两个因子都由 免费给出; key 侧的两个 max 是 block 固化时顺手算好的标量: 前者是最大坐标范数, 后者是被丢掉的正交残差的最大范数, 两者都只占一个标量的缓存. 当这个上界低于当前候选集的 logit 门槛时, 整个 block 可以被安全跳过. 注意残差项不能省: 低秩假设不成立的那部分信息恰好都在 里, 只用子空间内那一项做排除, 会把藏在残差里的 needle 连同 block 一起剪掉.

这就把 在路由栈里的位置钉死了, 如果结合 blockwise Sparse Attention一起使用, 于是 token 轴与通道轴各管一半, 接口很干净:

设路由最终选出的 block 集合为 , 对每个选中 block 计算

这里分母仍然保留原始 attention 的 , 因为目标是近似 ; 直接换成 会额外把 logit 放大 , 如果训练中确实需要校准, 应该显式引入可学习 temperature. 然后仍然在所有选中 block 的 token 上做一次联合 Softmax:

其中 . 最后一个式子比较关键: 每个 block 先在自己的 维空间内对 Value 做聚合, 得到一个 维向量, 然后每个 block 只上投影一次再相加, 而不是把 个 Value 全部恢复到 维. 对任意 , 变换 , 时局部坐标相应旋转, logit 与最终输出都不变. 因此这个计算依赖的是子空间与局部坐标共同组成的等价类 , 而不是碰巧依赖某一组基; Grassmann 点本身并不包含 token 内容.

上式还顺带省掉了一个麻烦: 每个 block 各自上投影, 不同 block 的 维坐标从不直接相加, 所以不需要在它们之间定义平行移动或联络. 只有想把多个 block 的 维向量搓进同一个坐标系再统一上投影时才需要, 而它引入的误差未必比省下的那几次投影划算.

如果 K/V 在 block 内有明显的非零均值, 可以各自多存一个中心 , 只压缩中心化残差, 即把线性子空间换成仿射子空间. 代价是 logit 要补一项 , block 输出要补一项 , 每个 block 多存 个元素.

3.2.4 到底能够省多少 ?

如果 Key 与 Value 分别使用一套 的基, 一个 block 的原始 KV 存储量与 Grassmann 坐标存储量分别为

整个长度为 的 cache 因而从 变成

压缩比为

例如取 , 仅按元素数计算, 一个 block 从 个 KV 元素降为 个, 理论压缩比约为 ; 如果 K/V 共享一套基则为 .

令 为每个 query 经 top-k 与安全剪枝后实际读取的 block 数, 低维读取的主成本为

对比恢复原始 KV 后的 . 只有当 且 kernel 能够直接消费局部坐标时, 这一项的收益才接近 . 此外, 前面已经指出, 对全部 个 block 逐个打分仍是 的线性扫描, 要靠分层索引或近似检索才能压下去; 另外需要注意, Grassmann 投影本身并没有让 token 轴变成 sublinear.

3.2.5 只让"适合压缩"的 block 进入低维空间

前面的推导隐含了一个很强的假设: 一个 block 的 K/V 确实接近某个 维子空间. 这对重复性较高的自然语言段落可能成立, 对包含多个主题, 稀有实体, 公式或代码的 block 未必成立. 记

可以用联合投影残差来决定一个 block 是否应该降维:

其中 用于平衡 Key 与 Value 的残差尺度.

实践中可以只提供 几个硬件友好的压缩档位, 再加一个 raw 维的 bypass. 选择满足 的最小 ; 每个压缩档位对应一个 , 相同 rank 的 block 放在一起批处理, 一次 decode step 里被 top-k 选中的 block 也天然按档位分桶, 每个桶做一次 batched gather 与 GEMM, 避免混 rank 的碎片化 kernel.

如果到 残差仍然很大, 就直接走 raw bypass, 拆分 block, 或额外保留少量高残差的 exception tokens; 时没有必要再额外存一份 Grassmann basis. 注意档位与选择的职责分离: indexer 打分决定一个 block 会不会被读, rank 档位决定读它要花多少带宽, 前者逐 query 变化, 后者在 block 固化时一次锁死. 丢掉正交残差后, 即便某个 block 后来被 top-k 频繁选中, 也不能临时从 升到 , 除非 cache 事先保存了嵌套 residual 或另一份高 rank backing store. 这才对应"让一些 block 在更低维空间内处理", 而不是强迫所有上下文接受同一个压缩比.

仅看平均重构误差仍然不够. Key 投影引入的单个 logit 误差满足

这说明一条方差很低但会被未来 query 精确询问的事实, 完全可能藏在正交残差中; 一旦残差在 block 固化时被扔掉, 后面的 query 同样无法恢复.

在因果推理中, 当前尚未结束的 block 保留 full-dim 热缓存, 并且不进入 top-k 候选池: 它与最近若干个 block 一起构成强制保留的 local 分支, 无条件进入注意力支持集, 与 Sparse 路线保留 local window 的惯例一致. block 固化时才一次完成入池动作: 构建 indexer 打分描述符, 按残差选 rank 档位并转成低维坐标 (或标记 raw bypass), 缓存排除用的两个 max 标量. 训练时也只能让早期 query 使用其之前已经固化的 block, 否则利用整个 block 的未来 token 构造 basis 或打分描述符, 都会造成隐蔽的因果泄漏. 一个更稳妥的 cache 因而是三层的:

  1. 最近窗口和当前 block 保留 full-dim, 不经打分, always-on;
  2. 低残差的历史 block 存 , 被 top-k 选中才读, 在局部坐标内消费;
  3. 高残差 block 与 exception tokens 保留 full-dim 逃生路径, 同样只在被选中时读, 只是读出来的格式是 raw.

三层只是存储格式与读取条件不同, 不是三套各自归一化的 attention. 一次 decode step 的候选集合是"强制 local 分支 + top-k 选中的 block", 其中 raw token 与低维 block 产生的标量 logit 必须拼在同一个集合里做一次 Softmax, 否则各分支自行归一化再相加会系统性改变原 attention 的概率质量. 如果某个 exception token 以 raw 形式保留, 还必须从对应 block 的低维坐标中移除或 mask, 避免同一 token 被计算两次.

这里还有一个容易被位置编码掩盖的问题: 如果直接在 RoPE Key 上估计子空间, 同一个语义方向在不同位置发生的旋转可能会人为抬高 block 的有效 rank; 同样的旋转也会污染 indexer 描述符, 让内容相同, 位置不同的两个 block 在打分侧显得不同. 因此更合理的消融是分别比较"压缩 content Key, 位置分支单独保留.

把整条管线串起来: token 轴负责选, 通道轴负责压. block 固化时一次决定 rank 档位 (或 raw bypass), 并构建三份元数据: indexer 的打分描述符, 子空间基与局部坐标, 排除用的两个 max 标量. 读取时 query 对全部 sealed block 打分, 经安全剪枝与 top-k 得到 , 支持集为强制 local 分支加 : 低残差 block 在局部坐标内完成 token-level attention, 高残差 block 与 exception tokens 走 raw 路径, 所有候选拼进同一个 Softmax.

3.3 Harness? 时间轴上的压缩?

前两节主要关注在模型内部的一次 forward 内如何在token轴 和hidden-dim通道轴 上处理压缩, 而另外还有一个维度是在时间轴上, 特别是在 Agent runtime 的场景下, 用户消息, Agent输出, 工具调用与结果, 检查点等不断累积后带来的超长上下文, 那么下一轮到底应将哪些历史重新送入模型?

DeepSeek Harness 这篇论文《A Programming Paradigm for Spatiotemporal Composability》[6]虽然是在讲 Agent Harness, 但实质上我认为对于模型本身的推理过程中的 Memory 也同样需要一些局部的时间/空间可组合性的结构.

这篇论文把 effect 与 coeffect 提升为运行时机制:

这恰好对应我们在时间轴上要问的两件事: 压掉的历史能不能撤回, 以及压掉之后谁该被通知重算 .

3.3.1 时间轴其实是三个时钟

把"时间轴"当成一条轴是这里最容易出错的地方. 实际上有三个彼此独立走动的时钟:

  1. token 时钟, 通常由一次 forward 内的位置序 决定, 本文独立出来为token轴进行分析
  2. 会话时钟, 通常由交互的轮次, compaction 代次, 分支, 这是本节需要讨论
  3. 模型时钟, 由于权重, tokenizer, engram 等动态变化引起, 这是最容易被忽略的, 通常为冻结权重推理, 但是未来逐渐演进到 Test-Time Training(TTT) 时,那么我们需要同时考虑这些变化.

例如一种常见的 TTT 做法是构建一个"inner loop" , 通过在先前观测到的 token 上优化一个神经网络(例如一个 MLP)来动态构建一个临时的 KV 映射. 随后的推理步骤被视为查询这份被存储的知识. 例如 Titans:Memory as context, 我们可以构建一个off-path 的参数, query 向量可以通过这些 off-path 参数产生一个新的 block 补充到 KV context 中, 然后再进入 Attention block, 并且也可以在输出和残差部分更新.

实质上 TTT 带来的结果宏观来看就是模型的参数权重存在时变性. 有了时钟这个视角, 一个直接推论是 KV 是 在一组配置下物化的结果:

其中 是与 TTT 权重的有效参数, 是存储与解码约定. 而从论文 《Test-Time Training with KV Binding Is Secretly Linear Attention》[7]可以看到, TTT-KVB 本质上就是一个Linear Attn, 因此可以纳入到一个相对统一但又对立的 Sparse / Linear Attn框架下讨论.

3.3.2 Compaction! 时间轴的压缩

其实对于KV Cache 的处理, 从数据库工程直觉出发, Decoding 输出的 token 以及投机解码验证所需要的 KV 很自然的就会想到逐个 token 的去构建类似于 WAL 的东西. 然后针对实际的存储层次化结构, 很多 Agent Harness 框架在较长的上下文时都会使用compaction 的处理逻辑, 这一点是否也类似于 LSM ?

从工程侧来看, 一个存储系统要同时满足两件事: 写入必须立刻持久化且不能写坏, 读取必须能快速定位某个key. 这两件事天然对立:

WAL 与 LSM 树是对这一冲突的两半回答: WAL 负责"这件事确实发生了", LSM 树负责"发生过的事怎么排列才便于读". 把它映射到 Agent 的会话历史, 实质上我们需要 Attention 算法支持如下能力:

  1. 根据交互, 逐个token的 WAL 对于 Sparse Attn 是显然存在的, 而对于 Linear Attn 需要根据交互轮次做 State 快照.
  2. 在构建 compaction 的时候, 是否能够类似于 LSM 树的方式进行构建. 例如是否构建一个专门的 tombstone token 或者也可以叫 anchor token 用于寻址, 然后能够很快的索引到相应的 block 上. 对于逐个 token 也按照 block 进行 compact.
  3. 整个 context 以block为粒度按需的组合(Composable ctx), 也就是说不同的 block 可以交换.

3.3.3 代数上到底谁可逆, 谁可交换?

一个常见的说法是 Full Attn / Sparse Attn 在时间轴上是追加写因而可逆可交换, 而 Linear Attn 为了换取遗忘能力付出了这两条性质. 方向是对的, 但把可逆, 可交换和能不能压这三件事分开看, 结论要有意思得多. 关键是要区分写入侧的代数与读取侧有没有与长度无关的精确摘要:

机制 写入侧代数结构 逆是否存在 是否可交换 存在长度无关的精确摘要
Full / Sparse KV 自由交换幺半群 (集合并) 有, 删条目或回退水位线 是 没有
无门控线性注意力 交换群 (向量加) 有, 减去该项 是 有,
KDA / GDN 仿射收缩半群 没有 不 有, 但合并必须保序

第一行和第二行恰好互补, 这是个值得关注的点:

  1. Softmax KV 的写入侧是完美的 WAL (追加, 可交换, 删得掉), 但它的读取侧不存在一个维数与 block 长度无关的 query-independent 精确摘要 , 这是有损 compaction 真正困难的根源.
  2. 无门控线性注意力的 既结合又交换, 两个 block 可以精确 merge, 这正是 LSM 式 compaction 需要的代数结构; 代价是它没有遗忘, 于是撞到容量与干扰的墙.

KDA 是唯一两样都丢的: 逆不存在是因为收缩非满射; 不可交换是因为 乘的是整个状态. 因此可能在整个模型Forward-path上如何处理 KDA 在 Agent 场景下还是存在一些复杂度的, 至少达到等效的短轮次间的可逆还需要保存一些快照来实现.

4. 小结

唐老师前段时间那篇 《Memory for Large Language Models》[8] 文章, 从持久性和可表示性的角度来看, 分类如下:

从survey 的文献统计来看, 以 Linear Attn 为代表的 隐式 + 长期 的路线最近几年快速发展:

但是需要注意的是它本质上是通过状态矩阵压缩, 买到了更好的时间跨度的表示能力, 天然的构成了某种状态摘要. 但是付出了逐token级别的寻址能力, 以及在context上的可交换/可逆能力. 另一方面 显式 + 长期 这条路径也在快速发展, 特别是 TTT 和未来 RSI 相关的发展也需要关注.

因此本文第三章探讨的问题就有两方面, 一方面是将 KDA 这类 Linear attn 从 隐式 + 长期 逐渐迁移到 显式 + 长期 . 另一方面是针对 Sparse Attn 如何做更有效的 Session / ctx compaction.

后面可能需要更详细的从体系结构的视角展开, 例如我们把 Linear Attn State 当作某种Register File 而把Sparse Attn 的 KV 当作某种 Dcache, 而Attention 本身作为某种 Control-Unit. 再配合 MoE / Engram 作为某种可寻址的外部存储.

本文篇幅也足够长了, 后面有空再详细展开吧...

参考资料
[1]
Compression Barriers in Autoregressive Transformers: https://arxiv.org/pdf/2502.15955
[2]
Log-Linear Attention: https://arxiv.org/abs/2506.04761
[3]
Hybrid Associative Memories: https://arxiv.org/abs/2603.22325
[4]
当 Prefix Cache 遇见 KDA：Mooncake 如何 Day-0 支持 Kimi K3: https://mp.weixin.qq.com/s/Yxmt-Foq2D7b46sYOk7WAg
[5]
Memory Caching: RNNs with Growing Memory: https://arxiv.org/abs/2602.24281
[6]
A Programming Paradigm for Spatiotemporal Composability: https://github.com/cordiverse/paper
[7]
Test-Time Training with KV Binding Is Secretly Linear Attention: https://arxiv.org/pdf/2602.21204
[8]
Memory for Large Language Models: https://arxiv.org/abs/2607.25380