Résumé technique : Routage par l'atteinte (Routing by Reaching)

Énoncé du problème

Les réseaux de flux génératifs (GFlowNets) sont efficaces pour échantillonner des candidats diversifiés proportionnellement à une fonction de récompense, une capacité cruciale pour les tâches de découverte scientifique comme la génération de molécules. Cependant, les applications réelles nécessitent souvent de équilibrer plusieurs objectifs fréquemment conflictuels (par exemple, l'affinité de liaison, l'accessibilité synthétique et la drug-likeness). Les approches existantes pour les GFlowNets multi-objectifs se classent généralement en deux catégories : la scalarisation (entraînement d'un modèle unique conditionné par des vecteurs de préférence) et les opérateurs logiques (combinaison de distributions via des classifieurs).

Ces méthodes existantes souffrent de deux limitations critiques :

1. Manque de réutilisabilité : Elles nécessitent un entraînement conjoint pour chaque combinaison spécifique d'objectifs. L'introduction d'un nouvel objectif ou d'une nouvelle préférence de compromis nécessite de réentraîner le modèle à partir de zéro ou d'entraîner un nouveau classifieur auxiliaire, ce qui engendre un surcoût computationnel substantiel.
2. Cadres fragmentés : Les méthodes actuelles traitent généralement la scalarisation ou la composition logique de manière isolée, manquant d'un cadre unifié capable de gérer diverses schémas de composition (allant des sommes linéaires aux opérateurs non linéaires complexes) sans réentraînement.

Méthodologie

Les auteurs proposent un cadre intitulé « Routing by Reaching » (Routage par l'atteinte), qui compose des GFlowNets pré-entraînés au moment de l'inférence en mélangeant leurs politiques directes (forward policies). L'idée centrale est que la probabilité d'atteinte $u(s)$ — la probabilité de visiter un état $s$ sous une politique directe spécifique — fournit des poids naturels pour combiner les processus de génération. Intuitivement, un modèle qui route plus de masse de probabilité à travers un état actuel est considéré comme plus pertinent pour l'objectif de cet état.

La règle de mélange

Étant donné $k$ GFlowNets pré-entraînés avec des politiques directes $p_{i,F}$, des probabilités d'atteinte $u_i$, et des distributions de terminaison $p_i$, la politique de mélange proposée $p_{M,F}$ pour une transition d'un état $s$ vers $s'$ est définie par :

$$p_{M,F}(s'|s) = \frac{G(u_1(s)p_{1,F}(s'|s), \dots, u_k(s)p_{k,F}(s'|s))}{N_M(s)}$$

où :

• $G$ est la fonction de composition (par exemple, la somme pondérée pour la scalarisation, la moyenne harmonique pour la conjonction).
• $N_M(s)$ est une constante de normalisation locale garantissant que la politique somme à 1.
• La probabilité d'atteinte $u_i(s)$ agit comme un poids, mettant à l'échelle la probabilité de transition du $i$-ème modèle.

Calcul des probabilités d'atteinte

Pour implémenter cela sans réentraînement, les auteurs fournissent deux voies pour obtenir $u_i(s)$ à partir de modèles pré-entraînés :

1. Modèle F : Pour les modèles entraînés avec le Flow Matching, le Detailed Balance ou le Sub-Trajectory Balance, le flux d'état $F_i(s)$ est explicitement paramétré. $u_i(s)$ est dérivé directement comme $F_i(s)/Z_i$.
2. DB F : Pour les modèles entraînés avec le Trajectory Balance (où le flux d'état n'est pas explicite), $u_i(s)$ est reconstruit à la volée pendant la génération de trajectoire en utilisant l'identité de l'équilibre détaillé (detailed balance) : $u_i(s_t) = \prod_{j=1}^t \frac{p_{i,F}(s_j|s_{j-1})}{p_{i,B}(s_{j-1}|s_j)}$.

Analyse théorique

• Scalarisation linéaire ($\beta=1$) : Les auteurs prouvent que pour les sommes pondérées linéaires de récompenses, la politique de mélange récupère exactement la distribution cible. Le facteur de distorsion entre les distributions induites et cibles devient constant.
• Opérateurs non linéaires : Pour les compositions non linéaires (par exemple, $\beta \neq 1$, moyenne harmonique, contraste), le facteur de distorsion n'est pas garanti constant. Cependant, l'analyse empirique montre que la distorsion reste approximativement constante dans les régions de haute densité (où la récompense cible est élevée), assurant des approximations précises là où elles sont les plus importantes pour l'échantillonnage.

Contributions clés

1. Composition sans entraînement : Une méthode pour composer des GFlowNets pré-entraînés pour des ensembles d'objectifs arbitraires et divers schémas de composition (linéaires et non linéaires) sans aucun entraînement ou ajustement fin (fine-tuning) supplémentaire.
2. Garanties théoriques : Preuve de la récupération exacte pour la scalarisation linéaire et quantification de la qualité d'approximation pour les opérateurs non linéaires via un facteur de distorsion.
3. Cadre unifié : Une approche unique qui gère à la fois la scalarisation conditionnée par préférence et les opérateurs logiques (conjonction/soustraction), surmontant la fragmentation de la littérature précédente.
4. Efficacité : Réduction significative du temps d'inférence pour les opérateurs logiques par rapport aux baselines guidées par classifieur, car elle évite d'énumérer et d'évaluer tous les états successeurs pour un terme de guidage.

Résultats expérimentaux

Le cadre a été évalué sur une grille 2D synthétique et sur des tâches de génération de molécules réelles (basées sur des fragments et basées sur des atomes/QM9).

• Grille 2D synthétique :
  • Scalarisation : La méthode a obtenu un erreur L1 nettement plus faible que les baselines conditionnées par préférence (MOGFN, HN-GFN) pour 2 à 5 objectifs. Contrairement aux baselines, les performances sont restées stables à mesure que le nombre d'objectifs augmentait.
  • Opérateurs logiques : La méthode a performé de manière comparable aux baselines guidées par classifieur pour les opérations de moyenne harmonique et de contraste, mais sans nécessiter d'entraînement supplémentaire. Elle a nettement surpassé une baseline d'« ensemble » qui mélangeait les politiques sans pondération par la probabilité d'atteinte.
• Génération de molécules :
  • Scalarisation : L'approche a égalé ou dépassé la qualité d'échantillonnage (récompense moyenne et diversité) et l'hypervolume des baselines qui nécessitent un entraînement conjoint pour chaque combinaison d'objectifs.
    • Opérateurs logiques : La méthode a obtenu des pourcentages élevés de cibles de type « bin » pour les opérations de moyenne harmonique et de contraste, surpassant souvent le guidage par classifieur. Crucialement, elle a maintenu une validité moléculaire comparable aux modèles de base.
  • Coût d'inférence : Pour la scalarisation, la méthode a engendré un surcoût minimal (8–40 ms par échantillon). Pour les opérateurs logiques, elle était 30 à 65 fois plus rapide que le guidage par classifieur car elle évite l'énumération exhaustive des états requise par ce dernier.

Signification et affirmations

L'article affirme que « Routing by Reaching » permet aux GFlowNets pré-entraînés de servir de blocs de construction réutilisables. En composant ces modèles au moment de l'inférence, les praticiens peuvent rapidement s'adapter à diverses combinaisons d'objectifs et explorer le front de Pareto sans le coût prohibitif du réentraînement.

Les auteurs soulignent que, bien que la méthode soit théoriquement exacte pour la scalarisation linéaire, elle fournit des approximations précises pour les opérateurs non linéaires complexes dans les régions de haute récompense. La principale signification réside dans le découplage de l'entraînement des objectifs individuels de l'étape de composition, facilitant ainsi une exploration multi-objective flexible dans les domaines de la découverte scientifique comme la conception de médicaments, où les exigences évoluent souvent dynamiquement. Le travail ne prétend pas introduire de nouvelles capacités génératrices fondamentales, mais plutôt un mécanisme pour exploiter efficacement les modèles entraînés existants pour des tâches multi-objectifs complexes.