DMin: Scalable Training Data Influence Estimation for Diffusion Models

1. Problemstellung

Diffusionsmodelle (DMs) sind leistungsstarke generative Modelle, die jedoch aufgrund ihrer breiten Trainingsdaten aus dem Internet oft unvorhergesehene oder schädliche Inhalte sowie Verzerrungen (Biases) generieren können. Ein zentrales Forschungsziel ist die Influenzschätzung (Influence Estimation): Die Fähigkeit, für ein generiertes Bild zu bestimmen, welche Trainingsdaten-Samples den größten Einfluss auf die Generierung hatten.

Bisherige Methoden zur Influenzschätzung stoßen bei großen Diffusionsmodellen an ihre Grenzen:

• Rechen- und Speicherkosten: Methoden zweiter Ordnung (z. B. DataInf, K-FAC) erfordern die Inversion der Hesse-Matrix oder das Laden aller Gradienten über mehrere Zeitschritte. Für ein Modell wie Stable Diffusion 3 Medium (ca. 2 Milliarden Parameter) würde das Speichern der Gradienten für 10.000 Samples über 10 Zeitschritte einen Speicherbedarf von ca. 800 TB erfordern.
• Skalierbarkeit: Erste-Ordnung-Methoden (z. B. D-TRAK, Journey-TRAK), die auf zufälligen Projektionen basieren, benötigen Projektionsmatrizen, die ebenfalls hunderte Terabyte Speicher belegen, wenn sie auf große Modelle angewendet werden.
• Fragilität: In extrem tiefen Modellen neigen Gradienten zu Instabilität (sehr große Werte), was die Genauigkeit der Schätzung beeinträchtigt.

Ziel der Arbeit ist es, einen skalierbaren Rahmen zu schaffen, der die Influenzschätzung für Diffusionsmodelle mit Milliarden von Paramern und großen Datensätzen ermöglicht, ohne die Speicher- oder Rechenanforderungen zu sprengen.

2. Methodik: DMin Framework

Die Autoren schlagen DMin (Diffusion Model influence) vor, ein Framework, das Gradienten komprimiert und effiziente Suchalgorithmen nutzt. Der Prozess gliedert sich in folgende Schritte:

A. Gradientenberechnung (Gradient Computation)

• Für jedes Trainings-Sample $(p_i, x_i)$ werden die Gradienten des Modells über einen ausgewählten Satz von Zeitschritten $t \in \{1, \dots, T\}$ berechnet, genau wie während des Trainings (Vorwärts- und Rückwärtsdurchlauf).
• Bei Modellen mit Adaptern (z. B. LoRA) werden nur die Gradienten der Adapter-Parameter gesammelt, bei Vollmodellen alle Parameter.
• Die Gradienten über alle Zeitschritte werden zu einem Vektor konkateniert.

B. Gradientenkompression (Gradient Compression)

Um den Speicherbedarf drastisch zu reduzieren, wenden die Autoren eine vierstufige Kompressionstechnik an, inspiriert von Vektor-Kompressionsverfahren:

1. Padding: Der Gradientenvektor wird auf eine Länge $L_{pad}$ aufgefüllt, die durch die Ziel-Dimension $v$ teilbar ist.
2. Permutation: Eine zufällige Permutation zerstört die inhärente Struktur des Vektors.
3. Random Projection: Elementweise Multiplikation mit einem zufälligen Vektor (Elemente sind -1 oder 1 mit gleicher Wahrscheinlichkeit), um Redundanz zu reduzieren.
4. Group Addition: Die Vektorelemente werden in Gruppen unterteilt und summiert, um den komprimierten Vektor der Dimension $v$ zu erhalten.

Speichereffizienz: Anstatt den vollen Gradientenvektor zu speichern, werden nur die Permutationsindizes (4 Bytes pro Element) und der binäre Projektionsvektor (1 Bit pro Element) gespeichert. Dies reduziert den Speicherbedarf pro Sample von Gigabytes auf Kilobytes oder Megabytes.

C. Normalisierung

Um die Instabilität von Gradienten in tiefen Modellen zu adressieren, wird vor der Kompression eine L2-Normalisierung auf die Gradientenvektoren angewendet. Dies verhindert, dass extrem große Gradientenwerte das Ergebnis der inneren Produkte dominieren.

D. Influenzschätzung und KNN-Suche

• Exakte Schätzung: Der Einfluss eines Trainings-Samples auf ein generiertes Bild wird durch das innere Produkt der komprimierten Gradientenvektoren geschätzt (basierend auf einer Taylor-Approximation der Verluständerung).
• KNN-Suche (K-Nearest Neighbors): Um die Top-$k$ einflussreichsten Samples schnell zu finden, wird ein KNN-Index (basierend auf dem HNSW-Algorithmus) auf den komprimierten Gradienten aufgebaut. Dies ermöglicht das Abrufen der Top-$k$ Samples in unter 1 Sekunde.

3. Wichtige Beiträge

1. Skalierbarkeit: DMin ist die erste Methode, die eine Influenzschätzung für Diffusionsmodelle mit Milliarden von Parametern (z. B. Stable Diffusion 3 Medium) durchführt.
2. Extreme Speicherreduktion: Durch die Kompressionstechnik wird der Speicherbedarf von hunderten Terabyte auf Megabyte oder Kilobyte pro Sample reduziert (z. B. von ~339 TB auf ~726 MB für den gesamten Datensatz bei SD 3 Medium Full).
3. Echtzeit-Retrieval: Die Kombination aus Kompression und KNN-Suche ermöglicht das Finden der Top-$k$ einflussreichsten Trainingsdaten in Sekundenbruchteilen.
4. Robustheit: Die Einführung der L2-Normalisierung löst das Problem der Gradienteninstabilität in sehr tiefen Modellen.
5. Open Source: Die Autoren stellen eine PyTorch-Implementierung mit Multi-Prozess-Unterstützung bereit.

4. Ergebnisse

Die Evaluation erfolgte auf verschiedenen Modellen (SD 1.4 LoRA, SD 3 Medium LoRA, SD 3 Medium Full) und Datensätzen (Blumen, Lego-Sets, Magic Cards, MNIST, CIFAR-10).

• Genauigkeit (Detektionsrate): DMin übertrifft alle Baselines (Random, SSIM, CLIP, LiSSA, DataInf, D-TRAK) signifikant.
  • Bei SD 3 Medium (Full) erreichte DMin (mit $v=2^{16}$) eine Detektionsrate von ~96–99 % für Top-50 und Top-100 Samples auf dem „Magic Cards"-Datensatz, während andere Methoden oft nicht anwendbar waren oder deutlich schlechter abschnitten.
  • Interessanterweise schnitt die KNN-Suche oft besser ab als die exakte innere Produktberechnung, da sie eine breitere Repräsentation der Nachbarn erfasst.
• Effizienz (Zeit):
  • Während Methoden wie LiSSA oder DataInf Stunden oder Tage benötigen, um die Influenz für einen Test-Sample zu berechnen, benötigt DMin mit KNN-Suche nur 0,004 bis 0,7 Sekunden.
  • Der Speedup gegenüber unkomprimierten Methoden liegt im Bereich von 1.000-fach bis über 100.000-fach.
• Speicher:
  • Für SD 3 Medium (Full) reduzierte sich der Speicherbedarf von 339,39 TB (unkomprimiert) auf 726 MB (mit $v=2^{12}$).

5. Bedeutung und Fazit

DMin löst das fundamentale Skalierungsproblem der Trainingsdaten-Influenzschätzung bei modernen Diffusionsmodellen. Es ermöglicht erstmals:

• Die Nachverfolgung der Herkunft spezifischer generierter Inhalte in riesigen Modellen.
• Die Identifizierung von Verzerrungen oder schädlichen Trainingsdaten in Echtzeit.
• Die praktische Anwendung von Influence Estimation in industriellen Szenarien, wo Speicher- und Rechenressourcen begrenzt sind.

Die Arbeit demonstriert, dass durch intelligente Gradientenkompression und Normalisierung die Lücke zwischen theoretischer Machbarkeit und praktischer Anwendbarkeit bei großen generativen KI-Modellen geschlossen werden kann. Dies ist ein wichtiger Schritt hin zu transparenteren, sichereren und besser kontrollierbaren Diffusionsmodellen.