GFlowNets for AI-Driven Scientific Discovery

Abstract: Tackling the most pressing problems for humanity, such as the climate crisis and the threat of global pandemics, requires accelerating the pace of scientific discovery. While science has traditionally relied on trial and error and even serendipity to a large extent, the last few decades have seen a surge of data-driven scientific discoveries. However, in order to truly leverage large-scale data sets and high-throughput experimental setups, machine learning methods will need to be further improved and better integrated in the scientific discovery pipeline. A key challenge for current machine learning methods in this context is the efficient exploration of very large search spaces, which requires techniques for estimating reducible (epistemic) uncertainty and generating sets of diverse and informative experiments to perform. This motivated a new probabilistic machine learning framework called GFlowNets, which can be applied in the modeling, hypotheses generation and experimental design stages of the experimental science loop. GFlowNets learn to sample from a distribution given indirectly by a reward function corresponding to an unnormalized probability, which enables sampling diverse, high-reward candidates. GFlowNets can also be used to form efficient and amortized Bayesian posterior estimators for causal models conditioned on the already acquired experimental data. Having such posterior models can then provide estimators of epistemic uncertainty and information gain that can drive an experimental design policy. Altogether, here we will argue that GFlowNets can become a valuable tool for AI-driven scientific discovery, especially in scenarios of very large candidate spaces where we have access to cheap but inaccurate measurements or to expensive but accurate measurements. This is a common setting in the context of drug and material discovery, which we use as examples throughout the paper.

• The paper introduces GFlowNets, which sample complex objects by sequential construction, addressing scalability, uncertainty, and diversity challenges.
• The methodology employs trajectory and detailed balance losses with amortized inference to generate diverse candidates in molecular and causal discovery.
• Results show that GFlowNets outperform traditional methods like MCMC, VI, and RL by achieving broader exploration and improved candidate diversity.

GFlowNets for AI-Driven Scientific Discovery: A Technical Overview

AI-Driven Scientific Discovery and Core Challenges

Modern scientific progress, particularly in urgent fields such as climate solutions and pharmaceuticals, is increasingly constrained by the combinatorial explosion of candidate hypotheses and experiment designs. Traditional approaches relying on direct optimization, Bayesian methods, or RL typically exhibit poor scalability, insufficient uncertainty handling, or lack of diversity in candidate generation. This hinders autonomous scientific discovery, especially when objective functions are vastly underspecified or multi-modal, and expensive evaluations preclude exhaustive search.

The experimental science loop—comprising data modeling, hypothesis generation, and experiment design—demands frameworks intrinsically suited for efficient exploration, epistemic uncertainty modeling, batch proposal generation, and amortized inference.

Figure 1: The experimental science loop, with GFlowNet-applicable stages highlighted.

Limitations of Standard Approaches

Bayesian Optimization, MCMC, and VI are foundational in automated discovery but have structural and computational limitations:

• MCMC: Prone to mode collapse in high-dimensional and highly multimodal posteriors, with exponential time mixing and inefficient sampling, particularly over combinatorial structure spaces.
• Variational Inference: Standard ELBO-based approaches are mode-seeking due to reverse KL minimization, resulting in poor posterior mode coverage and diversity.
• RL and BO: Predisposed to policy or acquisition functions optimizing mean or maximum reward, leading to loss of candidate variety—a severe limitation for underdetermined or high-noise scientific objectives.

The GFlowNet Framework

GFlowNets (Generative Flow Networks) address the above by learning policies that sample complex objects (e.g., molecules, graphs) in a sequential, constructive fashion, with the terminal state probability proportional to a user-defined reward (often an unnormalized Bayesian posterior or proxy utility). The constructive paradigm enables the exploitation of compositional structure, Markovian state representations, and explicit stochasticity in candidate synthesis.

Figure 2: Schematic illustration of compositional object generation with a GFlowNet, e.g., sequence/graph construction for molecular design.

Key formalisms:

• Detailed Balance and Trajectory Balance Losses: Formulated over compositional state-action trajectories; the network is trained to ensure the flow into each state matches expected reward-driven distributions, aligning terminal state probabilities with the target.
• Amortized Inference: Parameterized policies and flow estimators are learned globally, enabling rapid test-time sampling and scalable posterior/marginal estimation (in contrast to runtime MC estimation).
• Multimodality and Diversity: By explicitly aligning sample frequencies with the (potentially multi-modal) reward landscape, GFlowNets recover a diversity of high-quality candidates—critical for underspecified objectives or exploration.

  Figure 3: Illustration of the potential for systematic generalization and mode discovery by amortized ML samplers (GFlowNets) versus MCMC.

Applications in Molecular and Material Discovery

Diverse Candidate Generation

In molecular design, binding affinity estimation and ADME properties are only partially observable proxies for true downstream effects. Direct maximization (as in RL/BO) commonly results in overfitting to artifacts of these proxies, while failing to cover the chemical diversity space.

• GFlowNets generate candidate molecules by sequentially composing fragments, with rewards derived from GNN-based binding affinity surrogates. Empirical results indicate that GFlowNets attain a significantly larger number of reward function modes compared to RL and MCMC, thus generating not only higher-scoring, but also structurally diverse, molecules.

Figure 4: GFlowNet-generated molecule batches show high diversity in chemical scaffolds, compared to RL and MCMC baselines.

Strong empirical findings include the discovery of more unique scaffolds, enhanced coverage of reward peaks, and improved exploration of the molecular space, especially under active learning regimes. In addition, multi-objective extensions (MOGFNs) demonstrate improved coverage of Pareto-optimal fronts in molecular and protein sequence design tasks, accommodating practical demands for trade-off between properties.

GFlowNets for Bayesian Posterior and Causal Structure Modeling

A central tenet of scientific inference is robust Bayesian uncertainty quantification. When the latent variable is a causal graph $G$ over random variables, posterior inference is particularly challenging due to the combinatorial size of the DAG space and possible multi-modality.

• Posterior over Causal Structures: GFlowNets are trained to sample DAGs (or more general cyclic graphs) with probability proportional to $P(D|G)P(G)$, where $P(G)$ encodes prior domain beliefs and $P(D|G)$ is the (possibly closed-form or approximated) marginal likelihood. The constructive nature of GFlowNets enables the sequential growth of the graph while upholding acyclicity and prior constraints.

  Figure 5: GFlowNet for Bayesian Causal Discovery—(left) sequential construction of DAGs by edge addition; (right) comparison of marginal posterior edge probabilities to exact baseline, highlighting accuracy.

The framework extends naturally to scenarios with unknown interventions, integration of biological priors, and supporting interpretability via Bayesian model averaging or marginalization over causal edge existence.

Theoretical and Practical Implications

Amortized Mutual Information Estimation

GFlowNets enable joint amortization of both distribution sampling and intractable expectation/information-theoretic utility computation (e.g., mutual information for experimental design), eliminating the need for computationally intensive nested MC estimators.

Active Experimental Design

As an acquisition engine, GFlowNet-sampled batches can optimize for epistemic uncertainty, predicted information gain, or direct utility, adapting to real-world multi-fidelity, parallelized settings, and leveraging learned marginal or posterior predictive networks.

Causality and Beyond

GFlowNets provide a unified approach to causal discovery and experimental design. The sequential construction process enables efficient exploration of structured latent spaces (causal graphs, reaction networks) and elegantly supports the incorporation of prior knowledge, constraints, and interpretable model averaging.

Open Problems and Future Directions

• Extension to Continuous and Mixed Spaces: Recent theoretical work has laid foundations for continuous GFlowNets, but scalable training in high-dimensional, hybrid domains remains open.
• Credit Assignment for Long/Heterogeneous Trajectories: While local credit schemes partially alleviate gradient diffusion, further research is required for efficient training over very long compositional sequences (e.g., proteins).
• Scalable Training/Exploration Strategies: Optimal off-policy or prioritized trajectory sampling strategies, beyond heuristic mixtures, remain an ongoing research challenge.
• Integration with Bayesian Optimal Experimental Design: Systematic amortization of MI-based utilities in large-scale, sequential closed-loop experimentation holds promise for automated science platforms.

Conclusion

GFlowNets constitute a principled, scalable, and flexible framework for AI-driven scientific discovery. By representing and sampling from complex, multimodal distributions over structured objects via amortized deep policies, they address core deficiencies of MCMC, VI, and RL in exploration and uncertainty modeling. Their deployment in molecule/material design, active learning, and Bayesian causal discovery demonstrates concrete advantages in diversity, robustness, and tractability. Future efforts should target scaling up to richer domains (continuous/hybrid spaces, larger graphs), optimizing training protocols, and deeper integration with information-theoretic and causal objectives—potentially advancing closed-loop, AI-augmented scientific workflows.

---

Reference:

GFlowNets for AI-Driven Scientific Discovery (2302.00615)