Source-linked AI summary
Biological Sequence Design with GFlowNets
Moksh Jain, Emmanuel Bengio, Alex-Hernandez Garcia, Jarrid Rector-Brooks, Bonaventure F. P. Dossou, Chanakya Ekbote, Jie Fu, Tianyu Zhang, Micheal Kilgour, Dinghuai Zhang, Lena Simine, Payel Das, Yoshua Bengio
TL;DR
De novo biological sequence design must select candidates through costly, imperfect multi-stage evaluations, making diverse proposals important. The paper combines GFlowNet generation with epistemic-uncertainty-based active learning and offline training data, reporting better diversity and relative novelty with high-scoring candidates across tasks, while noting GFP difficulties and added learner complexity.
Problem
Biological sequence design involves costly, multi-stage evaluations with imperfect early screening, so the design phase needs diverse useful and informative candidate batches.
Method
The method uses a GFlowNet generator with epistemic uncertainty in predicted reward and additional training trajectories from existing labeled oracle data.
Results
The method obtained better results especially in diversity and relative novelty of generated candidates, while producing high-scoring candidates across biological sequence design tasks.
Takeaways & Limitations
GFlowNets can serve as candidate generators in active learning for biological sequence design, with uncertainty and oracle-derived training data supporting the proposed objective.
Takeaways & Limitations
The approach uses two learners with separate hyperparameters and performed poorly on GFP, where de novo generation struggled to find good candidates.
Abstract
from arXiv · showhide
Design of de novo biological sequences with desired properties, like protein and DNA sequences, often involves an active loop with several rounds of molecule ideation and expensive wet-lab evaluations. These experiments can consist of multiple stages, with increasing levels of precision and cost of evaluation, where candidates are filtered. This makes the diversity of proposed candidates a key consideration in the ideation phase. In this work, we propose an active learning algorithm leveraging epistemic uncertainty estimation and the recently proposed GFlowNets as a generator of diverse candidate solutions, with the objective to obtain a diverse batch of useful (as defined by some utility function, for example, the predicted anti-microbial activity of a peptide) and informative candidates after each round. We also propose a scheme to incorporate existing labeled datasets of candidates, in addition to a reward function, to speed up learning in GFlowNets. We present empirical results on several biological sequence design tasks, and we find that our method generates more diverse and novel batches with high scoring candidates compared to existing approaches.
1. Introduction
Biological sequence design uses iterative candidate generation and increasingly precise evaluation, making diversity important when early oracles are imperfect. The paper proposes GFlowNet-based active learning with uncertainty estimation, offline training data, and validation across three design tasks.
- Biological sequence design spans applications from medicine to material design, including anti-microbial peptide discovery motivated by antimicrobial resistance.
- Diversity helps candidates cover distinct modes of an imperfect goodness function, increasing the chance that at least one survives downstream evaluation.
- Multi-stage evaluation filters candidates using increasingly precise and costly oracles, from simulations and wet-lab experiments toward clinical trials.
- Active learning iteratively proposes batches, receives oracle feedback, and uses a surrogate's acquisition function to select candidates balancing expected reward and epistemic uncertainty.
- GFlowNets can sample candidates proportionally to a reward function, producing sets that cover reward modes while retaining high scores and diversity.
- The proposed method combines a GFlowNet generator with active learning, off-policy training from static datasets, and epistemic uncertainty, then evaluates it on three protein and DNA tasks.
2. Problem Setup
The problem is sequential black-box optimization over discrete biological objects with limited, batched oracle queries and an available initial labeled dataset. The objective is to generate top candidates that are useful, diverse, and novel relative to known data.
- The algorithm searches a discrete object space for candidates maximizing a non-negative usefulness oracle under N queries arranged in batches of size b.
- Each round proposes a new batch, evaluates it with the oracle, and augments the current dataset with candidate-score pairs for the next round.
- An initial labeled dataset D0 provides a practical historical starting point, distinguishing this setup from standard black-box optimization and connecting it to multi-round offline model-based optimization.
- The target is a diverse set of K top-scoring candidates drawn from generated candidates rather than the initial dataset.
- Candidate quality is assessed through usefulness, diversity across oracle modes, and novelty relative to the initial dataset.
- All three metrics are computed on the TopK candidates and should be considered jointly because any single metric can give a misleading evaluation.
3. GFlowNets For Sequence Design
GFlowNet-AL learns a stochastic generator for diverse sequence candidates, guided by proxy predictions and epistemic uncertainty across active-learning rounds. It also incorporates labeled data during GFlowNet training to improve learning while retaining exploration.
- GFlowNet foundations: GFlowNets learn policies that sample discrete objects with probability proportional to a non-negative reward, supporting diverse candidate generation.For sequence generation, objects are constructed autoregressively by appending tokens to partial sequences.
- Active-learning loop: GFlowNet-AL trains a proxy on the current dataset, uses its mean and uncertainty through an acquisition function as the reward, samples candidates, and evaluates them with an oracle.The evaluated batch is added to the dataset before the next active-learning round.
- Exploration: Exploratory training mixes the learned forward policy with uniform actions, helping prevent training from becoming stuck in one or a few modes.This mixture is analogous to ϵ-greedy exploration in reinforcement learning.
- GFlowNet training: Trajectory balance is the paper’s GFlowNet objective, with a learnable log Zθ intended to improve credit assignment, robustness, and learning speed.The paper uses Equation 11 for GFlowNet training.
- Leveraging data during training: Offline trajectories from the available dataset are mixed with on-policy trajectories during GFlowNet training, while relying exclusively on the dataset can be sub-optimal because it may not cover X.The offline proportion γ interpolates between online and offline training settings.
- Leveraging data during training: Mixing offline trajectories is hypothesized and experimentally verified to improve convergence speed and guarantee exploration around examples in the dataset.These benefits are presented specifically for active-learning settings.
4. Related Work
Biological sequence design has been studied with reinforcement learning, Bayesian optimization, deep generative models, and deep model-based optimization, among other approaches.
- Prior approaches to biological sequence design include reinforcement learning, Bayesian optimization, deep generative models, and deep model-based optimization.
5. Experiments
Experiments evaluate GFlowNet-AL across AMP, TF-Bind-8, and GFP sequence-design tasks using performance, diversity, and novelty criteria. Results show strong performance on AMP and TF-Bind-8, while GFP exposes a proxy-model failure mode; ablations examine offline data and uncertainty estimation.
- Tasks and Evaluation Criteria: Experiments cover AMP, TF-Bind-8, and GFP design tasks with varying sequence lengths, vocabularies, rounds, batch sizes, and starting datasets.AMP uses 10 rounds and batches of 1000; TF-Bind-8 and GFP use single-round settings with batches of 128.
- Tasks and Evaluation Criteria: Performance, diversity, and novelty are evaluated on the highest-scoring generated candidates, with additional percentile metrics for TF-Bind-8 and GFP.The additional metrics are the 100th and 50th percentile scores on generated batches.
- Anti-Microbial Peptide Design: GFlowNet-AL generates more diverse and novel AMP sequences than the baselines while achieving better final TopK performance.COMs performs poorly for large diverse batches, while AmortizedBO generated nonsensical peptides dominated by W and C.
- Anti-Microbial Peptide Design: The amino-acid distribution of GFlowNet-AL’s generated peptides closely matches that of known AMPs, while the average instability index is 26.5 with a maximum of 36.An instability score above 40 indicates instability.
- TF-Bind-8: On TF-Bind-8, GFlowNet-AL outperforms other baselines in TopK performance and novelty, although MINs achieves the best diversity score.Considering all metrics, MINs has lower performance and novelty, and GFlowNet-AL also outperforms evaluated methods on the 100th- and 50th-percentile metrics.
- GFP: GFlowNet-AL does not outperform baselines on GFP because a proxy trained on biased data can create spurious reward modes absent from the true reward.De novo generation by GFlowNet-AL and DynaPPO consequently struggles to find good candidates in this setting.
- Training with the Oracle Data: Using offline data improves GFlowNet training speed and performance up to γ = 0.50, whereas larger γ reduces exploration and generalization.The ablation samples 4096 examples every 1000 training steps and evaluates average Top100 reward.
- Effect of Uncertainty Estimates: Ablations compare Deep Ensembles, MC Dropout, and a single proxy model to study how uncertainty estimation affects GFlowNet-AL.The supplied results passage states that any uncertainty estimate can provide an advantage over having none.
6. Conclusion and Future Work
The paper introduces a GFlowNet-based generative active-learning algorithm for biological sequence design, combining epistemic uncertainty and acquisition scoring to produce diverse, novel candidates. It accelerates GFlowNet training with oracle-derived trajectories, while noting limitations involving two learners and poor de novo GFP performance.
- Conclusion: The proposed algorithm uses GFlowNets as candidate generators and Bayesian-optimization principles to produce diverse and novel sequence batches.It estimates epistemic uncertainty and uses an acquisition function to score candidates.
- Conclusion: Oracle-labeled training sequences are used to construct additional GFlowNet training trajectories, accelerating training.
- Conclusion: The method improves diversity and relative novelty of generated candidates, especially when epistemic uncertainty and oracle-derived empirical distributions are used.
- Limitations and Future Work: Using both a proxy model and a generative policy introduces two separate learners, each with its own hyperparameters.
- Limitations and Future Work: The method performs poorly on the GFP task, where de novo generation struggles to find good candidates.
A. Task Details
The experiments cover anti-microbial peptides, TF-Bind-8 DNA sequences, and GFP proteins, using task-specific datasets and oracle constructions. Dataset sizes, sequence spaces, and evaluation settings differ across these biological design tasks.
- A.1. Anti-Microbial Peptides: The AMP dataset contains 6,438 positive AMPs and 9,522 non-AMPs after filtering peptides by length and modification type.Selected peptides have lengths between 12 and 60 and target Gram-positive bacteria.
- A.1. Anti-Microbial Peptides: The AMP dataset is split into D1 for algorithm use and D2 for training the oracle that simulates wet-lab feedback.The split follows a strict grouping principle based on shared targets, titles, or clusters.
- A.2. TF-Bind-8: The TF-Bind-8 dataset contains 65,792 samples covering every possible length-8 nucleotide string, with 32,898 samples used for model training.Because the dataset contains all possible sequences, its oracle is exact.
- A.3. GFP: The GFP task uses 56,086 protein sequences of length 237 and evaluates maximum fluorescence in GFP derivatives.The dataset and oracle come from the Design-Bench implementation.
B. Implementation Details
The implementation evaluates published baselines and configures the proposed method with neural proxy models, acquisition functions, and a trajectory-balance GFlowNet generator. Baseline and generator settings vary across tasks where needed.
- Baselines: Baseline implementations include published AmortizedBO and adapted DynaPPO code, plus COMs, MINs, CbAS, BO-qEI, and CMA-ES implementations.
- Baselines: AmortizedBO varies mutation count, acquisition function, and generations between proposals, while DynaPPO varies trajectories, learning rate, exploration penalty, and radius.
- Software: The proposed GFlowNet-AL algorithm is implemented in PyTorch.
- Proxy Model: The proxy uses a two-hidden-layer MLP of dimension 2048 with ReLU activations across all three tasks.Ensembles contain five members, while MC Dropout uses 25 samples with dropout rate 0.1.
- GFlowNet Generator: The GFlowNet generator uses a two-hidden-layer MLP, trajectory-balance training, Adam optimization, and an offline-trajectory proportion γ of 0.5.The log Z learning rate is 10^-3 across experiments.
- Proposed Method: The proposed GFlowNet-AL results use an ensemble proxy model with UCB as the acquisition function.
C.1. AMP Generation: Additional Results
Additional AMP results show that AmortizedBO generated overwhelmingly poor and unnatural sequences, including outputs dominated by repeated amino-acid patterns. These sequences lacked characteristics commonly found in antimicrobial peptides.
- AMP Generation: All top 1,000 AmortizedBO sequences reached the AMP task’s maximum allowed length.
- AMP Generation: The generated sequences were unnatural and lacked important AMP characteristics, including the dominance of amino acid K.
- AMP Generation: Dynamic-length AmortizedBO was implemented by adding a stop token, but its outputs still overwhelmingly used the maximum sequence length.
C.2. TF-Bind-8 and GFP: Additional Results
Additional results on GFP and TF-Bind-8 report both high-end and typical sequence scores, with GFlowNets outperforming baselines under these percentile metrics.
- GFlowNets outperform baselines on the 100th-percentile GFP and 50th-percentile TF-Bind-8 results.These evaluations correspond to maximum GFP scores and median TF-Bind-8 scores.
- The GFP evaluation reports maximum and median scores for proposed sequences.
- The TF-Bind-8 evaluation reports maximum and median scores for proposed sequences.
C.3. Effect of Uncertainty: Additional Results
Across additional experiments, uncertainty estimation consistently affects performance, while results are organized across acquisition functions and multiple biological sequence-design tasks.
- The uncertainty model is consistently the key factor affecting performance across the three tasks.
- Additional results include AMP Generation, TF-Bind-8, and GFP tasks, including GFlowNet-AL-None without proxy uncertainty.
- The experiments compare UCB and EI acquisition functions with different uncertainty-estimation methods.