Technische Samenvatting: Latent-Kernel Discrete Flow Maps voor Few-Step Generatie

Probleemstelling

Discrete diffusie- en flow-matchingmodellen genereren tekst door een sequentie over vele stappen te denoisen. Om de computationele efficiëntie te behouden, factoriseren deze modellen de transitiekans over tokenposities doorgaans, waarbij ze elke token onafhankelijk beslissen op basis van de huidige sequentietoestand. Hoewel deze factorisatie individuele stappen goedkoop maakt, creëert het een fundamentele bottleneck voor few-step generatie.

Wanneer een doelsequentie sterke correlaties tussen posities vereist (bijv. onderwerp-werkwoordovereenkomst of bijpassende haakjes), leiden onafhankelijke updates tot het vastleggen van tokens afzonderlijk. Dit resulteert vaak in lokaal consistente maar globaal inconsistente combinaties (bijv. "dog barks" vs. "dogs bark"). Het corrigeren van deze mismatches vereist veel extra denoising-stappen, wat de efficiëntievoordelen van few-step generatie tenietdoet.

Bestaande oplossingen om generatie te versnellen vallen uiteen in twee categorieën:

1. Likelihood-modellen (bijv. MDLM, SEDD): Trainen sterke gefactoriseerde denoisers vanaf nul, maar vereisen veel stappen tijdens inferentie om correlaties op te lossen.
2. Few-step accelerators (bijv. Di4C, SDTT, ReDi): Destilleren of rectificeren een getraind "teacher"-model naar een few-step student. Deze methoden erven echter de kwaliteitscel van de teacher en kunnen de prestaties van het model dat ze imiteren niet overtreffen.

De kernuitdaging is of een model gecorreleerde stappen inherent kan uitdrukken zonder afhankelijk te zijn van een teacher, waardoor de factorisatie-aanname bij de wortel wordt doorbroken terwijl de lage computationele kosten behouden blijven.

Methodologie: Latent-Kernel Discrete Flow Maps (LKF)

De auteurs stellen Latent-Kernel Discrete Flow Maps (LKF) voor, een from-scratch flow-map model dat intra-step correlatie introduceert via een mengeling van gefactoriseerde componenten die verbonden zijn door één enkele gedeelde latente variabele.

De Modelarchitectuur

De eindige-tijd transitiekernel $\Psi_{t \to s}(x_s | x_t)$ wordt gedefinieerd als een mengeling van $M$ gefactoriseerde componenten:
$$\Psi_{t \to s}(x_s | x_t) = \sum_{k=1}^{M} w_k(x_t, t, s) \prod_{i=1}^{L} P_i(x_s^i | x_t, k, t, s)$$
Waar:

• $k \in [M]$ een gedeelde latente componentindex is die één keer per sequentie wordt gesampled.
• $w_k$ router-gewichten zijn die de waarschijnlijkheid van component $k$ voorspellen.
• $P_i$ gefactoriseerde tokenverdelingen zijn voor positie $i$, geconditioneerd op de latente $k$.

Kernmechanisme:

1. Gedeelde Latente Variabele: Aan het begin van een denoising-traject (of elke stap) wordt een enkele latente $k$ getrokken uit de router $w(\cdot | x_t)$.
2. Conditionele Onafhankelijkheid: Gegeven $k$, worden tokens nog steeds onafhankelijk over posities gesampled. Omdat $k$ echter gedeeld wordt over de gehele sequentie, induceert het marginaliseren over $k$ correlaties tussen posities.
3. Efficiëntie: Het model gebruikt een late latent injection architectuur. Een gedeelde transformer-trunk verwerkt de input om een latent-agnostische representatie te creëren. Alleen de laatste $L_k$ blokken worden gerepliceerd voor elk van de $M$ componenten, geconditioneerd op $k$. Dit houdt de computationele kosten van marginalisering laag ($O(M \cdot L_k)$ in plaats van $O(M \cdot \text{depth})$).

Training en Inferentie

• Training: Het model wordt getraind door de latente $k$ exact te marginaliseren. De doelstelling is de log-likelihood van de mengeling:
  $$\log \Psi = \log \sum_{k} \exp(\log w_k + \ell_k)$$
  waarbij $\ell_k$ de log-waarschijnlijkheid is die door component $k$ aan de doeltokens wordt toegewezen. Dit dwingt de router om een mengeling te leren waarbij de latente variabele sequentie-brede correlaties vastlegt. Een router-regularisator zorgt ervoor dat de mengelingcomponenten worden benut (om collaps naar een enkele component te voorkomen) en dat de latente variabele voldoende informatie draagt.
• Inferentie (Best-of-M): Om de heterogeniteit van de mengeling te benutten, passen de auteurs een Best-of-M decoding strategie toe. $M$ parallelle rollouts worden gegenereerd, waarbij elke rollout een andere component $k$ bevriest. De uiteindelijke output wordt geselecteerd op basis van de eigen sequentie-niveau mengeling negative ELBO (NELBO) van het model. Deze selectie vereist geen extern reward-model of teacher.

Belangrijkste Bijdragen

1. Inherent Gecorreleerde Stappen: Introductie van LKF, een flow-map kernel die gecorreleerde stappen inherent uitdrukt via een mengeling van $M$ gefactoriseerde componenten verbonden door een enkele gedeelde latente variabele, waardoor de noodzaak voor een teacher-model wordt geëlimineerd.
2. Theoretische Analyse: Een exacte informatie-decompositie die laat zien dat de totale correlatie (TC) die een rank-$M$ stap kan dragen, begrensd wordt door $(L-1)H(k|xt)$. De auteurs bewijzen dat de benaderingsfout van het model ondergrensd wordt door de niet-gevangen totale correlatie van de doelverdeling.
3. Synthetische Validatie: Experimenten op synthetische datasets (hidden-agreement en parity) laten zien dat LKF informatie vastlegt en de likelihood monotoon verbetert met $M$ voor representeerbare correlaties, terwijl het correct faalt om correlaties te vatten die een exponentiële mengeling-rank vereisen (parity), wat overeenkomt met theoretische hardheidsresultaten.
4. Empirische Prestaties: Op de LM1B en WikiText-103 benchmarks bereikt LKF state-of-the-art few-step generatieve perplexiteit. Bij $M=8$ verbetert LKF de generatieve perplexiteit met 2.1× tot 3.3× ten opzichte van likelihood-baselines (MDLM, SEDD) en overtreft het gedestilleerde/gerectificeerde accelerators (Di4C, ReDi) zonder een teacher te gebruiken.

Experimentele Resultaten

• Benchmarks: Geëvalueerd op LM1B en WikiText-103.
• Prestaties:
  • LM1B: Het $M=8$ LKF-model bereikt een generatieve perplexiteit (gen-PPL) van 105 bij 32 function evaluations (NFE), vergeleken met 304 voor MDLM en 117 voor de beste gerectificeerde accelerator (ReDi).
  • WikiText-103: Het $M=8$ model bereikt een gen-PPL van 75, wat aanzienlijk beter is dan MDLM (228) en SEDD (247).
• Diversiteit: In tegen tegenstelling tot gerectificeerde accelerators, die vaak lijden aan entropie-collaps (het produceren van teksten met lage diversiteit), behoudt LKF een hoge sample-entropie terwijl de likelihood wordt verbeterd.
• Ablaties:
  • De winst wordt toegeschreven aan de mengelingstructuur zelf, niet alleen aan reranking. Het vervangen van de $M$ componenten door $M$ ruis-samples van een enkele component resulteert in een aanzienlijke prestatiedaling (bijv. 35% hogere PPL bij $M=8$).
  • Best-of-M selectie is cruciaal; een single-commitment policy (commit-M) levert resultaten op die vergelijkbaar zijn met de baseline, wat bevestigt dat de mengeling de nodige kandidaten-diversiteit biedt die selectie exploiteert.
  • Router Efficiëntie: De router fungeert als een menggewicht in plaats van een kwaliteitsignaal; de selectie is het meest effectief wanneer deze gebaseerd is op de eigen likelihood (running evidence) van het model, in plaats van de initiële voorspelling van de router.

Betekenis en Claims

Het artikel claimt dat LKF een verschuiving vertegenwoordigt van "teacher-afhankelijke" versnelling naar "from-scratch" few-step generatie. Door de transitiekernel als een low-rank mengeling te modelleren, stelt LKF een enkele stap in staat om massa te plaatsen op gecorreleerde voltooiingen (bijv. onderwerp-werkwoordovereenkomst) met dezelfde sampling-complexiteit als een gefactoriseerd model.

De auteurs benadrukken dat:

• De methode de kwaliteitscel van een teacher niet overneemt.
• De winst schaalt met de grootte van de mengeling $M$.
• De aanpak theoretisch onderbouwd is, met grenzen aan de correlatiecapaciteit en benaderingsfout.
• Het de beste few-step generatieve perplexiteit bereikt onder de baseline-methoden op standaard taalmodellering-benchmarks, wat een levensvatbaar pad biedt voor efficiënte, hoogwaardige discrete generatie zonder de overhead van distillatie of rectificatie.