JLT: Clean-Latent Prediction in Latent Diffusion Transformers

Technische Samenvatting: JLT: Schone-Latente Voorspelling in Latente Diffusietransformators

Probleemstelling

Denoising-diffusiemodellen keren traditioneel een corruptieproces om, maar het specifieke doel dat het neurale netwerk moet voorspellen, verschilt per architectuur. Veelvoorkomende doelen zijn ruis ($\epsilon$), snelheid ($v = x - \epsilon$) of schone data ($x$). Hoewel deze doelen algebraïsch omkeerbaar zijn voor een vaste corruptietijd, suggereert recent onderzoek (bijv. JiT) dat onder eindige netwerkkapaciteit en eindige data de keuze van het voorspeldoel de moeilijkheidsgraad van het regressieprobleem aanzienlijk beïnvloedt. Specifiek stelt het voorspellen van schone data het model in staat zich te richten op gestructureerde, laagdimensionale variaties op het manifold, in plaats van omgevingsruis buiten het manifold te reconstrueren.

De centrale vraag die dit artikel adresseert, is of dit voordeel van "schone-voorspelling" behouden blijft wanneer het model opereert in een gecomprimeerde latente ruimte in plaats van ruwe pixelruimte. Bij latente diffusie heeft een Variational Autoencoder (VAE) al veel van de variabiliteit in de ruwe pixels verwijderd. De auteurs onderzoeken of de geometrische eigenschappen van het voorspeldoel ($x$ versus $v$) een kritieke factor blijven, zelfs wanneer de representatieruimte vaststaat en gecomprimeerd is.

Methodologie

De auteurs introduceren JLT (Joint Latent Transformer), een gecontroleerd experimenteel raamwerk ontworpen om het effect van het voorspeldoel binnen een vaste latente ruimte te isoleren.

1. Architectuur en Opzet:

   • Ruggengraat: Een Base-scale Transformer (130M parameters) met 12 blokken, een verborgen dimensie van 768 en 12 attention-heads, wat de configuratie van JiT-B/16 weerspiegelt.
   • Representatie: Het model werkt op bevroren FLUX.2 VAE latente codes. Dit zorgt ervoor dat de representatieruimte identiek is voor alle vergeleken varianten.
   • Training: Modellen worden getraind voor 250K stappen (200 epochs) op ImageNet 256 $\times$ 256.
   • Gecontroleerde Variabelen: Het onderzoek houdt strikt constant de representatie (FLUX.2 VAE), de schaal van de architectuur, de optimizer-instellingen, de batchgrootte en de sampling-protocollen.

2. Experimentele Vergelijking:
   De auteurs vergelijken twee primaire varianten onder dezelfde instellingen:

   • JLT (Schone-Latente Voorspelling): Het netwerk wordt direct gesuperviseerd om de schone latente vector $x$ te voorspellen.
   • DiT (Snelheidsvoorspelling): Een匹配的 baseline waarbij het netwerk gesuperviseerd wordt om de snelheidsvector $v = x - \epsilon$ te voorspellen.
   • Varianten: Experimenten worden uitgevoerd bij twee VAE-grid patchgroottes, aangeduid als /1 en /2 (bijv. JLT-B/1 versus DiT-B/1).

3. Theoretische Analyse:
   Het artikel biedt een lokale Gaussische analyse om de empirische kloof te verklaren. Uitgaande van een lokale lineaire-Gaussische benadering van de latente data-verdeling ($x \sim \mathcal{N}(0, \Sigma)$) en onafhankelijke ruis ($\epsilon \sim \mathcal{N}(0, I)$):

   • Covariantiestructuur: De covariantie van het snelheidsdoel is $Cov(y_v) = \Sigma + I$, terwijl die van het schone doel $Cov(y_x) = \Sigma$ is. Dit impliceert dat snelheidsvoorspelling een isotrope eenheidsvloer toevoegt aan elke richting. In richtingen waar de schone data een lage variantie heeft (zwak ondersteund door het manifold), dwingt snelheidsvoorspelling het model om een doel met eenheidsvariantie te leren, wat effectief ruis in die richtingen versterkt.
   • Conditionele Ambiguïteit: De analyse toont aan dat de Bayes-residuvariantie voor snelheidsvoorspelling wordt geschaald met een factor $1/(1-t)^2$ ten opzichte van schone voorspelling. Naarmate de eigenwaarde $\lambda_i$ van de data-covariantie naar nul nadert, neigt de coëfficiënt voor schone voorspelling naar 0 (dempend voor de richting), terwijl de snelheidscoëfficiënt neigt naar $-1/(1-t)$ (versterkend voor de richting).

Belangrijkste Resultaten

De experimenten op ImageNet 256 $\times$ 256 tonen een consistente en significante prestatiekloof aan die gunstig is voor schone-latente voorspelling:

• Kwantitatieve Prestaties:
  • JLT-B/1 (Schoon) bereikt een FID-50K van 2.50 (met classifier-free guidance), vergeleken met 6.56 voor de匹配的 DiT-B/1 (Snelheid).
  • JLT-B/2 (Schoon) bereikt een FID-50K van 14.81, vergeleken met 28.71 voor DiT-B/2.
  • De Inception Score (IS) toont eveneens een aanzienlijke verbetering (bijv. 232.51 versus 132.12 voor de /1-varianten).
• Trainingsdynamiek: Trainingscurven geven aan dat schone-latente modellen eerder het regime met lage FID betreden (rond 100K stappen) en tijdens de hele training een duidelijke marge behouden ten opzichte van snelheidsmodellen.
• Robuustheid: Het voordeel geldt voor verschillende tokenisatiegranulariteiten (/1 en /2), wat suggereert dat het voordeel geen artefact is van een specifieke patchgrootte.

Betekenis en Claims

Het artikel stelt dat de keuze van het voorspeldoel in latente diffusie een geometrische modelkeuze is en niet slechts een algebraïsche parametrisering.

1. Doelgeometrie boven Compressie: De resultaten weerleggen het idee dat latente compressie op zichzelf de efficiëntie van latente diffusie verklaart. De prestatiekloof bestaat binnen dezelfde vaste latente ruimte, wat bewijst dat de "doelgeometrie" (de verdeling van het gesuperviseerde signaal) een onderscheidende factor is ten opzichte van de "representatiegeometrie" (de VAE-tokenizer).
2. Mechanisme van Verbetering: De auteurs betogen dat schone voorspelling superieur is omdat het lage-variantie latente richtingen dempt (die vaak ruis of zwak ondersteunde manifold-kenmerken zijn), terwijl snelheidsvoorspelling deze richtingen versterkt door een isotrope covariantievloer te erven.
3. Reikwijdte van Claims: De auteurs zijn bescheiden over hun theoretische claims. Zij stellen niet dat schone voorspelling globaal optimaal is voor alle tokenizers, ruisschema's of samplers. In plaats daarvan bieden zij een lokaal Gaussisch mechanisme dat consistent is met de gemeten empirische kloof. Het onderzoek wordt gepresenteerd als een gecontroleerde ablatie om het doeleffect te isoleren, en niet als een voorstel voor een nieuwe ruggengraat-architectuur of een universele vervanging voor alle diffusiedoelen.

Concluderend stelt het artikel dat zelfs in gecomprimeerde latente ruimtes het direct voorspellen van de schone latente toestand een geometrisch voordeel biedt ten opzichte van snelheidsvoorspelling, wat leidt tot aanzienlijk hogere generatiekwaliteit op ImageNet.