Tensor-Train Joint Modeling for Few-Step Discrete Diffusion

Technische Samenvatting: Tensor-Train Joint Modeling voor Few-Step Masked Diffusion

Probleemstelling
Discrete diffusiemodellen, in het bijzonder Masked Diffusion Models (MDM's), bieden een veelbelovend alternatief voor autoregressieve (AR) modellen voor het genereren van sequentiële discrete data (bijv. tekst, moleculen) door parallelle token-generatie mogelijk te maken. Echter, hun potentieel voor generatie in enkele stappen (few-step generation) blijft onverløten vanwege een fundamentele structurele beperking: de aanname van conditionele onafhankelijkheid. Huidige MDM's modelleren de conditionele zuivere distributie $p_\theta(x|x_t)$ als een product van onafhankelijke marginalen ($\prod p_\theta(x_i|x_t)$). Deze aanname introduceert een systematische "parallellisatie-bias" die cumuleert met het aantal tokens dat per stap wordt ontmaskerd (unmasked). In de few-step regimes, waar veel tokens tegelijkertijd worden ontmaskerd, degradeert deze bias de kwaliteit van de samples ernstig.

Het expliciet modelleren van de gezamenlijke distributie (joint distribution) $p_\theta(x|x_t)$ om deze bias op te lossen is theoretisch wenselijk, maar computationeel onhandelbaar, omdat het vereist is een $N$-dimensionale tensor met $V^N$ entries te representeren (waarbij $V$ de vocabulairegrootte is en $N$ de sequentielengte). Bestaande benaderingen proberen dit te omzeilen door gebruik te maken van hulp-AR-modellen of latente variabelen, maar deze erven ofwel externe inferentiekosten, ofwel vertrouwen op impliciete correlaties in plaats van expliciete gezamenlijke parametrisatie.

Methodologie
De auteurs stellen een framework voor dat de gezamenlijke distributie $p_\theta(x|x_t)$ expliciet parametriseert binnen het discrete diffusiemodel zelf met behulp van tensor-decompositie. De kernmethodologie bestaat uit twee hoofdcomponenten:

1. Tensor-decompositie voor Joint Modeling:
   Het framework representeert de conditionele distributie als een low-rank tensor, waarbij ondersteuning wordt geboden aan twee specifieke decomposities:

• Canonical Polyadic Decomposition (CPD): Benadert de tensor als een som van rank-1 tensoren. Hoewel dit de standaard MDM's generaliseert (het rank-1 geval), behandelt het alle tokenposities symmetrisch.
• Tensor-Train Decomposition (TTD): Benadert de tensor als een sequentie van cores (matrices) die verbonden zijn in een keten. De auteurs identificeren een cruciale structurele bias in TTD: de TT-rank op elk splitsingspunt komt overeen met de rank van de unfolding-matrix van de gezamenlijke distributie. Voor sequentiële data die gedomineerd worden door lokale afhankelijkheden (bijv. natuurlijke taal, moleculaire line-notaties), blijven deze unfolding-ranks klein, waardoor TTD de afhankelijkheden tussen nabijgelegen tokens efficiënt kan vangen met een lage rank. Dit is theoretisch onderbouwd door de stelling van Oseledets.

2. Efficiënte Sampling via Iteratieve Marginale Inferentie:
   Directe sampling uit de gezamenlijke distributie is onhandelbaar. De auteurs presenteren een kettingregel-gebaseerde samplingprocedure die iteratieve marginale inferentie uitvoert.

• Algemene Procedure: De methode samplet tokens één voor één (of in batches) door te conditioneren op eerder gesamplede posities. Het maakt gebruik van caching en parallelle prefix sums om marginale waarschijnlijkheden efficiënt te berekenen zonder de volledige gezamenlijke distributie te evalueren.
• Voorafbepaalde Schema's: Voor vaste unmasking-schema's optimaliseert de methode de inferentie door ongeconditioneerde cores te contracteren en een hulp-head te gebruiken om gecontracteerde cores direct te voorspellen, waardoor redundante berekeningen worden vermeden.

3. Integratie en Fine-Tuning:
   Het framework is ontworpen om via lichtgewicht fine-tuning te integreren in gepre-trainde MDM's. Architecturale wijzigingen omvatten het vervangen van de uiteindelijke output-head door een tensor-gedecomponeerde structuur. Om de marginale voorspellingen van het gepre-trainde model te behouden, worden de gewichten geïnitialiseerd als kopieën van de originele head (met een kleine hoeveelheid ruis om symmetrie te breken), waardoor het model gezamenlijke afhankelijkheden kan leren met minimale kosten vergeleken met training vanaf nul.

Belangrijkste Bijdragen

• Eerste Expliciete Joint Modeling in Discrete Diffusion: Dit werk introduceert het eerste framework dat de conditionele zuivere distributie van discrete diffusie parametriseert via tensor-decompositie, wat de standaard MDM's strikt generaliseert als het rank-één geval.
• Structurele Bias van TTD voor Sequentiële Data: De auteurs formaliseren de geschiktheid van Tensor-Train Decomposition voor sequentiële data. Ze demonstreren dat de structuur van TTD aansluit bij de lokaliteit van afhankelijkheden in natuurlijke taal en moleculaire notaties, waarbij TTD beter presteert dan CPD in deze domeinen.
• Efficiënt Sampling Algoritme: Een nieuw iteratief marginaal inferentieproces wordt gepresenteerd, wat een hanteerbare sampling uit de gezamenlijke distributie mogelijk maakt met minimale overhead door middel van caching en parallelle prefix sums.
• Lichtgewicht Fine-Tuning: Er wordt een recept geboden voor het integreren van dit framework in gepre-trainde modellen, wat substantiële verbeteringen in few-step regimes oplevert tegen een fractie van de kosten van training vanaf nul.

Experimentele Resultaten
Het framework werd geëvalueerd op tekstgeneratie (OpenWebText, LM1B) en molecuulgeneratie (QM9) met behulp van de basismodellen MDLM en VADD.

• Tekstgeneratie: TTD-gebaseerde fine-tuning verminderde de generatieve perplexiteit aanzienlijk ten opzien van vanilla modellen, met name in het few-step regime. Bijvoorbeeld, op OpenWebText verminderde TTD de 8-step perplexiteit met 32,7% ten opzichte van vanilla VADD. CPD vertoonde een bescheiden of geen verbetering ten opzichte van de baselines.
• Molecuulgeneratie: Op QM9 behaalde de TTD-aanpak consistent hogere scores voor validiteit, uniciteit en nieuwheid vergeleken met de baselines, waarbij de grootste winsten werden waargenomen bij links-naar-rechts generatie waar lokale afhankelijkheden cruciaal zijn.
• Efficiëntie: Ondanks de theoretische complexiteit was de sampling-overhead minimaal. Op OpenWebText was de TTD-verbeterde VADD slechts 1,7% trager dan de originele VADD bij 128 timesteps.

Betekenis en Claims
Het paper claimt de eerste te zijn die succesvol tensor-decompositie heeft geïntroduceerd voor expliciete gezamenlijke waarschijnlijkheidsmodellering binnen discrete diffusiemodellen. Er wordt betoogd dat door verder te gaan dan de aanname van conditionele onafhankelijkheid via een hanteerbare low-rank parametrisatie, de systematische parallellisatie-bias kan worden gemitigeerd. De auteurs benadrukken dat TTD bijzonder geschikt is voor sequentiële data vanwege de theoretische aansluiting bij lokale token-afhankelijkheden, wat een praktisch pad biedt naar hoogwaardige, few-step generatie zonder te vertround op externe autoregressieve modellen. Het werk demonstreert dat zelfs low-rank benaderingen voldoende zijn om de noodzakelijke gezamenlijke structuur te vangen voor significante prestatiewinsten in few-step regimes.