Multi-Fidelity Active Learning with GFlowNets

TL;DR

This work tackles the challenge of efficiently discovering diverse, high-scoring candidates in combinatorially large, high-dimensional spaces under limited budgets. It introduces MF-GFN, a framework that extends GFlowNets with multi-fidelity active learning by learning a joint policy over inputs and fidelity levels and optimising a cost-aware acquisition, MF-MES, via a multi-fidelity GP with deep kernel learning. Through extensive experiments on DNA aptamers, antimicrobial peptides, and small molecules, MF-GFN achieves substantial cost reductions over single-fidelity baselines while retaining diversity and discovering multiple high-scoring modes, outperforming RL-based and some BO baselines. The approach demonstrates practical potential to accelerate scientific discovery and materials/drug design by efficiently allocating computational or experimental resources across fidelity levels. The work also discusses limitations, such as simulated costs, and outlines future work including more complex design spaces and multi-objective extensions.

Abstract

In the last decades, the capacity to generate large amounts of data in science and engineering applications has been growing steadily. Meanwhile, machine learning has progressed to become a suitable tool to process and utilise the available data. Nonetheless, many relevant scientific and engineering problems present challenges where current machine learning methods cannot yet efficiently leverage the available data and resources. For example, in scientific discovery, we are often faced with the problem of exploring very large, structured and high-dimensional spaces. Moreover, the high fidelity, black-box objective function is often very expensive to evaluate. Progress in machine learning methods that can efficiently tackle such challenges would help accelerate currently crucial areas such as drug and materials discovery. In this paper, we propose a multi-fidelity active learning algorithm with GFlowNets as a sampler, to efficiently discover diverse, high-scoring candidates where multiple approximations of the black-box function are available at lower fidelity and cost. Our evaluation on molecular discovery tasks shows that multi-fidelity active learning with GFlowNets can discover high-scoring candidates at a fraction of the budget of its single-fidelity counterpart while maintaining diversity, unlike RL-based alternatives. These results open new avenues for multi-fidelity active learning to accelerate scientific discovery and engineering design.

Figures (11)

• Figure 1: Illustration of multi-fidelity active learning with GFlowNets (Algorithm \ref{['alg:mfal']}). Given a set of $M$ oracles $f_1, \ldots, f_M$ (center left) with varying fidelities and costs $\lambda < \ldots < \lambda_M$, respectively, we can construct a data set $\mathcal{D}$ (top left) with annotations from the oracles. With this data, we fit a multi-fidelity surrogate (center), modelling the posterior $p(f_m(x) | x, m, \mathcal{D})$. The surrogate is used to evaluate a multi-fidelity acquisition function---max-value entropy search in our experiments--- which makes the reward to train a GFlowNet (right). The GFlowNet is trained to sample both an object $x$ and the fidelity $m$ proportionally to this reward. Once the GFlowNet is trained, we sample $N$ tuples $(x, m)$ and select the top $B$ according to the acquisition function (bottom left). Finally, we annotate each new candidate with the selected oracle, add them to the data set and repeat the process until the budget is exhausted.
• Figure 2: Results on the DNA aptamers and AMP tasks. The curves indicate the mean energy $f_M$ within the top-100 samples computed at the end of each active learning round and plotted as a function of the budget used. The colour of the round markers indicates the diversity within the batch (darker colour indicating higher diversity), computed as the average pairwise sequence identity distance (see \ref{['sup:metrics']}). In both the DNA and the AMP tasks, MF-GFN outperforms all baselines in terms of cost efficiency, while obtaining great diversity in the final batch of top-$K$ candidates.
• Figure 3: Results on the molecular discovery tasks: (a) ionisation potential (IP), (b) electron affinity (EA). These visualisations are analogous to those in \ref{['fig:dna_amp']}. The diversity of molecules is computed as the average pairwise Tanimoto distance (see \ref{['sup:metrics']}). Results generally show MF-GFN's faster convergence in discovering diverse molecules with desirable properties.
• Figure 4: Results on the synthetic tasks---Branin and Hartmann functions. The curves indicate the mean score $f_M$ within the top-50 and top-10 samples (for Branin and Hartmann, respectively) computed at the end of each active learning round and plotted as a function of the budget used. The random baseline is omitted from this plot to facilitate the visualisation since the results were significantly worse in these tasks. We observe that MF-GFN clearly outperforms the single-fidelity counterpart (SF-GFN) and slightly improves upon the GFlowNet baseline that samples random fidelities. On Hartmann, MF-PPO initially outperforms the other methods.
• Figure 5: Mean scores (energy) of diverse top-$K$ candidates on the DNA (top left), AMP (top right) and molecular (bottom) tasks. The mean energy is computed across the top-$K$ candidates at each active learning round that also satisfy the criteria of diversity. Consistent with the diversity metrics observed in \ref{['fig:dna_amp']}, we here see that GFlowNet-based methods, and especially MF-GFN, obtain good results according to this metric, while MF-PPO achieves comparatively much lower mean energy.