DiT
Jak działa
Wejściem DiT jest zaszumiony latent z_t (np. 32×32×4 z VAE Stable Diffusion). (1) Patchify: latent jest dzielony na patche p×p (p=2,4,8) i każdy jest linearnie projektowany na token o wymiarze d, tworząc sekwencję T = (h/p)·(w/p) tokenów. (2) Positional embedding: dodawane są (zwykle sinusoidalne lub uczone) embeddingi pozycji 2D. (3) Conditioning: krok t i warunek c są kodowane i wstrzykiwane do każdego bloku. Najlepszy wariant adaLN-Zero parametryzuje skalowanie i przesunięcie LayerNorm (γ, β) oraz bramki rezydualne (α) jako funkcje (t,c), z α zainicjalizowanym na zero (każdy blok startuje jako identyczność). (4) Bloki Transformera: standardowe self-attention + MLP nad tokenami patchy. (5) Decode: finalna LayerNorm + linear projektuje tokeny z powrotem na patche szumu/kowariancji, które są rearanżowane w latent ε. Inferencja działa jak w każdym modelu dyfuzyjnym — iteracyjne odszumianie schedulerem, z CFG dla warunkowania. MMDiT (SD3) rozszerza to o osobne strumienie tokenów tekstu i obrazu łączone w blokach attention.
Rozwiązany problem
Konwolucyjna U-Net w modelach dyfuzyjnych ma ograniczoną skalowalność — wzrost jakości nasyca się przy zwiększaniu parametrów, a indukcyjne biasy konwolucji utrudniają wykorzystanie ogromnych budżetów obliczeniowych. DiT pokazuje, że czysty Transformer skaluje się znacznie lepiej (jakość rośnie monotonicznie z FLOPs), korzystając z tych samych praw skalowania co LLM-y.
Komponenty
Dzieli latent na patche p×p i linearnie projektuje każdy na token wymiaru d. Mniejszy p → więcej tokenów → wyższa jakość ale większy koszt.
Adaptive LayerNorm regresujący skalę/przesunięcie (γ,β) oraz bramki rezydualne (α) z embeddingu (t,c). α inicjalizowane na 0 → każdy blok startuje jako identyczność, stabilizując trening.
Oficjalna
Standardowe bloki Transformera (multi-head self-attention + feed-forward). Globalny receptive field nad wszystkimi patchami od pierwszej warstwy.
Sinusoidalne lub uczone embeddingi pozycji dodawane do tokenów patchy (Transformer sam w sobie jest permutacyjnie niezmienniczy).
LayerNorm + linear mapujący każdy token na patch przewidywanego szumu (i opcjonalnie kowariancji), rearanżowany w latent ε.
Implementacja
Zmniejszenie p z 8 do 2 zwiększa liczbę tokenów 16×, a koszt attention 256× — łatwo przekroczyć pamięć.
Naiwne warunkowanie (in-context, cross-attn) trenuje się gorzej i mniej stabilnie niż adaLN-Zero.
DiT wymaga więcej danych i obliczeń niż U-Net by nauczyć się lokalnych korelacji, których konwolucja dostarcza za darmo.
Ewolucja
Peebles & Xie pokazują, że Transformer zastępujący U-Net skaluje się monotonicznie z FLOPs i pobija U-Net na ImageNet.
Chen i in. trenują DiT text-to-image przy ułamku kosztu SD, z cross-attention dla tekstu.
OpenAI opisuje Sorę jako diffusion transformer operujący na czasoprzestrzennych patchach latentu wideo.
Stability AI wprowadza MMDiT z osobnymi strumieniami tokenów tekstu i obrazu oraz rectified flow.
Black Forest Labs wydaje Flux, czołowy otwarty model oparty na DiT/MMDiT (12B).
Hiperparametry (konfigurowalne osie)
Rozmiar patcha (2, 4, 8). Mniejszy p → więcej tokenów → wyższa jakość, kwadratowo większy koszt.
Konfiguracje DiT-S/B/L/XL analogiczne do ViT. Skalowanie poprawia FID monotonicznie.
adaLN-Zero vs in-context vs cross-attention. adaLN-Zero osiągnął najlepsze wyniki.
Wynika z rozdzielczości latentu i patch size. Koszt attention rośnie O(T²).
Złożoność obliczeniowa
Złożoność czasowa: O(T² · d) per krok odszumiania, T = liczba patchy, d = wymiar modelu.
Paradygmat wykonania
Cały Transformer jest aktywny na każdy krok odszumiania (z wyjątkiem wariantów MoE-DiT).
Równoległość
Trening i pojedynczy krok inferencji są w pełni równoległe po tokenach (jak w każdym Transformerze). Sekwencyjne pozostają tylko kroki odszumiania dyfuzji — wspólne dla wszystkich modeli dyfuzyjnych.
Wymagania sprzętowe
Czysty Transformer to matmul-heavy obciążenie idealne dla tensor cores; korzysta z flash-attention i optymalizacji LLM-owych.
Transformer dobrze paralelizuje się na TPU (XLA); duże treningi DiT/MMDiT korzystają z TPU/wielu GPU.