Het artikel beschrijft de overstap van Linum v2 (een Latent Diffusion Model) naar de nieuwe JiT-DDT (Decoupled Diffusion Transformer) architectuur om de trainingsefficiëntie van text-to-image modellen te verhogen.
Kernpunten van de innovatie:
- Van Latent naar Pixel-space: Waar traditionele LDM's vertrouwen op een Variational Autoencoder (VAE) voor compressie, elimineert JiT-DDT de VAE. Dit voorkomt de 'compressiemuur' en beperkingen in reconstructie.
- x-prediction vs v-prediction: Door over te stappen op x-prediction wordt het probleem van de 'vloek van dimensionaliteit' omzeild, waardoor agressieve patchificatie mogelijk wordt zonder instabiliteit.
- Decoupled Architectuur (DDT): Om het verlies van fijne details in pixel-ruimte tegen te gaan, is het model gesplitst in een Conditionele Encoder (voor laag-frequente structuur) en een Velocity Decoder (voor hoog-frequente details).
- Architecturale verbeteringen: Er is gebruikgemaakt van een single-stream DiT, SwiGLU activaties, de Muon optimizer en een aangepast noise schedule om details zoals haartextuur en belichting te verbeteren.
Resultaten: De JiT-DDT traint 3,6× sneller (in GPU-uren) dan de baseline, terwijl de outputresolutie is gestegen van 256×256 naar 512×512 pixels. Hoewel het model superieur is in artistieke stijlen en belichting, blijft de weergave van kleine menselijke gezichten een verbeterpunt.
Training van Text-to-Image Modellen 3,6× Sneller
Samenvatting (TL;DR)
Linum v2 werd beperkt door de enorme omvang van het attention context window. Een clip van 5 seconden in 720p kostte maar liefst 110K tokens. Ter vergelijking: LLM's zien tijdens 97% van hun pretraining samples met minder dan 8K tokens. Omdat de kosten van attention kwadratisch stijgen, is het inkorten van het contextvenster de belangrijkste hefboom om de training van modellen te versnellen.
De meeste generatieve beeld- en videosystemen zijn Latent Diffusion Models (LDM's). Zij splitsen compressie en generatie in onafhankelijk getrainde modules: de Variational Autoencoder (VAE) en de Diffusion Transformer (DiT). Recentelijk zijn pixel-space modellen, zoals de JiT, een veelbelovend alternatief gebleken. Deze reduceren twee modellen tot één en stellen het diffusiemodel in staat om een latente ruimte specifiek voor generatie te construeren, in plaats van te vertrouwen op een ruimte die is gebouwd voor reconstructie.
Bij training op onze (afbeelding, bijschrift) dataset bleek de JiT moeite te hebben met het produceren van fijne details. Wij stellen een nieuwe encoder-decoder architectuur voor (JiT-DDT) die deze details herstelt en veel efficiënter traint dan zijn LDM-tegenhanger. Vergeleken met onze Linum v2-baseline traint de JiT-DDT een text-to-image model met 3,6× minder GPU-uren, ondanks dat het afbeeldingen genereert met 4× zoveel pixels.
Vergelijking: Snelheid en Resolutie
| Kenmerk | Linum v2 (Vorige) | JiT-DDT (Nieuw) |
|---|
| Resolutie | 256×256 | 512×512 |
| Architectuur | 2.0B latent-space DiT + VAE | 2.5B active pixel-space DiT |
| Tokens | 256 latente tokens | 320 pixel tokens (64 enc + 256 dec) |
| GPU-uren | Basislijn | 3,6× minder |
| Gezien samples | Basislijn | 4,2× minder |
---
De VAE-compressiemuur
Bijna alle generatieve beeld- en videomodellen zijn Latent Diffusion Models (LDM's). Deze bestaan uit twee hoofdcomponenten: een Variational Auto Encoder (VAE) voor compressie en een Diffusion Transformer (DiT) voor generatie.
Werken met ruwe pixels is te duur (vooral voor video). Daarom moeten RGB-pixels worden gereduceerd tot een kleiner aantal tokens voor de DiT. De VAE is hiervoor getraind: hij stuurt pixel-samples door een probabilistische encoder naar tokens in een latente ruimte, waarna een probabilistische decoder deze weer terugbrengt naar de pixel-ruimte.
Bij het bouwen van een LDM train je de VAE apart en bevries je deze vervolgens (er is geen gradiëntstroom van de DiT naar de VAE). Hierdoor blijft de latente ruimte statisch tijdens de training van de DiT.
We willen zoveel mogelijk token-compressie uit de VAE halen om de kosten van attention in onze DiT te beperken. Echter, populaire open-source modellen zoals FLUX, Ideogram en Z-Image beperken zich tot een reductie van 16×16 tokens. Er lijkt een empirisch plafond te liggen voor de hoeveelheid compressie die we uit een standaard CNN VAE kunnen halen zonder dat de reconstructie verslechtert.
Aggressieve compressie via een unified model
Recentelijk is in een paper (JiT) aangetoond dat een token-reductie van 32×32 mogelijk is door de VAE volledig weg te laten en de compressietaak in de DiT zelf te plaatsen via patchification.
Het proces van patchificatie:
- 4×4-pixel patches worden omgezet in 48-dimensionele tokens.
- Een lineaire bottleneck reduceert dit naar 12 dimensies.
(Bij 512px gebruikt JiT 32×32 patches, waardoor een 512×512 afbeelding 256 tokens wordt).
Waarom dit voorheen mislukte: De vloek van dimensionaliteit
DiT's hebben moeite om te leren van hoog-dimensionale inputs. Aggressieve patchificatie duwt informatie in de kanaaldimensie, wat instabiliteit veroorzaakt. Dit komt echter niet door de architectuur, maar door het v-prediction, v-loss flow matching doel dat de afgelopen jaren standaard is gebruikt.
Flow matching herinnering: In flow matching construeren we een recht pad tussen een sample uit de data-distributie ($x$) en een sample van Gaussische ruis ($\epsilon$). We trainen een netwerk om de snelheid (velocity) langs dat pad te benaderen. Dit noemen we v-prediction.
Het probleem met v-prediction: Bij v-prediction moet het neurale netwerk impliciet zowel het signaal ($x$) als de ruis ($\epsilon$) leren. Ruis is een willekeurige Gaussische variabele die de volledige hoog-dimensionale ruimte vult. Naarmate de dimensie ($D$) toeneemt, wordt het fitten van deze ruis exponentieel moeilijker.
Door over te stappen op x-prediction, omzeilen we dit probleem. Omdat afbeeldingen en video's intrinsiek laag-dimensionaal zijn (ze liggen op een dunne 'manifold' in de hoog-dimensionale ruimte), kan het model effectief $x$ voorspellen, zelfs als $D$ groot is.
Conclusie: Door x-prediction, v-loss te gebruiken, kunnen we agressieve patchificatie toepassen, de kanaaldimensie vergroten en het compressieprobleem volledig in de DiT onderbrengen.
---
JiT uitbreiden voor text-to-image modellen
Hoewel JiT veelbelovend is, worstelt het met fijne details. Wij geloven dat een simpelere, coherente architectuur beter werkt dan twee losse modellen.
De JiT-architectuur configuratie
Om van LDM naar pixel-space te gaan, zijn we overgestapt op x-prediction, v-loss en hebben we de volgende aanpassingen doorgevoerd:
- Single-stream DiT: In plaats van afwisselende blokken van self-attention (beeld) en cross-attention (tekst), concateneren we visuele en tekstuele tokens in één stroom. Dit verhoogt de expressiviteit van de relatie tussen tekst en beeld.
- Breder in plaats van dieper: We zijn overgestapt van een 40-laags transformer (2048 hidden size) naar een 23-laags transformer met een bredere 2944 hidden size.
- Perceptuele verliezen: Omdat we geen VAE-decoder hebben, gebruiken we LPIPS en P-DINO loss om de resultaten menselijker te maken (geen wazige beelden).
- SwiGLU: Standaard SiLU activaties zijn vervangen door gated SiLUs (SwiGLU).
- Muon optimizer: 2D-matrices worden geoptimaliseerd met Muon, terwijl input/output lagen en biases AdamW gebruiken.
- PixelREPA: Een auxiliary loss die de verborgen toestand van laag 8 aligneert met DINOv3-representaties van het schone beeld, waarbij 20% van de tokens wordt gemaskeerd.
- Sigmoid attention gating: Om attention sinks te voorkomen in de single-stream architectuur.
- Qwen tekst-embeddings: We gebruiken hidden states van Qwen3.5-4B in plaats van T5-XXL.
Het JiT-recept: Loss en Schedule
- Noise Schedule: We scheven de timesteps ($t$) naar waarden waar de structuur van het beeld wordt bepaald. De $\sigma$-term is gekoppeld aan de pixelcount ($\sigma=1$ voor 256px, $\sigma=2$ voor 512px).
- Loss: MSE (x-prediction met velocity weighting), gecombineerd met LPIPS, P-DINO en PixelREPA.
---
Herstellen van fijne details in pixel-ruimte
JiT-modellen hebben vaak een "airbrushed" of oververzadigd uiterlijk. Om dit op te lossen, introduceren we de JiT-DDT (Decoupled Diffusion Transformer).
DDT: Decoupled Diffusion Transformer
Het probleem bij standaard DiT's is dat dezelfde modules zowel de laag-frequente semantiek (structuur) moeten coderen als de hoog-frequente details moeten decoderen. Dit creëert een optimalisatie-dilemma.
DDT lost dit op door het model te splitsen:
- Conditionele Encoder: Focus op laag-frequente structuur.
- Velocity Decoder: Focus op hoog-frequente details.
De JiT-DDT Architectuur
Wij passen dit principe toe in pixel-ruimte:
- Encoder: Voorspelt een versie van het beeld op 64×64 resolutie (8× downsampled). Hierdoor leert de encoder expliciet de structuur.
- Decoder: Gebruikt de output van de encoder, het noised beeld op volledige resolutie en de tekstprompt om het uiteindelijke 512×512 beeld te genereren.
Belangrijke afwijkingen van originele DDT:
- Symmetrische grootte: Encoder en decoder zijn ongeveer even groot (elk ~1.1B parameters), omdat de taak van de encoder (64x64) aanzienlijk eenvoudiger is.
- Tekstconditie in decoder: We geven de decoder toegang tot de tekstprompt, omdat veel details in de tekst beschreven staan.
Optimalisatie van het Noise Schedule
We ontdekten dat architectuur alleen niet genoeg was. Door het noise schedule tijdens de training aan te passen, konden we details zoals sproeten en haartextuur herstellen:
- Fase 1: Focus op globale structuur.
- Fase 2: Verschuiving naar "schonere" beelden (lagere $t$), waardoor het model meer tijd besteedt aan het verfijnen van details.
---
Architecturale Verfijningen (Refiners)
Geïnspireerd door Z-Image hebben we Refiners toegevoegd. Dit zijn modaliteits-specifieke blokken vóór de gedeelde DiT-trunk.
- De encoder en decoder krijgen elk vier specifieke blokken (twee voor beeld, twee voor tekst).
- Dit helpt het model om de verschillende modaliteiten te "masseren" voordat ze worden gecombineerd.
- Dit loste problemen met overmatige sproeten en kleurgradatie op.
---
Van Linum v2 naar JiT-DDT: Conclusie
Onze reis laat zien dat we de training van text-to-image modellen drastisch kunnen versnellen door de VAE te elimineren en over te stappen op pixel-space training met x-prediction.
Eindoordeel:
- Snelheid: 3,6× minder GPU-uren.
- Kwaliteit: 4× meer pixels (512px vs 256px).
- Resultaat: Betere getrouwheid aan kunststijlen en realistischere belichting, hoewel menselijke gezichten die niet het hoofdonderwerp zijn, nog een verbeterpunt vormen.
---
Waarom werkt de JiT-DDT?
De rol van structuur vs. holistisch begrip
Onderzoek (zoals iREPA) suggereert dat alignment met vision-encoders (zoals DINO) werkt omdat het de vroege lagen van een netwerk leert om lokale structuur te herkennen, niet noodzakelijkerwijs holistische concepten. Onze encoder-prediction taak (het voorspellen van een 64x64 beeld) doet in feite hetzelfde: het stript alle detail weg en dwingt het model te focussen op structuur.
Gradiënt-problemen in DiT's
Het lijkt erop dat DiT's een probleem hebben met gradiënt-propagatie. Modellen zoals Self-Flow, RAE v2 en onze JiT-DDT verbeteren allemaal door een extra loss-term (een "gradiënt-snelweg") eerder in het netwerk te introduceren. Dit versnelt het leerproces aanzienlijk.
---
Appendix: Visuele Vergelijkingen (Linum v2 vs. JiT-DDT)
Uit tests op 26 verschillende prompts blijken de volgende trends:
- Artistieke stijlen: JiT-DDT is veel trouwer aan stijlen zoals houtskooltekeningen, Romeïsche mozaïeken en olieverfschilderijen.
- Belichting: De belichting is realistischer; Linum v2 had een sterke neiging tot oververzadiging.
- Zwakke punten: JiT-DDT heeft nog steeds moeite met realistische menselijke gezichten wanneer deze klein zijn in het frame (bijv. vervormde ogen of monden).
Training van Text-to-Image Modellen 3,6× Sneller
Samenvatting (TL;DR)
Linum v2 werd beperkt door de enorme omvang van het attention context window. Een clip van 5 seconden in 720p kostte maar liefst 110K tokens. Ter vergelijking: LLM's zien tijdens 97% van hun pretraining samples met minder dan 8K tokens. Omdat de kosten van attention kwadratisch stijgen, is het inkorten van het contextvenster de belangrijkste hefboom om de training van modellen te versnellen.
De meeste generatieve beeld- en videosystemen zijn Latent Diffusion Models (LDM's). Zij splitsen compressie en generatie in onafhankelijk getrainde modules: de Variational Autoencoder (VAE) en de Diffusion Transformer (DiT). Recentelijk zijn pixel-space modellen, zoals de JiT, een veelbelovend alternatief gebleken. Deze reduceren twee modellen tot één en stellen het diffusiemodel in staat om een latente ruimte specifiek voor generatie te construeren, in plaats van te vertrouwen op een ruimte die is gebouwd voor reconstructie.
Bij training op onze (afbeelding, bijschrift) dataset bleek de JiT moeite te hebben met het produceren van fijne details. Wij stellen een nieuwe encoder-decoder architectuur voor (JiT-DDT) die deze details herstelt en veel efficiënter traint dan zijn LDM-tegenhanger. Vergeleken met onze Linum v2-baseline traint de JiT-DDT een text-to-image model met 3,6× minder GPU-uren, ondanks dat het afbeeldingen genereert met 4× zoveel pixels.
Vergelijking: Snelheid en Resolutie
| Kenmerk | Linum v2 (Vorige) | JiT-DDT (Nieuw) |
|---|
| Resolutie | 256×256 | 512×512 |
| Architectuur | 2.0B latent-space DiT + VAE | 2.5B active pixel-space DiT |
| Tokens | 256 latente tokens | 320 pixel tokens (64 enc + 256 dec) |
| GPU-uren | Basislijn | 3,6× minder |
| Gezien samples | Basislijn | 4,2× minder |
---
De VAE-compressiemuur
Bijna alle generatieve beeld- en videomodellen zijn Latent Diffusion Models (LDM's). Deze bestaan uit twee hoofdcomponenten: een Variational Auto Encoder (VAE) voor compressie en een Diffusion Transformer (DiT) voor generatie.
Werken met ruwe pixels is te duur (vooral voor video). Daarom moeten RGB-pixels worden gereduceerd tot een kleiner aantal tokens voor de DiT. De VAE is hiervoor getraind: hij stuurt pixel-samples door een probabilistische encoder naar tokens in een latente ruimte, waarna een probabilistische decoder deze weer terugbrengt naar de pixel-ruimte.
Bij het bouwen van een LDM train je de VAE apart en bevries je deze vervolgens (er is geen gradiëntstroom van de DiT naar de VAE). Hierdoor blijft de latente ruimte statisch tijdens de training van de DiT.
We willen zoveel mogelijk token-compressie uit de VAE halen om de kosten van attention in onze DiT te beperken. Echter, populaire open-source modellen zoals FLUX, Ideogram en Z-Image beperken zich tot een reductie van 16×16 tokens. Er lijkt een empirisch plafond te liggen voor de hoeveelheid compressie die we uit een standaard CNN VAE kunnen halen zonder dat de reconstructie verslechtert.
Aggressieve compressie via een unified model
Recentelijk is in een paper (JiT) aangetoond dat een token-reductie van 32×32 mogelijk is door de VAE volledig weg te laten en de compressietaak in de DiT zelf te plaatsen via patchification.
Het proces van patchificatie:
- 4×4-pixel patches worden omgezet in 48-dimensionele tokens.
- Een lineaire bottleneck reduceert dit naar 12 dimensies.
(Bij 512px gebruikt JiT 32×32 patches, waardoor een 512×512 afbeelding 256 tokens wordt).
Waarom dit voorheen mislukte: De vloek van dimensionaliteit
DiT's hebben moeite om te leren van hoog-dimensionale inputs. Aggressieve patchificatie duwt informatie in de kanaaldimensie, wat instabiliteit veroorzaakt. Dit komt echter niet door de architectuur, maar door het v-prediction, v-loss flow matching doel dat de afgelopen jaren standaard is gebruikt.
Flow matching herinnering: In flow matching construeren we een recht pad tussen een sample uit de data-distributie ($x$) en een sample van Gaussische ruis ($\epsilon$). We trainen een netwerk om de snelheid (velocity) langs dat pad te benaderen. Dit noemen we v-prediction.
Het probleem met v-prediction: Bij v-prediction moet het neurale netwerk impliciet zowel het signaal ($x$) als de ruis ($\epsilon$) leren. Ruis is een willekeurige Gaussische variabele die de volledige hoog-dimensionale ruimte vult. Naarmate de dimensie ($D$) toeneemt, wordt het fitten van deze ruis exponentieel moeilijker.
Door over te stappen op x-prediction, omzeilen we dit probleem. Omdat afbeeldingen en video's intrinsiek laag-dimensionaal zijn (ze liggen op een dunne 'manifold' in de hoog-dimensionale ruimte), kan het model effectief $x$ voorspellen, zelfs als $D$ groot is.
Conclusie: Door x-prediction, v-loss te gebruiken, kunnen we agressieve patchificatie toepassen, de kanaaldimensie vergroten en het compressieprobleem volledig in de DiT onderbrengen.
---
JiT uitbreiden voor text-to-image modellen
Hoewel JiT veelbelovend is, worstelt het met fijne details. Wij geloven dat een simpelere, coherente architectuur beter werkt dan twee losse modellen.
De JiT-architectuur configuratie
Om van LDM naar pixel-space te gaan, zijn we overgestapt op x-prediction, v-loss en hebben we de volgende aanpassingen doorgevoerd:
- Single-stream DiT: In plaats van afwisselende blokken van self-attention (beeld) en cross-attention (tekst), concateneren we visuele en tekstuele tokens in één stroom. Dit verhoogt de expressiviteit van de relatie tussen tekst en beeld.
- Breder in plaats van dieper: We zijn overgestapt van een 40-laags transformer (2048 hidden size) naar een 23-laags transformer met een bredere 2944 hidden size.
- Perceptuele verliezen: Omdat we geen VAE-decoder hebben, gebruiken we LPIPS en P-DINO loss om de resultaten menselijker te maken (geen wazige beelden).
- SwiGLU: Standaard SiLU activaties zijn vervangen door gated SiLUs (SwiGLU).
- Muon optimizer: 2D-matrices worden geoptimaliseerd met Muon, terwijl input/output lagen en biases AdamW gebruiken.
- PixelREPA: Een auxiliary loss die de verborgen toestand van laag 8 aligneert met DINOv3-representaties van het schone beeld, waarbij 20% van de tokens wordt gemaskeerd.
- Sigmoid attention gating: Om attention sinks te voorkomen in de single-stream architectuur.
- Qwen tekst-embeddings: We gebruiken hidden states van Qwen3.5-4B in plaats van T5-XXL.
Het JiT-recept: Loss en Schedule
- Noise Schedule: We scheven de timesteps ($t$) naar waarden waar de structuur van het beeld wordt bepaald. De $\sigma$-term is gekoppeld aan de pixelcount ($\sigma=1$ voor 256px, $\sigma=2$ voor 512px).
- Loss: MSE (x-prediction met velocity weighting), gecombineerd met LPIPS, P-DINO en PixelREPA.
---
Herstellen van fijne details in pixel-ruimte
JiT-modellen hebben vaak een "airbrushed" of oververzadigd uiterlijk. Om dit op te lossen, introduceren we de JiT-DDT (Decoupled Diffusion Transformer).
DDT: Decoupled Diffusion Transformer
Het probleem bij standaard DiT's is dat dezelfde modules zowel de laag-frequente semantiek (structuur) moeten coderen als de hoog-frequente details moeten decoderen. Dit creëert een optimalisatie-dilemma.
DDT lost dit op door het model te splitsen:
- Conditionele Encoder: Focus op laag-frequente structuur.
- Velocity Decoder: Focus op hoog-frequente details.
De JiT-DDT Architectuur
Wij passen dit principe toe in pixel-ruimte:
- Encoder: Voorspelt een versie van het beeld op 64×64 resolutie (8× downsampled). Hierdoor leert de encoder expliciet de structuur.
- Decoder: Gebruikt de output van de encoder, het noised beeld op volledige resolutie en de tekstprompt om het uiteindelijke 512×512 beeld te genereren.
Belangrijke afwijkingen van originele DDT:
- Symmetrische grootte: Encoder en decoder zijn ongeveer even groot (elk ~1.1B parameters), omdat de taak van de encoder (64x64) aanzienlijk eenvoudiger is.
- Tekstconditie in decoder: We geven de decoder toegang tot de tekstprompt, omdat veel details in de tekst beschreven staan.
Optimalisatie van het Noise Schedule
We ontdekten dat architectuur alleen niet genoeg was. Door het noise schedule tijdens de training aan te passen, konden we details zoals sproeten en haartextuur herstellen:
- Fase 1: Focus op globale structuur.
- Fase 2: Verschuiving naar "schonere" beelden (lagere $t$), waardoor het model meer tijd besteedt aan het verfijnen van details.
---
Architecturale Verfijningen (Refiners)
Geïnspireerd door Z-Image hebben we Refiners toegevoegd. Dit zijn modaliteits-specifieke blokken vóór de gedeelde DiT-trunk.
- De encoder en decoder krijgen elk vier specifieke blokken (twee voor beeld, twee voor tekst).
- Dit helpt het model om de verschillende modaliteiten te "masseren" voordat ze worden gecombineerd.
- Dit loste problemen met overmatige sproeten en kleurgradatie op.
---
Van Linum v2 naar JiT-DDT: Conclusie
Onze reis laat zien dat we de training van text-to-image modellen drastisch kunnen versnellen door de VAE te elimineren en over te stappen op pixel-space training met x-prediction.
Eindoordeel:
- Snelheid: 3,6× minder GPU-uren.
- Kwaliteit: 4× meer pixels (512px vs 256px).
- Resultaat: Betere getrouwheid aan kunststijlen en realistischere belichting, hoewel menselijke gezichten die niet het hoofdonderwerp zijn, nog een verbeterpunt vormen.
---
Waarom werkt de JiT-DDT?
De rol van structuur vs. holistisch begrip
Onderzoek (zoals iREPA) suggereert dat alignment met vision-encoders (zoals DINO) werkt omdat het de vroege lagen van een netwerk leert om lokale structuur te herkennen, niet noodzakelijkerwijs holistische concepten. Onze encoder-prediction taak (het voorspellen van een 64x64 beeld) doet in feite hetzelfde: het stript alle detail weg en dwingt het model te focussen op structuur.
Gradiënt-problemen in DiT's
Het lijkt erop dat DiT's een probleem hebben met gradiënt-propagatie. Modellen zoals Self-Flow, RAE v2 en onze JiT-DDT verbeteren allemaal door een extra loss-term (een "gradiënt-snelweg") eerder in het netwerk te introduceren. Dit versnelt het leerproces aanzienlijk.
---
Appendix: Visuele Vergelijkingen (Linum v2 vs. JiT-DDT)
Uit tests op 26 verschillende prompts blijken de volgende trends:
- Artistieke stijlen: JiT-DDT is veel trouwer aan stijlen zoals houtskooltekeningen, Romeïsche mozaïeken en olieverfschilderijen.
- Belichting: De belichting is realistischer; Linum v2 had een sterke neiging tot oververzadiging.
- Zwakke punten: JiT-DDT heeft nog steeds moeite met realistische menselijke gezichten wanneer deze klein zijn in het frame (bijv. vervormde ogen of monden).
Groetjes,