PaperlayerView exact revision on arXiv

Source-linked AI summary

MeSH: Memory-as-State-Highways for Recursive Transformers

Chengting Yu, Xiaobo Shu, Yadao Wang, Yizhen Zhang, Haoyi Wu, Jiaang Li, Rujiao Long, Ziheng Chen, Yuchi Xu, Wenbo Su, Bo Zheng

arXiv:2510.07739v2cs.LGcs.AI

TL;DR

Under matched compute, recursive models with fewer parameters often lag non-recursive counterparts. MeSH uses an explicit memory buffer and iteration-specific routing to manage state and diversify computation; its recursive models consistently outperform recursive baselines and surpass a larger Vanilla model at 1.4B scale.

  • Problem

    Under matched compute, recursive models with fewer parameters often lag non-recursive counterparts, motivating ways to improve their performance without matching their parameter counts.

    Source: §1 · p. 1

  • Method

    MeSH externalizes state management into an explicit memory buffer and uses step-wise routers to dynamically compose information across iterations.

    Source: §1 · p. 2, §3.3 · p. 6

  • Results

    At Pythia-1.4B, +mesh improves 0-shot average accuracy by +1.06% over Vanilla while using approximately 33% fewer non-embedding parameters.

    Source: §4.2 · p. 8

  • Takeaways & Limitations

    Iteration-specific routing lets MeSH dynamically switch recurrence behaviors, enabling functional specialization across recursive steps.

    Source: §3.3 · p. 6

    Experiments validate MeSH up to Pythia-6.9B on the deduplicated Pile, leaving its parameter-efficiency gains at larger scales and under different training regimes unclear.

    Source: G LIMITATIONS AND FUTURE WORK · p. 28

Abstract

from arXiv · show

Recursive transformers reuse parameters and iterate over hidden states multiple times, decoupling compute depth from parameter depth. However, under matched compute, recursive models with fewer parameters often lag behind non-recursive counterparts. By probing hidden states, we trace this performance gap to two primary bottlenecks: undifferentiated computation, where the core is forced to adopt a similar computational pattern at every iteration, and information overload, where long-lived and transient information must coexist in a single hidden state. To address the issues, we introduce a Memory-as-State-Highways (MeSH) scheme, which externalizes state management into an explicit memory buffer and employs lightweight routers to dynamically diversify computation across iterations. Probing visualizations confirm that MeSH successfully resolves the pathologies by inducing functional specialization across iterations. On the Pythia suite (160M-6.9B), MeSH-enhanced recursive transformers consistently improve over recursive baselines and outperforms its larger non-recursive counterpart at the 1.4B scale, improving average downstream accuracy by +1.06% with 33% fewer non-embedding parameters. Our analysis establishes MeSH as a scalable and principled architecture for building stronger recursive models. Our code is available at https://github.com/LivingFutureLab/MeSH/ .

1 INTRODUCTION

Recursive transformers offer a parameter-efficient way to increase computational depth, but under matched compute they often lag non-recursive models. MeSH addresses this gap with explicit memory and step-wise routing.

1 INTRODUCTION

  • Recursive transformers repeatedly apply a compact, weight-shared core, decoupling computational depth from parameter depth.This can enable adaptive computation and add computational depth as a scaling axis.Source: §1 · p. 1
  • Under matched compute, recursive models with fewer parameters often have higher perplexity or lower accuracy than standard Transformers.The comparison holds for models with equivalent FLOPs but unique parameters per layer.Source: §1 · p. 1
  • MeSH replaces the overloaded hidden state with an explicit memory buffer and lightweight, step-wise routers to manage state and diversify computation.The introduction identifies undifferentiated computation and information overload as the bottlenecks it targets.Source: §1 · p. 2

2 WHY NAIVE RECURSION FAILS: A DIAGNOSTIC ANALYSIS

Naive recursion lacks explicit information about loop progress and forces one hidden state to serve multiple roles. The resulting diagnostics point to underused iterations and reduced representational capacity.

2 WHY NAIVE RECURSION FAILS: A DIAGNOSTIC ANALYSIS

  • A shared core without explicit loop-progress information cannot specialize its function across iterations, while a single hidden state must carry conflicting information.The paper names these bottlenecks undifferentiated computation and information overload.Source: §2 · p. 2
  • The first core loop performs most of the computational work, later updates approach zero, and consecutive loop states show high CKA similarity.These observations indicate that later loops are poorly used and representations stagnate near a fixed point.Source: §2 · pp. 2–3
  • Loop-state singular-value spectra decay faster than the initial-state spectrum, indicating lower effective rank and a shared low-dimensional subspace.The paper links this pattern to a stability–plasticity trade-off: persistent context is favored over transient, stepwise processing.Source: §2 · p. 3

3 METHODOLOGY: ALLEVIATING INFORMATION OVERLOAD AND ENABLING FUNCTIONAL SPECIALIZATION

The methodology moves from fixed recurrence rules toward MeSH, which combines an explicit memory buffer with learnable, step-wise routing to address information overload and enable functional specialization.

3 METHODOLOGY: ALLEVIATING INFORMATION OVERLOAD AND ENABLING FUNCTIONAL SPECIALIZATION

  • The methodology reviews fixed additive recurrence schemes before introducing MeSH as a framework targeting both information overload and limited functional specialization.The fixed schemes supplement context but do little to address undifferentiated computation.Source: §3 · p. 3

3.1 PRELIMINARIES: ARCHITECTURE OF RECURSIVE TRANSFORMERS

  • The Prelude-Recurrent-Coda architecture prepares an initial state, repeatedly applies a shared core block, then passes the final loop state to a coda.The core refines the state over K iterations.Source: §3.1 · pp. 3–4

3.2 MESH: MEMORY–AS-STATE-HIGHWAYS FOR RECURSIVE TRANSFORMERS

  • MeSH replaces fixed state passing with a memory buffer initialized from token embeddings and step-wise routers that control memory writes and reads.The routers use iteration-specific parameters, and the next state is synthesized from the updated buffer.Source: §3.2 · pp. 4–5

3.3 HOW MESH ADDRESSES THE DIAGNOSED PATHOLOGIES

  • Iteration-specific routers dynamically compose context from historical states, allowing the model to switch recurrence behaviors and specialize computation across steps.The routers replace rigid additive updates with a learnable read-write cycle.Source: §3.3 · p. 6
  • Externalizing persistent memory frees the hidden state to serve as a workspace for transient computation, while the buffer preserves long-lived information for retrieval.MeSH presents this as a multi-slot, state-aware alternative to a single overloaded information channel.Source: §3.3 · p. 6

4 EXPERIMENTS

Experiments diagnose three pathologies in naive recursion and show that MeSH improves recursive-model performance while retaining fewer non-embedding parameters. Additional analyses examine training dynamics and performance across layer allocations.

4.1 A COMPARATIVE DIAGNOSTIC ANALYSIS OF RECURRENCE SCHEMES

  • Across 500 Pythia-410M samples, MeSH balances computational effort across three core loops, unlike the sharply decaying effort in base, residual, and anchor variants.The state-update magnitude is used as a proxy for each block’s computational effort.Source: §4.1 · p. 6, Figure 3 · p. 7
  • MeSH reduces similarity between consecutive loop states while retaining a strong connection to the initial context, breaking the base model’s representational stagnation.Figure 4 measures pairwise hidden-state similarity with CKA.Source: §4.1 · p. 7, Figure 4 · p. 7
  • MeSH preserves high-dimensional hidden-state structure across iterations, whereas base-model loop states collapse into a lower-dimensional subspace.The comparison uses the normalized singular-value spectrum of hidden states.Source: §4.1 · pp. 7–8, Figure 5 · p. 7

4.2 MAIN RESULTS

  • +1.06% (0-shot) and +0.86% (5-shot) average accuracy: Pythia-1.4B +mesh beats Vanilla with approximately 33% fewer non-embedding parameters.The authors also report state-of-the-art perplexity across all evaluated datasets.Source: §4.2 · p. 8, Table 1 · p. 8

4.3 FURTHER ANALYSIS AND ABLATION STUDIES

  • For the 1.4B recursive variants, +mesh has lower training loss throughout pretraining and stronger, more consistent 0-shot accuracy gains across checkpoints.Figure 6 compares training loss over 120k steps and downstream accuracy over checkpoints.Source: §4.3 · p. 9, Figure 6 · p. 9
  • The 805M +mesh model reaches 50.6% 0-shot and 52.8% 5-shot accuracy, exceeding the 1.2B-parameter Vanilla model’s 49.5% and 51.9%.The authors characterize this as a 1.46x improvement in parameter efficiency.Source: §4.3 · p. 9, Figure 7 · p. 9
  • With compute matched to the 24-layer Vanilla model, +mesh has lower perplexity than base recursion across layer allocations and approaches Vanilla using approximately 30% fewer non-embedding parameters.The control study uses Pythia-410M and varies the prelude-loop-coda layer distribution.Source: §4.3 · pp. 9–10, Figure 8 · p. 10

5 CONCLUSION

The conclusion presents explicit, routed memory as a scalable approach to stronger recursive models. Related work situates MeSH among recursive transformers, iteration-differentiation methods, connectivity schemes, latent reasoning, and memory-augmented transformers.

5 CONCLUSION

  • The paper concludes that explicit, routed state management is a scalable and effective principle for building stronger recursive models.Experiments are reported to validate that MeSH addresses diagnosed pathologies while improving recursive-backbone performance.Source: §5 · p. 10

A RELATED WORK

  • Prior looped-transformer work explores shared-layer iteration, adaptive computation, and latent reasoning; MeSH contributes an improved recursive foundation for that line of research.The cited work includes Universal Transformer [Dehghani et al., 2018] and later studies [Saunshi et al., 2025; Zhu et al., 2025b; Geiping et al., 2025].Source: A RELATED WORK · p. 16
  • Iteration differentiation has used loop-index embeddings, per-iteration LoRA, and MoE gating; MeSH instead routes information through memory while keeping the core weight-shared.The related-work discussion notes limited gains from element-wise loop embeddings [Geiping et al., 2025; Zhu et al., 2025a].Source: A RELATED WORK · p. 16
  • Unlike layer-to-layer skip and dense-connectivity schemes, MeSH manages loop-to-loop information flow with an external buffer and step-wise routers.The discussion contrasts this design with Hyper-connections [Zhu et al., 2024] and MUDDformer [Xiao et al., 2025].Source: A RELATED WORK · pp. 16–17
  • Recursive models can process reasoning steps internally rather than outputting them as natural-language chain-of-thought, and MeSH targets the recursive architecture underlying that approach.The paper describes latent chain-of-thought and explicit chain-of-thought as potentially complementary.Source: A RELATED WORK · p. 17
  • Where many memory-augmented transformers extend context or store long-term knowledge, MeSH manages intermediate states across iterations on one fixed-length input.Source: A RELATED WORK · p. 17

B EXPERIMENTAL DETAILS

The experimental details specify pretraining, architecture, optimization, and downstream evaluation. The pseudocode contrasts fixed state-passing variants with MeSH’s dynamic memory-based recurrence.

B EXPERIMENTAL DETAILS

  • Models are pretrained from scratch for one epoch on a 250B-token deduplicated subset of the Pile, following the Pythia methodology.They use the GPT-NeoX tokenizer with a 50,257-token vocabulary.Source: B EXPERIMENTAL DETAILS · p. 17
  • Recursive models use a prelude-loop-coda structure, with MeSH adding a state buffer and step-wise routers at block boundaries.Each router is a linear layer followed by softmax; buffer size is set to B = Nloop + 3.Source: B EXPERIMENTAL DETAILS · p. 17
  • Training uses AdamW, cosine learning-rate decay, global batch size 512, and sequence length 4096, with BF16 and FlashAttention-2.Peak learning rates range from 6.0 × 10^-4 for 160M to 2.0 × 10^-4 for 1.4B models.Source: B EXPERIMENTAL DETAILS · p. 18
  • Evaluation covers perplexity on Pile, Wikitext, and Lambada, plus 0-shot and 5-shot accuracy on nine downstream benchmarks.The benchmark suite and scoring conventions are specified in the evaluation description.Source: B EXPERIMENTAL DETAILS · p. 18

C PSEUDOCODE

  • The pseudocode contrasts common recursive variants’ fixed, heuristic state-passing with MeSH’s dynamic memory-based recurrence.Algorithm 1 outlines common variants; Algorithm 2 details MeSH.Source: C PSEUDOCODE · p. 18

D DISCUSSION: EXPRESSIVE POWER OF MESH AS A GENERAL RECURRENCE

MeSH generalizes fixed recursive updates by dynamically retrieving and combining historical states with current computation. Its benefits persist at larger Pythia scales, while buffer ablations identify a preferred capacity.

D DISCUSSION: EXPRESSIVE POWER OF MESH AS A GENERAL RECURRENCE

  • MeSH’s compute-write-read recurrence dynamically retrieves and combines historical states to form the next state, generalizing additive residual and anchor updates.The state combines a retrieved historical summary with a gated output of the current core block; routers can adapt the recurrence rule across iterations.Source: D DISCUSSION: EXPRESSIVE POWER OF MESH… · pp. 20–21

E.1 SCALING TO LARGER RECURSIVE MODELS

  • At 2.8B and 6.9B, MeSH lowers perplexity on all four datasets and boosts 0-shot and 5-shot average accuracy, with non-embedding parameters at 68.75% of non-recursive counterparts.Both models use a 6+10R2+6 recursive configuration; no vanilla counterparts were trained.Source: §E.1 · p. 21, Table 3 · p. 22

E.2 ABLATION STUDY: MESH BUFFER LENGTH

  • For Pythia-410M with Nloop = 2, k = 2 (buffer length B = 5) gives the lowest perplexity on all four datasets; performance slightly degrades at k = 3.The authors adopt B = Nloop + 3 for main experiments, after modeling capacity as Nloop + 1 essential states plus k auxiliary slots.Source: §E.2 · p. 22, Table 4 · p. 22

E.3 ABLATION STUDY: HEURISTIC STATE-PASSING SCHEMES

The ablation finds that heuristic state combinations are brittle, while MeSH outperforms the tested explicit schemes. Learnable linear combinations avoid the worst degradation but do not match the best hand-picked heuristic.

E.3 ABLATION STUDY: HEURISTIC STATE-PASSING SCHEMES

  • 9.24 PPL for +residual+anchor+anchor* is worse than +anchor at 9.19 PPL, while +anchor+anchor* is the best explicit scheme at 9.17 PPL.The hand-picked combination requires manual tuning; learnable static and dynamic combinations avoid the worst degradation but do not match the best heuristic.Source: §E.3 · pp. 22–23
  • 9.09 PPL for +mesh beats every tested explicit state-passing scheme on the Pile dataset.The best heuristic, +anchor+anchor*, reaches 9.17 PPL.Source: §E.3 · p. 23, Table 5 · p. 23
  • MeSH uses nonlinear retrieval and composition from its memory buffer rather than a linear combination of predefined supplementary states.The passage attributes its advantage to discovering implicit, high-dimensional combinations.Source: §E.3 · p. 23

E.4 COMPLEXITY ANALYSIS

MeSH adds minimal parameter and compute overhead in the reported Pythia-1.4B recursive configuration. The accompanying training-dynamics figure tracks accuracy across nine downstream tasks and their average at checkpoints during pre-training.

E.4 COMPLEXITY ANALYSIS

  • 61,470 router parameters add 0.005% relative to non-embedding parameters, alongside a 33.33% reduction in those parameters versus vanilla.The reported setup is Pythia-1.4B with a 4+8R2+4 configuration.Source: §E.4 · pp. 23–24, Table 6 · p. 23
  • +0.014% FLOPs is MeSH’s reported single-forward-pass overhead relative to the base recursive model.For Nloop = 2, batch size 1, sequence length 4096, hidden size 2048, and buffer length 5, the estimate is approximately 0.755 GFLOPs.Source: §E.4 · p. 24, Table 7 · p. 24

E.5 DETAILED TRAINING DYNAMICS ON DOWNSTREAM TASKS

  • Figure 9 reports 1.4B-parameter models’ performance on nine individual downstream tasks and their average accuracy, evaluated at various checkpoints throughout pre-training, providing a more granular view of training dynamics presented in Section 4.3.Source: §E.5 · p. 24

E.6 APPLYING MESH TO NON-RECURSIVE ARCHITECTURES

On Pythia-1.4B, MeSH improves the reported non-recursive Vanilla baseline’s Pile perplexity, and the paper reports that recursive+mesh also surpasses a larger Vanilla counterpart despite fewer parameters.

E.6 APPLYING MESH TO NON-RECURSIVE ARCHITECTURES

  • 7.26 PPL for vanilla+mesh is lower than 7.44 PPL for the standard Vanilla baseline on the Pile.The comparison applies MeSH at boundaries between conceptual blocks in the non-recursive Pythia-1.4B model.Source: §E.6 · p. 24, Table 8 · p. 26
  • The recursive+mesh Pythia-1.4B model surpasses its larger Vanilla counterpart with 33.3% fewer parameters.The passage presents this as a main-experiment result and says the non-recursive experiment shows MeSH can provide an uplift without parameter sharing.Source: §E.6 · pp. 24, 26
  • The authors identify general architectural use of explicit, routed state management as a direction for future research.The supplied non-recursive evidence is described as indicating broader applicability, while exploration beyond recursive transformers remains future work.Source: §E.6 · p. 26

E.7 APPLYING MESH TO MOE-BASED ARCHITECTURES

On an OLMoE recursive backbone, MeSH mitigates the performance drop from naive recursion and recovers performance comparable to the vanilla counterpart under matched compute. A supplemental subspace analysis contrasts redundant naive-loop states with distinct MeSH states.

E.7 APPLYING MESH TO MOE-BASED ARCHITECTURES

  • MeSH recovers OLMoE recursive performance to a level comparable to the vanilla counterpart under matched compute, mitigating the drop from naive recursion.The comparison uses an OLMoE model with 2.6B total and 512M activated parameters and a 4+8R2+4 recursive configuration.Source: §E.7 · p. 26, Table 9 · p. 26
  • MoE routing selects a subset of feed-forward experts for computational sparsity, whereas MeSH routes information across iterations; the paper describes the approaches as orthogonal and potentially complementary.The experiment keeps MoE-specific routing settings consistent across its models.Source: §E.7 · p. 26

E.8 SUPPLEMENTAL ANALYSIS: SHARED SUBSPACE COLLAPSE

  • In Figure 10, naive recursive states collapse onto a shared subspace, while MeSH states retain distinct, high-dimensional structures.The analysis projects loop states onto a shared basis formed from concatenated iteration states; the figure reports means and standard deviations across 500 samples.Source: Figure 10 · p. 27, §E.8 · p. 27
  • The authors interpret MeSH’s distinct states as visual evidence that external memory helps avoid the representational bottleneck and supports progressive refinement across iterations.The reported interpretation links memory-based information preservation with high-dimensional computation in the loop.Source: §E.8 · p. 27

F CASE STUDY: UNPACKING THE INTERNAL DYNAMICS OF MESH

A single-sequence case study illustrates how MeSH routes information differently across iterations and tokens, then combines initial and intermediate states in later loop inputs.

F CASE STUDY: UNPACKING THE INTERNAL DYNAMICS OF MESH

  • Figure 11 shows iteration- and token-specific routing, with average Write weights directing different iterations to distinct memory slots.The example assigns Iteration 1 to Slot 1 and Iteration 2 to Slot 4; token-level weights differ from sequence averages.Source: F CASE STUDY: UNPACKING THE INTERNAL… · p. 27, Figure 11 · p. 28
  • The input to the final loop, h(3), combines the prelude output with core outputs from previous loops.The composition analysis describes the resulting state as retaining a connection to the initial state while incorporating intermediate information.Source: F CASE STUDY: UNPACKING THE INTERNAL… · p. 28

G LIMITATIONS AND FUTURE WORK

The authors identify limits in the tested scale and training setting, and leave broader evaluation on non-recursive backbones for future work. They also describe the limited role of LLMs in producing the paper.

G LIMITATIONS AND FUTURE WORK

  • MeSH was evaluated up to Pythia-6.9B on the deduplicated Pile; persistence of its parameter-efficiency gains at larger scales or under different training regimes remains unclear.A comprehensive investigation beyond recursive backbones is outside this paper’s scope, despite ablation evidence of benefits to non-recursive transformers.Source: G LIMITATIONS AND FUTURE WORK · p. 28

H STATEMENT ON LARGE LANGUAGE MODELS USAGE

  • LLMs were used only for language editing, not research ideation, methodology, or experimental implementation.The authors state that substantive content is original and was verified by them.Source: H STATEMENT ON LARGE LANGUAGE MODELS… · p. 28

Related papers

Loading 2510.07739v2…