Reference-guided Policy Optimization for Molecular Optimization via LLM Reasoning

Ce papier présente RePO, une nouvelle méthode d'optimisation de politique guidée par des références qui combine l'apprentissage par renforcement pour l'exploration et un apprentissage supervisé pour l'exploitation, afin de surmonter les limitations des approches actuelles dans l'optimisation moléculaire basée sur les grands modèles de langage en l'absence de trajectoires de raisonnement détaillées.

L'essentiel

🧪 Le Problème : L'Architecte et le Dessin Manquant

Imaginez que vous avez un architecte très intelligent (c'est notre Modèle de Langage, ou LLM) et que vous lui demandez de rénover une maison (une molécule chimique).

• L'objectif : Vous voulez que la maison soit plus confortable (meilleure propriété chimique) tout en restant très similaire à l'originale (pour ne pas la détruire).
• Le problème : Dans les manuels d'apprentissage actuels, on donne à l'architecte un seul exemple : "Voici la maison de départ, et voici la maison finie idéale."
• Ce qui se passe :
  1. Si on lui demande de juste copier le résultat final (SFT), il arrête de réfléchir. Il devient paresseux, donne une réponse courte et directe, et oublie comment construire étape par étape. C'est comme un élève qui apprend par cœur la réponse sans comprendre le cours.
  2. Si on lui dit "Essaie, et je te donne un point si c'est bien" (RLVR), il essaie au hasard. Mais comme trouver la bonne modification est comme chercher une aiguille dans une botte de foin, il ne reçoit presque jamais de points. Il finit par ne rien faire de peur de se tromper, ou il fait des modifications trop timides qui n'améliorent rien.

💡 La Solution : RePO (L'Architecte Guidé)

Les auteurs proposent une nouvelle méthode appelée RePO (Optimisation de Politique Guidée par Référence). Voici comment ça marche avec une analogie simple :

Imaginez que vous apprenez à cuisiner un plat complexe.

• L'approche classique (SFT) : On vous donne juste la photo du plat fini. Vous essayez de le copier aveuglément sans comprendre les techniques.
• L'approche classique (RL) : On vous laisse cuisiner seul dans une cuisine vide. Si le plat est bon, on vous félicite. Si c'est raté, on ne dit rien. Vous avez peur de cuisiner et vous ne progressez pas.

L'approche RePO, c'est comme avoir un Chef de Cuisine (la référence) qui vous regarde :

1. L'Exploration (Le RL) : Vous êtes libre de réfléchir, de tester des ingrédients, de dessiner des croquis de recettes (c'est le "raisonnement"). Le système vous récompense si le résultat final est bon. Cela vous encourage à explorer de nouvelles idées.
2. La Guidance (La Référence) : Mais attention ! À la fin, le Chef regarde votre plat fini et dit : "Tiens, ton plat ressemble beaucoup à celui que je voulais. C'est bien !".
   • Le secret : Le Chef ne vous dit pas comment vous avez pensé (il ne vous force pas à copier vos croquis). Il valide seulement le résultat final en le comparant à son exemple.
   • Cela permet à l'architecte de rester créatif dans sa réflexion, mais de savoir qu'il est sur la bonne voie grâce à l'exemple de référence.

🚀 Pourquoi c'est génial ?

• Équilibre parfait : RePO évite deux pièges : celui de l'architecte qui ne réfléchit plus (copie aveugle) et celui de l'architecte qui perd son temps à chercher dans le vide.
• Résultats concrets : Sur les tests, cette méthode a permis de créer des molécules (des "maisons") qui sont à la fois plus performantes (meilleures propriétés) et plus sûres (ressemblent bien à l'original) que les méthodes précédentes.
• Adaptabilité : Même si on change la façon de donner les instructions (un nouveau style de recette), l'architecte guidé par RePO s'adapte mieux que les autres.

En résumé

C'est comme apprendre à conduire une voiture de course :

• Les anciennes méthodes vous donnaient soit la photo de l'arrivée (vous copiez sans comprendre), soit vous lâchaient sur une piste sans instructeur (vous avez peur de tourner).
• RePO, c'est avoir un instructeur à côté de vous qui vous laisse conduire et essayer des virages (exploration), mais qui vous dit "Bravo, tu es sur la bonne trajectoire" quand vous arrivez au bon endroit (guidance par la référence).

Le résultat ? Une conduite plus sûre, plus rapide, et des pilotes qui comprennent vraiment la route ! 🏎️✨

Résumé technique

1. Problématique et Contexte

L'article aborde le défi de l'optimisation moléculaire basée sur des instructions, où un modèle de langage (LLM) doit modifier une molécule d'entrée ($m_0$) pour améliorer une propriété cible (ex: QED, LogP) tout en maintenant une similarité structurelle élevée avec la molécule originale.

Le problème central réside dans la « mismatch de supervision » (inadéquation de la supervision) des méthodes actuelles :

• Données limitées : Les ensembles de données fournissent généralement une seule molécule de référence optimisée ($m_{ref}$) sans trajectoire de raisonnement intermédiaire (étapes de modification).
• Échec du SFT (Supervised Fine-Tuning) : Un SFT standard sur les réponses finales tend à effondrer le processus de raisonnement. Le modèle apprend à sauter directement à la réponse (« answer-only »), supprimant les étapes de raisonnement multi-étapes nécessaires pour explorer l'espace chimique.
• Échec du RLVR (Reinforcement Learning with Verifiable Rewards) : Des méthodes comme GRPO, lorsqu'elles sont appliquées à partir d'un modèle de base ou d'un SFT, souffrent de récompenses clairsemées. Dans un espace chimique contraint, les molécules qui satisfont à la fois l'amélioration de la propriété et la contrainte de similarité sont rares. Cela conduit à une exploration conservatrice (modifications minimes) ou à un échec de l'apprentissage.

2. Méthodologie : RePO (Reference-guided Policy Optimization)

Les auteurs proposent RePO, une approche d'optimisation qui combine l'exploration guidée par la récompense et l'exploitation de références, sans nécessiter de données de trajectoire étiquetées.

Objectif d'Optimisation

L'objectif de RePO ($J_{RePO}$) intègre trois termes pour chaque requête $q = (x, m_0)$ et un ensemble de $G$ réponses échantillonnées $\{o_i\}$ :

1. Terme d'Exploration (RLVR) : Basé sur GRPO, il met à jour la politique sur l'ensemble des tokens (raisonnement + réponse) en utilisant un avantage relatif de groupe ($\hat{A}_{i,k}$) basé sur une récompense vérifiable. Cela encourage l'exploration de nouvelles molécules.

   • Récompense ($r$) : Combinaison de la similarité structurelle (Tanimoto sur les empreintes digitales moléculaires) et d'une récompense binaire pour l'amélioration de la propriété cible.

2. Terme de Guidance par Référence (Answer-Level) : C'est l'innovation clé. Au lieu d'imiter la trajectoire de raisonnement de la référence, RePO augmente la probabilité de la molécule de référence ($m_{ref}$) conditionnée par la trajectoire de raisonnement générée par le modèle ($t_i$).

   • Formellement : $\log \pi_\theta(m_{ref} | q, t_i)$.
   • Cela ancre la réponse finale à une solution valide connue tout en laissant le modèle libre d'explorer différents chemins de raisonnement pour y parvenir.

3. Régularisation KL : Pour stabiliser l'entraînement et empêcher la dérive de la politique par rapport à la politique de référence initiale.

Mécanismes Clés

• Masquage des gradients : Les gradients du terme de guidance ne sont appliqués qu'aux tokens de la réponse finale, pas aux tokens de raisonnement intermédiaires. Cela empêche le modèle de copier aveuglément le raisonnement de la référence, préservant ainsi sa capacité d'exploration.
• Pas de trajectoires étiquetées : La méthode n'a pas besoin de données d'entraînement avec des étapes de raisonnement explicites, seulement de la molécule de référence finale.

3. Contributions Principales

1. Révélation du problème de supervision : Les auteurs démontrent empiriquement que le SFT « réponse seule » effondre le raisonnement multi-étapes, et que le RLVR pur souffre de signaux d'apprentissage trop clairsemés sous contraintes de similarité.
2. Proposition de RePO : Une nouvelle formulation d'optimisation qui couple l'exploration RL (sur les trajectoires) et la guidance de référence (sur la réponse finale), résolvant le compromis entre exploration et exploitation.
3. Validation Expérimentale : Démonstration que RePO surpasse systématiquement les baselines (SFT, GRPO, GRPO initialisé par SFT) sur des benchmarks standards, tout en généralisant mieux aux instructions non vues.

4. Résultats Expérimentaux

Les expériences ont été menées sur deux benchmarks : TOMG-Bench (optimisation mono-objectif) et MuMOInstruct (optimisation multi-objectif).

• Performance Globale : RePO obtient les meilleurs scores sur la métrique composite Success Rate × Similarité (SR × Sim).
  • Sur TOMG-Bench, RePO améliore le taux de succès de jusqu'à 17,4 % par rapport à GRPO sur certaines tâches.
  • Il surpasse les méthodes SFT et GRPO sur 4 des 6 tâches mono-objectif.
• Équilibre des Objectifs : Sur les tâches multi-objectifs (MuMOInstruct), RePO parvient mieux à équilibrer des objectifs concurrents (ex: augmenter la perméabilité BBB tout en maintenant la toxicité faible) que les méthodes SFT (qui sacrifient souvent la similarité) ou GRPO (qui sont trop conservatrices).
• Généralisation : RePO maintient ses avantages sur des styles d'instructions non vus lors de l'entraînement, prouvant sa robustesse.
• Qualité du Raisonnement : L'évaluation via un « LLM-as-a-judge » montre que RePO génère des trajectoires de raisonnement chimiquement valides et cohérentes, contrairement à GRPO qui produit souvent des erreurs chimiques ou des réponses non valides.
• Échelle d'Inférence : Les performances de RePO s'améliorent avec l'augmentation du budget de calcul (échantillonnage multiple), indiquant une capacité à exploiter efficacement les ressources de calcul.

5. Signification et Impact

Ce travail est significatif pour plusieurs raisons :

• Résolution du dilemme Exploration/Exploitation en Chimie : Il offre une solution élégante au problème de la rareté des récompenses dans l'espace chimique contraint, en utilisant des références finales pour guider l'exploration sans la restreindre excessivement.
• Efficacité des Données : RePO fonctionne sans nécessiter de données coûteuses de trajectoires de raisonnement (step-by-step), ce qui est crucial car de telles données sont rares en chimie computationnelle.
• Validité Chimique : En évitant l'effondrement du raisonnement (collapse) typique du SFT, RePO permet aux LLM de raisonner correctement sur les principes chimiques (stérique, électronégativité) pour proposer des modifications valides.
• Généralisation des LLM : Cela démontre que les LLM généraux, correctement optimisés, peuvent surpasser les modèles spécialisés pré-entraînés sur des corpus chimiques pour des tâches d'optimisation complexe.

En conclusion, RePO établit un nouvel état de l'art pour l'optimisation moléculaire assistée par LLM, en combinant la puissance du raisonnement génératif avec la rigueur des contraintes chimiques via une guidance de référence intelligente.