Pareto-Optimal Offline Reinforcement Learning via Smooth Tchebysheff Scalarization

Ce papier présente STOMP, un nouvel algorithme d'apprentissage par renforcement hors ligne multi-objectif basé sur la scalarisation Tchebysheff lisse qui surpasse les méthodes existantes en optimisant simultanément des attributs contradictoires, tels que l'activité et la spécificité, dans des tâches d'ingénierie des protéines.

L'essentiel

🧬 Le Dilemme du Chef Cuisinier : Comment trouver l'équilibre parfait ?

Imaginez que vous êtes un chef cuisinier de génie (c'est votre Intelligence Artificielle). Votre travail consiste à créer de nouveaux plats (des protéines).

Le problème, c'est que vos clients ont des demandes contradictoires :

1. Le plat doit être délicieux (très actif).
2. Il ne doit pas être toxique (sûr).
3. Il doit être facile à préparer (stable).

Souvent, un plat très délicieux est difficile à préparer, ou un plat très sûr manque de saveur. C'est ce qu'on appelle un conflit d'objectifs.

Dans le passé, les chercheurs utilisaient une méthode simple pour résoudre ce problème : ils donnaient une note globale en faisant une moyenne pondérée.

Exemple : "Je veux 50% de goût et 50% de sécurité."

Le problème de l'ancienne méthode : C'est comme si vous essayiez de tracer une ligne droite sur une carte pour trouver le meilleur chemin. Mais si le terrain est vallonné (avec des pics et des creux), une ligne droite ne passera jamais par les plus beaux points de vue cachés dans les vallées. Vous ratez les solutions de compromis parfaites.

🚀 La Solution : STOMP (Le GPS de l'Équilibre)

Les auteurs de cet article, travaillant chez Profluent Bio, ont inventé une nouvelle méthode appelée STOMP. Pour comprendre comment ça marche, utilisons une analogie.

1. Le problème de l'échelle (La règle de mesure)

Imaginez que vous devez comparer deux choses très différentes :

• La température d'un four (qui va de 0 à 300 degrés).
• Le prix d'un ingrédient (qui va de 1 à 10 euros).

Si vous essayez de les additionner directement, le prix (10) n'aura aucune influence sur la température (300). C'est injuste. Il faut d'abord normaliser les mesures pour qu'elles soient comparables.

2. La technique magique : "Lissage Tchebysheff"

Au lieu de faire une simple moyenne, STOMP utilise une technique mathématique sophistiquée appelée lissage Tchebysheff.

L'analogie du "Seuil de tolérance" :
Imaginez que vous êtes un juge dans un concours de cuisine. Au lieu de dire "Je veux le plat le plus savoureux ET le plus sûr", vous dites :

"Je vais regarder le plat le plus mauvais que vous pouvez faire pour chaque critère. Ensuite, je vais chercher le plat qui s'éloigne le moins de ce 'pire scénario' pour TOUS les critères en même temps."

C'est comme si vous cherchiez le point le plus proche du centre d'une cible, même si la cible a une forme bizarre et irrégulière. Cette méthode permet de trouver des solutions que les anciennes méthodes (les lignes droites) ne voyaient tout simplement pas.

🧪 L'Expérience : Des protéines sur mesure

Pour tester leur idée, les chercheurs ont utilisé cette méthode sur des modèles de langage pour les protéines (des IA qui "lisent" et "écrivent" des séquences d'acides aminés, comme un texte).

Ils ont demandé à l'IA de créer des protéines avec trois objectifs différents selon les expériences :

• DHFR : Une protéine qui doit être active, même en présence d'un poison (un antibiotique).
• PbrR : Une protéine qui doit attraper le plomb (utile) mais repousser le zinc (inutile).
• α-Amylase : Une protéine qui doit être active, stable et facile à produire.

Le résultat ?
STOMP a été un succès retentissant.

• Là où les anciennes méthodes trouvaient 5 ou 6 bonnes solutions, STOMP en a trouvé 8 ou 9.
• Il a réussi à explorer des zones "interdites" ou difficiles où les compromis étaient les plus intéressants.
• C'est comme si STOMP avait trouvé des ingrédients cachés dans le garde-manger que les autres chefs ne voyaient pas.

💡 Pourquoi c'est important pour tout le monde ?

Ce n'est pas juste pour les biologistes. Cette méthode change la façon dont on utilise l'IA pour résoudre des problèmes complexes :

1. Plus de compromis intelligents : Au lieu de choisir entre "bon marché" et "bon", on peut trouver des options qui sont "très bonnes" sur les deux plans.
2. Robustesse : La méthode fonctionne aussi bien sur des petits jeux de données que sur des grands, et avec différents types d'IA.
3. Au-delà de la biologie : Bien que l'article parle de protéines, cette méthode pourrait servir à :
   • Créer des chatbots qui sont à la fois utiles et sûrs.
   • Générer des images qui sont à la fois belles et respectueuses des droits d'auteur.
   • Optimiser des voitures autonomes pour être à la fois rapides et économes.

En résumé

Les chercheurs ont remplacé la vieille règle du "moyenne simple" par une boussole plus fine (STOMP) capable de naviguer dans des terrains complexes. Grâce à cela, ils peuvent guider les intelligences artificielles vers des solutions de compromis parfaites qui étaient invisibles jusqu'alors, que ce soit pour créer de nouveaux médicaments ou pour améliorer nos assistants numériques.

C'est passer de la recherche d'un "bon moyen terme" à la découverte de l'équilibre optimal.

Résumé technique

1. Problématique et Contexte

L'alignement des grands modèles de langage (LLM) avec les préférences humaines repose souvent sur l'apprentissage par renforcement (RL) hors-ligne. Bien que l'alignement à objectif unique soit bien maîtrisé, de nombreuses applications réelles exigent l'optimisation simultanée de plusieurs objectifs conflictuels.

• Exemples : Pour les chatbots, il faut maximiser l'utilité tout en minimisant les risques (sécurité). Pour l'ingénierie des protéines, il faut optimiser simultanément l'activité catalytique, la spécificité et la stabilité.
• Le défi : Aucun modèle unique ne peut optimiser tous les objectifs en même temps. L'objectif est donc de trouver le front de Pareto, c'est-à-dire l'ensemble des solutions non dominées où l'amélioration d'un objectif ne se fait qu'au détriment d'un autre.
• Limitation des méthodes actuelles : La méthode standard consiste à scalariser linéairement les récompenses (moyenne pondérée). Cependant, il a été prouvé mathématiquement que cette approche échoue à récupérer les solutions situées dans les régions non convexes du front de Pareto, qui sont souvent les compromis les plus intéressants.

2. Méthodologie : STOMP

Les auteurs proposent une nouvelle approche nommée STOMP (Smooth Tchebysheff Optimization of Multi-Objective Preferences). Au lieu de scalariser directement les récompenses, ils reformulent le problème de RL multi-objectif lui-même comme un problème d'optimisation à scalariser.

A. Scalarisation Tchebysheff Lisse (STS)

Le papier s'appuie sur la scalarisation Tchebysheff lisse (Smooth Tchebysheff Scalarization - STS), une technique qui surmonte les limites de la scalarisation linéaire en permettant de trouver l'ensemble complet du front de Pareto, y compris les régions non convexes.

• La scalarisation classique de Tchebysheff utilise une fonction min-max non différentiable.
• La version lisse remplace le max par une fonction log-sum-exp, rendant l'optimisation par gradient possible.

B. Reformulation du RL Multi-Objectif

L'innovation clé de STOMP réside dans la manière dont la scalarisation est appliquée :

1. Formulation du problème : Le RL multi-objectif est vu comme la minimisation simultanée de la divergence KL entre la politique apprise $\pi$ et plusieurs politiques optimales théoriques $\pi^*_i$ (une pour chaque objectif).
2. Déduction de la récompense scalarisée : En appliquant la STS à ce problème d'optimisation, les auteurs dérivent une nouvelle fonction de récompense scalarisée $R_{ST}$.
3. Standardisation dynamique : Une difficulté majeure de la STS est sa sensibilité aux échelles des récompenses. STOMP résout cela en standardisant dynamiquement les récompenses individuelles basées sur leurs distributions observées dans l'ensemble de données hors-ligne.
   • Cela évite les hyperparamètres de mise à l'échelle manuelle.
   • La récompense est centrée par rapport à une fonction de partition estimée ($Z_i(x)$), ce qui permet de pénaliser plus sévèrement les séquences rares avec de faibles récompenses dans des distributions asymétriques.

C. Algorithme STOMP

STOMP étend l'algorithme DPO (Direct Preference Optimization) au cadre multi-objectif :

• Il utilise un jeu de données hors-ligne associant chaque séquence à plusieurs récompenses scalaires.
• La fonction de perte (Loss) est une adaptation de OffsetDPO, utilisant la différence de récompense scalarisée $R_{ST}$ pour déterminer les paires "gagnant/perdant".
• La perte intègre une régularisation de vraisemblance négative (NLL) pour éviter la dégradation de la politique par rapport à la distribution de référence, surtout lorsque les données ne sont pas échantillonnées selon la politique de référence.

3. Contributions Clés

1. Reformulation théorique : Première application de la scalarisation Tchebysheff lisse directement au problème d'optimisation du RL multi-objectif, plutôt qu'aux récompenses brutes.
2. Algorithme STOMP : Un algorithme d'apprentissage hors-ligne robuste qui étend le DPO pour gérer plusieurs objectifs conflictuels de manière principielle.
3. Standardisation automatique : Une méthode pour normaliser les récompenses basée sur la distribution empirique, éliminant le besoin de réglage manuel des échelles et permettant de couvrir le front de Pareto complet.
4. Validation empirique rigoureuse : Évaluation sur des tâches complexes d'ingénierie de protéines avec des modèles de langage protéiques (PLM) de différentes tailles.

4. Résultats Expérimentaux

Les auteurs ont évalué STOMP sur trois jeux de données biologiques (DHFR, PbrR, $\alpha$-Amylase) en alignant trois modèles de langage protéiques (ProGen3-3B, ProGen-RA-3B, ProGen-RA-10B).

• Métrique principale : Le hypervolume (volume couvert par le front de Pareto), qui mesure la qualité et la diversité des solutions trouvées.
• Comparaison : STOMP a été comparé à des baselines de pointe utilisant la scalarisation linéaire (DPO-Lin, ODPO-Lin) et une implémentation naive de la scalarisation Tchebysheff (ODPO-STZ).
• Performance :
  • STOMP a atteint le hypervolume le plus élevé dans 8 cas sur 9 (ou a égalé le meilleur résultat).
  • Il a surpassé les autres méthodes de plus de 10 % relatifs dans plusieurs configurations (notamment sur PbrR et $\alpha$-Amylase).
  • Contrairement aux baselines dont les performances étaient incohérentes selon le dataset et le modèle, STOMP s'est montré robuste et capable de couvrir systématiquement le front de Pareto complet.
• Évaluation générative : Lors de la génération de nouvelles protéines et de l'estimation de leurs récompenses via des modèles de régression gaussienne, les modèles alignés avec STOMP ont produit des séquences avec les meilleurs hypervolumes attendus, confirmant que l'amélioration n'est pas seulement artificielle (sur le jeu de test) mais se généralise à la génération.

5. Signification et Impact

• Au-delà de la biologie : Bien que validé sur l'ingénierie des protéines, la méthode est agnostique au domaine. Elle est applicable à tout alignement multi-objectif, y compris l'optimisation des chatbots (aide vs sécurité) ou la génération d'images (qualité vs fidélité au prompt).
• Supériorité sur la scalarisation linéaire : Le papier démontre empiriquement et théoriquement que la scalarisation linéaire est insuffisante pour les compromis complexes (régions non convexes), et que STOMP offre une solution mathématiquement fondée pour explorer l'ensemble des compromis possibles.
• Futur de l'alignement : STOMP ouvre la voie à des modèles post-entraînés capables de naviguer dynamiquement dans des espaces de compromis multi-attributes, rendant les IA plus adaptables à des contraintes du monde réel souvent contradictoires.

En résumé, STOMP représente une avancée majeure dans l'apprentissage par renforcement hors-ligne multi-objectif, offrant un cadre robuste pour aligner les modèles de langage sur des préférences complexes et conflictuelles sans sacrifier la couverture de l'espace des solutions optimales.