PaperlayerView exact revision on arXiv

Source-linked AI summary

Protein Design with Guided Discrete Diffusion

Nate Gruver, Samuel Stanton, Nathan C. Frey, Tim G. J. Rudner, Isidro Hotzel, Julien Lafrance-Vanasse, Arvind Rajpal, Kyunghyun Cho, Andrew Gordon Wilson

arXiv:2305.20009v2cs.LGq-bio.BM

TL;DR

Protein design needs efficient optimization over vast discrete sequence spaces, but structure-based methods face inverse-folding and structural-data limitations, while discrete sequences hinder direct gradient guidance. The paper introduces NOS for hidden-state-guided discrete diffusion and combines it with saliency-based edits and ensemble uncertainty in LaMBO-2. LaMBO-2 produced high-value feasible solutions in silico and, in exploratory antibody experiments, generated multiple submicromolar binders across four targets with a median of five edits.

  • Problem

    Protein design must optimize vast discrete sequence spaces under expensive validation, while structure-based approaches require inverse folding and are limited by scarce structural data.

  • Method

    NOS guides categorical and continuous diffusion through denoising-network hidden-state gradients, while LaMBO-2 adds saliency-selected edits and ensemble-based uncertainty to Bayesian sequence optimization.

  • Results

    Multiple submicromolar binders were generated for all 4 antibody targets in exploratory in vitro experiments, with a median of 5 edits to the seed.

  • Takeaways & Limitations

    LaMBO-2 can support antibody library enrichment while optimizing expression yield and binding affinity under locality and developability constraints.

  • Takeaways & Limitations

    The wet-lab validation differs substantially from prior work, preventing a true apples-to-apples comparison.

Abstract

from arXiv · show

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.

1 Introduction

Protein design must search a vast discrete sequence space efficiently, yet structure-based approaches require difficult inverse folding and scarce structural data. NOS addresses discrete optimization directly in sequence space and is combined with LaMBO to form LaMBO-2.

  • Motivation: Protein design seeks high-function sequences while limiting experiments to small enriched candidate libraries because the sequence space is vast and validation is costly.These libraries are commonly obtained through high-throughput screening or antibody-immunization workflows.
  • Motivation: Structure-space design requires inverse folding, may produce unrealizable structures, and is constrained by limited high-quality structural data.Direct sequence search avoids sequence recovery and can exploit much larger sequence datasets.
  • Motivation: Discrete sequence optimization limits direct use of gradients, while gradient-free methods can be computationally expensive.Continuous latent-space methods address this challenge but require learned decoders to recover discrete sequences.
  • Contributions: NOS provides gradient-guided sampling for discrete diffusion, and LaMBO-2 combines NOS with Bayesian optimization for protein design.The paper also uses saliency-based edit selection to improve optimization under limited edits.
  • Contributions: NOS outperforms state-of-the-art structure- and sequence-based baselines on unguided and guided infilling tasks in in silico validation.The supplied passage reports this as the paper’s validation finding without specifying individual benchmark values.

2 Related Work

Prior work provides discrete diffusion models and gradient-guided generation, but general categorical diffusions have received limited guidance research. The paper positions NOS against autoregressive, embedding-based, and evolutionary alternatives.

  • Discrete diffusion: Discrete diffusion approaches either apply categorical noise directly to sequences or apply Gaussian corruption to token-vector embeddings.These approaches have competed with autoregressive models in text generation tasks.
  • Guided generation: Gradient guidance is established for continuous diffusion and has been adapted to autoregressive and embedding-based language models.PPLM and related methods use gradient steps in continuous representations or decoder activations.
  • Guided generation: General categorical diffusions that predict denoised categorical distributions have received comparatively little gradient-guidance research.NOS explores guidance for this setting and is evaluated against PPLM and DiGress.
  • Protein design baselines: Evolutionary optimizers such as AdaLead and PEX provide black-box baselines for protein sequence optimization.The paper compares them with NOS in protein-design experiments involving in silico fitness proxies.

3 Background

Protein design optimizes one or more sequence objectives under expressibility and edit constraints, while diffusion models generate samples by reversing corruption processes. Categorical and continuous-noise formulations provide alternative sequence-generation mechanisms.

  • Protein design formulation: Protein design searches fixed-length sequences for single or multiple objectives, optionally under a fixed edit budget relative to a seed sequence.Expressibility is required because non-expressing proteins have undefined measured properties.
  • Diffusion models: Continuous-noise diffusion embeds categorical sequences before applying Gaussian noise, whereas categorical-noise diffusion corrupts tokens directly.These two formulations are contrasted in Figure 2.
  • Protein design formulation: Naturalness metrics proxy expressibility, but naturalness can conflict with the objectives of designing previously unobserved sequences.The paper represents this trade-off through an unnormalized density proportional to p(w) exp(v(w)).
  • Diffusion models: Denoising diffusion reverses a corruption process from clean data to a prior by repeatedly predicting denoised variables and sampling reverse transitions.The background describes Gaussian corruption and learned reverse distributions for continuous representations.
  • Diffusion models: Categorical diffusion applies discrete transitions directly to sequence tokens, using [MASK] absorbing states and categorical reverse sampling.The denoiser predicts clean-token distributions from corrupted sequences.

4 Methods

NOS guides discrete diffusion through denoising-network hidden states while preserving likelihood, and LaMBO-2 integrates NOS with saliency-based edit selection and ensemble uncertainty. The resulting framework supports iterative, constrained protein sequence design.

  • 4.1 NOS: diffusioN Optimized Sampling: NOS modifies denoising transitions by applying value-function gradients to continuous hidden-state representations of corrupted discrete sequences.The method supports both categorical and Gaussian diffusion formulations.
  • 4.1 NOS: diffusioN Optimized Sampling: The modified transition combines objective guidance with a likelihood-preserving regularization term.The KL term keeps transitions close to the original denoiser prediction while gradients favor higher objective values.
  • 4.1 NOS: diffusioN Optimized Sampling: NOS differs from PPLM by refining the entire sequence iteratively, allowing tokens across the sequence to change at each optimization step.This distinction is identified as particularly important for protein design.
  • 4.2 LaMBO-2: function-guided protein design: The antibody saliency example concentrates edits on CDRH3 while retaining some allocation to framework and other CDR positions.CDRs are marked on the variable-heavy-region visualization.
  • 4.2 LaMBO-2: function-guided protein design: LaMBO-2 converts value-function gradients at sequence embeddings into an edit-position distribution, while preserving fixed positions by assigning them zero saliency.The approach allocates more editing resources to positions that most affect the value estimate.

5 Experiments

The experiments evaluate sequence diffusion, NOS-guided optimization, and LaMBO-2 for antibody design across unguided infilling, in silico objectives, and in vitro lead optimization. Results favor sequence-only diffusion for CDR recovery, show LaMBO-2 preserves naturalness during optimization, and report high expression and binding rates in the final experimental round.

  • 5.1 Unguided antibody CDR infilling: Sequence diffusion methods consistently outperform structure-based methods on unguided antibody CDR infilling when only predicted structures are available.Diffusion infills are often on-par or better than IgLM, while DiffAb and RFDiffusion frequently underperform sequence-only diffusion.
  • 5.2 Optimizing antibodies for in silico objectives: NOS produces higher likelihood samples at similar or substantially improved objective values than alternative guided and structure-based generation methods.The comparison includes DiffAb, RFDiffusion, PPLM, and DiGress across multiple guidance strengths.
  • 5.3 Antibody lead optimization: in silico evaluation: LaMBO-2 preserves naturalness while improving acquisition value, whereas AdaLead and PEX rapidly leave the well-supported natural-sequence region.PEX produces infeasible samples beyond 4 iterations, AdaLead violates the constraint after 2 iterations, and LaMBO-2 remains feasible through 16 steps.
  • 5.3 Antibody lead optimization: in silico evaluation: Salient edit-position selection has a much larger effect on objective value than NOS guidance in the few-edit regime.The ablation varies edit budgets from 8 to 32 and compares saliency-based position selection with uniform random selection and unguided edits.

6 Discussion

The paper outlines extensions to LaMBO-2 beyond protein design and details the categorical diffusion processes used for masked-sequence generation.

  • Future work: LaMBO-2 could extend to small-molecule design because its improvements are not protein-specific.Such applications may require padding or variable-length corruption because small molecules do not admit sequence alignments.
  • Future work: Optimizing sequences longer than 20K tokens may require implicit convolutions or modified self-attention.The paper identifies gene perturbations as an example of this longer-sequence setting.
  • Future work: The authors identify classifier-free guidance and class-conditional generation as directions for goals difficult to express as black-box functions or constraints.
  • Categorical diffusion: The reverse process uses a denoising model to reconstruct less-corrupted sequences from masked inputs.
  • Categorical diffusion: The categorical diffusion forward process gradually corrupts sequence tokens with [MASK] tokens according to transition matrices and a noise schedule.The cumulative transition matrix is defined as ¯Q_t = Q_1Q_2...Q_t, with α_t controlling the schedule.

B Methodological Details

The methodological details formulate guided sampling as posterior-like sequence generation, using hidden-state Langevin updates while regularizing diffusion transitions toward the unguided model.

  • Infilling: Infilling preserves selected residues through a mask while sampling the remaining positions with categorical denoising diffusion.The same procedure is extended with a guidance Langevin step for guided infilling.
  • Discrete guidance: The discrete posterior regularizer compares guided and unguided reverse transition distributions, rather than relying on a continuous score function.
  • Guided sampling: NOS decomposes the denoiser into a language-model head and trunk, then guides sampling in shared hidden representations.The modified denoising distribution weights candidate reconstructions by the value function.
  • Guided sampling: Langevin updates combine value gradients, a KL regularizer, and optional Gaussian noise before the next reverse diffusion step.The authors set the additional stochasticity weight λ3 to zero in experiments because reverse transitions already introduce randomness.
  • Training data: The diffusion models are trained on paired antibody heavy- and light-chain sequences from OAS and SAbDab aligned with ANARCI.

C.1 Infilling experiment

The infilling experiment uses unguided categorical diffusion with masks determined by the selected complementarity-determining regions.

  • Infilling experiment: The trained diffusion models generate antibody infills using Algorithm 1 without guidance and Chothia numbering for consistent CDR specification.Baseline methods receive wrapper scripts that convert the selected CDR identifiers into their native formats.

C.2 MCMC comparison

The comparison examines MCMC and diffusion-based sampling, including iterative refinement, guidance settings, and hidden-state optimization details.

  • MCMC comparison: Hundreds of thousands of burn-in steps are required in practice for the Metropolis-Hastings chain to obtain good samples.The chain uses uniform random mutations with a temperature-dependent acceptance correction.
  • MCMC comparison: Two orders of magnitude faster convergence is reported for the diffusion models while reaching similar energy values to Metropolis-Hastings-adjusted MCMC.
  • Energy construction: The energy combines IgLM sequence probabilities with a beta-sheets objective trained on IgLM representations.The guidance strength is tuned to produce sequences with approximately 40% beta sheets.
  • PPLM details: Larger PPLM regularization values prevent samples from differing substantially from the unguided model, while moderate larger step sizes can increase objective values.
  • Iterative refinement: Iterative refinement noticeably improves beta-sheets objective values compared with contrived autoregressive sampling.The effect is much more limited for the SASA objective, and iterative refinement alone is not considered sufficient.

C.4 Model Architecture and Training

The experiments compare NOS variants, guidance choices, and baselines while integrating guided diffusion into LaMBO-2's constrained sampling procedure. Results emphasize likelihood–fitness trade-offs and the differing effects of guidance location and regularization.

  • Sampling: Iterative refinement consistently improves NOS-D objective values, but not enough to establish refinement alone as sufficient for strong sampling.The comparison contrasts diffusion-style and autoregressive-style sampling.
  • Hyperparameters: λ is the most important hyperparameter because it controls how far samples move from seed sequences.η affects hidden-state optimization when chosen effectively.
  • Hyperparameters: Guidance in the first hidden state tends to work better for NOS-C, whereas final-hidden-state guidance tends to work better for NOS-D.The guidance layer is either the token embeddings or the final representations before the linear head.
  • Architecture: NOS uses Gaussian or discrete corruptions with learned embeddings, while DiGress uses discrete corruptions with fixed one-hot encodings.All models use the same backbone transformer and regression heads for comparison.
  • Baseline comparison: Both NOS methods improve objective values without sacrificing likelihood, unlike baselines that either remain near seeds or reduce likelihood substantially.DiffAb and IgLM cluster near starting seeds, while RFDiffusion explores more diverse but often lower-likelihood samples.
  • LaMBO-2 sampling: LaMBO-2 performs projected SGLD updates in denoiser hidden states, decodes candidate sequences, and returns the best feasible sample under value and constraint checks.The step begins with corruption and hidden-state initialization, then iterates stochastic projected updates and decoding.

D.3 Architecture and Hyperparameters

LaMBO-2 conditions antibody sequence generation on antigen and unmasked antibody context while jointly predicting generative token distributions and protein-property distributions. Its architecture supports multi-task prediction and uncertainty estimation under edit and developability constraints.

  • Inputs: LaMBO-2 takes antibody VH and VL regions plus the unaligned antigen sequence as inputs.Antigen concatenation makes generation conditional on both antigen and unmasked antibody sequence.
  • Model outputs: The model jointly predicts antigen-conditional token distributions for corrupted positions and discriminative distributions over protein properties.Predictions intended to be antigen-independent are made invariant using random-antigen augmentation.
  • Architecture: A shared encoder connects a generative discrete-diffusion head with discriminative heads for expression and affinity.Ten ensemble heads per discriminative task provide uncertainty estimates.
  • Training: Training uses random task-head and task-minibatch selection in a jointly trained multi-task model.This design addresses partially missing measurements across protein properties.
  • Generation: Generation samples 1K designs from expert-selected antibody seeds using a shared edit budget of B = 16 and 16 diffusion steps.An inverse linear noise schedule was used because it performed better in this experiment.
  • Constraints: Designs must satisfy locality constraints and sequence-liability constraints, including canonical cysteine conservation and no unpaired cysteines.These constraints target antibody developability and structural integrity.

D.4 Training Data, Class Imbalance, and Label Smoothing

LaMBO-2 trains its property predictors on expression and affinity datasets with different coverage and uses label discretization to address severe biological-data imbalance. Corrupted-input training aligns discriminator behavior with the diffusion process.

  • Training data: Expression heads use 10K measurements augmented to 160K rows, while binding heads use 10K SPR measurements augmented to 12K rows.Expression augmentation pairs measurements with random antigens to encourage antigen invariance.
  • Training data: Expression and affinity datasets only partially overlap, motivating the model's multi-task architecture.The generative diffusion head is trained only on binding antibody–antigen pairs from the SPR data.
  • Scope: The LaMBO-2 models were not pretrained, although the authors state that suitable pretraining might improve performance.They note that the usefulness of available antibody corpora for this purpose remains unclear.
  • Label discretization: Continuous properties such as expression yield and binding affinity are discretized so minority classes can be upsampled.The scheme treats antibody–antigen pairs with pKD below 4 as non-binders and uses pKD deciles for remaining measurements.
  • Noisy-input training: Discretized labels also enable discriminator training on corrupted inputs using a label-smoothing-inspired procedure.The discriminators follow the diffusion noise schedule and use cross-entropy loss.
  • Noisy-input training: As diffusion inputs become more defined, the value-gradient strength grows, while highly corrupted inputs drive the discriminator toward a uniform prior.This training avoids evaluating value gradients on out-of-distribution inputs during generation.

D.5 Baselining LaMBO-2 Against Unguided Sequence and Structure-Based Diversification:

LaMBO-2 is evaluated against unguided sequence and structure-based diversification, including in silico objectives and experimentally validated antibody designs. The comparisons show a trade-off between predicted optimization and naturalness, while structure-based infilling favors naturalness with limited predicted functional shift.

  • Structure-based diversification: Unguided structure-based infilling produces high-likelihood samples, but antigen conditioning yields only a slight shift toward better predicted function.The comparison includes DiffAb and RFDiffusion infilling of CDR H2 and H3 in hu4D5.
  • Saliency: Independent saliency models trained with different corruption processes substantially agree on important antibody regions, especially CDRH3 and CDRL1.Some framework-region attribution may relate to the fitness of antibody germlines.
  • Saliency: Saliency maps are used online to choose antibody positions for intervention, but their reliability in decision-making contexts remains an open question.The authors distinguish this use from the offline feature-attribution setting studied in much prior work.
  • In vitro validation: Experimental validation expresses and purifies LaMBO-2 designs, measures affinity with SPR, and relates yield and affinity to edit distance from the seed.Figure 18 reports experimentally validated yield for expressing designs and affinity for binding designs.
Loading 2305.20009v2…