MemDLM: Gedächtnisgestütztes DLM-Training
MemDLM: Memory-Enhanced DLM Training
March 23, 2026
Autoren: Zehua Pei, Hui-Ling Zhen, Weizhe Lin, Sinno Jialin Pan, Yunhe Wang, Mingxuan Yuan, Bei Yu
cs.AI
Zusammenfassung
Diffusions-Sprachmodelle (DLMs) bieten gegenüber autoregressiven (AR) Modellen attraktive Vorteile, wie paralleles Decodieren mit Voll-Attention und flexibles Generieren. Allerdings leiden sie unter einer bemerkenswerten Trainings-Inferenz-Diskrepanz: DLMs werden mit einem statischen, einstufigen Masked-Prediction-Ziel trainiert, aber über eine mehrstufige, progressive Denoising-Trajektorie eingesetzt. Wir schlagen MemDLM (Memory-Enhanced DLM) vor, das diese Lücke verringert, indem es einen simulierten Denoising-Prozess via Bi-level-Optimierung in das Training einbettet. Eine innere Schleife aktualisiert einen Satz schneller Gewichte, die einen parametrischen Speicher (Parametric Memory) bilden und die lokale Trajektorie-Erfahrung jeder Stichprobe erfassen, während eine äußere Schleife das Basismodell konditioniert auf diesen Speicher aktualisiert. Indem der Merkdruck von Token-Repräsentationen auf Parameter verlagert wird, erzielt MemDLM schnellere Konvergenz und niedrigere Trainingsverluste. Darüber kann die innere Schleife zur Inferenzzeit als Adaptionsschritt reaktiviert werden, was zusätzliche Gewinne beim Langzeitkontext-Verständnis bringt. Wir stellen fest, dass dieser parametrische Speicher bei Aktivierung zur Inferenzzeit als emergenter Abrufmechanismus innerhalb der Gewichte wirkt und MemDLM dabei hilft, Token-Attention-Engpässe bei anspruchsvollen „Nadel-im-Heuhaufen“-Abfrageaufgaben weiter zu reduzieren. Code: https://github.com/JarvisPei/MemDLM.
English
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, MemDLM yields faster convergence and lower training loss. Moreover, the inner loop can be re-enabled at inference time as an adaptation step, yielding additional gains on long-context understanding. We find that, when activated at inference time, this Parametric Memory acts as an emergent in-weight retrieval mechanism, helping MemDLM further reduce token-level attention bottlenecks on challenging Needle-in-a-Haystack retrieval tasks. Code: https://github.com/JarvisPei/MemDLM.