💬 NLP

Characterizing Memorization in Diffusion Language Models: Generalized Extraction and Sampling Effects

This paper presents a theoretical and empirical framework demonstrating that while diffusion language models (DLMs) can memorize training data with probability increasing alongside sampling resolution, they exhibit substantially lower leakage of personally identifiable information compared to autoregressive models under aligned prefix-conditioned evaluations.

Original authors: Xiaoyu Luo, Wenrui Yu, Qiongxiu Li, Johannes Bjerva

Published 2026-03-04

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Xiaoyu Luo, Wenrui Yu, Qiongxiu Li, Johannes Bjerva

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

The Big Picture: Two Ways to Write a Story

Imagine you have two different writers trying to finish a story based on a few starting words (a "prefix").

The Old Writer (Autoregressive Model - ARM): This writer is like a strict line of dominoes. They write one word, then the next, then the next, moving strictly from left to right. They can't go back and change what they wrote earlier. This is how most current AI (like the original versions of ChatGPT) works.
The New Writer (Diffusion Model - DLM): This writer is like a sculptor working with a block of clay. They start with a block of "noise" (or a page full of blank masks ⟨M⟩). They don't write word-by-word. Instead, they look at the whole page, guess what a few words might be, fill them in, then look at the whole page again to refine those guesses. They keep polishing the whole page until it makes sense. This is the new "Diffusion" technology.

The Problem: The "Photocopy" Risk

Both writers have a secret flaw: Memorization.
If you train them on a massive library of books, they might accidentally memorize specific pages. If you ask them to finish a sentence that starts a famous book, they might just "regurgitate" (copy-paste) the rest of the book, including private info like phone numbers or email addresses.

We know the "Old Writer" does this a lot. But we didn't know if the "New Writer" (Diffusion) was safer or more dangerous.

The Paper's Discovery: The "Zoom Lens" Effect

The researchers built a new way to test the New Writer. They realized that the New Writer has a special setting called Sampling Resolution. Think of this as a Zoom Lens or the Number of Steps the writer takes to finish the job.

Low Resolution (Fast/Coarse): The writer guesses big chunks of text at once. It's fast, but it's like looking at a photo from far away. You might get the general idea, but you won't see the tiny details.
High Resolution (Slow/Fine): The writer guesses one tiny word at a time, refining the whole picture step-by-step. It's like zooming in until you can read the individual pixels.

The Big Finding:
The paper proves a surprising rule: The more steps (higher resolution) you take, the more likely the AI is to accidentally copy the training data word-for-word.

Analogy: Imagine trying to guess a secret password.
- If you guess the whole password in one giant leap (Low Resolution), you might get it wrong.
- If you guess one letter, check it, then guess the next letter, and so on (High Resolution), you are much more likely to eventually reconstruct the exact original password.
The Twist: The researchers found that the "Old Writer" (who writes one word at a time) is actually just the "New Writer" turned up to Maximum Resolution. Because the Old Writer always writes one word at a time, it is naturally more prone to memorizing exact secrets than the New Writer, which usually guesses in chunks.

The Experiments: Testing the "Leakage"

The team ran tests to see how much "Personal Identifiable Information" (PII) like emails and phone numbers these models leaked.

The Setup: They took two models of similar size (one Old, one New) and gave them a prompt that started with a fake email address.
The Result:
- The Old Writer (ARM) was a "leaky faucet." It frequently finished the sentence with the exact real email address it had memorized.
- The New Writer (DLM) was much more secure. Even when asked to finish the sentence, it rarely copied the exact private data.
The Catch: If they forced the New Writer to use the "Maximum Resolution" (guessing one letter at a time, just like the Old Writer), it did start leaking more. But in its normal, faster mode, it was safer.

Why Does This Matter?

Privacy: If you are worried about an AI accidentally spilling your private data (like your home address or credit card number), the new Diffusion models might be a safer choice if you don't force them to work in "slow-motion" (high resolution).
The Trade-off: The new models are faster and can look at the whole sentence at once (which helps them understand context better), but this "blurry" approach actually helps hide memorized secrets.
The Warning: If you tell a Diffusion model to be super precise and take many tiny steps to generate text, you are essentially turning it into a "photocopier" that is more likely to leak secrets.

Summary in a Nutshell

Old AI: Writes word-by-word. Very good at copying secrets exactly.
New AI: Polishes the whole text at once. Generally better at keeping secrets safe because it's "guessing" rather than "reciting."
The Rule: The more precise and step-by-step the New AI is forced to be, the more it starts acting like the Old AI and leaking secrets.
Conclusion: The new Diffusion models are promising for privacy, but we need to be careful not to force them to work in a way that makes them too precise, or they might start memorizing our private data just like the old ones.

1. Problem Statement

Large Language Models (LLMs), particularly Autoregressive Models (ARMs), are known to memorize and verbatim reproduce training data, posing significant risks to privacy (PII leakage) and copyright. While Diffusion Language Models (DLMs) have emerged as a competitive alternative with bidirectional generation capabilities, their memorization behavior remains largely unexplored.

The Gap: Existing memorization frameworks (e.g., prefix-suffix extraction) are designed for the unidirectional, sequential nature of ARMs. They do not naturally extend to DLMs, which recover masked tokens through stochastic, non-sequential denoising trajectories.
The Question: How does the fundamental shift from unidirectional next-token prediction to bidirectional masking and reverse denoising affect a model's tendency to memorize and leak training data?

2. Methodology

The authors propose a unified theoretical and empirical framework to characterize memorization in DLMs.

A. Generalized Probabilistic Extraction Framework

The paper introduces a Generalized (n, p)-discoverable Extraction framework that unifies prefix-conditioned decoding (ARMs) and diffusion-based generation under arbitrary masking patterns and stochastic sampling trajectories.

Definition: A sequence is $(n, p)$ -extractable if, given observed tokens (unmasked context), the model recovers the exact original masked tokens within $n$ independent queries with probability at least $p$ .
Probability Estimation: Unlike ARMs where decoding is deterministic or fixed-order, DLM recovery paths are stochastic. The authors estimate the single-trial recovery probability ( $p_z$ ) by averaging results over multiple mask patterns and sampling trajectories.

B. Theoretical Analysis: Sampling Resolution

The core theoretical contribution is the analysis of Sampling Resolution ( $N$ ), defined as the number of steps taken to recover masked tokens.

Assumption 4.2 (Monotonicity): The probability of correctly recovering masked tokens increases as the set of observed (correctly recovered) tokens expands.
Theorem 4.3: Under the monotonicity assumption, increasing the sampling resolution (i.e., using more, finer-grained denoising steps) strictly increases the probability of exact training data extraction.
Proposition 4.4: Autoregressive decoding is identified as a limiting case of diffusion generation where the sampling resolution is maximal ( $N = |M|$ , recovering one token at a time in a fixed sequence). This implies ARMs represent the upper bound for memorization probability in this framework.

C. Experimental Setup

Models: The authors trained DLMs (170M, 690M, 1.1B parameters) and an ARM baseline (1.1B) on the SlimPajama dataset using identical compute budgets and tokenizers to isolate architectural effects. They also evaluated the large-scale LLaDA-8B model.
Datasets:
- Validation: Random windows from SlimPajama to test theoretical vs. empirical recovery probabilities.
- PII Leakage: An aligned evaluation using the Enron email dataset (emails and phone numbers) to test prefix-conditioned PII completion.
- Generalization Check: The TREC 2007 Spam dataset (disjoint from training but same domain) to distinguish memorization from generalization.

3. Key Contributions

Unified Framework: A generalized probabilistic formulation for memorization that applies to both ARMs and DLMs, treating prefix-conditioned extraction as a special case of diffusion generation.
Sampling Resolution Control: A theoretical proof (Theorem 4.3) demonstrating that finer-grained sampling resolution (more denoising steps) monotonically increases the risk of exact memorization.
Architectural Comparison: An aligned analysis showing that, under identical prefix-conditioned settings, DLMs exhibit substantially lower PII memorization compared to ARMs of similar scale.

4. Key Results

Validation of Theory: Empirical results confirm that the theoretical probability estimates derived from the generalized framework align closely with observed extraction rates (Figure 2).
Resolution Impact: Experiments across 1-step, 2-step, 5-step, 10-step, and "Max" (per-token) resolutions show a clear, approximately monotonic increase in exact-recovery success rates as the number of denoising steps increases (Figure 3). This validates that ARMs (max resolution) are more prone to verbatim leakage than coarse-grained diffusion.
PII Leakage Comparison:
- In the aligned Enron PII task, DLMs leaked significantly fewer memorized samples than ARMs.
- For the 1.1B models, the ARM extracted 213 emails at $p=50\%$ , while the DLM extracted only 16 (at Max resolution) and 0 (at 1-step).
- Even the LLaDA-8B (8B parameters), despite being trained on significantly more data (15.6B tokens vs. 2.3B for the 1.1B ARM), showed memorization rates comparable to or lower than the much smaller 1.1B ARM.
Memorization vs. Generalization: Comparisons between training data (Enron) and disjoint test data (TREC Spam) showed a clear separation in reconstruction likelihood, confirming that the extraction metric captures true memorization rather than domain generalization.

5. Significance and Implications

Privacy Implications: The findings suggest that DLMs may be inherently more robust against verbatim training data leakage than ARMs, particularly when using coarser sampling resolutions (fewer steps). This offers a potential mitigation strategy for privacy-sensitive applications.
Theoretical Insight: The work establishes a formal link between generation dynamics and memorization, proving that the "unidirectional" nature of ARMs is a specific, high-risk configuration of the broader diffusion paradigm.
Future Directions: The paper highlights the need to investigate how post-training techniques (SFT, RLHF) might shift DLMs toward more prefix-suffix behaviors, potentially increasing their memorization risks to match ARMs.

In summary, this paper provides the first systematic characterization of memorization in Diffusion Language Models, proving that sampling resolution is a critical control knob for memorization risk and demonstrating that DLMs currently offer a lower leakage profile for PII compared to their autoregressive counterparts.