MemDLM: Treinamento de DLM Aprimorado por Memória
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
Resumo
Os Modelos de Linguagem de Difusão (DLMs) oferecem vantagens atrativas em relação aos modelos auto-regressivos (AR), como a decodificação paralela com atenção completa e geração flexível. No entanto, eles sofrem com uma notável incompatibilidade entre treinamento e inferência: os DLMs são treinados com um objetivo estático de previsão mascarada em etapa única, mas são implantados por meio de uma trajetória progressiva de desruído em múltiplas etapas. Propomos o MemDLM (DLM Aprimorado por Memória), que reduz essa lacuna incorporando um processo simulado de desruído no treinamento via Otimização Bi-nível. Um laço interno atualiza um conjunto de pesos rápidos, formando uma Memória Paramétrica que captura a experiência da trajetória local de cada amostra, enquanto um laço externo atualiza o modelo base condicionado a essa memória. Ao transferir a pressão de memorização das representações de tokens para os parâmetros, o MemDLM resulta em convergência mais rápida e menor perda de treinamento. Além disso, o laço interno pode ser reativado no momento da inferência como uma etapa de adaptação, gerando ganhos adicionais na compreensão de contexto longo. Descobrimos que, quando ativada no momento da inferência, esta Memória Paramétrica atua como um mecanismo emergente de recuperação em pesos, ajudando o MemDLM a reduzir ainda mais os gargalos de atenção a nível de token em tarefas desafiadoras de recuperação do tipo "Agulha no Palheiro". 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.