💻 computer science

DMin: Scalable Training Data Influence Estimation for Diffusion Models

The paper introduces DMin, a scalable framework that enables efficient influence estimation for billion-parameter diffusion models by leveraging gradient compression to drastically reduce storage requirements and retrieval time while maintaining high accuracy.

Original authors: Huawei Lin, Yingjie Lao, Weijie Zhao

Published 2026-04-10

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Huawei Lin, Yingjie Lao, Weijie Zhao

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 magical artist named Diffusion. This artist can paint anything you ask for: a "cute fluffy baby giraffe-lion," a "cosmic horror Cthulhu," or a "maple forest in autumn."

But here's the mystery: Where does this artist get their ideas?

The artist was trained on a massive library of millions of images and descriptions from the internet. When you ask for a specific painting, the artist doesn't just pull a single image from the shelf; they mix and match concepts from thousands of different training examples to create something new.

The Problem:
If you want to know exactly which pictures from that massive library inspired a specific painting, it's like trying to find a single grain of sand on a beach.

The Old Way: Previous methods tried to solve this by taking a "photocopy" of every single page in the library, highlighting every word the artist read, and then comparing those photocopies to your new painting.
The Catch: For a modern, super-smart artist (like Stable Diffusion 3), the "photocopy" of just one training image is huge. If you tried to keep photocopies of all the images the artist learned from, you would need hundreds of Terabytes of hard drive space. That's like filling a warehouse with paper just to do a little math. It's impossible for regular computers.

The Solution: DMin (Diffusion Model Influence)
The authors of this paper invented a new tool called DMin. Think of it as a super-smart, ultra-compact librarian.

Here is how DMin works, using simple analogies:

1. The "Fingerprint" Instead of the "Photo"

Instead of keeping a giant, high-resolution photo of every training image (which takes up too much space), DMin takes a fingerprint of the artist's reaction to that image.

The Old Way: "Here is the whole book the artist read." (Too big!)
The DMin Way: "Here is a tiny, 1KB note that says exactly how the artist's brain changed when they saw this image."

2. The "Magic Compression" Trick

How do they shrink a massive fingerprint into a tiny note? They use a clever trick called Gradient Compression.
Imagine you have a giant, messy pile of Lego bricks (the data).

Step 1: You shake the pile so the bricks are in a random order (Permutation).
Step 2: You ask a friend to flip a coin for every brick. If it's heads, you keep the brick; if tails, you flip it upside down (Random Projection).
Step 3: You group the bricks into small buckets and just count how many are in each bucket (Group Addition).

The result? You don't have the original bricks anymore, but you have a tiny list of numbers that still tells you everything you need to know about the pile's shape. This shrinks the storage from Gigabytes to Kilobytes.

3. The "Speedy Search" (KNN)

Once the librarian has these tiny fingerprints for millions of images, they can answer your question instantly.

The Question: "Which training images inspired this 'Cosmic Horror Cthulhu'?"
The Search: The librarian takes the fingerprint of your new Cthulhu and runs a super-fast "match" against the tiny fingerprints in the database.
The Result: In less than one second, the librarian points to the top 25 images that were most influential.

Why This Matters

Transparency: If an artist accidentally paints something offensive or biased, we can now trace it back to the specific bad image in the training data that caused it.
Efficiency: We can do this on the world's biggest, most powerful AI models without needing a supercomputer the size of a city.
Safety: It helps us understand and fix the "memories" of AI, making them safer and more reliable.

In a Nutshell:
Before, trying to find the source of an AI's idea was like trying to find a needle in a haystack while the haystack was on fire and you were running out of oxygen. DMin is like turning that haystack into a tiny, organized card catalog, allowing you to find the needle in a blink of an eye.

1. Problem Statement

Diffusion Models (DMs) are powerful generative tools, but their training data is often scraped from the internet, leading to potential biases, inaccuracies, or the generation of harmful content. Understanding which specific training samples influenced a generated image is crucial for model transparency, bias detection, and debugging.

However, existing influence estimation methods face critical scalability barriers when applied to large-scale diffusion models (e.g., Stable Diffusion 3 with billions of parameters):

Storage Bottleneck: Second-order methods (requiring Hessian inversion) and first-order methods (requiring gradient storage) demand massive storage. For a 2B parameter model, storing gradients for just 10,000 samples across 10 timesteps requires ~800 TB of storage, which is impractical.
Computational Cost: Calculating influence scores for large datasets is computationally prohibitive, often taking hours or days.
Fragility: Influence estimation in extremely deep models is unstable due to large gradient magnitudes, leading to inaccurate results.

Current solutions are limited to small models or LoRA-tuned (adapter) models, failing to scale to full-parameter billion-scale models.

2. Methodology: DMin Framework

The authors propose DMin (Diffusion Model influence), a scalable framework that estimates influence by compressing gradients and utilizing efficient search algorithms. The workflow consists of three main stages (illustrated in Figure 2 of the paper):

A. Gradient Computation

Process: For a training pair (prompt $p_i$ , image $x_i$ ), the model performs a forward and backward pass at specific timesteps $t$ to compute the loss gradient $\nabla_\theta L$ .
Adapters vs. Full Model: For LoRA-tuned models, gradients are collected only from the adapter parameters. For full models, gradients are collected from all parameters.
Sampling: Instead of computing gradients for all $T$ timesteps (e.g., 1,000), a subset of timesteps is sampled to reduce computational load.

B. Gradient Compression (The Core Innovation)

To overcome the storage bottleneck, DMin employs a four-step compression technique that reduces gradient storage from GBs to KBs without significant accuracy loss:

Padding: The gradient vector is padded to a length divisible by a compression factor $v$ .
Permutation: A random permutation is applied to disrupt inherent vector structures.
Random Projection: The vector is element-wise multiplied by a random projection vector (elements are -1 or 1 with equal probability). This projects the gradient onto a randomized basis.
Group Addition: The permuted vector is divided into groups of size $v$ , and elements within each group are summed to produce a compressed vector of dimension $v$ .

Storage Efficiency: Only the permutation indices (4 bytes/element) and the binary projection vector (1 bit/element) need to be stored alongside the compressed vector.
Normalization: To address gradient instability in deep models, L2 normalization is applied to gradient vectors before compression. This prevents outliers from dominating the inner product calculations.

C. Influence Estimation & Retrieval

Exact Estimation: The influence score $I_\theta(X_s, X_i)$ is calculated as the sum of inner products between the compressed gradients of the generated image and the cached compressed gradients of training samples across timesteps.
KNN Retrieval: For tasks requiring only the top- $k$ most influential samples, DMin constructs a K-Nearest Neighbors (KNN) index (using HNSW algorithm) on the compressed gradients. This allows for on-the-fly retrieval of top- $k$ samples in under 1 second.

3. Key Contributions

Scalability to Billion-Parameter Models: DMin is the first method capable of influence estimation for diffusion models with billions of parameters (e.g., SD 3 Medium Full), a feat previously impossible due to storage constraints.
Extreme Storage Reduction: By combining gradient compression and normalization, DMin reduces storage requirements from hundreds of Terabytes to mere Megabytes or Kilobytes (e.g., reducing 339 TB to 726 MB for SD 3 Full).
Real-Time Retrieval: Leveraging KNN search on compressed gradients, DMin retrieves the top- $k$ influential samples in under 1 second, a massive speedup compared to existing methods.
Stability: The introduction of L2 normalization mitigates the fragility of influence estimation in deep networks.
Open Source: The authors provide a PyTorch implementation with multiprocessing support.

4. Experimental Results

The authors evaluated DMin on Stable Diffusion 1.4 (LoRA), Stable Diffusion 3 Medium (LoRA), and Stable Diffusion 3 Medium (Full), as well as unconditional DDPMs on MNIST and CIFAR-10.

Accuracy (Detection Rate):
- DMin significantly outperformed baselines (Random, SSIM, CLIP, LiSSA, DataInf, D-TRAK).
- On the "Magic Cards" subset with SD 3 Medium (Full), DMin achieved a Top-5 detection rate of ~96%, whereas baselines like SSIM and CLIP scored near 0%.
- On MNIST, DMin achieved a Top-5 detection rate of 0.8006, compared to 0.2560 for Journey-TRAK.
- Normalization is critical: Removing normalization caused a substantial drop in performance, confirming its role in stabilizing deep model gradients.
Efficiency (Time & Memory):
- Memory: For SD 3 Medium (Full), uncompressed gradients require 339 TB. DMin with compression ( $v=2^{12}$ ) reduced this to 726 MB (a compression ratio of 0.00017%).
- Time: Retrieving the top-5 influential samples for a test image took 0.004 seconds with DMin (KNN), compared to 2,939 seconds for LiSSA and 345 seconds for D-TRAK. This represents a speedup of over 100,000x in some configurations.

5. Significance

Democratizing Interpretability: DMin makes it feasible to analyze and interpret the behavior of state-of-the-art, billion-parameter diffusion models, which was previously restricted to small-scale research.
Safety and Compliance: By enabling rapid identification of the training origins of generated content, DMin aids in detecting copyright infringement, bias, and safety violations in generative AI.
Methodological Advancement: The paper demonstrates that careful gradient compression and normalization can overcome the "curse of dimensionality" in influence estimation, offering a blueprint for scalable interpretability in other large-scale deep learning domains.