[Paper Review] Biological Sequence Design with GFlowNets

Moksh Jain, Emmanuel Bengio+11 more|arXiv (Cornell University)|Mar 2, 2022

vaccines and immunoinformatics approachesBiochemistry, Genetics and Molecular Biology20 citations

TL;DR

The paper proposes GFlowNet-based active learning (GFlowNet-AL) to generate diverse, high-scoring biological sequences by coupling a GFlowNet generator with an uncertainty-aware proxy and offline data augmentation.

ABSTRACT

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.

Motivation & Objective

• Motivate the need for diverse candidate design in multi-stage, costly wet-lab pipelines for biological sequences.
• Introduce an active learning framework (GFlowNet-AL) that uses GFlowNets to generate diverse candidate sequences.
• Leverage epistemic uncertainty and offline data to speed up training and improve exploration.
• Evaluate the approach on multiple protein and DNA design tasks and compare with baselines.

Proposed method

• Use a GFlowNet generator to sample candidates with probability proportional to a reward R(x).
• Define R(x) as a function of a proxy model’s mean μ(x) and epistemic uncertainty σ(x) via an acquisition function F(μ,σ).
• Train a proxy model M to output μ and σ for sequences x; update M using the current dataset D_i.
• Optionally augment GFlowNet training with offline trajectories from the dataset to speed learning and improve exploration around known examples.
• Incorporate epistemic uncertainty using MC Dropout or ensembles, combined with acquisition functions (Upper Confidence Bound or Expected Improvement).
• Train the GFlowNet using flow-matching or trajectory balance objectives to align terminal flow with rewards.

Experimental results

Research questions

• RQ1 Can a GFlowNet-based generator produce highly scoring biological sequences while maintaining diversity and novelty?
• RQ2 Does incorporating offline data and epistemic uncertainty improve learning speed, exploration, and quality of generated batches in active learning for sequence design?
• RQ3 How does GFlowNet-AL perform across different sequence design tasks (protein and DNA) compared with established baselines?
• RQ4 What is the impact of offline trajectories and uncertainty-aware rewards on the diversity, novelty, and performance of TopK selections?

Key findings

• GFlowNet-AL achieves higher diversity and novelty in AMP design than baselines and shows strong TopK performance.
• On TF-Bind-8, GFlowNet-AL delivers the best TopK and novelty, though MINs shows higher diversity (overall evaluation favors GFlowNet-AL when considering all metrics).
• In GFP tasks, GFlowNet-AL underperforms relative to some baselines in performance but maintains competitive diversity and novelty; results noted caveats related to proxy bias and evaluation.
• Across tasks, GFlowNet-AL generally provides improvements in diversity and novelty metrics compared to baselines, with competitive or superior performance on several tasks.
• The study demonstrates that combining offline data and uncertainty-aware rewards can accelerate learning and improve exploration around known examples.