💬 NLP

Extracting Training Data from Diffusion Language Models via Infilling

This paper introduces "infilling extraction" to demonstrate that diffusion language models are significantly more vulnerable to training data memorization than autoregressive models, as their bidirectional denoising capabilities allow adversaries to extract up to three times more verbatim sequences, including redacted personally identifiable information, using edge-conditioned masks.

Original authors: Yihan Wang, N. Asokan

Published 2026-05-26

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Yihan Wang, N. Asokan

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

Imagine you have a giant, super-smart library assistant (a Large Language Model) that has read millions of books, emails, and documents. You want to know: If someone asks this assistant to repeat a specific sentence from its memory, how likely is it to spill the beans?

For a long time, researchers only tested this by asking the assistant to "finish the sentence" starting from the very beginning. It's like showing someone the first half of a story and asking them to guess the rest. This works well for standard AI models that read left-to-right, like a human reading a book.

But this paper introduces a new type of AI called a Diffusion Language Model (DLM). Think of a DLM not as a reader, but as a restorer of damaged paintings. If you take a masterpiece and cover random parts of it with white paint (masks), the DLM's job is to look at the whole picture—what's to the left, what's to the right, and what's in the middle—and guess what the hidden parts should be.

The authors of this paper say: "Hey, we've been testing these painting restorers the wrong way!"

Here is the breakdown of their findings using simple analogies:

1. The Old Way vs. The New Way

The Old Way (Prefix-Conditioned): Imagine you show the restorer only the left edge of a painting and ask, "What comes next?" The paper found that if you only do this, the DLM seems pretty safe. It doesn't leak much secret info.
The New Way (Infilling Extraction): The authors realized that because DLMs can look at both sides of a gap, a sneaky attacker could cover up the middle of a sentence and ask, "Fill in the blank." Or, they could cover the ends and ask, "What's in the middle?"
- The Analogy: It's like playing a game of "Mad Libs." If you give the AI the beginning and the end of a sentence, it might be able to guess the middle word perfectly, even if it couldn't guess the end if you only gave it the beginning.

2. The Big Discovery: "Edge" is the Danger Zone

The researchers tested different ways of hiding parts of the text (masks). They found that how you hide the text matters more than you think.

The "Edge" Effect: If you reveal the very beginning and the very end of a sentence (hiding the middle), the DLM is three times more likely to leak the exact words than if you only revealed the beginning.
Why? Because the DLM uses clues from both sides to reconstruct the middle. It's like having two people whispering the start and end of a secret to a third person; they can figure out the whole secret much easier than if only one person whispered the start.

3. The "Redacted Document" Scenario

The paper looks at a very realistic scary scenario:

Imagine a company trains an AI on private emails but then blacks out (redacts) all the phone numbers and names before releasing the data or the model.
The Attack: A hacker takes the model and says, "I have the email with the name blacked out. Please fill in the blank."
The Result: Even though the AI was trained on redacted data, the "painting restorer" (DLM) is so good at looking at the surrounding context that it can often guess the blacked-out name or number better than a standard "book-reading" AI could.
The Takeaway: Just because you black out the sensitive info in the training data doesn't mean the AI won't remember it. In fact, for this type of AI, it might be easier to guess the secret if you give it the context on both sides.

4. Tuning the "Restoration" Settings

The researchers also found that the settings used to "restore" the text (like how many steps the AI takes to guess) change how much it leaks.

Slow and Steady: If you make the AI take many small, careful steps to fill in the blanks, it leaks more secret info.
Fast and Loose: If you make it guess quickly, it leaks less.
The Lesson: The "knobs" you turn to make the AI faster or smarter also control how much it remembers. You can't just turn one knob for speed without affecting safety.

5. Fine-Tuning Doesn't Fix It

Finally, they asked: "If we teach the AI to be a polite chatbot later on (a process called SFT), does it forget the secrets?"

The Answer: No. It's like trying to teach a parrot to stop saying a swear word by teaching it new words. The parrot might stop saying it sometimes, but if you ask the right question (using the "edge" trick), it will still blurt out the old secret. The memory is still there; it's just been rearranged.

Summary

This paper warns us that Diffusion Language Models are like painting restorers, not just book readers. If you only test them by asking them to finish a sentence from the start, you think they are safe. But if you test them by asking them to fill in the middle of a sentence (using clues from both sides), they are much more likely to leak private information, even if that information was redacted from their training data.

The main takeaway: We need to stop testing these models with just one eye open (looking left-to-right) and start testing them with both eyes open (looking at the whole picture), or we will vastly underestimate how much private data they are holding onto.

Technical Summary: Extracting Training Data from Diffusion Language Models via Infilling

Problem Statement

While the memorization of training data in autoregressive language models (ARMs) has been extensively studied through prefix-conditioned extraction (where an adversary provides a prefix and the model generates the suffix), this approach fails to capture the full risk profile of Diffusion Language Models (DLMs). DLMs operate via iterative denoising and possess a bidirectional inductive bias, allowing them to reconstruct tokens at arbitrary positions given surrounding context. Relying solely on prefix-conditioned probing, a protocol inherited from ARM research, significantly underestimates the risk of training data extraction in DLMs because it ignores the model's ability to "fill in" masked tokens using both preceding and following context.

Methodology

The authors introduce Infilling Extraction, a generalized data-extraction protocol parameterized by an arbitrary binary mask $m$ . In this protocol, a training sequence is partially masked, and the DLM is tasked with reconstructing the masked tokens based on the unmasked context. The sequence is considered "extracted" if the reconstruction matches the original masked text.

Experimental Setup

The study evaluates this protocol on two large-scale DLMs (LLaDA-8B and Dream-7B) and compares them against a scale-matched ARM baseline (LLaMA-2-7B). The experiments cover:

Five Extraction Modes:
- Prefix-conditioned: First $k$ tokens revealed (standard ARM baseline).
- Suffix-conditioned: Last $k$ tokens revealed.
- Edge-conditioned: First and last $k/2$ tokens revealed (bracketing context).
- Middle-conditioned: A contiguous span in the middle revealed.
- Random: A uniformly random subset of $k$ tokens revealed.
- Targeted: Only sensitive entities (PII) are masked, with the entire surrounding sequence revealed.
Three Training Pipelines: Full Fine-Tuning (FT), Low-Rank Adaptation (LoRA), and Supervised Fine-Tuning (SFT) under a chat template.
Three Datasets:
- MIMIR Pile: For verbatim sequence extraction.
- Enron Email: For Personally Identifiable Information (PII) extraction (email addresses and phone numbers).
- Tulu3: For prompt/response leakage in instruction-tuned settings.
Decoding Parameters: The study analyzes the impact of denoising steps ( $T$ ), semi-autoregressive block size ( $b$ ), sampling temperature ( $\tau$ ), and remasking strategies.

Key Contributions

Formalization of Infilling Extraction: The paper defines a mask-parameterized protocol that generalizes prefix-conditioned extraction, explicitly accounting for the bidirectional nature of DLMs.
Systematic Evaluation: It provides the first systematic study of memorization in large DLMs across diverse access patterns, exposure regimes, and inference-time configurations.
Identification of New Leakage Channels: The work empirically demonstrates that bidirectional access opens leakage channels structurally inaccessible to ARMs, such as recovering user prompts from model responses.

Key Results

1. Mask Geometry Governs Extractability

The placement of revealed context significantly impacts extraction success.

Edge-Conditioned Superiority: Edge-conditioned masks (revealing both start and end) extract up to three times more verbatim sequences than prefix-conditioned masks. For example, on LLaDA-8B-FT, the Edge mode achieved 0.41% ASR compared to 0.18% for Prefix.
Bidirectional Bias: This confirms that DLMs leverage context from both directions to reconstruct masked tokens, a capability not utilized in standard prefix probing.

2. Bidirectional Access Opens New Leakage Channels

Prompt Recovery: In instruction-tuned models (SFT), adversaries can recover user prompts from known model responses (Response-conditioned extraction). While prompt-conditioned extraction (standard) showed ~~10% ASR, response-conditioned extraction still achieved non-trivial ASR (~~1.78–3.36%), a channel impossible for ARMs.
Token vs. Sequence: While token-level recovery (TER) can be high (e.g., 56% for random masks), exact verbatim extraction (ASR) remains low, suggesting that revealed tokens provide scaffolding that aids partial recovery but does not guarantee full sequence reproduction.

3. Realistic Threat Model: Redacted Data and PII

The study simulates a realistic scenario where an adversary obtains a redacted version of the training data (e.g., via regulatory release or partial leak) and uses the model to recover the redacted PII.

Targeted Infilling: When only PII tokens are masked and the surrounding context is revealed, DLMs show high extractability. On the Enron dataset, fine-tuned LLaDA-8B achieved 40.5% recall for email addresses and 28.9% for phone numbers.
Comparison with ARMs: Under a strict prefix-conditioned attack, DLMs appeared safer than ARMs (LLaDA: 16.3% vs. LLaMA: 35.9% recall for emails). However, under targeted infilling, DLMs matched or exceeded the ARM baseline (LLaDA: 40.5% vs. LLaMA: 35.9%). This indicates that prefix-only evaluations severely underestimate DLM risks in realistic scenarios involving redacted data.

4. Influence of Decoding Parameters

Denoising Steps ( $T$ ): Higher $T$ (more refinement steps) increases recall, while throughput-oriented low $T$ settings obscure memorized content.
Block Size ( $b$ ): Larger block sizes (fully parallel denoising) significantly increase extraction success compared to small blocks (AR-like behavior), identifying bidirectional parallelism as a primary driver of leakage.
Remasking Strategy: Confidence-based remasking (committing high-confidence predictions) amplifies extraction compared to random remasking.
Temperature ( $\tau$ ): Greedy decoding ( $\tau \to 0$ ) generally yields the strongest extraction.

5. Persistence Through SFT

A follow-up Supervised Fine-Tuning (SFT) stage on instruction data redistributes but does not eliminate memorization. While email recall decreased slightly after SFT, phone number recall remained largely unchanged, suggesting that short, idiosyncratic numeric strings are orthogonal to instruction-following objectives and survive the alignment process.

Significance and Claims

The paper argues that the current safety evaluation of DLMs is incomplete because it relies on prefix-conditioned protocols designed for ARMs. By introducing infilling extraction, the authors demonstrate that:

DLMs pose a greater risk of training data leakage than previously thought, particularly when adversaries can leverage bidirectional context or access redacted training data.
Mask geometry is a critical factor in memorization, with edge-conditioned and targeted infilling revealing vulnerabilities invisible to prefix-only probing.
Decoding parameters act as levers for both amplification and mitigation; practitioners can tune these to balance performance and leakage resistance.

The authors conclude that while DLMs may appear safer under traditional prefix-based metrics, their bidirectional nature creates unique and potent channels for data extraction that must be addressed in future safety research and deployment strategies.