Technisch Overzicht: Reward Transport – Eigenschapcontrole in Flow Matching via Ruisruimte-uitlijning

1. Probleemstelling

In Flow Matching (FM) en verwante diffusie-gebaseerde generatieve modellen wordt de koppeling—de regel die bepaalt welk ruisvector wordt gekoppeld aan welk datapunt tijdens de training—doorgaans behandeld als een computationele heuristiek. Standaard benaderingen gebruiken onafhankelijke koppeling voor eenvoud of minibatch optimal transport (OT) om trajecten te verstevigen en variantie te verminderen. In deze paradigma's wordt de koppeling uitsluitend geoptimaliseerd om training te faciliteren, niet om specifieke structurele eigenschappen in het geleerde flowveld te coderen.

Gevolgelijk vereist het bereiken van controleerbare generatie (bijv. het genereren van moleculen met specifieke logP- of QED-scores) meestal expliciete mechanismen tijdens de inferentie, zoals:

• Classifier-free guidance: Vereist een doelwaarde van een eigenschap als input.
• Gradient guidance: Het berekenen van gradiënten van een eigenschapvoorspeller bij elke bemonsteringsstap.
• Reinforcement Learning (RL): Het finetunen van de generator tegen een oracle, wat vaak oracle-queries tijdens de sampling vereist.

Deze methoden introduceren overhead tijdens de inferentie, afhankelijkheid van externe beloningsmodellen of oracles, of complexe trainingspipelines. Het artikel onderzoekt of de koppeling zelf kan dienen als een alignment interface, die controleerbare structuur direct in het flowveld injecteert tijdens de training, zodat controle beschikbaar is tijdens de inferentie zonder extra berekeningen of oracle-aanroepen.

2. Methodologie: Reward Transport

De auteurs stellen Reward Transport voor, een methode die de ruis-data koppeling herdefinieert als een mechanisme om taakstructuur in het flowveld te injecteren.

2.1 Eigenschap-uitgelijnde Optimal Transport Koppeling

In plaats van ruis en data onafhankelijk te koppelen, construeert Reward Transport een monotone koppeling gebaseerd op een doelmoleculaire eigenschap $y(x)$ (bijv. logP, QED).

1. Scalaire Sleutels: Voor elke ruisvector $z$ (een sequentie van token-embeddings) wordt een scalaire sorteersleutel $s(z)$ berekend. De primaire keuze is de Euclidische norm van de gemiddeld gepoolde ruis: $s(z) = \|\bar{z}\|_2$.
2. Ranggebaseerde Koppeling: Gegeven een batch ruisvectoren $\{z_n\}$ en moleculen $\{x_n\}$ met eigenschappen $\{y_n\}$, worden beide sets gesorteerd op hun respectieve sleutels ($s$ en $y$). Ze worden vervolgens rang-voor-rang gekoppeld: de ruisvector met de $k$-de kleinste $s$ wordt gekoppeld aan het molecuul met de $k$-de kleinste $y$.
3. Trainingsdoelstelling: Het flow matching-model $v_\theta$ wordt getraind om deze uitgelijnde paren te transporteren. De loss combineert een standaard MSE op de eindpuntvoorspelling (in continue embedding-ruimte) en een Cross-Entropy (CE) loss om ervoor te zorgen dat de gedecodeerde tokens geldige moleculen vormen.

2.2 Richting-conditionering

Om ervoor te zorgen dat het model de ruisruimte-coördinaat $s$ tijdens de inferentie kan benutten, wordt een Direction MLP geïntroduceerd. Deze twee-laagse MLP brengt de genormaliseerde scalaire $s$ in kaart naar een $d$-dimensionale embedding, die samen met de tijd- en positionele embeddings aan de input van de transformer wordt toegevoegd. Dit stelt het model in staat om zijn velocity field te conditioneren op de specifiek gebruikte ruiscoördinaat tijdens de monotone koppeling.

2.3 Inferentiemechanisme

Tijdens de inferentie specificeert de gebruiker een scalaire richting-signaal $s^*$ (een ruisruimte-coördinaat, geen doelwaarde van een eigenschap). Het model genereert moleculen door het flowveld te integreren beginnend vanaf ruisvectoren waar $s(z) \geq s^*$.

• Theoretische Garantie (Propositie 1): In het limiet waar het geleerde flow de rangcorrespondentie tussen $s$ en $y$ perfect behoudt (de coupling-preserving limit), induceert conditionering op $s \geq s^*$ een distributie die equivalent is aan één selectiestap van de Cross-Entropy Method (CEM). Specifiek truncheert het de gegenereerde distributie naar de top-$\epsilon$ fractie van de eigendistributie, waarbij $\epsilon$ wordt bepaald door de drempelwaarde $s^*$.
• Zero Overhead: Deze controle vereist geen oracle, geen beloningsmodel, geen gradiëntberekening en geen rejection sampling tijdens de inferentie.

3. Belangrijkste Bijdragen

1. Coupling als Alignment Interface: Het artikel herkadert de ruis-data koppeling in flow matching van een training-heuristiek naar een distributie-niveau alignment interface. Het demonstreert dat het kiezen van de koppeling op basis van een eigenschap de geleerde flow organiseert zodanig dat ruisruimte-coördinaten overeenkomen met eigenschapswaarden.
2. Single-Knob Moleculaire Sturing: De methode bereikt monotone controle over moleculaire eigenschappen (logP en QED) met behulp van een enkele scalaire ruiscoördinaat. Cruciaal is dat dezelfde knop tegenovergestelde structurele responsen produceert voor verschillende targets: het verhogen van $s$ genereert grotere moleculen voor logP (lipofiliciteit) maar kleinere moleculen voor QED (drug-likeness). Dit sluit een generieke grootte-bias uit en bevestigt dat de koppeling target-specifieke structuur in de flow schrijft.
3. Reward Transport voor Discrete Sequenties: De auteurs implementeren dit voor variabel lengte SELFIES token-sequenties, waarbij zij twee specifieke faalmodi identificeren en aanpakken in flow-modellen die discrete data genereren:
   • Impliciete Signaalvervaging: In Pre-LayerNorm architecturen wordt de magnitude-informatie (gecodeerd in $\|\bar{z}\|$) gewist. De Direction MLP lost dit op door een geleerd richtingpatroon te injecteren.
   • Length Confound: Zonder unmasked MSE loss op padding-tokens leert het model een shortcut waarbij een hoge $s$ simpelweg de sequentielengte vergroot om aan de norm-constraint te voldoen. Unmasked MSE dwingt het model om de ruis op padding-posities te regresseren, waardoor de magnitude-signaal verankerd wordt en deze shortcut wordt voorkomen.

4. Experimentele Resultaten

De methode werd geëvalueerd op de ZINC-250K en GuacaMol datasets.

• Controleprestaties:
  • logP: Bij het doorlopen van $s$ van -3 tot +7 op ZINC-250K verschoof de gemiddelde logP van 1.50 naar 5.44 (+137%), met een Spearman correlatie $\rho=1.000$ tussen $s$ en de gemiddelde eigenschap. De effectgrootte (Cohen's $d$) bereikte 1.63.
  • QED: Dezelfde knop verschoof QED monotoon (bijv. +0.059 bij $s=5$ op ZINC), met $\rho=1.000$.
  • Structurele Specificiteit: Voor logP verhoogde het verhogen van $s$ het aantal zware atomen (12 $\to$ 23). Voor QED verlaagde het de hoeveelheid zware atomen (23 $\to$ 16).
• Generalisatie: De resultaten werden gereproduceerd op GuacaMol (1.6M moleculen), waarbij nog sterkere effectgroottes voor logP (d=2.31) werden gezien en een consistente monotoniciteit voor QED binnen het geldige bereik.
• Vergelijking met Baselines:
  • Vergeleken met Conditional Flow Matching (waar eigenschapswaarden expliciete inputs zijn), bereikt Reward Transport een vergelijkbare monotoniciteit ($\rho=1.000$) maar opereert het via distributie-truncatie in plaats van expliciete target-hitting.
  • Vergeleken met $\epsilon$-prediction diffusion, laat de methode zien dat alignment op koppeling-niveau structureel afwezig is in $\epsilon$-prediction, omdat de Bayes-optimale target onafhankelijk is van de koppeling, wat het gradiëntsignaal verzwakt.
• Trainingsdynamiek: Rangbehoud (de correlatie tussen $s$ en gegenereerde $y$) ontstaat snel tijdens fine-tuning (binnen één epoch) en vlakt af onder de 1.0 (bijv. 0.57 voor logP), wat consistent is met de theoretische "distributie-benadering" kloof beschreven in Propositie 1.

5. Betekenis en Claims

Het artikel claimt dat Reward Transport de koppeling in flow matching vestigt als een onderscheidend, distributie-niveau controlekanaal dat op zichzelf alignment-structuur kan dragen.

• Efficiëntie: Het elimineert de noodzaak voor oracle-calls, beloningsmodellen of iteratieve optimalisatie tijdens de inferentie, wat de computationele kosten aanzienlijk verlaagt vergeleken met RL of guidance-gebaseerde methoden.
• Mechanisme: Het demonstreert dat de "stuurbaarheid" van generatieve modellen niet hoeft te vertrouchten op expliciete conditionering-inputs; het kan in de geometrie van het model worden ingebakken via de training-koppeling.
• Beperkingen en Omvang: De auteurs zijn bescheiden over de omvang. De methode biedt een 1D distributie-knop (het trunceren van de distributie) in plaats van exacte per-sample waarde-toewijzing. Het werkt goed voor scalaire eigenschappen die gecorreleerd zijn met globale moleculaire kenmerken (zoals grootte of polariteit), maar faalt voor topologische eigenschappen zoals Synthetic Accessibility (SA) waar een 1D scalaire sleutel complexe structuren van gelijke grootte niet kan onderscheiden. Verder is aangetoond dat de methode incompatibel is met $\epsilon$-prediction targets, wat benadrukt dat de interface afhankelijk is van de specifieke predictie-target van het netwerk.

Samenvattend verschuift dit werk het paradigma van controleerbare generatie van "het sturen van het traject tijdens de inferentie" naar "het uitlijnen van de trainingsdistributie", wat een principieel, zero-overhead mechanisme biedt voor eigenschapcontrole in flow matching.