让大模型推理不占内存,循环计算也能高效运行
Memory-Efficient Looped Transformer: Decoupling Compute from Memory in Looped Language Models

- 用共享缓存替代重复存储,让推理深度与内存解耦
- 相同规模下比Ouro内存降低90%以上,性能更优
- 仅需轻量后训练,适合部署在资源受限设备
循环语言模型通过多步迭代更新内部表示实现推理,但传统方法在每轮迭代中保留标准键值缓存(KV cache),导致内存随推理深度线性增长。本文提出内存高效的循环变换器(MELT),每个层仅维护一个共享的KV缓存,并通过可学习门控机制随时间更新。为支持稳定训练,采用两阶段微调:插值过渡+注意力对齐蒸馏,从原始的LoopLM模型迁移而来。实验表明,基于预训练Ouro参数微调的MELT模型,在同等规模下性能优于标准LLM,内存占用接近标准模型,远低于Ouro。MELT实现了恒定内存的迭代推理,无需牺牲性能,仅需轻量级后训练过程。
原文摘要 · Abstract (English)
Recurrent LLM architectures have emerged as a promising approach for improving reasoning, as they enable multi-step computation in the embedding space without generating intermediate tokens. Models such as Ouro perform reasoning by iteratively updating internal representations while retaining a standard Key-Value (KV) cache across iterations, causing memory consumption to grow linearly with reasoning depth. Consequently, increasing the number of reasoning iterations can lead to prohibitive memory usage, limiting the practical scalability of such architectures. In this work, we propose Memory-Efficient Looped Transformer (MELT), a novel architecture that decouples reasoning depth from memory consumption. Instead of using a standard KV cache per layer and loop, MELT maintains a single KV cache per layer that is shared across reasoning loops. This cache is updated over time via a learnable gating mechanism. To enable stable and efficient training under this architecture, we propose to train MELT using chunk-wise training in a two phase procedure: interpolated transition, followed by attention-aligned distillation, both from the LoopLM starting model to MELT. Empirically, we show that MELT models fine-tuned from pretrained Ouro parameters outperform standard LLMs of comparable size, while maintaining a memory footprint comparable to those models and dramatically smaller than Ouro's. Overall, MELT achieves constant-memory iterative reasoning without sacrificing LoopLM performance, using only a lightweight post-training procedure.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。