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

TL;DR

MetaState introduces a persistent working memory to discrete diffusion language models, significantly improving reasoning capabilities by enabling information flow across denoising steps.

Contribution

It proposes MetaState, a lightweight recurrent module that enhances reasoning in frozen dLLMs without retraining the entire model.

Findings

MetaState improves reasoning performance by an average of 4.5% across benchmarks.

It adds only about 0.6% trainable parameters to the backbone.

MetaState effectively propagates intermediate reasoning states across denoising steps.

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 \textbf{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 \textbf{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 \textbf{Mixer} that reads backbone activations into memory slots, a GRU-style \textbf{Updater} that integrates information across steps, and a cross-attention \textbf{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\%$ across all evaluations.