Riepilogo Tecnico: Reward Transport – Controllo delle Proprietà nel Flow Matching tramite Allineamento nello Spazio del Rumore

1. Definizione del Problema

Nel Flow Matching (FM) e nei modelli generativi correlati basati sulla diffusione, l'accoppiamento (coupling)—ovvero la regola che determina quale vettore di rumore viene associato a quale punto dei dati durante l'addestramento—viene tipicamente trattato come un'euristica computazionale. Gli approcci standard utilizzano un accoppiamento indipendente per semplicità o l'Optimal Transport (OT) su minibatch per raddrizzare le traiettorie e ridurre la varianza. In questi paradigmi, l'accoppiamento è ottimizzato esclusivamente per facilitare l'addestramento, non per codificare proprietà strutturali specifiche nel campo di flusso appreso.

Di conseguenza, ottenere una generazione controllabile (ad esempio, generare molecole con specifici punteggi di logP o QED) richiede solitamente meccanismi espliciti al momento dell'inferenza, quali:

• Classifier-free guidance: che richiede il valore di una proprietà target come input.
• Gradient guidance: che calcola i gradienti di un predittore di proprietà ad ogni passo di campionamento.
• Reinforcement Learning (RL): che affina il generatore contro un oracle, richiedendo spesso query all'oracle durante il campionamento.

Questi metodi introducono un overhead all'inferenza, la dipendenza da modelli di ricompensa o oracle esterni, o pipeline di addestramento complesse. Il paper si chiede se l'accoppiamento stesso possa servire come un interfaccia di allineamento, incorporando la struttura controllabile direttamente nel campo di flusso durante l'addestramento, in modo che il controllo sia disponibile all'inferenza senza ulteriori computazioni o chiamate all'oracle.

2. Metodologia: Reward Transport

Gli autori propongono il Reward Transport, un metodo che ridefinisce l'accoppiamento rumore-dati come un meccanismo per iniettare la struttura del task nel campo di flusso.

2.1 Accoppiamento di Optimal Transport Allineato alle Proprietà

Invece di associare rumore e dati indipendentemente, il Reward Transport costruisce un accoppiamento monotono basato su una proprietà molecolare target $y(x)$ (ad esempio, logP, QED).

1. Chiavi Scalari: Per ogni vettore di rumore $z$ (una sequenza di embedding di token), viene calcolata una chiave di ordinamento scalare $s(z)$. La scelta primaria è la norma euclidea del mean-pooling del rumore: $s(z) = \|\bar{z}\|_2$.
2. Associazione basata sul Rango: Dato un batch di vettori di rumore $\{z_n\}$ e molecole $\{x_n\}$ con proprietà $\{y_n\}$, entrambi i set sono ordinati in base alle rispettive chiavi ($s$ e $y$). Vengono quindi accoppiati rango per rango: il vettore di rumore con il $k$-esimo $s$ più piccolo è accoppiato con la molecola che ha il $k$-esimo $y$ più piccolo.
3. Obiettivo di Addestramento: Il modello di flow matching $v_\theta$ è addestrato per trasportare questi accoppiamenti allineati. La perdita combina un MSE standard sulla predizione dell'endpoint (in spazio di embedding continuo) e una perdita di Cross-Entropy (CE) per garantire che i token decodificati formino molecole valide.

2.2 Condizionamento della Direzione

Per garantire che il modello possa accedere alla coordinata dello spazio del rumore $s$ durante l'inferenza, viene introdotto un Direction MLP. Questo MLP a due strati mappa lo scalare normalizzato $s$ in un embedding a $d$ dimensioni, che viene aggiunto all'input del transformer insieme agli embedding temporali e posizionali. Ciò consente al modello di condizionare il proprio campo di velocità alla specifica coordinata del rumore utilizzata durante l'accoppiamento monotono.

2.3 Meccanismo di Inferenza

All'inferenza, l'utente specifica un segnale di direzione scalare $s^*$ (una coordinata dello spazio del rumore, non un valore di proprietà target). Il modello genera molecole integrando il campo di flusso partendo da vettori di rumore dove $s(z) \geq s^*$.

• Garanzia Teorica (Proposizione 1): Nel limite in cui il flusso appreso preserva perfettamente la corrispondenza del rango tra $s$ e $y$ (limite di conservazione dell'accoppiamento), il condizionamento su $s \geq s^*$ induce una distribuzione equivalente a un passo di selezione del Metodo della Cross-Entropy (CEM). Nello specifico, esso trunca la distribuzione generata al top-$\epsilon$ frazione della distribuzione della proprietà, dove $\epsilon$ è determinato dalla soglia $s^*$.
• Overhead Zero: Questo controllo non richiede oracle, modelli di ricompensa, calcolo di gradienti o campionamento per rifiuto (rejection sampling) al momento dell'inferenza.

3. Contributi Chiave

1. L'Accoppiamento come Interfaccia di Allineamento: Il paper inquadra l'accoppiamento rumore-dati nel flow matching da euristica di addestramento a interfaccia di allineamento a livello di distribuzione. Dimostra che la scelta dell'accoppiamento basata su una proprietà organizza il campo di flusso appreso in modo che le coordinate dello spazio del rumore corrispondano ai valori delle proprietà.
2. Steering Molecolare a Singola Manopola: Il metodo ottiene un controllo monotono sulle proprietà molecolari (logP e QED) utilizzando una singola coordinata del rumore. Fondamentalmente, la stessa manopola produce risposte strutturali opposte per target diversi: aumentare $s$ genera molecole più grandi per logP (lipofilia) ma molecole più piccole per QED (drug-likeness). Ciò esclude un generico bias di dimensione e conferma che l'accoppiamento scrive la struttura specifica del target nel flusso.
3. Reward Transport per Sequenze Discrete: Gli autori implementano questo metodo per sequenze di token SELFIES a lunghezza variabile, identificando e affrontando due specifici modi di fallimento nei modelli di flow che generano dati discreti:
   • Erasure Implicita del Segnale: Nelle architetture Pre-LayerNorm, l'informazione sull'ampiezza (codificata in $\|\bar{z}\|$) viene cancellata. Il Direction MLP risolve questo problema iniettando un pattern di direzione appreso.
   • Confounding della Lunghezza: Senza una perdita MSE non mascherata sui token di padding, il modello apprende una scorciatoia dove un alto $s$ aumenta semplicemente la lunghezza della sequenza per soddisfare il vincolo della norma. L'MSE non mascherato forza il modello a regredire il rumore grezzo nelle posizioni di padding, ancorando il segnale di ampiezza e prevenendo questa scorciatoia.

4. Risultati Sperimentali

Il metodo è stato valutato sui dataset ZINC-250K e GuacaMol.

• Performance di Controllo:
  • logP: Su ZINC-250K, far scorrere $s$ da -3 a +7 ha spostato il logP medio da 1.50 a 5.44 (+137%), con una correlazione di Spearman $\rho=1.000$ tra $s$ e la proprietà media. L'effetto (Cohen's $d$) ha raggiunto 1.63.
  • QED: La stessa manopola ha spostato il QED in modo monotono (ad esempio, +0.059 a $s=5$ su ZINC), con $\rho=1.000$.
  • Specificità Strutturale: Per logP, aumentare $s$ ha aumentato il numero di atomi pesanti (12 $\to$ 23). Per QED, aumentare $s$ ha diminuito il numero di atomi pesanti (23 $\to$ 16).
• Generalizzazione: I risultati sono stati replicati su GuacaMol (1.6M molecole), mostrando effetti ancora più forti per logP ($d=2.31$) e una monotonia costante per QED entro il range valido.
• Confronto con i Baseline:
  • Rispetto al Conditional Flow Matching (dove i valori delle proprietà sono input espliciti), il Reward Transport raggiunge una monotonia simile ($\rho=1.000$) ma opera tramite troncamento della distribuzione piuttosto che tramite l'individuazione esplicita di un target.
  • Rispetto alla diffusione con predizione di $\epsilon$, il metodo mostra che l'allineamento a livello di accoppiamento è strutturalmente assente nella predizione di $\epsilon$, poiché il target Bayes-ottimale è indipendente dall'accoppiamento, attenuando il segnale del gradiente.
• Dinamiche di Addestramento: La preservazione del rango (la correlazione tra $s$ e la proprietà $y$ generata) emerge rapidamente durante il fine-tuning (entro un epoch) e si stabilizza sotto 1.0 (ad esempio, 0.57 per logP), in linea con il gap di "approssimazione distributiva" descritto nella Proposizione 1.

5. Significato e Rivendicazioni

Il paper sostiene che il Reward Transport stabilisce l'accoppiamento nel flow matching come un canale di controllo distinto, a livello di distribuzione, che può trasportare una struttura di allineamento da solo.

• Efficienza: Elimina la necessità di chiamate all'oracle, modelli di ricompensa o ottimizzazione iterativa all'inferenza, riducendo significativamente il costo computazionale rispetto ai metodi basati su RL o guidance.
• Meccanismo: Dimostra che la "steerability" (capacità di guida) dei modelli generativi non deve necessariamente dipendere da input di condizionamento espliciti; può essere incorporata nella geometria del modello tramite l'accoppiamento di addestramento.
• Limiti e Ambito: Gli autori sono prudenti riguardo all'ambito. Il metodo fornisce una manopola distributiva 1D (troncando la distribuzione) piuttosto che un assegnamento esatto del valore per singolo campione. Funziona bene per proprietà scalari correlate con caratteristiche molecolari globali (come dimensione o polarità), ma fallisce per proprietà topologiche come la Sintetizzabilità (SA), dove una chiave scalare 1D non può distinguere strutture complesse simili in termini di dimensione. Inoltre, il metodo si è dimostrato incompatibile con i target di predizione $\epsilon$, evidenziando che l'interfaccia dipende dal target di predizione specifico della rete.

In sintesi, il lavoro sposta il paradigma della generazione controllabile dalla "guida della traiettoria all'inferenza" all' "allineamento della distribuzione di addestramento", offrendo un meccanismo principato e a overhead zero per il controllo delle proprietà nel flow matching.