54.0
Best overall average
MetaState + Soft Diffusion reaches the strongest aggregate result, supporting that within-step soft representations and cross-step memory are complementary.
vs. 52.6 Soft Diffusion · 53.6 MetaState
COLM 2026
arXiv
MetaState: Persistent Working Memory Enhances Reasoning in Discrete Diffusion Language Models
A lightweight recurrent interface gives frozen diffusion language models a continuous workspace that survives sampling and remasking.
1Georgia Institute of Technology · 2University of Massachusetts Amherst · 3Harvard University
+4.5pp
average reasoning gain
~0.6%
trainable parameters
2
dLLM families
4
reasoning benchmarks
The core problem
Reasoning is lost between diffusion steps.
Standard dLLMs compress rich hidden activations into sparse discrete tokens after every forward pass. MetaState adds a parallel continuous path so useful intermediate state can persist.
In brief
Abstract
Discrete diffusion language models generate text by iteratively denoising a masked sequence, but each step sees only the current hard-masked tokens. The richer continuous representation computed at the previous step is discarded. We call this recurring bottleneck the Information Island issue.
MetaState equips a frozen dLLM with persistent, fixed-size working memory. A Mixer reads backbone activations, a gated Updater carries selected information forward, and an Injector writes the state back into the next step. A dedicated K-step unrolling objective teaches the memory to evolve across the trajectory.
Across Dream-7B and LLaDA-8B, MetaState improves mathematical reasoning and code generation in every main comparison—an average gain of 4.5 percentage points—while training only about 0.6% of parameters.
01
Identify the Information Island bottleneck at the discrete inter-step interface.
02
Introduce constant-size recurrent memory around an otherwise frozen backbone.
03
Validate consistent gains across two model families, four tasks, and varied decoding settings.
Method
A recurrent loop around a frozen backbone.
The state contains a fixed number of slots independent of sequence length. Every component operates in a compact bottleneck dimension, leaving the pretrained model weights untouched.
-
01
Read · Mixer
Memory slots cross-attend to the final-layer hidden activations, compressing variable-length context into a fixed-size state.
-
02
Retain · Updater
A time-conditioned GRU-style update gate balances prior state against new evidence at each denoising step.
-
03
Write · Injector
Position-aware cross-attention adds memory-conditioned modulation to token embeddings through a zero-initialized residual path.
Memory64 slots × 1,024 dims
Training4-step unrolling
Data50K Tülu-3 sequences
BackboneFrozen throughout
Main results
Consistent gains in math and code.
MetaState improves all 16 backbone–benchmark comparisons in the main evaluation, with especially large gains for base models and tasks vulnerable to long-horizon drift.
Dream-7B
Accuracy / Pass@1| Model | GSM8K | MATH-500 | HumanEval | MBPP |
|---|---|---|---|---|
| Base | 73.7 | 37.6 | 54.9 | 52.6 |
| + MetaState | 76.7 | 46.4 | 59.2 | 53.6 |
| Gain | +3.0 | +8.8 | +4.3 | +1.0 |
| Instruct | 74.8 | 45.0 | 56.1 | 51.0 |
| + MetaState | 78.1 | 46.6 | 59.8 | 55.0 |
| Gain | +3.3 | +1.6 | +3.7 | +4.0 |
LLaDA-8B
Accuracy / Pass@1| Model | GSM8K | MATH-500 | HumanEval | MBPP |
|---|---|---|---|---|
| Base | 67.4 | 28.8 | 33.5 | 25.6 |
| + MetaState | 77.9 | 37.0 | 39.6 | 33.0 |
| Gain | +10.5 | +8.2 | +6.1 | +7.4 |
| Instruct | 78.5 | 36.8 | 37.2 | 26.0 |
| + MetaState | 79.5 | 37.8 | 39.6 | 32.2 |
| Gain | +1.0 | +1.0 | +2.4 | +6.2 |
GSM8K and MATH-500 measure mathematical reasoning; HumanEval and MBPP report Pass@1 functional correctness for code generation.
Evidence beyond the headline
The memory signal is robust and complementary.
Additional experiments evaluate compatibility, robustness, transfer, and practical cost.
28 / 28
Positive threshold comparisons
MetaState improves every matched GSM8K setting across confidence thresholds from 0.60 to 0.95.
Across Dream and LLaDA · Base and Instruct+7.4pp
Transfer to LLaDA 1.5
The same interface improves all four metrics on a newer RL-optimized backbone, with the largest gain on MBPP.
Also +1.82 GSM8K · +0.61 HumanEval · +0.60 MATHWhat the ablations show
Persistence—not extra parameters alone—drives the gain.
- Recurrence is essential. Removing prior state drops the eight-task average from 53.6 to 47.6 and is especially harmful for LLaDA.
- Trajectory-aware training matters. Detaching state between steps reduces the average to 52.1, with the largest losses on MATH-500.
- Selective read and write matter. Parameter-matched MLP replacements underperform the attention-based Mixer and Injector.
Practical overhead
Small memory footprint, measurable latency.
The persistent state stays constant with sequence length. The recurrent modules add modest but non-zero work at every denoising step.
- Module footprint
- 89–91 MiB
- State per sequence
- 0.125 MiB
- Sample latency
- +13.8–15.5%
- Allocated memory
- ~+110 MiB
Current limitation. Training uses one warmup plus K sequential backbone passes, and inference executes all three recurrent modules at each denoising step. Kernel fusion and hardware-aware scheduling are promising directions for reducing this overhead.
Reference
Citation
@misc{xia2026metastate,
title = {MetaState: Persistent Working Memory Enhances Reasoning in
Discrete Diffusion Language Models},
author = {Kejing Xia and Mingzhe Li and Lixuan Wei and Zhenbang Du and
Xiangchi Yuan and Dachuan Shi and Qirui Jin and Wenke Lee},
year = {2026},
eprint = {2603.01331},
archivePrefix = {arXiv},
primaryClass = {cs.CL},
url = {https://arxiv.org/abs/2603.01331}
}