MemDLM: Memory-Enhanced DLM Training

Zehua Pei; Hui-Ling Zhen; Weizhe Lin; Sinno Jialin Pan; Yunhe Wang; Mingxuan Yuan; Bei Yu

cs.CLMar 23, 2026

MemDLM: Memory-Enhanced DLM Training

Authors: Zehua Pei, Hui-Ling Zhen, Weizhe Lin, Sinno Jialin Pan, Yunhe Wang, Mingxuan Yuan, Bei Yu

Organizations: The Chinese University of Hong Kong · Huawei Technologies Co., Ltd

Abstract

Diffusion Language Models (DLMs) offer attractive advantages over Auto-Regressive (AR) models, such as full-attention parallel decoding and flexible generation. However, standard DLM training uses a static, single-step masked prediction objective that never exposes the model to the progressive denoising dynamics of inference, and forces all contextual information to be maintained purely through token-space attention, which becomes increasingly diluted as context length grows. We propose MemDLM (Memory-Enhanced DLM), which introduces a second memory channel by embedding a simulated denoising trajectory 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, while an outer loop updates the base model conditioned on this memory. By offloading part of the memorization burden from token-space attention to parameter space, MemDLM yields faster convergence, stronger long-context representations, and lower training loss, even when the fast weights are discarded at inference time. Re-enabling the inner loop at inference provides an additional prompt-specific adaptation effect, where the Parametric Memory acts as an emergent in-weight retrieval mechanism on challenging Needle-in-a-Haystack tasks. Code: https://github.com/JarvisPei/MemDLM.

Figures & tables

Figure 1: Needle-in-a-Haystack results overview. Gray bars denote Standard MDLM and blue bars denote MemDLM. Left: detailed results on RULER-MV, RULER-VT, RULER-CWE, and BABILong for the LLaDA-MoE-7B-A1B-Base and LLaDA2.1-mini backbones. Right: mean absolute improvement of MemDLM over Standard MDLM for each task, averaged across the evaluated context lengths within each backbone.

Figure 2: Overview of MemDLM. Left: standard MDLM training uses a static single-step denoising objective from $x_{t}$ to $x_{0}$ . Right: MemDLM uses Bi-level Optimization in which an inner loop updates fast weights $\phi$ along an anchor-consistent local trajectory ( $x_{t_{\textit{pre}}}\rightarrow x_{t}\rightarrow x_{0}$ ), and the outer loop updates the base model $\theta$ on the anchor state $x_{t}$ conditioned on this parametric memory. Legend: dark tokens denote mask tokens, light tokens denote observed tokens, straight arrows denote forward or reverse prediction flow, and blue curved arrows denote inner-loop fast-weight updates.

Backbone	Method	RULER-MV		RULER-VT		RULER-CWE		BABILong
Backbone	Method	4K	8K	4K	8K	4K	8K	2K	4K	8K
LLaDA-MoE	Standard MDLM	$87.90$	$92.50$	$99.40$	$78.84$	$69.10$	$55.18$	$46.00$	$45.60$	$41.80$
	MemDLM (Train-Only)	$95.85$	$96.95$	$99.52$	$94.84$	$75.72$	$58.20$	$59.40$	$50.60$	$45.60$
	MemDLM (Train & Inference)	$\mathbf{96.80}$	$\mathbf{97.00}$	$\mathbf{99.60}$	$\mathbf{95.80}$	$\mathbf{83.70}$	$\mathbf{61.96}$	$\mathbf{61.20}$	$\mathbf{51.20}$	$\mathbf{47.00}$
LLaDA2.1	Standard MDLM	$98.65$	$98.00$	$59.32$	$43.72$	$72.34$	$67.08$	$56.00$	$48.80$	$47.40$
	MemDLM (Train-Only)	$99.65$	$99.00$	$88.56$	$59.15$	$78.96$	$68.02$	$63.50$	$58.50$	$55.00$

Table 1: Needle-in-a-Haystack results on RULER and BABILong. We report Standard MDLM, MemDLM (Train-Only), and MemDLM (Train & Inference). Bold indicates the best result within each backbone block.

Method	RULER-MV		RULER-VT		RULER-CWE		BABILong
Method	16K	32K	16K	32K	16K	32K	16K	32K
Standard MDLM	$22.35$	$14.30$	$52.56$	$9.48$	$44.20$	$17.82$	$19.20$	$6.80$
MemDLM (Train-Only)	$25.48$	$15.20$	$55.30$	$10.35$	$54.25$	$22.48$	$21.00$	$8.50$
MemDLM (Train & Inference)	$\mathbf{29.40}$	$\mathbf{15.35}$	$\mathbf{56.84}$	$\mathbf{11.44}$	$\mathbf{57.24}$	$\mathbf{24.48}$	$\mathbf{22.20}$	$\mathbf{9.00}$

Table 2: Length extrapolation on Needle-in-a-Haystack tasks with the LLaDA-MoE backbone. MemDLM continues to outperform Standard MDLM at 16K and 32K across RULER and BABILong.

Method	TriviaQA	PR-en	PR-zh	MF-en	MF-zh	2WikiQA	DuReader	MultiNews	TREC	SAMSum	LCC	RB-P
Standard MDLM	$55.29$	$50.32$	$74.50$	$29.22$	$42.47$	$21.85$	$22.13$	$23.52$	$70.00$	$13.59$	$62.55$	$54.65$
MemDLM (Train-Only)	$87.74$	$54.36$	$86.29$	$31.44$	$42.87$	$22.43$	$22.19$	$26.35$	$70.50$	$15.80$	$64.20$	$55.05$
MemDLM (Train & Inference)	$\mathbf{87.77}$	$\mathbf{54.69}$	$\mathbf{87.38}$	$\mathbf{31.97}$	$\mathbf{43.28}$	$\mathbf{22.61}$	$\mathbf{22.34}$	$\mathbf{26.70}$	$\mathbf{71.00}$	$\mathbf{16.14}$	$\mathbf{64.25}$	$\mathbf{55.23}$

Table 3: LongBench results on LLaDA-MoE. Train-Only uses Parametric Memory during training but disables it at inference; Train & Inference reactivates the inner loop on the prompt.

Figure 3: Comparison with the untuned pretrained LLaDA-MoE-7B-A1B-Base model across context lengths.

Figure 4: Training dynamics on LLaDA-MoE and LLaDA2.1. Faint train-loss curves show raw values and bold curves show smoothed trends. Across both backbones, MemDLM converges faster and reaches lower train and evaluation loss.

Figure 7: Sensitivity to the inference anchor ratio.

Figure 8: Consistency of the trajectory.

Appendix figures & tables2 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 11: Exposure Bias Ratio (

\mathcal{R}_{\text{EB}}

) across denoising steps. Standard MDLM degrades rapidly, while MemDLM remains substantially flatter.

Adaptation scope variants (LLaDA-MoE)
Backbone	Configuration	Wall Time (h)	s/step	Overhead	BABILong-1K
LLaDA-MoE	Standard MDLM	$0.47$	$1.78$	$1.0\times$	$0.538$
LLaDA-MoE	MemDLM (default)	$0.97$	$3.67$	$2.1\times$	$0.684$
LLaDA2.1	Standard MDLM	$1.60$	$6.08$	$1.0\times$	$0.646$
LLaDA2.1	MemDLM (default)	$3.40$	$12.87$	$2.1\times$	$0.706$
	FFN, last 5% layers	$0.92$	$3.51$	$2.0\times$	$0.616$

Appendix

Table 4: Training workload comparison. Wall-clock time and per-step cost are measured from TensorBoard logs. Overhead is relative to the Standard MDLM baseline of each backbone.

Diffusion Language Model Training

See all 40 papers

ALoDLM: Adaptively Looped Diffusion Language Models
Oct 3, 2026Liancheng Fang, Zhuowei Li, Youngeun Kim +10Adaptive Computation for Diffusion Models Diffusion Language Models
University of Illinois Chicago · Amazon AGI · Korea University
Data-Efficient Autoregressive-to-Diffusion Language Models via On-Policy Distillation
Jun 4, 2026Xingyu Su, Jacob Helwig, Shubham Parashar +6Language Model Distillation Diffusion Language Models
Department of Computer Science and Engineering, Texas A&M University · Department of Bioinformatics and Systems Medicine,2026 University of Texas Health Science Center at Houston · Department of Electrical and Computer Engineering, Texas A&M UniversityJun 4
DyLLM: Efficient Diffusion LLM Inference via Saliency-based Token Selection and Partial Attention
Mar 9, 2026Younjoo Lee, Seungkyun Dan, Junghoo Lee +2Diffusion Language Model Inference Diffusion Language Model Decoding
Seoul National University, Seoul, Republic of Korea

Oct 3, 2026J/K move · Enter open · S save

ALoDLM: Adaptively Looped Diffusion Language Models

Liancheng Fang, Zhuowei Li, Youngeun Kim, Tianchen Zhao+9

University of Illinois Chicago · Amazon AGI · Korea University

Adaptive Computation for Diffusion Models Diffusion Language Models Diffusion Language Model Inference Diffusion Language Model Decoding Diffusion Language Model Training Parallel Decoding for Diffusion Language Models

Open paper AskPDF arXiv

Diffusion language models (DLMs) enable fast generation by predicting multiple tokens in parallel, but their practical adoption remains limited by a persistent quality gap relative to comparably sized autoregressive (AR) models. We attribute this gap to a computation-difficulty mismatch: within a partially observed sequence, some unknown tokens are easy to predict, while others require substantially more computation. Existing DLMs nevertheless apply uniform computational depth to all unknown positions at each denoising step. We introduce ALoDLM, which replaces uniform computation with token-adaptive latent recurrence. At each denoising step, ALoDLM iteratively refines latent representations and allocates computation according to token difficulty. Tokens ready to commit are fed back as discrete context, while unresolved tokens retain and further refine their latent states through additional recurrent passes. To learn token prediction and computation allocation jointly, we formulate token-wise computation schedules as latent variables and derive a conditional negative evidence lower bound (NELBO). We train ALoDLM at 1.7B and 8B parameter scales. Across eleven benchmarks, ALoDLM outperforms all evaluated DLMs and the corresponding AR baselines in average benchmark score at both scales. ALoDLM also retains fast parallel decoding, yielding a strong quality-efficiency trade-off among evaluated autoregressive and diffusion models under optimized inference engines.

Abstract

Diffusion Language Model Training

Explore similar work

ALoDLM: Adaptively Looped Diffusion Language Models

Data-Efficient Autoregressive-to-Diffusion Language Models via On-Policy Distillation

DyLLM: Efficient Diffusion LLM Inference via Saliency-based Token Selection and Partial Attention