Why Pool When You Can Flow? Active Learning with GFlowNets

2025-08-31

AI Summary

Traditional pool-based active learning methods (e.g., BALD) face severe computational bottlenecks in ultra-large-scale virtual screening (billions of molecules), as they require exhaustive evaluation of the entire unlabeled pool—rendering them inefficient and poorly scalable. To address this, we propose BALD-GFlowNet: the first active learning framework that integrates generative flow networks (GFlowNets) to directly sample high-value molecules proportional to their BALD information gain, thereby eliminating pool-wide scoring entirely. Its computational complexity is independent of pool size, enabling high-throughput molecular generation. Empirically, BALD-GFlowNet achieves screening performance on par with standard BALD while markedly improving structural diversity among selected compounds. Experiments on drug discovery virtual screening tasks demonstrate its superior efficiency, scalability, and chemical validity, establishing a novel paradigm for active learning at billion-scale.

Technology Category

• Machine Learning: Active Learning
• Search and Optimization: Sampling/Simulation-based Search
• Natural Language Processing: Generation

Application Category

• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs
• Economics, Online Markets and Human Computation: Economic ramifications for generative AI infrastructure and applications

Abstract

The scalability of pool-based active learning is limited by the computational cost of evaluating large unlabeled datasets, a challenge that is particularly acute in virtual screening for drug discovery. While active learning strategies such as Bayesian Active Learning by Disagreement (BALD) prioritize informative samples, it remains computationally intensive when scaled to libraries containing billions samples. In this work, we introduce BALD-GFlowNet, a generative active learning framework that circumvents this issue. Our method leverages Generative Flow Networks (GFlowNets) to directly sample objects in proportion to the BALD reward. By replacing traditional pool-based acquisition with generative sampling, BALD-GFlowNet achieves scalability that is independent of the size of the unlabeled pool. In our virtual screening experiment, we show that BALD-GFlowNet achieves a performance comparable to that of standard BALD baseline while generating more structurally diverse molecules, offering a promising direction for efficient and scalable molecular discovery.

Problem

• Scalable active learning for large unlabeled datasets
• Computational efficiency in virtual drug screening
• Generative sampling to replace pool-based acquisition

Innovation

• Generative Flow Networks for active learning
• Direct sampling proportional to BALD reward
• Scalability independent of unlabeled pool size