Beyond the Proxy: Trajectory-Distilled Guidance for Offline GFlowNet Training
Abstract
Generative Flow Networks (GFlowNets) excel at sampling diverse, high-reward objects. In many practical applications where active reward queries are infeasible, these models must be trained using static offline datasets. Prevailing training methods typically rely on a proxy model to provide reward feedback for online sampled trajectories. However, constructing a reliable proxy is often challenging due to data scarcity or high evaluation costs. While existing proxy-free approaches attempt to address this, they often impose coarse constraints that limit the model's ability to explore effectively. To overcome these limitations, we propose Trajectory-Distilled GFlowNet (TD-GFN), a novel proxy-free training framework. TD-GFN utilizes inverse reinforcement learning (IRL) to extract dense, transition-level edge rewards from offline trajectories, providing rich structural guidance for efficient exploration. Crucially, to ensure robustness, these rewards guide the policy indirectly through DAG pruning and prioritized backward sampling. This design ensures that gradient updates rely exclusively on ground-truth terminal rewards from the dataset, thereby preventing error propagation. Empirical results demonstrate that TD-GFN significantly outperforms a broad range of existing baselines in both convergence speed and sample quality, establishing a more robust and efficient paradigm for offline GFlowNet training.
Lay Summary
Many scientific breakthroughs — discovering a new drug, designing a useful protein, or finding a better material — come down to searching an unimaginably large pool of candidates for the few that score well. A class of machine learning models called Generative Flow Networks is well suited to this task because, instead of returning a single winner, it proposes many distinct, promising options. The catch is that training them normally requires repeatedly asking "how good is this candidate?" — and in real labs each answer can take weeks of experiments or expert review. To avoid this, researchers usually train a "stand-in" scoring model from past data, but the stand-in can be badly wrong on unfamiliar candidates and steer the search astray. We propose TD-GFN, a training method that does away with the stand-in entirely. From a fixed log of past attempts, it learns which intermediate decisions tended to lead somewhere worthwhile, uses that signal to prune dead-end paths and concentrate on promising ones, and only ever updates itself using the real, measured outcomes already recorded. On standard benchmarks, TD-GFN finds higher-quality and more diverse candidates faster than prior methods, making such models practical in settings where every real evaluation is precious.