Technische Samenvatting: Parallel Decoding Distillation voor Snelle Beeld- en Videogeneratie

1. Probleemstelling

Grootschalige diffusie- en flow matching-modellen hebben state-of-the-art (SOTA) resultaten behaald in tekst-naar-beeld, tekst-naar-video en multi-modale generatie. Echter, hun inherente iteratieve bemonsteringsprocessen vereisen honderden netwerkevaluaties (NFEs), wat leidt tot aanzienlijke computationele kosten en latentie. Deze bottleneck belemmert toepassingen zoals real-time videogeneratie, contentbewerking en interactieve wereldmodellering.

Bestaande versnellingsmethoden vallen over het algemeen uiteen in twee categorieën:

• Trajectgebaseerde methoden: Deze distilleren het sequentiële bemonsteringsproces van de docent (teacher) naar minder stappen. Hoewel effectief voor afbeeldingen, verslechteren ze vaak de videokwaliteit of vereisen ze kostbare trainingsalgoritmen (bijv. Jacobian-vectorproducten of eindige verschillen) wanneer ze op videomodellen worden toegepast.
• Distributiegebaseerde methoden: Deze lijnen de marginale distributies van de student en de docent uit in plaats van specifieke trajecten te volgen. Hoewel dominant in video-distillatie (bijv. door gebruik te maken van Variational Score Distillation of adversariële verliezen), lijden ze aan mode collapse, wat leidt tot een verlies van videodiversiteit en statische beweging. Ze vereisen ook vaak afwisselende trainingsdoelstellingen en zijn gevoelig voor hyperparameters.

Het artikel identificeert de behoefte aan een schaalbare, trajectgebaseerde distillatiemethode die de optimalisatieproblemen van distributiegebaseerde methoden (mode collapse) en de computationele overhead van bestaande trajectgebaseerde methoden (JVPs, eindige verschillen) vermijdt, specifiek voor hoogresolutie videogeneratie.

2. Methodologie: Parallel Decoding Distillation (PDD)

De auteurs introduceren Parallel Decoding Distillation (PDD), een trajectgebaseerd distillatiekader ontworpen om de inferentie van diffusie- en flow matching-modellen te versnellen.

Kernconcept

In tegenstelling tot traditionele methoden die meerdere denoising-stappen samenvoegen tot één grote stap, leert PDD een parallel decoder die de gemiddelde snelheden voor meerdere opeenvolgende tijdsintervallen in een enkele netwerkevaluatie voorspelt.

• Discretisering: Het tijdsdomein $[0, 1]$ wordt gediscretiseerd in $N$ intervallen. Deze worden gegroepeerd in blokken van grootte $L$.
• Parallel Voorspellen: Gegeven een staat $X_n$ op tijd $t_n$, voorspelt de parallelle decoder $\bar{u}_\theta$ de gemiddelde snelheden voor alle $L$ intervallen binnen het blok $\{n, \dots, n+L-1\}$ gelijktijdig:
  $$\bar{u}_\theta(n|X_n) \approx u_k(X_k), \quad k = n, \dots, n+L-1$$
  waarbij $u_k$ de gemiddelde snelheid is van het $k$-de interval.
• Sampling: Tijdens de inferentie beweegt het model $L$ intervallen voor in één stap met behulp van de voorspelde snelheden, waardoor het totaal aantal NFEs wordt verminderd tot $N/L$. Door de blokgrootte $L$ tijdens de training te variëren, ondersteunt het model variabele NFEs tijdens de inferentie zonder extra tijdsconditionering.

Trainingsdoelstelling

PDD maakt gebruik van een vereenvoudigde, regressiegebaseerde trainingsdoelstelling die VSD, adversariële verliezen, JVPs en eindige verschillen vermijdt.

• Docent Benadering: De doelgemiddelde snelheden worden benaderd met behulp van een Runge-Kutta solver (Euler of Midpoint) toegepast op het vooraf getrainde docentmodel.
• On-Policy Training: Het verlies wordt berekend op de eigen outputs van de student (on-policy). De student voorspelt snelheden, simuleert het traject binnen het blok, en het verlies wordt berekend door de voorspelde snelheid van de student bij een gesamplede stap $k$ te vergelijken met de benaderde gemiddelde snelheid van de docent bij die staat.
• Verliesfunctie:
  $$\mathcal{L}_{PD}(\theta) = \mathbb{E} \left[ \| \bar{u}_\theta(n|X_n) - u_k(\text{sg}(\bar{X}_k)) \|^2 \right]$$
  waarbij $\text{sg}$ de stop-gradient operator is, en $\bar{X}_k$ de staat is die bereikt wordt door de voorspelde snelheden van de student te volgen.

Architectuur

• Backbone: PDD gebruikt exact dezelfde backbone-architectuur als het vooraf getrainde docentmodel.
• Outputlaag: De laatste lineaire laag wordt uitgebreid door de kanaaldimensie $N$ keer te herhalen (waarbij $N$ de roostergrootte is). Dit stelt het model in staat om $N$ verschillende snelheidsprognoses uit te voeren die overeenkomen met elke tijdstap.
• Laagfusie: Tijdens de generatie, in plaats van $L$ aparte lineaire lagen toe te passen, worden de gewichten gefuseerd in een enkele lineaire laag die de gewogen gemiddelde snelheid voor het blok uitvoert. Dit zorgt ervoor dat er geen extra computationele kosten zijn tijdens de inferentie vergeleken met de enkelvoudige stap-evaluatie van de docent.

Data-vrije Training

Voor grootschalige modellen waarbij data niet beschikbaar is, ondersteunt PDD een data-vrij trainingsschema. Het model samplet een initiële ruis en wisselt tussen solver-stappen (het voortstuwen van de staat met behulp van de eigen voorspellingen van de student) en trainingsstappen, waardoor het effectief leert van zijn eigen gegenereerde trajecten.

of 3. Belangrijkste Bijdragen

1. Formulering van PDD: Een schaalbare, trajectgebaseerde distillatiemethode die gemiddelde snelheden voor meerdere intervallen voorspelt in een enkele evaluatie, wat snelle inferentie mogelijk maakt voor flow matching en diffusiemodellen.
2. Vereenvoudigde Training: Een enkele regressiegebaseerde doelstelling die de noodzaak voor JVPs, eindige verschillen, meerfasige training, VSD of GAN-verliezen elimineert, terwijl het nog steeds hoogwaardige en diverse samples produceert.
3. Flexibele Architectuur: Een ontwerp dat compatibel is met elk vooraf getraind model dat variabele NFEs ondersteunt tijdens de inferentie zonder dat extra tijdsconditionering of een tweede tijdscoördinaat vereist is.
4. Empirische Validatie: Demonstratie van SOTA-prestaties op ImageNet-256, Qwen-Image (Tekst-naar-Beeld), Wan2.1 (Tekst-naar-Video 1.3B/14B) en LTX-2.3 (Tekst-naar-Video/Audio), met significante verbeteringen in generatiediversiteit vergeleken met distributiegebaseerde baselines.

4. Experimentele Resultaten

De auteurs evalueerden PDD over drie taken:

• Class-Conditional Beeldgeneratie (ImageNet-256):

  • PDD bereikte concurrerende FID-scores in de single-step setting (FID 2.69 met Midpoint), waarbij het SOTA-methoden zoals FreeFlow en Pi-Flow evenaart of overtreft, terwijl het variabele NFEs ondersteunt (1, 2, 4, 8).
  • De methode vertoonde stabiele convergentie over verschillende hyperparameters.

• Tekst-naar-Beeld (Qwen-Image 20B):

  • Geëvalueerd op OneIG, DPG-Bench en GenEval benchmarks met 2, 4 en 8 NFEs.
  • PDD (Midpoint) behaalde de beste algemene scores op OneIG en DPG-Bench bij 4 en 8 NFEs, waarmee het baselines zoals DMD2, TwinFlow en Pi-Flow overtrof.
  • Cruciaal is dat PDD een aanzienlijk hogere diversiteit (gemeten door OneIG diversiteitsmetrieken) behield vergeleken met DMD2, dat leed aan mode collapse.

• Tekst-naar-Video (Wan2.1 1.3B & 14B):

  • Geëvalueerd op VBench met 4 en 8 NFEs.
  • PDD bereikte SOTA videokwaliteit en algemene scores, waarbij het eerste werd op het 1.3B model en tweede op het 14B model (achter AnyFlow in algemene metrieken, maar met betere diversiteit).
  • Diversiteit: PDD demonstreerde superieure diversiteitsscores vergeleken met distributiegebaseerde baselines (DMD2, AnyFlow).
  • Beweging: Visuele analyse gaf aan dat door PDD gegenereerde video's hogere graden van beweging en dynamische consistentie vertoonden vergeleken met baselines, die vaak statische of repetitieve video's produceerden.

• Tekst-naar-Video/Audio (LTX-2.3):

  • PDD destilleerde het 22B LTX-2.3 model voor 10s 720p videogeneratie met audio met slechts 8 NFEs.
  • Na slechts 250 trainingsiteraties presteerde PDD op gelijke hoogte met of beter dan het officiële 8-staps gedistilleerde model, ondanks het feit dat het geen toegang had tot trainingsdata.
  • Menselijke voorkeursevaluatie (via Gemini 3.1 Pro) liet zien dat PDD won of gelijk was aan het officiële gedistilleerde model op het gebied van prompt alignment, visuele kwaliteit, bewegingskwaliteit en audiokwaliteit.

5. Betekenis en Claims

Het artikel claimt dat PDD de eerste pure trajectgebaseerde distillatiemethode is die in staat is tot hoogresolutie, few-step videogeneratie met verbeterde diversiteit.

• Overwinnen van Mode Collapse: Door vast te houden aan een trajectgebaseerde aanpak zonder te vertround op distributiegebaseerde verliezen (VSD/GANs), vermijdt PDD de mode collapse en het gebrek aan beweging dat vaak wordt gezien bij huidige video-distillatiemethoden.
• Computationele Efficiëntie: De methode elimineert de noodzaak voor dure JVPs en eindige verschillen, waardoor het schaalbaar is naar grote modellen (tot 22B parameters) en compatibel is met data-vrije training.
• Flexibiliteit: Het vermogen om variabele NFEs tijdens de inferentie te ondersteunen via blokgrootte-selectie biedt een praktisch voordeel voor diverse toepassingsscenario's.
• Beweging en Diversiteit: De auteurs benadrukken dat PDD niet alleen de generatie versnelt, maar ook de "rijkdom" van de output van de docent behoudt, specifiek met betrekking tot bewegingsdynamiek en de diversiteit van samples, wat cruciaal is voor videogeneratie.

De auteurs concluderen dat hoewel PDD momenteel beperkt is tot flow- en diffusiemodellen, het principe van parallelle decodering in de toekomst mogelijk kan worden gegeneraliseerd naar discrete autoregressieve modellen. Ze merken ook op dat het onderzoeken van PDD in data-afhankelijke settings buiten ImageNet-256 een open gebied blijft voor toekomstig onderzoek.