Rooted Absorbed Prefix Trajectory Balance with Submodular Replay for GFlowNet Training

Cet article propose RapTB, une nouvelle fonction objectif ancrée à la racine pour améliorer l'assignation de crédit aux préfixes, couplée à une stratégie de replay submodulaire (SubM) pour atténuer les biais de distribution, permettant ainsi aux GFlowNets d'entraîner des modèles de langage générant des molécules valides et diversifiées tout en évitant l'effondrement des modes.

L'essentiel

Imaginez que vous essayez d'enseigner à un robot très intelligent (une Intelligence Artificielle) comment créer des choses nouvelles et utiles, comme de nouvelles molécules pour des médicaments ou des phrases intelligentes. Le robot utilise une méthode appelée GFlowNet.

Le problème, c'est que ce robot a tendance à devenir paresseux et répétitif. Au lieu d'explorer toutes les possibilités, il s'entête sur une seule idée qui lui a plu au début et ne fait que la copier encore et encore. C'est ce qu'on appelle un "effondrement de mode" (mode collapse). Il produit soit des phrases trop courtes, soit des molécules qui ne fonctionnent pas, car il a oublié d'explorer le reste du monde.

Les auteurs de cette paper ont inventé deux solutions magiques pour réveiller le robot : RapTB et SubM.

Voici comment ça marche, avec des analogies simples :

1. Le Problème : Le Robot qui s'endort

Imaginez que le robot est un écrivain qui doit écrire un roman.

• Le problème actuel : Il écrit les premiers mots ("Il était une fois..."), puis il reçoit une récompense seulement à la toute fin du livre. S'il écrit un livre court et ennuyeux qui finit bien, il reçoit son prix.
• La conséquence : Il se dit : "Ah, écrire court et simple rapporte des points !" Alors, il arrête d'essayer d'écrire des histoires complexes et longues. Il s'arrête toujours après 3 mots. C'est le biais de longueur.
• Autre problème : Il s'assoit sur la même chaise (le même début de phrase) pour écrire tous ses livres. Personne ne veut lire la même histoire 1000 fois. C'est l'effondrement du préfixe.

2. La Solution 1 : RapTB (Le Guide de Voyage Intérieur)

Pour corriger cela, les auteurs ont créé RapTB.

• L'analogie : Imaginez que le robot est un randonneur dans une forêt. Avant, il ne recevait un prix qu'en arrivant au sommet de la montagne (la fin de la phrase). S'il tombait en chemin, il ne savait pas s'il avait bien fait ou non.
• La magie de RapTB : Avec RapTB, le randonneur reçoit des indices de récompense à chaque pas.
  • Si le randonneur fait un pas vers la droite, le système lui dit : "Attends, si tu continues dans cette direction, tu pourrais atteindre un trésor plus gros plus tard !"
  • C'est comme si le randonneur pouvait voir un peu plus loin dans le brouillard (le "suffixe absorbé"). Il sait que même s'il n'est pas encore au sommet, son chemin actuel est prometteur.
  • Résultat : Le robot n'arrête plus trop tôt. Il ose écrire des phrases plus longues et plus complexes, car il comprend que chaque mot compte pour la récompense finale.

3. La Solution 2 : SubM (Le Collectionneur de Trésors Divers)

Même avec de bons indices, le robot risque de s'entêter sur les mêmes chemins s'il ne garde en mémoire que les meilleurs exemples. C'est là qu'intervient SubM.

• L'analogie : Imaginez un musée (la mémoire du robot).
  • L'ancienne méthode : Le musée ne gardait que les 100 tableaux les plus beaux (ceux avec la plus haute récompense). Résultat ? Le musée était rempli de 100 copies presque identiques du même tableau. C'était ennuyeux et ne servait pas à apprendre la diversité.
  • La méthode SubM : Le nouveau conservateur (SubM) utilise une règle intelligente. Il dit : "Je veux garder les tableaux les plus beaux, MAIS je veux aussi qu'ils soient tous différents les uns des autres, et qu'ils aient des tailles variées."
  • Il utilise une mathématique spéciale (appelée "submodulaire") pour choisir un échantillon parfait : un peu de beauté, beaucoup de diversité, et des longueurs variées.
  • Résultat : Le robot apprend à partir d'une collection riche et variée, pas juste d'une seule idée répétée.

En Résumé

En combinant RapTB (qui donne des indices précis à chaque étape pour ne pas s'arrêter trop tôt) et SubM (qui s'assure que le robot apprend d'une variété d'exemples et non d'une seule), les auteurs ont réussi à :

1. Éviter que le robot ne produise des phrases trop courtes.
2. Éviter qu'il ne répète toujours la même chose.
3. Générer des molécules chimiques plus efficaces et des textes plus créatifs.

C'est comme passer d'un élève qui recopie bêtement la même phrase pour avoir une bonne note, à un artiste qui explore tout un tableau de possibilités pour créer une œuvre d'art unique et complète.

Résumé technique

1. Problématique : L'Effondrement des Modes dans les GFlowNets pour les LLM

Les Réseaux de Flux Génératifs (GFlowNets) sont conçus pour apprendre des politiques stochastiques qui échantillonnent des objets (comme des séquences de tokens) avec une probabilité proportionnelle à leur récompense. Contrairement à l'apprentissage par renforcement classique qui cherche un optimum unique, les GFlowNets visent à explorer l'ensemble des modes à haute récompense.

Cependant, lorsqu'ils sont appliqués aux Grands Modèles de Langage (LLMs) via une formulation d'arbres de préfixes terminables, les GFlowNets souffrent de deux modes d'échec majeurs, qualifiés d'effondrement des modes :

1. Effondrement des préfixes (Prefix Collapse) : L'entropie des premiers tokens chute drastiquement. Différentes terminaisons partagent des préfixes quasi identiques, réduisant la diversité de l'exploration.
2. Biais de longueur (Length Bias) : Le modèle favorise systématiquement des séquences trop courtes ou trop longues, souvent en raison d'une mauvaise attribution du crédit.

Causes identifiées :

• Attribution de crédit faible : Les récompenses ne sont données qu'à la fin de la trajectoire (terminal), ce qui crée un bruit élevé et une attribution ambiguë pour les étapes intermédiaires.
• Biais de la mémoire de replay (Replay Bias) : L'utilisation de buffers de replay standard (priorisés par la récompense) concentre l'entraînement sur un sous-ensemble étroit de trajectoires, induisant un décalage de distribution et renforçant l'effondrement.

---

2. Méthodologie Proposée

Les auteurs proposent une approche double combinant un nouvel objectif d'apprentissage (RapTB) et une stratégie de mise à jour du buffer de replay (SubM).

A. RapTB : Rooted Absorbed Prefix Trajectory Balance

RapTB vise à densifier les signaux d'apprentissage sans introduire de contraintes contradictoires qui déstabiliseraient l'arrêt de la génération.

1. Préfixes Racines (Rooted Prefixes) : Contrairement au Subtrajectory Balance (SubTB) qui impose des contraintes sur toutes les fenêtres de sous-trajectoires (créant des conditions aux limites hétérogènes et conflictuelles), RapTB restreint la supervision dense aux préfixes racines (partant de l'état initial $s_0$). Cela élimine les conflits de conditions aux limites.
2. Récompenses Suffixes Absorbées (Absorbed Suffix Rewards) : Pour stabiliser l'attribution du crédit, RapTB propage la récompense terminale vers les préfixes intermédiaires via des mécanismes de "backups" :
   • Max Backup : La récompense d'un préfixe est au moins égale à la meilleure récompense observée dans son suffixe.
   • Soft Backup : Une agrégation lissée (log-sum-exp) des récompenses du suffixe, pondérée par la distance.
   • Cible Mixte : Une combinaison des deux pour fournir un signal de crédit dense et à faible variance.
3. Découplage des Gradients d'Arrêt : Une innovation cruciale est l'arrêt des gradients (stop-gradient) sur les logits de terminaison dans la branche auxiliaire. Cela empêche le modèle de satisfaire les contraintes de sous-trajectoires en déformant arbitrairement la probabilité d'arrêt (ce qui cause le biais de longueur), forçant ainsi l'objectif global Trajectory Balance (TB) à rester le seul ancrage pour la calibration de l'arrêt.

B. SubM : Submodular Replay

Pour contrer le biais induit par le replay classique (qui favorise les "riches" modes), les auteurs introduisent SubM, une stratégie de rafraîchissement du buffer basée sur l'optimisation submodulaire.

• Objectif : Sélectionner un sous-ensemble de trajectoires du buffer actuel et d'un nouveau lot généré, en maximisant une fonction submodulaire monotone.
• Fonction Objectif : Elle combine trois termes :
  1. Qualité/Récompense : Privilégier les trajectoires à haute récompense.
  2. Diversité (Facility Location) : Maximiser la couverture de l'espace des solutions (en utilisant des métriques de similarité comme les empreintes moléculaires Tanimoto ou la similarité Jaccard).
  3. Couverture de Longueur : Assurer une répartition équilibrée des longueurs de séquences pour éviter le biais de longueur.
• Efficacité : L'algorithme utilise une approche gloutonne (greedy) avec une garantie de sous-optimalité, permettant une mise à jour rapide du buffer.

---

3. Contributions Clés

1. Caractérisation de l'effondrement : Identification et preuve empirique que l'effondrement des modes dans les LLM-GFlowNets est une combinaison reproductible d'effondrement de préfixe et de biais de longueur, pilotés par une attribution de crédit à haute variance et des décalages de distribution dus au replay.
2. RapTB : Un nouvel objectif qui améliore le Trajectory Balance (TB) en ajoutant une supervision de préfixe racine avec des cibles absorbées. Cela fournit des signaux d'apprentissage denses tout en évitant les contraintes de terminaison instables du SubTB.
3. SubM : Une stratégie de replay qui équilibre récompense, diversité et couverture de longueur via l'optimisation submodulaire, améliorant la robustesse de l'exploration.
4. Résultats Empiriques : Démonstration que la combinaison RapTB + SubM surpasse les méthodes existantes (TB, SubTB) sur des tâches de génération de molécules (SMILES), d'expressions arithmétiques (Expr24) et de génération de texte (CommonGen).

---

4. Résultats Expérimentaux

Les expériences ont été menées sur trois tâches principales :

• Génération de Molécules (SMILES) :

  • Performance : RapTB + SubM atteint le meilleur compromis qualité-diversité tout en maintenant une validité chimique élevée (>98%).
  • Diversité : Contrairement au TB standard qui se concentre sur des séquences courtes, RapTB + SubM explore efficacement des longueurs plus longues et maintient une diversité moléculaire (FPDiv) supérieure.
  • Stabilité : Réduction significative de l'effondrement des préfixes (mesuré par la masse Top-1 et l'entropie des préfixes).

• Génération d'Expressions Arithmétiques (Expr24) :

  • Couverture : RapTB + SubM double la couverture normalisée des solutions correctes par rapport aux meilleures bases de référence, tout en conservant une précision proche de 100%.
  • Calibration : Élimination du "dérive de terminaison" (termination drift) observée avec SubTB, où le modèle apprenait à arrêter la génération à des moments inappropriés pour satisfaire les contraintes locales.

• CommonGen (Génération de phrases) :

  • Robustesse : Le modèle préserve les priors linguistiques naturels (longueur de phrase, arrêt naturel) tout en optimisant la récompense, là où SubTB échoue catastrophiquement en produisant des phrases excessivement longues et incohérentes.

---

5. Signification et Impact

Ce travail adresse un problème fondamental dans l'application des GFlowNets aux modèles de langage : la tension entre l'exploration de modes multiples et la stabilité de l'entraînement.

• Théorique : Il propose une solution élégante au problème des contraintes conflictuelles dans les arbres de préfixes terminables, en montrant que la supervision dense doit être ancrée à la racine et que les gradients de terminaison doivent être protégés dans les branches auxiliaires.
• Pratique : La méthode RapTB + SubM offre un cadre robuste pour l'optimisation de LLMs dans des domaines nécessitant une exploration diverse (découverte de médicaments, conception de matériaux, raisonnement logique), permettant de trouver des solutions de haute qualité sans sacrifier la diversité ni la validité.
• Futur : Cela ouvre la voie à des objectifs d'apprentissage qui couplent explicitement le replay conscient de la couverture (coverage-aware) avec l'apprentissage de sous-structures adaptatives pour un entraînement plus stable et efficace.

En résumé, le papier démontre que pour entraîner efficacement des GFlowNets sur des LLMs, il est impératif de corriger à la fois la signalisation interne (via RapTB) et la distribution externe des données d'entraînement (via SubM) pour éviter l'effondrement des modes.