跨模型复用 KV 缓存:ridge 回归换掉 7 秒 prefill

一个 14B 模型刚读完 32,000 token 的代码库,回答了几个简单问题。现在轮到难题了,得换 32B 上场。

问题是:32B 对刚才那 32,000 token 一无所知。它必须自己从头再读一遍。

在 Qwen3 32B 上,这一遍要 6,975 毫秒——将近 7 秒,读的还是完全相同的文本。用户看着光标闪。而 14B 那边,刚才辛苦算出来的东西正静静躺在显存里。

先说结论

同一家族内的两个模型,它们对同一段文本算出的 KV 缓存之间,关系基本是线性的。可以用一个闭式最小二乘拟合(ridge regression)直接把小模型的缓存换算成大模型的缓存,不做反向传播,不训练。
换算比重算快 2.7 到 25 倍。Qwen3 14B → 32B 在 32K token 上:换算 277.6 ms,重算 6975.3 ms,25 倍。
六对模型里有四对能保住目标模型 73%~98% 的准确率(五个基准平均),另外两对掉到 42%~44%。KV 形状对得上并不保证换算成功。
最反直觉的一条:决定成败的不是误差多大,而是误差落在哪。拟合质量指标 $R^2$ 跟最终准确率几乎不相关( $r=-0.20$ ),而「注意力输出余弦」相关到 $r=+0.57$ 。
这是 NVIDIA 的论文,2026 年 8 月 4 日上 arXiv,只测了家族内、稠密全注意力模型。跨家族(Qwen3 → Llama)明确列为未解决。

一段文本,两份笔记

先说清 KV 缓存是什么,以及为什么它不能直接拿来用。

Transformer 读一段文本时,每一层、每个注意力头都会为每个 token 算出两个向量:一个 Key,一个 Value。可以把它们理解为这一层给这个 token 做的笔记——Key 是「我是什么,谁该来找我」,Value 是「找到我之后你能拿走什么」。后面生成新 token 时,模型拿当前的 Query 去跟所有历史 Key 比对,按相似度加权取出对应的 Value。

这堆笔记就是 KV 缓存。它存在的意义是:算过的不用再算。生成第 1001 个 token 时,前 1000 个 token 的笔记直接读缓存。

麻烦在于笔记是私人的。14B 的第 20 层写下的 Key,和 32B 的第 20 层写下的 Key,是两套不同权重矩阵的产物。数值不同,含义的编码方式不同,连层数都不一样(论文给的深度比:Qwen3 这几对是 1.6~1.8 倍,Llama 3.1 8B→70B 是 2.5 倍)。把 14B 的缓存直接塞给 32B,等于把别人的速记本推到你面前让你接着往下讲。

那有没有可能,这两套笔记之间存在某种系统性的对应关系?

三个朴素猜想,逐个撞墙

猜想一:直接复制。 只要两个模型的 KV 维度一样就行。

Qwen3 14B 和 32B 恰好都是 8 个 KV 头、每头 128 维——论文称这种情况为 matched-KV。形状对得上,张量塞进去不会报错。

但这只保证了不崩,不保证有意义。论文的消融实验里有个数字说明了这一点:去掉跨层选择、只用单层对应( $k=1$ ),GSM8K 从 90.98% 掉到 0.38%,困惑度从 7.33 涨到 22.73。形状对齐距离语义对齐还差得远。

猜想二:训一个小网络来翻译。 这是自然的下一步,也是先前工作的路子。

问题是成本。每一对模型都要训一个,每加一个新模型就要重训。而且要有梯度、要有训练数据、要调超参。对于「今天线上换了个模型」这种运维事件,这个反应速度不够。

论文的发现是:大部分情况下不需要非线性。 跨模型 KV 关系已经足够线性,一个闭式解就够了。四对成功的模型里,ridge 回归打平甚至略胜 MLP(Qwen3 14B→32B:ridge 97.6%,MLP 97.3%)。

猜想三:那就每层对每层做个线性回归。 目标模型第 $l$ 层的 Key,用源模型第 $l$ 层的 Key 去拟合。

这就是 $k=1$ 的那一栏,0.38%。原因是信息不在同一层。目标层 $l$ 需要的东西,散落在源模型的多个层里。论文的做法是对每个目标层 $l$ ,按 $R^2$ 挑出最相关的 top- $k$ 个源层,把它们的 KV 特征横向拼接起来一起做回归:

$\mathbf{X}_{K}^{l}=[\bar{\mathbf{K}}_{s}^{l_{1}}\,\|\,\bar{\mathbf{K}}_{s}^{l_{2}}\,\|\,\cdots\,\|\,\bar{\mathbf{K}}_{s}^{l_{k}}]$

$k$ 按模型对扫描选取(Qwen3 14B→32B 选 8,Llama 3.1 8B→70B 选 20)。 $k$ 从 1 提到 8,Key 的 $R^2$ 从 0.56 升到 0.79。论文明确说,跨层源选择是三个组件里贡献最大的一个。

机制:位置要先剥掉

还有一道坎:位置编码。

现代模型用 RoPE(旋转位置编码)标记 token 的位置。做法是按位置把 Key 向量旋转一个角度:第 5 个 token 转 5 份,第 500 个转 500 份。关键在于,KV 缓存里存的是已经转过的 Key:

$\mathbf{k}_{\text{RoPE}}(t)=\mathbf{R}_{\Theta}(t)\,\mathbf{k}_{\text{content}}$

如果直接在转过的 Key 上拟合,回归权重就把位置信息一起学进去了——而且只学到了拟合时见过的那个位置分布(论文用 1024 token 标定)。上线跑 32K,就超出了权重见过的范围。

解法是先转回去。因为旋转矩阵 $\mathbf{R}_{\Theta}$ 是正交的,逆运算精确且几乎免费:

$\hat{\mathbf{K}}_{t}=(\mathbf{K}_{s}\,\mathbf{R}_{\Theta_{s}}^{-1}(t)\,\mathbf{W}_{K}+\mathbf{b}_{K})\,\mathbf{R}_{\Theta_{t}}(t)$

读法是从右往左三步:剥掉源模型的旋转,在无位置的「内容空间」里做线性换算,再按目标模型的规则重新旋转。Value 不带位置编码,直接换算。

这一步的重要性有实验支撑,而且很不均匀。消融表里有一栏叫 - inference RoPE——在内容空间拟合了,推理时却忘了重新旋转:

配置	ARC-C	HellaSwag	WinoGrande	MMLU	GSM8K
完整	61.60	80.70	68.98	78.09	90.98
少了推理时 RoPE	44.97	75.39	56.59	25.79	4.17

MMLU 掉到 25.79%——四选一的随机水平就是 25%。GSM8K 掉到 4.17%。但 HellaSwag 只掉了 5 个百分点。同一个 bug,在不同任务上一个致命一个几乎看不出来。只测常识题的话,这个错误会被漏掉。

有意思的是,如果拟合和推理两边都保留 RoPE(完全耦合),短上下文上几乎无损(61.09 / 80.73 / 68.59 / 77.70 / 90.98)。剥离 RoPE 的价值不在短文本精度,而在于让权重与位置解耦,从而按构造外推到长上下文。

标定成本很低:500 条 FineWeb-Edu 序列,每条 1024 token,步长 4 下采样,约 12.8 万个 token 级观测。正则系数 $\lambda$ 扫四个数量级都有宽平坦区,只在 $\lambda=1$ 才崩;样本数 $N=200$ 之后就平了, $N=50$ 也只差 1.6 个百分点。唯一有真实代价的轴是标定语料的领域:换成 CodeAlpaca,HellaSwag 掉 5.24 个百分点;换成 Wikipedia 在噪声内。

六对模型,分成两档

结果分层很清楚。「保留率」= 换算后准确率 / 目标模型自己跑的准确率:

模型对	$k$	五基准均值	ARC-C	GSM8K
Qwen3 14B → 32B	8	97.6%	101.0%	95.6%
Qwen3 8B → 32B	12	87.5%	94.0%	68.8%
Llama 3.1 8B → 70B	20	72.8%	90.9%	18.2%
Ministral 3 3B → 8B	全部	76.2%	90.6%	36.6%
Ministral 3 3B → 14B	20	44.2%	43.6%	3.2%
Ministral 3 8B → 14B	12	41.6%	40.7%	1.6%

三件事值得注意。

一,GSM8K 那一列崩得最狠。 即使在成功档里,Llama 8B→70B 的数学题只剩 18.2%,而它的 ARC-C 有 90.9%。多步算术对缓存保真度的要求远高于常识判断。同一个换算器,看常识题像成功,看数学题像失败。

二,参数比不是障碍。 Llama 3.1 8B → 70B 是全表最悬殊的比例(近 9 倍),却排在成功档。而 Ministral 8B → 14B 只差 1.75 倍,是全表最差。跨的参数量差距不预测成败。

三,六对全是 matched-KV。也就是说,形状对齐是这套方法的入场条件,而不是成功的保证——两个失败对同样满足 matched-KV。论文把 mismatched-KV 明确列为未测。

为什么这样设计:换掉判据

上面的两档分裂催生了论文最有价值的部分。

如果想上线,你需要一个事前指标:给定一对模型,拟合完之后能不能判断这次换算靠不靠谱?自然的候选是 $R^2$ ,即回归拟合得多好。

$R^2$ 不管用。论文给了一组对照:Llama 3.1 8B→70B 的 Key $R^2$ 是 0.84,小到大方向保住 94%,大到小方向只有 37%。Ministral 3B→8B 的 $R^2$ 也是 0.84,两个方向都保住 93%。同一个 $R^2$ ,截然不同的下游结果。跨 12 组评测统计,标定域 $R^2$ 与 HellaSwag 保留率的相关系数是 $r=-0.20$ ——基本无关。

原因在于 $R^2$ 衡量的是平均重建每一个 K/V 通道的好坏,所有维度等权。而注意力不等权。注意力拿 Key 去和目标模型的 Query 打分,再按打出来的权重去加权 Value。有些方向 Query 根本不看,那里的误差无关紧要;有些方向 Query 盯得很紧,那里差一点就全乱。

于是论文换了个量:注意力输出余弦——直接测「用换算后的 KV 算出的注意力输出」与「用真 KV 算出的注意力输出」之间的余弦相似度,跨层跨头取平均。它与保留率的相关系数是 $r=+0.57$ 。

再往下一层是「误差集中度」:把每 token 的 K 误差投影到目标模型每头 Query 矩阵 $\mathbf{Q}_h$ 的右奇异向量上,按对应奇异值的平方加权,再除以全体分量的平均误差。集中度 > 1 表示误差正好落在注意力要读的地方,< 1 表示落在注意力忽略的地方。

这解释了 MLP 为什么在失败对上有用。把 ridge 换成 MLP,在两个失败对上:

模型对	ridge	MLP	$\Delta$	$\Delta$ K-集中度	$\Delta$ 余弦
Qwen3 14B → 32B	97.6%	97.3%	$-0.3$ pp	$-0.03$	$-0.03$
Ministral 3 3B → 8B	93.3%	91.8%	$-1.5$ pp	$-0.45$	$+0.07$
Ministral 3 3B → 14B	68.0%	92.3%	+24.3 pp	$-2.31$	$+0.41$
Ministral 3 8B → 14B	58.7%	95.5%	+36.8 pp	$-2.71$	$+0.48$

MLP 并没有把总误差压到多小——它把误差挪到了注意力不看的方向。集中度降约 2.5,余弦升约 0.45,HellaSwag 保留率涨 24 到 37 个百分点。

而且这不是「MLP 更强」。在 ridge 已经成功的对上,MLP 略微更差。更能说明问题的是 Ministral 3B→8B:集中度和余弦两个指标都改善了,HellaSwag 反而掉了 1.5 个百分点。重新分布误差不是充分条件,只在原本错放的误差大到构成瓶颈时才起作用。

失败对的另一个特征:ridge 在 HellaSwag token 上的 $R^2_K$ 是深度负值( $-7.81$ 和 $-3.22$ )。也就是说标定集上拟合得好的线性映射,到了评测 token 上根本外推不了。MLP 把它拉回接近零(仍为负)。

边界

家族内,稠密全注意力。 跨家族(Qwen3 → Llama 3.1)是开放问题。滑窗/局部注意力等混合架构,以及 Nemotron 3 这类带 SSM 状态的注意力-循环混合体,都在范围外。
余弦是事后指标。 它要先把 mapper 拟合出来才能算,所以只能筛掉坏的,不能省掉拟合。论文把「拟合前就能估计可迁移性的信号」列为待做。
标定只用了一个领域。 FineWeb-Edu。附录在一对模型上量了换 Wikipedia / CodeAlpaca 的代价,但两者都没有把「学科」和「语体」分开,所以不能推断标定限制在医学、法律这类单一领域会怎样。
$k$ 是在报告用的同一批基准上选的。 附录测出这个影响最多 2.49 个百分点并补了留出评测,但都不等于样本外选 $k$ 。
多轮只测了一对。 Qwen3 14B ↔ 32B,CoQA,100 段约 15 轮的对话。小到大的差距 10 轮累积 1.7 个百分点;大到小是线性漂移,每轮 0.33 个百分点。十轮内不会雪崩,但线性项在很长的会话里仍会累积。
单篇预印本,一周内。 上面每个数字都来自这一篇作者自己的测量,没有独立复现。