Technische Samenvatting: Distillering van Drifting Transformers met Representation Autoencoders

Probleemstelling

Hoewel Representation Autoencoders (RAEs) de prestaties van flow-gebaseerde generatieve modellen aanzienlijk hebben verbeterd door gebruik te maken van semantisch rijke, label-specifieke geclusterde latente ruimtes afgeleid van gepre-trainde DINO-features, blijft het distilleren van deze modellen naar efficiënte one-step of few-step generatoren een uitdaging.

De auteurs identificeren een fundamentele geometrische mismatch: RAE latente ruimtes vertonen extreme anisotropie en grote krommingen in hun onderliggende Ordinary Differential Equation (ODE) trajecten vergeleken met traditionele Variational Autoencoder (VAE) ruimtes. Bestaande traject-gebaseerde distillatiemethoden vertrouwen op de assumptie van gladde, bijna rechte teacher-trajecten. De sterk gekromde trajecten in RAE-ruimtes zorgen ervoor dat deze methoden instabiel of inefficiënt worden. Bovendien vereist de recent voorgestelde Drifting Model, die een distributie-gebaseerde aanpak hanteert om gekromde trajecten te vermijden, in de oorspronkelijke formulering een hulp-module in de vorm van een apart getrainde Masked Autoencoder (MAE) om geconcentreerde features te leveren, wat zorgt voor aanzienlijke computationele overhead.

Methodologie

Het artikel stelt Drift-RAE voor, een distillatiekader dat het Drifting Model-paradigma specifiek aanpast voor RAE latente ruimtes, waardoor de noodzaak voor hulp-feature extractoren wordt geëlimineerd.

1. Theoretische Analyse van Geometrie en Dynamica

De auteurs analyseren eerst kwantitatief de latente ruimtes en tonen aan dat RAE-features aanzienlijk anisotroper zijn (lagere Participation Ratio en Spectral Entropy) dan SD-VAE-features. Ze bewijzen theoretisch (Theorem 1) dat in hoog-dimensionale ruimtes, indien positieve samples te ver verspreid zijn, het aantrekkende component van het drifting veld de neiging heeft te verdwijnen, wat leidt tot een gedegenereerd veld.

• Kerninzicht: Waar de oorspronkelijke Drifting Model een MAE nodig heeft om verspreide features te concentreren, zijn RAE latente ruimtes natuurlijk semantisch geconcentreerd. Dit stelt de Drifting Model in staat om direct op RAE latents te opereren zonder de hulp-MAE module.

2. Het Drift-RAE Kader

De methode distilleert een gepre-trainde flow model $v_\theta$ naar een one-step generator $G_\theta(z) = z - v_\theta(z, 1)$. In plaats van ODE-trajecten te matchen, wordt de generator getraind om een drifting objectief te minimaliseren dat gegenereerde samples richting de echte datadistributies duwt en weg van andere gegenereerde samples.

• Objectief: De loss wordt token-wise berekend op de RAE latente representaties, waarbij gebruik wordt gemaakt van een genormaliseerd drifting veld $\tilde{V}$ afgeleid van positieve (echte) en negatieve (gegenereerde) samples.

3. Praktische Modificaties voor Stabiliteit

Om de Drifting Model af te stemmen op het Diffusion-GAN framework en de trainingsstabiliteit te verbeteren, introduceren de auteurs drie specifieke modificaties:

• Softmax Dimensie: De auteurs verwijderen de bi-directionele softmax trick die in de oorspronkelijke Drifting Model wordt gebruikt. Ze behouden een enkele softmax over sample indices, wat beter aansluit bij de gradiëntrichting van de onderliggende log-sum-exp potentiaal.
• Input Perturbatie: Om gradiëntverdwijning veroorzaakt door niet-snijdende manifolds tussen echte en gegenereerde data te voorkomen, worden de gegenereerde samples geperturbeerd met ruis gesampled uit een Laplace-distributie voordat het drifting veld wordt berekend.
• Partiële Detachment van Negatieve Samples: Om het aantal negatieve samples te vergroten dat wordt gebruikt voor distributie-approximatie zonder het GPU-geheugenbudget te overschrijden, stellen de auteurs voor om een grotere pool van negatieve samples ($N_{total\_neg}$) te gebruiken waarbij slechts een subset ($N_{neg}$) gradiënten draagt, terwijl de rest wordt 'detached'. Dit verbetert de empirische benadering van de gegenereerde distributie.

Belangrijkste Bijdragen

1. Geometrische Analyse: Een kwantitatieve demonstratie dat RAE latente ruimtes een hoge kromming en anisotropie bezitten, waardoor standaard traject-gebaseerde distillatie instabiel wordt.
2. Theoretische Rechtvaardiging: Een bewijs dat Drifting Models falen op te ver verspreide ruimtes (wat de MAE in de oorspronkelijke formulering noodzakelijk maakt) maar floreren in de natuurlijk geconcentreerde semantische ruimtes van RAEs.
3. Drift-RAE: Een nieuwe distillatiemethode die state-of-the-art one-step generatie in RAE ruimtes bereikt zonder een hulp-MAE feature extractor te vereisen.
4. Stabiliteitsverbeteringen: Theoretische en empirische inzichten in het modificeren van de Drifting objectief (softmax dimensie, input perturbatie en sample detachment) om de Drifting Model af te stemmen op Diffusion-GAN principes en convergentie te verbeteren.

Experimentele Resultaten

De methode is geëvalueerd op de ImageNet 256 $\times$ 256 dataset met DiTDH-XL en DiTDH-L backbones.

• Prestaties: Drift-RAE bereikt een FID van 1.77 (met DiTDH-XL) en 2.12 (met DiTDH-L) met slechts één generatiestap.
• Efficiëntie: Het model bereikt deze resultaten in slechts 10.000 distillatiestappen (ca. 16 epochs).
• Vergelijking:
  • Het presteert beter dan vorige distillatiemethoden in RAE ruimtes (bijv. MF-RAE).
  • Het bereikt een vergelijkbare FID en verbeterde FDDINOv2 scores vergeleken met de oorspronkelijke Drifting Model, ondanks het niet gebruiken van de computationeel dure hulp-MAE.
  • Het blijft competitief met one-step en few-step generatoren getraind in andere latente ruimtes (bijv. Pixel Space, SD-VAE).

Betekenis en Claims

Het artikel claimt dat Drifting een effectief en veelbelovend kader biedt voor het distilleren van generatieve modellen in representatieruimtes. De primaire betekenis ligt in het oplossen van de instabiliteit bij het distilleren van RAE-gebaseerde modellen door gebruik te maken van hun inherente semantische concentratie, waardoor de afhankelijkheid van hulp-feature extractoren wordt weggenomen.

De auteurs merken bescheiden de beperkingen op, waarbij zij erkennen dat hun theoretische analyse steunt op vereenvoudigde hoog-dimensionale modellen en dat het vanaf nul trainen van Drifting Models (zonder distillatie) in RAE ruimtes een openstaand probleem blijft. Ze wijzen ook op het feit dat de vereiste voor veel positieve samples van dezelfde klasse schaalbaarheidsuitdagingen kan vormen voor tekst-naar-beeld generatie, wat zij als een richting voor toekomstig onderzoek voorstellen.