🤖 machine learning

Regularize or Localize: When Training-Time KV-Cache Geometry Pays Under Quantization

本文表明，虽然仅将 LeJEPA 抗崩溃目标应用于隐藏状态无法改善 KV 缓存量化，但在训练期间直接对 KV 缓存几何结构进行正则化能显著提升在粗粒度逐通道量化方案下的性能，尽管这种优势在采用如 KIVI 等更复杂的量化配置时会减弱。

原作者： Libo Sun, Po-Wei Harn, Zewei Zhang, Peixiong He, Xiao Qin

发布于 2026-07-21

📖 1 分钟阅读☕ 轻松阅读

原作者： Libo Sun, Po-Wei Harn, Zewei Zhang, Peixiong He, Xiao Qin

原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

想象一下，你正试图教一个机器人写故事。为了做到这一点，你给了它一个庞大的图书库，让它通过一遍又一遍地预测句子中的下一个词来学习。这就是现代 AI 语言模型的工作方式。但有一个问题：随着机器人的阅读，它会建立起对目前故事内容的“记忆”，这个记忆被称为 KV-cache（键值缓存）。把这个缓存想象成机器人桌上堆叠的便签纸。每当机器人写下一个新词时，它就会增加一张便签。如果机器人需要记住一个很长的故事，这叠便签就会变得非常庞大，存储它们会占用大量的计算机内存。

为了节省空间，工程师们尝试通过一种叫做**量化（quantization）**的过程来缩小这些便签。这就像是将一张高清晰度的照片压缩成一张微小的缩略图。通常情况下，这样做效果很好，但有时机器人的记忆中的“便签”是以一种奇怪、歪斜的方式书写的。它们全都向同一个方向聚集，就像一堆铅笔都向同一个方向倾斜一样。这使得它们很难在不丢失重要细节的情况下进行压缩。科学家们面临的大问题是：我们能否从一开始就教机器人以一种更平衡、更“圆润”的方式来书写它的便签，从而让它们在稍后更容易被缩小？这篇论文正是通过研究一种特殊的训练技巧，来调查是否能改变机器人记忆的形状，并帮助它装进更小的盒子里。

实验：塑造机器人的记忆

论文中的研究人员决定测试一种特定的训练技巧，叫做 SIGReg。想象一下，机器人的大脑是一个巨大的、多层结构的工厂。通常，这个工厂只致力于得到正确答案（预测下一个词）。SIGReg 增加了第二条规则：“嘿，确保你的思想形状不要全都向同一个方向倾斜。”它推动机器人的内部表示变得更像一个完美的、圆形的云团，而不是一个扁平且歪斜的煎饼。

团队训练了一个中等规模的机器人（1.1 亿参数），使用了 100 亿词汇的海量数据集。他们测试了三个主要想法，扮演着像侦探一样试图解开记忆形状来源之谜的角色。

1. 这个技巧对机器人的“思想”有效吗？
是的，有效。当他们将 SIGReg 应用于机器人的隐藏思想（内部处理步骤）时，其“歪斜度”下降了 38%。机器人的写作能力并没有变差太多；它的困惑度（perplexity）上升不到 0.35%，这几乎可以忽略不计。因此，这个技巧成功地重塑了机器人的内部几何结构，而没有破坏它的表达能力。

2. 这个技巧能修复“便签”（KV-cache）吗？
剧情在这里发生了转折。研究人员发现，仅仅修复机器人的“思想”并不能自动修复“便签”（缓存）。缓存依然保持着歪斜的状态，和以前一样。

失败的修复方案： 他们尝试通过冻结机器人的大脑并仅仅添加新的“头”来读取便签进行“补救”。但这并没有奏效。
真正的修复方案： 要想真正改变便签的形状，他们必须保持机器人的大脑处于解锁状态，并在机器人继续学习的同时，将 SIGReg 技巧直接应用于便签本身。当他们这样做时，缓存的歪斜度大幅下降了 94%。
教训： 你不能仅仅修复大脑并指望记忆随之改变。你必须直接训练记忆。

3. 圆润的缓存是否能做出更好的缩略图（量化）？
这是最令人兴奋的部分。团队尝试将便签缩小成 3-bit 的微小缩略图（一个非常小的尺寸），以观察故事质量会受到多大的影响。

好消息： 对于使用直接缓存技巧训练的机器人，缩小便签带来的“损伤”比普通机器人小了 4.3 到 7.9 倍。事实上，经过训练的机器人是唯一一个更偏好特定测量方式（逐通道缩放/per-channel scaling）的模型，而这种方式在压缩方面表现得更好。
坏消息（代价）： 这种巨大的优势仅出现在他们使用一种简单的、“粗粒度”的缩小方式时。当他们使用一种更先进、更真实的设置（称为“KIVI 式”配置，它混合了不同的缩放方法并加入了零点）时，这种优势消失了。在这种复杂的场景下，经过训练的机器人和普通的机器人表现得几乎完全一样。

结论

那么，他们学到了什么？论文表明，你可以通过在训练期间将特定的正则化技巧直接应用于 KV-cache，来训练语言模型拥有一个更“圆润”且平衡的记忆。如果使用简单的、粗粒度的压缩方法，这会让记忆更容易被压缩。

然而，论文非常谨慎地指出，这并不是解决所有压缩问题的灵丹妙药。如果你使用的是现实世界系统中使用的那些复杂的、现代化的压缩技术（带有混合缩放和零点），这种训练技巧带来的益处就会消失。在这些复杂的场景中，记忆的“圆润度”并没有起到作用。

简而言之：研究人员找到了一种重塑机器人记忆的方法，并且在处理简单的缩小任务时效果极佳。但对于现实生活中使用的那些高级、高科技的缩小工作，额外的训练并没有带来额外的速度或空间节省。这种优势是真实存在的，但它有一个非常明确的边界：它只在压缩尺度是“粗粒度”时才有效，并且无法在现代、混合配置的量化复杂性中幸存下来。

技术摘要：正则化还是局部化：训练时 K/V 缓存几何结构在量化中的价值

问题陈述

使用交叉熵（CE）训练的 Transformer 语言模型通常会产生各向异性的隐藏状态几何结构，即表示向量会集中在少数几个共享方向上。虽然联合嵌入预测架构（JEPA）引入了如 SIGReg（一种描绘式 Epps–Pulley 正态性检验）之类的目标函数，旨在将嵌入推向各向同性（ $N(0, I)$ ），但以往将 SIGReg 应用于大语言模型（LLM）的尝试均未取得显著效果。这些失败发生在将 SIGReg 作为非标准流的辅助项，或仅在微调期间应用于最终层状态时。

一个关键的空白在于：SIGReg 是否能重塑标准自回归（AR）预训练的几何结构，更重要的是，这种重塑是否能在事后量化过程中使 键/值（K/V）缓存 获益？现有工作主要关注事后压缩或架构变更，而对于“训练时分布正则化”能否改善标准 K/V 缓存几何结构以提升量化鲁棒性的问题，仍很大程度上未被探索。

研究方法

作者采用三阶段实验设计，使用在 10B FineWeb Token 上预训练的 110M 参数 Llama 风格模型 进行研究。

第一阶段：剂量-反应与隐藏状态特征表征

目标： 确定在预训练期间将 SIGReg 添加到 CE 损失函数中是否能降低隐藏状态的各向异性。
设置： 作者对正则化权重 $\lambda$ 进行扫描（$0, 0.001, 0.01, 0.1, 1.0$）。
指标： 成对余弦各向异性（Pairwise-cosine anisotropy）和困惑度（Perplexity）。
发现： 确定 $\lambda=0.01$ 为最佳运行点，该值在显著降低各向异性的同时，对困惑度的影响极小。

第二阶段：桥接探测（适配 vs. 改造）

为了确定几何结构的变化如何传播到 K/V 缓存以及如何实现，作者从 $\lambda=0.01$ 的检查点开始测试了三种特定的探测方式：

A1 (适配/Adaptation)： 解冻主干网络，并在所有层中直接对隐藏状态 ( $h$ ) 和 K/V 投影 ( $K_\ell, V_\ell$ ) 应用 SIGReg 进行持续训练。
A2 (冻结预测/Frozen Prediction)： 在冻结的主干网络上训练逐层预测器，以观察 K/V 几何结构是否可以仅从隐藏状态中恢复。
A3 (冻结解码/Frozen Decoding)： 训练解码器以从冻结主干网络的最终隐藏状态中重建 K/V。
对照组 (Control)： 一个“匹配延续”（matched continuation）实验，其中主干网络在相同的 Token 预算下继续训练，但不添加 K/V 正则化项（ $\lambda_{KV}=0$ ）。

第三阶段：量化评估

作者评估了由此产生的模型在 事后伪量化（量化 $\to$ 反量化）下的 K/V 缓存表现。

测试方案： 对称无组量化（每 Token 或每通道）、非对称网格（带有零点）以及混合排列（KIVI 风格：每通道 K，每 Token V）。
指标： 相对于模型自身全精度参考值的 $\Delta$ NLL（负对数似然）。

核心结果

1. 隐藏状态几何结构

在 $\lambda=0.01$ 时，SIGReg 使三个配对种子中的隐藏状态成对余弦各向异性降低了 38%。困惑度增加不足 0.35%，且零样本能力（如 LAMBADA、HellaSwag 等）保持完好，没有出现一致性的退化迹象。

2. K/V 缓存几何结构与“改造失败”

无自动传播： 隐藏状态中各向异性的降低并未自动传播到 K/V 缓存。冻结主干探测（A2 和 A3）无法恢复几何结构，这表明 K/V 各向异性不仅仅是最终隐藏状态的函数。
需要直接正则化： 只有 A1（通过直接 K/V 正则化进行持续训练）成功降低了缓存的各向异性，实现了平均缓存各向异性 94% 的降幅。
机制隔离： 未带 K/V 项的匹配延续实验（ $\lambda_{KV}=0$ ）未能复现此效应，从而证明了 K/V 项是实现该效果的必要原因。

3. 量化收益

正则化几何结构的收益高度依赖于量化方案：

粗粒度尺度（无组）： 在未经转换的对称无组量化下，A1 适配模型是唯一倾向于使用**每通道缩放（per-channel scaling）**而非每 Token 缩放的模型。在 3-bit 对称每通道量化下，基准模型的 $\Delta$ NLL 损伤比 A1 模型高出 4.3–7.9 倍。
细粒度尺度（KIVI 风格）： 在完整的模拟 KIVI 风格配置（混合排列、零点和分组缩放 $g=32$ ）下，优势消失了。包括基准模型在内的所有模型都达到了近乎持平的性能。
存储开销： 即使考虑到分组零点带来的元数据开销（使得 3-bit 分组存储在容量上与 4-bit 无组存储相当），在 Token 本地分组和混合缩放配置下，训练时的优势也会消失。

意义与主张

本文声称是首个评估针对标准 K/V 缓存几何结构的训练时分布正则化对事后缓存量化影响的研究。

几何控制： 作者证明了可以通过训练时的直接正则化来显式控制 K/V 缓存的几何结构，但这种效应并不会仅通过隐藏状态正则化而产生，也无法通过冻结主干探测进行“改造”。
条件效用： 这种训练干预仅在量化缩放尺度较粗（跨越长 Token 范围）时才会产生切实的收益。其优势特定于正则化几何结构与量化方案之间的相互作用；它并不适用于现代的细粒度配置（如带有零点和本地分组的 KIVI 风格）。
指标解释： 本文澄清了成对余弦各向异性指标测量的是平均方向对齐（由共享均值主导），而非协方差白化。结果显示了二者的脱节：虽然均值对齐度下降，但谱属性（如 Top-1 占比）在 K 层和 V 层表现不同，这表明正则化主要作用于均值/异常方向。

局限性

作者明确界定了其研究范围：

范畴： 结果仅限于 110M 参数和 FineWeb 语料库；尚未测试具有不同缓存布局（如 GQA/MQA）的更大规模模型。
模拟： 量化是通过模拟（伪量化）进行的，未进行实际的运行时间或内存测量。
探测范围： 针对冻结探测的负面结果是基于特定的容量和预算得出的；中间地带（例如 LoRA 规模的适配）仍未经过测试。
并非“首次各向同性”主张： 作者并未声称自己是第一个实现各向同性的研究者（引用了 CosReg, SimCTG 等），而是第一个专门针对 K/V 缓存几何结构进行量化鲁棒性研究的研究。