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

• URL: http://arxiv.org/abs/2605.07721v1
• Date: Fri, 08 May 2026 13:25:27 GMT
• Title: Memory-Efficient Looped Transformer: Decoupling Compute from Memory in Looped Language Models
• Abstract summary: Memory-Efficient Looped Transformer (MELT) is a novel architecture that decouples reasoning depth from memory consumption.<n>We show that MELT achieves constant-memory iterative reasoning without sacrificing LoopLM performance.
• Score: 8.430043498727617
• License: http://creativecommons.org/licenses/by-nc-sa/4.0/
• Abstract: 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.

Related papers

• Linear-Time and Constant-Memory Text Embeddings Based on Recurrent Language Models [3.740517574718329]
  We introduce a vertically chunked inference strategy that enables fast embedding generation with memory usage that becomes constant in the input length once it exceeds the vertical chunk size.<n>By fine-tuning Mamba2 models, we demonstrate their viability as general-purpose text embedders.<n>We empirically validate the applicability of our inference strategy to Mamba2, RWKV, and xLSTM models, confirming consistent runtime-memory trade-offs across architectures.
  arXiv  Detail & Related papers  (2026-04-20T12:50:15Z)
• MemDLM: Memory-Enhanced DLM Training [56.40248490616793]
  Diffusion Language Models (DLMs) offer attractive advantages over Auto-Regressive (AR) models.<n>They suffer from a notable train-inference mismatch.<n>We propose MemDLM (Memory-Enhanced DLM), which embeds a simulated denoising process into training.
  arXiv  Detail & Related papers  (2026-03-23T17:39:56Z)
• Adaptive Loops and Memory in Transformers: Think Harder or Know More? [2.141762982579701]
  Chain-of-thought prompting enables reasoning in language models but requires explicit verbalization of intermediate steps.<n>Looped transformers offer an alternative by iteratively refining representations within hidden states.<n>We investigate transformer models that feature both adaptive per-layer looping and gated memory banks.
  arXiv  Detail & Related papers  (2026-03-09T13:49:03Z)
• NextMem: Towards Latent Factual Memory for LLM-based Agents [58.35585202907478]
  NextMem is a latent factual memory framework that utilizes an autoregressive autoencoder to efficiently construct latent memory.<n>Extensive experiments demonstrate that NextMem achieves superior performance.
  arXiv  Detail & Related papers  (2026-02-26T14:35:27Z)
• Out of the Memory Barrier: A Highly Memory Efficient Training System for LLMs with Million-Token Contexts [68.79341332280062]
  Training Large Language Models (LLMs) on long contexts is severely constrained by prohibitive GPU memory overhead, not training time.<n>We introduce OOMB, a highly memory-efficient training system that directly confronts this barrier.<n>Our approach employs a chunk-recurrent training framework with on-the-fly activation recomputation, which maintains a constant activation memory footprint.
  arXiv  Detail & Related papers  (2026-02-02T13:52:40Z)
• CoMeT: Collaborative Memory Transformer for Efficient Long Context Modeling [40.705016911274]
  We introduce a novel architecture that enables LLMs to handle arbitrarily long sequences with constant memory usage and linear time complexity.<n>CoMeT can be integrated into pre-trained models with only minimal fine-tuning.<n>A model equipped with CoMeT and fine-tuned on 32k contexts can accurately retrieve a passkey from any position within a 1M token sequence.
  arXiv  Detail & Related papers  (2026-02-02T07:49:44Z)
• EdgeInfinite: A Memory-Efficient Infinite-Context Transformer for Edge Devices [3.739419555718102]
  Transformer-based large language models (LLMs) encounter challenges in processing long sequences on edge devices.<n>We present EdgeInfinite, a memory-efficient solution for infinite contexts that integrates compressed memory into Transformer-based LLMs.
  arXiv  Detail & Related papers  (2025-03-28T07:26:37Z)
• Memory Layers at Scale [67.00854080570979]
  This work takes memory layers beyond proof-of-concept, proving their utility at contemporary scale.<n>On downstream tasks, language models augmented with our improved memory layer outperform dense models with more than twice the budget, as well as mixture-of-expert models when matched for both compute and parameters.<n>We provide a fully parallelizable memory layer implementation, demonstrating scaling laws with up to 128B memory parameters, pretrained to 1 trillion tokens, comparing to base models with up to 8B parameters.
  arXiv  Detail & Related papers  (2024-12-12T23:56:57Z)
• LoRC: Low-Rank Compression for LLMs KV Cache with a Progressive Compression Strategy [59.1298692559785]
  Key-Value ( KV) cache is crucial component in serving transformer-based autoregressive large language models (LLMs) Existing approaches to mitigate this issue include: (1) efficient attention variants integrated in upcycling stages; (2) KV cache compression at test time; and (3) KV cache compression at test time. We propose a low-rank approximation of KV weight matrices, allowing plug-in integration with existing transformer-based LLMs without model retraining. Our method is designed to function without model tuning in upcycling stages or task-specific profiling in test stages.
  arXiv  Detail & Related papers  (2024-10-04T03:10:53Z)
• Mini-Sequence Transformer: Optimizing Intermediate Memory for Long Sequences Training [78.93900796545523]
  Mini-Sequence Transformer (MsT) is a methodology for highly efficient and accurate LLM training with extremely long sequences. MsT partitions input sequences and iteratively processes mini-sequences to reduce intermediate memory usage. integrated with the huggingface library, MsT successfully extends the maximum context length of Qwen, Mistral, and Gemma-2 by 12-24x.
  arXiv  Detail & Related papers  (2024-07-22T01:52:30Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.