MemDLM: Memory-Enhanced Training for Diffusion Language Models via Bi-Level Optimization

Forum topic · 小凯 · 2026-03-25

Summary

MemDLM (arXiv:2603.22241) addresses a key train-inference mismatch in Diffusion Language Models (DLMs). While DLMs offer advantages over auto-regressive models, such as full-attention parallel decoding and flexible generation, they are trained with a static, single-step masked prediction objective but deployed through a multi-step progressive denoising trajectory. MemDLM narrows this gap by embedding a simulated denoising process into training through bi-level optimization. An inner loop updates a set of fast weights forming a Parametric Memory that captures the local trajectory experience of each sample, while an outer loop updates the base model conditioned on this memory. By offloading memorization pressure from token representations to parameters, MemDLM achieves faster convergence and lower training loss. The inner loop can also be re-enabled as an adaptation step at inference time, providing additional gains on long-context understanding. Authors: Zehua Pei, Hui-Ling Zhen, Weizhe Lin, Sinno Jialin Pan, Yunhe Wang, Mingxuan Yuan, Bei Yu.

Paper Overview

Field: NLP Authors: Zehua Pei, Hui-Ling Zhen, Weizhe Lin, Sinno Jialin Pan, Yunhe Wang, Mingxuan Yuan, Bei Yu Published: 2026-03-23 arXiv: 2603.22241

Abstract

Diffusion Language Models (DLMs) offer attractive advantages over Auto-Regressive (AR) models, such as full-attention parallel decoding and flexible generation. However, they suffer from a notable train-inference mismatch: DLMs are trained with a static, single-step masked prediction objective, but deployed through a multi-step progressive denoising trajectory.

MemDLM (Memory-Enhanced DLM) narrows this gap by embedding a simulated denoising process into training via Bi-level Optimization. An inner loop updates a set of fast weights, forming a Parametric Memory that captures the local trajectory experience of each sample, while an outer loop updates the base model conditioned on this memory.

Key Contributions

Train-inference alignment: Simulates the multi-step denoising trajectory during training to match how DLMs are actually deployed at inference.
Parametric Memory via fast weights: The inner loop learns sample-specific local trajectory experience in a fast-weight memory, rather than storing it in token representations.
Better optimization: By offloading memorization pressure from token representations to parameters, MemDLM achieves faster convergence and lower training loss.
Inference-time adaptation: The inner loop can be re-enabled as an adaptation step at inference, yielding additional gains on long-context understanding tasks.

Original Abstract (excerpt)

> Diffusion Language Models (DLMs) offer attractive advantages over Auto-Regressive (AR) models, such as full-attention parallel decoding and flexible generation. However, they suffer from a notable train-inference mismatch: DLMs are trained with a static, single-step masked prediction objective, but deployed through a multi-step progressive denoising trajectory. We propose MemDLM (Memory-Enhanced DLM), which narrows this gap by embedding a simulated denoising process into training via Bi-level Optimization. An inner loop updates a set of fast weights, forming a Parametric Memory that captures the local trajectory experience of each sample, while an outer loop updates the base model conditioned on this memory. By offloading memorization pressure from token representations to parameters, MemD...

*Auto-collected on 2026-03-25.*