🤖 machine learning

$f$ -Trajectory Balance: A Loss Family for Tuning GFlowNets, Generative Models, and LLMs with Off- and On-Policy Data

Ce papier introduit l'équilibre de trajectoire $f$ , une famille de fonctions de perte qui étend l'approche de l'erreur quadratique moyenne à toutes les $f$ -divergences, permettant l'entraînement de modèles génératifs et de LLM avec des données hors politique tout en préservant les propriétés d'optimisation spécifiques (telles que la couverture des modes) de leurs gradients de $f$ -divergence correspondants en politique.

Auteurs originaux : Jake Fawkes, Jason Hartford

Publié 2026-05-18

📖 5 min de lecture🧠 Analyse approfondie

CC BY 4.0

Auteurs originaux : Jake Fawkes, Jason Hartford

Article original sous licence CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ Ceci est une explication générée par l'IA de l'article ci-dessous. Elle n'a pas été rédigée ni approuvée par les auteurs. Pour une précision technique, consultez l'article original. Lire la clause de non-responsabilité complète

Imaginez que vous enseigniez à un robot à explorer un vaste système de grottes sombres. La grotte possède de nombreuses chambres cachées (appelées « modes »), dont certaines contiennent de l'or (récompenses élevées), tandis que d'autres sont simplement vides ou dangereuses. Votre objectif est d'enseigner au robot de trouver le plus grand nombre possible de chambres à or différentes, plutôt que de se contenter de découvrir une seule grande mine d'or et d'ignorer tout le reste.

Ce papier présente un nouvel ensemble d'« outils pédagogiques » (fonctions de perte) pour aider à entraîner ces robots, qu'ils génèrent des molécules pour de nouveaux médicaments, créent de l'art ou écrivent du code.

Voici la décomposition des idées du papier à l'aide d'analogies simples :

1. Le Problème : La « Ruée vers l'or » contre l'« Explorateur »

Par le passé, la méthode standard pour enseigner à ces robots ressemblait à une Ruée vers l'or. Le robot trouvait un endroit avec un peu d'or, s'enthousiasmait, puis passait tout son temps à creuser là-bas. Il ignorait toutes les autres mines d'or de la grotte.

En termes techniques : Cela s'appelle la « recherche de modes ». Le modèle se collapse sur quelques réponses à haute probabilité et ignore le reste de la diversité.
L'objectif du papier : Nous voulons une « couverture des modes ». Nous voulons que le robot se disperse et explore toutes les mines d'or, même les plus petites, pour obtenir une carte complète de la grotte.

2. L'Ancien Outil : La Perte « KL au Carré »

Récemment, des chercheurs ont découvert une astuce ingénieuse pour enseigner aux robots en utilisant une méthode d'« erreur quadratique » (mesurant la distance entre ce que le robot pense et ce qu'il devrait penser).

L'Analogie : Imaginez un professeur notant un élève. Si l'élève se trompe, le professeur ne dit pas simplement « Faux ». Il calcule le carré de l'erreur.
L'Avantage : Cette méthode est très stable. Elle fonctionne même si l'élève apprend à partir d'anciennes notes (données hors politique) plutôt que de leçons en direct.
Le Défaut : Même si elle est stable, cette méthode spécifique « au carré » agit toujours comme une Ruée vers l'or. Elle pousse naturellement le robot à se concentrer sur quelques réponses seulement, manquant ainsi la diversité de la grotte.

3. La Nouvelle Solution : La Famille « f-Trajectory Balance »

Les auteurs ont réalisé que la méthode « au carré » n'était qu'une saveur spécifique d'une famille beaucoup plus vaste d'outils pédagogiques. Ils appellent cette famille f-Trajectory Balance.

Pensez à cette famille comme à un cadran ou à un bouton de volume sur une radio.

Tourner le cadran dans un sens (Nombres bas) : Vous obtenez des outils qui agissent comme un Super Explorateur. Ces outils forcent le robot à se disperser et à trouver chaque mine d'or possible, même les plus difficiles d'accès. C'est idéal pour la découverte de médicaments où vous voulez trouver n'importe quelle molécule qui pourrait fonctionner, pas seulement la plus évidente.
Tourner le cadran dans l'autre sens (Nombres élevés) : Vous obtenez des outils qui agissent comme un Chercheur d'or. Ces outils disent au robot d'ignorer les petites mines et de se concentrer intensément sur le plus gros tas d'or, le plus évident. Cela est utile si vous ne voulez que la meilleure réponse unique.
Le Milieu de Terrain : Vous pouvez régler le cadran n'importe où entre les deux pour obtenir un mélange d'exploration et de concentration.

4. Pourquoi C'est Important

Le papier prouve deux choses principales :

Le « Bouton Magique » : Ils ont montré que vous pouvez remplacer l'outil « au carré » standard par n'importe lequel de ces nouveaux outils « à cadran », et les mathématiques fonctionnent toujours parfaitement. Le robot apprend aussi bien, mais avec une personnalité différente (plus exploratrice ou plus concentrée).
Stabilité : Tout comme l'ancien outil, ces nouveaux outils fonctionnent même si le robot apprend à partir d'anciennes données (hors politique). Ceci est crucial pour les applications réelles comme les Modèles de Langage à Grande Échelle (LLM), où l'entraînement se produit souvent de manière asynchrone (le robot apprend à partir de données générées il y a un moment).

5. Tests Réels (Les Cartes de Grottes)

Les auteurs ont testé ces outils dans trois « grottes » différentes :

Le Jeu de Grille : Un labyrinthe informatique simple avec quatre coins remplis d'or. L'outil standard n'a trouvé qu'un seul coin. Les nouveaux outils « Explorateur » ont trouvé les quatre coins rapidement.
Découverte de Molécules (SynFlowNet) : Ils ont tenté de générer de nouvelles molécules chimiques. En tournant le cadran sur « Explorateur », ils ont généré une variété beaucoup plus large de molécules uniques qui pourraient potentiellement être des médicaments, plutôt que de simples variations des mêmes quelques produits chimiques.
Modèles de Langage (LLM) : Ils ont ajusté des modèles d'IA pour résoudre des problèmes de mathématiques. Les nouveaux outils ont permis à l'IA d'explorer différentes façons de résoudre des problèmes sans rester coincée dans une boucle répétant la même réponse, même lorsque les données d'entraînement étaient décalées (asynchrones).

Résumé

Ce papier n'invente pas un nouveau type de robot. Au contraire, il invente un nouvel ensemble d'instructions sur la façon de les entraîner.

Ancienne méthode : « Trouvez la plus grande mine d'or et creusez là-bas. » (Stable, mais ennuyeux).
Nouvelle méthode : « Voici un cadran. Vous pouvez choisir d'être un Chercheur d'or, un Super Explorateur, ou n'importe quoi entre les deux. Et peu importe votre choix, l'entraînement sera stable et fonctionnera avec d'anciennes données. »

Cela donne aux ingénieurs un bouton simple pour contrôler à quel point leurs modèles d'IA doivent être créatifs ou concentrés, sans casser le processus d'entraînement.

Résumé technique : f-Trajectory Balance

Énoncé du problème

L'alignement et le fine-tuning des modèles génératifs, y compris les grands modèles de langage (LLM) et les réseaux de flux génératifs (GFlowNets), reposent généralement sur la minimisation de l'écart entre la distribution d'un modèle et une distribution cible. Dans le fine-tuning par apprentissage par renforcement (RL), cela est souvent formulé comme la maximisation de la récompense espérée soumise à une pénalité de divergence de Kullback-Leibler (KL).

Un défi majeur dans ce domaine est le compromis entre les comportements de recherche de modes (mode-seeking) et de couverture de modes (mode-covering). Les objectifs standards, tels que la divergence KL inverse (souvent mise en œuvre via Trajectory Balance ou la perte KL au carré), ont tendance à être de type recherche de modes, ce qui provoque l'effondrement des modèles sur quelques modes à haute probabilité. Cela est sous-optimal pour les tâches nécessitant une exploration, telles que la découverte de médicaments ou la génération de candidats diversifiés, où la « couverture de modes » (capturer toute la diversité de la distribution cible) est préférée.

De plus, les méthodes existantes ont souvent du mal avec la validité hors politique (off-policy). Bien que des estimateurs comme REINFORCE ou PPO soient courants, ils souffrent d'une variance élevée. Les approches récentes utilisant la perte KL au carré ( $KL_{sq}$ ) offrent des gradients à faible variance en politique et une validité hors politique, mais elles sont intrinsèquement liées aux propriétés de la divergence KL inverse, limitant ainsi la capacité à contrôler le compromis exploration-exploitation.

Méthodologie

Les auteurs proposent f-Trajectory Balance, une famille de fonctions de perte de substitution qui généralise la perte KL au carré à l'ensemble de la famille des divergences f.

Fondement théorique

La contribution théorique centrale est l'établissement d'une correspondance univoque entre les fonctions de perte invariantes par translation sur les log-probabilités et les divergences f.

De la divergence à la perte : Pour toute fonction convexe $f$ définissant une divergence f $D_f(p_\theta \| p^*)$ , les auteurs dérivent une fonction de perte $L_f$ définie sur la différence des log-probabilités $\Delta_\theta(y) = \log p_\theta(y) - \log p^*(y)$ :
$L_f(\Delta_\theta(y)) = \int_0^{\Delta_\theta(y)} (f'(e^t) - f'(1)) dt$
Cette perte possède deux propriétés critiques :
- Équivalence des gradients en politique : Lorsque les échantillons sont tirés de la politique du modèle ( $\mu = p_\theta$ ), les gradients auto-différenciés espérés de $L_f$ correspondent aux gradients de la divergence f correspondante $D_f$ .
- Validité hors politique : Lorsqu'elle est évaluée sur des données hors politique (échantillonnées à partir d'une distribution de comportement $\mu$ ), la perte reste un objectif valide avec le même minimiseur global ( $p_\theta = p^*$ ), à condition que le support de la cible soit contenu dans le support de la politique de comportement.
De la perte à la divergence : Inversement, toute perte $\ell$ strictement convexe et invariante par translation sur les log-probabilités correspond à la minimisation d'une divergence f spécifique. L'article prouve que ces applications sont l'inverse l'une de l'autre.

Implémentations pratiques

Pour répondre aux défis pratiques liés à l'entraînement, les auteurs introduisent plusieurs extensions :

Perte DevGrad (Normalisation par lot) : Pour gérer les fonctions de partition intraitables (courantes dans les GFlowNets et les cibles non normalisées), les auteurs généralisent l'estimateur VarGrad. Au lieu d'utiliser la variance par lot, ils utilisent la déviation généralisée par lot des différences de log-probabilités. Cela implique d'estimer la constante de normalisation $\log Z$ en minimisant la perte sur le lot, ce qui centre efficacement les coefficients de la fonction de score et réduit la variance.
Perte tempérée : Pour éviter les débordements numériques lors de l'utilisation de petits paramètres de température $\beta$ (où la distribution cible est $p^* \propto \pi_{ref} \exp(r/\beta)$ ), les auteurs proposent une perte tempérée. Cela consiste à mettre l'échelle de la perte par $1/\beta$ et à l'appliquer à la distribution tempérée, garantissant que les gradients restent bornés et stables.
f-Trajectory Balance pour les GFlowNets : Le cadre est appliqué aux GFlowNets en imposant des contraintes d'équilibre de trajectoire à l'aide de la perte $L_f$ . Cela généralise le Trajectory Balance standard (qui correspond à la divergence KL inverse) à des divergences f arbitraires.

Contributions clés

Famille de pertes généralisée : Dérivation d'une famille générale de pertes de substitution invariantes par translation ( $L_f$ ) dont les gradients en politique correspondent à ceux de la divergence f correspondante, tout en conservant la validité hors politique.
Équivalence théorique : Preuve de la relation inverse entre les pertes de log-probabilité invariantes par translation et les divergences f, permettant la conception systématique de nouvelles fonctions de perte aux propriétés souhaitées (par exemple, couverture de modes).
Extension aux cibles non normalisées : Généralisation de la technique VarGrad à des divergences f arbitraires via la perte DevGrad, résolvant le problème de la fonction de partition intraitable dans le fine-tuning par RL.
Contrôle de l'exploration : Démonstration que le choix de la divergence f (paramétré par $\alpha$ dans les $\alpha$ -divergences) contrôle directement le comportement de recherche de modes par rapport à la couverture de modes. Des valeurs de $\alpha$ plus faibles (par exemple, KL direct, Hellinger) favorisent la couverture de modes, tandis que des valeurs plus élevées (par exemple, $\chi^2$ de Pearson) favorisent la recherche de modes.

Résultats expérimentaux

Les auteurs évaluent la famille de pertes sur quatre tâches distinctes :

Grille synthétique (GFlowNets) : Sur une tâche de grille 2D conçue pour tester la couverture de modes, les pertes avec un $\alpha$ plus faible (KL direct, Hellinger) ont découvert avec succès les quatre modes de la distribution cible beaucoup plus rapidement et plus fièrement que le Trajectory Balance standard (KL inverse). À l'inverse, le $\chi^2$ de Pearson ( $\alpha=2$ ) a montré un comportement fort de recherche de modes, restant coincé dans le premier mode trouvé.
Découverte de molécules (SynFlowNet) : Appliqué au cadre SynFlowNet pour générer des molécules synthétiquement accessibles, les auteurs ont constaté que le réglage de $\alpha$ permettait de contrôler la diversité des molécules générées. Un calendrier de recuit (commençant par un $\alpha$ plus faible pour l'exploration et augmentant jusqu'à $\alpha=1$ pour l'exploitation) a produit le meilleur compromis, générant des molécules à plus haute récompense avec une plus grande diversité par rapport au Trajectory Balance standard.
Modèles de diffusion : Dans l'échantillonnage conditionnel de chiffres MNIST (impairs vs pairs), le Trajectory Balance standard a échoué à échantillonner uniformément à travers les modes cibles (suréchantillonnant des chiffres spécifiques comme 0 et 6). Les pertes f-Trajectory Balance, en particulier celles avec un $\alpha$ plus faible, ont permis un échantillonnage plus uniforme à travers les modes cibles.
Ajustement asynchrone de LLM : Les pertes ont été appliquées pour affiner des LLM (familles Qwen et OLMo) sur des problèmes de mathématiques (GSM8k, MATH) avec un délai asynchrone de 50 étapes. Les pertes de divergence f ont démontré un entraînement stable et des compromis cohérents entre entropie et récompense, tandis que les implémentations PPO standard ont montré une instabilité dans les mêmes conditions asynchrones.

Importance et affirmations

L'article affirme que l'efficacité de la perte KL au carré n'est pas un phénomène isolé, mais un cas spécifique d'une famille plus large de pertes. En étendant la propriété de « correspondance des gradients en politique » et de « validité hors politique » à l'ensemble de la famille des divergences f, les auteurs fournissent un cadre unifié pour le réglage des modèles génératifs.

L'importance principale réside dans la capacité à contrôler explicitement le compromis exploration-exploitation simplement en sélectionnant une fonction de perte de substitution différente (ou en réglant un paramètre $\alpha$ ), sans modifier l'algorithme d'entraînement sous-jacent. Cela permet aux praticiens de choisir entre des comportements de recherche de modes (pour la maximisation de la récompense) et de couverture de modes (pour la diversité et l'exploration), tout en maintenant la stabilité et la compatibilité hors politique requises pour l'entraînement asynchrone à grande échelle moderne des LLM et des GFlowNets. Ce travail suggère que ce cadre est particulièrement précieux pour des applications telles que la découverte de médicaments et les agents exploratoires où la couverture de l'espace des candidats à haute récompense est critique.