Technische Samenvatting: Token Time Continuous Diffusion voor Taalmodellering

Probleemstelling

Huidige diffusie-taalmodellen (LM's) opereren voornamelijk in de discrete tokenruimte, waarbij gebruik wordt gemaakt van iteratieve verfijningen zoals het ontmaskeren van null-tokens of het denoisen vanuit uniforme willekeur. Hoewel deze modellen competitieve benchmarks behalen, lijden ze aan een kritieke beperking die bekend staat als het factorisatieprobleem in scenario's met een hoge versnelling (weinig stappen) tijdens de generatie. Wanneer meerdere tokens simultaan worden gedenoised in één stap, samplet het model uit een productmarginale distributie in plaats van de werkelijke onderliggende gezamenlijke distributie. Dit leidt tot een significante degradatie van de tekstkwaliteit naarmate het aantal generatiestappen afneemt.

Bovendien kampen bestaande diffusiemodellen in de continue ruimte met het "plotselinge transitieprobleem", waarbij token-embeddings oninformatief blijven bij hoge ruisniveaus en zeer zeker worden bij lage ruisniveaus, met een zeer smalle tussenliggende transitiezone. Eerdere pogingen om dit aan te pakken via globale tijdwarping faalden omdat ze geen rekening hielden met de inherente variabiliteit in de moeilijkheid van het afleiden van verschillende tokens binnen een sequentie. Daarnaast hebben globale timemodellen moeite met het onderscheid tussen een schone input-prompt en een ruisend canvas dat generatie vereist, wat leidt tot suboptimale conditionele generatie.

Methodologie: Token Time Continuous Diffusion (TTCD)

De auteurs stellen Token Time Continuous Diffusion (TTCD) voor, een nieuw framework dat opereert in de continue embeddingruimte en het concept van per-token tijden introduceert.

Kernmechanismen

1. Evolutie in de Continue Ruimte: TTCD brengt Gaussische ruis (tijd $t=0$) deterministisch naar schone token-embeddings (tijd $t=1$) zonder tussenliggende samplingstappen, waardoor de factorisatieproblemen die inherent zijn aan de discrete ruimte worden vermeden.
2. Per-Token Tijden: In tegen tegenstelling tot eerdere methoden die een enkel globaal ruisniveau ($t_g$) opleggen aan alle tokens, wijst TTCD een unieke lokale tijd $t_l$ toe aan elk token $l$.
   • Rangtoewijzing: Elk token krijgt een rang $u_l \in [0, 1]$.
   • Tijdfunctie: De lokale tijd wordt bepaald door een functie $t_l = F(t_g, u_l)$. De auteurs selecteren $F$ zodanig dat $t_l$ een Beta-distributie volgt, wat ervoor zorgt dat hoger gerangschikte tokens (die gemakkelijker af te leiden zijn) sneller van ruis naar data progresseren dan lager gerangschikte tokens.
   • Eigenschappen: De functie $F$ zorgt ervoor dat bij globale tijd $t_g=0$, alle tokens ruis zijn; bij $t_g=1$, alle tokens data zijn; en dat het gemiddelde van de lokale tijden gelijk is aan de globale tijd.
3. Dynamisch Denoising: Tijdens de inferentie worden tokens met verschillende snelheden gedenoised op basis van hun toegewezen rangen. "Zekere" tokens (lage entropie) krijgen hogere rangen en worden eerder gedenoised, terwijl "onzekere" tokens later worden gedenoised. Dit stelt het model in staat om de generatie van moeilijke tokens te conditioneren op de reeds bepaalde "zekere" tokens, wat de interactie tussen token-afhankelijkheden beter modelleert.
4. Conditionele Generatie: Voor prefix-geconditioneerde taken krijgen input-prompt tokens een vaste tijd van $t=1$ (schoon) gedurende het proces, terwijl canvas-tokens evolueren van $t=0$ naar $t=1$. Dit onderscheidt de prompt expliciet van de gegenereerde inhoud.
5. Architecturale Modificatie: Het model past de Diffusion Transformer (DiT) architectuur aan door de globale tijd-conditionering in de adaptive LayerNorm (adaLN) te vervangen door per-token tijd-conditionering. Deze modificatie introduceert geen extra parameters.
6. Distillatie (Shortcut Modellen): Om de inferentie te versnellen, gebruiken de auteurs een zelf-distillatiestrategie met behulp van "shortcut"-modellen. Het model wordt gefinetuned om de gemiddelde flow-snelheid te voorspellen tussen willekeurige start- en eind-globale tijden, geconditioneerd op token-tijden, wat hoogwaardige generatie mogelijk maakt in zeer weinig stappen (bijv. 1-4 stappen).

Belangrijkste Bijdragen

1. Continue Ruimte met Per-Token Snelheden: De introductie van een continue-ruimte diffusie-LM waarbij tokens met verschillende snelheden worden gedenoised op basis van toegewezen rangen, wat de gedachte formaliseert dat de identiteit van tokens verschillende niveaus van afleidbaarheid heeft.
2. Parameter-vrije Architecturale Verandering: Een modificatie van de standaard DiT-architectuur om per-token tijd-conditionering te ondersteunen zonder het aantal parameters te verhogen.
3. State-of-the-Art Prestaties bij Weinig Stappen: Demonstratie dat TTCD bestaande discrete en continue diffusiemodellen overtreft in scenario's met een hoge versnelling (generatie met weinig stappen).
4. Effectieve Distillatie: Een succesvolle toepassing van op shortcuts gebaseerde zelf-distillatie op continue flow-modellen, wat resulteert in state-of-the-art resultaten voor generatie met weinig stappen.

Experimentele Resultaten

De auteurs evalueerden TTCD op drie taken:

1. Sudoku Oplossen (Algoritmisch Redeneren):

   • Bij 9x9 Sudoku oplossen met een generatiebudget van 2 stappen, behaalde TTCD een nauwkeurigheid van 31,51%.
   • Dit presteerde aanzienlijk beter dan discrete masking-baselines (11,65% met entropie-gebaseerde unmasking) en naïeve continue methoden (bijna 0%).
   • De resultaten benadrukken het vermogen van TTCD om tokens correct te rangschikken op zekerheid en de factorisatiefout te vermijden die de discrete parallelle sampling teistert.

2. OpenWebText Taalgeneratie:

   • Een TTCD-model met 160 miljoen parameters werd getraind en zelf-gedistilleerd.
   • Onconditionele Generatie: In 1-4 stappen generatie was de gedistilleerde TTCD vergelijkbaar met of presteerde het beter dan discrete baselines (Duo, MDLM) en de continue baseline (FLM) in termen van generatieve perplexiteit (Gen PPL) en tekstentropie. Opvallend genoeg behield TTCD een hogere entropie (diversiteit) vergeleken met FLM, dat leed aan entropie-instorting.
   • Conditionele Generatie: In prefix-geconditioneerde taken (het genereren van 32 of 128 tokens gegeven een prompt), presteerde TTCD significant beter dan alle baselines, inclus van de gedistilleerde versies van Duo en FLM. Het per-token tijdmechanisme bleek van nature geschikt voor het afhandelen van het verschil tussen schone prompts en ruisende canvas-sequenties.

3. Molecuulgeneratie (QM9):

   • In classifier-free guidance experimenten genereerden continue-ruimte methoden met token-tijd (TTCD) meer valide moleculen en bereikten ze een betere frontier-prestatie vergeleken met discrete baselines.

Betekenis en Claims

Het artikel beweert dat TTCD het fundamentele "factorisatieprobleem" van discrete diffusiemodellen aanpakt door in een continue ruimte te opereren, en het "plotselinge transitieprobleem" van bestaande continue modellen door per-token tijden te introduceren.

De auteurs stellen dat:

• Haalbaarheid bij Hoge Versnelling: TTCD is het eerste diffusie-taalmodel dat competitieve of superieure prestaties levert aan discrete modellen in scenario's met een hoge versnelling (weinig stappen), een regime waarin discrete modellen doorgaans falen vanwege sampling-onnauwkeurigheden.
• Superioriteit in Conditionele Generatie: De expliciete afhandeling van per-token tijden maakt superieure conditionele generatie mogelijk, aangezien het model natuurlijk het onderscheid kan maken tussen vaste input-prompts en evoluerende output-sequenties.
• Schaalbaarheid: De methode schaalt effectief, wat wordt aangetoond door de training van een 160M parameter TTCD-model dat vergelijkbare zelf-gedistilleerde baselines op dezelfde data overtreft.

De auteurs blijven bescheiden over de beperkingen en merken op dat de huidige resultaten beperkt zijn tot de schaal van 100M-160M parameters en dat opschaling naar 1B+ parameters noodzakelijk is om volledig te kunnen vergelijken met multi-token predictor autoregressieve modellen. Ze suggereren ook dat het token-tijdconcept potentieel uitgebreid kan worden naar uniforme noising discrete modellen in toekomstig werk.