MemDLM: Geheugenverbeterde DLM-training
MemDLM: Memory-Enhanced DLM Training
March 23, 2026
Auteurs: Zehua Pei, Hui-Ling Zhen, Weizhe Lin, Sinno Jialin Pan, Yunhe Wang, Mingxuan Yuan, Bei Yu
cs.AI
Samenvatting
Diffusie-taalmodellen (DLM's) bieden aantrekkelijke voordelen ten opzichte van auto-regressieve (AR) modellen, zoals parallelle decodering met volledige aandacht en flexibele generatie. Ze lijden echter onder een opmerkelijk train-inferentie-mismatch: DLM's worden getraind met een statisch, enkelstaps doel van gemaskeerde voorspelling, maar worden ingezet via een meerstaps progressief denoisetraject. Wij stellen MemDLM (Memory-Enhanced DLM) voor, dat deze kloof verkleint door een gesimuleerd denoiseproces in de training in te bedden via Bi-level Optimalisatie. Een binnenste lus werkt een set snelle gewichten bij, die een Parametrisch Geheugen vormen dat de lokale trajectervaring van elk sample vastlegt, terwijl een buitenste lus het basismodel update, conditioneel op dit geheugen. Door de memorisatiedruk van tokenrepresentaties naar parameters te verplaatsen, levert MemDLM een snellere convergentie en een lager trainingsverlies op. Bovendien kan de binnenste lus tijdens inferentie opnieuw worden geactiveerd als een aanpassingsstap, wat extra winst oplevert bij begrip van lange context. Wij ontdekken dat, wanneer geactiveerd tijdens inferentie, dit Parametrisch Geheugen fungeert als een emergent retrievalmechanisme in de gewichten, waardoor MemDLM token-level aandachtsknelpunten verder vermindert bij uitdagende Needle-in-a-Haystack-retrievaltaken. 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.