BALD-GFlowNet Framework

• The paper introduces BALD-GFlowNet, which combines BALD’s mutual information-based reward with a generative GFlowNet policy to sample informative objects without exhaustively scoring massive pools.
• It leverages Monte Carlo dropout for BALD score estimation and uses trajectory balance techniques to train the GFlowNet, reducing evaluation costs in virtual screening.
• Empirical results show that BALD-GFlowNet achieves near-pool-based BALD performance with significantly fewer surrogate evaluations while generating structurally diverse molecules.

BALD-GFlowNet is a generative active learning framework that combines Bayesian Active Learning by Disagreement (BALD) with Generative Flow Networks (GFlowNets) so that objects are sampled with probability proportional to a BALD-based reward, rather than being selected by exhaustively scoring a finite unlabeled pool (Zhang et al., 31 Aug 2025). It was introduced to address the scalability limits of pool-based active learning in settings such as high-throughput virtual screening, where the unlabeled library can contain up to billions of molecules and exhaustive acquisition becomes computationally dominant (Zhang et al., 31 Aug 2025). In its canonical formulation, BALD provides the informativeness criterion through mutual information, while the GFlowNet supplies a generative policy that amortizes acquisition and naturally samples across multiple high-reward modes (Zhang et al., 31 Aug 2025).

1. Problem setting and motivation

BALD-GFlowNet is motivated by the limitations of pool-based active learning when the unlabeled set is extremely large. In the standard pool-based setting, one has a labeled training set $\mathcal{D}_{\text{train}} = \{(x_i, y_i)\}$, a large unlabeled pool $\mathcal{D}_{\text{pool}} = \{x_j\}$, and an oracle that provides labels for selected points. At each iteration, an acquisition function is evaluated on all unlabeled samples, the top-$K$ points are queried, and the model is retrained (Zhang et al., 31 Aug 2025). This becomes problematic in virtual screening, where libraries such as ZINC or Enamine REAL may contain $10^7$–$10^9$ molecules, and where docking every molecule is infeasible; the cited example is that 1.3 billion ligands can take 28 days on 8,000 GPUs (Zhang et al., 31 Aug 2025).

The central computational bottleneck is that BALD requires exhaustively computing mutual information for every sample in the unlabeled pool. In the formulation reported for BALD-GFlowNet, the cost of pool-based BALD scales linearly with the pool size in virtual screening, and in the 2D grid setup the oracle calls scale as $O(n^2)$ in their specific setup because of repeated active-learning iterations with exhaustive scoring (Zhang et al., 31 Aug 2025). BALD-GFlowNet removes this dependence by replacing pool-wide scoring with generative sampling from a learned reward-proportional policy (Zhang et al., 31 Aug 2025).

This replacement changes the acquisition paradigm. Traditional BALD asks which existing pool elements should be labeled next, whereas BALD-GFlowNet trains a generator to sample directly from regions of high informativeness (Zhang et al., 31 Aug 2025). The framework is therefore particularly suited to cases where the unlabeled pool is either gigantic or effectively conceptual, as in de novo molecular generation (Zhang et al., 31 Aug 2025).

2. Constituent principles: BALD and GFlowNets

The BALD component is defined by the mutual information between the unknown label $y$ and the model parameters $\theta$ given $x$ and the current dataset $D$: $\mathcal{D}_{\text{pool}} = \{x_j\}$0 Using entropy notation, the acquisition function is written as

$\mathcal{D}_{\text{pool}} = \{x_j\}$1

It is high when the average prediction is uncertain but individual parameter samples are confident and disagree, so it targets epistemic uncertainty (Zhang et al., 31 Aug 2025). In the reported experiments, BALD is approximated with Monte Carlo dropout: several dropout masks are sampled for each input, each mask yields a different prediction, and the predictive entropy and expected per-model entropy are combined to estimate BALD$\mathcal{D}_{\text{pool}} = \{x_j\}$2 (Zhang et al., 31 Aug 2025).

The GFlowNet component is a generative model over trajectories or discrete objects that aims to sample terminal objects proportionally to a non-negative reward function. The target relation is

$\mathcal{D}_{\text{pool}} = \{x_j\}$3

or equivalently

$\mathcal{D}_{\text{pool}} = \{x_j\}$4

for an unknown normalization constant $\mathcal{D}_{\text{pool}} = \{x_j\}$5 (Zhang et al., 31 Aug 2025). Objects are generated through trajectories of states and actions, and flow conservation constraints enforce consistency between incoming and outgoing flows (Zhang et al., 31 Aug 2025). In the 2D grid experiment, the method uses Trajectory Balance (TB), with

$\mathcal{D}_{\text{pool}} = \{x_j\}$6

while the virtual screening experiment uses Relative Trajectory Balance (RTB), a variant used for fine-tuning pre-trained GFlowNets with regularization toward feasible molecules (Zhang et al., 31 Aug 2025).

The combination of BALD and GFlowNets is conceptually direct: BALD supplies the reward signal, and the GFlowNet learns a policy whose terminal distribution is proportional to that reward (Zhang et al., 31 Aug 2025). This preserves a central GFlowNet property emphasized across the literature: unlike reward-maximizing RL, GFlowNets are designed to cover the reward landscape rather than collapse onto a single mode. In mathematical reasoning with LLMs, for example, GFlowNet fine-tuning was reported to achieve accuracy similar to reward-maximizing RL while producing the highest number of distinct correct solutions on GSM8K and MATH, which is evidence for multi-modal reward-proportional sampling in another domain (Takase et al., 2024).

3. BALD-GFlowNet method

BALD-GFlowNet integrates the BALD score directly into the GFlowNet reward. Let $\mathcal{D}_{\text{pool}} = \{x_j\}$7 denote an object and $\mathcal{D}_{\text{pool}} = \{x_j\}$8 its BALD score,

$\mathcal{D}_{\text{pool}} = \{x_j\}$9

In the 2D grid experiment, the terminal reward is simply

$K$0

In the virtual screening case study, the reward is a product of BALD with molecular feasibility and drug-likeness terms: $K$1 Here $K$2 denotes topological polar surface area, $K$3 the quantitative estimate of drug-likeness, $K$4 the synthetic accessibility score, and $K$5 a ring-related structural property (Zhang et al., 31 Aug 2025). The target distribution is then

$K$6

In virtual screening, $K$7 is normalized into $K$8 and shaped via a task slope parameter to avoid numerical instability (Zhang et al., 31 Aug 2025).

The active learning loop has five operational stages. First, an initial training set $K$9 is sampled from a large unlabeled pool $10^7$0, an initial surrogate model is trained, and a GFlowNet policy is initialized. Second, BALD-based rewards are computed from the surrogate, optionally together with $10^7$1, $10^7$2, $10^7$3, and $10^7$4. Third, the GFlowNet is trained for a fixed number of episodes using TB or RTB so that its policy matches the reward-proportional target. Fourth, the trained policy generates a batch of objects,

$10^7$5

Fifth, the oracle is queried on those objects, the labeled set is updated, and the surrogate is retrained or fine-tuned (Zhang et al., 31 Aug 2025).

A defining feature is that the generated batch is not selected by ranking a finite pool. In the molecular setting, the GFlowNet constructs molecules directly, atom by atom, using a graph-based Transformer policy (Zhang et al., 31 Aug 2025). As a result, the acquisition cost per active learning iteration is determined primarily by GFlowNet training and reward evaluation on sampled objects, not by exhaustive scanning of the entire unlabeled set (Zhang et al., 31 Aug 2025).

4. Experimental realization in virtual screening

The primary case study is virtual screening for Janus Kinase 2 (JAK2). The library is drawn from Enamine REAL, with a target-specific dataset constructed using Deep Docking to pre-enrich potential hits. The data split uses 10,000 molecules for initial training, 100,000 molecules as a held-out test set, and the remaining approximately 1.74M molecules as the conceptual unlabeled pool (Zhang et al., 31 Aug 2025). The oracle is QuickVina2 docking, and binary labels are defined using a 1% quantile cutoff of the most negative docking scores in the initial training data, yielding a highly imbalanced classification problem (Zhang et al., 31 Aug 2025).

The surrogate classifier is a MoLFormer-based model that takes SMILES strings, uses attention pooling, DenseNet-like blocks, and self-attention, and is trained with cross-entropy loss and mixup on molecular embeddings to mitigate class imbalance. Uncertainty estimation is performed with Monte Carlo dropout using 3 models and dropout rate 0.1 (Zhang et al., 31 Aug 2025). The molecular GFlowNet is an atom-level, graph-based Transformer policy with 15 GNN layers and 4 heads per layer, pre-trained and then fine-tuned with RTB while regularizing against an offline dataset of feasible molecules (Zhang et al., 31 Aug 2025).

The main reported results are as follows.

Method	Peak F1	Surrogate evaluations at peak
BALD	0.32	12.2M
BALD-GFlowNet	0.30	1.74M
Random	0.28	0

BALD-GFlowNet therefore reaches near-BALD performance at approximately $10^7$6 fewer surrogate evaluations (Zhang et al., 31 Aug 2025). Runtime scaling is also contrasted explicitly: BALD exhibits $10^7$7 scaling with fitted line $10^7$8, whereas BALD-GFlowNet has $10^7$9 complexity and runtime essentially flat across pool sizes. At 50M molecules, the estimated wall-clock speedup is approximately $10^9$0 in favor of BALD-GFlowNet (Zhang et al., 31 Aug 2025).

The method is also reported to generate more structurally diverse molecules. The stated explanation is twofold: the underlying GFlowNet explores multiple modes of the reward landscape, and the composite reward promotes chemical viability and structural diversity through properties such as $10^9$1 and ring count (Zhang et al., 31 Aug 2025). Exact numeric diversity metrics were not given in the main text, so the diversity claim is qualitative at that level of detail (Zhang et al., 31 Aug 2025).

5. Computational and methodological comparison with pool-based BALD

The central contrast is between exhaustive acquisition and generative acquisition. Standard pool-based BALD computes BALD$10^9$2 for every $10^9$3 through multiple stochastic forward passes, then selects the top-$10^9$4 samples (Zhang et al., 31 Aug 2025). Its per-iteration complexity is

$10^9$5

where $10^9$6 is the number of Monte Carlo samples (Zhang et al., 31 Aug 2025). By design, BALD-GFlowNet does not evaluate the entire pool; instead, it amortizes the acquisition function into a generative policy (Zhang et al., 31 Aug 2025).

This distinction has practical implications beyond asymptotic complexity. BALD-GFlowNet is described as suitable for pools whose size is effectively independent of acquisition cost, because the runtime depends on the number of GFlowNet training episodes and the reward evaluations used during training, not on the number of candidate molecules sitting in an external catalog (Zhang et al., 31 Aug 2025). The paper further reports that, on the synthetic grid task, BALD-GFlowNet requires far fewer oracle calls than exhaustive BALD to locate high-uncertainty regions (Zhang et al., 31 Aug 2025).

From a methodological perspective, the framework also changes the geometry of exploration. Pool-based BALD is constrained by what is already enumerated, whereas BALD-GFlowNet learns what an informative object looks like and samples from that induced distribution (Zhang et al., 31 Aug 2025). A plausible implication is that this generative view becomes increasingly advantageous as the candidate space shifts from “very large pool” to “combinatorial or de novo design space,” because the distinction between selection and generation becomes operationally important.

The relationship to broader GFlowNet training theory is explicit in subsequent work. “Towards Understanding and Improving GFlowNet Training” analyzes how practical GFlowNets can underfit high-reward regions and proposes prioritized replay training, relative edge flow policy parametrization, and guided trajectory balance to improve sample efficiency on skewed reward landscapes (Shen et al., 2023). “GFlowNet Training by Policy Gradients” recasts GFlowNet objectives through policy-dependent rewards and derives policy-gradient and TRPO-style training schemes with theoretical guarantees (Niu et al., 2024). These developments were not part of the original BALD-GFlowNet formulation, but they indicate concrete training directions for reward landscapes defined by acquisition functions such as BALD (Shen et al., 2023, Niu et al., 2024).

6. Relation to adjacent GFlowNet variants and prospective extensions

BALD-GFlowNet belongs to a broader line of work in which GFlowNets are used as reward-proportional samplers in domains where exhaustive evaluation is difficult. Several later or parallel formulations clarify how the idea may be extended.

“Bifurcated Generative Flow Networks” introduces a factorized parameterization

$10^9$7

separating state flow from edge allocation in order to improve data efficiency and scalability in large action spaces (2406.01901). This is relevant because BALD-based active learning also operates on expensive rewards and potentially large combinatorial decision spaces; the paper explicitly notes that such factorization may inform designs like BALD-GFlowNet (2406.01901).

“AbFlowNet” shows how a diffusion model can be interpreted as a GFlowNet trajectory and optimized jointly with a terminal reward through Trajectory Balance: $10^9$8 Although its application is antibody-antigen design rather than active learning, it demonstrates that terminal rewards can be propagated through diffusion trajectories without an online RL loop (Abir et al., 18 May 2025). This suggests a possible route for extending BALD-GFlowNet to diffusion-based generators in molecular or structural domains, where the terminal reward would be informativeness rather than binding energy.

A related line of evidence concerns diversity. In mathematical reasoning with LLMs, GFlowNet fine-tuning with a SubTB objective was reported to produce the highest number of distinct correct solutions while maintaining accuracy similar to reward-maximizing RL (Takase et al., 2024). That result does not concern BALD-GFlowNet directly, but it reinforces a property that is central to its use in active learning: reward-proportional sampling can preserve multiple high-value modes instead of collapsing to a single solution style or structure (Takase et al., 2024).

The framework also has clearly identified limitations. The BALD-GFlowNet paper notes dependence on surrogate model quality, sensitivity to reward design, domain-specific implementation, and the engineering cost of training GFlowNets with TB or RTB objectives (Zhang et al., 31 Aug 2025). In virtual screening, balancing BALD-based informativeness against drug-likeness and synthesizability is described as nontrivial; mis-calibrated weights may bias the policy toward informative but chemically implausible molecules, or chemically plausible but uninformative ones (Zhang et al., 31 Aug 2025). Future extensions suggested there include other acquisition functions, improved Bayesian uncertainty estimators, broader scientific domains such as biological sequence design and materials discovery, and hybrid pool-generative schemes (Zhang et al., 31 Aug 2025).

In that sense, BALD-GFlowNet is best understood as a specific instantiation of a more general program: replacing exhaustive acquisition over massive pools with reward-proportional generative sampling, where the reward is an information-theoretic criterion derived from the current surrogate model (Zhang et al., 31 Aug 2025). Subsequent GFlowNet work on training objectives, scalable parameterizations, and diffusion fusion suggests that this program is extensible, but those extensions remain prospective rather than intrinsic to the original BALD-GFlowNet formulation (2406.01901, Abir et al., 18 May 2025, Shen et al., 2023, Niu et al., 2024).