Large Scale Diffusion Distillation via Score-Regularized Continuous-Time Consistency

Dit paper introduceert rCM, een schaalbaar framework voor large-scale diffusion-distillatie dat via een FlashAttention-2 JVP-kern en score-regularisatie hoge-kwaliteit afbeeldingen en video's in slechts 1 tot 4 stappen genereert, terwijl het de diversiteit verbetert en modusinstorting voorkomt.

De kern

Stel je voor dat je een kunstenaar hebt die fantastische schilderijen kan maken, maar het duurt uren voordat hij één afbeelding klaar heeft. Dat is precies het probleem met moderne AI-afbeeldingsgeneratoren (zoals die voor foto's en video's). Ze zijn geweldig, maar ze zijn traag. Ze moeten stap voor stap werken, net als iemand die een schilderij langzaam opbouwt van een ruwe schets naar een fijn detail.

De onderzoekers van dit paper (uit Tsinghua University en NVIDIA) wilden een oplossing vinden: Hoe maken we deze kunstenaars supersnel, zonder dat de kwaliteit daalt?

Hier is hun verhaal, vertaald naar alledaagse taal:

1. Het Probleem: De "Snelle" Kunstenaar die Vergeet

Er was al een nieuwe techniek bedacht, genaamd sCM (een soort "Consistency Model"). De gedachte was simpel: in plaats van stap voor stap te werken, leer je de kunstenaar om direct van de ruwe schets naar het eindresultaat te springen. Dit zou de snelheid met 50 keer verhogen!

Maar er was een groot probleem:

• De "Wazige" Kunstenaar: Als je dit op grote schaal toepaste (op enorme modellen die video's maken), werd de kunstenaar wel snel, maar hij begon te "wazig" te worden. Details verdwenen, tekst op borden werd onleesbaar, en in video's zagen objecten eruit alsof ze door elkaar heen liepen.
• De Oorzaak: Het bleek dat de kunstenaar tijdens het "springen" kleine foutjes maakte. Omdat hij zo snel ging, stapelden deze foutjes zich op, net als een toren van kaarten die langzaam instort.

2. De Oplossing: Twee Leraren in Eén

De onderzoekers bedachten een slimme truc. Ze realiseerden zich dat er twee soorten "leraren" zijn in de AI-wereld, en dat ze elkaars zwakke punten kunnen opvangen:

1. De "Mode-Covering" Leraar (De Snelle, maar Wazige): Deze leert de kunstenaar om alles te zien. Hij zorgt dat er veel variatie is (diversiteit), maar negeert soms de fijne details. Dit is de basis van de snelle methode.
2. De "Mode-Seeking" Leraar (De Kwaliteitsbewaker): Deze leert de kunstenaar om perfecte afbeeldingen te maken. Hij let op elk detail, maar heeft de neiging om steeds hetzelfde te tekenen (geen variatie).

De Innovatie (rCM):
De onderzoekers hebben deze twee leraren samengevoegd in één systeem, dat ze rCM noemen.

• Ze gebruiken de Snelle Leraar om de kunstenaar te leren hoe hij in één klap van A naar B moet gaan (voor snelheid en variatie).
• Ze gebruiken de Kwaliteitsbewaker als een "controleur" die af en toe kijkt: "Hé, wacht even, die tekst op het bord is niet leesbaar, maak het scherp!"

Dit werkt als een tandemfiets: één persoon trapt hard voor snelheid, de ander zorgt dat je niet van de weg afrijdt en dat je de weg goed ziet.

3. De Technische "Magie" (Maar dan simpel)

Om dit te laten werken op enorme schaal (modellen met 14 miljard parameters, dat is groter dan de bevolking van China!), moesten ze de onderliggende techniek aanpassen:

• De Rekenmachine: De oude manier om de "sprong" te berekenen was te traag en onnauwkeurig voor zulke grote modellen. Ze bouwden een nieuwe, super-snelle rekenmethode (een "FlashAttention-kernel") die het werk als een goed georganiseerd team doet, in plaats van als een eenzame rekenaar.
• De Stabiliteit: Ze ontdekten dat de "controleur" soms te streng werd en de kunstenaar in de war bracht. Ze stelden een "balansknop" in (de $\lambda$-parameter). Als je deze goed afstelt, krijg je het perfecte evenwicht: snel, scherp, en met genoeg variatie.

4. Het Resultaat: Van Uren naar Seconden

Wat hebben ze bereikt?

• Snelheid: Video's en foto's worden nu gegenereerd in 1 tot 4 stappen in plaats van 50 of meer. Dat is een versnelling van 15 tot 50 keer!
• Kwaliteit: De afbeeldingen zijn scherp. Je kunt tekst lezen op borden, en in video's bewegen objecten natuurlijk zonder te vervormen.
• Variatie: In tegenstelling tot andere snelle methoden (die vaak steeds dezelfde saaie beelden maken), blijft rCM creatief en divers.

Samenvattend

Stel je voor dat je eerder een traag, maar perfect werkend schilderij had. Nu heb je een supersnelle kunstenaar die net zo goed werkt, maar die ook nog eens een controleur aan zijn zijde heeft die zorgt dat hij niet slordig wordt.

Dit paper laat zien dat je AI niet hoeft te kiezen tussen snelheid en kwaliteit. Met hun nieuwe methode (rCM) kun je beide hebben, zelfs voor de grootste en meest complexe video's die er bestaan. Het is een grote stap naar AI die in real-time kan denken en creëren, net als een mens.

Technische samenvatting

Probleemstelling

Hoewel continue-tijd consistentiemodellen (zoals sCM en MeanFlow) theoretisch sterk zijn en snelheidswinst bieden voor academische schalen, bleef hun toepasbaarheid op grootschalige tekst-naar-beeld (T2I) en tekst-naar-video (T2V) taken onduidelijk. Er waren twee hoofdbelemmeringen:

1. Infrastructuuruitdagingen: De berekening van Jacobiaan-vector producten (JVP), essentieel voor sCM, is moeilijk te schalen met moderne parallelle trainingsarchitecturen (zoals FlashAttention-2, FSDP en Context Parallelism) en leidt tot numerieke fouten bij lage precisie (BF16).
2. Kwaliteitsbeperkingen: Bestaande sCM-modellen vertoonden significante kwaliteitsproblemen bij het genereren van fijne details (zoals tekst) en temporale consistentie in video's. Dit werd toegeschreven aan de "mode-covering" aard van het voorwaartse divergentiedoelwit, wat leidt tot error-accumulatie en wazige of vervormde resultaten.

Bestaande state-of-the-art methoden voor grootschalige distillatie (zoals DMD2) gebruiken vaak adversarial training of score-distillatie, maar lijden vaak onder "mode collapse" (verminderde diversiteit) of vereisen complexe GAN-tuning.

Methodologie: rCM

De auteurs stellen rCM (Score-Regularized Continuous-Time Consistency Model) voor, een raamwerk dat de snelheid van consistentiemodellen combineert met de kwaliteitsvoordelen van score-distillatie.

1. Infrastructuur voor Schaalbaarheid:

   • Ontwikkeling van een FlashAttention-2 JVP-kernel (geschreven in Triton) die JVP-berekeningen integreert in de voorwaartse pass van attention-laagjes. Dit maakt training mogelijk op modellen met meer dan 10 miljard parameters en hoge-resolutie video.
   • Compatibiliteit met FSDP (Fully Sharded Data Parallel) en Context Parallelism door de netwerklagen te refactoren zodat JVP binnen elke laag kan worden uitgevoerd.
   • Stabilisatie van de JVP-berekening door het gebruik van FP32 precisie voor tijds-embeddings en het toepassen van "semi-continuous time" benaderingen om numerieke instabiliteit te voorkomen.

2. Score-Regularisatie (De Kerninnovatie):

   • De auteurs identificeren dat sCM (voorwaartse divergentie) diversiteit biedt maar lage kwaliteit, terwijl score-distillatie (reverse divergentie) hoge kwaliteit biedt maar neigt naar mode collapse.
   • rCM lost dit op door score-distillatie (specifiek DMD-loss) te integreren als een long-skip regularizer voor het sCM-objectief.
   • De totale loss functie is: $L_{rCM} = L_{sCM} + \lambda L_{DMD}$.
   • Hierbij zorgt $L_{sCM}$ voor snelheid en diversiteit, terwijl $L_{DMD}$ (gebaseerd op reverse KL-divergentie) de visuele kwaliteit en details verbetert zonder de diversiteit te verliezen.

3. Trainingsschema:

   • Het model wordt getraind zonder complexe multi-stadia of uitgebreide hyperparameter-zoekopdrachten.
   • De student kan samples genereren in 1 tot 4 stappen, wat een versnelling van 15x tot 50x oplevert ten opzichte van de leraar.

Belangrijkste Bijdragen

• Eerste grootschalige implementatie: Dit is het eerste werk dat continue-tijd consistentiedistillatie succesvol toepast op modellen tot 14 miljard parameters en video-taken tot 5 seconden.
• Infrastructuurdoorbraak: De ontwikkeling van een FlashAttention-2 JVP-kernel die JVP-computatie mogelijk maakt in grote, parallelle trainingssessies, wat eerder een bottleneck was.
• Theoretische en Praktische Combinatie: Het introduceren van een hybride doelwit dat de "mode-covering" eigenschap van consistentiemodellen combineert met de "mode-seeking" eigenschap van score-distillatie, wat resulteert in modellen met zowel hoge kwaliteit als hoge diversiteit.
• Geen GAN-tuning: In tegenstelling tot DMD2, vereist rCM geen discriminators of complexe GAN-tuning, wat het makkelijker te implementeren en te stabiliseren maakt.

Resultaten

De methode is gevalideerd op twee grote modellen: Cosmos-Predict2 (T2I, tot 14B parameters) en Wan2.1 (T2V, tot 14B parameters).

• Kwaliteit: rCM presteert op gelijke hoogte met of zelfs beter dan DMD2 op kwaliteitsmetrieken (GenEval voor T2I, VBench voor T2V). Het lost de wazigheid en vervormingen van puur sCM op, zelfs bij uitdagende prompts zoals het renderen van kleine tekst.
• Diversiteit: In tegenstelling tot DMD2, dat vaak neigt tot mode collapse (objecten op dezelfde posities/oriëntaties), behoudt rCM de hoge diversiteit van sCM. Video's tonen diverse bewegingen en objectposities.
• Snelheid: De gedistilleerde modellen genereren hoogwaardige samples in slechts 1 tot 4 stappen.
  • T2I: Competitieve resultaten in 1 stap.
  • T2V: Hoge kwaliteit in 2 stappen.
  • Dit resulteert in een versnelling van 15x tot 50x ten opzichte van de oorspronkelijke diffusion-modellen.
• Ablatie: Een balansparameter $\lambda = 0.01$ bleek de "sweet spot" te zijn voor het afwegen van kwaliteit versus diversiteit.

Betekenis

Dit werk positioneert rCM als een praktisch en theoretisch onderbouwd raamwerk voor de toekomst van grootschalige generatieve AI. Het overbrugt de kloof tussen de snelheid van consistentiemodellen en de kwaliteit van geavanceerde distillatiemethoden. Door de noodzaak van GAN-tuning en complexe scheduling te elimineren, maakt het grootschalige, snelle generatie van video en beelden toegankelijker. De resultaten suggereren dat het combineren van voorwaartse en reverse divergenties een unificerend paradigma kan zijn voor de volgende generatie generatieve modellen, met toepassingen die reiken tot interactieve wereldmodellen en autoregressieve video-generatie.