Memory-Efficient Looped Transformer: Decoupling Compute from Memory in Looped Language Models

循环式大语言模型（LLM）架构作为一种提升推理能力的新兴方法备受关注，其核心优势在于：无需生成中间文本标记（intermediate tokens），即可在嵌入空间（embedding space）中执行多步计算。以Ouro为代表的此类模型，通过迭代更新内部表征来实现推理，同时在各次迭代中沿用标准的键值（KV）缓存；然而，这种设计导致内存占用随推理深度线性增长。因此，增加推理迭代次数将引发难以承受的内存开销，严重制约了该类架构在实际应用中的可扩展性。 本文提出一种新型架构——内存高效循环Transformer（Memory-Efficient Looped Transformer, MELT），首次实现了推理深度与内存消耗的解耦。与传统方案为每一层、每一次循环单独维护一套KV缓存不同，MELT仅在每一层维护一个共享的KV缓存，该缓存被所有推理循环复用；缓存内容则通过一个可学习的门控机制（learnable gating mechanism）随时间动态更新。为保障该架构下训练过程的稳定性与高效性，我们设计了一种分两阶段的“分块训练”（chunk-wise training）策略：第一阶段采用插值过渡（interpolated transition），第二阶段采用注意力对齐蒸馏（attention-aligned distillation），二者均以LoopLM的初始模型为教师模型，指导MELT的训练。 实验结果表明：基于预训练Ouro参数微调得到的MELT模型，在性能上全面超越同规模的标准大语言模型，同时其内存占用与这些标准模型相当，且显著低于Ouro模型。总体而言，MELT仅需一个轻量级的后训练流程，即可实现常数级内存开销的迭代式推理，且完全不牺牲LoopLM原有的推理性能。

现有循环式大语言模型（如Ouro）通过多步嵌入空间推理提升推理能力，但其每轮迭代均扩展标准KV缓存，导致内存消耗随推理深度线性增长，严重制约实际部署的可扩展性。这是一个在高效迭代推理架构中尚未被充分解决的新工程-理论交叉问题。

MELT提出‘缓存解耦’核心思想：每层仅维护一个跨循环共享的KV缓存（而非每轮新建），并通过可学习门控机制动态更新该缓存；结合两阶段轻量级后训练策略（插值过渡 + 注意力对齐蒸馏），实现从LoopLM/Ouro到MELT的稳定迁移——首次在保持常数级内存复杂度（O(1) per layer）的同时，不损失循环推理性能。

实验基于Ouro预训练权重进行微调，在数学推理（GSM8K）、常识推理（ARC-Challenge）等基准上超越同尺寸标准LLM，内存占用与基线LLM相当、仅为Ouro的~1/5（深度×5时）；采用chunk-wise训练缓解长序列梯度不稳定；暂未开源代码，但方法完全兼容Hugging Face生态；值得深入的方向包括：门控机制的理论收敛性分析、缓存共享对长程依赖建模的边界研究、以及向MoE循环架构的自然扩展。

LoopLM: Reasoning with Language Models via Iterative Refinement (ICML 2024); Ouro: Recurrent LLMs for Step-by-Step Reasoning (NeurIPS 2023); ReAct: Synergizing Reasoning and Acting in Language Models (ICLR 2023); Self-Refine: Iterative Refinement with Self-Feedback (ACL 2024); Memory Transformer: Compressing KV Cache for Long Context (EMNLP 2023)