Adaptive Loops and Memory in Transformers: Think Harder or Know More?

• URL: http://arxiv.org/abs/2603.08391v3
• Date: Wed, 11 Mar 2026 09:07:47 GMT
• Title: Adaptive Loops and Memory in Transformers: Think Harder or Know More?
• Abstract summary: 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.
• Score: 2.141762982579701
• License: http://creativecommons.org/licenses/by-sa/4.0/
• Abstract: Chain-of-thought (CoT) prompting enables reasoning in language models but requires explicit verbalization of intermediate steps. Looped transformers offer an alternative by iteratively refining representations within hidden states. This parameter efficiency comes at a cost, as looped models lack the storage capacity of deeper models which use unique weights per layer. In this work, we investigate transformer models that feature both adaptive per-layer looping, where each transformer block learns to iterate its hidden state via a learned halting mechanism, and gated memory banks, that provide additional learned storage. We find that looping primarily benefits mathematical reasoning, while memory banks help recover performance on commonsense tasks compared to parameter and FLOP matched models. Combining both mechanisms yields a model that outperforms an iso-FLOP baseline, with three times the number of layers, across math benchmarks. Analysis of model internals reveals layer specialization: early layers learn to loop minimally and access memory sparingly, while later layers do both more heavily.

Related papers

• Memory-Efficient Looped Transformer: Decoupling Compute from Memory in Looped Language Models [8.430043498727617]
  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.
  arXiv  Detail & Related papers  (2026-05-08T13:25:27Z)
• Memory Caching: RNNs with Growing Memory [56.25483647131372]
  We introduce Memory Caching (MC), a technique that enhances recurrent models by caching checkpoints of memory states (a.k.a. hidden states)<n>We propose four variants of MC, including gated aggregation and sparse selective mechanisms, and discuss their implications on both linear and deep memory modules.<n>The results indicate that while Transformers achieve the best accuracy, our MC variants show competitive performance, close the gap with Transformers, and performs better than state-of-the-art recurrent models.
  arXiv  Detail & Related papers  (2026-02-27T18:53:41Z)
• Echo State Transformer: Attention Over Finite Memories [2.118933003468525]
  We introduce Echo State Transformers (EST), a hybrid architecture that elegantly resolves the challenge of sequential data processing.<n>EST integrates the Transformer attention mechanisms with principles from Reservoir Computing to create a fixed-size window distributed memory system.<n>EST ranks first overall in two of five categories, outperforming strong state-of-the-art baselines on classification and anomaly detection tasks.
  arXiv  Detail & Related papers  (2025-06-25T09:56:25Z)
• 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)
• MemoryFormer: Minimize Transformer Computation by Removing Fully-Connected Layers [43.39466934693055]
  We present MemoryFormer, a novel transformer architecture which significantly reduces the computational complexity (FLOPs) from a new perspective. This is made possible by utilizing an alternative method for feature transformation to replace the linear projection of fully-connected layers. We conduct extensive experiments on various benchmarks to demonstrate the effectiveness of the proposed model.
  arXiv  Detail & Related papers  (2024-11-20T02:41:53Z)
• Learning-to-Cache: Accelerating Diffusion Transformer via Layer Caching [56.286064975443026]
  We make an interesting and somehow surprising observation: the computation of a large proportion of layers in the diffusion transformer, through a caching mechanism, can be readily removed even without updating the model parameters. We introduce a novel scheme, named Learningto-Cache (L2C), that learns to conduct caching in a dynamic manner for diffusion transformers. Experimental results show that L2C largely outperforms samplers such as DDIM and DPM-r, alongside prior cache-based methods at the same inference speed.
  arXiv  Detail & Related papers  (2024-06-03T18:49:57Z)
• Memory-efficient Stochastic methods for Memory-based Transformers [3.360916255196531]
  Memory-based transformers can require a large amount of memory and can be quite inefficient. We propose a novel two-phase training mechanism and a novel regularization technique to improve the training efficiency of memory-based transformers.
  arXiv  Detail & Related papers  (2023-11-14T12:37:25Z)
• Recurrent Action Transformer with Memory [39.58317527488534]
  This paper proposes a novel model architecture that incorporates a recurrent memory mechanism designed to regulate information retention. We conduct experiments on memory-intensive environments (ViZDoom-Two-Colors, T-Maze, Memory Maze, Minigrid-Memory), classic Atari games, and MuJoCo control environments. The results show that using memory can significantly improve performance in memory-intensive environments, while maintaining or improving results in classic environments.
  arXiv  Detail & Related papers  (2023-06-15T19:29:08Z)
• Blockwise Parallel Transformer for Large Context Models [70.97386897478238]
  Blockwise Parallel Transformer (BPT) is a blockwise computation of self-attention and feedforward network fusion to minimize memory costs. By processing longer input sequences while maintaining memory efficiency, BPT enables training sequences 32 times longer than vanilla Transformers and up to 4 times longer than previous memory-efficient methods.
  arXiv  Detail & Related papers  (2023-05-30T19:25:51Z)
• A Model or 603 Exemplars: Towards Memory-Efficient Class-Incremental Learning [56.450090618578]
  Class-Incremental Learning (CIL) aims to train a model with limited memory size to meet this requirement. We show that when counting the model size into the total budget and comparing methods with aligned memory size, saving models do not consistently work. We propose a simple yet effective baseline, denoted as MEMO for Memory-efficient Expandable MOdel.
  arXiv  Detail & Related papers  (2022-05-26T08:24:01Z)
• IOT: Instance-wise Layer Reordering for Transformer Structures [173.39918590438245]
  We break the assumption of the fixed layer order in the Transformer and introduce instance-wise layer reordering into the model structure. Our method can also be applied to other architectures beyond Transformer.
  arXiv  Detail & Related papers  (2021-03-05T03:44:42Z)

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.