Résumé Technique : Reward Transport – Contrôle de Propriété via l'Alignement dans l'Espace de Bruit

1. Énoncé du Problème

Dans le Flow Matching (FM) et les modèles génératifs apparentés basés sur la diffusion, le couplage — la règle déterminant quel vecteur de bruit est associé à quel point de donnée pendant l'entraînement — est généralement traité comme une heuristique de calcul. Les approches standards utilisent un couplage indépendant pour des raisons de simplicité ou un transport optimal (OT) par mini-lots pour linéariser les trajectoires et réduire la variance. Dans ces paradigmes, le couplage est optimisé uniquement pour faciliter l'entraînement, et non pour encoder des propriétés structurelles spécifiques dans le champ de flux appris.

Par conséquent, obtenir une génération contrôlable (par exemple, générer des molécules avec des scores logP ou QED spécifiques) nécessite généralement des mécanismes explicites au moment de l'inférence, tels que :

• Le guidage sans classificateur (classifier-free guidance) : Nécessite une valeur de propriété cible comme entrée.
• Le guidage par gradient (gradient guidance) : Calcule les gradients d'un prédicteur de propriété à chaque étape de l'échantillonnage.
• L'apprentissage par renforcement (RL) : Affine le générateur contre un oracle, ce qui nécessite souvent des requêtes à l'oracle pendant l'échantillonnage.

Ces méthodes introduisent une surcharge de calcul à l'inférence, une dépendance vis-à-vis de modèles de récompense ou d'oracles externes, ou des pipelines d'entraînement complexes. L'article se demande si le couplage lui-même peut servir d'interface d'alignement, en encodant une structure contrôlable directement dans le champ de flux lors de l'entraînement, de sorte que le contrôle soit disponible à l'inférence sans calcul supplémentaire ni appel à un oracle.

2. Méthodologie : Reward Transport

Les auteurs proposent le Reward Transport, une méthode qui redéfinit le couplage bruit-donnée comme un mécanisme d'injection de la structure de la tâche dans le champ de flux.

2.1 Couplage de Transport Optimal Aligné sur les Propriétés

Au lieu d'associer le bruit et les données de manière indépendante, le Reward Transport construit un couplage monotone basé sur une propriété moléculaire cible $y(x)$ (par exemple, logP, QED).

1. Clés Scalaires : Pour chaque vecteur de bruit $z$ (une séquence d'embeddings de tokens), une clé de tri scalaire $s(z)$ est calculée. Le choix principal est la norme euclidienne de la moyenne pondérée du bruit : $s(z) = \|\bar{z}\|_2$.
2. Appariement par Rang : Étant donné un lot de vecteurs de bruit $\{z_n\}$ et de molécules $\{x_n\}$ avec leurs propriétés $\{y_n\}$, les deux ensembles sont triés selon leurs clés respectives ($s$ et $y$). Ils sont ensuite appariés rang par rang : le vecteur de bruit ayant le $k$-ième plus petit $s$ est associé à la molécule ayant le $k$-ième plus petit $y$.
3. Objectif d'Entraînement : Le modèle de flow matching $v_\theta$ est entraîné pour transporter ces paires alignées. La perte combine une erreur quadratique moyenne (MSE) standard sur la prédiction de l'extrémité (dans l'espace d'embedding continu) et une perte d'entropie croisée (CE) pour garantir que les tokens décodés forment des molécules valides.

.2 Conditionnement de Direction

Pour garantir que le modèle puisse accéder à la coordonnée de l'espace de bruit $s$ lors de l'inférence, un MLP de Direction est introduit. Ce MLP à deux couches projette le scalaire normalisé $s$ vers un embedding de dimension $d$, qui est ajouté à l'entrée du transformer aux côtés des embeddings temporels et positionnels. Cela permet au modèle de conditionner son champ de vitesse à la coordonnée de bruit spécifique utilisée lors du couplage monotone.

2.3 Mécanisme d'Inférence

À l'inférence, l'utilisateur spécifie un signal de direction scalaire $s^*$ (une coordonnée de l'espace de bruit, et non une valeur de propriété cible). Le modèle génère des molécules en intégrant le champ de flux à partir de vecteurs de bruit où $s(z) \geq s^*$.

• Garantie Théorique (Proposition 1) : Dans la limite où le flux appris préserve parfaitement la correspondance de rang entre $s$ et $y$ (limite de conservation du couplage), le conditionnement sur $s \geq s^*$ induit une distribution équivalente à une étape de sélection de la méthode de l'entropie croisée (CEM). Plus précisément, cela tronque la distribution générée pour ne garder la fraction supérieure de $\epsilon$ de la distribution de la propriété, où $\epsilon$ est déterminé par le seuil $s^*$.
• Surcharge Nulle : Ce contrôle ne nécessite aucun oracle, aucune récompense, aucun calcul de gradient, ni échantillonnage par rejet au moment de l'inférence.

3. Contributions Clés

1. Le Couplage comme Interface d'Alignement : L'article recadre le couplage bruit-donnée dans le flow matching, passant d'une heuristique d'entraînement à une interface d'alignement au niveau de la distribution. Il démontre que le choix du couplage basé sur une propriété cible organise le champ de flux appris de telle sorte que les coordonnées de l'espace de bruit correspondent aux valeurs de propriété.
2. Pilotage Moléculaire à Bouton Unique : La méthode permet un contrôle monotone des propriétés moléculaires (logP et QED) à l'aide d'une seule coordonnée scalaire de bruit. Crucialement, le même bouton produit des réponses structurelles opposées pour différentes cibles : augmenter $s$ génère des molécules plus grandes pour le logP (lipophilie) mais plus petites pour le QED (similitude avec un médicament). Cela exclut un biais de taille générique et confirme que le couplage écrit une structure spécifique à la cible dans le flux.
3. Reward Transport pour les Séquences Discrètes : Les auteurs implémentent cette méthode pour des séquences de tokens SELFIES de longueur variable, identifiant et traitant deux modes de défaillance spécifiques aux modèles de flux générant des données discrètes :
   • Érasement Implicite du Signal : Dans les architectures Pre-LayerNorm, l'information de magnitude (encodée dans $\|\bar{z}\|$) est effacée. Le MLP de Direction résout cela en injectant un motif de direction appris.
   • Confusion de Longueur (Length Confound) : Sans une perte MSE non masquée sur les tokens de padding, le modèle apprend un raccourci où un $s$ élevé augmente simplement la longueur de la séquence pour satisfaire la contrainte de norme. La MSE non masquée force le modèle à régresser le bruit brut aux positions de padding, ancrant ainsi le signal de magnitude et empêchant ce raccourci.

4. Résultats Expérimentaux

La méthode a été évaluée sur les jeux de données ZINC-250K et GuacaMol.

• Performance de Contrôle :
  • logP : Sur ZINC-250K, balayer $s$ de -3 à +7 a déplacé le logP moyen de 1,50 à 5,44 (+137 %), avec une corrélation de Spearman $\rho=1,000$ entre $s$ et la propriété moyenne. La taille de l'effet (Cohen's $d$) a atteint 1,63.
  • QED : Le même bouton a déplacé le QED de manière monotone (ex: +0,059 à $s=5$ sur ZINC), avec $\rho=1,000$.
  • Spécificité Structurelle : Pour le logP, augmenter $s$ a augmenté le nombre d'atomes lourds (12 $\to$ 23). Pour le QED, augmenter $s$ a diminué le nombre d'atomes lourds (23 $\to$ 16).
• Généralisation : Les résultats ont été répliqués sur GuacaMol (1,6M de molécules), montrant des tailles d'effet encore plus fortes pour le logP ($d=2,31$) et une monotonie constante pour le QED dans la plage valide.
• Comparaison aux Références (Baselines) :
  • Comparé au Flow Matching Conditionnel (où les valeurs de propriété sont des entrées explicites), le Reward Transport atteint une monotonie similaire ($\rho=1,000$) mais opère via une troncature de distribution plutôt que par une atteinte de cible explicite.
  • Comparé à la diffusion par prédiction de $\epsilon$, la méthode montre que l'alignement au niveau du couplage est structurellement absent dans la prédiction de $\epsilon$, car la cible de Bayes-optimale est indépendante du couplage, ce qui atténue le signal de gradient.
• Dynamique d'Entraînement : La préservation du rang (la corrélation entre $s$ et la propriété générée $y$) émerge rapidement pendant l'affinage (en moins d'un époque) et plafonne en dessous de 1,0 (ex: 0,57 pour logP), ce qui est cohérent avec l'écart d'approximation distributionnelle décrit dans la Proposition 1.

5. Signification et Revendications

L'article affirme que le Reward Transport établit le couplage dans le flow matching comme un canal de contrôle distinct au niveau de la distribution capable de porter une structure d'alignement par lui-même.

• Efficacité : Il élimine le besoin d'appels d'oracle, de modèles de récompense ou d'optimisation itérative lors de l'inférence, réduisant considérablement le coût de calcul par rapport aux méthodes de RL ou de guidage.
• Mécanisme : Il démontre que la "capacité de pilotage" des modèles génératifs ne doit pas dépendre d'entrées de conditionnement explicites ; elle peut être intégrée dans la géométrie du modèle via le couplage d'entraînement.
• Limites et Portée : Les auteurs font preuve de modestie quant à la portée de la méthode. Elle fournit un bouton de distribution unidimensionnel (tronquant la distribution) plutôt qu'une assignation de valeur exacte par échantillon. Elle fonctionne bien pour les propriétés scalaires corrélées à des caractéristiques moléculaires globales (comme la taille ou la polarité), mais échoue pour les propriétés topologiques comme la Synthétic Accessibility (SA) où une clé scalaire 1D ne peut distinguer des structures complexes de taille similaire. De plus, la méthode est montrée comme incompatible avec les cibles de prédiction de $\epsilon$, soulignant que l'interface dépend de la cible de prédiction du réseau.

En résumé, ce travail déplace le paradigme de la génération contrôlable de "guider la trajectoire à l'inférence" vers "aligner la distribution d'entraînement", offrant un mécanisme de contrôle de propriété fondé et sans surcharge pour le flow matching.