Abstract

Bei Diffusion-Transformer-(DiT-)Modellen, insbesondere im Bereich der Videogenerierung, stellt die Aufmerksamkeit (Attention) aufgrund der langen Sequenzlänge und der quadratischen Komplexität eine zentrale Leistungsbremse dar. Wir beobachten, dass sich die Aufmerksamkeitsgewichte in zwei Teile aufteilen lassen: einen kleinen Anteil großer Gewichte mit hoher Rangordnung und den restlichen Anteil mit sehr niedriger Rangordnung. Dies legt nahe, für den ersten Teil eine spärliche Beschleunigung und für den zweiten Teil eine niedrigrangige Beschleunigung einzusetzen. Auf dieser Erkenntnis aufbauend stellen wir SLA (Sparse-Linear Attention) vor, eine lernbare Aufmerksamkeitsmethode, die spärliche und lineare Aufmerksamkeit kombiniert, um die Geschwindigkeit von Diffusionsmodellen zu erhöhen. SLA klassifiziert die Aufmerksamkeitsgewichte in drei Kategorien – kritisch, marginal und vernachlässigbar – und wendet auf kritische Gewichte eine O(N²)-Aufmerksamkeit, auf marginale Gewichte eine O(N)-Aufmerksamkeit und überspringt vernachlässigbare Gewichte. SLA integriert diese Berechnungen in einen einzigen GPU-Kernel und unterstützt sowohl den Vorwärts- als auch den Rückwärtsdurchlauf. Durch lediglich wenige Feinabstimmungsschritte mit SLA erreichen DiT-Modelle eine 20-fache Reduktion der Aufmerksamkeitsberechnung, was zu einer erheblichen Beschleunigung führt, ohne die Qualität der Generierung zu beeinträchtigen. Experimente zeigen, dass SLA die Aufmerksamkeitsberechnung um 95 % reduziert, ohne die End-to-End-Generationsqualität zu verschlechtern, und damit die Baseline-Methoden übertrifft. Zudem implementieren wir einen effizienten GPU-Kernel für SLA, der eine 13,7-fache Beschleunigung der Aufmerksamkeitsberechnung und eine 2,2-fache Beschleunigung der End-to-End-Videogenerierung auf dem Wan2.1-1.3B-Modell ermöglicht.