Diffusion Transformer (DiT)

Ein Diffusion Transformer (DiT) ist ein Diffusionsmodell, dessen Entrauschungsnetz ein Transformer statt eines faltungsbasierten U-Net ist. Er teilt ein verrauschtes Bildlatent in Patches auf, verarbeitet diese als Tokenfolge mit Selbstaufmerksamkeit – so wie ein Vision Transformer (ViT) ein Bild verarbeitet – und sagt bei jedem Entrauschungsschritt das zu entfernende Rauschen voraus. William Peebles und Saining Xie stellten das Modell in Scalable Diffusion Models with Transformers (ICCV 2023) vor. Dort erreichte DiT-XL/2 auf dem klassenkonditionalen ImageNet mit 256×256 Pixeln einen Fréchet-Inception-Abstand (FID) von 2,27 und übertraf damit den bisherigen Bestwert von 3,60 aus der latenten Diffusion. Die Arbeit zeigte, dass sich die Bildqualität mit zunehmendem Rechenaufwand eines DiT stetig verbessert. Seitdem haben Transformer-Entrauschungsnetze U-Nets in den meisten großen Bild- und Videogeneratoren ersetzt, darunter Stable Diffusion 3, FLUX.1 und Sora von OpenAI.

So funktioniert ein Diffusion Transformer#

Ein DiT arbeitet im latenten Raum des vortrainierten variationalen Autoencoders (VAE), den Stable Diffusion verwendet. Dessen Encoder verkleinert die Abmessungen um den Faktor 8, sodass ein Bild mit 256×256×3 Pixeln zu einem Latent mit 32×32×4 Werten wird. Eine Patchify-Schicht teilt das verrauschte Latent in Patches der Größe p×p und bildet jeden Patch linear auf ein Token ab. Patchgrößen von 2, 4 und 8 ergeben 256, 64 beziehungsweise 16 Token. Sinus-Cosinus-Positions-Embeddings kennzeichnen die ursprüngliche Position jedes Tokens. Eine Reihe von Transformer-Blöcken verarbeitet die Sequenz, und eine abschließende lineare Schicht wandelt jedes Token wieder in einen Patch aus vorhergesagtem Rauschen und Varianz um. Diese Patches werden anschließend zur Form des Latents zusammengesetzt. Jeder Block wird durch den Diffusionszeitschritt und die Klassenbezeichnung – oder bei einem Text-zu-Bild-Modell durch den Textprompt – konditioniert. Bei der Generierung wird diese Vorhersage ausgehend von reinem Rauschen bis zum letzten Schritt wiederholt. Der VAE-Decoder wandelt das bereinigte Latent in Pixel um:

Die Arbeit verglich vier Möglichkeiten, die Bedingung in jeden Block einzuspeisen: zusätzliche Token in der Sequenz, Kreuzaufmerksamkeit, adaptive Layer-Norm (adaLN) und adaLN-Zero. adaLN-Zero erreichte den niedrigsten FID und war zugleich am rechen effizientesten. Ein kleines Netz verarbeitet die Summe der Zeitschritt- und Klassen-Embeddings und schätzt eine Verschiebung und Skalierung, die nach jeder Layer-Norm angewendet werden, sowie ein Gate, das jeden Residualzweig skaliert. Die Gates beginnen bei null, sodass jeder Block das Training zunächst als Identitätsfunktion startet. Dieser PyTorch-Code teilt ein Latent in Patches auf und wendet den Aufmerksamkeitszweig eines Blocks an:

import torch
from torch import nn

torch.manual_seed(0)
# A 256x256 image becomes a 4x32x32 latent; 2x2 patches turn it into 256 tokens of width 384 (DiT-S/2)
latent = torch.randn(1, 4, 32, 32)
tokens = nn.Conv2d(4, 384, kernel_size=2, stride=2)(latent).flatten(2).transpose(1, 2)

# adaLN-Zero: the timestep and class embedding regresses a shift, scale, and gate for the attention branch
condition = torch.randn(1, 384)
modulation = nn.Linear(384, 3 * 384)
nn.init.zeros_(modulation.weight)
nn.init.zeros_(modulation.bias)
shift, scale, gate = modulation(condition).unsqueeze(1).chunk(3, dim=-1)

# Self-attention branch of a DiT block, modulated by the condition and scaled by the gate
attention = nn.MultiheadAttention(384, num_heads=6, batch_first=True)
h = nn.functional.layer_norm(tokens, (384,)) * (1 + scale) + shift
out = tokens + gate * attention(h, h, h)[0]
print(f"tokens {tuple(tokens.shape)}, block is identity at initialization: {torch.equal(out, tokens)}")

Die Ausgabe lautet tokens (1, 256, 384), block is identity at initialization: True. Die 256 Token mit einer Breite von 384 entsprechen dem kleinsten Modell, DiT-S/2. Das Gate mit dem Wert null leitet die Eingabe unverändert weiter, bis das Training seinen Wert verändert. In einem vollständigen DiT-Block werden für den MLP-Zweig eine zweite Verschiebung, Skalierung und ein zweites Gate geschätzt.

Skalierung mit Token und Rechenaufwand#

Die Arbeit trainierte vier Modellgrößen, von DiT-S (12 Schichten, Breite 384) bis DiT-XL (28 Schichten, Breite 1.152), jeweils mit Patchgrößen von 8, 4 und 2. Bei allen 12 Modellen entsprach der FID dem Rechenaufwand des Vorwärtsdurchlaufs in Gflops, unabhängig davon, ob dieser durch ein tieferes oder breiteres Netz oder durch mehr Token entstand. Eine Halbierung der Patchgröße vervierfacht die Anzahl der Token und den Rechenaufwand, fügt aber fast keine Parameter hinzu. DiT-XL/2 hat 675 Millionen Parameter, den VAE mit 84 Millionen Parametern nicht mitgerechnet, und benötigt bei 256×256 Pixeln 119 Gflops pro Vorwärtsdurchlauf. Bei 512×512 Pixeln verarbeitet das Modell 1.024 Token aus einem Latent mit 64×64×4 Werten bei 525 Gflops und erreicht einen FID von 3,04. Da Selbstaufmerksamkeit jedes Token mit jedem anderen vergleicht, steigen ihre Kosten bei höheren Auflösungen und längeren Videos quadratisch. Ultralytics YOLO26 vermeidet diesen Aufwand, indem es Aufmerksamkeit nur auf der Stufe mit Maßstab 1/32 einsetzt, die bei einem Bild mit 640×640 Pixeln 400 Positionen umfasst, und überall sonst Faltungen für die Echtzeiterkennung verwendet.

DiT im Vergleich zu U-Net und ViT#

• U-Net-Entrauschungsnetze: Die U-Nets in DDPM sowie Stable Diffusion 1.x und 2.x sind faltungsbasierte Encoder-Decoder, die auf mehreren Auflösungsstufen arbeiten. Skip-Verbindungen verknüpfen übereinstimmende Stufen, und auf niedrigeren Auflösungen kommt zusätzlich Aufmerksamkeit zum Einsatz. Ein DiT behält eine einzige Tokenauflösung bei, nutzt in jedem Block globale Aufmerksamkeit und hat weniger bildspezifische Designentscheidungen. Daher lässt er sich ähnlich wie Transformer in großen Sprachmodellen skalieren.
• Vision Transformer: ViT kodiert ein Bild für Aufgaben wie die Klassifizierung. DiT übernimmt dessen Patch-Embeddings und Transformer-Blöcke, ergänzt eine Konditionierung des Zeitschritts mit adaLN-Zero und gibt für jeden Patch eine Vorhersage aus.
• Latente Diffusion und Flow Matching: Ein latentes Diffusionsmodell bezeichnet den allgemeinen Ansatz, im latenten Raum eines VAE zu entrauschen. Diffusion oder Flow Matching ist das Trainingsziel. DiT ist das Entrauschungsnetz innerhalb dieses Ansatzes. Aktuelle Generatoren kombinieren DiT-Grundmodelle mit Flow-Matching-Zielen.

DiT für Text-zu-Bild- und Videogenerierung#

• Multimodaler DiT (MMDiT): Stable Diffusion 3 verarbeitet Bild- und Text-Token mit separaten Gewichten für jede Modalität und führt sie in einer einzigen Aufmerksamkeitsoperation zusammen. So können Informationen in beide Richtungen zwischen Prompt und Bild fließen. Dadurch verbesserten sich Typografie und die Umsetzung von Prompts. Stable Diffusion 3.5 Large skaliert den Ansatz auf rund 8 Milliarden Parameter.
• FLUX.1: FLUX.1 von Black Forest Labs ist ein rectified-flow-Transformer mit 12 Milliarden Parametern für die Text-zu-Bild-Generierung.
• Video: Der technische Bericht zu Sora von OpenAI beschreibt einen Diffusion Transformer, der Videos in einen latenten Raum komprimiert und in Raum-Zeit-Patches aufteilt. So werden Clips mit unterschiedlicher Dauer, Auflösung und unterschiedlichen Seitenverhältnissen zu Tokenfolgen verschiedener Länge. Für die Videogenerierung ist das größte Movie-Gen-Modell von Meta ein Transformer mit 30 Milliarden Parametern. Sein Kontext von 73.000 Video-Token entspricht 16 Sekunden Video mit 16 Bildern pro Sekunde.