通过低秩潜变量压缩循环注意力缓存,显著提升大模型推理效率。
Looped Latent Attention: Cross-Loop KV Compression for Looped Transformers

- 用潜变量重构循环中的键值对,实现缓存的高效压缩。
- 在相同缓存预算下,性能超越主流压缩方法,最高压缩比达32倍。
- 适用于长推理任务,可提升数学推理准确率并减少无效输出。
循环共享权重的Transformer通过复用单个块减少参数,但解码时仍需为每个循环步骤单独存储键值(K/V)缓存。我们发现该循环索引缓存具有高度结构化特性:固定标记、层和头的K/V向量在循环中呈现短低秩轨迹,而头与层维度变化较平缓。为此提出环形潜在注意力(Looped Latent Attention, \\lla{}),一种后训练缓存编码器,仅存储紧凑的K/V潜变量,并在注意力读取时重建特定循环的K/V。默认每头编码器压缩循环信息,\\lla{}-2D进一步将头维度合并为单一潜变量,实现极致压缩。编码器基于教师激活的SVD初始化,并通过逻辑值与注意力输出蒸馏优化。在相同缓存预算下,每头\\lla{}优于头轴多级注意力(MLA)、跨层共享、KV量化及最终循环重用,表明循环缓存虽低秩,但不可安全坍缩为单一状态。该轴向优势在Ouro-2.6B-Thinking上成立,并迁移至Huginn-3.5B,在独立解码评估中,SVD编码器在32倍压缩下几乎无损。缓存压缩为精确缩减。在单张H200上,潜变量路径使Ouro-1.4B在4k上下文下的批量容量从32提升至768,压缩比达21.3倍。对于如MATH-500的长推理轨迹,学生生成前缀的在线优化使4倍压缩下的准确率从0.43提升至0.66,且相比逐标记离策略蒸馏,显著降低无答案生成。
原文摘要 · Abstract (English)
Looped, weight-tied Transformers reduce parameters by reusing a single block, but decoding still stores a separate K/V cache for every recurrence step. We show that this loop-indexed cache is highly structured. For a fixed token, layer and head, K/V vectors trace a short low-rank trajectory across loops, while the head and layer axes remain much flatter. We introduce Looped Latent Attention (\lla{}), a post-training cache codec that stores compact K and V latents and reconstructs loop-specific K/V vectors only when attention reads them. The default per-head codec compresses recurrence, while \lla{}-2D also folds heads into one latent for the extreme-compression regime. The codec is initialized from the SVD of teacher activations and refined with logit and attention-output distillation. At matched cache budget, per-head \lla{} outperforms head-axis MLA, cross-layer sharing, KV quantization and final-loop reuse, showing that the recurrent cache is low-rank but not safely collapsible to a single state. The same axis advantage holds on Ouro-2.6B-Thinking and transfers to Huginn-3.5B, where an SVD codec remains near-lossless to $32\times$ compression in decoder-independent evaluation. The cache reduction is exact. On one H200, the latent-store path increases measured Ouro-1.4B batch capacity at 4k context from 32 to 768 sequences at $21.3\times$ compression. Lastly, for long reasoning rollouts such as in MATH-500, on-policy refinement on student-generated prefixes raises accuracy at $4\times$ compression from 0.43 to 0.66 and reduces no-answer generations when compared to token-level off-policy distillation.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。