Relational Feature Caching for Accelerating Diffusion Transformers

Die Arbeit stellt Relational Feature Caching (RFC) vor, ein neuartiges Framework, das durch die Nutzung der Korrelation zwischen Eingabe- und Ausgabe-Features die Vorhersagegenauigkeit bei Diffusion-Transformern verbessert und so die Rechenzeit im Vergleich zu rein zeitbasierten Caching-Ansätzen signifikant reduziert.

Das Wesentliche

Das große Problem: Der müde Maler

Stell dir vor, du hast einen genialen, aber sehr langsamen Künstler (das ist das Diffusions-Modell). Dieser Künstler soll ein Bild malen, indem er von einem riesigen, chaotischen Farbspritzer (Rauschen) beginnt und Schritt für Schritt das Bild klarer macht.

Um ein einziges perfektes Bild zu erstellen, muss dieser Künstler 50 Mal über den gesamten Leinwand streichen. Bei jedem Streich muss er extrem viel rechnen, um zu entscheiden, welche Farbe wohin gehört. Das dauert ewig und kostet viel Energie.

Die alte Lösung: Das „Abklatsch"-Verfahren

Um schneller zu sein, haben andere Forscher eine Idee gehabt: „Warum müssen wir bei jedem Schritt neu rechnen? Die Farben ändern sich doch von Schritt zu Schritt nur ganz wenig!"

Also haben sie einen Trick angewandt:

1. Sie lassen den Künstler alle 5 Schritte das Bild wirklich neu berechnen.
2. Bei den dazwischenliegenden Schritten (Schritt 6, 7, 8) sagen sie: „Mach mal Pause! Wir nehmen einfach das Bild von Schritt 5 und sagen: 'Schritt 6 sieht fast genauso aus'."

Das Problem dabei: Manchmal ändert sich das Bild plötzlich stark (z. B. wenn eine neue Form entsteht). Wenn der Künstler dann einfach nur das alte Bild kopiert, entstehen Fehler. Das Bild wird unscharf oder verzerrt.

Die neue Lösung: RFC (Der kluge Assistent)

Die Autoren dieses Papiers haben sich das genauer angesehen und zwei Dinge bemerkt:

1. Die Änderungen im Bild sind oft unregelmäßig. Man kann sie nicht einfach linear vorhersagen (wie bei einer geraden Linie).
2. Aber! Es gibt eine starke Verbindung zwischen dem, was der Künstler sieht (Eingabe), und dem, was er malt (Ausgabe).

Ihre neue Methode, RFC, funktioniert wie ein kluger Assistent, der dem Künstler hilft:

1. Der „Verhältnis-Check" (Relational Feature Estimation)

Statt nur zu raten, wie sich das Bild ändert, schaut der Assistent auf die Eingabe.

• Die Analogie: Stell dir vor, der Künstler malt eine Mauer aus Ziegeln. Wenn der Assistent sieht, dass der Künstler gerade einen neuen Ziegel in die Hand nimmt (Änderung in der Eingabe), weiß er sofort: „Aha, das bedeutet, das Bild wird sich jetzt stärker ändern als beim letzten Mal!"
• Frühere Methoden sagten nur: „Es wird sich wahrscheinlich ein bisschen ändern."
• RFC sagt: „Da du einen neuen Ziegel hast, wird sich das Bild genau so stark ändern wie beim letzten Mal, als du einen neuen Ziegel hattest."
• Ergebnis: Die Vorhersage ist viel genauer, weil sie auf dem Zusammenhang zwischen Ursache (Eingabe) und Wirkung (Ausgabe) basiert, nicht nur auf der Zeit.

2. Der „Warnschalter" (Relational Cache Scheduling)

Der Assistent weiß auch, wann er aufhören soll zu raten.

• Die Analogie: Der Assistent überwacht den Künstler. Solange die Änderungen vorhersehbar sind, spart er Zeit. Aber sobald er merkt: „Hoppla, hier wird es chaotisch, die Vorhersage wird schlecht!", sagt er: „Stopp! Rechne jetzt das Bild wirklich neu aus!"
• Frühere Methoden rechneten in festen Abständen neu (z. B. alle 5 Schritte).
• RFC rechnet genau dann neu, wenn es nötig ist. Wenn das Bild stabil ist, spart er Zeit. Wenn es komplex wird, investiert er Rechenleistung.

Warum ist das so toll?

Stell dir vor, du fährst Auto:

• Die alten Methoden fahren mit Tempomaten auf einer kurvigen Straße. Sie halten die Geschwindigkeit konstant, auch wenn die Kurven enger werden. Das ist unsicher und führt zu Abdrift.
• RFC ist wie ein Fahrer mit einem Navigationsgerät, das die Kurven erkennt. Er bremst rechtzeitig, wenn die Kurve kommt, und beschleunigt, wenn die Straße gerade ist.

Das Ergebnis:

• Die Bilder sind viel schärfer und genauer (weniger Fehler).
• Es ist trotzdem viel schneller als das vollständige Berechnen jedes Schrittes.
• Es funktioniert bei Text-zu-Bild, Text-zu-Video und vielen anderen Aufgaben.

Zusammenfassung in einem Satz

RFC ist wie ein intelligenter Assistent, der nicht einfach nur das Alte kopiert, sondern genau hinschaut, wie sich die Eingabe verändert, um vorherzusagen, wie sich das Bild ändern wird – und genau dann neu rechnet, wenn es wirklich nötig ist. So wird das Malen von KI-Bildern schneller und besser.

Technische Zusammenfassung

1. Problemstellung

Diffusion Transformer (DiT) Modelle haben sich als State-of-the-Art-Architekturen für generative Aufgaben (z. B. Text-zu-Bild und Text-zu-Video) etabliert. Allerdings sind diese Modelle rechenintensiv, da sie viele Denoising-Schritte (Timesteps) erfordern, um hochwertige Ergebnisse zu erzielen.

Um die Inferenzgeschwindigkeit zu erhöhen, wurden Feature-Caching-Methoden entwickelt. Diese speichern die Ausgaben rechenintensiver Module (wie Attention und MLP) zu bestimmten Zeitpunkten und nutzen sie für nachfolgende Schritte, um redundante Berechnungen zu vermeiden.

• Herausforderung: Frühere Methoden nutzten die zwischengespeicherten Features direkt, was zu einer Akkumulation von Fehlern und Qualitätsverlust führte.
• Aktueller Stand: Neuere, vorhersagebasierte Ansätze (z. B. TaylorSeer, FasterCache) versuchen, fehlende Features durch temporale Extrapolation (z. B. Taylor-Reihenentwicklung) vorherzusagen, unter der Annahme, dass sich Features glatt über die Zeit entwickeln.
• Kernproblem: Die Analyse zeigt, dass die Magnitude der Änderungen in den Ausgabe-Features über die Timesteps hinweg irregulär ist. Reine zeitbasierte Extrapolation führt daher zu signifikanten Vorhersagefehlern, besonders bei großen Intervallen zwischen vollständigen Berechnungen, was die Generierungsqualität stark beeinträchtigt.

2. Methodik: Relational Feature Caching (RFC)

Die Autoren schlagen RFC vor, ein Framework, das nicht nur auf temporaler Extrapolation basiert, sondern die starke Korrelation zwischen den Eingabe- und Ausgabe-Features eines Moduls nutzt. Das Framework besteht aus zwei Hauptkomponenten:

A. Relational Feature Estimation (RFE)

RFE zielt darauf ab, die Größe der Änderungen in den Ausgabe-Features genauer zu schätzen.

• Beobachtung: Die Magnitude der Änderungen in den Ausgabe-Features ($\Delta O$) ist hochkorreliert mit der Magnitude der Änderungen in den Eingabe-Features ($\Delta I$).
• Annahme: Das Verhältnis $s_k = \|\Delta O\|_2 / \|\Delta I\|_2$ ist über die Zeit relativ konstant (lokal linearer Zusammenhang).
• Umsetzung: Anstatt die Ausgabe nur basierend auf vergangenen Ausgaben zu extrapolieren, nutzt RFE die leicht zu berechnenden Eingabe-Änderungen, um die erwartete Größe der Ausgabe-Änderung zu skalieren.
  • Formel: $\|\Delta O(t-k)\|_2 \approx s_N(t) \cdot \|\Delta I(t-k)\|_2$.
  • Dies ermöglicht eine präzisere Schätzung der Features, selbst wenn die zeitliche Dynamik unregelmäßig ist.

B. Relational Cache Scheduling (RCS)

RCS ist eine dynamische Strategie, um zu bestimmen, wann eine vollständige Berechnung (Full-Compute) durchgeführt werden muss, um die Vorhersagefehler zu kontrollieren.

• Problem: Die direkte Messung des Ausgabe-Fehlers erfordert eine vollständige Berechnung, was den Zweck des Caching untergräbt.
• Lösung: RCS nutzt den Vorhersagefehler der Eingabe-Features als effizienten Proxy für den Ausgabe-Fehler. Da Eingabe- und Ausgabe-Änderungen korreliert sind, steigt der Ausgabe-Fehler, wenn der Eingabe-Fehler groß wird.
• Mechanismus: Der kumulierte relative L1-Fehler der Eingabe-Vorhersage wird überwacht. Sobald dieser einen Schwellenwert $\tau$ überschreitet, wird eine vollständige Berechnung ausgelöst. Dies führt zu einem adaptiven, nicht-uniformen Intervall zwischen den Berechnungen (häufiger in späteren, komplexeren Phasen des Denoising-Prozesses).

3. Hauptbeiträge

1. Analyse der Feature-Dynamik: Nachweis, dass die Magnitude der Feature-Änderungen in DiTs irregulär ist, aber stark mit den Eingabe-Änderungen korreliert.
2. RFE (Relational Feature Estimation): Eine neue Vorhersagemethode, die Eingabe-Varianzen nutzt, um Ausgabe-Änderungen genauer zu modellieren als reine zeitliche Extrapolation.
3. RCS (Relational Cache Scheduling): Eine adaptive Scheduling-Strategie, die auf Eingabe-Fehlern basiert, um den Trade-off zwischen Effizienz und Qualität zu optimieren.
4. Umfassende Evaluation: Demonstration der Überlegenheit von RFC gegenüber State-of-the-Art-Methoden (FORA, TaylorSeer, ToCa) über verschiedene Modelle (DiT-XL/2, FLUX.1, HunyuanVideo) und Aufgaben (Klassen-bedingt, Text-zu-Bild, Text-zu-Video).

4. Ergebnisse

Die Experimente zeigen konsistent bessere Ergebnisse von RFC im Vergleich zu bestehenden Methoden:

• Quantitative Ergebnisse:
  • Auf ImageNet (DiT-XL/2) erreicht RFC bei ähnlichen FLOPs (Rechenkosten) deutlich niedrigere sFID-Werte (bessere Bildqualität) als TaylorSeer. Beispielsweise erreicht RFC bei ca. 3,37 TFLOPs einen sFID von 3,40, während TaylorSeer bei 4,76 TFLOPs einen sFID von 3,12 hat (besserer Wert bei höherem Aufwand), aber bei gleichen Kosten (ca. 3,37 TFLOPs) ist RFC überlegen.
  • Bei Text-zu-Bild (FLUX.1 dev) und Text-zu-Video (HunyuanVideo) zeigt RFC signifikant geringere Fehler (FID2FC, sFID2FC, LPIPS) im Vergleich zu den Baselines, was auf eine präzisere Feature-Wiederherstellung hindeutet.
• Qualitative Ergebnisse:
  • Generierte Bilder und Videos behalten Details besser bei (z. B. Textur von Ziegelsteinen, Struktur von Objekten) und weisen weniger Artefakte auf als Methoden, die nur auf temporaler Extrapolation basieren.
• Effizienz:
  • Der Overhead von RFE ist minimal, da Eingabe-Features nur leichte Operationen (LayerNorm, Scaling) erfordern.
  • RCS ermöglicht eine flexible Anpassung der Rechenlast, ohne die Qualität stark zu beeinträchtigen.

5. Bedeutung und Fazit

Das Paper adressiert eine fundamentale Schwäche bestehender Caching-Methoden für Diffusion Transformer: die Annahme einer glatten zeitlichen Entwicklung von Features. Durch die Einführung der Beziehung zwischen Eingabe und Ausgabe als Leitprinzip für die Vorhersage und das Scheduling, überwindet RFC die Grenzen der reinen temporalen Extrapolation.

• Innovation: Der Ansatz verschiebt den Fokus von rein zeitbasierten Modellen hin zu relationalen Modellen innerhalb des Netzwerks.
• Praktischer Nutzen: RFC ermöglicht eine signifikante Beschleunigung der Inferenz bei DiTs, ohne die Generierungsqualität zu opfern, was für den praktischen Einsatz von großen Diffusionsmodellen entscheidend ist.
• Generalität: Die Methode ist nicht auf DiTs beschränkt, sondern wurde auch erfolgreich auf U-Net-Architekturen angewendet.

Zusammenfassend stellt RFC einen robusten und effizienten Rahmen dar, der durch die Nutzung interner Netzwerkkorrelationen die Grenzen der Feature-Caching-Technologie für Diffusionsmodelle erweitert.