Technische Samenvatting: Verbeterde Baselines met Representation Autoencoders (RAEv2)

Probleemstelling

Representation Autoencoders (RAE) hebben zich ontwikkeld als een veelbelovend kader voor het vervangen van traditionele Variational Autoencoders (VAE's) bij het trainen van diffusion transformers, met als doel tokenisatie te uniformeren voor zowel visueel begrip als generatie. De oorspronkelijke RAE-formulering stuit echter op drie kritieke knelpunten die praktische adoptie belemmeren:

1. Suboptimale Reconstructie: De reconstructieprestaties blijven achter bij gespecialiseerde VAE's, met name wanneer getraind zonder gespecialiseerde data (bijvoorbeeld tekst of gezichten).
2. Incompatibiliteit met Guidance: RAE heeft moeite met traditionele Classifier-Free Guidance (CFG), wat noodzaakt tot het trainen van een secundair, zwakker diffusion-model voor AutoGuidance (AG). Dit voegt aanzienlijke rekenkosten en complexiteit toe.
3. Onderbenutte Encoder Representaties: Eerdere werken gaan standaard uit van het gebruik van alleen de laatste laag van een voorgetrainde visuele encoder als latente representatie, en negeren hierbij de rijke, diverse abstracties die in tussenliggende lagen aanwezig zijn.

Methodologie

Het artikel stelt RAEv2 voor, een verbeterde baseline die deze problemen aanpakt via drie kerninzichten en ontwerpveranderingen, terwijl tegelijkertijd een trainingsvrije aanpak voor de encoder wordt behouden.

1. Generaliseerde Representatie-encoder

In plaats van uitsluitend te vertrouwen op de laatste laag ($K=1$) van een voorgetrainde visuele encoder, stellen de auteurs een generaliseerde formulering voor waarbij de latente representatie de som is van de laatste $K$ encoder-lagen.

• Mechanisme: De output wordt gedefinieerd als $\mathbf{x} = \sum_{\ell=L-K+1}^{L} \mathbf{z}_\ell$, waarbij $\mathbf{z}_\ell$ de feature-map in laag $\ell$ voorstelt.
• Redenering: Verschillende lagen vullen elkaar aan; eerdere tot middelste lagen behouden fijnere ruimtelijke structuren, terwijl latere lagen globale semantiek benadrukken. Het aggregeren ervan behoudt de geometrische structuur zonder de latent-dimensie te vergroten of extra geleerde parameters te vereisen.
• Implementatie: De auteurs vergelijken eenvoudige optelling (MLS) met projectie via een willekeurige matrix (MLR) en constateren dat eenvoudige optelling consequent betere generatieprestaties oplevert, terwijl de reconstructie verbetert.

2. Complementaire Mechanismen van RAE en REPA

Het artikel daagt de heersende aanname uit dat RAE de noodzaak voor Representation Alignment (REPA) elimineert. REPA distilleert dezelfde voorgetrainde representatie naar tussenliggende diffusion-lagen.

• Vinding: Empirische analyse over 27 visuele encoders onthult dat RAE en REPA werken via complementaire mechanismen.
  • RAE biedt een semantisch rijke latente ruimte (gebruikmakend van globale semantiek, gemeten door Linear Probing-accuratie).
  • REPA regulariseert de ruimtelijke structuur van tussenliggende diffusion-features (gemeten door Local Distance Similarity, LDS).
• Resultaat: Het gebruik van dezelfde representatie als zowel de encoder (RAE) als het doel voor tussenliggende lagen (REPA) verbetert consequent de generatieprestaties. Deze synergie stelt sterkere encoders (bijvoorbeeld DINOv3-L), die uitblinken in zowel globale als ruimtelijke metrieken, in staat superieure generatieresultaten te behalen.

3. REPA als Self-Guidance (x-prediction)

De oorspronkelijke RAE vereist een apart zwak model voor guidance. De auteurs herformuleren de REPA-head om interne guidance mogelijk te maken.

• Inzicht: In de RAE-latente ruimte voert de REPA-projectie-head (die features van vroege lagen afbeeldt op de schone latent) effectief x-prediction uit.
• Mechanisme: Door de volledige modeloutput te herparameteriseren tot ook x-prediction, dient de voorspelling van de REPA-head ($\hat{\mathbf{x}}_{repa}$) als een "zwakkere" baseline vergeleken met de voorspelling van het volledige model ($\hat{\mathbf{x}}_{full}$).
• Guidance-formule: Guidance wordt intern toegepast binnen één forward pass: $\hat{\mathbf{x}}_{guided} = \hat{\mathbf{x}}_{full} + w \cdot (\hat{\mathbf{x}}_{full} - \hat{\mathbf{x}}_{repa})$.
• Voordeel: Dit elimineert de noodzaak voor een apart AutoGuidance-model en vermijdt de extra forward pass die door CFG wordt vereist, waardoor het aantal function evaluations (NFEs) voor guidance effectief wordt gehalveerd.

Belangrijkste Resultaten

RAEv2 toont aanzienlijke verbeteringen in convergentiesnelheid, generatiekwaliteit en reconstructie op ImageNet, tekst-naar-beeld en navigatie-wereldmodellen.

• Trainings-efficiëntie: RAEv2 bereikt 10$\times$ snellere convergentie vergeleken met de oorspronkelijke RAE. Het bereikt een state-of-the-art guided gFID van 1.06 in slechts 80 epochs op ImageNet-256, terwijl de oorspronkelijke RAE 800 epochs nodig had om 3.26 te bereiken.
• Nieuwe Efficiëntiemetriek (EPFID): Het artikel introduceert EPFID@k (epochs om een ongeleide gFID $\le k$ te bereiken) als een praktischer metriek dan absolute gFID. RAEv2 bereikt een EPFID@2 van 35 epochs, een enorme verbetering ten opzichte van de oorspronkelijke RAE (177 epochs) en eerdere werken (>480 epochs).
• Generatiekwaliteit: Op de recent voorgestelde Representation Fréchet Distance ($FD_r$) bereikt RAEv2 een score van 2.17 in 80 epochs, en overtreft hiermee de vorige beste van 3.26 (getraind voor 800 epochs) zonder enige post-training.
• Reconstructie: De generaliseerde formulering (specifiek met $K=23$) bereikt een Pareto-optimale afweging tussen reconstructie (rFID 0.18) en generatie, en presteert beter dan propriëtaire VAE's (bijvoorbeeld Flux-VAE, SD3.5), ondanks dat het uitsluitend is getraind op ImageNet.
• Generalisatie:
  • Tekst-naar-beeld: RAEv2 toont consistente verbeteringen ten opzichte van Flux-VAE en de oorspronkelijke RAE op GenEval, DPG-Bench en GenAI-Bench.
  • Navigatie-wereldmodellen: Bij taken voor voorspelling van toekomstige frames (RECON-dataset) bereikt RAEv2-NWM een FVD van 105.61, wat significant beter is dan NWM (200.97) en RAE (312.01), met verminderd flickeren bij autoregressieve rollouts.

Betekenis en Claims

Het artikel beweert dat RAEv2 een vereenvoudigde, zeer efficiënte baseline biedt voor representation autoencoders. De primaire bijdragen zijn:

1. Vereenvoudiging: Het verwijdert de noodzaak voor encoder-finetuning, gespecialiseerde trainingsdata of aparte guidance-modellen.
2. Efficiëntie: Het reduceert drastisch de trainingsduur die nodig is om state-of-the-art prestaties te bereiken, en verschuift de focus van verbeteringen in absolute metrieken naar trainings-efficiëntie (EPFID).
3. Gefuseerde Tokenisatie: Door het behoud van de globale semantische begrijpingscapaciteiten van de bevroren encoder (geverifieerd via linear probing) terwijl generatie en reconstructie worden verbeterd, versterkt het de levensvatbaarheid van het gebruik van voorgetrainde visuele encoders als een gefuseerd tokenisatieschema voor zowel begrijpings- als generatietaken.

De auteurs hopen dat deze inzichten de praktische adoptie van representation autoencoders in diverse generatieve modeleringsscenario's faciliteren.