🤖 machine learning

Looped Latent Attention: Cross-Loop KV Compression for Looped Transformers

本文介绍了循环潜变量注意力（Looped Latent Attention, LLA），这是一种后训练编解码器，它利用了权重共享 Transformer 中循环索引 KV 缓存的低秩结构，通过 SVD 初始化和蒸馏技术，在保持近乎无损性能的同时，实现了极高的压缩率（高达 32 倍）并显著提升了批处理容量。

原作者： James O' Neill, Fergal Reid

发布于 2026-07-20

📖 1 分钟阅读☕ 轻松阅读

CC BY 4.0

原作者： James O' Neill, Fergal Reid

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

AI 大脑中的记忆问题

想象一下，你正在试图教一个机器人写故事。这个机器人的大脑由数百万个微小的开关（参数）组成，这些开关知道词语是如何组合在一起的。但问题在于：为了写一个长篇故事，机器人不能只靠看自己的大脑；它需要一个“草稿本”。每当它写下一个新词时，它都必须记住它刚刚写下的所有内容，以免重复自己或跑题。在人工智能的世界里，这个草稿本被称为 KV 缓存（KV cache）。它是为每一次对话准备的私人笔记本，并且随着机器人每说出一个词，它就会变得越来越大。

现在，想象一种设计得极其高效的特殊机器人。这个机器人没有拥有许多不同房间（层）来处理不同任务的大型大脑，而是只有一个聪明的小房间，它会反复多次地访问这个房间。它重复使用同一套指令，通过循环执行这些指令来进行更深层的思考。这被称为循环 Transformer（Looped Transformer）。这节省了大脑本身的空间，这很棒。但有一个隐蔽的问题：尽管机器人重复使用同一个房间，但它每次循环时仍然会写一张新的草稿本页面。如果机器人循环四次来解决一个数学问题，它最终会为故事中的同一个时刻产生四页独立的笔记。这意味着，“循环”机器人仍然需要大量的内存来运行，这抵消了其高效设计的初衷。科学家们面临的重大问题是：我们能否在不让机器人忘记如何思考的前提下，缩小那个草稿本？

论文的发现：笔记中的秘密捷径

这篇论文引入了一种被称为**循环潜变量注意力（Looped Latent Attention, LLA）**的聪明技巧，来解决那个记忆问题。研究人员发现，机器人的笔记并不像看起来那么混乱。当机器人循环进行思考过程时，它为同一个词编写的笔记并不会发生剧烈的变化；它们遵循一条平滑、可预测的路径，就像一颗在大地缓坡上滚动的弹珠。LLA 不再保存机器人草稿本的每一页，而是仅保存那条路径的一个微小的“摘要”，以及一套只有在机器人真正需要阅读时才如何重建特定页面的指令。

可以这样理解：想象你正在看一部电影，主角正走在一条走廊里。在普通的计算机中，你会记录角色每走一步时的全高清照片。这占用了大量的存储空间。但有了 LLA，你意识到角色只是在走直线。所以，与其保存数千张照片，不如保存一张照片并附上一条简单的笔记，比如：“第 2 步向前移动 1 英寸，第 3 步向前移动 2 英寸。”当计算机需要查看第 3 步时的角色时，它可以使用这条笔记快速画出那张图。这正是论文中所做的工作：它压缩了记忆中的“循环”部分，将四页纸变成了一个单一的、微小的摘要，这个摘要可以瞬间被扩展回完整的页面。

研究人员在多个 AI 模型上测试了这种方法，其中包括一个名为 Ouro-1.4B 的模型。他们发现这种方法效果极佳。事实上，在名为 H200 的强大计算芯片上，他们能够同时在内存中容纳 768 场不同的对话，而以前只能容纳 32 场。这在支持更多人同时使用 AI 方面是一个巨大的飞跃。他们还展示了即使在 AI 解决非常长的、困难的数学问题时，这个技巧依然有效，尽管它需要额外的训练，以确保它在自己编写故事而不是仅仅模仿老师时不会感到困惑。

这篇论文对什么说“不”

了解这种方法不是什么非常重要。研究人员首先测试了一个非常简单的想法：“如果我们直接扔掉前三页笔记，只保留最后一页会怎样？”他们认为也许机器人在经过几次循环后会稳定下来并得出最终答案，因此早期的笔记并不重要。他们尝试了这样做，结果却是一场灾难。AI 完全丧失了做数学的能力，在测试中得分竟然为零。这证明了笔记所走的“路径”是至关重要的；你不能只抓取终点，你需要捕捉整个旅程。

他们还将这种方法与其他缩减内存的方法进行了比较，例如将笔记挤压在一起（量化/quantization），或者在不同的大脑部分之间共享笔记（头轴压缩/head-axis compression）。虽然这些方法也有所帮助，但论文表明，压缩“循环”轴才是真正的秘诀。如果你尝试压缩其他部分而不是循环部分，AI 就会失去推理能力。论文明确指出：循环是记忆中最冗余的部分，也是唯一应该进行压缩的地方。

他们的结论有多可靠？

作者对这些结果非常有信心，因为他们是直接进行测量的。他们不仅仅是在猜测，而是让 AI 在解决数学问题（GSM8K）和编写代码（HumanEval）等真实任务上运行。他们展示了通过 4 倍压缩（使内存缩小 4 倍），AI 的表现几乎与原始未压缩版本一样出色。当他们将其推向 21.3 倍的极端压缩时，AI 仍能处理许多任务，尽管在面对最长、最复杂的推理步骤时开始显得吃力。

他们还在另一种不同类型的 AI 模型 Huginn-3.5B 上测试了该方法，该模型具有不同的脑结构，并且会循环多达 32 次。即便如此，该方法依然奏效，证明了这不仅仅是针对某一个特定机器人的偶然现象。论文指出，这种“低秩轨迹”（notes 的平滑路径）是这些循环型机器人思考方式的一个基本属性。虽然他们已经证明了这在节省内存和增加容量方面的有效性，但他们也指出，为了在非常长的任务上获得最佳表现，AI 需要进行第二轮训练，即在它使用自己的压缩笔记进行练习，而不是仅仅模仿老师。这种“策略内”（on-policy）训练修复了当机器人开始编写自己的故事时产生的微小误差。

简而言之，这篇论文表明，通过意识到其记忆具有特定的重复性，我们可以让循环型 AI 模型变得更加高效。通过存储这种重复性的微小摘要而非完整的笔记，我们可以在单个计算机芯片上容纳数百倍的对话，从而在不需要超昂贵硬件的情况下，让更多人能够使用强大的 AI。

技术摘要：循环潜变量注意力 (Looped Latent Attention, LLA)

问题陈述

在循环权重共享 Transformer (Looped, weight-tied Transformers) 中，通过重复应用单个块（循环）来减少参数量。虽然这有效地压缩了模型的权重，但推理阶段的内存瓶颈仍然在于 键/值 (KV) 缓存。在标准的循环推理中，模型在每个循环步骤都会重新计算注意力，并为每个步骤、层和头存储独立的 KV 缓存。因此，一个具有 $T$ 次循环和 $D$ 个层的模型，其携带的 KV 缓存相当于一个 $T \times D$ 层的解码器，这抵消了权重共享带来的内存收益。

论文指出，这种以循环为索引的缓存具有高度的结构性。对于给定的 token、层和头，K/V 向量在不同循环步骤之间的变化轨迹是短促且低秩的，而非探索高维空间。这种冗余性是针对循环轴 (loop axis) 特有的；而在头轴和层轴上的频谱则保持显著平坦（高秩）。挑战在于如何压缩这个循环轴，同时避免将轨迹坍缩为一个固定的点，否则会导致性能下降。

方法论：循环潜变量注意力 (LLA)

作者引入了 循环潜变量注意力 (LLA)，这是一种后训练 (post-training) 缓存编解码器，它用紧凑的潜变量表示和特定于循环的重构映射，取代了 $T$ 个逐循环的 K/V 向量。

核心机制

潜变量存储： LLA 不再为特定的 token、层和头存储 $T$ $T$ 个向量 ( $x_1, \dots, x_T$ $x_{1}, \dots, x_{T}$ )，而是存储一个单一的中心化低秩潜变量 $c \in \mathbb{R}^r$ $c \in R^{r}$ 。
- $c = (x - \mu)W_{\text{down}}$
- 这里， $\mu$ 是校准集的均值， $W_{\text{down}}$ 是下投影矩阵。
特定于循环的重构： 循环步骤 $t$ $t$ 的特定 K/V 向量通过一个特定于循环的上投影进行重构：
- $\hat{x}_t = c W_{\text{up},t}^\top + \mu_t$
非对称秩： 该方法对键 (Keys) 和值 (Values) 使用不同的秩 ( $r_k$ 和 $r_v$ )。经验分析表明，Value 向量在循环过程中的稳定速度比 Key 向量慢，因此促使 $r_v > r_k$ 。
LLA-2D 变体： 为了实现极端压缩，LLA-2D 将头轴 (head axis) 折叠进潜变量中，存储每个 token 和层对应的单一潜变量，以重构所有 $H$ 个头。这使缓存大小减少了 $H$ 倍，但也导致了头之间的耦合。

训练与初始化

SVD 初始化： 编解码器通过对校准集上教师模型激活值的奇异值分解 (SVD) 进行初始化。下投影矩阵被设置为前 $r$ 个右奇异向量。这种初始化至关重要；随机初始化会导致显著更高的 KL 散度。
蒸馏： 通过冻结教师模型，并最小化教师分布与替换后的模型分布之间的前向 KL 散度（辅以注意力输出匹配项）来优化编解码器。
策略内优化 (On-Policy Refinement)： 为了解决长自回归展开中的暴露偏差问题（即模型基于自身生成的序列前缀而非教师的序列进行条件化），采用了第二阶段的策略内蒸馏。编解码器使用从学生模型自身采样的前缀进行更新，从而稳定长文本生成性能。

核心贡献

定义 LLA： 一种针对权重共享 Transformer 的跨循环潜变量 K/V 编解码器，其压缩率 $\rho = \frac{2Td_{\text{head}}}{r_k + r_v}$ 。它为每个 token 和每层存储 $(r_k + r_v)H$ 个标量，而不是 $2Td_{\text{head}}$ 个。
轴向特定压缩： 论文证明了在循环模型（如 Ouro-1.4B, Ouro-2.6B-Thinking, 和 Huginn-3.5B）中，循环轴是最可压缩的维度。压缩循环轴的效果优于压缩头轴 (MLA)、层轴（跨层共享）或在匹配缓存预算下应用量化。
拒绝固定点坍缩： 作者展示了循环轨迹不能被单一状态（最终循环）所取代。使用仅重用最终循环 K/V 的零参数控制方法会导致生成性能崩溃（GSM8K 准确率降至 0.000），这证明了“短促低秩轨迹”的必要性。
部署影响： 确切的标量减少直接转化为更高的服务能力。在单张 H200 GPU 上，LLA 使 Ouro-1.4B 在 4k 上下文下的最大 Batch Size 从 32 提升至 768 个序列（实现了 21.3 倍压缩）。

实验结果

压缩 vs. 准确率

GSM8K 性能： 在 Ouro-1.4B 上，轻度压缩（1.33× 至 2×）可以匹配教师模型的准确率（约 0.79–0.80）。在 4× 压缩时，LLA 维持了 0.59 的准确率，显著优于头轴压缩 (0.52) 和最终循环重用 (0.00)。
规模泛化性： 在 Ouro-2.6B-Thinking 上，LLA 在 4× 到 10.67× 压缩下仍保持接近无损的水平 (0.81–0.83)，而头轴基准模型则出现了大幅度的准确率下降。
迁移至 Huginn： 该方法可迁移至 Huginn-3.5B（32 次循环）。无需训练的 SVD 编解码器在解码器无关指标上达到了近乎无损的保真度，最高可达 32× 压缩。

长文本生成稳定性

MATH-500： 长数学推导过程对前缀漂移非常敏感。使用教师前缀的离策 (Off-policy) 训练会导致随着生成长度增加，准确率下降且“无答案”率升高。
策略内优化的增益： 应用策略内优化将 4× 压缩下的 MATH-500 准确率从 0.43 提升至 0.66，并显著减少了无答案生成的现象。

服务能力

内存减少： 缓存减少是确切的。在 21.3× 压缩下，Ouro-1.4B 在 1M 上下文下的 KV 缓存从 768 GiB 降至 36 GiB。
吞吐量： “潜变量存储路径”（仅缓存潜变量）是计算密集型的，但它使得原本因显存溢出 (OOM) 而无法实现的 Batch Size 和上下文长度成为可能。“快速路径”（重构并缓存完整的 K/V）虽然保留了吞吐量，但并未实现内存节省。

意义与主张

论文声称，在权重共享的循环 Transformer 中，循环是压缩性最强的缓存轴。其意义在于：

机制洞察： 它识别出循环轴在频谱上与头轴和层轴截然不同（更低秩），从而允许高效压缩，同时不会将模型的动态行为坍缩为一个静态的固定点。
实际部署： 它提供了一种后训练解决方案，可以在不重新训练基础模型架构的情况下，直接提高现有硬件上的 Batch Size 和上下文长度能力。
未来方向： 结果表明，原生循环架构可以受益于解耦的潜变量缓存，即循环计算的是共享内容加上廉价的逐循环增量，而非重新实例化完整的 K/V 状态。

作者强调，其准确性主张并不依赖于吸收式注意力 (absorbed attention) 的实现；其核心结论是“低秩轨迹假设”的有效性，以及编解码器在严格内存约束下保持教师模型行为的能力。