Beyond Single Tokens: Distilling Discrete Diffusion Models via Discrete MMD

Probleemstelling

Discrete diffusiemodellen (voor tekst en afbeeldingen) vereisen momenteel een groot aantal steppings (iteraties) om hoogwaardige data te genereren. In tegenstelling tot continue diffusiemodellen, waar er geavanceerde distillatiemethoden bestaan om het aantal stappen drastisch te verminderen, ontbreekt dit voor discrete modellen.

• Huidige beperkingen: Bestaande methoden voor discrete distillatie (zoals SDTT) leiden vaak tot instorting van de kwaliteit en diversiteit van de generaties.
• Efficiëntie: Hoewel discrete diffusiemodellen (zoals Diffusion LLMs) betere acceleratorgebruik hebben dan causale LLMs (die token-voor-token werken), zijn ze inefficiënt vanwege de hoge rekentijd die nodig is voor convergentie.
• Doel: Het ontwikkelen van een methode om discrete diffusiemodellen te distilleren naar "few-step" generators (weinig stappen) zonder in te leveren op kwaliteit of diversiteit, en waarbij de student zelfs de leraar kan overtreffen.

Methodologie: Discrete Moment Matching Distillation (D-MMD)

De auteurs introduceren D-MMD, een generalisatie van de Moment Matching Distillation (MMD) uit het continue domein, specifiek ontworpen voor discrete processen.

1. Kernconcept:

   • In plaats van de log-likelihood direct te maximaliseren (wat moeilijk is bij discrete steppen), probeert D-MMD de momenten (verwachtingen) van de verdeling van de leraar na te bootsen.
   • Het algoritme gebruikt een min-max spel tussen drie componenten:
     • De Student-generator ($\hat{x}_\eta$): Probeert de loss onder de leraar te minimaliseren, maar de loss onder een hulpmodel te maximaliseren.
     • De Leraar ($\hat{x}_\theta$): Een vast model dat de doelverdeling vertegenwoordigt.
     • Het Hulpmodel ($\hat{x}_\phi$): Wordt getraind om de verwachting van de student te voorspellen, maar wordt geregulariseerd om dicht bij de leraar te blijven.

2. Aanpassing voor Discrete Ruimte:

   • Bij continue modellen worden gradients direct berekend op de uitkomsten. Bij discrete modellen (categorische verdelingen) is dit lastig vanwege de niet-differentieerbare steppen.
   • Soft Targets: In plaats van harde steppen (hard samples) te gebruiken, gebruikt D-MMD de zachte waarschijnlijkheidsvectoren (soft probability vectors) van de student.
   • Log-probabiliteit: De update-regel voor de student gebeurt in de ruimte van log-probabiliteiten. De gradient is evenredig met het verschil in log-probabiliteit tussen de leraar en het hulpmodel:
     $$\nabla_\eta L \propto \log \hat{x}_\theta - \log \hat{x}_\phi$$
   • Dit zorgt ervoor dat de student de output-entropie verlaagt en correlaties tussen tokens leert, ondanks dat de outputverdeling factorized (onafhankelijk) is.

3. Trainingsschema:

   • Het algoritme wisselt af tussen het updaten van de student (minimiseren van de leraar-afwijking) en het updaten van het hulpmodel (maximaliseren van de overeenkomst met de leraar en de student).
   • Voor Masked Diffusion kan de student zachte targets gebruiken; voor Uniform Diffusion zijn harde steppen nodig voor het trainen van het hulpmodel.

Belangrijkste Bijdragen

• D-MMD Algoritme: Een nieuwe, robuuste distillatiemethode die werkt voor zowel gemaskerde als uniforme discrete diffusieprocessen.
• Overtreffen van de Leraar: Het is voor het eerst dat een gedistilleerde discrete generator (met weinig stappen) betere kwaliteit bereikt dan de oorspronkelijke leraar (die veel stappen nodig heeft).
• Nieuwe Evaluatiemetric (Gradient Moment): De auteurs introduceren de Gradient Moment (GM) van een referentiemodel (bijv. GPT-2) als nieuwe metriek voor tekst.
  • Waarom? Bestaande metrics zoals "Generative Perplexity" zijn onbetrouwbaar omdat ze kunnen worden "gegamd" door lage temperaturen (wat leidt tot repetitieve, maar lage perplexiteit teksten).
  • Hoe werkt het? De metric meet de grootte van de gradient van een getraind LLM op de gegenereerde samples. Als de samples goed zijn, is de gradient klein (het model ziet ze als trainingsdata). Als de samples slecht zijn, is de gradient groot.

Resultaten

De methode is getest op zowel afbeeldingen (CIFAR-10) als tekst (Open Web Text).

• Afbeeldingen (CIFAR-10):

  • D-MMD-generatoren behalen een veel lagere FID (Fréchet Inception Distance) dan de leraars met een fractie van het aantal stappen.
  • Voorbeeld: Een Masked D-MMD-generator bereikt een FID van 3.5 met slechts 64 stappen, terwijl de leraar (1024 stappen) een FID van 6.4 haalt. De leraar zelf haalt met 1024 stappen een FID van 6.4.
  • De Pareto-front (stappen vs. kwaliteit) van D-MMD ligt aanzienlijk beter dan die van de leraar.

• Tekst:

  • D-MMD overtreft de leraar bij tekstgeneratie met slechts 16 stappen (Masked D-MMD: 0.236 GPT-2 GM vs. leraar 0.275 bij 256 stappen).
  • Bestaande methoden zoals SDTT (Progressive Distillation) degradeerden naarmate het aantal distillatiestappen toenam, terwijl D-MMD stabiel bleef en beter werd.
  • Noise Conditioning: Voor masked distillatie bleek het toevoegen van extra ruisinvoer essentieel voor de kwaliteit, wat de generator in staat stelt zijn outputverdeling effectiever te "instorten" (collapse) zonder kwaliteit te verliezen.

Betekenis en Impact

• Efficiëntie: D-MMD maakt het mogelijk om discrete diffusiemodellen (voor tekst en beeld) te gebruiken in scenario's waar snelheid cruciaal is, zonder de kwaliteit op te offeren.
• Paradigmaverschuiving: Het weerlegt de aanname dat gedistilleerde modellen per definitie slechter zijn dan hun leraars. Door het gebruik van een adversariale instelling (student vs. hulpmodel) kan de student de "mode-seeking" eigenschappen van de leraar verbeteren.
• Evaluatie: De introductie van de Gradient Moment-metric biedt een betrouwbaarder alternatief voor perplexiteit bij het evalueren van generatieve modellen, wat essentieel is voor de toekomstige ontwikkeling van discrete diffusiemodellen.
• Toepasbaarheid: De methode is breed toepasbaar, van onvoorwaardelijke beeldgeneratie tot block-wise autoregressieve tekstgeneratie, en biedt een solide basis voor verdere optimalisatie van discrete generatieve AI.