MetaState: Persistent Working Memory Enhances Reasoning in Discrete Diffusion Language Models

Kejing Xia; Mingzhe Li; Lixuan Wei; Zhenbang Du; Xiangchi Yuan; Dachuan Shi; Qirui Jin; Wenke Lee

cs.CLMar 2, 2026

MetaState: Persistent Working Memory Enhances Reasoning in Discrete Diffusion Language Models

Authors: Kejing Xia, Mingzhe Li, Lixuan Wei, Zhenbang Du, Xiangchi Yuan, Dachuan Shi, Qirui Jin, Wenke Lee

Organizations: Georgia Institute of Technology · University of Massachusetts Amherst · Harvard University

Abstract

Discrete diffusion language models (dLLMs) generate text by iteratively denoising a masked sequence. However, standard dLLMs condition each denoising step solely on the current hard-masked sequence, while intermediate continuous representations are discarded after sampling and remasking. We term this bottleneck the Information Island issue: continuous information remains isolated within individual denoising steps and fails to propagate across the trajectory. This bottleneck is especially harmful for reasoning, which requires intermediate reasoning state to be preserved and updated across many denoising steps. To address this limitation, we introduce MetaState, a lightweight recurrent augmentation that equips a frozen dLLM backbone with persistent, fixed-size working memory. MetaState comprises three modules with a shared time conditioner: a cross-attention Mixer that reads backbone activations into memory slots, a GRU-style Updater that integrates information across steps, and a cross-attention Injector that writes the updated memory back into the backbone. We train these modules with a dedicated $K$ -step unrolling pipeline to learn multi-step dynamics. MetaState adds only ${\sim}0.6\%$ trainable parameters while keeping the backbone frozen, and consistently improves reasoning performance over frozen baselines on mathematical reasoning and code generation benchmarks, with an average gain of 4.5 percentage points across all evaluations.

Figures & tables

Figure 1: The Information Island issue in discrete diffusion: sampling and remasking compress continuous hidden activations into discrete tokens, imposing a lossy bottleneck between denoising steps. MetaState addresses this issue by maintaining a persistent state across steps.

Figure 2: Performance comparison between MetaState and frozen baselines on reasoning benchmarks for LLaDA-8B and Dream-7B in both Instruct and Base versions.

Figure 3: Overview of the MetaState architecture. The three modules (Injector, Mixer, Updater) and the shared time conditioner form a recurrent loop around the frozen backbone, propagating a persistent state across denoising steps.

Model	GSM8K	MATH	HE	MBPP	GSM8K	MATH	HE	MBPP	Avg.
	Dream Backbone (7B)				LLaDA Backbone (8B)
Base	73.7	37.6	54.9	52.6	67.4	28.8	33.5	25.6	46.8
+ MetaState	76.7	46.4	59.2	53.6	77.9	37.0	39.6	33.0	52.9
$\Delta$ vs. Base	+3.0	+8.8	+4.3	+1.0	+10.5	+8.2	+6.1	+7.4	+6.2
Instruct	74.8	45.0	56.1	51.0	78.5	36.8	37.2	26.0	50.7
+ MetaState	78.1	46.6	59.8	55.0	79.5	37.8	39.6	32.2	53.6

Table 1: Main results on 4 different benchmarks (generation length 256, block size 32, dual cache). MATH = MATH-500, HE = HumanEval.

\Delta

denotes improvement over the corresponding baseline. Bold marks the best result per column within each backbone group.

Method	GSM8K	MATH	HE	MBPP	GSM8K	MATH	HE	MBPP	Avg.
	Dream-Instruct (7B)				LLaDA-Instruct (8B)
Instruct baseline	74.8	45.0	56.1	51.0	78.5	36.8	37.2	26.0	50.7
+ Soft Diffusion	79.1	45.8	60.4	52.0	78.7	37.8	37.8	29.0	52.6
+ MetaState	78.1	46.6	59.8	55.0	79.5	37.8	39.6	32.2	53.6
+ MetaState †	79.4	46.4	59.2	55.6	80.3	37.6	40.9	32.4	54.0

Table 2: Compatibility with Soft Diffusion ( Zhu et al., 2025c ) . MetaState and Soft Diffusion target different levels of the pipeline and can be combined.

\dagger

denotes ‘+ Soft Diffusion’. Bold marks the best result per column. Hyperparameter details are provided in Appendix A.7 .

	Dream-Instruct (7B)				LLaDA-Instruct (8B)
Model Variants	GSM8K	MATH	HE	MBPP	GSM8K	MATH	HE	MBPP	Avg.
Backbone only	74.8	45.0	56.1	51.0	78.5	36.8	37.2	26.0	50.7
MetaState (full)	78.1	46.6	59.8	55.0	79.5	37.8	39.6	32.2	53.6
w/o recurrence	77.6	45.0	61.0	53.4	62.6	32.6	23.8	24.6	47.6
w/o BPTT	76.7	43.4	59.2	53.2	78.2	34.4	38.4	33.0	52.1
w/o time cond.	77.2	44.8	57.9	54.6	78.9	36.2	39.0	31.8	52.6

Table 3: Ablation studies. Each row modifies one component while keeping the rest of MetaState intact. MLP variants are parameter-matched. MATH = MATH-500, HE = HumanEval. Bold : full model. Gray : backbone without MetaState (from Table 1 ).

Appendix figures & tables23 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 4: A step-by-step denoising trajectory. Each denoising step shows both the full argmax model prediction ( out ) and the remasked sequence ( in ) that is actually passed to the next step.

Figure 5: Comparison of Pass@1

(t)

and EverPass@1

(t)

across denoising steps on GSM8K with LLaDA-Instruct-8B.

Steps	Orig. PASS	Our PASS	Orig. EVER	Our EVER	gain	Orig. gap	Our gap
1–256	49.58	54.80	65.88	69.05	+5.23	16.30	14.25
256	63.00	66.00	79.00	79.00	+3.00	16.00	13.00
1–85	31.59	40.42	45.96	54.84	+8.84	14.38	14.41
86–170	56.34	59.87	73.89	74.95	+3.53	17.55	15.08
171–256	60.67	64.01	77.65	77.28	+3.34	16.98	13.27

Appendix

Table 4: GSM8K trajectory diagnostic for LLaDA-Instruct. EverPass@1 counts an example as correct if any intermediate full prediction along the denoising trajectory is correct. The gap is EverPass@1 minus Pass@1 .

Subset	Metric	Dream	LLaDA
1	Original response-step cosine, mean	0.9749	0.9775
2	Original response-step cosine, mean	0.9742	0.9769
1	MetaState state Frobenius norm, mean	31.42	52.04
2	MetaState state Frobenius norm, mean	31.21	51.98
1	MetaState mean slot norm, mean	3.93	6.50
2	MetaState mean slot norm, mean	3.90	6.50

Appendix

Table 5: GSM8K state and memory diagnostics. Two 128-example subsets show the same pattern: the original consecutive-step hidden-state cosine is nearly identical across backbones, while MetaState memory activity is larger and more variable on LLaDA.

Backbone	Target	Slice	MetaState MRR	Best control MRR	Gap
LLaDA	gold_intermediate	all	0.663	0.434	+0.229
LLaDA	gold_intermediate	not_visible	0.588	0.363	+0.225
LLaDA	gold_intermediate	not_in_question	0.646	0.410	+0.235
LLaDA	model_final	all	0.308	0.220	+0.088
LLaDA	model_final	not_visible	0.318	0.220	+0.099
LLaDA	model_final	not_in_question	0.311	0.220	+0.092

Appendix

Table 6: GSM8K numeric semantic retention probe. The probe reads the incoming MetaState memory before injection and ranks candidate numbers. Controls include random ranking, constant-state features, question-number preference, and visible-text-number preference.

Symbol	Description	Value
$M$	Number of memory slots	$64$
$D_{s}$	State dimension per slot	$1024$
$d_{c}$	Time conditioner output dimension	$1024$
$d_{m}$	Mixer bottleneck dimension	$768$
$d_{b}$	Injector bottleneck dimension	$768$
$n_{q}$	Query heads (Mixer & Injector)	$8$

Appendix

Table 7: Architecture hyperparameters for the recurrent modules. All symbols correspond to notation introduced in § 4 .

Optimization
Parameter	Value
Optimizer	AdamW
Learning rate	$2\times 10^{-5}$
LR schedule	cosine, 5% linear warmup
Gradient clipping	max norm $1.0$
$\beta_{1},\beta_{2}$	$0.9,0.95$

Appendix

Table 8: Training and inference hyperparameters (shared across LLaDA and Dream backbones).

Module / Component	LLaDA	Dream
SharedTimeConditioner	2,102,016	2,102,016
Sinusoidal MLP ( $256{\to}1024{\to}1024$ )	1,312,768	1,312,768
Content path (norm + proj + gate)	789,248	789,248
Mixer	17,899,264	17,506,048
Slot self-attention ( $D_{s}$ space, GQA)	3,146,752	3,146,752
Down-projections † ( $d_{m}{\times}D_{s}+d_{m}{\times}D$ )	3,932,160	3,538,944

Appendix

Table 9: Per-module parameter breakdown of the MetaState recurrent modules for each backbone (

D_{s}{=}1024

d_{m}{=}d_{b}{=}768

M{=}64

). Rows marked † are the only backbone-dependent components.

Backbone	Method	GSM8K	MATH	HE	MBPP
Dream-Instruct	Soft Diffusion	(0.9, 0.15)	(0.9, 0.1)	(0.9, 0.2)	(0.9, 0.2)
Dream-Instruct	+ MetaState	(0.2, 0.1)	(0.2, 0.2)	(0.9, 0.1)	(0.2, 0.1)
LLaDA-Instruct	Soft Diffusion	(0.2, 0.1)	(0.2, 0.15)	(0.9, 0.2)	(0.2, 0.2)
LLaDA-Instruct	+ MetaState	(0.2, 0.1)	(0.9, 0.2)	(0.2, 0.15)	(0.2, 0.2)

Appendix

Table 10: Best Soft Diffusion hyperparameters

(p,r_{f})

for each method, backbone, and benchmark. All values selected by grid search over

p\in\{0.2,0.9\}

and

r_{f}\in\{0.1,0.15,0.2\}

Backbone	Metric	Soft Diff. avg	Soft Diff. + Meta avg	$\Delta$
LLaDA-Instruct	GSM8K flexible EM	78.14	79.48	+1.34
LLaDA-Instruct	HumanEval pass@1	36.89	39.12	+2.23
LLaDA-Instruct	MBPP pass@1	28.27	31.77	+3.50
LLaDA-Instruct	MATH-500 math verify	37.23	37.20	-0.03
Dream-Instruct	GSM8K flexible EM	76.25	78.27	+2.02
Dream-Instruct	HumanEval pass@1	57.52	56.91	-0.61

Appendix

Table 11: Soft Diffusion average over all six grid configurations. Averaging removes dependence on a single selected

(p,r_{f})

pair.

Method	GSM8K	HumanEval	GSM8K	HumanEval	Avg.
	Dream-Instruct (7B)		LLaDA-Instruct (8B)
Instruct	74.8	56.1	78.5	37.2	61.7
+ LoRA	76.0	55.5	79.2	39.0	62.4
+ MetaState	78.1	59.8	79.5	39.6	64.3
+ MetaState †	79.4	59.2	80.3	40.9	65.0

Appendix

Table 12: Comparison with LoRA fine-tuning on Instruct backbones. LoRA (

r{=}32

\alpha{=}64

, all-linear) is trained on the same 50k Tülu-3 data with matched optimization hyperparameters.

\dagger

denotes ‘+ Soft Diffusion’. Bold marks the best result per column.

Backbone	Block size	Original	MetaState	Gain
Dream-Instruct	16	78.24	78.85	+0.61
Dream-Instruct	32	74.75	78.09	+3.34
Dream-Instruct	64	65.43	77.18	+11.75
Dream-Instruct	128	52.01	74.83	+22.82
Dream-Instruct	256	50.34	53.07	+2.73
LLaDA-Instruct	16	77.94	78.62	+0.68

Appendix

Table 13: GSM8K block-size sweep. Block size changes the per-block denoising trajectory length; MetaState is positive at all tested sizes.

Model	Threshold	Original	MetaState	Gain	Orig. fwd/seq	Meta fwd/seq
Dream-Base	0.95	73.16	76.42	+3.26	156.34	78.50
Dream-Base	0.85	70.43	75.28	+4.85	139.10	67.14
Dream-Base	0.80	70.58	74.15	+3.57	131.05	63.10
Dream-Base	0.75	69.60	73.16	+3.56	123.63	58.82
Dream-Base	0.70	66.26	70.20	+3.94	117.02	54.75
Dream-Base	0.65	64.37	67.48	+3.11	110.39	51.78

Appendix

Table 14: GSM8K confidence-threshold sweep at generation length 256, block size 32, and dual cache. Lower thresholds reduce the effective number of denoising updates. MetaState improves all matched cells.

Length	GSM8K Orig.	GSM8K Meta	Gain	MBPP Orig.	MBPP Meta	Gain
128	69.98	71.87	+1.89	49.80	52.60	+2.80
256	74.75	78.09	+3.34	51.80	55.00	+3.20
384	75.89	78.85	+2.96	50.20	55.00	+4.80
512	75.59	79.83	+4.24	50.40	54.20	+3.80
768	74.30	79.53	+5.23	51.40	55.20	+3.80
1024	75.51	78.70	+3.19	52.20	54.60	+2.40

Appendix

Table 15: Dream-Instruct generation-length sweep. MetaState remains positive on GSM8K and MBPP from 128 to 1024 tokens.

Model	Original	MetaState	Gain	Gain / pooled std
Dream-Base	$73.92\pm 0.80$	$76.00\pm 0.88$	+2.08	2.47
Dream-Instruct	$75.19\pm 0.88$	$78.64\pm 0.67$	+3.45	4.41
LLaDA-Base	$67.87\pm 0.75$	$77.67\pm 0.67$	+9.80	13.78
LLaDA-Instruct	$78.16\pm 0.86$	$78.92\pm 0.98$	+0.76	0.82

Appendix

Table 16: GSM8K multi-seed robustness. MetaState improves all four settings across seeds, with large gains on Base backbones and smaller improvements on Instruct backbones.

Figure 6: Sensitivity to the dense-reveal loss mixing ratio

\lambda_{d}

. Increasing

\lambda_{d}

from 0.6 to 0.9 yields modest drops in average scores.

Figure 7: Sensitivity to the number of memory slots

M

. Increasing

M

from 32 to 48 yields modest gains on most benchmarks, while further increasing to 64 produces modest drops.

Figure 8: Sensitivity to training data size. MetaState achieves competitive performance even with 30k training examples, and accuracy remains stable as the dataset increases to 50k.

Figure 9: Sensitivity to the unroll depth

K

. Performance is broadly stable across

K\in\{3,4,5\}

, with average accuracy varying by less than one point.

Family	Data	Runtime
Dream-Instruct	50k	2.78 h (10012.8 s)
LLaDA-Instruct	50k	1.89 h (6787.9 s)
LLaDA-1.5	50k	1.87 h (6742.1 s)

Appendix

Table 17: Training wall-clock for the 50k Tulu-3 subset. The backbone remains frozen; the listed time is the offline cost of training the external MetaState interface.

Family	Raw tok/s	Drop	Sample latency	Overhead
Dream-Instruct	162.66 $\rightarrow$ 134.03	17.60%	3096.82 $\rightarrow$ 3664.79 ms	+15.49%
LLaDA-Instruct	150.20 $\rightarrow$ 127.47	15.13%	3815.96 $\rightarrow$ 4427.21 ms	+13.82%

Appendix

Table 18: GSM8K inference throughput and latency. The current implementation adds about 14–15% latency and 15–18% raw throughput overhead under the main decoding setup.

Family	Module size	State / seq	Allocated memory	Delta	Peak delta
Dream	89.18 MiB	0.125 MiB	14526.15 $\rightarrow$ 14635.34 MiB	+109.19 MiB	+181.50 MiB
LLaDA	91.43 MiB	0.125 MiB	15292.51 $\rightarrow$ 15402.32 MiB	+109.81 MiB	+117.84 MiB

Appendix

Table 19: Memory overhead. The bf16 module footprint is about 90 MiB, and peak allocated memory rises by less than 1% for both backbones.

Benchmark	Original	MetaState	Gain
GSM8K	80.74	82.56	+1.82
HumanEval	37.80	38.41	+0.61
MBPP	28.20	35.60	+7.40
MATH-500	37.40	38.00	+0.60

Appendix

Table 20: MetaState on LLaDA 1.5, a newer RL-optimized dLLM backbone. MetaState improves all four metrics under the same protocol, supporting transfer beyond the original LLaDA and Dream checkpoints.

Diffusion Language Models

Diffusion language models generate and refine text through iterative denoising, enabling bidirectional context and parallel token updates.

See all 471 papers

Register Tokens for Bounded-State Reasoning in Diffusion Language Models
Sep 14, 2026Albert Ge, Chandan Singh, Yufan Zhuang +3Memory-Augmented Language Models Diffusion Language Models
University of Wisconsin–Madison · Microsoft Research · UC San Diego
Trace-Based On-Policy Distillation for Masked Diffusion Language Models
Jul 18, 2026Haolin Ren, Ziyang Huang, Chenhao Yuan +2Masked Diffusion Models Diffusion Model Distillation
Institute of Automation, Chinese Academy of Sciences, Beijing, China · University of Chinese Academy of Sciences, Beijing, China
Large Language Continuous Diffusion Models
Oct 2, 2026Zhihan Yang, Wei Guo, Jean-Marie Lemercier +14Continuous Diffusion Language Models Diffusion Language Model Inference
NVIDIA · Cornell University · Georgia Institute of Technology

Sep 14, 2026J/K move · Enter open · S save

Register Tokens for Bounded-State Reasoning in Diffusion Language Models

Albert Ge, Chandan Singh, Yufan Zhuang, Xiaodong Liu+2

University of Wisconsin–Madison · Microsoft Research · UC San Diego

Memory-Augmented Language Models Diffusion Language Models

Open paper AskPDF arXiv

Masked diffusion language models (dLLMs) generate text by iteratively denoising masked tokens with bidirectional attention. Extending reasoning across generation chunks normally requires keeping earlier generated text in context. We ask whether a dLLM can instead continue reasoning after that text is cleared, using only a fixed-size carried state. We implement this state as a small number of register tokens: dedicated fixed-position tokens whose continuous hidden states are trained to carry reasoning progress across generation chunks. We post-train dLLMs to decode a chunk of text, clear it while preserving the register values, and continue decoding from the prompt and carried state. In our main comparisons on LLaDA and Dream, registers outperform discrete-text carry on every benchmark, with gains of up to 8.5 points on math and 19.5 points on code. Registers are especially effective for bounded code generation, where correct programs usually span several chunks. Finally, registers can be further refined with reinforcement learning on long-horizon reasoning tasks.

Abstract

Diffusion Language Models

Explore similar work

Register Tokens for Bounded-State Reasoning in Diffusion Language Models

Trace-Based On-Policy Distillation for Masked Diffusion Language Models

Large Language Continuous Diffusion Models