MemDLM: Entrenamiento de DLM Mejorado con Memoria
MemDLM: Memory-Enhanced DLM Training
March 23, 2026
Autores: Zehua Pei, Hui-Ling Zhen, Weizhe Lin, Sinno Jialin Pan, Yunhe Wang, Mingxuan Yuan, Bei Yu
cs.AI
Resumen
Los Modelos de Lenguaje de Difusión (DLM) ofrecen ventajas atractivas frente a los modelos Auto-Regresivos (AR), como la decodificación paralela con atención completa y una generación flexible. Sin embargo, adolecen de una notable discrepancia entre el entrenamiento y la inferencia: los DLM se entrenan con un objetivo estático de predicción enmascarada en un solo paso, pero se despliegan a través de una trayectoria de desruido progresivo de múltiples pasos. Proponemos MemDLM (DLM Mejorado con Memoria), que reduce esta brecha incrustando un proceso de desruido simulado en el entrenamiento mediante Optimización Bi-nivel. Un bucle interno actualiza un conjunto de pesos rápidos, formando una Memoria Paramétrica que captura la experiencia de la trayectoria local de cada muestra, mientras que un bucle externo actualiza el modelo base condicionado por esta memoria. Al descargar la presión de memorización de las representaciones de tokens a los parámetros, MemDLM logra una convergencia más rápida y una pérdida de entrenamiento menor. Además, el bucle interno puede reactivarse en el momento de la inferencia como un paso de adaptación, generando ganancias adicionales en la comprensión de contexto largo. Encontramos que, cuando se activa durante la inferencia, esta Memoria Paramétrica actúa como un mecanismo emergente de recuperación en los pesos, ayudando a MemDLM a reducir aún más los cuellos de botella de atención a nivel de token en tareas de recuperación desafiantes tipo "Aguja en un Pajar". Código: 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.