DiG: Skalierbare und effiziente Diffusionsmodelle mit Gated Linear Attention

Diffusionsmodelle mit umfangreichem Pre-Training haben bedeutende Erfolge im Bereich der visuellen Inhaltserstellung erzielt, insbesondere verkörpert durch Diffusion Transformers (DiT). Allerdings standen DiT-Modelle vor Herausforderungen hinsichtlich der Skalierbarkeit und der quadratischen Komplexitätseffizienz. In diesem Papier zielen wir darauf ab, die Fähigkeit zur Modellierung langer Sequenzen der Gated Linear Attention (GLA) Transformers zu nutzen und deren Anwendbarkeit auf Diffusionsmodelle auszudehnen. Wir führen die Diffusion Gated Linear Attention Transformers (DiG) ein, eine einfache, anpassbare Lösung mit minimalem Parameteraufwand, die dem DiT-Design folgt, aber überlegene Effizienz und Effektivität bietet. Neben einer besseren Leistung als DiT weist DiG-S/2 eine um 2,5 höhere Trainingsgeschwindigkeit als DiT-S/2 auf und spart 75,7\% GPU-Speicher bei einer Auflösung von 1792 x 1792. Darüber hinaus analysieren wir die Skalierbarkeit von DiG über eine Vielzahl von rechnerischer Komplexität. DiG-Modelle zeigen mit erhöhtem Tiefen-/Breitengrad oder einer Erweiterung der Eingabetoken konstant abnehmende FID. Wir vergleichen DiG weiter mit anderen subquadratischen Diffusionsmodellen. Bei gleicher Modellgröße ist DiG-XL/2 4,2-mal schneller als das neueste Mamba-basierte Diffusionsmodell bei einer Auflösung von 1024 und 1,8-mal schneller als DiT mit CUDA-optimierter FlashAttention-2 bei der Auflösung von 2048. All diese Ergebnisse zeigen seine überlegene Effizienz unter den neuesten Diffusionsmodellen. Der Code wurde unter https://github.com/hustvl/DiG veröffentlicht.