Technische Samenvatting: Een Continuous-Time Markov Chain Framework voor Insertion Language Models

Probleemstelling

Autoregressieve taalmodellen (ARMs) genereren sequenties strikt van links naar rechts, wat hun flexibiliteit beperkt voor taken die niet-sequentiële generatie vereisen, zoals planning en sub-doel gebaseerde infilling. Masked Diffusion Models (MDMs) maken generatie in willekeurige volgorde en met multi-token stappen mogelijk, maar hebben vaak moeite met sequenties van variabele lengte omdat ze padding-tokens behandelen als gewone vocabulaire-items, wat de kwaliteit van de generatie beïnvloedt en sampling beperkt tot blokgebaseerde procedures van links naar rechts.

Bestaande Insertion Language Models (ILMs) pakken deze beperkingen aan door iteratief tokens te invoegen op willekeurige posities, waarbij variabele lengtes en relatieve ordeningsbeperkingen op natuurlijke wijze worden afgehandeld. De huidige ILM-formuleringen vertrouwen echter op ad-hoc trainingsdoelen en samplingprocedures die een rigoureuze probabilistische basis missen. Bov�ens hebben bestaande ILMs vaak moeite met de uitdaging om te bepalen wanneer de generatie moet stoppen.

Methodologie

De auteurs stellen Diffusion-based Insertion Language Models (DILMs) voor, een principieel framework afgeleid van eerste principes met behulp van een Continuous-Time Markov Chain (CTMC) op de ruimte van variabelen-lengte sequenties.

1. Het Noising Proces

Het forward noising proces wordt geformuleerd als een CTMC die tokens verwijdert uit een grondwaarheid-sequentie $x_0$ uniform willekeurig over een tijdsinterval $[0, T]$.

• Toestandsruimte: De toestandsruimte bevat zowel de sequentie-inhoud als de lengte ervan.
• Verwijderingsmechanisme: Tokens worden één voor één verwijderd (of via onafhankelijke rates) met een tijdsafhankelijke rate $\sigma_t$.
• Twee Varianten:
  • Joint Noising: Verwijderingen vinden plaats als een enkel event waarbij het aantal verwijderde tokens een Poisson-distributie volgt, en posities worden uniform gekozen.
  • Independent Noising: Elke positie heeft een onafhankelijke verwijderingsrate, gemodelleerd via een Binomial-distributie.

2. Het Generatieve Proces (Reverse CTMC)

Het generatieve proces is de tijdsomkering van het noising proces, waarbij een sequentie wordt geconstrueerd door tokens in te voegen. De auteurs leiden een Evidence Lower Bound (ELBO) af voor de data log-likelihood onder dit reverse proces. Ze stellen twee parametrisaties voor voor de reverse rate matrix $\hat{R}^\theta_t$:

• DILM-S (Single Insertion): Modelleert de gezamenlijke waarschijnlijkheid van een insertie-locatie en een vocabulaire-item, evenals de "length-to-go" (het aantal tokens dat nog ingevoegd moet worden). Deze parametrisatie voorspelt één insertie per keer.
• DILM-M (Multiple Insertion): Modelleert onafhankelijke, per-positie insertie-rates. Dit stelt het model in staat om meerdere inserties simultaan te samplen over verschillende gaten in een enkele stap.

Beide varianten maken gebruik van een transformer-backbone met adaptive layer normalization (AdaLN) die geconditioneerd is op tijd $t$. Om de stopconditie af te handelen, voorspellen de modellen de "length-to-go" (of gap-groottes), waardoor de generatie kan eindigen wanneer de voorspelde insertie-rate naar nul daalt, in plaats van te vertrouwen op een aparte stopping classifier of hyperparameters.

3. Sampling

Sampling wordt uitgevoerd via Euler discretisatie (voor DILM-S) of een geconstreerde $\tau$-leaping procedure (voor DILM-M). Deze methoden maken een afweging mogelijk tussen samplekwaliteit en het aantal forward passes (sampling stappen), een flexibiliteit die niet aanwezig is bij standaard ILMs.

Belangrijkste Bijdragen

1. Principieel Framework: Het artikel leidt een diffusion-stijl denoising doel af voor insertie-gebaseerde generatie vanuit eerste principes met behulp van een CTMC-framework, waarmee het verder gaat dan ad-hoc formuleringen.
2. Unificatie: De voorgestelde samplingprocedures verenigen bestaande formuleringen van ILMs, door aan te tonen dat zij speciale gevallen zijn van dit bredere framework.
3. Twee Parametrisaties: De auteurs introduceren twee varianten (DILM-S en DILM-M) die verschillen in hoe ze insertie-rates modelleren (joint versus onafhankelijk), wat flexibiliteit biedt in sampling-efficiëntie en multi-token generatie.
4. Theoretische Onderbouwing: Het werk biedt een rigoureuze afleiding van de ELBO en bespreekt computationele overwegingen voor training en sampling.

Empirische Resultaten

De auteurs evalueren DILMs op synthetische planningstaken en taalmodellering benchmarks.

• Synthetische Planning (Star Graph Taken): Op variabelen-lengte planningstaken waar ARMs en MDMs moeite mee hebben, bereiken DILM-S en DILM-M een bijna perfecte nauwkeurigheid (vergelijkbaar met bestaande ILMs), wat aantoont dat de diffusion-gebaseerde aanpak de voordelen van insertie-gebaseerde generatie behoudt voor complexe structurele taken.
• Taalmodellering (LM1B en OpenWebText):
  • Op de LM1B benchmark zijn DILM-S en DILM-M competitief met bestaande ILMs en MDMs in termen van per-token negative log-likelihood (NLL).
  • Op OpenWebText presteren DILM varianten beter dan zowel ILM als MDM baselines.
  • Sampling Flexibiliteit: Een belangrijke bevinding is de mogelijkheid om generatiekwaliteit af te wegen tegen het aantal sampling stappen. In tegen tegenstelling tot bestaande ILMs, die vlakke of grillige perplexity-curves vertonen, vertonen DILMs een vloeiende verbetering in generatieve perplexiteit naarmate het aantal sampling stappen toeneemt.

Betekenis en Claims

Het artikel claimt dat DILMs een probabilistisch geprincipeerd alternatief bieden voor bestaande insertiemodellen, door een rigoureus trainingsdoel en een samplingprocedure te bieden. Het framework behoudt de kernvoordelen van ILMs — flexibele generatievolgorde en natuurlijke afhandeling van variabelen-lengte — terwijl het de beperkingen van ad-hoc stopcriteria en het gebrek aan theoretische onderbouwing aanpakt.

De auteurs merken op dat hoewel DILM-M meerdere inserties toestaat, het momenteel beperkt is tot één token per gat, wat de modellering van opeenvolgende token-correlaties beperkt. Zij identificeren multi-token span insertie en het opschalen naar grotere modellen als richtingen voor toekomstig werk. Het werk legt een fundament voor het behandelen van insertie-gebaseerde generatie als een continuous-time diffusion proces, waarbij eerdere heuristische benaderingen onder één enkele theoretische paraplu worden verenigd.