MetaState: Persistent Working Memory
for Discrete Diffusion Language Models

Abstract

Discrete diffusion language models (dLLMs) generate text by iteratively denoising a masked sequence. Compared with autoregressive models, this paradigm naturally supports parallel decoding, bidirectional context, and flexible generation patterns. However, standard dLLMs condition each denoising step only on the current hard-masked sequence, while intermediate continuous representations are discarded after sampling and remasking. We refer to this bottleneck as the Information Island problem. It leads to redundant recomputation across steps and can degrade cross-step consistency.

We address this limitation with MetaState, a lightweight recurrent augmentation that equips a frozen dLLM backbone with a persistent, fixed-size working memory that remains independent of sequence length. MetaState consists of three trainable modules: a cross-attention Mixer that reads backbone activations into memory slots, a GRU-style Updater that integrates information across denoising steps, and a cross-attention Injector that feeds the updated memory back into backbone activations. We train these modules with K-step unrolling to expose them to multi-step denoising dynamics during fine-tuning.

On LLaDA-8B and Dream-7B, MetaState introduces negligible trainable parameters while keeping the backbone frozen, and it consistently improves accuracy over frozen baselines. These results demonstrate that persistent cross-step memory is an effective mechanism for bridging denoising steps and improving generation quality in discrete diffusion language models.

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