GFlowState: Visualizing the Training of Generative Flow Networks Beyond the Reward

1. Problématique

Les Réseaux de Flux Génératifs (GFlowNets) sont un cadre probabiliste puissant conçu pour échantillonner des objets (comme des molécules ou des matériaux) proportionnellement à une fonction de récompense. Contrairement aux méthodes d'apprentissage par renforcement classiques qui cherchent une seule solution optimale, les GFlowNets visent à générer un ensemble diversifié de solutions à haute valeur.

Cependant, le entraînement des GFlowNets reste difficile à interpréter et à diagnostiquer. Les outils d'analyse de machine learning standards (comme TensorBoard ou Weights & Biases) permettent de suivre des métriques agrégées (pertes, récompenses moyennes), mais ils échouent à révéler :

• Comment le modèle explore l'espace des échantillons.
• La construction des trajectoires d'échantillonnage.
• L'évolution des probabilités de transition entre les états.
• Les dynamiques de "collapse de mode" (où le modèle se concentre excessivement sur une seule région à haute récompense).

La complexité combinatoire de l'espace d'états et la nature séquentielle des décisions rendent l'inspection manuelle impossible, créant un besoin critique pour des méthodes d'analyse visuelle adaptées.

2. Méthodologie : GFlowState

Les auteurs proposent GFlowState, un système d'analyse visuelle interactive conçu pour rendre observables les dynamiques structurelles de l'entraînement des GFlowNets. Le système repose sur quatre vues principales qui interagissent entre elles :

1. Sample Ranking (Classement des échantillons) :

   • Un graphique en "bump chart" montrant l'évolution du classement des objets générés (par récompense, perte, ou métrique personnalisée) au fil des itérations.
   • Permet d'identifier visuellement les découvertes de nouveaux modes à haute récompense et les phénomènes de stagnation (lignes horizontales).

2. State Projection (Projection d'état) :

   • Réduit la dimensionnalité des échantillons générés (et d'un jeu de données de validation si disponible) en 2D via des techniques comme UMAP.
   • Les points sont agrégés dans des hexagones (hexbins) pour éviter la surcharge visuelle.
   • Le codage couleur permet d'analyser localement la corrélation entre la probabilité d'échantillonnage et la récompense, la perte, ou la couverture par rapport à un jeu de validation. Cela aide à détecter les régions sous-explorées ou mal modélisées.

3. DAG View (Vue du graphe acyclique dirigé) :

   • Représente les trajectoires d'échantillonnage sous forme de graphe acyclique dirigé (DAG).
   • Utilise une exploration interactive "bottom-up" : l'utilisateur part de l'état source et peut développer (expand) les nœuds pour inspecter les sous-graphes.
   • Les chaînes linéaires sont agrégées pour simplifier la visualisation. Cette vue permet d'analyser les décisions du modèle étape par étape (Traj 1 et Traj 2).

4. Transition Heatmap (Carte thermique des transitions) :

   • Affiche les transitions (arêtes du DAG) classées par une métrique spécifique (probabilité, fréquence, variance) sur l'axe des itérations.
   • Permet de visualiser l'évolution temporelle des probabilités de transition, aidant à détecter les changements de politique ou les modes de collapse (Dyn 1 et Dyn 3).

Flux de travail :
Le système fonctionne en deux phases :

• Pendant l'entraînement : Les échantillons et leurs attributs sont enregistrés dans une base de données SQL. Les vues de classement et de projection sont disponibles en temps réel.
• Après l'entraînement : Le DAG complet est calculé à partir de toutes les trajectoires enregistrées pour permettre l'analyse approfondie des transitions et des structures de trajectoires.

3. Contributions Clés

• Identification de tâches d'analyse : Définition rigoureuse des besoins des développeurs de GFlowNets, notamment l'analyse des trajectoires, l'exploration de l'espace d'échantillonnage et l'analyse des dynamiques d'entraînement.
• Développement de GFlowState : Un prototype interactif intégrant les quatre vues décrites ci-dessus, conçu pour être intégré facilement dans les flux de travail des développeurs (via Python/Plotly Dash).
• Évaluation par études de cas : Validation du système sur deux domaines distincts :
  1. Un environnement de grille (2D) pour la démonstration conceptuelle.
  2. Un environnement de cristaux (composition chimique, groupe d'espace, paramètres de maille) pour une application scientifique réelle.

4. Résultats

Les études de cas menées avec des chercheurs de l'Institut Mila (Québec) ont démontré l'efficacité de GFlowState :

• Débogage accéléré : Dans l'environnement de grille, les experts ont pu identifier instantanément les itérations où de nouveaux modes de récompense étaient découverts (vers l'itération 4 500), un processus qui nécessitait auparavant un chargement manuel et fastidieux de nombreux points de contrôle.
• Analyse de la diversité : Dans l'étude sur les cristaux, le système a permis de visualiser comment le modèle passait d'une exploration initiale limitée à une couverture plus large de l'espace d'états. La comparaison avec un jeu de données de validation a révélé des régions sous-représentées.
• Compréhension des trajectoires : La vue DAG a été jugée particulièrement utile par les experts pour comprendre les dépendances entre les sous-environnements (ex: comment le choix d'un groupe d'espace influence la composition chimique) et pour vérifier la plausibilité physique des états générés.
• Détection de pathologies : Le système a permis de détecter visuellement le "mode collapse" (le modèle se focalisant sur un seul mode) et d'observer comment les probabilités de transition évoluaient pour corriger ce biais au fil de l'entraînement.

5. Signification et Impact

Ce travail comble un vide important dans l'écosystème de l'apprentissage automatique pour la découverte scientifique.

• Interprétabilité : Il transforme les GFlowNets, souvent considérés comme des "boîtes noires" complexes, en systèmes dont la dynamique structurelle est observable.
• Accélération du développement : En fournissant des outils de diagnostic visuel, GFlowState permet aux chercheurs d'identifier plus rapidement les échecs d'entraînement, d'ajuster les hyperparamètres et de valider la qualité des modèles.
• Généralisabilité : Bien que conçu pour les GFlowNets, l'approche démontre la valeur des solutions d'analyse visuelle spécifiques à un modèle pour les paradigmes d'apprentissage complexes. Le code est open-source et le système est en cours d'intégration dans la bibliothèque officielle des GFlowNets, facilitant son adoption par la communauté scientifique.

En résumé, GFlowState n'est pas seulement un outil de visualisation, mais un instrument de diagnostic essentiel qui améliore la fiabilité et l'efficacité de l'utilisation des GFlowNets dans des applications critiques comme la découverte de médicaments et de matériaux.