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

Victor Conchello Vendrell; Arnau Padres Masdemont; Niccolò Grillo; Jordi Ros-Giralt; Arash Behboodi; Fabio Valerio Massoli

Source-linked AI summary

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

Victor Conchello Vendrell, Arnau Padres Masdemont, Niccolò Grillo, Jordi Ros-Giralt, Arash Behboodi, Fabio Valerio Massoli

arXiv:2605.07721v2May 8, 2026cs.CLcs.AIcs.LG

TL;DR

Looped transformers support latent multi-step reasoning but their KV-cache memory grows with reasoning depth. MELT shares a gated KV cache across loops and, when integrated into Ouro, outperforms similarly sized standard transformers across reasoning benchmarks while maintaining constant memory.

Problem
Looped transformers enable latent reasoning but their memory grows linearly with the number of reasoning loops because of per-loop KV states.
Method
MELT shares a single KV cache per layer across loops, updates it with learnable gating, and adapts pretrained LoopLMs through two-phase chunk-wise training.
Results
MELT consistently outperforms similarly sized standard transformers across mathematical and general reasoning benchmarks while maintaining a comparable constant memory footprint.
Takeaways & Limitations
MELT shows that looped iterative computation can improve reasoning performance over standard models without increasing memory under the evaluated settings.
Takeaways & Limitations
MELT inherits Ouro’s fixed inference-time loop count, so it does not adapt reasoning steps to the varying needs of different inputs and tokens.

Abstract

from arXiv · show

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.

1 Introduction

Looped transformers enable latent multi-step reasoning but their per-loop KV caches make memory grow with reasoning depth. MELT addresses this by sharing and gating one KV cache per layer, then adapting pretrained LoopLMs through a two-phase chunk-wise procedure while retaining performance and reducing memory.

Motivation: Looped transformers repeatedly pass hidden states through one transformer stack, enabling latent reasoning without generating intermediate tokens.LoopLM has shown that looped models can match or surpass transformers nearly twice their size.
Motivation: Per-loop KV-cache retention makes memory grow linearly with reasoning iterations, limiting the scalability of looped architectures.Chain-of-Thought similarly couples reasoning depth to output length, increasing latency and memory usage.
Empirical findings: MELT outperforms similarly sized standard transformers on reasoning benchmarks while preserving originating LoopLM performance with dramatically lower memory than looped baselines.The comparison is against looped baselines that retain per-loop KV growth.
MELT architecture: MELT decouples reasoning depth from memory by sharing a single KV cache per layer across loops and updating it with a learnable gate.This architecture is designed to avoid per-loop KV growth.
Training procedure: MELT adapts pretrained LoopLMs through chunk-wise training with interpolated transition followed by attention-aligned distillation using the frozen LoopLM.The procedure is presented as data-efficient and uses MELT initialized from pretrained Ouro weights.

2 Related work

Related work frames looped transformers as a vertical-reasoning paradigm built on iterative computation, while highlighting KV-cache growth and limitations of existing compression methods. MELT’s training transition and representation-level distillation relate to prior gradual adaptation and knowledge-distillation approaches.

Looped transformers: Looped transformers extend reasoning vertically through recurrent architectures, adaptive layer repetition or skipping, and iterative computation.This line of work includes HRM, TRM, and adaptive-depth methods, and is described as a strong architectural paradigm.
KV cache compression and vertical sharing: Existing KV-cache methods reuse or compress states across heads, layers, and recurrence steps to reduce memory growth.Examples include MQA/GQA, hybrid global–local attention, recursion-aware caching and sharing, and untrained reuse.
KV cache compression and vertical sharing: Prior KV-cache reduction methods often degrade performance on stronger models and longer reasoning traces in complex long-context tasks.The passage notes that memory savings can be effective in constrained settings but remain limited for demanding reasoning workloads.
Training transitions and representation-level distillation: MELT’s training procedure relates to progressive growing, gradual training and adaptation, architectural modification, and knowledge distillation for stabilizing pretrained-model transitions.These approaches motivate interpolating existing and newly introduced components and using distillation during architectural adaptation.

3 Memory-Efficient Looped Transformer

MELT decouples looped reasoning depth from KV-cache memory by sharing and updating one cache per layer through a gated latent state. Chunk-wise training, interpolated transition, and attention-aligned distillation adapt pretrained LoopLMs to these dynamics while preserving alignment and stability.

Memory-efficient architecture: MELT keeps one fixed-size KV cache per layer, making total cache memory O(N × L) instead of O(N × L × T) for LoopLM.Each loop updates cached token states rather than appending new states, so depth-wise memory is O(1) per layer.
Memory-efficient architecture: The gating mechanism lets each token attend to keys and values integrating information across all preceding-token time steps, rather than only the current step.A single shared cache adds one row per token and updates it across loops.
Memory-efficient architecture: A latent state evolves across iterations and is projected into keys and values, preserving query–key alignment while separating memory updates from attention retrieval.The learnable gated momentum mechanism determines which information is retained or overwritten over time.
Chunk-wise training: Chunk-wise training processes chunks sequentially while parallelizing tokens within each chunk, balancing autoregressive fidelity against training efficiency.Smaller chunks more closely match inference dynamics, while larger chunks improve efficiency.
Two-phase adaptation: Interpolated transition linearly increases α from 0 to 1 between LoopLM and MELT KV pairs, followed by frozen-teacher distillation with attention alignment at every layer and loop.The second phase stabilizes training and reduces the gap to the original LoopLM after unconstrained MELT continuation degrades performance.

4 Experimental results

MELT outperforms non-looped baselines across broad reasoning benchmarks while using substantially less memory than Ouro. Ablations show that gated aggregation and the complete two-phase training procedure are important for performance.

Benchmark setup: MELT is evaluated on six mathematical and four general reasoning benchmarks against similarly sized non-looped models and Ouro-1.4B-Thinking.All models use LightEval v0.8.1 with default benchmark prompts, extraction procedures, and evaluation settings.
Benchmark results: MELT consistently outperforms all non-looped baselines across mathematical and general reasoning benchmarks, with only limited exceptions.It leads on AIME24, AIME26, MATH500, OlympiadBench, MMLU, and HumanEval; Qwen3-1.7B leads on AIME25 and AMC23, while Gemma4-E2B leads on GPQA.
Benchmark results: MELT slightly trails Ouro on most benchmarks but surpasses it on HumanEval, achieving a strong performance–efficiency trade-off.Ouro benefits from substantially higher memory usage because it retains a full per-loop KV cache.
Memory efficiency: MELT reduces KV-cache memory by ∼3-4× relative to Ouro by maintaining a constant-size latent state instead of appending loop-specific KV entries.For a 32k-token generation, Ouro exceeds Qwen by ∼20 GB, whereas MELT exceeds Qwen by only ∼2.5 GB.
Ablations: Element-wise gating performs best after Phase 1, while the full two-phase training recipe improves performance consistently over configurations with progressively removed components.Removing attention-aligned distillation, interpolation training, knowledge distillation, or chunk-wise training degrades results across benchmarks.

5 Limitations and future work

MELT inherits fixed inference-time loop depth, lacks current exploration of GQA, and requires sequential KV updates during training. Future work targets adaptive loop depth, GQA integration, and more parallelizable training strategies.

Fixed loop depth: MELT inherits Ouro’s fixed recurrent-loop count at inference, limiting adaptation to inputs and tokens with different reasoning requirements.Its constant-size latent state supports future adaptive loop-depth extensions that dynamically allocate reasoning steps by input complexity.
GQA integration: The current implementation does not explore GQA, whose integration could reduce inference memory bandwidth and KV-cache overhead.MELT’s constant-memory design makes combining it with GQA promising for further efficiency improvements.
Training parallelism: Sequential KV updates during training constrain parallelism relative to standard Transformer training.More parallelizable training strategies are needed to scale MELT to larger models, longer reasoning horizons, and broader application domains.

6 Conclusion · A Extended Related Work

MELT decouples latent reasoning depth from memory growth by replacing append-only KV caching with a gated, constant-size latent state, while retaining competitive reasoning performance. The related work situates MELT among recurrent looped transformers, KV-cache sharing and compression, progressive transition strategies, and activation-level knowledge distillation.

6 Conclusion: MELT enables deep latent reasoning with inference-time compute scaling without linear memory growth.It replaces the append-only KV cache with a gated, constant-size latent state.
6 Conclusion: Training MELT uses chunk-wise processing because sequential KV-cache computation prevents fully parallel token processing.Smaller chunks improve inference fidelity, whereas larger chunks improve training efficiency.
6 Conclusion: MELT surpasses similarly sized standard Transformer baselines across mathematical and general reasoning benchmarks under the same constant memory budget.The results demonstrate improved reasoning capability without increasing memory.
A Extended Related Work: Recurrent reasoning work spans simple architectures such as HRM and TRM and transformer modifications with looping mechanisms.These approaches represent a complementary direction to chain-of-thought and other iterative token computation techniques.
A Extended Related Work: Compared with similarly sized vanilla transformers, looped transformers show stronger capacity for multi-hop reasoning, length generalization, and learning algorithms.Scaling looped architectures over many steps remains challenged by optimization instability and vanishing gradients.
A Extended Related Work: KV-cache efficiency research includes head-level sharing through MQA and GQA and vertical sharing across layers or recurrence steps.Cross-Layer Attention exploits relatively stable KV representations across adjacent layers so multiple layers can read from a shared cache.
A Extended Related Work: Prior looped-transformer methods address growing KV caches through global-plus-local caching and recursion-wise selective updates for active thinking tokens.These approaches respectively combine first-loop global information with a sliding recent-token window or restrict attention to selected KV pairs at each depth.
A Extended Related Work: MELT’s interpolated transition and activation-level distillation relate to progressive growing and knowledge distillation methods that align intermediate representations.Progressive growing smoothly interpolates old and newly added pathways, while activation-level distillation transfers structural and semantic information beyond final logits.

B Analysis of existing KV-Cache sharing methods on long reasoning · C Attention Alignment Loss · D Hyperparameters

The paper tests whether untrained KV-cache sharing extends to Ouro under long reasoning, finding severe failures, and documents attention alignment plus reproducibility and compute settings for MELT experiments.

B Analysis of existing KV-Cache sharing methods on long reasoning: Existing KV-cache sharing methods had mainly been evaluated on weaker looped models and short or constrained generations, leaving long-reasoning applicability unclear.The analysis therefore examines state-of-the-art Ouro models.
B Analysis of existing KV-Cache sharing methods on long reasoning: The study evaluates first-loop and last-loop KV-cache reuse, each with or without preserving the original prompt KV cache.These variants follow previously proposed sharing strategies and prompt-cache preservation settings.
B Analysis of existing KV-Cache sharing methods on long reasoning: All four untrained KV-sharing variants obtain zero performance on several reasoning benchmarks.Full results are reported in Table 5 for Ouro-1.4B-Thinking.
B Analysis of existing KV-Cache sharing methods on long reasoning: The findings contrast with reported last-loop KV sharing with prefill, which performed comparably to Ouro on GSM8K and MATH-500.The paper hypothesizes that few-shot CoT evaluation constrained generation and reduced reasoning drift.
C Attention Alignment Loss: Figure 5 illustrates that Attention Alignment Loss matches MELT attention outputs to the frozen LoopLM teacher at each layer and reasoning loop.The figure supports the explanation of the auxiliary alignment loss in Subsection 3.3.
D Hyperparameters: Table 6 reports training hyperparameters for MELT-1.6B and ablation studies, while Table 7 summarizes benchmark evaluation settings.These tables provide the settings required to reproduce the training and evaluation runs.
D Hyperparameters: For HLE, evaluation uses a 300-sample subset of the original dataset.This qualification applies to the benchmark evaluation settings summarized in Table 7.
D Hyperparameters: 130 hours on 8 H100 GPUs correspond to 1,040 GPU-hours for the main MELT-1.6B run, while the three ablations required 1,440 GPU-hours total.Evaluation used approximately 500 GPU-hours, preliminary experiments roughly 15,000 GPU-hours, and the overall project approximately 20,000 GPU-hours.

E Theoretical analysis

The theoretical analysis specializes its proofs to MELT’s update of only the last row of the state matrix, while noting that the results generalize to full-matrix updates.

E Theoretical analysis: MELT’s proofs explicitly account for updating only the last row of the state matrix.The analysis states that the results also generalize to full-matrix updates.

E.1 Spectral stability of the gated update · E.2 Gradient Superhighway

MELT’s saturated gate stabilizes recurrence by making the local Jacobian converge to identity, preserving gradients across arbitrarily many loops. This yields a Gradient Superhighway that supports optimization of deeper looped transformers without vanishing gradients.

E.1 Spectral stability of the gated update: In the saturated-gate regime, gradients are preserved across loops, providing the foundation for MELT’s Gradient Superhighway.The result applies when the gate saturates element-wise toward 1 in selected latent dimensions.
E.1 Spectral stability of the gated update: When z_t →1, the restricted Jacobian J_t converges to the identity matrix and its spectral radius satisfies ρ(J_t) →1.This ensures gradient magnitude is preserved over any arbitrary number of iterations.
E.1 Spectral stability of the gated update: As z_t →1, the sigmoid derivative vanishes, causing the gate-state derivative ∂z_t/∂h_{t−1} to approach zero.This removes one contribution to the recurrent Jacobian in the shielding regime.
E.1 Spectral stability of the gated update: As z_t →1, the factor (1 − z_t) approaches zero, nullifying the recurrent weight matrix’s contribution to ∂h̃_t/∂h_{t−1}.The gate therefore suppresses this recurrent contribution in the shielding regime.
E.1 Spectral stability of the gated update: The saturated-gate limit gives lim_z→1 J_t = I + 0 + 0, so the Jacobian is approximately identity and ρ(J_t) = 1.The conclusion follows because all eigenvalues of the identity matrix are 1.
E.1 Spectral stability of the gated update: The gate acts as a structural stabilizer by keeping recurrence spectral radii near unity, while its boundedness prevents exploding-gradient instability.The passage also states that this permits gradients to propagate through long sequences without vanishing.
E.2 Gradient Superhighway: For shielding-regime dimensions with Z_t ≈1, each local Jacobian acts effectively as the identity, simplifying the product of Jacobians across recurrent updates.This follows from Proposition E.1’s analysis of backpropagation through the recurrent updates.
E.2 Gradient Superhighway: The resulting Gradient Superhighway provides a direct path for error signals through arbitrary depths, alleviating vanishing gradients and enabling deeper looped-transformer optimization.Standard recurrent dynamics instead exhibit gradient norms scaling as O(λ^T), with exponential decay when |λ| < 1.

F Notes on reproducibility and inference efficiency in Ouro

This appendix records technical observations from reproducing and analyzing Ouro, with the aim of clarifying practical comparisons of memory usage and inference efficiency against MELT.

Purpose: The appendix documents technical observations made while reproducing and analyzing claims from Ouro.These observations are presented as part of the paper’s reproducibility analysis.
Purpose: The authors explicitly state that these observations are not intended to diminish Ouro’s contributions.The appendix frames its discussion as clarification rather than criticism of Ouro’s contributions.
Comparison focus: The observations address practical aspects relevant to comparing memory usage and inference efficiency between Ouro and MELT.The appendix focuses on comparison-relevant implementation and analysis considerations.

F.1 Reproducibility

The authors could not fully reproduce Ouro’s reported results because key implementation details and evaluation settings were underspecified or mismatched, producing non-trivial performance discrepancies despite released code and checkpoints.

F.1 Reproducibility: The reported Ouro results could not be fully reproduced using the released code and pretrained checkpoints.Key implementation details and evaluation settings were either underspecified or differed from those needed to match the reported numbers.
F.1 Reproducibility: Non-trivial discrepancies emerged between the paper’s reported performance and results obtained from the released artifact.

F.2 Early-exit gating and effective compute · G Existing assets

Ouro’s early-exit gating is intended to enable adaptive computation, but the released implementation generally does not reduce recurrent computation or memory. The appendix catalogs the benchmarks, models, datasets, codebases, and licenses used for transparency and reproducibility.

F.2 Early-exit gating and effective compute: Ouro’s learned gating mechanism is designed to let tokens exit early when further recurrent steps seem unnecessary.This mechanism is presented as a way to enable adaptive computation.
F.2 Early-exit gating and effective compute: The default threshold effectively disables early exiting, and the paper does not explain how to choose the threshold in practice.The released configuration therefore does not realize the intended early-exit behavior under typical settings.
F.2 Early-exit gating and effective compute: Even when triggered, early exit selects logits but still computes every recurrent loop through the maximum depth.Thus, gating changes output selection rather than terminating subsequent loop computation.
F.2 Early-exit gating and effective compute: Code inspection confirms that Ouro’s gating does not reduce inference-time compute or memory under typical settings.This behavior conflicts with its conceptual framing as an adaptive-compute mechanism.
F.2 Early-exit gating and effective compute: The limitation is hypothesized to arise from KV-cache dependencies, because later tokens need final-loop KV states for autoregressive consistency.Per-token early termination could therefore break access to the states required by later tokens.
F.2 Early-exit gating and effective compute: MELT’s constant-memory KV updates may avoid this specific limitation, but early-exit strategies for MELT remain future work.The passage identifies a potential architectural advantage without reporting a completed investigation.
G Existing assets: The existing-assets appendix documents evaluation benchmarks, comparison models, training datasets, and implementation codebases with their corresponding licenses and source links.These resources are listed in Tables 8–11 to support transparency and reproducibility.