PaperlayerView exact revision on arXiv

Source-linked AI summary

MemDLM: Memory-Enhanced DLM Training

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

arXiv:2603.22241v2cs.CL

TL;DR

Standard DLM training does not expose models to inference-time progressive denoising and relies on increasingly diluted token-space attention for long contexts. MemDLM simulates denoising with Bi-level Optimization, writing local trajectory information into fast-weight Parametric Memory and conditioning the base model on it. It improves long-context results, including from 78.8% to 95.8% on 8K RULER Variable Tracking, while inference-time re-enablement adds prompt-specific adaptation interpreted as in-weight retrieval.

  • Problem

    Standard DLM training uses a static single-step objective that omits progressive denoising dynamics and leaves long-context information to increasingly diluted token-space attention.

  • Method

    MemDLM uses Bi-level Optimization: an inner loop updates fast weights as Parametric Memory of local denoising trajectories, while an outer loop updates the base model conditioned on that memory.

  • Results

    MemDLM improves long-context performance across reported settings, including RULER Variable Tracking at 8K from 78.8% to 95.8% and BABILong at 8K from 47.4% to 57.0%.

  • Takeaways & Limitations

    Parameter-space memory strengthens long-context DLM representations even when fast weights are discarded, while re-enabling the inner loop adds prompt-specific adaptation interpreted as in-weight retrieval.

  • Takeaways & Limitations

    Experiments cover only two DLM backbones and one instruction-tuning dataset, while inference-time adaptation adds prompt-side latency.

Abstract

from arXiv · show

Diffusion Language Models (DLMs) offer attractive advantages over Auto-Regressive (AR) models, such as full-attention parallel decoding and flexible generation. However, standard DLM training uses a static, single-step masked prediction objective that never exposes the model to the progressive denoising dynamics of inference, and forces all contextual information to be maintained purely through token-space attention, which becomes increasingly diluted as context length grows. We propose MemDLM (Memory-Enhanced DLM), which introduces a second memory channel by embedding a simulated denoising trajectory 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, while an outer loop updates the base model conditioned on this memory. By offloading part of the memorization burden from token-space attention to parameter space, MemDLM yields faster convergence, stronger long-context representations, and lower training loss, even when the fast weights are discarded at inference time. Re-enabling the inner loop at inference provides an additional prompt-specific adaptation effect, where the Parametric Memory acts as an emergent in-weight retrieval mechanism on challenging Needle-in-a-Haystack tasks. Code: https://github.com/JarvisPei/MemDLM.

1 Introduction

MemDLM addresses static single-step training and long-context attention dilution by adding a parameter-space memory channel through simulated denoising trajectories. It improves long-context performance, including when fast weights are discarded, while inference-time re-enablement adds prompt-specific adaptation.

  • Standard DLM training misses progressive denoising dynamics and relies on token-space attention to preserve information in long contexts.
  • MemDLM improves RULER Variable Tracking at 8K from 78.8% to 95.8% on LLaDA-MoE and BABILong at 8K from 47.4% to 57.0% on LLaDA2.1.
  • MemDLM uses Bi-level Optimization to simulate progressive denoising and induce Parametric Memory as a second channel alongside attention.
  • Training improves optimization and long-context performance even when fast weights are discarded at inference.
  • Re-enabling the inner loop at inference provides prompt-specific adaptation interpreted as in-weight retrieval.

2 Preliminaries and Motivation

MDLMs corrupt text with masking and train a single-step reconstruction objective, whereas inference performs iterative denoising. This mismatch and increasingly diluted attention motivate a complementary parameter-space memory channel.

  • Discrete diffusion maps clean sequences x0 through progressively noisier states until x1, and trains the reverse process to recover x0.
  • Absorbing-state masking replaces tokens with a dedicated mask token, with masking probability t under a linear noise schedule.
  • Standard MDLM reconstructs currently masked tokens from xt using uniformly sampled timesteps and a time-dependent loss weight ω(t).
  • Inference iteratively unmasks high-confidence predictions, producing a progressive trajectory from a fully masked sequence toward t = 0.
  • Longer contexts dilute dense attention across tokens, making task-relevant information harder to preserve through token-space representations alone.
  • A second memory channel can write local denoising trajectories into parameter space, enriching training while reducing attention's memorization burden.

3 Methodology

MemDLM aligns training with iterative denoising by using a bi-level optimization framework that builds transient parametric memory from local trajectories. The outer loop then updates the base model conditioned on this memory, reducing reliance on token-space attention.

  • Bi-level Optimization: MemDLM partitions parameters into base weights θ and parameter-efficient fast weights ϕ, embedding a simulated denoising trajectory through bi-level optimization.The fast weights can be implemented with low-rank adapters.
  • Inner Loop: The inner loop unrolls a K-step denoising trajectory, updating fast weights that accumulate sample-specific contextual details into a final Parametric Memory state ϕK.The fast weights start from zero and are updated by gradient descent.
  • Inner Loop: The inner loop uses an anchor-consistent trajectory that first denoises a noisier state toward anchor state xt, then predicts clean state x0 from xt.For the described implementation, K = 2 and ϕ0 = 0.
  • Inner Loop: The two inner-loop stages make the final fast-weight state ϕ2 encode transitions from a noisier local state through xt toward x0.This produces an anchor-centered local trajectory in parameter space.
  • Outer Loop: After the inner loop, the outer objective evaluates the same anchor timestep and masked state while conditioning predictions on Parametric Memory ϕ2.The outer objective mirrors standard MDLM training but uses the adapted fast weights.
  • Outer Loop: A first-order approximation updates θ without computing second-order Hessian matrices, while fast weights absorb batch-specific trajectory information.This may reduce pressure on the base model to memorize local context purely through token-space representations.

4 Experiments

Experiments show that MemDLM improves long-context retrieval, training efficiency, and generalization across backbones and context lengths. Ablations further identify trajectory consistency, restricted adaptation scope, and two-stage inner-loop design as important choices.

  • 4.2 Main Results: Long-Context Information Retrieval: MemDLM consistently improves over Standard MDLM across both backbones, especially at longer contexts.Train-Only already provides large gains, while inference-time adaptation adds further prompt-specific improvement.
  • 4.2 Main Results: Long-Context Information Retrieval: 78.8% to 95.8%: MemDLM improves RULER Variable Tracking at 8K on LLaDA-MoE.On LLaDA2.1, BABILong at 8K improves from 47.4% to 57.0%.
  • 4.2 Main Results: Long-Context Information Retrieval: MemDLM continues to outperform Standard MDLM at 16K and 32K across RULER and BABILong.Performance drops for all methods as context grows, but MemDLM retains an advantage in the extrapolation regime.
  • 4.3 Long-Context Generalization: MemDLM improves LongBench performance in zero-shot Train-Only mode, with Train & Inference yielding consistent further gains.The two settings respectively disable and reactivate the inner loop at inference.
  • 4.4 Understanding MemDLM During Training: MemDLM descends more rapidly in training loss and maintains lower evaluation loss on both LLaDA-MoE and LLaDA2.1.Across 1K–32K, it also improves over both the pretrained base and Standard MDLM-trained models.
  • 4.4 Understanding MemDLM During Training: Restricting inner-loop updates to FFN modules in the last 10% of layers yields the best BABILong-1K score, 0.684.Full-parameter updates achieve lower train loss but underperform this restricted adaptation.
  • 4.5 Understanding MemDLM During Inference: Inference performance is relatively insensitive to anchor ratios from 0.2 to 0.8 across 1K–16K contexts.At 16K, scores remain between 0.212 and 0.232; the default 0.2 is chosen as a simple operating point.
  • 4.6 Ablations: A consistent trajectory improves BABILong-1K from 0.604 to 0.684, while additional pre-anchor steps lower downstream scores despite lowering training loss.The 3-step and 4-step variants score 0.644 and 0.590, respectively.

5 Related Work

MemDLM connects diffusion-language-model training, long-context attention, fast-weight memory, bi-level adaptation, and test-time learning. Its distinctive focus is using inner-loop updates to encode each sample’s local denoising trajectory rather than adapting across few-shot task episodes.

  • Diffusion language models and the training-inference gap: Related diffusion-language-model work addresses the training-inference discrepancy through progressive remasking, trajectory-aware learning, and other inference-aligned objectives.
  • Long-context attention and its limitations: Long-context studies examine attention dilution using length-aware normalization, dynamic rescaling, attention sinks, and positional methods.
  • Fast weights and parametric memory: Fast-weight literature stores short-lived, sample-specific information in parameters, while related adaptation methods use local weight updates as parametric memory.
  • Meta-learning and Bi-level Optimization: Unlike few-shot meta-learning, MemDLM’s inner loop internalizes each sample’s local denoising trajectory.
  • Test-time training: Test-time training methods update model behavior online using unlabeled or self-supervised signals, including weight-based long-context compression.

6 Limitations

MemDLM’s scope is limited by added training and inference costs, narrow experimental coverage, and an inference procedure confined to prompt adaptation.

  • 6 Limitations: The inner loop adds approximately 2× default wall-clock training time, despite parameter-efficient adapters and restricted update scope.
  • 6 Limitations: Experiments cover only two DLM backbones and one instruction-tuning dataset, leaving broader architectural and pretraining generalization untested.
  • 6 Limitations: Inference-time adaptation adds prompt-side latency, and extending the inner loop into decoding remains future work.

7 Conclusion

MemDLM uses Bi-level Optimization and fast weights as Parametric Memory to add a parameter-space memory channel for diffusion language models. Simulated denoising trajectories improve optimization and long-context performance, while inference-time inner-loop re-enabling adds prompt-specific adaptation.

  • MemDLM combines Bi-level Optimization with fast weights as Parametric Memory for memory-aware diffusion language model training.
  • Simulating denoising trajectories during training offloads part of the memorization burden from token-space attention to parameter space.
  • MemDLM improves optimization and long-context performance even in the Train-Only setting, when fast weights are not used at inference.
  • Re-enabling the inner loop at inference provides prompt-specific adaptation interpreted as in-weight retrieval.

A Additional Experimental Details

The experiments use instruction-tuned LongAlpaca data, two LLaDA backbones, and evaluations spanning synthetic retrieval, reasoning-in-a-haystack, and realistic long-context tasks. Training and decoding configurations are adapted to each backbone while shared evaluation options remain fixed where applicable.

  • Implementation and Baselines: MemDLM is evaluated with LLaDA-MoE-7B-A1B-Base and LLaDA2.1-mini (16B), with most training-stage analyses and core ablations using LLaDA-MoE.
  • Training Data and Processing: Instruction tuning uses LongAlpaca sequences of at most 4,096 tokens, keeping prompt tokens unmasked and applying noise only to response tokens.
  • Evaluation Benchmarks: Evaluation covers RULER, BABILong, and LongBench, spanning retrieval, multi-hop reasoning, aggregation, question answering, summarization, classification, and code-completion tasks.
  • Generation and Decoding Settings: LLaDA-MoE uses single-block generation and LLaDA2.1 uses multi-block generation; both models generate one token per iterative denoising step.

C Empirical Analysis of Exposure Bias

The analysis measures train–inference mismatch by comparing losses under static masking and self-generated sequential denoising. Exposure bias rises sharply for Standard MDLM, whereas MemDLM remains substantially flatter even when its inner loop is disabled at inference.

  • Static Condition: The static condition predicts masked response tokens from a pristine context masked according to the true forward process.
  • Sequential Condition: The sequential condition starts from a fully masked response and iteratively unmasks tokens using the model’s own predictions until timestep t.
  • Exposure Bias Measurement: Exposure Bias Ratio is defined as R_EB = L_seq/L_static, with higher values indicating more severe exposure bias during self-denoising.
  • Empirical Findings: Standard MDLM shows a steeply rising exposure-bias curve, while MemDLM has substantially flatter degradation even in Train-Only mode.

D Additional Training Ablations

Additional ablations examine inner-loop gradient normalization and the pre-anchor state design. Local normalization is more important than the precise clipping threshold, while a pre-anchor scale of 1.5 performs best without making results highly fragile nearby.

  • Gradient Normalization: 0.684 is the best BABILong-1K score, achieved with local per-parameter gradient normalization and clipping at 1.0; global normalization scores 0.632.
  • Gradient Normalization: 0.630, 0.640, and 0.682 result from local normalization with clipping thresholds of 0.5, 2.0, and no clipping, respectively.
  • Pre-anchor Design: A pre-anchor scale of 1.5 performs best at 0.684, while 1.75 and 2.0 remain competitive at 0.674 and 0.678; 1.25 scores 0.624.

E Training Workload Analysis

MemDLM’s default configuration incurs approximately 2.1× the wall-clock overhead of Standard MDLM, with overhead varying predictably according to the number of inner-loop parameters updated.

  • Wall-clock time and per-step cost are measured from TensorBoard logs, with overhead reported relative to each backbone’s Standard MDLM baseline.The comparison uses the workload measures defined for Table 4.
  • 2.1× wall-clock overhead is added by MemDLM compared with Standard MDLM under the default configuration on both backbones.The default uses a 2-step inner loop and FFN-only LoRA on the last 10% of layers.
  • ∼1 hour versus ∼0.5 hours is the training time for MemDLM versus Standard MDLM on LLaDA-MoE.
  • ∼3.4 hours versus ∼1.6 hours is the training time for MemDLM versus Standard MDLM on LLaDA2.1.
  • Near-2× overhead is maintained when inner-loop updates are restricted to FFN modules in the last 5–10% of layers.The per-step cost scales with the number of parameters updated in the inner loop.
Loading 2603.22241v2…