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

# 循环潜变量注意力：面向循环Transformer的跨循环KV压缩

**来源**：https://arxiv.org/html/2607.15456

**James O’Neill**\texttt{james\.oneill@intercom\.io} Fin AI Research  
**Fergal Reid** Fin AI Research  

###### 摘要

循环、权重共享的Transformer通过重复使用一个块来减少参数，但解码时仍为每个递归步骤存储独立的K/V缓存。我们证明，这种以循环索引的缓存具有高度结构性。对于固定token、层和注意力头，K/V向量在循环间呈现短的低秩轨迹，而头和层轴则平坦得多。我们提出循环潜变量注意力（Looped Latent Attention, LLA），这是一种训练后缓存编解码器，它存储紧凑的K和V潜变量，仅在注意力读取时才重建循环特定的K/V向量。默认的每头编解码器压缩循环维度，而LLA‑2D还将头折叠到单个潜变量中，用于极端压缩场景。该编解码器通过教师激活的SVD初始化，并使用KL散度和注意力输出蒸馏进行精调。在匹配的缓存预算下，每头LLA优于头轴MLA、跨层共享、KV量化和最终循环复用，表明循环缓存是低秩的，但不能安全地折叠为单一状态。相同的轴优势在Ouro‑2.6B‑Thinking上成立，并迁移到Huginn‑3.5B，其中SVD编解码器在解码器无关评估中在32×压缩下仍接近无损。缓存缩减是精确的。在一台H200上，潜变量存储路径在4k上下文下将Ouro‑1.4B的批次容量从32序列增加到768序列，实现了21.3×压缩。对于长数学展开，基于学生生成前缀的在线精调在4×压缩下将MATH‑500从0.43提升到0.66，并减少了无答案生成。

## 1 引言

在推理时，限制Transformer内存的是其KV缓存，而非权重。权重一次性加载并在整个批次中共享。缓存对每个请求私有，并随上下文和每个生成token增长。超过一定上下文长度后，每个请求的状态（而非参数数量）决定了设备能容纳多少个请求以及批量服务是否经济。

循环、权重共享的Transformer通过重复使用一个块来应对这种参数成本：Universal Transformers (Dehghani et al., 2019) 将深度重新表述为递归，最近的循环语言模型如Ouro/LoopLM (Zhu et al., 2025) 将该理念扩展到数十亿参数。然而，缓存是非共享的。每次递归步骤重新计算注意力并写入自己的键和值，因此缓存除了层、头和token外，还以递归步骤索引。具有 *T*=4 个循环和 *D*=24 层的模型因此拥有96层解码器的KV缓存，同时共享24层解码器的参数。

这个缓存携带的自由度远低于其标称大小。对于固定token、层和头，每循环的K/V向量轨迹短且低秩，且冗余特定于循环轴：在相同预算下，头和层谱平坦得多。循环潜变量注意力（LLA）通过将 *T* 个每循环K/V向量替换为单个潜变量和循环特定的重建映射来利用这一点。

真正的问题是压缩哪个轴，以及压缩时缓存必须保留什么。循环解码器引入了一种退化捷径：如果递归只是收敛到一个不动点，可以存储最终循环并在各处重用。它还继承了通常的Transformer选项——压缩头、共享层、量化值。我们的实验将这些区分开来。我们冻结教师，保持注意力模式和递归不变，只改变缓存存储的内容。在这种受控交换下，最终循环复用获得了相同的4×压缩，但使得GSM8K生成崩溃为零，而一个小的低秩轨迹代码在同一预算下保持了性能。

### 贡献点

- • 我们定义了LLA，一种面向权重共享循环Transformer的跨循环潜变量K/V编解码器。它每token每层存储 \((r_k + r_v)H\) 个标量，而非 \(2THd_{\rm head}\) 个，压缩比 \(\rho = 2Td_{\rm head}/(r_k + r_v)\)。LLA‑2D通过将头折叠到单个潜变量中扩展了同一家族，用于极端压缩场景。
- • 我们证明循环是Ouro‑1.4B、Ouro‑2.6B‑Thinking和Huginn‑3.5B中最可压缩的缓存轴。K和V遵循不同的轨迹，因此 \(r_v > r_k\)，并且与头轴MLA、跨层共享、KV量化和最终循环复用在匹配预算下的比较有利于循环轴。
- • 我们将表示结果与部署联系起来：一个明确的在线精调目标使用学生生成的前缀来稳定长展开，精确的标量缩减转化为实际的推理增益，在单台H200上将Ouro‑1.4B的最大批次从32序列提高到768序列，压缩比为21.3×。

## 2 相关工作

### 循环与通用Transformer

通用Transformer (Dehghani et al., 2019) 将权重共享块应用于多个步骤，将深度重新表述为递归。最近的循环语言模型将此扩展到十亿参数的语言模型，并在推理任务中用计算换取参数效率 (Zhu et al., 2025)。对这些模型的机制研究报道，每次循环的计算会沉淀到低维轨迹上 (Blayney et al., 2026)，而LLA正是基于这种结构构建缓存编解码器。递归并未解决服务成本，因为在注意力机制中，每个循环仍存储独立的K/V状态。

### 潜变量和压缩KV缓存

DeepSeek‑V2中引入的多头潜变量注意力（MLA）通过跨头/通道轴存储低秩潜变量并重建每头键和值，压缩了标准Transformer的KV缓存 (DeepSeek‑AI, 2024)。跨层共享针对层轴 (Brandon et al., 2024; Sun et al., 2024; Liu et al., 2024)，量化针对精度 (Hooper et al., 2024)，token逐出针对序列轴 (Zhang et al., 2023; Xiao et al., 2024)。所有这些都作用于单次Transformer。而LLA则压缩循环，这是权重共享循环模型独有的轴。由于相同的块和投影在每一步被复用，循环轴的冗余性具有头和层轴所缺乏的原因。头轴（由多头查询和分组查询注意力所针对 (Shazeer, 2019; Ainslie et al., 2023)）在LLA之上组合（LLA‑2D），序列轴逐出则与其相乘（第7节）。

### 内存高效循环Transformer

MELT (Vendrell et al., 2026) 是最接近的并行工作。它蒸馏每层单个门控K/V状态，而LLA保留循环轨迹并压缩其秩。这种区别并非表面上的。我们的零参数最终循环复用对照（将循环坍缩为一个状态）在相同缓存预算下失败了。两种方法也覆盖了不同的区间。坍缩状态将内存缩减固定在大约循环数附近，而可调潜变量可以通过选择 \(r_k\) 和 \(r_v\) 位于该点之上或之下。

## 3 循环潜变量注意力

循环模型节省参数，但仍在每个递归步骤实现化K/V状态。我们用一种编解码器压缩该递归轴，该编解码器改装到冻结教师上，除了内存中表示的存储外不改变注意力计算，因此实验可以隔离循环索引的K/V堆栈是否为低秩。我们首先定义未压缩的递归缓存，然后介绍跨循环潜变量及其精确压缩比。相同的构造用于K和V，但秩不同。然后我们描述SVD初始化和用于测试循环轴是否为正确压缩轴的匹配基线。

> **图1：循环潜变量注意力。**
> (a) LLA位于权重共享循环Transformer块的KV缓存路径上：块在每次递归步骤输出每循环的K、V，编解码器仅压缩该缓存的循环轴，保持注意力模式和块的其余部分不变。
> (b) LLA‑2D（循环×头），家族变体，将循环和头一起分解为每token每层的单个潜变量，从单个向量以MLA风格解码所有 \(H\) 个头。默认的每头LLA保留 \(H\) 个独立的每头潜变量，仅压缩（更冗余的）循环轴。
> (c) 头轴基线（每循环MLA）在每次递归步骤独立压缩头，没有跨循环共享。LLA则通过每循环上投影 \(W_t^\uparrow\) 从单个共享跨循环潜变量重建每个循环。

### 教师缓存

考虑一个权重共享循环Transformer，具有 \(D\) 层、\(H\) 个宽度为 \(d_{\rm head}\) 的注意力头和 \(T\) 个递归步骤。对于一个层、一个头、一个token和一个轴（K或V），教师输出 \(T\) 个向量。我们将其堆叠为

\[
x = [x_1; \ldots; x_T] \in \mathbb{R}^{T d_{\rm head}}. \tag{1}
\]

未压缩的缓存存储每个循环、头和层的K和V，每token每层成本为 \(2THd_{\rm head}\) 个标量。

### 跨循环编解码器

对于每个层、头和轴，LLA存储一个中心化的低秩潜变量

\[
c = (x - \mu) W_{\rm down},\qquad c \in \mathbb{R}^r, \tag{2}
\]

其中 \(\mu \in \mathbb{R}^{T d_{\rm head}}\) 是校准集的均值，\(W_{\rm down} \in \mathbb{R}^{T d_{\rm head} \times r}\)。循环 \(t\) 通过循环特定的上投影重建：

\[
\hat{x}_t = c W_{{\rm up},t}^\top + \mu_t,\qquad W_{{\rm up},t} \in \mathbb{R}^{d_{\rm head} \times r}. \tag{3}
\]

K和V有独立的潜变量和秩 \(r_k\) 和 \(r_v\)，且 \(r_v > r_k\)，因为值缓存在循环间比键缓存更慢地稳定。压缩后的缓存每token每层存储 \((r_k + r_v)H\) 个激活标量，因此

\[
\rho = \frac{2THd_{\rm head}}{(r_k + r_v)H} = \frac{2Td_{\rm head}}{r_k + r_v}. \tag{4}
\]

该比率计算的是每请求缓存。固定的编解码器矩阵在不同token、上下文和批次间共享，不随序列长度扩展。

### RoPE与重建

潜变量适用于RoPE之前的内容投影。在循环解码器中，一个token在每次递归步骤具有相同的位置，因此旋转相位 (Su et al., 2021) 在循环间共享，可以在重建后应用。这比标准的MLA更简单，后者需要显式的解耦RoPE分支以沿头/通道轴保留位置信息。

### LLA‑2D变体（循环×头）

公式(2)的每头编解码器可以通过将头轴折叠到同一潜变量中进一步推进。LLA‑2D每token每层存储一个潜变量，并从中MLA风格地重建 **所有** \(H\) 个头的K/V，此外还加上公式(3)的每循环重建。这使缓存降至每token每层 \((r_k + r_v)\) 个标量，比每头LLA小一个因子 \(H\)，代价是将头通过一个共享代码耦合在一起。LLA‑2D是该家族中极端压缩的成员。其下投影输入是整个 \(THd_{\rm head}\) 堆栈，因此其参数数量随头数和递归长度增长。这在大的 \(T\) 时成为绑定约束（Huginn迁移，第5节）。第5节刻画了两个成员之间的精度/压缩权衡。不加限定的LLA指每头变体。

### 初始化与转换

对于每个层、头和轴，我们在校准集上收集教师K/V激活，通过每循环均值中心化，并从堆叠循环矩阵的前 \(r\) 个右奇异向量初始化 \(W_{\rm down}\)。初始上投影是对应的循环切片。然后使用冻结教师对SVD初始化进行微调。目标函数是从教师分布到交换模型的KL散度加上注意力输出匹配项：

\[
\mathcal{L} = \mathrm{KL}\left(p_{\rm teacher} \| p_{\rm swap}\right) + \frac{\lambda_{\rm attn}}{TD} \sum_{t,d} \|a_{t,d} - \hat{a}_{t,d}\|_2^2, \tag{5}
\]

在报告的实验中 \(\lambda_{\rm attn}=0.5\)。仅更新编解码器参数。SVD先验比其他任何设计选择都更重要。在4×压缩下，SVD初始化在短时间转换后达到保留集KL 0.105，而随机初始化在相同预算下仍为0.771（表10，图10）。

### 在线精调

公式(5)是教师强制的，在教师访问的前缀上拟合编解码器。对于长自回归展开，交换模型则基于由其自身编解码器产生的前缀进行条件判断。因此我们在离策略拟合之后添加一个可选的第二阶段。交换模型采样补全，冻结教师对这些采样前缀进行评分，仅更新编解码器参数。不使用奖励或正确性过滤。第6节并排写出了离策略和在线策略的目标。在线策略版本将教师前缀替换为学生生成的前缀，并停止通过采样的梯度。

### 匹配轴基线

主要消融实验在相同标量预算下将可比的潜变量编解码器置于其他缓存轴上，测试循环轴是否正确。头轴基线是每循环MLA（无跨循环共享）。层基线跨相邻层复用K/V。精度基线是低位KV量化。我们还包含一个零参数对照，它只缓存最终循环的K/V并用于每个循环。该对照在 \(T=4\) 时与LLA具有相同的4×缓存占用，直接测试循环轨迹是否可以被其终点替代。

## 4 教师缓存中的跨循环结构

在比较训练过的编解码器之前，我们先问冻结教师是否已经包含LLA假设的结构。这是比下游精度更严格的测试，因为它直接检查缓存张量，无需优化、解码启发式或任务特定效应。如果循环轴在谱上不区别于头或层轴，那么跨循环编解码器将是一个任意设计选择。如果教师在循环轴上秩更低，那么后续的等缓存排序就有了机制解释。

> **图2：跨循环K/V是低秩的，因为递归使...**