内存高效型循环Transformer：循环语言模型中的计算与内存解耦

Hugging Face Daily Papers 论文
memory-efficient looped-transformer recurrent-llm kv-cache reasoning chunk-wise-training

摘要

提出内存高效型循环Transformer（MELT），这是一种新型循环大语言模型架构，通过跨循环共享单一KV缓存，并结合插值过渡与注意力对齐蒸馏的分块训练方法，实现了推理深度与内存消耗的解耦。

循环大语言模型架构已成为提升推理能力的一种极具前景的方法，因为它们能够在嵌入空间中进行多步计算，而无需生成中间 token。以 Ouro 等模型为例，它们通过迭代更新内部表示来执行推理，同时在各次迭代中保留标准的键值（KV）缓存，这导致内存消耗随推理深度呈线性增长。因此，增加推理迭代次数会导致内存占用过高，限制了此类架构的实际可扩展性。在本文中，我们提出了一种新型架构——内存高效型循环Transformer（MELT），该架构实现了推理深度与内存消耗的解耦。与传统为每一层和每个循环分配独立KV缓存的方式不同，MELT每层仅维护一个单一KV缓存，并在所有推理循环中共享。该缓存通过可学习的门控机制随时间进行更新。为使该架构下的训练稳定且高效，我们提出采用分块训练的两阶段流程来训练MELT：首先是插值过渡，随后是注意力对齐蒸馏，两个阶段均将知识从起始的LoopLM模型迁移至MELT。实验表明，从预训练的Ouro参数微调而来的MELT模型，在保持与同规模标准LLM相当的内存占用的同时，性能优于后者，且内存占用远低于Ouro。总体而言，MELT仅通过轻量级的后训练流程，便在不牺牲LoopLM性能的前提下，实现了恒定内存占用的迭代推理。
查看原文
查看缓存全文

缓存时间: 2026/05/12 10:53

论文页面 - 内存高效循环 Transformer：在循环语言模型中将计算与内存解耦

Source: https://huggingface.co/papers/2605.07721

摘要

本文提出了一种名为 MELT 的新型循环 LLM 架构，该架构通过在推理循环间共享单一 KV Cache，并结合带有插值过渡（interpolated transition）与注意力对齐蒸馏（attention-aligned distillation）的分块训练（chunk-wise training），实现了推理深度与显存消耗的解耦。

循环 LLM 架构 (https://huggingface.co/papers?q=Recurrent%20LLM%20architectures) 已成为提升推理能力的一项极具前景的方案，因为它们能够在嵌入空间 (https://huggingface.co/papers?q=embedding%20space) 内实现多步计算 (https://huggingface.co/papers?q=multi-step%20computation)，而无需生成中间 token。诸如 Ouro 之类的模型通过迭代更新内部表示来进行推理，同时在多次迭代间保留标准的键值 (KV) 缓存，这导致显存消耗 (https://huggingface.co/papers?q=memory%20consumption) 随推理深度 (https://huggingface.co/papers?q=reasoning%20depth) 呈线性增长。因此，增加推理迭代次数会导致内存占用过高，从而限制了此类架构的实际可扩展性。在本研究中，我们提出了内存高效循环 Transformer (MELT)，这是一种新型架构，能够将推理深度 (https://huggingface.co/papers?q=reasoning%20depth) 与显存消耗 (https://huggingface.co/papers?q=memory%20consumption) 解耦。MELT 不为每一层和每次循环单独分配标准 KV 缓存，而是为每层维护单一的 KV 缓存，并使其在所有推理循环间共享。该缓存通过一种可学习的门控机制 (https://huggingface.co/papers?q=learnable%20gating%20mechanism) 随时间进行更新。为了在该架构下实现稳定高效的训练，我们提出采用分块训练 (https://huggingface.co/papers?q=chunk-wise%20training) 策略，并设计了一个两阶段流程：先是插值过渡 (https://huggingface.co/papers?q=interpolated%20transition)，随后是注意力对齐蒸馏 (https://huggingface.co/papers?q=attention-aligned%20distillation)。这两个阶段均以 LoopLM (https://huggingface.co/papers?q=LoopLM) 作为起始模型，最终训练得到 MELT。实验结果表明，基于预训练 Ouro 参数微调的 MELT 模型性能优于同等规模的常规 LLM，同时其显存占用与这些模型相当，且远小于 Ouro 的占用。总体而言，仅凭借轻量级的后训练流程，MELT 便实现了恒定显存的迭代推理 (https://huggingface.co/papers?q=iterative%20reasoning)，且未牺牲 LoopLM (https://huggingface.co/papers?q=LoopLM) 的性能。

View arXiv page (https://arxiv.org/abs/2605.07721) View PDF (https://arxiv.org/pdf/2605.07721) Add to collection (https://huggingface.co/login?next=%2Fpapers%2F2605.07721)

Get this paper in your agent:

hf papers read 2605.07721

Don’t have the latest CLI? curl -LsSf https://hf.co/cli/install.sh | bash

引用此论文的模型 0

暂无模型引用此论文

请在模型的 README.md 中引用 arxiv.org/abs/2605.07721，以便从此页面建立链接。

引用此论文的 Dataset 0

暂无数据集引用此论文

请在数据集的 README.md 中引用 arxiv.org/abs/2605.07721，以便从此页面建立链接。

引用此论文的 Space 0

暂无 Space 引用此论文

请在 Space 的 README.md 中引用 arxiv.org/abs/2605.07721，以便从此页面建立链接。

收录此论文的 Collection 0

暂无 Collection 收录此论文

将此论文添加至 合集 以便从此页面建立链接。

相似文章

Mela：基于转化假说的测试时记忆巩固

Hugging Face Daily Papers

本文介绍了 Mela，这是一种受人类记忆巩固机制启发的、增强记忆的 Transformer 架构，其核心特征为层级记忆模块（HMM），能够有效提升长上下文语言建模的性能。

语言模型需要睡眠

Hacker News Top

本文提出了一种类似睡眠的巩固机制，适用于基于Transformer的大语言模型，该机制定期将最近上下文转换为SSM块中的持久快速权重，清除KV缓存，从而在不增加推理延迟的情况下提升长期推理能力。

Δ-Mem：大型语言模型的高效在线记忆

Hacker News Top

提出 delta-Mem，一种轻量级在线记忆机制，利用紧凑状态矩阵并通过增量规则学习进行更新，以提升冻结大型语言模型的长上下文性能，无需全量微调或上下文扩展。