Diff-ES: Stage-wise Structural Diffusion Pruning via Evolutionary Search

Die Arbeit stellt Diff-ES vor, einen neuartigen Pruning-Ansatz für Diffusionsmodelle, der mittels evolutionärer Suche ein optimales stufenweises Sparsitäts-Schema ermittelt und durch effizientes Weight-Routing eine signifikante Beschleunigung bei minimalen Qualitätsverlusten ermöglicht.

Das Wesentliche

Stell dir vor, ein Diffusions-Modell (die KI, die Bilder malt) ist wie ein Künstler, der ein riesiges Gemälde erstellt.

Normalerweise arbeitet dieser Künstler sehr langsam: Er beginnt mit einem komplett verschmierten, verrauschten Bild und entfernt schrittweise das Rauschen, bis das Bild klar wird. Das Problem ist, dass er für jeden dieser Schritte fast die gleiche Menge an Energie und Zeit aufwendet – egal, ob er gerade erst den groben Umriss zeichnet oder bereits die feinen Details der Augen verfeinert.

Die bisherigen Methoden, um diesen Künstler schneller zu machen, waren wie ein starrer Bauplan: Man hat ihm gesagt, er soll in jedem Schritt genau gleich viel "schummeln" (also Teile seines Gehirns abschalten), um Zeit zu sparen. Das Problem dabei: Das funktioniert nicht gut. Manchmal ist der Künstler in einem Schritt sehr wichtig (z. B. beim Zeichnen der Konturen), und in einem anderen Schritt weniger wichtig (beim Verfeinern der Textur). Wenn man ihn überall gleich stark drosselt, wird das Bild am Ende oft schief oder unscharf.

Was ist Diff-ES?

Diff-ES ist wie ein genialer Manager, der diesen Künstler nicht starr regelt, sondern ihm einen intelligenten, dynamischen Arbeitsplan gibt.

Hier ist die Idee in drei einfachen Schritten:

1. Die Reise in Etappen (Stages)

Statt den gesamten Malprozess als einen langen, eintönigen Weg zu sehen, teilt Diff-ES die Reise in Etappen auf.

• Frühe Etappen: Hier wird das Grundgerüst gelegt. Das ist wie das Skizzieren der groben Formen.
• Späte Etappen: Hier werden die Details hinzugefügt. Das ist wie das Hinzufügen von Lichtreflexen oder Haartexturen.

2. Der evolutionäre Sucher (Evolutionary Search)

Früher haben Forscher versucht, den besten Arbeitsplan selbst zu erraten (wie "Vielleicht ist Schritt 1 wichtig, Schritt 2 weniger"). Das war oft falsch.
Diff-ES nutzt stattdessen einen evolutionären Algorithmus. Stell dir das wie eine Schule für Künstler-Assistenten vor:

• Der Manager erstellt viele verschiedene Arbeitspläne (z. B. "Plan A: In der Mitte stark sparen", "Plan B: Am Anfang stark sparen").
• Jeder Plan wird getestet: Der Künstler malt ein Bild.
• Ein Bewerter (ein einfacher, schneller Check) schaut sich das Bild an: "Sieht das gut aus? Ist das Gesicht noch erkennbar?"
• Die besten Pläne überleben, die schlechten werden verworfen.
• Die Überlebenden werden leicht verändert (mutiert), um neue, noch bessere Pläne zu finden.
• Nach vielen Runden hat der Manager den perfekten, maßgeschneiderten Plan gefunden, der genau weiß, wann er sparen darf und wann nicht.

3. Der clevere Umzug (Weight Routing)

Ein großes Problem bei früheren Methoden war, dass man für jede Etappe einen komplett separaten Künstler (Modell) im Speicher halten musste. Das war wie ein riesiger Lagerhallen-Overhead.
Diff-ES ist schlauer: Es hat einen einzigen Künstler, aber es tauscht ihm seine Werkzeuge dynamisch aus.

• Wenn der Künstler in der "Skizzen-Etappe" ist, bekommt er Werkzeuge, die für grobe Formen optimiert sind (wenige Details).
• Wenn er in die "Detail-Etappe" kommt, werden ihm sofort die feinen Pinsel gegeben.
• Das passiert so schnell und effizient, dass man keine riesigen Lagerhallen (Speicher) braucht. Es ist wie ein Schweizer Taschenmesser, das je nach Aufgabe das richtige Werkzeug ausklappt, ohne dass man das ganze Messer wechseln muss.

Warum ist das so toll?

• Bessere Bilder: Weil der Plan perfekt auf die Bedürfnisse jedes Schrittes abgestimmt ist, sieht das Endergebnis fast genauso gut aus wie das des langsamen, ungeschnittenen Künstlers.
• Viel schneller: Da in den weniger wichtigen Etappen massiv gespart wird, ist das Bild viel schneller fertig.
• Kein Neulernen: Der Künstler muss nicht neu trainiert werden. Diff-ES findet einfach den besten Weg, ihn während des Arbeitens zu optimieren.

Zusammenfassend:
Diff-ES ist wie ein perfekter Choreograf für einen Tanz. Anstatt den Tänzer den ganzen Tanz über gleich schnell tanzen zu lassen (was ihn erschöpft und ungenau macht), sagt er: "Hier, in der langsamen Passage, sei sehr vorsichtig und gib alles! Aber in der schnellen Drehung, spare Energie und mach es einfach!" Das Ergebnis ist ein schnellerer Tanz, der trotzdem wunderschön aussieht.

Technische Zusammenfassung

1. Problemstellung

Diffusionsmodelle haben sich als Standard für die hochauflösende Bildgenerierung etabliert, leiden jedoch unter hohem Rechenaufwand aufgrund ihrer mehrstufigen Denoisings-Prozesse und großer Modellgrößen. Bisherige Ansätze zur Effizienzsteigerung konzentrierten sich entweder auf die Reduzierung der Sampling-Schritte oder auf die Komprimierung der Modellparameter.

Ein zentrales Problem bestehender strukturierter Pruning-Methoden (das Entfernen ganzer Kanäle oder Schichten) ist die Annahme, dass alle Denoising-Schritte gleich wichtig sind. Neuere Arbeiten zeigen jedoch, dass die Rechenressourcen je nach Schritt variieren sollten. Ein aktueller Ansatz, MosaicDiff, versucht dies durch manuell festgelegte, heuristische Schedules (z. B. Aufteilung in drei grobe Phasen) zu lösen. Dabei werden jedoch mehrere unabhängig voneinander pruned Modelle während der Inferenz „gestickt" (zusammengefügt), was den Speicherverbrauch erhöht. Zudem sind diese manuellen Schedules oft nicht generalisierbar und führen zu suboptimalen Ergebnissen, da die Wichtigkeit der Schritte stark modellabhängig ist.

2. Methodik: Diff-ES

Das Paper stellt Diff-ES vor, einen Framework zur stufenweisen strukturellen Pruning von Diffusionsmodellen mittels evolutionärer Suche. Das Ziel ist die automatische Optimierung des Sparsity-Schedules (Dichte der Parameter) über den Denoising-Trajektorie hinweg unter Einhaltung eines globalen Budgets.

Die Methode besteht aus folgenden Kernkomponenten:

• Stufenweise Aufteilung: Der Denoising-Prozess wird in $n$ Stufen unterteilt. Jede Stufe erhält einen spezifischen Sparsity-Level (Dichte), der bestimmt, wie stark das Modell in dieser Phase komprimiert wird.
• Evolutionäre Suche (Level-Switch Mutation):
  • Anstatt einen festen Schedule zu verwenden, wird eine Population von Kandidaten-Schedules initialisiert.
  • Ein Level-Switch-Mutations-Operator tauscht Sparsity-Level zwischen zwei zufällig ausgewählten Stufen aus, während das globale Sparsity-Budget strikt eingehalten wird.
  • Die Fitness jedes Kandidaten wird durch das Generieren einer kleinen Bildmenge und die Berechnung eines leichtgewichtigen Metriks (z. B. CLIP-IQA, TOPIQ oder SSIM) bewertet.
  • Die besten Individuen überleben und bilden die Basis für die nächste Generation. Dieser Prozess wiederholt sich, bis ein optimaler Schedule gefunden ist.
• SNR-bewusste Kalibrierung: Um die Pruning-Entscheidungen zu treffen, wird ein Kalibrierungsdatensatz erstellt, der die Signal-zu-Rausch-Verhältnisse (SNR) der jeweiligen Diffusionsstufe widerspiegelt. Dies ermöglicht eine präzise Schätzung der lokalen Krümmung (Hessian-Matrix) für das Pruning.
• Second-Order Strukturiertes Pruning: Der Framework nutzt fortschrittliche Pruning-Methoden wie OBS (Optimal Brain Surgeon), Wanda oder LayerDrop. OBS nutzt zweite Ordnungsinformationen, um redundante Parameter zu entfernen und die verbleibenden Gewichte so anzupassen, dass der Fehler minimiert wird.
• Effizientes Weight Routing (Schlüsselinnovation):
  • Ein Hauptproblem bei der Kombination von evolutionärer Suche und OBS ist der hohe Rechenaufwand, da OBS iterativ Gewichte aktualisiert.
  • Diff-ES löst dies, indem es die vollständige Pruning-Trajektorie für jede Stufe einmalig vorab berechnet und die aktualisierten Projektionsgewichte in einer kompakten Datenbank speichert.
  • Während der Suche wird kein neues Pruning durchgeführt. Stattdessen wird ein Weight-Routing-Mechanismus verwendet, der die benötigten, stufenspezifischen Gewichte dynamisch aus der Datenbank lädt und im Speicher des Hauptmodells austauscht.
  • Dies vermeidet die Notwendigkeit, mehrere vollständige Modelle im GPU-Speicher zu halten (wie bei MosaicDiff) und eliminiert redundante Berechnungen.

3. Wichtige Beiträge

1. Identifikation einer Limitierung: Die Autoren zeigen auf, dass heuristische, manuell getunte Schedules (wie bei MosaicDiff) nicht generalisieren und oft schlechter abschneiden als adaptive Lösungen.
2. Neuer Optimierungsrahmen: Einführung eines evolutionären Suchframeworks, das den Sparsity-Schedule automatisch an die spezifischen Anforderungen des Modells (z. B. DiT vs. SDXL) anpasst.
3. Plug-and-Play Kompatibilität: Diff-ES ist ein generischer Rahmen, der mit verschiedenen strukturierten Pruning-Methoden (LayerDrop, Wanda, OBS) kombiniert werden kann.
4. Speichereffizienz: Durch das Weight-Routing wird der Speicherbedarf im Vergleich zum „Stitching"-Ansatz (mehrere Modelle laden) drastisch reduziert, ohne die Genauigkeit zu opfern.

4. Ergebnisse

Die Methode wurde auf zwei führenden Architekturen evaluiert: DiT (Transformer-basiert) und SDXL (U-Net-basiert).

• Qualität vs. Geschwindigkeit: Diff-ES erreicht konsistent bessere Ergebnisse als starke Baselines (Diff-Pruning, DeepCache, OBS-Diff, MosaicDiff).
  • Auf SDXL bei 30% Sparsity: Diff-ES erreicht einen FID von 25,87 (verglichen mit 59,09 bei MosaicDiff und 28,49 bei OBS-Diff). Dies zeigt eine drastische Verbesserung der Bildqualität bei gleicher Kompression.
  • Auf DiT bei 50% Sparsity: Diff-ES erreicht einen FID von 12,86 (verglichen mit 22,29 bei MosaicDiff).
• Anpassungsfähigkeit: Die gefundenen Schedules unterscheiden sich je nach Modell. Während MosaicDiff einen festen Heuristik-Schedule verwendet, findet Diff-ES für DiT und SDXL unterschiedliche Muster (z. B. höhere Dichte in frühen/späten Phasen bei SDXL), was zu besseren Ergebnissen führt.
• Effizienz: Diff-ES erzielt signifikante Beschleunigungen (Wall-clock Speedups) bei minimalem Qualitätsverlust. Der Overhead durch die einmalige Suche ist gering im Vergleich zum Nutzen.
• Speichernutzung: Im Vergleich zum „Model Stitching" von MosaicDiff spart Diff-ES durch Weight Routing ca. 42,7% GPU-Speicher (beim DiT-Modell).

5. Bedeutung und Fazit

Diff-ES adressiert das fundamentale Problem, dass Diffusionsmodelle in verschiedenen Phasen ihrer Generierung unterschiedlich empfindlich auf Kompression reagieren. Durch die Automatisierung der Ressourcenallokation mittels evolutionärer Suche und die Einführung eines effizienten Weight-Routing-Mechanismus überwindet die Methode die Nachteile manueller Heuristiken und speicherintensiver Ansätze.

Die Arbeit etabliert einen neuen State-of-the-Art für strukturiertes Pruning von Diffusionsmodellen. Sie zeigt, dass eine intelligente, datengetriebene Verteilung der Rechenressourcen über die Denoising-Schritte hinweg sowohl die Generierungsqualität als auch die Effizienz signifikant steigern kann, ohne dass ein Nachtrainieren (Retraining) des Modells erforderlich ist. Dies macht Diff-ES zu einer robusten und allgemeinen Lösung für die Beschleunigung von Diffusionsmodellen in der Praxis.