Information-Geometric Forward Policy Training in GFlowNets

2026-08-04

AI Summary

This work addresses the challenges of optimization efficiency and structure awareness in Generative Flow Networks (GFlowNets) when performing inference in discrete and hybrid spaces based on unnormalized rewards. It introduces information geometry into forward policy training for the first time, proposing a natural gradient update mechanism grounded in the Fisher–Rao metric. By deriving a sequential decomposition of trajectory Fisher information, the study characterizes conditions under which temporal score interactions vanish or become coupled, and develops three Fisher information estimation paradigms tailored to different target distribution structures—namely, exact marginalization, separable subroutines, and belief propagation. Empirical results demonstrate that this structure-aware Riemannian optimization approach significantly outperforms conventional Euclidean methods, exhibiting superior convergence speed, enhanced exploration capability, and greater scalability.

Abstract

Generative Flow Networks (GFlowNets) have emerged as a flexible framework for amortised inference over discrete and mixed discrete-continuous objects, requiring only an unnormalised target density specified through a reward. In this work, we formulate forward-policy training in GFlowNets through the information geometry of the induced trajectory sampler. Treating the forward policy as an induced trajectory sampler, we show that its intrinsic first-order geometry is given by the Fisher-Rao metric of the trajectory family, and that the associated natural gradient provides the canonical local update whenever the corresponding Fisher information is computable or accurately approximable. We derive an exact decomposition of the trajectory Fisher into per-step conditional second moments, which clarifies when temporal score interactions vanish and when dense couplings remain under shared parameterisation. This leads to three computational regimes: settings with tractable exact Fisher information, settings where Monte Carlo estimators of the expected Fisher are sufficient, and structure-exploitable settings in which target locality or factorisation yields accurate approximations of the Fisher expectation. In the latter case, graphical-model tools such as exact marginalisation, separator methods, and belief propagation provide principled surrogates for natural-gradient updates. The resulting framework turns target structure into optimisation geometry and yields a tractable route to structure-aware forward-policy training in GFlowNets. We illustrate the framework empirically through examples comparing convergence and exploration behaviour under Riemannian and Euclidean optimisation.

Problem

• GFlowNets
• forward policy training
• information geometry
• Fisher information
• structure-aware optimization

Innovation

• Information Geometry
• GFlowNets
• Natural Gradient
• Fisher Information
• Structure-aware Optimization