2025-11-23T17:13:17.428108

Token Pruning for Caching Better: 9 Times Acceleration on Stable Diffusion for Free

Zhang, Xiao, Tang et al.
Stable Diffusion has achieved remarkable success in the field of text-to-image generation, with its powerful generative capabilities and diverse generation results making a lasting impact. However, its iterative denoising introduces high computational costs and slows generation speed, limiting broader adoption. The community has made numerous efforts to reduce this computational burden, with methods like feature caching attracting attention due to their effectiveness and simplicity. Nonetheless, simply reusing features computed at previous timesteps causes the features across adjacent timesteps to become similar, reducing the dynamics of features over time and ultimately compromising the quality of generated images. In this paper, we introduce a dynamics-aware token pruning (DaTo) approach that addresses the limitations of feature caching. DaTo selectively prunes tokens with lower dynamics, allowing only high-dynamic tokens to participate in self-attention layers, thereby extending feature dynamics across timesteps. DaTo combines feature caching with token pruning in a training-free manner, achieving both temporal and token-wise information reuse. Applied to Stable Diffusion on the ImageNet, our approach delivered a 9$\times$ speedup while reducing FID by 0.33, indicating enhanced image quality. On the COCO-30k, we observed a 7$\times$ acceleration coupled with a notable FID reduction of 2.17.
academic

Token Pruning for Caching Better: 9× Beschleunigung auf Stable Diffusion kostenlos

Grundinformationen

  • Paper-ID: 2501.00375
  • Titel: Token Pruning for Caching Better: 9× Acceleration on Stable Diffusion for Free
  • Autoren: Evelyn Zhang, Bang Xiao, Jiayi Tang, Qianli Ma, Chang Zou, Xuefei Ning, Xuming Hu, Linfeng Zhang
  • Klassifizierung: cs.CV (Computervision), cs.LG (Maschinelles Lernen)
  • Veröffentlichungsdatum: 31. Dezember 2024
  • Paper-Link: https://arxiv.org/abs/2501.00375
  • Code-Link: github.com/EvelynZhang-epiclab/DaTo

Zusammenfassung

Stable Diffusion hat bedeutende Erfolge im Bereich der Text-zu-Bild-Generierung erzielt, doch sein iterativer Entrauschungsmechanismus führt zu hohen Rechenkosten und langsamer Generierungsgeschwindigkeit. Obwohl Methoden wie Feature-Caching aufgrund ihrer Effektivität und Einfachheit Aufmerksamkeit erhalten, führt die einfache Wiederverwendung von Features aus vorherigen Zeitschritten dazu, dass Features zwischen benachbarten Zeitschritten ähnlich werden, was die Dynamik der Features über die Zeit verringert und letztendlich die Qualität der generierten Bilder beeinträchtigt. Dieses Paper präsentiert eine dynamikbewusste Token-Pruning-Methode (DaTo), um die Einschränkungen des Feature-Caching zu überwinden. DaTo beschneidet selektiv Tokens mit niedriger Dynamik und ermöglicht nur hochdynamischen Tokens, an Self-Attention-Schichten teilzunehmen, wodurch die Feature-Dynamik zwischen Zeitschritten erweitert wird. Bei Anwendung auf Stable Diffusion auf ImageNet erreicht die Methode eine 9×-Beschleunigung, während die FID um 0,33 sinkt; auf COCO-30k wird eine 7×-Beschleunigung mit signifikantem FID-Rückgang von 2,17 beobachtet.

Forschungshintergrund und Motivation

Problemhintergrund

Diffusionsmodelle haben bedeutende Fortschritte in der generativen Modellierung erzielt und werden weit verbreitet in Text-zu-Bild-Generierung, Videogenerierung und anderen Aufgaben eingesetzt. Allerdings führt der iterative Entrauschungsmechanismus von Diffusionsmodellen zu enormen Rechenkosten und langsamer Generierungsgeschwindigkeit, was ihre breitere Anwendung einschränkt.

Einschränkungen bestehender Methoden

Die aktuellen Methoden zur Beschleunigung von Diffusionsmodellen umfassen hauptsächlich:

  1. Reduzierung der Sampling-Schritte: wie schnelle Sampler wie DDIM
  2. Reduzierung der Rechenkosten pro Schritt: einschließlich Wissensdestillation, Strukturpruning, Quantisierung, Token-Pruning und Feature-Caching

Unter diesen ist Feature-Caching aufgrund seiner Effektivität und Einfachheit weit verbreitet. Es speichert Features aus vorherigen Zeitschritten und verwendet sie in nachfolgenden Zeitschritten erneut. Allerdings zwingt die Feature-Wiederverwendung Features verschiedener Zeitschritte, ähnliche Werte zu haben, was die Dynamik der Features über Zeitschritte hinweg verringert, den ursprünglichen Diffusionsprozess beschädigt und somit die Generierungsqualität beeinträchtigt.

Forschungsmotivation

Das Paper beobachtet durch Experimente, dass die Feature-Unterschiede zwischen benachbarten Zeitschritten bei Modellen mit Feature-Caching im Vergleich zum ursprünglichen Stable Diffusion signifikant abnehmen. Dies wirft eine kritische Frage auf: Ist es möglich, Feature-Caching durchzuführen und gleichzeitig die korrekte Feature-Dynamik zu bewahren?

Kernbeiträge

  1. Vorschlag der dynamikbewussten Token-Pruning-Methode (DaTo): Durch das Beschneiden von Tokens, deren Dynamik durch Feature-Caching in verschiedenen Zeitschritten verringert wird, und deren Wiederherstellung durch Tokens mit großer Dynamik, wird die Qualitätsverschlechterung vermieden, die durch Feature-Caching verursacht wird.
  2. Entwurf einer evolutionären Suchstrategie: Vorschlag einer evolutionären Methode zur Suche nach optimalen Feature-Caching- und Token-Pruning-Strategien, um das volle Potenzial von DaTo freizusetzen.
  3. Erreichung signifikanter Leistungsverbesserungen: Umfangreiche Experimente auf Stable Diffusion und SDXL zeigen, dass ohne Training und zusätzliche Daten auf Stable Diffusion bis zu 9×-Beschleunigung mit verlustfreier Generierungsqualität erreicht werden kann.

Methodische Details

Aufgabendefinition

Die Aufgabe dieses Papers besteht darin, den Inferenzprozess des Stable Diffusion-Modells erheblich zu beschleunigen, während die Bildgenerierungsqualität erhalten bleibt. Die Eingabe ist ein Textprompt, die Ausgabe ist das entsprechende hochwertige Bild, und die Einschränkung besteht darin, dass das Modell nicht neu trainiert werden muss.

Modellarchitektur

1. Dynamikbewusstes Token-Pruning (DaTo)

Basis-Token-Auswahl:

  • Zeitliche Rausch-Differenz-Bewertung: Für den t-ten Zeitschritt wird die absolute Differenz der Ausgaben der beiden benachbarten vorherigen Zeitschritte berechnet: 
    DiffScore = (1/C) * Σ|f_up_0(x_{t+2})_c - f_up_0(x_{t+1})_c|
  • Patch-basierte Token-Auswahl: Das Bild wird in nicht überlappende s×s-Patches unterteilt, und in jedem Patch wird der Token mit dem höchsten DiffScore als Basis-Token ausgewählt.

CFG-Ausrichtung: Um die klassifiziererfreie Anleitung (CFG) zu handhaben, werden die Basis-Token-Positionen der bedingten Generierung in die unbedingte Generierung kopiert:

X_base,i,j[k] = X_base,i,j[k - B/2], k ∈ {B/2, B/2+1, ..., B-1}

Pruning-Token-Auswahl: Basierend auf der Kosinusähnlichkeit werden die K Tokens ausgewählt, die den Basis-Tokens am ähnlichsten sind, um sie zu beschneiden:

X_prune = arg topK max Cosine_Similarity(X_i, X_j)

Pruning-Token-Wiederherstellung: Die beschnittenen Tokens werden durch direkte Kopie ihres ähnlichsten Basis-Tokens wiederhergestellt.

2. Zeitschritt-bewusstes Feature-Caching

Suchraum-Beschneidung:

  • Caching-Tiefe d ist auf {0, 1, 1/2} beschränkt
  • Pruning-Verhältnis r ist auf {0,3, 0,4, 0,5, 0,6, 0,7} beschränkt

Evolutionärer Suchalgorithmus: Der NSGA-II-Multiziel-Optimierungsalgorithmus wird verwendet, mit Optimierungszielen einschließlich:

  • Inferenz-Latenz
  • Generierungsqualität (FID)

Der Suchprozess umfasst Standard-Evolutionsoperationen wie Selektion, Crossover und Mutation, um letztendlich die optimale zeitschrittbewusste Strategie F(t) zu erhalten.

Technische Innovationspunkte

  1. Dynamik-Wiederherstellungsmechanismus: Durch selektives Beschneiden von Tokens mit niedriger Dynamik und deren Wiederherstellung mit hochdynamischen Tokens wird die durch Feature-Caching beschädigte Feature-Dynamik-Verteilung erfolgreich wiederhergestellt.
  2. Einheitliches Caching-Pruning-Framework: Feature-Caching und Token-Pruning werden in einem trainingsunabhängigen Framework kombiniert, um Informationswiederverwendung auf Zeit- und Token-Ebene zu erreichen.
  3. Adaptive Strategiesuche: Für die unterschiedlichen Redundanzeigenschaften verschiedener Zeitschritte wird eine Methode zur automatischen Suche nach optimaler Caching-Tiefe und Pruning-Verhältnis vorgeschlagen.

Experimentelle Einrichtung

Datensätze

  • ImageNet-1k: Generierung von 2000 Bildern mit 512×512 Pixeln (2 pro Klasse)
  • COCO-30k: Generierung von 30000 Bildern (1 pro Bildunterschrift)
  • MS COCO-Validierungssatz: Für SDXL-Bewertung, Generierung von 5k Bildern mit 1024×1024 Pixeln

Bewertungsmetriken

  • FID (Fréchet Inception Distance): Messung der Generierungsqualität
  • CLIP Score: Bewertung der Text-Bild-Ausrichtung
  • Inception Score: Bildqualitätsbewertung
  • Latenz und Beschleunigungsverhältnis: Effizienzbewertung

Vergleichsmethoden

  • DDIM/DPM: Schnelle Sampler
  • ToMeSD: Token-Merge-Methode
  • DeepCache: Feature-Caching-Methode
  • DeepCache & ToMeSD: Naive Kombinationsmethode

Implementierungsdetails

  • Verwendung des NSGA-II-Evolutionsalgorithmus mit Populationsgröße 20, 100 Generationen
  • CFG-Skala: 7,5 (SD v1.5), 9,0 (SD v2), 7,0 (SDXL)
  • Sampling-Schritte: 50 PLMS-Schritte
  • Tests auf einer einzelnen 4090-GPU

Experimentelle Ergebnisse

Hauptergebnisse

Stable Diffusion v1.5 (ImageNet):

  • Konfiguration e1: 9,01×-Beschleunigung, FID sinkt von 27,64 auf 27,31
  • Übertrifft alle Vergleichsmethoden in allen Konfigurationen

Stable Diffusion v2 (ImageNet):

  • Konfiguration e2: 7,25×-Beschleunigung, FID von 28,20
  • Im Vergleich zum ursprünglichen Modell FID von 29,8 auf 28,20 gesunken

COCO-30k-Datensatz:

  • SD v1.5: 7×-Beschleunigung, FID sinkt von 12,15 auf 9,98 (Rückgang von 2,17)
  • SD v2: 7,25×-Beschleunigung, FID von 13,68 auf 13,88

SDXL (MS COCO):

  • 2,32×-Beschleunigung, FID sinkt von 24,25 auf 23,10
  • Signifikant besser als DeepCache (1,75×) und DeepCache&ToMeSD (1,78×)

Ablationsstudien

Effektivität von DiffScore: Bei verschiedenen Caching-Einstellungen und Pruning-Verhältnissen verbessert die Verwendung von DiffScore konsistent die FID-Werte und beweist die Effektivität der zeitlichen Rausch-Differenz-Bewertung.

Auswirkung der CFG-Ausrichtung: Mit zunehmendem Pruning-Verhältnis nimmt der Nutzen der CFG-Ausrichtungskonfiguration allmählich zu, wobei die FID-Verbesserung bei hohem Pruning-Verhältnis (0,7) zwischen 13 und 30 Punkten liegt.

Fallstudien

Visuelle Vergleichsergebnisse zeigen, dass DaTo in mehreren Aspekten hervorragende Leistungen erbringt:

  • Inhaltstreu: Hohe Ähnlichkeit mit Originalbildinhalten
  • Detailbewahrung: Beibehaltung feiner Texturen in hochdetaillierten Szenen
  • Stilanpassung: Ausgewogene Inhaltsbewahrung und Stilgenauigkeit bei Bild-zu-Bild-Aufgaben
  • Prompt-Ausrichtung: Genaue Generierung aller Elemente in komplexen Textprompts

Experimentelle Erkenntnisse

  1. Feature-Dynamik-Wiederherstellung: DaTo stellt die Feature-Differenz-Verteilung erfolgreich auf ein Niveau nahe dem ursprünglichen Stable Diffusion wieder her
  2. Sparsame Kodierungseffekt: Moderates Token-Pruning und Feature-Caching können die Modellleistung durch Konzentration auf Schlüsselfeatures verbessern
  3. Strategieverallgemeinerung: Die auf SD v1.5 gesuchte Strategie zeigt gute Leistung auf SDXL und anderen Datensätzen

Verwandte Arbeiten

Effiziente Diffusionsmodelle

  • Reduzierung der Sampling-Schritte: DDIM, Konsistenzmodelle usw.
  • Netzwerkkompression: Quantisierung, Pruning, Destillation usw.
  • Architekturoptimierung: U-Net-Verbesserungen, Transformer-Optimierungen usw.

Token-Reduktionsstrategien

  • Lernmethoden: DynamicViT, A-ViT usw., die Hilfsmodelle zum Sortieren und Beschneiden verwenden
  • Heuristische Methoden: Token Pooling, Token Merging usw., trainingsunabhängige Methoden
  • Anwendung auf Diffusionsmodelle: ToMeSD, AT-EDM usw., an Generierungsaufgaben angepasste Methoden

Caching-Mechanismen

  • U-Net-Caching: DeepCache nutzt zeitliche Redundanz zum Caching von Features
  • DiT-Caching: Δ-DiT-Caching-Strategie für Diffusion Transformer
  • Optimierungsherausforderungen: Ausgewogenheit zwischen Effizienzsteigerung und Qualitätserhaltung

Schlussfolgerungen und Diskussion

Hauptschlussfolgerungen

  1. DaTo löst erfolgreich das Problem des Feature-Dynamik-Verlusts, der durch Feature-Caching verursacht wird
  2. Die durch evolutionäre Suche erhaltene adaptive Strategie ist signifikant besser als feste Konfigurationen
  3. Die Methode erreicht auf mehreren Modellen und Datensätzen signifikante Beschleunigung und Qualitätsverbesserung

Einschränkungen

  1. Suchkosten: Obwohl ≤20 GPU-Stunden akzeptabel sind, ist dennoch zusätzliche Rechenressource erforderlich
  2. Hardware-Abhängigkeit: Leistungsverbesserungen können je nach Hardware-Konfiguration variieren
  3. Einschränkungen bei extremen Einstellungen: Zu hohes Pruning-Verhältnis oder zu niedrige Caching-Aktualisierungsfrequenz beeinträchtigen die Leistung

Zukünftige Richtungen

  1. Adaptive Strategielernverfahren: Entwicklung intelligenterer adaptiver Caching- und Pruning-Strategien
  2. Anpassung an andere Architekturen: Erweiterung auf mehr Diffusionsmodell-Architekturen
  3. Theoretische Analyse: Tieferes Verständnis der Rolle des sparsamen Kodierungsprinzips in Diffusionsmodellen

Tiefgreifende Bewertung

Stärken

  1. Hohe Innovativität: Erste systematische Lösung des Feature-Dynamik-Verlustproblems bei Feature-Caching
  2. Praktische Methode: Trainingsunabhängig, einfach zu implementieren und zu integrieren
  3. Umfassende Experimente: Vollständige Bewertung auf mehreren Modellen und Datensätzen
  4. Theoretische Unterstützung: Bietet theoretische Erklärung durch sparsame Kodierung
  5. Open-Source-freundlich: Vollständige Code-Implementierung bereitgestellt

Schwächen

  1. Unzureichende theoretische Analyse: Theoretische Erklärung für die Verbesserung der FID ist relativ einfach
  2. Abhängigkeit vom Suchalgorithmus: Benötigt evolutionäre Suche zur Findung optimaler Strategien, erhöht Komplexität der Verwendung
  3. Einzelne Bewertungsmetriken: Hauptsächlich auf FID-Bewertung angewiesen, mangelnde Vielfalt bei Qualitätsmetriken
  4. Fehlende Benutzerstudien: Keine menschliche Bewertung zur Validierung der Generierungsqualität

Auswirkungen

  1. Akademischer Wert: Bietet neue Ideen und Methoden für die Beschleunigung von Diffusionsmodellen
  2. Praktischer Wert: Kann direkt auf bestehende Stable Diffusion-Modelle angewendet werden
  3. Reproduzierbarkeit: Detaillierte Implementierungsdetails und Open-Source-Code bereitgestellt
  4. Inspirationswert: Bietet Beispiel für Token-Level-Optimierung in generativen Modellen

Anwendungsszenarien

  1. Ressourcenbegrenzte Umgebungen: Mobile Geräte, Edge-Computing-Szenarien
  2. Echtzeitanwendungen: Interaktive Anwendungen, die schnelle Generierung erfordern
  3. Batch-Generierung: Großflächige Bildgenerierungsaufgaben
  4. Forschungsprototypen: Forschungsprojekte, die schnelle Iteration erfordern

Literaturverzeichnis

Das Paper zitiert 46 verwandte Literaturquellen, die Diffusionsmodelle, Token-Reduktion, Caching-Mechanismen und andere verwandte Bereiche abdecken und eine solide theoretische Grundlage und Vergleichsbenchmarks für diese Forschung bieten.

Gesamtbewertung: Dies ist ein hochqualitatives Computervisions-Paper, das eine innovative Lösung für das wichtige Problem der Diffusionsmodell-Beschleunigung bietet. Das Methodendesign ist elegant, die experimentelle Bewertung umfassend und der praktische Wert hervorragend. Obwohl die theoretische Analysentiefe etwas zu wünschen übrig lässt, sind die praktischen Beiträge und Auswirkungen bemerkenswert.