Constrained Flow Optimization via Sequential Fine Tuning for Molecular Design

Résumé Technique : Optimisation de Flux Contrainte par Fine-Tuning Séquentiel pour la Conception Moléculaire

1. Énoncé du Problème

L'article traite du défi consistant à adapter les modèles de fondation génératifs, spécifiquement les modèles de flux et de diffusion, pour des applications de découverte scientifique réelle telles que la conception moléculaire et l'ingénierie des protéines. Bien que ces modèles excellent dans la capture de distributions de données complexes, l'échantillonnage standard est insuffisant pour les tâches de découverte où l'objectif est de générer des candidats qui maximisent des récompenses spécifiques (ex. : affinité de liaison, moment dipolaire) tout en respectant strictement des contraintes spécifiques au domaine (ex. : accessibilité synthétique, limites de toxicité ou stabilité énergétique).

Les méthodes de fine-tuning existantes, souvent basées sur l'apprentissage par renforcement (RL) ou la théorie du contrôle, optimisent généralement des fonctions de récompense sous régularisation KL pour éviter la dérive distributionnelle. Cependant, elles peinent à gérer les contraintes strictes de manière fiable. Les approches naïves, telles que l'incorporation des contraintes sous forme de termes pondérés dans l'objectif, sont jugées peu fiables car le poids optimal entre récompenses et contraintes varie selon les tâches et les phases d'entraînement, nécessissant un réglage par essais et erreurs inefficace. De plus, ces méthodes ne peuvent garantir que les échantillons générés respectent des contraintes explicites, produisant souvent des échantillons à haute récompense qui violent des propriétés physiques ou chimiques critiques.

La question centrale posée est la suivante : Comment un modèle de flux ou de diffusion pré-entraîné peut-il être ajusté pour optimiser de manière fiable et prévisible le compromis entre l'optimisation de la récompense et la satisfaction des contraintes ?

2. Méthodologie

Les auteurs proposent un cadre formel et une solution algorithmique à ce problème.

Cadre d'Optimisation Générative Contrainte :
Le problème est formulé comme la recherche d'une politique $\pi^*$ qui maximise l'espérance de la récompense $r(x)$ tout en satisfaisant une contrainte attendue $c(x) \leq B$, sous réserve d'une régularisation KL par rapport à un modèle pré-entraîné $\pi_{pre}$. L'objectif est :
$$\max_{\pi} \mathbb{E}_{x \sim p^{\pi}_1}[r(x)] - \alpha D_{KL}(p^{\pi}_1 || p^{pre}_1) \quad \text{sous contrainte de} \quad \mathbb{E}_{x \sim p^{\pi}_1}[c(x)] \leq B$$
Cette formulation généralise à la fois l'optimisation guidée par la récompense et la génération contrainte (où $r$ est constante).

Optimisation de Flux Contrainte (CFO - Constrained Flow Optimization) :
Pour résoudre cela, les auteurs introduisent la Constrained Flow Optimization (CFO), un algorithme basé sur le schéma du Lagrangien Augmenté (AL). La CFO transforme le problème contraint en une séquence de sous-problèmes de fine-tuning non contraints et régularisés par l'entropie. L'algorithme fonctionne de manière itérative avec les étapes suivantes :

1. Construction de l'Objectif Augmenté : À chaque itération $k$, une fonction d'objectif augmentée $f_k(x)$ est construite. Cette fonction combine la récompense $r(x)$ avec un terme de pénalité dérivé de la violation de la contrainte, ajusté par un multiplicateur de Lagrange $\lambda_k$ et un paramètre de pénalité $\rho_k$ :
   $$f_k(x) = r(x) - \frac{\rho_k}{2} \left[ \max\left(0, c(x) - B - \frac{\lambda_k}{\rho_k}\right) \right]^2$$
2. Étape de Fine-Tuning : Un solveur de fine-tuning standard (ex. : Adjoint Matching, DiffusionNFT) est utilisé pour optimiser la politique $\pi_k$ par rapport à l'objectif augmenté $f_k$, régularisé par la divergence KL par rapport au modèle pré-entraîné.
3. Mises à jour des Variables Duales :
   • Multiplicateur de Lagrange ($\lambda$) : Mis à jour en fonction de l'écart de contrainte empirique $G_k = \mathbb{E}[c(x)] - B$. Si la contrainte est violée ($G_k > 0$), $\lambda$ est diminué pour renforcer la pénalité ; si elle est satisfaite, elle est relâchée.
   • Paramètre de Pénalité ($\rho$) : Mis à jour en fonction d'une statistique de contraction $V_k$. Si l'écart de contrainte ne se contracte pas suffisamment, $\rho$ est augmenté pour imposer une faisabilité plus stricte.

Cette approche séquentielle permet à l'algorithme d'équilibrer automatiquement la maximisation de la récompense et la satisfaction des contraintes sans réglage manuel des poids.

3. Principales Contributions

L'article présente quatre contributions majeures :

• Cadre Formel : Il formule l'optimisation générative contrainte via le fine-tuning de flux, capturant la tâche pratique d'adaptation guidée par la récompense sous des contraintes arbitraires.
• Algorithme CFO : Il introduit la CFO, une méthode basée sur le Lagrangien Augmenté qui résout de manière prouvée le problème contraint via un fine-tuning séquentiel.
• Garanties Théoriques : Sous des hypothèses réalistes d'un solveur approximatif et de conditions d'oracle spécifiques, les auteurs fournissent des garanties de convergence. Ils prouvent que la CFO renvoie une politique qui minimise l'infaisabilité et, sous des conditions plus fortes (erreur d'approximation nulle), atteint l'optimalité globale.
• Validation Empirique : La méthode est évaluée à la fois sur des contextes illustratifs de faible dimension et sur des tâches de conception moléculaire de haute dimension, démontrant sa capacité à échanger efficacement récompenses et contraintes.

4. Résultats Expérimentaux

Les auteurs évaluent la CFO sur deux fronts principaux :

Contextes Illustratifs de Faible Dimension :

• Optimisation Contrainte : En utilisant un mélange de Gaussiennes, la CFO a réussi à diriger le modèle vers des régions de haute récompense (proche d'un point cible) tout en restant dans des régions de contraintes valides (triangles). Comparée à l'Adjoint Matching standard (qui ignore les contraintes), la CFO a considérablement réduit les violations de contraintes (de 2,47 à 0,12) tout en maintenant une récompense élevée.
• Génération Contrainte : En fixant la récompense à une constante, la CFO a récupéré des capacités de génération contrainte. La variation du seuil de contrainte $B$ a permis un relâchement contrôlé de la contrainte, démontrant la flexibilité de la méthode.
• Agnosticisme du Solveur : La CFO a performé de manière cohérente lorsqu'elle était couplée à la fois à des solveurs de premier ordre (Adjoint Matching) et à des solveurs sans gradient (DiffusionNFT), confirmant sa modularité.

Tâches de Conception Moléculaire :

• Maximisation du Moment Dipolaire sous Contrainte d'Énergie : En utilisant le modèle FlowMol sur le jeu de données GEOM Drugs, la CFO a optimisé le moment dipolaire (récompense) tout en imposant une limite supérieure sur l'énergie totale (contrainte, proxy de la stabilité).
  • La CFO a augmenté le moment dipolaire de 6,55 D (pré-entraîné) à 8,39 D tout en respectant la contrainte d'énergie (énergie attendue -82,28 Ha vs limite -80 Ha).
  • En revanche, le fine-tuning non contraint (Adjoint Matching) a atteint des récompenses similaires mais a violé la contrainte d'énergie (énergie attendue -78,31 Ha).
• Comparaison aux Baselines : La CFO a surpassé les baselines à poids fixe ajustés manuellement. Le réglage manuel a nécessité une recherche exhaustive sur 18 ordres de grandeur de poids pour trouver une seule configuration égalant la performance de la CFO, soulignant l'inefficacité du réglage manuel.
• Validité et Efficacité : Bien que le fine-tuning réduise le pourcentage de molécules chimiquement valides (un défi connu de la conception générative), la CFO a maintenu une validité plus élevée que les baselines non contraintes à des niveaux de récompense comparables. La CFO a également démontré son efficacité computationnelle, ne nécessitant qu'un seul passage de fine-tuning par rapport aux coûts d'échantillonnage élevés des méthodes de contrainte au moment de l'inférence comme DiffOpt.

5. Signification et Revendications

L'article affirme que la CFO fournit un mécanisme fiable et prévisible pour l'optimisation générative contrainte. Sa signification réside dans :

• Compromis Automatisé : Elle élimine le besoin de réglage manuel et inefficace par essais et erreurs des poids de contrainte, ce qui constitue un goulot d'étranglement majeur dans les flux de travail de conception générative actuels.
• Rigueur Théorique : Elle offre des garanties de convergence tant pour la satisfaction des contraintes que pour l'optimalité de la récompense, ancrant la méthode dans la théorie de l'optimisation établie (Lagrangien Augmenté).
• Utilité Pratique : En s'appliquant avec succès à des tâches de conception moléculaire de haute dimension, les auteurs démontrent que la méthode passe à l'échelle pour des problèmes scientifiques réels où les contraintes (comme la stabilité chimique) sont non négociables.
• Modularité : Le cadre est agnostique vis-à-vis du solveur de fine-tuning sous-jacent, permettant son intégration avec diverses stratégies d'optimisation (basées sur le gradient ou sans gradient) et son application à différentes modalités de données (continues, discrètes ou mixtes).

Les auteurs concluent que la CFO permet d'adapter les modèles de flux pré-entraînés pour maximiser les récompenses attendues tout en imposant des contraintes arbitraires, offrant ainsi un outil robuste pour les applications de découverte scientifique.