MemDLM: Memory-Enhanced DLM Training

Diffusion Language Models (DLMs) have emerged as a promising alternative to traditional Auto-Regressive (AR) models due to their parallel generation and flexible text manipulation capabilities. AR models generate text by predicting the next token sequentially, which can be slow and struggle to capture global context. In contrast, DLMs leverage full-attention parallel decoding, overcoming these limitations. However, despite their architectural advantages, DLMs face a significant train-inference mismatch. During training, DLMs optimize a static Masked Diffusion Language Modeling (MDLM) objective, receiving heavily masked text and predicting the clean sequence in a single step. In contrast, during inference, DLMs generate text through an iterative, progressive denoising trajectory, conditioning predictions on their own intermediate, noisy outputs. This mismatch leads to a discrepancy between training optimization and actual deployment, affecting model performance.

The core problem addressed by this research is the train-inference mismatch in diffusion language models (DLMs). Specifically, DLMs are trained with a static, single-step masked prediction objective, but deployed through a multi-step progressive denoising trajectory. This mismatch results in a misalignment between the optimization landscape during training and the model's actual deployment, leading to compounded errors during generation. The base model is never trained on these progressive, sequential trajectories, making it vulnerable to its own noisy predictions during inference. Addressing this mismatch is crucial for improving DLMs' training efficiency and long-context understanding capabilities.

MemDLM introduces several core innovations to address the train-inference mismatch in diffusion language models:

• �� Bi-level Optimization Framework: MemDLM employs a bi-level optimization framework to embed a simulated denoising process into training. The inner loop updates fast weights, forming a parametric memory that captures local trajectory experience, while the outer loop updates the base model conditioned on this memory.
• �� Parametric Memory Mechanism: By offloading memorization pressure from token representations to parameters, MemDLM reduces the burden on the base model, enhancing training efficiency and long-context understanding.
• �� Inference-Time Adaptation: MemDLM provides an additional adaptation pathway during inference by re-enabling the inner loop, further improving long-context understanding and reducing token-level attention bottlenecks.

MemDLM employs a bi-level optimization framework to embed a simulated denoising process into training, with the following steps:

• �� Inner Loop: Updates a set of fast weights, forming a parametric memory that captures the local trajectory experience of each sample. Fast weights dynamically accumulate sample-specific contextual details through gradient descent, resulting in a stable parametric state.
• �� Outer Loop: Updates the base model conditioned on the memory formed by the inner loop. By offloading part of the local memorization burden to fast weights, the base model no longer relies solely on vulnerable token-space representations to preserve context.
• �� Inference-Time Adaptation: The inner loop can be re-enabled at inference time as an adaptation step, further enhancing long-context understanding. The parametric memory acts as an emergent in-weight retrieval mechanism, helping MemDLM reduce token-level attention bottlenecks in complex tasks.

The experimental design includes long-context information retrieval tasks on the LLaDA-MoE and LLaDA2.1 backbones. Datasets used include RULER and BABILong, focusing on evaluating model performance in 'Needle-in-a-Haystack' tasks. In experiments, MemDLM improves RULER Variable Tracking accuracy at 8K length from 78.8% to 95.8% on the LLaDA-MoE backbone, and BABILong accuracy at 8K from 54.0% to 61.0% on LLaDA2.1. Ablation studies are conducted to validate the effectiveness of the parametric memory mechanism and the inference-time adaptation effect of the inner loop.

Experimental results show that MemDLM excels in long-context information retrieval tasks, significantly reducing token-level attention bottlenecks in challenging 'Needle-in-a-Haystack' tasks. On the LLaDA-MoE backbone, MemDLM improves RULER Variable Tracking accuracy at 8K length from 78.8% to 95.8%. On LLaDA2.1, it improves BABILong accuracy at 8K from 54.0% to 61.0%. These results demonstrate that the parametric memory mechanism significantly improves the base model's long-context representation capability even without enabling the inner loop during training.

MemDLM has wide-ranging applications in long-context information retrieval and understanding tasks. Its outstanding performance in complex tasks makes it suitable for high-precision information retrieval scenarios such as legal document analysis, scientific literature retrieval, and large database querying. Additionally, MemDLM's flexibility and adaptability make it valuable for natural language processing tasks requiring long-context processing, such as multi-document QA, summarization, and code completion.

Despite MemDLM's excellent performance in long-context tasks, it still faces performance degradation in extremely long texts. Additionally, the method increases computational complexity, particularly during inference when the inner loop is enabled, potentially slowing down inference speed. Future research directions include optimizing MemDLM's computational efficiency, particularly the inner loop during inference, and exploring its application on larger datasets.