研究
我的研究关注长上下文模型应该如何记住读过的内容。注意力保留一切,也为此付出代价;循环模型和快权重模型则在压缩的同时遗忘。我的大部分工作都围绕这一取舍展开:记忆是什么,如何写入,又如何读出。
其中贯穿着三个研究方向。
- 测试时训练——记忆是前向传播过程中在线优化的快权重。我研究谁拥有这些权重,用什么信号监督它们,以及读取它们是否一定要是确定性的。
- 高效架构——记忆是 KV 缓存或残差流。我研究如何去除其中的冗余维度:缓存中的深度,以及层间路由中的累积重叠。
- KV 缓存压缩与长上下文推理——记忆是一笔必须分配的固定预算。我研究如何分配它:在层之间、注意力头之间,以及模型自身的思维链之中。
以下每个项目都按研究假设、具体方法、评测设置与实验结果展开。
测试时训练
快权重是一种在模型阅读过程中不断优化的记忆。围绕这种记忆,我关注三个问题:谁拥有它、什么监督它、以及如何读取它。
Universal Test-Time Training(通用测试时训练)
Adobe Research · 投稿中
研究假设。 现有每一种 TTT 设计都让每层拥有一个其他层无法访问的私有快权重张量。这一惯例来自 Transformer 的参数布局,而非测试时优化本身的任何性质,并带来三个代价:固定预算被分成 L 份,使每层只能得到一个较窄的算子;沿深度的划分在设计时便已固定;实现需要执行 L 次小操作,而不是一次分组更新。我们的假设是,快权重应该是全局通用的。而且,共享对象是循环状态,而非静态参数,因此共享改变的是模型能记住什么,不仅仅是参数如何布局。
方法。 一个快权重模块库归整个模型所有,而不属于某个网络块;每个块保留自己的慢权重和局部注意力,但只拥有一个地址,即指向模块库的路由函数。读取时,对所寻址的模块做门控组合;到达分块边界时,将所有路由到同一模块的层与注意力头的贡献累加,用于更新该模块。权重归属由此变成一种路由约束,让三种设计处于同一条阶梯上:uTTT-Dense(所有层使用同一个共享稠密算子)、TTT-MoE(路由模块池,但仍局限于各层内部;作为控制变量的中间方案,将路由的影响与全局共享的影响分开),以及 uTTT-MoE(一个全局共享的模块池)。另一个结果表明,通用 TTT 模型是一个在(分块,层)联合索引上完全循环的模型:深度成为第二个循环维度,提供额外的内循环步骤;某个分块的深层写入的状态,可以在下一个分块的浅层读取。每层私有的 TTT 则是严格采用固定分区的特殊情况。
评测设置。 在 Long-Data-Collections 上,以 32K 上下文训练 124M 和 760M 参数的大分块自回归语言模型,评测 Books3 上的逐 token 损失,以及 4K/8K/16K/32K 长度下全部 13 项 RULER 检索任务,对比 LaCT、DeltaNet-SWA、Gated DeltaNet-SWA 和完整注意力。另评测了 LVSM 风格的新视角合成,包括物体级(Objaverse → Google Scanned Objects)和场景级(DL3DV)任务。
结果。 从每层私有的 LaCT,到层内路由池(先不做负载均衡,再做负载均衡),再到全局共享池,RULER 13 项任务的平均分在 124M 规模上依次为 8.50 → 9.99 → 12.82 → 15.46,在 760M 规模上为 23.21 → 24.86 → 25.73 → 27.85;模型规模扩大约 6 倍后,这一排序依然成立。只有最后一步真正改变了权重归属,因此控制变量后的全局共享收益是 124M 上 +2.6 分,760M 上 +2.1 分;其余提升来自路由与负载均衡。渲染方面,物体级 PSNR 从 24.63 dB(每层私有的稠密算子),提升到 24.79(层内模块池),再到 25.52(全局模块池)。其中,在容量相同的条件下,比较将 64 个专家沿深度分配与全局共享,所对应的归属变化是 24.79 → 25.52 这一步。规模匹配的稠密变体达到 25.93 dB,扩大激活专家集合后可达到 26.51 dB。uTTT-MoE 在两个规模上的逐 token 损失都低于完整注意力,但检索分数仍低于后者。
结论的适用边界。 优势主要集中于较短上下文:按长度拆分 760M 的结果,uTTT-MoE 在 4K 上提升 8.9 分,却在 16K 上下降 1.8 分;32K 上表现最强的快权重模型则是另一个变体。场景级渲染中,不同配置相差 0.19 dB,作者认为这不足以区分优劣。uTTT-Dense 虽然有定义,但没有进行语言建模评测,因此只出现在渲染结果中。此外,关于效率的那一半论证尚未得到测量支持:论文没有报告吞吐量或峰值内存,而能够加强深度循环的顺序更新调度,又会失去作为设计动机之一的分组执行优势。
Test-Time Generation(测试时生成)
Adobe Research · 投稿中
| 研究假设。 现有快权重记忆通过拟合确定性目标来更新,并通过所得的端点映射读取状态。平方损失下的贝叶斯预测器是条件均值,因此具有多峰分布的绑定会坍缩到它们的平均值;在容量有限时,共享的值空间几何、键路由和实例间残余差异,都要在这一个端点映射中竞争。我们的假设是,记忆应该建模条件分布 *p(v | k, h)*,而不只是它的均值——潜在噪声恰好可以承载回归所平均掉的那部分残余不确定性。 |
方法。 将记忆读出重新定义为以键为条件、从潜在噪声到目标值的传输过程。把快权重状态视为条件关联记忆,并建立噪声与值之间的桥梁后,检索间隔可以分解为键先验项,以及一个由信噪比(SNR)放大的值空间去噪项;也就是说,迭代回忆能够带来一次性查找所无法提供的区分度。状态始终不被显式展开为表格,而是调制一个门控 MLP 去噪器。写入时,对观测到的绑定加入噪声进行监督,每个绑定使用多次噪声采样;读取时从高斯噪声出发,按所选步数进行积分,从而让读出深度成为测试时计算量的调节旋钮,而不是训练时就要固定的选择。
评测设置。 增量式新视角合成,包括物体级(Objaverse → GSO)和场景级(DL3DV-10K),对比 LaCT 风格的确定性读出。在 Long-Data-Collections 上,以 124M 和 760M 规模进行长上下文自回归语言建模,在 32K token 序列上用 Book-3 的逐 token 损失评测。此外,还进行测试时计算量扩展实验、一项参数量可比的公平性对照,以及目标函数消融。
结果。 物体级平均 PSNR 从 25.09 → 26.00,最后一个评测视角从 26.31 → 27.23。场景级提升更大:平均值 16.27 → 17.46,最终视角 16.80 → 18.64。这一现象符合检索间隔分析:仅靠位姿键无法消除纹理和遮挡的歧义,因此额外区分度来自值空间去噪。语言建模在完整上下文范围内保持了持续但较小的损失优势。两个测试时调节项(更新重复次数与读出深度)都能在无需重新训练的情况下提升可达到的质量上限。一项参数量可比的对照,在比主要结果更低的配置下运行两种 TTG 变体,PSNR 为 24.92 与 24.95,表明记忆模块的参数分配并非推动该数值变化的关键;但该对照没有包含参数量匹配的确定性基线,因此它提供的是关于 TTG 自身敏感性的证据,而非完整的归因。目标函数消融也需要明确指出:在相同协议下,基于能量的变体优于默认的流匹配方案,平均 PSNR 为 24.16 → 24.99。作者将其解读为:贡献在于生成式记忆接口,而不在于流匹配这一特定选择。
成本与证据基础。 TTG 用计算量换取质量:随机重复会增加写入成本,更深的读出会提高每次查询的延迟,合适的运行配置取决于任务与预算。论文也指出,其证据基础仍较有限:数据集种类不多,大多只进行了一次训练,且没有误差条或多随机种子的显著性检验。因此,这些趋势被用作机制的证据,而不是一条缩放定律。
Test-Time Training with Next-Token Prediction(以下一 token 预测监督测试时训练)
2026 年 arXiv 预印本 · 共同第一作者 · PDF
研究假设。 原位 TTT 已经基本固定了快权重放在哪里(MLP 的下投影)以及如何写入(分块并行的秩一更新),剩下的开放设计维度是监督目标,而先前方案在这里存在问题。它们训练快权重去重建学习到的局部激活代理;这一工程手段与真正训练骨干模型的下一 token 预测信号脱节。我们的假设是:测试时,每个提示词 token 本身就携带 NTP 信号;让写入与模型自身的因果轨迹对齐,可以在不重新设计骨干网络的情况下改善长上下文利用能力。
方法。 在选定的 MLP 块中,将下投影原位复用为快权重。值目标是同一层在下一个位置的上下文隐藏状态,即真正观测到下一个 token 后产生的状态,并通过一个以恒等映射初始化的轻量投影。写入以分块并行的独占前缀和方式累积,且分块最后一个 token 不参与配对,因此更新严格满足因果性,未来信息不会进入某个 token 自身的预测。推理时,由于完整提示词已被观测,可基于缓存的激活,将写入闭式求解为岭回归;替换权重后,基于已有的提示词 KV 缓存进行解码,并在每个样本结束后恢复原权重。只有这个推理时的写入步骤可以即插即用;快权重写入规则及其投影需要先经过持续预训练学习,依骨干模型不同,训练量为 0.1B–2B token。
评测设置。 在四个骨干模型——Llama-3.1-8B、Mistral-7B-v0.3、Qwen3-4B 和 Qwen3-0.6B——上,评测 4K/8K/16K/32K 长度下的 RULER Full-13,以及通过 lm-evaluation-harness 衡量的通用能力。另外,在 Llama-3.1-8B 和 Mistral-7B-v0.3 上评测 LongBench-v2 的中等长度子集:基于 33K–128K 词文档的 215 道选择题。
结果。 这是唯一在全部四个骨干模型上都比原始模型提高了 RULER 分数的方法:Full-13 平均提升分别为 +3.90 / +3.03 / +4.06 / +2.88;单项最大提升是 Qwen3-4B 在 16K 上的 +11.74。计算量匹配的基线反而出现退步:持续预训练让 Llama-3.1-8B 下降 5.9 分,Mistral 下降 7.7 分。在 LongBench-v2 上,该方法在两个骨干模型上均取得最高的总体准确率:Llama 为 25.6 → 31.2,Mistral 为 26.5 → 30.2(持续预训练和原位 TTT 在这里也有提升,但幅度较小)。两组消融支持了核心论点:在 Qwen3-4B 上,NTP 目标在全部四个长度上都优于所有卷积代理目标,4K 约高 5 分,8K 约高 3 分,16K 约高 9 分,32K 高 5–8 分;用简单的 Hebbian 内积替换岭回归求解后,方法完全失效(12.61 对比 59.70),表明 Gram 白化起到了关键作用。通用能力基本保持不变。
高效架构
同一个问题,不同形式的记忆。这里的冗余是结构性的:缓存或残差流的某个维度,承载的信息少于它的体积所暗示的程度。
LastKV: Last Layer Key-Value Cache is All You Need(只需最后一层的键值缓存)
Adobe Research · 投稿中
研究假设。 长上下文 KV 缓存同时沿两个维度增长:序列长度和网络深度。token 维度的方法(淘汰、压缩、滑动窗口)已经得到广泛研究,深度维度却没有。现有跨层共享要么没有走到极限(只在少量相邻层间共享),要么在其他地方付出代价(在每个 token 位置进行凝聚,破坏严格的并行训练;或者要求重新设计 decoder-decoder 架构)。我们的假设是:共享不必发生在每个 token 位置,只需发生在承载长程记忆的地方。 将缓存写入推迟到分块边界后,最后一层的缓存就可以合法地供所有层读取,从而在不修改 Transformer 块的情况下实现最大 L 倍共享。
方法。 每个分块都以标准的逐层 KV Transformer 方式运行,保留完整的分块内部并行性。分块结束时,只提交其最后一层的 KV 对,释放另外 L−1 层缓存;之后的分块在每一层拼接这份共享上下文,并通过普通注意力读取它。保留缓存量从 O(TCLd) 降为 O(TCd)。由于语言和视频都需要跨越分块接缝的连续性,方法加入相邻分块细化:仅为前一个分块保留逐层缓存,再通过可学习的逐层、逐头权重,将其与该分块的最后一层缓存插值融合;760M 模型只增加 576 个标量。默认构造复用最后一层已经算好的 K/V,因此不增加任何参数或 FLOPs。
评测设置。 四组设置,刻意覆盖不同模态。在 32K 上下文上从零预训练 124M/760M/3B 语言模型,评测 Books3 验证损失和全部 13 项 RULER 任务。通过持续预训练转换 Marin 8B,评测 12 项知识、推理、数学基准及训练吞吐量。自回归文本生成视频:在 teacher forcing 和 Self-Forcing 蒸馏下微调 1.3B 与 14B 的 Wan 2.1 T2V,采用完整 VBench 协议评测。另外,还评测物体级与场景级新视角合成。
结果。 在 760M 上,三个 LastKV 变体的验证损失均低于完整缓存基线(2.5373 / 2.5366 / 2.5375,对比 2.5423)。最关键的是持续预训练结果:一个公开发布的 8B Transformer 被转换为在 32 层中只保留一层 KV;在 12 项基准中的 8 项上,与原始基线的差距约在 1 分以内,平均下降 0.6 分;分块大小为 1024 时,维持了 99.6% 的前向吞吐量和 95.6% 的前向加反向吞吐量。这一比较是正在进行中的训练快照:在总计 2T token 的计划中,选取约训练到 120B token 的匹配检查点,因此它主张的是在较早训练预算阶段达到接近的表现,而非取得优势。渲染质量提升,物体级为 25.38 → 25.95 dB,场景级为 17.46 → 17.81。视频在 teacher forcing 下保持表现:14B 的 VBench 总分为 80.29 → 80.45,1.3B 为 82.09 → 81.84;但在 Self-Forcing 蒸馏下,总分比完整缓存低 0.4–0.6 分,差距几乎全部来自动态程度(dynamic degree)。边界消融结果清晰:没有相邻分块细化时,每个边界之后的损失尖峰为 +0.103 nats;加入后为 −0.003。
代价落在哪里。 检索表现有升有降,并非始终保持:在使用门控注意力的 760M 模型上,13 项 RULER 的平均分在 16K 时从 43.6 降到 29.6,在 32K 时从 34.5 降到 31.6,多键任务更有利于完整缓存;但去掉门控后,8K–32K 范围内的排序会反转。此外,尚无端到端部署基准:没有在优化过的软件栈上报告解码延迟、每秒 token 数,或固定内存预算下可支持的最大上下文。
Delta Attention Residuals(增量注意力残差)
2026 年 arXiv 预印本 · 共同第一作者 · PDF
研究假设。 Attention Residuals 用对先前层输出的可学习 softmax 注意力,替代固定的加法残差;但它关注的是累积隐藏状态。随着深度增加,这些累积和共享越来越长的公共前缀。我们的假设是:被路由状态中的来源冗余,加上每次重置都会丢弃残差流的替换式路由,共同导致方法失效。高度相似来源的 softmax logits 缺乏对比度,路由趋向均匀并发生坍缩,最终退化为平均。
方法。 两处修改。增量来源:对每个子层的输出(相邻隐藏状态之差)进行路由,而非对累积状态路由;注意力与 MLP 输出处于不同子空间,因此来源在结构上就具有多样性。加法路由:把路由组合加到残差流上,而不是替换它,确保残差流始终保留,并且在块边界不会丢失任何贡献。
评测设置。 在 FineWeb-Edu 上从零预训练五种规模的语言模型:220M、533M、1044M、一个 Qwen3-0.6B 配置,以及一个 7.57B 配置。另外,对预训练 Qwen3-0.6B 进行微调转换,评测 8 项下游基准,并直接测量沿深度变化的路由尖锐程度。
结果。 一次优劣反转验证了诊断:在 1044M 规模上,基线的验证困惑度为 29.70,AttnRes 比不做修改更差(31.76),而增量方案改善至 29.13。在 7.57B 上,差距扩大:基线 17.43,AttnRes 18.58,Delta Block 为 16.00,改善 8.2%;为此增加 589.8K 个路由参数(占模型的 0.008%),代价是训练吞吐量下降 35%,内存增加 3%。机制检验也与解释一致:AttnRes 的最大路由权重从浅层约 1.0 衰减到深层约 0.2,而增量版本保持在约 0.6,平均值为其 1.8 倍。
成本。 最细粒度的变体在 1044M 上达到最佳困惑度,但吞吐量下降 69%,内存为原来的 3.5 倍;推荐默认使用更粗粒度的 Delta Block,正是因为它以约三分之一的额外开销保留了大部分收益。
KV 缓存压缩与长上下文推理
固定的内存预算,以及如何分配它的问题。这是我博士研究的起点,也是目前得到最广泛采用的研究方向。
PyramidKV: Dynamic KV Cache Compression based on Pyramidal Information Funneling(基于金字塔式信息汇聚的动态 KV 缓存压缩)
COLM 2025 · Spotlight(24 / 418) · PDF · 项目主页 · 代码
研究假设。 先前每种淘汰方法都给各层相同的预算。但注意力并非沿深度均匀分布:在底层广泛分散,在中层汇集成局部模式,到顶层则集中在少数关键 token 上。如果信息确实呈漏斗式汇聚,统一预算就会在顶层浪费内存,同时让底层资源不足。
方法。 将固定总预算按递减等差数列分配给各层:底层多、顶层少,总量保持不变。在每一层、每一个注意力头内部,保留从末尾指令窗口获得最多注意力权重总量的 token。方法无需训练,在 prefill 阶段应用,并且可与量化独立结合。
评测设置。 在 LLaMA-3-8B/70B-Instruct 和 Mistral-7B-Instruct 上评测 LongBench(17 个数据集,覆盖单文档与多文档问答、摘要、小样本学习、合成检索和代码补全)与 Needle-in-a-Haystack,将缓存预算从 64 扫描到 2048。
结果。 只保留 12% 的 KV 缓存即可匹配完整缓存质量。在仅保留 0.7% 的极端设置下,当预算为 64 时,LLaMA-3-8B-Instruct 上的 TREC 分数比次优方法高 19.5 个绝对分数点(58.00 对比 38.50)。在 Needle-in-a-Haystack 上,LLaMA-3-70B-Instruct 只需保留 128 个条目,就能在最长 8K 的上下文中达到 100.0 的准确率。
R-KV: Redundancy-aware KV Cache Compression for Reasoning Models(面向推理模型的冗余感知 KV 缓存压缩)
NeurIPS 2025 · PDF · 项目主页 · 代码 · 视频
研究假设。 KV 压缩原本针对长输入提示词设计。推理模型却反转了问题:膨胀来自模型逐 token 生成、且高度自我重复的思维链。纯粹基于注意力重要性的淘汰此时会失效,因为它忠实保留了许多重复推理步骤的近似副本;实际上,在 10% 预算下只保留约 60% 的完整缓存准确率。我们的假设是:加入显式冗余惩罚,就能消除这部分损失。更进一步的主张是:移除冗余推理 token 是一种甚至可能优于完整缓存的去噪过程。
方法。 每个 token 的得分等于其注意力重要性减去它与其他已保留键的平均余弦相似度,因此既被关注、又不重复的 token 会胜出。该方法在解码过程中实时运行,每当解码缓冲区填满就重新评分,而非只在 prefill 时应用。无需训练。
评测设置。 在 DeepSeek-R1-Distill-Llama-8B 和 DeepSeek-R1-Distill-Qwen-14B 上,评测长思维链数学推理任务 MATH-500 和 AIME 2024,并在 16K 序列长度下测量吞吐量和批量大小。
结果。 在使用 R1-Llama-8B 的 AIME-24 上,R-KV 在 10% 预算下达到无损表现,而只按注意力淘汰的方案仅达到约 60%;在 16% 预算下,达到完整缓存准确率的 105%,验证了去噪效果。达到无损压缩所需的预算取决于具体设置:同一模型在 MATH-500 上需要 34%,R1-Qwen-14B 则分别需要 25% / 54%。系统表现方面:在压缩后保留 10% 缓存时,节省 90% 内存,吞吐量为原来的 6.6 倍,批量大小为原来的 9 倍;固定为 1024 token 预算时,批量大小为原来的 13.4 倍,吞吐量为 9.2 倍。
Not All Heads Matter: Head-Level KV Cache Compression(并非所有注意力头都同等重要:头级 KV 缓存压缩)
研究假设。 包括 PyramidKV 在内的层级分配方法,隐含地将同一层内所有注意力头视为同等重要。实际上,各头的重要性高度不均,因此按头分配是更细的控制维度。更关键的假设涉及如何衡量重要性:标准的“检索头”分数通过纯粹复制粘贴式的大海捞针检索来估计,奖励的是逐字复制;但真实长上下文问答需要检索,以及对检索结果进行推理。
方法。 通过一个探针离线估计头的重要性:其中的“针”包含推理步骤、干扰项和正确答案;给某个头计分时,计算它落在整个正确答案上的注意力权重总量,而非只看单个最大值位置。固定这些分数,并据此在所有头之间重新分配全局预算。
评测设置。 在 Llama-3-8B-Instruct 和 Mistral-7B-Instruct 上,以 64 到 1024 的预算评测 LongBench 上下文问答、LooGLE 长依赖问答、Needle-in-a-Haystack,以及 Reasoning-in-a-Haystack 探针。
结果。 在上下文问答上,保留 1.5% 的 KV 缓存即可达到完整缓存表现的 97%。收益恰好集中在假设所预测的低预算区间:预算为 64 时,在 Llama-3-8B-Instruct 上领先最强基线 Ada-SnapKV 约 3.4 分,在 Mistral-7B-Instruct 上约 3.6 分,比层级分配高出 5 分以上。
HeadInfer: Memory-Efficient LLM Inference by Head-wise Offloading(通过逐头卸载实现内存高效的大语言模型推理)
ICML 2025 长上下文基础模型研讨会(Workshop on Long-Context Foundation Models) · PDF · 代码
研究假设。 这是与压缩对应的无损方案:不去决定丢弃什么,而是把内存瓶颈视为一个数据放置问题。现有卸载按层进行,仍要求 GPU 同时容纳整层缓存,因此显存占用仍随隐藏维度增长。但各注意力头在计算上独立,因此常驻显存占用还能再缩小一个头数的倍数,而且不引入任何近似。
方法。 每次只让一个头的 KV 缓存常驻 GPU,其余从 CPU 内存中流式传入;结合分块 prefill、随上下文长度自适应的头分组,以及双缓冲异步预取,让 PCIe 传输与计算重叠。
评测设置。 在消费级硬件上,用 LongBench v2、SCBench(9 项任务,平均 227K token)、4K–128K 长度下的 RULER,以及 Needle-in-a-Haystack,测量可达到的最大上下文和吞吐量。
结果。 对于 Llama-3-8B 的 1M token 序列,GPU KV 缓存从 128 GB 降到 1 GB,GPU 总内存占用从 207 GB 降到 17 GB,减少 92%。借助约 512GB 主机内存,这足以在单张 24GB RTX 4090 上为 8B 模型运行 4M token 推理;4M 上限来自 CPU 内存,而非 GPU。相比之下,标准推理支持 25K token,按层卸载支持 45K。在中等上下文长度下,prefill 吞吐量基本不变:20K 时为 7210 对比 7235 token/秒。
实际局限。 它扩展的是可处理范围,速度并未提高,而且远未达到极限长度时就会付出解码代价:20K 时已是 6 对比 33 token/秒,到 1M token 时为 0.15 token/秒。
开源项目
KVCache-Factory——面向自回归模型的 KV 缓存压缩方法统一实现,也是 PyramidKV 的参考实现。
R-KV——面向推理模型、在解码阶段运行的 KV 缓存压缩。
