Single-Stream DiT
Jak działa
1) Wejścia z różnych modalności są tokenizowane: obraz jest kodowany do latentów (zwykle przez VAE) i dzielony na łaty, a warunkowanie tekstowe pochodzi z enkodera tekstu. 2) Tokeny tekstu i obrazu są konkatenowane w jedną sekwencję. 3) Sekwencja przechodzi przez stos identycznych bloków transformera o współdzielonych wagach; w każdym bloku stosowana jest pełna self-attention po całej wspólnej sekwencji, dzięki czemu tokeny obrazu i tekstu oddziałują bezpośrednio. 4) Warunkowanie globalne (krok czasowy dyfuzji, ewentualnie wektor klasy/pooled text) jest wstrzykiwane przez modulację adaLN (skala/przesunięcie/bramkowanie) wspólną dla wszystkich tokenów. 5) Wyjście bloków służy do predykcji szumu/prędkości w procesie odszumiania dyfuzyjnego, wykonywanym iteracyjnie przez wiele kroków.
Rozwiązany problem
Dwustrumieniowe architektury (MMDiT) utrzymują osobne wagi dla każdej modalności, co zwiększa liczbę parametrów i komplikuje projekt bloku. Single-Stream DiT upraszcza architekturę: jeden komplet wag przetwarza wszystkie tokeny, co redukuje liczbę parametrów na blok, ujednolica przepływ informacji i umożliwia pełną, dwukierunkową interakcję między tokenami wszystkich modalności w ramach jednej operacji self-attention.
Komponenty
Konkatenacja tokenów warunkowania tekstowego i łat latentnych obrazu w jedną sekwencję, która jest wspólnym wejściem dla całego stosu bloków.
Stos identycznych bloków (self-attention + MLP) o jednym komplecie wag, stosowanym do wszystkich tokenów niezależnie od modalności — kluczowa różnica wobec dwustrumieniowego MMDiT.
Jedna operacja uwagi obejmuje tokeny tekstu i obrazu jednocześnie, umożliwiając dwukierunkową interakcję bez osobnej warstwy cross-attention.
Oficjalna
Adaptacyjna normalizacja warstwowa generująca parametry skali, przesunięcia i bramkowania z wektora warunkującego (krok czasu dyfuzji, ewentualnie pooled text/klasa), wspólnie dla sekwencji.
Oficjalna
Implementacja
Tekst i obraz o różnej strukturze dzielą jedną sekwencję, więc schemat kodowania pozycji (np. RoPE, osobne przesunięcia dla modalności) musi poprawnie rozróżniać i lokalizować tokeny.
Dołączenie tokenów tekstu do sekwencji obrazu wydłuża n i zwiększa koszt O(n²) self-attention — istotne przy długich promptach i wysokiej rozdzielczości.
Jeden komplet wag dla obu modalności może uprzywilejować dominującą modalność bez odpowiedniego warunkowania i strojenia, w przeciwieństwie do dedykowanych ścieżek w MMDiT.
Ewolucja
Peebles i Xie zastępują U-Net transformerem operującym na łatach latentnych — pojedyncza sekwencja tokenów obrazu z warunkowaniem przez adaLN (pierwotna, jednostrumieniowa postać DiT).
Wprowadzenie osobnych wag dla modalności tekstu i obrazu z połączoną uwagą — architektura dwustrumieniowa, wobec której single-stream stanowi prostszą alternatywę.
FLUX łączy bloki dwustrumieniowe (DoubleStreamBlock) z blokami jednostrumieniowymi (SingleStreamBlock), w których wszystkie tokeny przepływają przez wspólne warstwy uwagi i MLP.
Unified Next-DiT traktuje tokeny tekstu i obrazu jako wspólną sekwencję, realizując w pełni jednostrumieniowe przetwarzanie multimodalne.
Hiperparametry (konfigurowalne osie)
Liczba współdzielonych bloków transformera w stosie.
Wymiar reprezentacji tokenów; główny czynnik skali modelu.
Liczba głowic w self-attention po wspólnej sekwencji.
Rozmiar łaty przy patchifikacji latentów — wpływa na liczbę tokenów obrazu, a więc na długość sekwencji i koszt uwagi.
Złożoność obliczeniowa
Złożoność czasowa: O(n² · d). Złożoność przestrzenna: O(n² + n · d).
Wąskie gardło obliczeniowe
Ponieważ tekst i obraz dzielą jedną sekwencję, długość n jest sumą tokenów obu modalności, co zwiększa kwadratowy koszt uwagi względem przetwarzania samych łat obrazu.
Paradygmat wykonania
Gęsty transformer: wszystkie wagi aktywne dla wszystkich tokenów, bez routingu per modalność.
Równoległość
Przetwarzanie tokenów w obrębie bloku jest w pełni równoległe, ale próbkowanie dyfuzyjne w inferencji wymaga sekwencyjnych kroków odszumiania.
Wymagania sprzętowe
Gęste mnożenia macierzy w self-attention i MLP wspólnej sekwencji efektywnie wykorzystują rdzenie tensorowe GPU.