Research map / Architectures and efficiency / Attention and transformers

MemDLM: Memory-Enhanced DLM Training

NeurIPS 2026 · Architectures and efficiency › Attention and transformers · approach: Sparse attention for long-context LLMs · 0 citations (OpenAlex)

Read the paperCheck your idea against 108,967 papers

AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.

Proposes MemDLM, which simulates a two-step denoising trajectory in an inner loop that writes fast LoRA weights (Parametric Memory) and uses them to condition an outer MDLM update, improving long-context retrieval of diffusion LMs.

Problem

Standard masked diffusion LM training uses a static single-step objective that does not expose the model to inference-time progressive denoising, and token-space attention becomes diluted over long contexts.

Approach

A bi-level optimization where an inner loop (two gradient steps on transient LoRA fast weights, pre-anchor alignment then anchor-to-target) builds a Parametric Memory of a local denoising trajectory, and the outer loop updates base weights conditioned on it using a first-order approximation; the fast weights can be discarded at inference or re-run on the prompt.

Key results

• On LLaDA-MoE, MemDLM raises RULER-VT at 8K from 78.84% to 95.80% (Train & Inference).
• On LLaDA2.1, BABILong at 8K improves from 47.40% to 57.00% with MemDLM (Train & Inference); Train-Only reaches 55.00%.
• Train-Only MemDLM already improves over Standard MDLM on LongBench for LLaDA-MoE (e.g., TriviaQA 55.29 to 87.74), with the inner loop disabled at inference.

Surprising findings

• Lower inner-loop training loss from more pre-anchor steps (3- and 4-step) does not translate to better downstream BABILong; accuracy drops from 0.684 to 0.590.
• Full-parameter inner-loop updates reach the lowest train loss but underperform a restricted FFN-only update on BABILong-1K.

Limitations

• About 2.1x wall-clock training overhead from the inner loop, and added latency when inner loop is re-run at inference.
• Experiments limited to two DLM backbones and one instruction-tuning dataset (LongAlpaca); no pretraining-stage or longer-context training tested.
• Inner-loop adaptation at inference is restricted to the prompt; extending to decoding is future work.

Experimental setup

Datasets: LongAlpaca

Benchmarks: RULER BABILong LongBench

Metrics: accuracy train loss eval loss exposure bias ratio

Baselines: Standard MDLM untuned LLaDA-MoE-7B-A1B-Base

Builds on: Masked diffusion language models (MDLM) LoRA LLaDA dllm

Closest papers by meaning

• Reinforcement Learning for Diffusion LLMs with Entropy-Guided Step Selection and Stepwise Advantages (NeurIPS 2026 · similarity 0.88)
• Looped Diffusion Language Models (NeurIPS 2026 · similarity 0.88)
• Efficient Test-Time Scaling via Hierarchical Search and Self-Verification for Discrete Diffusion Language Models (ICML 2026 · similarity 0.88)
• dKV-Cache: The Cache for Diffusion Language Models (NeurIPS 2025 · similarity 0.88)
• Denoising Time Matters: Diverse Generation in Diffusion Language Models (NeurIPS 2026 · similarity 0.88)
• SemanticDLM+: Improving Diffusion LLMs through Bias-variance Trade-off in Transition Kernel Design (NeurIPS 2026 · similarity 0.88)