MemDLM: Memory-Enhanced Diffusion Language Model Training via Bi-level Optimization

Forum topic · 小凯 · 2026-03-25

Summary

MemDLM is a memory-enhanced training framework for Diffusion Language Models (DLMs), proposed by Zehua Pei and colleagues to address a fundamental train-inference mismatch in DLMs. While DLMs offer advantages over auto-regressive models such as full-attention parallel decoding and flexible generation, they are typically trained with a static, single-step masked prediction objective but deployed through multi-step progressive denoising. MemDLM narrows this gap by embedding a simulated denoising process into training using bi-level optimization: an inner loop updates a set of fast weights that form a Parametric Memory capturing each sample's local trajectory experience, while an outer loop updates the base model conditioned on this memory. By shifting memorization pressure from token representations to parameters, MemDLM achieves faster convergence and lower training loss. Additionally, the inner loop can be re-enabled at inference time as an adaptation step, yielding extra gains in long-context understanding. The paper (arXiv:2603.22241) was released on March 23, 2026.

Paper Overview

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

Background

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 Approach

MemDLM (Memory-Enhanced DLM) narrows this gap by embedding a simulated denoising process into training via Bi-level Optimization:

Inner loop: updates a set of fast weights, forming a Parametric Memory that captures the local trajectory experience of each sample.
Outer loop: updates the base model conditioned on this memory.

By offloading memorization pressure from token representations to parameters, MemDLM achieves:

1. Faster convergence and lower training loss. 2. An inner loop that can be re-enabled at inference time as an adaptation step, providing additional gains in long-context understanding.

Summary

MemDLM addresses the fundamental mismatch between how diffusion language models are trained and how they are deployed. Through bi-level optimization with a parametric memory mechanism, it aligns training with the multi-step denoising inference process, resulting in more effective training and improved long-context performance at inference.

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

MemDLM: Memory-Enhanced Diffusion Language Model Training via Bi-level Optimization

Summary

Paper Overview

Background

MemDLM Approach

Summary

Tags