Key-Value Means：可扩展块循环压缩记忆的 Transformer Key-Value Means
块循环压缩记忆统一Transformer与线性RNN架构
前置知识
Transformer 自注意力
标准 Transformer 注意力 $A = \text{softmax}(QK^\top/\sqrt{d_h})V$，$Q,K,V \in \mathbb{R}^{N \times d_h}$，时间复杂度 $\Theta(N^2 d_h)$。每生成一个 token 都把键值对追加到 KV 缓存。
KVM 的全部动机正源于自注意力的二次方与线性缓存开销，本文提出的可扩展压缩状态正是要在这两个维度上同时降阶。
线性循环神经网络 (LRNN)
用固定大小的矩阵状态 $S_t = f(S_{t-1}, x_t)$ 替代注意力，时间和显存复杂度为 $\Theta(1)$ 每 token，但牺牲长程检索能力。代表工作有 RWKV-7、Gated DeltaNet (GDN)、Kimi Delta Attention (KDA)。
KVM 通过允许状态尺寸在 $\Theta(1)$ 与 $\Theta(\sqrt{N})$ 之间连续调节，明确把自身定位为 LRNN 与全注意力之间的可插值架构。
块滑动窗口注意力 (BSWA)
一种稀疏注意力变体，每层仅保留最近 $L = n_{\text{bswa}} \cdot C$ 个 token 的键值对，超出窗口的旧 token 被丢弃。TransformerFAM 与本工作都依赖 BSWA 作为未压缩的近上下文区域，从而区分近邻和远距信息。
KVM 的近上下文路径就是 BSWA，超出窗口的块在落入状态前必须先被压缩，因此 BSWA 是状态压缩的触发器。
RoPE 旋转位置编码
通过对查询和键施加与位置相关的旋转矩阵注入相对位置信息，常只作用于部分通道。本工作对状态键的 RoPE 子空间（前 $r$ 个通道）置零并对键做 LayerNorm，使状态与绝对位置解耦。
状态可能合并来自不同位置窗口的 token，强行 RoPE 会引入冲突；作者通过部分置零加可学习 $\tau_{\text{state}}$ 温度回避了该问题。
winner-take-all 硬分配
在多个候选中只保留最大响应者（1-of-K 硬分配），常用于聚类与向量量化。KVM 的核心合并规则即为 argmax 余弦相似度的硬分配，比传统 softmax 加权更稀疏。
实验显示把温度推向极限、收敛到 1-of-K 分配能显著提升效果，这是 KVM 区别于传统注意力加权机制的关键设计选择。
研究动机
长上下文建模面临两类架构的尖锐矛盾：标准 Transformer 自注意力随序列长度 $N$ 呈 $\Theta(N^2)$ 的 prefill 时间与 $\Theta(N)$ 的每 token 显存开销，难以支撑超长上下文；线性 RNN（LRNN）类架构虽然每 token 时间和显存均为 $\Theta(1)$，但固定状态在长程检索中迅速衰减。例如在 32K NIAH-S2 任务上，RWKV-7 (120M) 仅得 1.0 分，KVM 自身固定状态版本也仅 2.4 分。Block-Recurrent Transformer (BRT) 与 TransformerFAM 等块循环架构需要额外投影参数或丢失关键信息；Compressive Transformer 用固定压缩比的卷积式压缩，状态扩张效率低；TokenFormer 虽允许动态扩展权重但未直接用于注意力；OVQ 维护固定大小质心字典并依赖逐质心计数跟踪。
本文的目标是本文目标是设计一种长上下文架构，能够在 LRNN 的 $\Theta(1)$ 显存与全注意力的 $\Theta(N)$ 显存之间连续可调，且在标准实现下无需自定义核。具体包含四个子目标：(1) 提出统一 BSWA 与压缩状态的单层 softmax 注意力；(2) 给出子线性状态增长策略而不牺牲早期上下文召回；(3) 设计 JIT 键值重归一化方案以避免向量范数塌缩；(4) 提出可兼容 RoPE 的部分维度置零机制。
与已有工作不同的是，作者把切入点放在'winner-take-all 风格的硬分配 + JIT 重归一化'这一组合上，并强调单次 softmax 联合关注 BSWA 与状态这一统一视角。这区别于 OVQ（独立维护量化字典，需逐质心计数，无法无界扩张）以及 Compressive Transformer（用卷积压缩、状态以 FIFO 方式增长），同时也区别于 BRT（需额外状态投影矩阵）和 TransformerFAM（仅压缩超出的块、不在统一注意力里融合）。KVM 明确把状态尺寸调度作为可调超参，使架构可在 $\Theta(1)$ 与 $\Theta(\sqrt{N})$ 之间滑动。
核心方法
KVM（Key-Value Means）把单层 softmax 注意力拆成三段并联：最前面的若干 sink token、长度 $L = n_{\text{bswa}} \cdot C$ 的 BSWA 窗口、以及一个被周期性压缩与可选扩展的'均值状态'。每处理完一个长度为 $C$ 的查询块，从 BSWA 窗口掉落的 $C$ 个 overflow token 会被送入状态更新阶段：先挑出最不冗余的 $n_{\text{append}}$ 个直接 append 扩容，再用 winner-take-all 余弦相似度把剩余 token 合并到现有状态行（每个 overflow 只能命中 1 个状态行），合并后通过 JIT 归一化恢复键值的原始范数。直观看，KVM 把'传统 transformer 不断 append KV 缓存'改成'BSWA 缓存 + 一个软聚类均值池'，用少量参数让状态在需要时主动扩张。
核心创新是三件事的组合：(1) 单一 softmax 联合 attend BSWA 与压缩状态，由可学习的逐头逆温度 $\tau_{\text{state}}, \tau_{\text{bswa}}$ 调节两边贡献；(2) 压缩规则采用 1-of-K 硬分配（argmax 余弦相似度），相比 LRNN 常用的 IPLR/DPLR 软更新更稀疏、区分度更高；(3) 状态值用首次入槽时的范数 $\rho_i = \|s^V_i\|_2$ 作为固定 readout 半径，避免合并导致范数塌缩。这三件事共同保证状态可压缩、可恢复、可扩展，且全部使用标准矩阵运算实现。
方法步骤详情
方法分四步。步骤一（权重预处理）：键的前 $r$ 个 RoPE 通道置零并过 LayerNorm 得到记忆键 $\bar{k}_t$，用 ELU 门 $g_t = 1 + \text{ELU}(x_t W_g)$ 计算合并门控。步骤二（读出）：对状态键再 LN，对状态值做 JIT 归一化，与 BSWA 拼接后施加 softmax 得 $y_u$。步骤三（扩容）：算 $s_j = \max_i \bar{k}_j \hat{s}^{K\top}_i$ 衡量冗余度，取最不冗余的 $n_{\text{append}}$ 个 overflow 直接 append 到状态行，记录其值范数为新槽位半径 $\rho_+$。步骤四（合并）：对剩余 overflow 用 argmax 余弦相似度找最相似状态行 $\pi_e(j)$，把 $\check{k}_j = g_j \bar{k}_j$ 和 $\check{v}_j = g_j v_j$ 累加到该行键值上。
技术新颖性
技术新颖性体现在四个方面：(1) 状态在 softmax 内部与 BSWA 联合打分，无需独立层，这一点上与 OVQ 的'层叠式'做法有本质区别；(2) 1-of-K 硬分配 + 门控加权吸收的组合，让 KVM 不需要 LRNN 中常见的自定义矩阵乘核（GDN/KDA 都依赖专用 kernel）；(3) 把每槽位的初始值范数作为永久 readout 半径，巧妙回避了 sink token 范数过小带来的偏差；(4) 通过部分 RoPE 置零 + 跨段共享分头温度，让状态既保持位置无关又无需重新合并 attention 输出。这些选择共同支撑了 subquadratic prefill 与 sublinear state growth。
实验结果
在 120M 与 350M 规模、Prolong 数据集 8k 上下文训练后，KVM 取得多项 SOTA 级长上下文成绩。NIAH-S1 在 16K 时 KVM-sqrt (120M) 拿到 99.8 vs GPTA-2 的 87.8、OVQ/SWA 的 46.2；32K 时仍保持 99.6 vs GPTA-2 的 29.6。NIAH-S2 (32K) 上 KVM-sqrt (350M) 拿到 17.0，显著超过同档压缩架构（OVQ/SWA 仅 4.2、GPTA NoPE/SWA 0.0）。RULER-4K 平均 KVM-sqrt (120M) 29.6、KVM-sqrt (350M) 38.6，与 GPTA-2 (34.0/47.2) 接近。LongBench 上 350M KVM-sqrt 25.0 与 GPTA-2 的 25.1 持平。消融显示去掉 value-length normalization 让 RULER 平均分从 25.2 跌到 13.4、去掉 sink 跌到 19.1、去掉 merge gate 跌到 20.3，证明这三组件对长程检索至关重要。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| NIAH-S1 (16K context) | Accuracy | KVM-sqrt 120M = 99.8 | GPTA-2 120M = 87.8, OVQ/SWA 120M = 46.2 | 相比 GPTA-2 提升 +12.0 分，相比 OVQ/SWA 提升 +53.6 分 |
| NIAH-S1 (32K context) | Accuracy | KVM-sqrt 120M = 99.6 | GPTA-2 120M = 29.6, KVM-256 120M = 98.4 | 相比 GPTA-2 提升 +70.0 分，相比固定状态 KVM-256 提升 +1.2 分 |
| NIAH-S2 (32K context, 350M) | Accuracy | KVM-sqrt 350M = 17.0 | GPTA-2 350M = 20.8, OVQ/SWA 350M = 4.2, GPTA NoPE/SWA 350M = 0.0 | 相比同档压缩架构 OVQ/SWA 提升 +12.8 分，相比固定 KVM-256 (2.4) 提升 +14.6 分 |
| NIAH-S3 (16K context, 350M) | Accuracy | KVM-sqrt 350M = 77.6 | GPTA-2 350M = 37.6, OVQ/SWA 350M = 8.6 | 相比 GPTA-2 提升 +40.0 分，相比 OVQ/SWA 提升 +69.0 分 |
| RULER-4K (350M avg) | Average score | KVM-sqrt 350M = 38.6, KVM/SWA 350M = 40.7 | GPTA-2 350M = 47.2, GPTA NoPE/SWA 350M = 45.1, OVQ/SWA 350M = 29.0 | 纯 KVM 距 GPTA-2 差 8.6 分但显著优于 OVQ/SWA (+9.6)；KVM/SWA 混合架构仅比 GPTA-2 少 6.5 分 |
| LongBench few-shot (350M avg) | Average score | KVM-sqrt 350M = 25.0, KVM-256 350M = 23.7 | GPTA-2 350M = 25.1, RWKV-7 350M = 23.2 | 与 GPTA-2 基本持平（-0.1），优于 RWKV-7 (+1.8) |
| Ablation: remove value-length normalization (120M KVM-256) | RULER-4K avg | baseline = 25.2 | no v-len normalization = 13.4 | v-len 归一化单独贡献 +11.8 分，是所有消融中影响最大的组件 |
局限与改进
作者在论文中明确指出了若干局限：(1) 训练阶段使用的是固定调度表（fixed schedule）确定状态规模与块长度，并未探索训练-测试调度切换或基于损失的自适应阈值；(2) 没有结合 GQA、MLA 等成熟 KV 缓存效率技术，理论上可直接叠加但尚未验证；(3) 缺乏 teacher distillation 实验，作者认为 KVM 的查询/键/值投影与标准 attention 高度对齐、未来可以蒸馏压缩既有 transformer；(4) 在 32K NIAH-S2 这类干扰文本为连续新颖长篇散文的任务上，KVM-256 (350M) 仅 2.4 分，说明固定状态在'全新内容'上仍存在根本性容量不足；(5) 论文未公开更大规模（如 1B+）或更长训练（>7.8B token）的实验结果，目前仅在 120M/350M 上验证。
独立分析的弱点
从独立视角看，本文有几个值得注意的弱点。首先，状态扩展由固定调度决定，缺乏内容感知能力：在长篇真实文档里段落信息密度不均，统一调度会造成稀疏段浪费、密集段不足；改进方向是引入基于状态相似度分布的 adaptive threshold。其次，winner-take-all 的硬分配虽然稀疏高效，但缺少纠错机制（与 EMA/soft attention 相比），改进方向是允许每个 overflow 命中 top-2 状态行并加衰减。第三，sink 数量 $S=1$ 是经验选择，改进方向是让其成为可学习参数。第四，作者在 BSWA 区用 RoPE、在状态区部分置零 RoPE，这种分段处理语义上不统一，跨段注意力可能丢失相对位置信息；改进方向是探索文末提到的'logsumexp 重新合并'机制。第五，所有实验仅在 8K 上下文训练，KVM 在 16K-32K 测试上的强表现更像是外推而非真正训练得到，改进方向是使用更长上下文训练并报告训练/测试一致性。
未来方向
作者明确提出的未来方向有四：(1) 训练-测试调度解耦、调度在微调时自适应、数据相关调度等扩展；(2) 引入 GQA、MLA 等 KV 缓存效率技术以进一步降低显存；(3) 把 KVM 用作 transformer 的蒸馏目标，借助 KVM 仍使用标准 attention 的特点在某一层或某几层替换以缩小既有大模型显存；(4) JIT 归一化在 sink 上的应用以及 'logsumexp 重新合并机制' 让 BSWA 与状态共享 RoPE 的方向。基于本文结果可延伸的方向包括：把 KVM 与 MoE 结合（每层只激活部分状态行）、把状态作为可检索的 episodic memory 用于 agent 场景、以及在多模态长视频/长音频任务中验证 KVM 的 sublinear 优势。
复现评估
复现性较好。代码与训练好的模型均以 Apache 2.0 协议开源（论文中给出'code here'与'models here'链接，对应作者所属机构 Featherless AI 的发布页）。实验使用 Prolong 数据集（Gao et al., 2025）、GPTAlpha-2 主干（HalfRoPE），KVM 块大小 $C=256$、$n_{\text{bswa}}\text{\_chunks}=2$ 全部固定，调度选择有明确公式。规模仅 120M/350M 训练 3B/7.8B tokens，单卡 A100/H100 级别即可复现，无 LRNN 类工作普遍需要的自定义 CUDA kernel。短上下文评测使用 LM Evaluation Harness（Gao et al., 2024），长上下文使用 RULER 与 LongBench 标准子集。整体复现门槛较低，缺点是 8K 训练上下文在今天已属偏小，若想验证 32K+ 上的强外推性能仍需要重新跑更长的训练。
论文图表
