---
title: BALD-GFlowNet Framework
url: https://www.emergentmind.com/topics/bald-gflownet
type: topic
---

# BALD-GFlowNet Framework

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 [2509.00704]. 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 [2509.00704]. 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 [2509.00704].

## 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 [2509.00704]. 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 [2509.00704].

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 [2509.00704]. BALD-GFlowNet removes this dependence by replacing pool-wide scoring with generative sampling from a learned reward-proportional policy [2509.00704].

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 [2509.00704]. The framework is therefore particularly suited to cases where the unlabeled pool is either gigantic or effectively conceptual, as in de novo molecular generation [2509.00704].

## 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$:
\[
\text{BALD}(x) = I(y; \theta \mid x, D).
\]
Using entropy notation, the acquisition function is written as
\[
I(y; \theta \mid x, D) = \mathbb{H}[y \mid x, D]
 - \mathbb{E}_{\theta \sim p(\theta \mid D)} \left[ \mathbb{H}[y \mid x, \theta] \right].
\]
It is high when the average prediction is uncertain but individual parameter samples are confident and disagree, so it targets epistemic uncertainty [2509.00704]. 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$(x)$ [2509.00704].

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
\[
P_\pi(x) \propto R(x),
\]
or equivalently
\[
P_\pi(x) = \frac{R(x)}{Z},
\]
for an unknown normalization constant $Z$ [2509.00704]. Objects are generated through trajectories of states and actions, and flow conservation constraints enforce consistency between incoming and outgoing flows [2509.00704]. In the 2D grid experiment, the method uses Trajectory Balance (TB), with
\[
\log P_F(\tau) + \log Z = \log R(x) + \log P_B(\tau),
\]
while the virtual screening experiment uses Relative Trajectory Balance (RTB), a variant used for fine-tuning pre-trained GFlowNets with regularization toward feasible molecules [2509.00704].

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 [2509.00704]. 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 large language models, 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 [2410.20147].

## 3. BALD-GFlowNet method

BALD-GFlowNet integrates the BALD score directly into the GFlowNet reward. Let $x$ denote an object and $MI(x)$ its BALD score,
\[
MI(x) = I(y; \omega \mid x, D).
\]
In the 2D grid experiment, the terminal reward is simply
\[
R(x) = MI(x).
\]
In the virtual screening case study, the reward is a product of BALD with molecular feasibility and drug-likeness terms:
\[
R(x) = MI(x) \cdot TPSA(x) \cdot QED(x) \cdot SAS(x) \cdot Rings(x).
\]
Here $TPSA(x)$ denotes topological polar surface area, $QED(x)$ the quantitative estimate of drug-likeness, $SAS(x)$ the synthetic accessibility score, and $Rings(x)$ a ring-related structural property [2509.00704]. The target distribution is then
\[
p(x) \propto R(x).
\]
In virtual screening, $MI(x)$ is normalized into $[0,1]$ and shaped via a task slope parameter to avoid numerical instability [2509.00704].

The active learning loop has five operational stages. First, an initial training set $\mathcal{D}_{\text{train}}$ is sampled from a large unlabeled pool $\mathcal{D}_{\text{pool}}$, an initial surrogate model is trained, and a GFlowNet policy is initialized. Second, BALD-based rewards are computed from the surrogate, optionally together with $TPSA$, $QED$, $SAS$, and $Rings$. 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,
\[
b_t \sim \text{SampleFromGFNPolicy}(\pi, k).
\]
Fifth, the oracle is queried on those objects, the labeled set is updated, and the surrogate is retrained or fine-tuned [2509.00704].

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 [2509.00704]. 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 [2509.00704].

## 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 [2509.00704]. 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 [2509.00704].

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 [2509.00704]. 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 [2509.00704].

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 $7\times$ fewer surrogate evaluations [2509.00704]. Runtime scaling is also contrasted explicitly: BALD exhibits $\mathcal{O}(n)$ scaling with fitted line $y = 803.11x + 789.80$, whereas BALD-GFlowNet has $\mathcal{O}(1)$ complexity and runtime essentially flat across pool sizes. At 50M molecules, the estimated wall-clock speedup is approximately $2.5\times$ in favor of BALD-GFlowNet [2509.00704].

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 $TPSA$ and ring count [2509.00704]. Exact numeric diversity metrics were not given in the main text, so the diversity claim is qualitative at that level of detail [2509.00704].

## 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$(x)$ for every $x \in \mathcal{D}_{\text{pool}}$ through multiple stochastic forward passes, then selects the top-$K$ samples [2509.00704]. Its per-iteration complexity is
\[
O(|\mathcal{D}_{\text{pool}}| \cdot M),
\]
where $M$ is the number of Monte Carlo samples [2509.00704]. By design, BALD-GFlowNet does not evaluate the entire pool; instead, it amortizes the acquisition function into a generative policy [2509.00704].

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 [2509.00704]. 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 [2509.00704].

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 [2509.00704]. 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 [2305.07170]. “GFlowNet Training by Policy Gradients” recasts GFlowNet objectives through policy-dependent rewards and derives policy-gradient and TRPO-style training schemes with theoretical guarantees [2408.05885]. 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 [2305.07170; 2408.05885].

## 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
\[
F(s \rightarrow s') = F(s)\, A(s' \mid s),
\]
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:
\[
L = \mathbb{E}_{t \sim \text{Uniform}(1,\ldots,T)}\Big[L^t_{\text{type}} + L^t_{\text{pos}} + L^t_{\text{ori}}\Big] + w\cdot L_{\text{TB}}.
\]
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 [2505.12358]. 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 large language models, 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 [2410.20147]. 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 [2410.20147].

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 [2509.00704]. 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 [2509.00704]. 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 [2509.00704].

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 [2509.00704]. 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; 2505.12358; 2305.07170; 2408.05885].

Source: https://www.emergentmind.com/topics/bald-gflownet