MemDLM: Обучение DLM с расширенной памятью

MemDLM: Memory-Enhanced DLM Training

March 23, 2026

Авторы: Zehua Pei, Hui-Ling Zhen, Weizhe Lin, Sinno Jialin Pan, Yunhe Wang, Mingxuan Yuan, Bei Yu

cs.AI

Аннотация

Диффузионные языковые модели (DLM) обладают привлекательными преимуществами по сравнению с авторегрессионными (AR) моделями, такими как полное параллельное декодирование с механизмом внимания и гибкость генерации. Однако они страдают от значительного несоответствия между обучением и выводом: DLM обучаются с помощью статической цели маскированного предсказания за один шаг, но развертываются через многошаговую траекторию прогрессивного шумоподавления. Мы предлагаем MemDLM (DLM с расширенной памятью), которая сокращает этот разрыв за счет внедрения смоделированного процесса шумоподавления в обучение с помощью двухуровневой оптимизации. Внутренний цикл обновляет набор быстрых весов, формируя Параметрическую Память, которая фиксирует локальный траекторный опыт каждого примера, в то время как внешний цикл обновляет базовую модель с учетом этой памяти. Перекладывая нагрузку по запоминанию с представлений токенов на параметры, MemDLM обеспечивает более быструю сходимость и меньшие потери при обучении. Более того, внутренний цикл может быть повторно активирован во время вывода в качестве шага адаптации, что дает дополнительные преимущества в понимании длинного контекста. Мы обнаружили, что при активации во время вывода эта Параметрическая Память действует как возникающий механизм поиска в весах, помогая MemDLM further снижать узкие места внимания на уровне токенов в сложных задачах поиска "иголки в стоге сена". Код: 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.