Your GFlowNet Secretly Learns an Optimal Transport Plan

Résumé Technique : Votre GFlowNet apprend secrètement un plan de transport optimal

Énoncé du Problème
Les réseaux de flux génératifs (GFlowNets) sont conçus pour échantillonner des objets discrets structurés via des trajectoires stochastiques dans un graphe orienté. Bien qu'initialement développés pour les graphes acycliques, des travaux récents ont étendu les GFlowNets aux environnements non acycliques. Dans ces contextes, minimiser la longueur attendue des trajectoires est crucial pour l'efficacité de l'échantillonnage, menant à la formulation des GFlowNets à « flux minimum ». Parallèlement, le transport optimal (OT) fournit un cadre rigoureux pour comparer des distributions de probabilité en minimisant un coût de transport, généralement résolu via la programmation linéaire (formulation de Kantorovich). Cependant, la résolution de l'OT sur de grands graphes est souvent informatiquement intraitable à l'aide de solveurs exacts.

Ce document traite de l'écart entre ces deux domaines en étudiant si la formulation à flux minimum des GFlowNets non acycliques résout intrinsèquement un problème de transport optimal.

Méthodologie
Les auteurs établissent une équivalence théorique entre les GFlowNets non acycliques et le transport optimal discret sur les graphes. La méthodologie centrale implique :

1. Reformulation Théorique : Les auteurs partent du problème d'optimisation de flux minimum pour les GFlowNets non acycliques, qui minimise le flux total (équivalent à minimiser la longueur attendue de la trajectoire) sous contraintes de respect du bilan détaillé et de correspondance des récompenses. En fixant la distribution initiale du flux d'arêtes (le flux sortant de l'état source $s_0$) pour correspondre à une distribution source spécifique $L(u)$, ils transforment l'objectif du GFlowNet en un problème de programmation linéaire (LP).
2. Preuve d'Équivalence : Ils démontrent que ce problème de programmation linéaire contraint est mathématiquement équivalent au problème de transport optimal de Kantorovich. Dans cette formulation :
   • La distribution source est $L(u)$ définie sur un ensemble d'états initiaux $U$.
   • La distribution cible est la distribution de récompense $R(x)$ définie sur les états terminaux $X$.
   • Le coût de transport $d(u, x)$ est défini comme la distance du plus court chemin entre $u$ et $x$ dans le graphe.
   • La politique optimale du GFlowNet apprend implicitement le couplage optimal $\Pi$ qui minimise le coût de transport total.
3. Algorithme d'Apprentissage : Pour résoudre cela sur de grands graphes où la programmation linéaire exacte est infaisable, les auteurs emploient une paramétrage neuronal du GFlowNet. Ils utilisent l'objectif de l'équilibre des trajectoires (Trajectory Balance - TB) avec une régularisation de flux. Une modification clé est l'inclusion d'un terme « leward » (récompense vers la gauche) dans la fonction de perte pour imposer la distribution initiale fixe $L(u)$, parallèlement à la correspondance de récompense standard pour la distribution cible $R(x)$.

Contributions Clés

1. Équivalence Théorique : Le papier prouve que fixer la distribution de flux initiale dans un GFlowNet à flux minimum réduit exactement son objectif à un problème de transport optimal de Kantorovich avec des coûts de plus court chemin induits par le graphe. Par conséquent, la politique directe optimale du GFlowNet encode un plan de transport optimal, et l'échantillonnage de trajectoires à partir de celle-ci permet de récupérer le couplage optimal.
2. Lien avec le Problème de Beckmann : Les auteurs montrent que cette formulation est équivalente à la formulation discrète du problème de Beckmann sur graphe, reliant les GFlowNets à la théorie classique des flux de réseau.
3. Solveur d'OT Scalable : Ce travail démontre que le cadre d'apprentissage des GFlowNets peut être appliqué pour approximer des solutions de problèmes d'OT sur graphes. En utilisant des réseaux de neurones et l'échantillonnage de trajectoires stochastiques, la méthode passe à l'échelle sur des espaces combinatoires où les solveurs d'OT exacts deviennent intraçables.
4. Récupération de la Théorie du Plus Court Chemin : En tant que corollaire, la construction récupère la théorie du plus court chemin précédemment établie pour les GFlowNets non acycliques, montrant que les flux optimaux ne sont supportés que sur le sous-graphe des plus courts chemins.

Résultats Expérimentaux
Les auteurs valident leur théorie et leur méthode sur deux environnements :

• Environnement Hypergrid : Sur de petites hypergrilles ($H=10, 15, 20$), la politique apprise du GFlowNet atteint des coûts de transport ($OT^\star$) qui correspondent aux solutions exactes obtenues via des solveurs de programmation linéaire classiques (POT). La méthode récupère avec succès le couplage optimal à travers différentes distributions sources (formes en balle et en croissant) sans biaiser l'échantillonnage.
• Environnement de Permutation : Sur le graphe de Cayley du groupe symétrique $S_n$ (permutations), la méthode récupère des solutions quasi-optimales pour $n=4$ et $n=8$, où les solutions exactes sont traitables. Pour $n=20$, où l'OT exact est intraçable, la méthode produit des approximations raisonnables avec une faible erreur dans les probabilités de point fixe.
• Ablation : Les expériences sur le coefficient de régularisation $\lambda$ confirment le compromis entre la longueur de la trajectoire (efficacité de l'échantillonnage) et l'optimalité de l'échantillonnage (précision du plan de transport). Des valeurs de $\lambda$ plus faibles produisent des plans de transport plus précis mais des trajectoires plus longues.

Signification et Revendications
Le papier prétend révéler une connexion « secrète » : les GFlowNets, lorsqu'ils sont formulés pour minimiser le flux, résolvent implicitement des problèmes de transport optimal. La signification de ce travail réside dans :

• Unification : Il fournit une nouvelle vue théorique unifiant l'apprentissage des GFlowNets et le transport optimal, suggérant que les GFlowNets peuvent être vus comme une méthode pour apprendre des cartes de transport.
• Application Pratique : Il permet l'application d'algorithmes d'apprentissage GFlowNet (qui passent à l'échelle via des réseaux de neurones et l'échantillonnage) à de grands problèmes d'OT sur graphes qui seraient autrement prohibitifs pour les solveurs exacts.
• Extension du Cadre : Il étend la généralité du cadre GFlowNet au-delà des tâches d'échantillonnage et de génération vers le domaine du transport optimal computationnel, offrant une politique stochastique qui détermine non seulement où la masse est transportée mais aussi comment elle est transportée via des mouvements locaux réalisables avec des coûts de chemin quasi-optimaux.

Les auteurs font preuve de modestie quant aux implications futures, notant que bien que la méthode soit efficace sur le plan computationnel, le coefficient de régularisation $\lambda$ nécessite un réglage minutieux pour équilibrer l'optimalité et l'efficacité. Ils suggèrent que les travaux futurs pourraient explorer l'application de ces méthodes pour apprendre des cartes de transport entre des distributions non normalisées.