MemDLM: Addestramento di DLM Potenziato con Memoria

MemDLM: Memory-Enhanced DLM Training

March 23, 2026

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

cs.AI

Abstract

I modelli linguistici di diffusione (DLM) offrono vantaggi interessanti rispetto ai modelli auto-regressivi (AR), come il decoding parallelo con attenzione completa e una generazione flessibile. Tuttavia, soffrono di un notevole disallineamento addestramento-inferenza: i DLM vengono addestrati con un obiettivo statico di predizione mascherata a singolo passo, ma vengono utilizzati attraverso una traiettoria di denoising progressivo multi-passo. Proponiamo MemDLM (DLM potenziato con memoria), che riduce questo divario incorporando un processo di denoising simulato nell'addestramento tramite Ottimizzazione Bi-livello. Un ciclo interno aggiorna un insieme di pesi rapidi, formando una Memoria Parametrica che cattura l'esperienza della traiettoria locale di ciascun campione, mentre un ciclo esterno aggiorna il modello base condizionato da questa memoria. Scaricando la pressione di memorizzazione dalle rappresentazioni dei token ai parametri, MemDLM produce una convergenza più rapida e una loss di addestramento inferiore. Inoltre, il ciclo interno può essere riattivato durante l'inferenza come passo di adattamento, producendo ulteriori miglioramenti nella comprensione di contesti lunghi. Scopriamo che, quando attivata durante l'inferenza, questa Memoria Parametrica agisce come un meccanismo emergente di retrieval nei pesi, aiutando MemDLM a ridurre ulteriormente i colli di bottiglia dell'attenzione a livello di token su impegnativi compiti di retrieval "Ago nel Pagliaio". Codice: 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.