Technische Samenvatting: Variational Learning voor Insertie-gebaseerde Generatie

1. Probleemstelling

Autoregressieve sequentiemodellen genereren doorgaans tokens in een vaste links-naar-rechts volgorde. Hoewel dit effectief is voor veel taken, is deze factorisatie niet inherent aan alle datastructuren en kan het niet goed aansluiten bij realistische generatieproblemen zoals biologisch sequentieontwerp, programmasynthese en gestructureerde planning, waarbij onderdelen van de output afhankelijk kunnen zijn van toekomstige context.

Bestaande alternatieven, zoals gemaskeerde diffusiemodellen (niet-monotone generatie), bieden flexibiliteit door subsets van gemaskeerde posities op een vaste rasterstructuur bij te werken. Deze benaderingen lijden echter onder twee primaire beperkingen:

1. Statisch Canvas: Ze opereren op voorgeschreven, vaste rasters, wat externe lengtespecificaties vereist en het genereren van variabele lengtes indirect maakt.
2. Orde-agnosticisme: Hun trainingsdoelstellingen marginaliseren vaak over of gaan impliciet uit van vele mogievelijke update-schema's. Dit dwingt het model om conditionele kansen te benaderen voor een factorial aantal generatieordes, zelfs als slechts een klein deel daarvan overeenkomt met de onderliggende structuur van de data.

Recente werken hebben deze problemen afzonderlijk aangepakt: insertie-gebaseerde modellen maken generatie met variabele lengte mogelijk via willekeurige inserties, terwijl methoden voor het leren van de orde distributies over permutaties introduceren. Het combineren hiervan is echter niet triviaal omdat in insertiemodellen de insertieposities geen vaste coördinaten zijn, maar meebewegen met de gedeeltelijke sequentie, waardoor de generatievolgorde en tussenliggende toestanden nauw aan elkaar gekoppeld zijn.

2. Methodologie: Het Insertieproces (IP)

De auteurs stellen het Insertion Process (IP) voor, een probabilistisch raamwerk dat gezamenlijk leert:

1. Waar te insert: Een datagedreven voorkeur voor insertiebeslissingen (slot selectie).
2. Wat te insert: Een tokenverdeling.
3. Wanneer te termineren: Een stopregel (beleidsgebaseerd of classifier-gebaseerd).

2.1. Reparametrisatie via Permutaties

Een kernbijdrage aan de theoretische vorming is de formalisering van een bijectieve correspondentie tussen insertietrajecten en permutaties van de doelsequentie-indices.

• De Uitdaging: Directe variationele inferentie over insertietrajecten is moeilijk omdat relatieve slot-indices ($z_i$) trajectafhankelijk zijn, wat het lastig maakt om een stabiele posterior over de orde te definiëren in de ruimte van relatieve acties.
• De Oplossing: De auteurs mappen elk geldig insertietraject dat een sequentie $y_L$ produceert naar een unieke permutatie $\sigma \in S_L$, waarbij $\sigma_i$ de index is van het element dat in stap $i$ wordt ingevuld.
• Deterministische Mapping: De relatieve insertie-slot bij stap $i$ wordt een deterministische functie van de permutatie-prefix ($\sigma_{\le i}$). Specifiek: $z_i = 1 + |\{j \in \sigma_{<i} : j < \sigma_i\}|$.
• Likelihood Factorisatie: Deze verandering van variabelen zet de marginalisatie over trajectafhankelijke acties om in een exacte som over permutaties, wat een exacte reparametrisatie van de data-likelihood mogelijk maakt:
  $$p(y_L) = \sum_{\sigma \in S_L} p(y_L, \sigma)$$

2.2. Variationele Inferentie en Training

Om het model te trainen, maximaliseren de auteurs een Evidence Lower Bound (ELBO) op de log-likelihood van de data, waarbij de generatievolgorde (permutatie $\sigma$) als een latente variabele wordt behandeld.

• Variational Posterior: Ze introduceren een geamortiseerde variationele posterior $q_\theta(\sigma | y_L)$, geparametriseerd als een Plackett-Luce (PL) model.
• Optimalisatie: Omdat de latente orde discreet is, is de doelstelling niet-differentieerbaar. De auteurs optimaliseren de ELBO met behulp van:
  • Gumbel-Top-k Trick: Om permutaties efficiënt te samplen uit de PL-distributie.
  • REINFORCE met Leave-One-Out (RLOO): Om gradiënten te schatten.
  • Variancereductie: Door de permutatie-trajectstructuur te exploiteren, wordt de binnenste verwachting over de volgende index $\sigma_i$ (geconditioneerd op een vaste prefix) exact in gesloten vorm geëvalueerd, wat de variantie van de Monte Carlo-estimator vermindert.

2.3. Netwerkarchitectuur

• Generatieve Decoder: Een Transformer-gebaseerde decoder verwerkt de gedeeltelijke sequentie (aangevuld met boundary tokens) om contextuele embeddings te produceren.
  • Een location head scoort kandidaat-insertie-slots (inclusief een terminatie-actie) op basis van de embeddings.
  • Een content head voorspelt de tokenverdeling geconditioneerd op de geselecteerde slot.
• Inference Encoder: Een bidirectionale Transformer encoder verwerkt de volledige geobserveerde sequentie om scores voor elke globale index te produceren, die de Plackett-Luce distributie parametriseren voor het samplen van de generatievolgorde.

3. Belangrijkste Bijdragen

1. Probabilistisch Raamwerk voor Variabele-Lengte Insertie: Het eerste raamwerk dat de insertievolgorde voor variabelen-lengte sequenties leert met behulp van een bijectieve mapping naar permutaties, wat een exacte likelihood-reparametrisatie mogelijk maakt.
2. Permutatie-gemarginaliseerde Likelihood: Een theoretische afleiding die aantoont dat de data-likelihood kan worden uitgedrukt als een som over permutaties, wat variatiele inferentie met een latente orde-variabele faciliteert.
3. Efficiënt Trainingsalgoritme: Een variancereduceerde trainingsmethode met behulp van Plackett-Luce modellen, Gumbel-Top-k sampling en RLOO-gradiënten om de discrete, niet-differentieerbare aard van de latente orde aan te pakken.
4. Natuurlijke Ondersteuning voor Variabele Lengte: In tegenstelling tot fixed-canvas diffusiemodellen ondersteunt IP van nature generatie met variabele lengte en leert het datagedreven voorkeuren voor insertievolgordes.

4. Experimentele Resultaten

De auteurs evalueerden IP op twee domeinen: doel-geconditioneerde planning en moleculaire stringgeneratie.

4.1. Doel-geconditioneerde Planning

• Benchmarks: Maze planning (Braided, Imperfect, Perfect) en Star Graph planning.
• Baselines: Fixed-Order ARMs (FO-ARM), Any-Order ARMs (AO-ARM), Masked Diffusion Models (MDM) en Learning-Order ARMs (LO-ARM).
• Resultaten: IP bereikte state-of-the-art prestaties over alle maze-families en graaf-benchmarks.
  • Op Perfect Mazes (hoogst beperkt), presteerden FO-ARMs aanzienlijk slechter, terwijl IP een bijna perfecte nauwkeurigheid behield (bijv. 99,9% op Hard Perfect Mazes versus 27,2% voor FO-ARM).
  • IP verbeterde de sequentie- en token-nauwkeurigheid aanzienlijk en verminderde de edit distances vergeleken met alle baselines, wat aantoont dat het leren van instantie-specifieke insertievolgordes cruciaal is voor taken met langetermijnbeperkingen.

4.2. Moleculaire String (SMILES) Generatie

• Benchmark: GuacaMol onconditionele en conditionele generatie.
• Resultaten:
  • Validiteit & Nieuwheid: IP bereikte een hoge validiteit (97,2–97,4%) en de hoogste nieuwheid (~95%) van alle modellen, waarmee het random-order insertie (AO-IP) die naar 88,2% validiteit daalde, significant versloeg. Dit geeft aan dat IP de chemische syntaxisregels succesvol heeft geïnternaliseerd.
  • Distributionele Metrieken: Hoewel FO-ARMs iets beter presteerden op distributionele gelijkenis (KL divergentie, FCD), bood IP een superieure afweging door de chemische ruimte effectiever te verkennen.
  • Conditionele Generatie: In taken die insertie tussen vaste fragmenten vereisen (bijv. linker-ontwerp), behield IP ~100% validiteit, terwijl fixed-order modellen moeite hebben met dergelijke beperkingen.
  • Interpretatieerbaarheid: Analyse van de geleerde insertie-schema's onthulde een interpreteerbare, consistente generatievolgorde: (1) het bouwen van het moleculaire skelet (ringen/haakjes), (2) het inserten van atomen, en (3) het termineren. Dit gebeurde zonder expliciete grammaticale priors.

5. Betekenis en Claims

Het artikel beweert dat het leren van datagedependente insertievolgordes de modelkwaliteit en generalisatie aanzienlijk verbetert in domeinen die geen canonieke links-naar-rechts structuur hebben. Door de verbinding tussen insertietrajecten en permutaties te formaliseren, bieden de auteurs een hanteerbare, onbevooroordeelde variationele methode voor het trainen van insertie-gebaseerde modellen.

De auteurs positioneren hun werk als een stap richting generatieve modellering in domeinen waar de "natuurlijke" generatievolgorde ambigu of niet-lineair is. Ze geven expliciet aan dat toekomstig werk deze methode zou kunnen toepassen op andere domeinen zoals eiwitontwerp. Het artikel behoudt een bescheiden toon door op te merken dat hoewel IP de generalisatie en validiteit in complexe structurele taken verbetert, het niet noodzakelijkerwijs beter presteert dan fixed-order modellen op eenvoudige metrieken voor distributionele gelijkenis, wat wijst op een afweging tussen flexibiliteit en strikte distributionele getrouwheid.