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
ICLR 2026
cs.LG, cs.AI
2025-10-09
MeSH adds a memory buffer and step-wise routers to recursive Transformers. At Pythia-1.4B it beats vanilla by +1.06 points 0-shot, with 33% fewer non-embedding parameters.
Recursive Transformers reuse a shared block across iterations, so compute depth is no longer tied to parameter depth. Under matched FLOPs, the cheaper models still tend to lose: higher perplexity, worse downstream scores. The parameter saving arrives with a quality tax.
Hidden-state probes on a Pythia-410M looped model pin the gap on two structural failures. Undifferentiated computation: the shared core has no notion of which iteration it is in, so the first loop does almost all the work, later updates collapse toward zero, and consecutive states stay highly similar under CKA, a fixed point. Information overload: long-lived input semantics and transient working memory share one hidden state. The singular-value spectrum of loop states decays fast, i.e. effective rank drops.
The backbone is Prelude-Recurrent-Coda. A prelude maps token embeddings into the loop, a shared core runs for K steps, and a coda produces the final representation. Naive recurrence is just h ← fcore(h). Heuristic patches add a fixed supplement: residual adds the previous state, anchor adds the loop-entry state, anchor adds raw embeddings. Those skips help with forgetting. The core remains blind to step index.
MeSH moves state management outside the core. A buffer of B slots, each the same shape as the hidden state, starts with embeddings in slot 0 and zeros elsewhere. Each iteration has its own write and read routers: one linear layer plus softmax, producing token-wise weights over the B slots. The core writes a soft combination into the buffer, then the next state is a weighted read. A transitional write-read sits between prelude and the first loop; a final read feeds the coda.
Buffer length follows B = Nloop + 3. For 4+8R2+4 that is B=5: one slot per major state plus two scratch slots. On Pythia-1.4B the routers add 61,470 parameters, about 0.005% of non-embedding weights.
Pretraining follows Pythia: one epoch on a 250B-token deduplicated Pile split, sequence length 4096, global batch 512. Downstream numbers are 0-shot and 5-shot averages over 9 tasks.
| Scale | Variant | Non-emb params | 0-shot | 5-shot | Pile PPL |
| 1.4B | Vanilla 24-layer | 100% | 49.50 | 51.93 | 7.44 |
| 1.4B | Recursive base | -33.3% | 48.89 | 50.99 | 7.63 |
| 1.4B | +anchor | -33.3% | 49.39 | 51.27 | 7.51 |
| 1.4B | +MeSH | -33.3% | 50.56 | 52.79 | 7.39 |
At 1.4B, MeSH beats vanilla by +1.06 0-shot and +0.86 5-shot, with Pile perplexity 7.39 vs 7.44. At 160M the 5-shot edge is only +0.06. A 50%-compressed 410M setup still gains +0.73 5-shot. At 2.8B and 6.9B (6+10R2+6, about 68.75% of vanilla non-embedding params) there is no vanilla control; MeSH vs recursive base is +2.22 and +2.16 0-shot.
On 410M, the best hand-tuned mix of residual/anchor/anchor lands at Pile PPL 9.17; MeSH is 9.09. Buffer-size ablation peaks at B=5. Dropping MeSH into a non-shared 1.4B vanilla (cut as 4+8+8+4) lowers Pile PPL from 7.44 to 7.26. On an OLMoE backbone (2.6B total, 512M active), MeSH beats recursive base and still sits slightly below vanilla.
Probes show more even update magnitudes across the three loops, lower inter-loop CKA, and slower singular-value decay.
The usual complaint about looped Transformers is that matched compute still loses to a vanilla stack. MeSH's claim is that the bottleneck is overloaded single-channel state, not weight sharing. External memory plus step-wise routers let the core stay fully shared. For anyone building Universal Transformer or Huginn-style recursive depth, this is a near-zero-parameter add-on. At 1.4B it already beats a larger non-recursive model; at 2.8B/6.9B the lift over recursive baselines holds.
All of this is from-scratch pretraining. There is no recipe for splicing the buffer into an existing Llama checkpoint.
The authors flag two limits. Results stop at Pythia-6.9B on the deduplicated Pile; larger models and other recipes are untested. The non-recursive case is a single 1.4B block-partition experiment, not a sweep.
6.9B has no vanilla counterpart, so the "beats the larger model" story is not shown at that scale. Routers are per-iteration parameters and loop count is fixed at train time; there is no Adaptive Computation Time halt. Slots match hidden-state shape, so buffer memory grows linearly with sequence length and B. On OLMoE, MeSH does not overtake vanilla, so the MoE dividend is thinner.