[Paper Review] Protein Design with Guided Discrete Diffusion

Nate Gruver, Samuel C. Stanton+7 more|arXiv (Cornell University)|May 31, 2023

Gene Regulatory Network Analysis18 citations

TL;DR

The paper introduces diffusioNOS, a gradient-guided sampling method for discrete diffusion in sequence space, and combines it with LaMBO-2 for multi-objective antibody design, achieving high expression and binding in exploratory in vitro experiments.

ABSTRACT

A popular approach to protein design is to combine a generative model with a discriminative model for conditional sampling. The generative model samples plausible sequences while the discriminative model guides a search for sequences with high fitness. Given its broad success in conditional sampling, classifier-guided diffusion modeling is a promising foundation for protein design, leading many to develop guided diffusion models for structure with inverse folding to recover sequences. In this work, we propose diffusioN Optimized Sampling (NOS), a guidance method for discrete diffusion models that follows gradients in the hidden states of the denoising network. NOS makes it possible to perform design directly in sequence space, circumventing significant limitations of structure-based methods, including scarce data and challenging inverse design. Moreover, we use NOS to generalize LaMBO, a Bayesian optimization procedure for sequence design that facilitates multiple objectives and edit-based constraints. The resulting method, LaMBO-2, enables discrete diffusions and stronger performance with limited edits through a novel application of saliency maps. We apply LaMBO-2 to a real-world protein design task, optimizing antibodies for higher expression yield and binding affinity to several therapeutic targets under locality and developability constraints, attaining a 99% expression rate and 40% binding rate in exploratory in vitro experiments.

Motivation & Objective

• Motivate efficient protein design in the discrete sequence space without relying on inverse folding.
• Develop a gradient-guided discrete diffusion sampling method (NOS) to optimize sequence-level objectives.
• Enhance LaMBO with NOS and saliency-driven edits for multi-objective antibody design.
• Demonstrate practical improvements in expression and binding for antibodies under realistic constraints.

Proposed method

• Propose NOS (diffusioN Optimized Sampling) to guide discrete diffusion by applying gradients in the denoising network’s hidden states.
• Formulate NOS in two variants: NOS-C (continuous latent space guidance) and NOS-D (discrete transition guidance).
• Combine NOS with Langevin dynamics to sample from the modified transition distributions while preserving sequence likelihood.
• Integrate NOS into LaMBO-2 by replacing the MLM sampler and using saliency maps to select edit positions for multi-objective optimization.
• Use partial deep ensembles to estimate posterior objectives and employ acquisition functions for multi-objective optimization.
• Perform saliency-based edit-position selection to focus edits on influential sequence sites.

Experimental results

Research questions

• RQ1 Can NOS improve the trade-off between sequence likelihood and objective value for guided discrete diffusion in protein design?
• RQ2 Does saliency-driven edit-position selection enhance efficiency and performance in Bayesian optimization for antibody design?
• RQ3 How does LaMBO-2 perform in real-world antibody lead optimization with multi-objective constraints (expression yield, binding) in silico and in vitro?

Key findings

• NOS outperforms several state-of-the-art structure- and sequence-based baselines on unguided and guided infilling tasks in silico.
• LaMBO-2, combining NOS with saliency-guided edits, enables stronger performance with limited edits and multi-objective optimization.
• In antibody design tasks, LaMBO-2 yields higher acquisition-value samples while maintaining naturalness constraints better than genetic baselines.
• Antibody lead optimization with LaMBO-2 achieved multiple submicromolar binders across four targets in vitro after iterative rounds.
• The exploratory in vitro experiments report a 99% expression rate and 40% binding rate for optimized antibodies.