12288 参数重读历史 KV,多轮解题 AIME 涨 11.67 分

Can Computation from Earlier Problems Help LLMs Solve New Ones?

Jipei He, Wenhui Tan, Xiaoyi Yu, Enver Sangineto, Fiorenzo Parascandolo, Rita Cucchiara, Ruihua Song

cs.AI, cs.CL, cs.LG

2026-09-30

冻结主干只训 12288 个参数,STAIR 让新题重读旧题的 KV 缓存,多轮 AIME 准确率较原生带历史提升 11.67 个百分点。

这篇在解决什么

同一个会话里连着问几道独立的题,是聊天产品的常态。对用户,前面的题只是留在上下文里的文字;对模型,算这些文字时的内部计算也留了下来,新问题的 query 会通过 attention 直接读到这些位置。

这种遗留计算的效果并不稳定。MATH-500 上 Qwen3-4B Instruct 的 Avg@4 从无历史的 95.00% 掉到第四轮的 93.75%,GPQA-Diamond 上 Qwen3.5-4B 反而从 63.64% 涨到 75.25%。逐题配对统计更直观:第四轮 35 道题从错改对,60 道从对改错。历史既帮也伤,模型自己控制不了读什么。

方法

分析部分先做受控重放:固定当前题目、更换前置历史,两组各 128 题 × 32 条三轮历史交叉,把历史造成的表示位移分解为题目效应、历史效应和交互项。当前题的 query 确实在读历史,第 3 层约 39% 的注意力落在早前回答上,第 19 层还剩约 12%。第 19 层的交互项只占总位移能量 0.69%,但几何结构跨历史复现:题目间 8 近邻的重合率 48.06%,打乱身份的对照只有 11.73%。历史影响里有一个与当前题目绑定、反复出现的成分,学习读取因此有依据。

STAIR 本体分四步:

差分读出是承重设计:模型只学「反射改变了什么」,天然贴着 Native 预测,训练损失里再加一项对 Native teacher 的 KL 散度拴住它。可训练的只有各层各 head 的反射法向量,三个模型都是 12,288 个参数。训练数据是各模型自己在 DAPO-Math-17k 上生成的回答,错的也保留,teacher forcing 回放成四轮会话,训一个 epoch。

结果

T2-T4 平均 Avg@4,各条件均带历史,以 Native 为对照:

设置NativeSTAIR差值
Qwen3.5-4B · AIME 202540.5652.22+11.67
Qwen3.5-9B · AIME 202543.3349.44+6.11
Qwen3-4B Instruct · AIME 202534.4438.06+3.61
Qwen3.5-4B · MATH-50094.7896.27+1.49
Qwen3.5-9B · GPQA-Diamond79.8875.93-3.95

AMC23 上 Qwen3.5-4B 涨 6.37 个点(68.14 到 74.51)。MATH-500 上两个 Qwen3.5 都超过了无历史的 Vanilla 参考(96.27 对 96.00;96.78 对 96.65)。AIME 是另一回事,Qwen3.5-4B 的 Vanilla 有 60.83,STAIR 只有 52.22。STAIR 修的是「带历史」场景,多数设置下还没修到清空历史的水平。

共享同一 T1 历史的对照最干净(30 道 AIME 题、匹配轨迹):Native 21.67%,STAIR 48.33%,Pass@4 从 56.67% 升到 80.00%。消融齐全:随机反射方向 37.22%、去掉差分直接用反射读数 41.67%、同参数量但不读 bank 的控制器 36.67%、打乱 bank 内 K-V 配对 33.06%,全部低于完整方法。与 LoRA 对比(Qwen3-4B Instruct, AIME):STAIR 用 12,288 个参数拿 +3.61,rank 2 的 LoRA 用 336 倍参数拿 +5.00,两者联合训练再比 LoRA 单独高 3.89,学习读取与权重适配可以叠加。

开销:全量读取中位数多 0.81 秒、6.60 GiB 峰值显存,bank 中位约 63.4K token,每轮建 bank 还需 7.29 秒精确重放;可读 bank 截到 2K/8K/32K token,准确率从 48.33% 掉到 35.00/43.33/39.17%,基本得读全量。

为什么重要

多轮长会话的精度退化是部署里的真问题,agent 连续解题、客服连续答疑都碰得到。已有 KV cache 复用工作(Prompt Cache、CacheBlend)省的是 prefill 计算,这篇把遗留计算当精度资源来学读取,切入角不同。12,288 个参数、主干全冻、可与 LoRA 叠加,工程上近乎零侵入。机制发现也独立成立:历史引起的表示变化里有跨历史复现的题目相关结构,「读什么」是可学的,不只是噪声。

局限与存疑

作者自述三条:反射方向只在数学数据上训练,GPQA 上两个 Qwen3.5 无增益、9B 还下降,是否过拟合训练分布未验证;证据只覆盖 Qwen 家族的四轮会话;会话变长后 bank 线性增长,选择性保留还没做。

读下来另有两处要泼冷水。AIME 2025 每轮只有 120 个回答,单个回答翻转就值 0.83 个点,11.67 的最大涨幅落在最小最噪的 benchmark 上;MATH-500 上最多 1.49 个点(4B),9B 和 Instruct 不到 0.6。对照锚点始终是带历史的 Native,不是清空历史的 Vanilla,诚实的定位是「如果必须带历史,这是目前最便宜的补救」,并非全面超过清空上下文。每轮 7.29 秒建 bank 重放加 6.6 GiB 显存,延迟敏感的服务用不起。

术语

原文与代码

相关论文

全部论文解读