Title: Protein Design with Guided Discrete Diffusion
Venue: NeurIPS 2023
Accepted: yes
Keywords: protein design

diffusion model

classifier guidance
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.

Protein Design with Guided Discrete Diffusion

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 diffusio N O ptimized S ampling (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 realworld 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

Optimizing protein sequences for improved function has the potential for widespread impact [63]. Among many potential applications in engineering and medicine, engineered antibodies can be used to create cancer therapeutics that are much less harmful to the patient than radiotherapy or chemotherapy. Because the space of possible proteins is vast and discrete, and experimental validation is slow and expensive, all practical methods for protein design must restrict themselves to a small enriched library of candidates to find a viable option in as few measurements as possible [44]. In practice these enriched libraries are usually obtained through massive high-throughput in vitro screening [67], or in the case of antibodies by injecting a live animal with the target antigen and sequencing the animal's immune response [52]. Generative protein models offer the tantalizing prospect of enriched libraries produced nearly instantly at a fraction of the cost. Success in real-world applications, however, has proven elusive, in part because naive generative models produce outputs that are similar to their training data and are therefore unlikely to improve target qualities [53].

There are many approaches to guided generation of proteins, but one broad and important distinction is between methods that search in sequence space and those that search in structure space. A basic tenet of molecular biology is 'sequence determines structure, structure determines function' [9]. Hence when optimizing a protein for a desired function, it may seem more direct to design the protein in structure space, where gradient-based sampling methods can be used in tandem with carefully

1

[HEAVY] EVQLVQ

S

GAEV           KGSGYSF

G

TKLT [LIGHT] VLTSEADYYCQ           FGGG

T

KLTVL [ANTIGEN] KL...

corrupt

diffusioN Optimized Sampling (NOS)

likelihood

likelihood

fitness

fitness

: seed sequences

[HEAVY] EVQLVQ

L

GAEV           KGSGYSF

E

TKLT [LIGHT] VLTSEADYYCQ           FGGG

L

KLTVL [ANTIGEN] KL...

engineered potentials [1, 45, 77]. One of the drawbacks of this approach is the optimized structure must still be converted back to an amino acid sequence in order to be synthesized, a task known as 'inverse-folding' [19]. There is no guarantee that the optimized structure can be realized by an actual sequence, and the inverse-folding procedure may not find the sequence if it exists. Structural models are also computationally intensive and limited by the scarcity of high-quality structural data. Searching directly in sequence space eliminates the need to recover sequence from structure. Protein sequence models are also comparatively fast, especially during inference, and can leverage sequence datasets that are often several orders of magnitude larger than their structural equivalents.

Although sequence models are arguably the most practical foundation for protein design, they have historically suffered from the challenges of optimizing discrete sequences, where gradient-based sampling is not directly applicable. As a result, many sequence search methods resort to gradient-free sampling methods like Metropolis-Hastings MCMC [78, 37], which are flexible but computationally expensive, eroding a key advantage over structure search. Several methods have been proposed that maintain gradient-based search by performing guidance in a continuous latent space, with a learned decoder to sample discrete sequences [33, 32]. Notably, Stanton et al. [73] proposed LaMBO ( La tent M ulti-Objective B ayesian O ptimization), an optimization method that uses masked language model (MLM) style denoising guided with Bayesian acquisition values to address the online, multi-objective nature of real-world protein design. While LaMBO can quickly sample sequences with improved acquisition value, it has two key limitations. First, LaMBO is built on top of MLMs which, while strong representation learners are not strong generative models. In particular, MLMs lag behind other methods in producing high likelihood samples or infills. Second, despite being designed to improve known sequences instead of designing them completely from scratch, LaMBO and related methods have no principled framework for simultaneously enforcing an edit budget and choosing optimal edit locations based on that budget.

To address the first issue we propose NOS (diffusio N O ptimized S ampling), a new method for controllable categorical diffusion (Figure 1). Diffusion models capture complex distributional dependencies by making iterative denoising steps, but there is relatively little previous work on how to control these processes. NOS generates sequences with both high likelihood and desirable qualities by taking many alternating steps between corruption, guidance, and denoising in the continuous latent space of the model. Our in silico validation shows that NOS outperforms many state-of-the-art structure and sequence-based baselines on both unguided and guided infilling tasks. 2 To address the second problem (choosing optimal edit locations) we propose using embedding-gradient feature attributions (i.e. saliency maps) to determine which positions on the sequence are most important to edit to improve the guidance objective value. We combine NOS with saliency-selected edits to create LaMBO-2, a more powerful variant of the original LaMBO algorithm. Exploratory in vitro experimental validation of our designs provides evidence that LaMBO-2 can be used to create enriched antibody libraries without the aid of additional in vitro high-throughput screening.

2 Related Work

Discrete diffusions Austin et al. [4] and Hoogeboom et al. [40] constructed diffusion models for discrete categorical data using a categorical noise process. Recently, categorical diffusion has shown promise as a competitor to autoregressive models in text generation for machine translation and summarization. The approaches can be roughly grouped into methods that apply categorical noise distributions directly to sequences (CMLM [31], SUNDAE [65]), and those that apply Gaussian corruptions to token-vector embeddings (SED [74], CDCD [21]). In this work we show that NOS can guide both types of categorical diffusions. Within the space of protein design, our method is also closely related to joint diffusion models over both sequence and structure [2, 51], which also circumvent inverse folding. Because these models still rely on structure information at training time, they can be limited by data availability in the same manner as pure structure models.

Discrete generative guidance Gradient guidance typically augments sampling from a generative model with gradient steps to increase the occurrence of a desired attribute [54]. Gradient guidance is natural within the framework of continuous diffusion models [20], and Li et al. [47] use this connection to perform gradient-guided sampling from a diffusion language model. To obtain a continuous space, they perform Gaussian diffusion [39] on word embeddings, decoding out to tokens using a linear head. The original method required many careful engineering interventions, e.g. clamping latent representations to the nearest word embedding, that have been improved by recent methods, such as CDCD [21], but gradient guidance has not been discussed for these more recent formulations.

To achieve a similar form of gradient guidance without carefully engineering a latent space, Dathathri et al. [17] and Yang and Klein [83] propose gradient-guided autoregressive models by using the decoder's activations as a gradient-friendly latent space. These methods alternate between sampling from logits and ascending the likelihood of a separately trained classifier model. Surprisingly, despite work on gradient guidance for continuous noise diffusions and autoregressive language models, there has been little work on gradient guidance for general categorical diffusions that predict denoised categorical distributions (e.g. CMLM, SUNDAE, CDCD), which is a topic we explore in detail. One closely related method proposed in the context of generative models of small molecules is DiGress [79], which performs gradient guidance on one-hot token embeddings to bias the categorical sampling distribution of a denoising model. In our setting we show that categorical and Gaussian discrete diffusions guided with NOS outperform PPLM and DiGress (Subsec. 5.2).

Genetic algorithms Evolutionary algorithms are a popular solution for black-box optimization in discrete spaces [3, 22]. These methods are often evaluated on their ability to optimize in silico proxy estimates of actual in vitro fitness, e.g. deep learning models trained on experimental datasets. In Subsec. 5.3 we baseline NOS against two genetic optimizers from protein design literature, AdaLead [70] and Proximal Exploration (PEX) [61]. We show these baselines rapidly degrade sequence likelihood as the proxy fitness is improved, limiting the effective number of edits that can be made before checking the actual fitness of the samples, ultimately limiting their sample efficiency and rate of convergence to optimal solutions. By contrast, NOS consistently improves proxy fitness while maintaining sequence likelihood.

3 Background

We pose protein design as the problem of finding sequences, w P A L with alphabet A and fixed length L , 3 which maximize a single objective f p w q (e.g., binding affinity) or multiple objectives f 1 p w q , . . . , f k p w q (e.g., expression yield, binding affinity, and aggregation tendency). Designs can be generated from random noise ( ab initio design) or by making a fixed number of edits B P t 1 , . . . , L ´ 1 u to a seed sequence s P A L . A protein is only useful if it can be synthesized (i.e. expressed), and the objective value of non-expressing proteins is undefined since their properties cannot be measured. Therefore we must introduce the constraint w P E Ă A L , where E is the set of all expressible proteins. Since naturally occurring sequences must express in order to be observed, p p w q , the likelihood of a protein with respect to an empirical distribution of natural protein sequences, is often taken as a proxy for the tendency of a protein to express. In protein design, these proxies are

Categorical Noise Process

typically called metrics of naturalness . Since we are looking for sequences that by definition have not yet been identified in nature, naturalness and our other objectives are often in tension.

We can trade off naturalness and objective value by drawing samples from the unnormalized density

where E p w q ' ´ log p p w q is a scalar energy function, and the value function v : A L Ñ R expresses the utility of a sequence with respect to our objectives. When designing proteins from primary sequence, sampling efficiently from the resulting energy function can be challenging. Simple approaches, such as the MCMC sampler used by Verkuil et al. [78] can require hundreds of thousands of steps to converge (Appendix C.2). Guided diffusion models are an appealing alternative because they construct a fixed-length Markov chain that quickly generates low-energy samples.

Diffusion models Denoising diffusion models construct samples by reversing a diffusion process that maps clean data points, x 0 , to samples from a prior π p x q (Figure 2). The forward process ( x 0 Ñ x T ) is composed of conditional distributions p p x t | x t ´ 1 q (i.e., the noise process) that admit closed forms for the conditional distributions p p x t | x 0 q and p p x t ´ 1 | x t , x 0 q (e.g., independent Gaussian corruption). The reverse process ( x T Ñ x 0 ) converts samples from the prior into samples from the learned data distribution p θ p x 0 q by repeatedly predicting the denoised variable ˆ x 0 from noisy values x t and using the conditional distribution p p x t ´ 1 | x t , ˆ x 0 q to derive a transition distribution, p θ p x t ´ 1 | x t q . The specific choice of noise process has been shown to significantly affect the likelihood and quality of image samples [71]. For categorical data there are two common approaches to constructing a diffusion generative model, depending on the nature of the noise process. We include brief descriptions below and a more detailed account in Appendix A.

Continuous noise To learn a distribution p p w q , one strategy is to first embed w to a continuous variable x 0 with embedding matrix U θ and apply Gaussian noise [21]. The prior is taken to be π p x q ' N p 0 , I q while the forward process is p p x t | x 0 q ' N p x t ; ? ¯ α t x 0 , p 1 ´ ¯ α t q I q for ¯ α t P r 0 , 1 s . The values of ¯ α t are determined by a user-specified corruption schedule. For the reverse process, we learn a function, p θ p ˆ w | x t , t q , to predict the sequence from noised points x t by minimizing the following loss with respect to θ :

Using p θ p ˆ w | x t , t q we can construct a distribution for the reverse process

where p p x t ´ 1 | x t , x 0 q is also a Gaussian distribution. At inference time, we can use the learned reverse process to convert samples from π p x q into samples from the learned distribution p θ p x 0 q by repeatedly sampling p θ p x t ´ 1 | x t q , followed by sampling w ' p θ p ˆ w | x 0 , 0 q .

Categorical noise Alternatively, Austin et al. [4] proposed a forward process which operates directly on w , by introducing an absorbing state for each token w p i q ' [MASK]. The forward process p w 0 Ñ w T q is defined by a discrete transition matrix, describing the probability of mutating a token into a [MASK], and the corresponding prior is simply a point mass on the sequence of all [MASK] tokens. To learn the parameters of the denoiser p θ p ˆ w 0 | w t , t q we maximize the likelihood of the denoising process on ground truth sequences

Then, as above, we can use the denoiser to construct the reverse process

where p p w t ´ 1 | w t , w 0 q is also a categorical distribution derived using Bayes' rule. To sample, the transition distribution is applied for t ' r T, ..., 0 s .

4 Methods

Now we present practical methods for efficiently sampling from ˜ p p w q 9 p p w q exp p v p w qq (Eq. 1) by modifying the learned transition distribution with a learned value function v θ p w q . We then show how this sampling method can be used to perform protein design through guided infilling in sequence space. As before, we provide the most salient information below and the full details in Appendix B.

4.1 NOS: diffusioN Optimized Sampling

We introduce a general form of gradient guidance (NOS) for discrete diffusions with categorical denoising models, i.e. diffusion models that predict logits over the ground truth tokens (e.g. [21, 4]). The key challenge in applying gradient guidance to categorical data is simply the lack of a continuous representation. Fortunately, in any denoising network, e.g. p θ p ˆ w | x t , t q , the discrete sequence w t has many corresponding continuous representations in the form of hidden states of the model h t ' g d p w t q for d P t 0 , . . . , D u , where D is the depth of the encoder network and g 0 p w t q ' U θ w t . Notably, for the Gaussian diffusion models in Sec. 3, we can equivalently have x t ' g 0 p w t q , as corruption and sampling are performed on the learned token embeddings. In the case of the categorical noise diffusion p θ p ˆ w 0 | w t q ' p θ p ˆ w 0 | h t q , and thus for the purpose of guidance, we can consider a general p θ p ˆ w | h t q for both forms of corruption.

To sample from ˜ p θ p w t q 9 p θ p w t q exp p v θ p w t qq , we construct a modified denoising model,

This formulation requires that the denoising model and the value function share hidden states up to depth d , and that the value function also be trained on corrupted inputs w t . In Appendix D.4 we propose a simple procedure for corrupted discriminative training inspired by label smoothing [76]. Using this modified denoising model we can construct modified transition distributions using Eq. 2 or Eq. 3. There is one key difference between these transition distributions: in the continuous case (Eq. 2), smooth steps are taken in the token embedding space, while in the discrete case (Eq. 3) the transition leads to large jumps from one token embedding to another. In either case, it is possible to sample a discrete sequence w at any point in the chain using the logits of the denoiser p θ p ˆ w | h t q . When using Eq. 2 to derive a continuous transition distribution, we call the method NOS-C , and when using Eq. 3 for discrete transitions, we call the method NOS-D .

To sample from the modified transition distribution at each diffusion step, we use Langevin dynamics with temperature τ ą 0 , with the update step,

where η is the step size and λ is the regularization strength, followed by sampling p θ p w t ´ 1 | h 1 t q or p θ p h t ´ 1 | h 1 t q . While the gradient ∇ h v θ guides towards high values of the objective, the KL term ensures the resulting transition distribution still maximizes the likelihood of the original prediction.

NOS is related to the popular method plug-and-play language model (PPLM), which can be used for gradient-guidance of autoregressive language models [17]. PPLM guides sampling by taking gradient steps similar to Eq. 4 for each autoregressive decoding step (details in Appendix B). Unlike PPLM, NOS is a form of iterative refinement, meaning that tokens across the entire sequence can be modified at each optimization step. This distinction is particularly important for protein design, because function can be determined by complex interactions between distant parts of the sequence. As we see in Sec. 5, NOS leads to better trade-offs between likelihood and objective value.

4.2 LaMBO-2: function-guided protein design

Many unique challenges arise when applying guided diffusion to real-world protein design tasks. Our approach builds on the LaMBO-1 algorithm proposed by Stanton et al. [73], which explicitly

EVQLVES -GGGLVQPGGSLRLSCAASG- FN I KD- - - - - TY I HWVRQAPGKGLEWVAR I YPT - - - NGYTRYADSVKGRFT I SADTSKNTAYLQMNSLRAEDTAVYYCSRWGGDG- - - - - - - - - - - - - - - - - - - FYAMDYWGQGTLVTVSS

0

2

4

6

accounts for the online, multi-objective nature of protein design by optimizing a multi-objective Bayesian acquisition function. LaMBO-2 replaces the guided MLM sampler with NOS, selects edit positions based on acquisition value saliency, and replaces the discriminative deep kernel Gaussian process (GP) with ensemble-based uncertainty quantification.

Architecture and value function In order to apply the methods discussed in Subsec. 4.1 we require a generative diffusion model p θ p w q and a discriminator ˆ f θ p w q which share hidden layers up to depth d . The discriminator is trained to predict the objective function f . Like LaMBO-1 our architecture consists of many task-specific feature extraction layers that share a bidirectional encoder. Bayesian optimization is an iterative cycle of inference and data acquisition. During the data acquisition phase of any iteration i we need to find sequences with maximal acquisition value v i p w q ' E r u p w,f, D i qs , where D i is the data already available and the expectation is taken with respect to a posterior p θ p f | D i q and u is some utility function. For multi-objective tasks u is the hypervolume improvement utility function [18], however we note that single-objective tasks are easily accommodated by a different choice of utility function [82]. To estimate the expectation we draw samples from p θ p f | D q with an approach we call partial deep ensembles , where the discriminative layers of the model above the shared encoder are replicated k times and randomly initialized [81]. We provide further details about partial deep ensembles and our learned discriminators in Appendix D.2 and D.3.

Choosing edit positions When B ! L encoder-only architectures allow very precise control of edit positions since we will only change positions we corrupt. However, this feature introduces the need for some procedure to choose those positions, ideally where edits will most improve our objective value. We automatically select edit positions by computing the gradient of the value function with respect to h 0 to determine which positions affect the value estimate the most (see Figure 3 for an illustration). This method is related to the use of saliency maps to explain the decisions of classifiers [5, 69]. We use input saliency to induce a distribution over edit positions. Specifically, given an embedded sequence h 0 we define s i p h 0 q , the saliency with respect to v of position i P t 1 , . . . , L u as

where τ ą 0 is a temperature hyperparameter and 0 ă ε ! 1 . As τ Ñ `8 , P r edit w p i q 0 s ' 1 { L for all i . For each sequence we draw B edit positions without replacement according to Eq. 5. We conserve parts of the input we cannot change (e.g. the antigen sequence) by setting the the saliency to 0 before computing the edit position distribution. Importantly, the diffusion sampling process can also preserve the original values of selected positions when appropriate. If we select a highly conserved position, then the predicted logits will be low entropy and the guidance will incur a large KL penalty for changes (Eq. 4).

5 Experiments

We evaluate our methods on three increasingly complex antibody design tasks. First we compare our trained diffusion models on unguided infilling tasks, showing that sequence diffusion methods

RFDiffusion

DiffAb

IgLM

NOS-C (Ours)

NOS-D (Ours)

HCDR1

HCDR2

HCDR1+HCDR2+HCDR3

LCDR3

Infill Region

0.0

0.5

1.0

Seq. Recovery

consistently outperform structure-based methods when only predicted structures are available 4 . We then evaluate NOS by optimizing two objectives that can be simulated effectively in silico . Lastly, we evaluate LaMBO-2 on antibody lead optimization, with both in silico and in vitro experiments.

5.1 Unguided antibody CDR infilling

We focus on immunoglobulin G (IgG) format antibodies, which are comprised of a heavy (H) chain and a light (L) chain. Each chain has three complimentarity determining regions (CDRs), which tend to have strong effects on binding affinity to a target antigen but limited effects on other structural properties of the protein. Antibody design methods traditionally focus on proposing mutations to CDRs while leaving the rest of the protein fixed, which can be viewed as an infilling task. We select 10 seeds at random from paired OAS [56] and infill each CDR individually as well as in combination. To evaluate the performance of each model, we measure the sequence recovery rate, which is simply the accuracy of the infilling predictions given the ground truth sequence. As baselines, we include IgLM [68], a GPT2 language model trained on OAS, and two structure-based methods: DiffAb [51], a joint sequence-structure diffusion model trained on SAbDab, and RFDiffusion [80], a structural diffusion model trained on the PDB [10] that uses inverse folding to derive sequences. Although IgLM is trained with fill-in-the-middle augmentations [7], it does not natively support infilling multiple non-contiguous regions, and we do so by replacing regions that are not yet sampled with [UNK] tokens. For the structure-based methods, we provide starting structures generated with IgFold [64].

In Figure 4, we find that diffusion models often generate infills that are on-par or better than that those returned by IgLM by default, especially when multiple regions must be filled simultaneously. We also see that DiffAb, while being capable of sequence-structure co-design out of the box, often underperforms sequence-only diffusion, most likely because our sequence-based approaches have access to a larger training dataset, while paired datasets with sequences and structures are much more limited. Lastly RFDiffusion tends to generate relatively low likelihood CDR infills. The gap between DiffAb and RFDiffusion may be explained by the relative scarcity of antibody structures in the PDB compared to SAbDab, which has an antibody in every structure. The poor performance of structural methods on CDR infilling could also be a result of poor sequence recovery from structure during inverse folding, a problem that could be amplified for relatively unstructured loop regions like CDRs.

5.2 Optimizing antibodies for in silico objectives

To test guided sampling from our model, we run experiments on two simple single-objective tasks:

The percentage of beta sheets, measured on primary sequence [15]

The solvent accessible surface area (SASA) of the protein's predicted structure [64]

Seeds

NOS-D (Ours)

NOS-C (Ours)

RFDiffusion

DiffAb

PPLM

DiGress

-500

-400

0.2

0.3

% Beta Sheets

-500

-400

1.1

1.2

SASA

1e4

-500

-400

0.2

0.3

% Beta Sheets

-500

-400

1.1

1.2

SASA

1e4

Since we want plausibly natural antibodies with high objective value we examine the Pareto front for samples optimized for each objective, with log-likelihood assigned by ProtGPT [29] (trained on Uniref50 [75]) plotted against the value of the objective. As an autoregressive guided baseline, we run PPLM, using IgLM as the base generative model (details in Appendix C.3). We use DiGress [79] as a guided diffusion baseline. DiGress uses gradients on one-hot representations to performing guidance in the embedding layer and is thus closely related to our approach. We discuss differences between the methods and the details of our DiGress experiments in Appendix C.5. For PPLM, DiGress, and NOS, we generate samples for many different setting of the control hyperparameters (Section 4.1), which yields samples across the spectrum from aggressively optimizing the objective to conservatively maintaining high likelihood. We also include DiffAb and RFDiffusion without guidance as baselines, as examples of popular 'diversification' procedures, in which new samples are generated for later ranking and filtering. In Figure 5, we see that for both continuous and discrete corruptions NOS offers better trade-offs between optimizing the objective and maintaining high likelihood, while also generating high values of the objective at the extreme.

5.3 Antibody lead optimization: in silico evaluation

Having established the performance of NOS on simpler benchmarks, we now turn to real-world antibody design with LaMBO-2. From this point forward in all experiments we jointly condition on the heavy chain, light chain, and antigen sequence, and we jointly optimize the heavy and light chains only for improved expression yield and binding affinity to the antigen. As we discussed in Subsec. 4.2, one of the critical subproblems in Bayesian optimization is the identification of high value additions to the existing dataset. In this section we show that LaMBO-2 effectively applies NOS to this subproblem in the antibody design setting by finding high acquisition value sequences while preserving naturalness (which we quantify with the metric proposed by Shanehsazzadeh et al. [67]). We focus on optimizing hu4D5, a therapeutic antibody targeting the HER2 antigen 5

Comparison with genetic algorithms We first compare LaMBO-2 to two discrete optimization baselines, AdaLead and PEX. We generated 32 designs from the hu4D5 seed with each method, optimizing the same acquisition function derived from the same model. To ensure a fair comparison we limited all methods to a total of 512 model evaluations and a maximum of 2 edits per sampling iteration. We evaluated both the sample acquisition value and naturalness after each iteration. We identified an empirical naturalness threshold below which expression became unreliable and treated this threshold as a simple inequality constraint. Note that this experiment evaluates how each method balances the tradeoff between acquisition value and naturalness as sampling progresses, and does not involve evaluations of the actual black-box objective.

10

20

30

# Sampling Steps

0.1

0.2

0.3

0.4

0.5

Naturalness

10

20

30

# Sampling Steps

-4

-2

0

2

2

4

6

8

Diffusion Step

200

400

600

Acq. Value

B

=8

0

10

20

30

Diffusion Step

0

1000

2000

Acq. Value

B

=32

Uniform Edits + No Guidance

Uniform Edits + Guidance

Salient Edits + No Guidance

Salient Edits + Guidance

AdaLead

LaMBO-2

PEX

WJS

LaMBO-2 strictly dominates PEX in terms of naturalness and acquisition value at every sampling iteration, with PEX producing infeasible samples beyond 4 iterations. AdaLead improved sample value the most rapidly of all methods in this experiment, but also degrades naturalness the fastest, violating the constraint after only 2 sampling iterations. In contrast LaMBO-2 samples satisfy the naturalness constraint out to 16 sampling steps, producing the highest value feasible solutions. This result highlights the importance of accounting for distributional constraints when optimizing empirical proxies of fitness, since the quality of the proxy signal degrades rapidly outside the support of the training data. Genetic algorithms easily hack empirical models by leaving the support of natural sequences, where training data is necessarily absent, leading to poor quality solutions that nevertheless attain high acquisition value. In Appendix D.5 we show that both sequence and structure-based unguided infilling (i.e. random hit diversification) has the opposite behavior, producing samples with reasonable naturalness but low acquisition value.

Effect of salient edits To separate and independently study the effects of guidance (NOS) and salient position selection, we present an ablation in Figure 6 for optimization with a relatively small edit budget B ( B ă' 10% of mutable positions). To isolate the effects of salient edits we baseline against edit positions chosen uniformly at random, and to isolate the effects of guidance we set the step size η (Eq. 4) to 0. Small edit distance constraints are common in antibody engineering because the goal is typically to increase binding affinity without altering the binding location on the antigen (i.e. the engineered antibody should bind to the same epitope) [43]. One heuristic way to constrain the design to the same epitope is to set B « 8 , (about 2.7% of the antibody sequence length) [43], precisely the range we consider in Figure 6.

In the few edit regime we find that while both interventions improve sample objective value, selecting positions using saliency has a much larger effect than guidance. Although gradient guidance is a reliable and generally applicable tool for improved sampling, the scale of the edit position search dominates the scale of the search over token replacements that guidance affects. With a vocabulary of 21 tokens the number of possible token combinations ( 21 8 ) is dwarfed by the combinations of possible edit positions ( C 300 8 ). Salient selection of edit positions is, therefore, key to any practical application of NOS in budget-constrained design. Interestingly, this facet of protein design differs significantly from guided sampling of images, where generation is typically limited to fixed locations [50, 14], not a fixed edit budget spread over any location that will optimize the objective. These additional degrees of freedom pose an extra challenge.

5.4 Antibody lead optimization: in vitro evaluation

As our final evaluation in Figure 7 we present results using LaMBO-2 (specifically with NOS-D) to optimize 20 seed antibodies distributed across 4 different therapeutic target antigens, including

LaMBO-2 (Ours)

Round 1

Round 2

Round 3

0.00

0.25

0.50

0.75

1.00

Expression Rate

Round 1

Round 2

Round 3

0.0

0.2

0.4

0.6

Yield Improved

Round 1

Round 2

Round 3

0.0

0.2

0.4

0.6

Binding Rate

Round 1

Round 2

Round 3

0.0

0.1

0.2

0.3

0.4

Binding Improved

hu4D5/HER2 6 . We tested 374 designs in total over three rounds, retraining the model after each round and varying a range of design choices and hyperparameters. While expression and binding performance varied from round to round across seeds and targets, by the final round we were able to generate multiple submicromolar binders for all 4 targets with a median of 5 edits to the seed. See Appendix D.7 for individual yield and affinity measurements and experimental details.

The improvements to yield and affinity over time can be attributed both to the data added to the training corpus and methodological insights gleaned after each round. For example, the sharp drop in expression in Round 2 can mainly be attributed to framework residue deletions that arose when λ (the KL penalty coefficient) was set too small. In the following round we tried a range of larger λ values and fixed the sequence lengths and expression immediately recovered.

Figure 7 also reports binding affinity results of a related experiment from Shanehsazzadeh et al. [67] for context, though we emphasize that there are substantial differences between our wetlab validation and that of Shanehsazzadeh et al. [67] which prevent a true apples-to-apples comparison. In the latter experiment 1M designs were generated for the HER2 target and screened with a high-throughput assay. After screening 421 designs were validated with a high-fidelity surface plasmon resonance (SPR) assay. In addition to wetlab screening, their experiment also restricted edits to specific antibody CDRs. We optimized antibodies for a range of targets including HER2 and relied exclusively on in silico screening before validating with SPR, while placing no explicit restrictions on the edit locations. Despite these differences, our results provide initial evidence that it is possible to generate enriched libraries of antibody designs exclusively with in silico methods operating only on primary sequence. While the experimental validation provided is preliminary, we are actively pursuing more rigorous experimental testing in the form of up-scaled and repeated expression and binding experiments and specificity assessment.

6 Discussion

There are many exciting directions for future work. The original LaMBO algorithm was used to optimize small molecules in addition to proteins, and applying LaMBO-2 to small molecule design is a fruitful direction, as LaMBO-2's improvements are not protein-specific. While sequence alignments are a convenient solution to the length change problem in protein design, padding methods [47] or diffusion with a variable-length corruption process (e.g. [60]) will be needed for applications like small molecules which do not admit alignments. We are also eager to consider optimizing much longer sequences, such as gene perturbations [42], which can exceed 20K tokens in length and may necessitate the use of recent advancements such as implicit convolutions [35, 57, 59] or clever modifications of self-attention [16, 13, 48]. More general notions of guidance such as classifier-free guidance [38] for text or class-conditional generation [62, 12] are another intriguing direction, since some goals are difficult to express as black-box functions or constraints [49, 58].

References

Rebecca F Alford, Andrew Leaver-Fay, Jeliazko R Jeliazkov, Matthew J O'Meara, Frank P DiMaio, Hahnbeom Park, Maxim V Shapovalov, P Douglas Renfrew, Vikram K Mulligan, Kalli Kappel, et al. The rosetta all-atom energy function for macromolecular modeling and design. Journal of chemical theory and computation , 13(6):3031-3048, 2017.

Namrata Anand and Tudor Achim. Protein structure and sequence generation with equivariant denoising diffusion probabilistic models. arXiv preprint arXiv:2205.15019 , 2022.

Charles Audet. A survey on direct search methods for blackbox optimization and their applications . Springer, 2014.

Jacob Austin, Daniel D Johnson, Jonathan Ho, Daniel Tarlow, and Rianne van den Berg. Structured denoising diffusion models in discrete state-spaces. Advances in Neural Information Processing Systems , 34:17981-17993, 2021.

David Baehrens, Timon Schroeter, Stefan Harmeling, Motoaki Kawanabe, Katja Hansen, and Klaus-Robert Müller. How to explain individual classification decisions. The Journal of Machine Learning Research , 11:1803-1831, 2010.

Jasmijn Bastings, Sebastian Ebert, Polina Zablotskaia, Anders Sandholm, and Katja Filippova. " will you find these shortcuts?" a protocol for evaluating the faithfulness of input salience methods for text classification. arXiv preprint arXiv:2111.07367 , 2021.

Mohammad Bavarian, Heewoo Jun, Nikolas Tezak, John Schulman, Christine McLeavey, Jerry Tworek, and Mark Chen. Efficient training of language models to fill in the middle. arXiv preprint arXiv:2207.14255 , 2022.

Prajjwal Bhargava, Aleksandr Drozd, and Anna Rogers. Generalization in nli: Ways (not) to go beyond simple heuristics, 2021.

Carl Ivar Branden and John Tooze. Introduction to protein structure . Garland Science, 2012.

Stephen K Burley, Helen M Berman, Gerard J Kleywegt, John L Markley, Haruki Nakamura, and Sameer Velankar. Protein data bank (pdb): the single global macromolecular structure archive. Protein crystallography: methods and protocols , pages 627-641, 2017.

Stephen Casper, Yuxiao Li, Jiawei Li, Tong Bu, Kevin Zhang, and Dylan Hadfield-Menell. Benchmarking interpretability tools for deep neural networks. arXiv preprint arXiv:2302.10894 , 2023.

Huiwen Chang, Han Zhang, Jarred Barber, AJ Maschinot, Jose Lezama, Lu Jiang, Ming-Hsuan Yang, Kevin Murphy, William T Freeman, Michael Rubinstein, et al. Muse: Text-to-image generation via masked generative transformers. arXiv preprint arXiv:2301.00704 , 2023.

Rewon Child, Scott Gray, Alec Radford, and Ilya Sutskever. Generating long sequences with sparse transformers. arXiv preprint arXiv:1904.10509 , 2019.

Hyungjin Chung, Byeongsu Sim, Dohoon Ryu, and Jong Chul Ye. Improving diffusion models for inverse problems using manifold constraints. arXiv preprint arXiv:2206.00941 , 2022.

Peter JA Cock, Tiago Antao, Jeffrey T Chang, Brad A Chapman, Cymon J Cox, Andrew Dalke, Iddo Friedberg, Thomas Hamelryck, Frank Kauff, Bartek Wilczynski, et al. Biopython: freely available python tools for computational molecular biology and bioinformatics. Bioinformatics , 25(11):1422-1423, 2009.

Tri Dao, Dan Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. Flashattention: Fast and memory-efficient exact attention with io-awareness. Advances in Neural Information Processing Systems , 35:16344-16359, 2022.

Sumanth Dathathri, Andrea Madotto, Janice Lan, Jane Hung, Eric Frank, Piero Molino, Jason Yosinski, and Rosanne Liu. Plug and play language models: A simple approach to controlled text generation. arXiv preprint arXiv:1912.02164 , 2019.

Samuel Daulton, Maximilian Balandat, and Eytan Bakshy. Differentiable Expected Hypervolume Improvement for Parallel Multi-Objective Bayesian Optimization. Advances in Neural Information Processing Systems , 33, 2020.

Justas Dauparas, Ivan Anishchenko, Nathaniel Bennett, Hua Bai, Robert J Ragotte, Lukas F Milles, Basile IM Wicky, Alexis Courbet, Rob J de Haas, Neville Bethel, et al. Robust deep learning-based protein sequence design using proteinmpnn. Science , 378(6615):49-56, 2022.

Prafulla Dhariwal and Alexander Nichol. Diffusion models beat gans on image synthesis. Advances in Neural Information Processing Systems , 34:8780-8794, 2021.

Sander Dieleman, Laurent Sartran, Arman Roshannai, Nikolay Savinov, Yaroslav Ganin, Pierre H Richemond, Arnaud Doucet, Robin Strudel, Chris Dyer, Conor Durkan, et al. Continuous diffusion for categorical data. arXiv preprint arXiv:2211.15089 , 2022.

Benjamin Doerr and Frank Neumann. Theory of evolutionary computation: Recent developments in discrete optimization. 2019.

John Duchi, Elad Hazan, and Yoram Singer. Adaptive subgradient methods for online learning and stochastic optimization. Journal of machine learning research , 12(7), 2011.

James Dunbar and Charlotte M Deane. Anarci: antigen receptor numbering and receptor classification. Bioinformatics , 32(2):298-300, 2016.

James Dunbar, Konrad Krawczyk, Jinwoo Leem, Terry Baker, Angelika Fuchs, Guy Georges, Jiye Shi, and Charlotte M Deane. Sabdab: the structural antibody database. Nucleic acids research , 42(D1):D1140-D1146, 2014.

Patrick Emami, Aidan Perreault, Jeffrey Law, David Biagioni, and Peter St John. Plug & play directed evolution of proteins with gradient-based discrete mcmc. Machine Learning: Science and Technology , 4(2):025014, 2023.

Michael Emmerich. Single-and multi-objective evolutionary design optimization assisted by gaussian random field metamodels. dissertation, Universität Dortmund , 2005.

Michael TM Emmerich, André H Deutz, and Jan Willem Klinkenberg. Hypervolume-based expected improvement: Monotonicity properties and exact computation. In 2011 IEEE Congress of Evolutionary Computation (CEC) , pages 2147-2154. IEEE, 2011.

Noelia Ferruz, Steffen Schmidt, and Birte Höcker. Protgpt2 is a deep unsupervised language model for protein design. Nature communications , 13(1):1-10, 2022.

Nathan C. Frey, Dan Berenberg, Joseph Kleinhenz, Isidro Hotzel, Julien Lafrance-Vanasse, Ryan Lewis Kelly, Yan Wu, Arvind Rajpal, Stephen Ra, Richard Bonneau, Kyunghyun Cho, Andreas Loukas, Vladimir Gligorijevic, and Saeed Saremi. Learning protein family manifolds with smoothed energy-based models. In ICLR 2023 Physics4ML Workshop , 2023. URL https://openreview.net/forum?id=IilnB8jfoP9 . Spotlight presentation.

Marjan Ghazvininejad, Omer Levy, Yinhan Liu, and Luke Zettlemoyer. Mask-predict: Parallel decoding of conditional masked language models. arXiv preprint arXiv:1904.09324 , 2019.

Vladimir Gligorijevic, Daniel Berenberg, Stephen Ra, Andrew Watkins, Simon Kelow, Kyunghyun Cho, and Richard Bonneau. Function-guided protein design by deep manifold sampling. bioRxiv , 2021.

Rafael Gómez-Bombarelli, Jennifer N Wei, David Duvenaud, José Miguel Hernández-Lobato, Benjamín Sánchez-Lengeling, Dennis Sheberla, Jorge Aguilera-Iparraguirre, Timothy D Hirzel, Ryan P Adams, and Alán Aspuru-Guzik. Automatic chemical design using a data-driven continuous representation of molecules. ACS central science , 4(2):268-276, 2018.

Alexandros Graikos, Nikolay Malkin, Nebojsa Jojic, and Dimitris Samaras. Diffusion models as plug-and-play priors. arXiv preprint arXiv:2206.09012 , 2022.

Albert Gu, Karan Goel, Ankit Gupta, and Christopher Ré. On the parameterization and initialization of diagonal state space models. Advances in Neural Information Processing Systems , 35:35971-35983, 2022.

Carolina Gutierrez and Rachel Schiff. Her2: biology, detection, and clinical implications. Archives of pathology & laboratory medicine , 135(1):55-62, 2011.

Brian Hie, Salvatore Candido, Zeming Lin, Ori Kabeli, Roshan Rao, Nikita Smetanin, Tom Sercu, and Alexander Rives. A high-level programming language for generative protein design. bioRxiv , pages 2022-12, 2022.

Jonathan Ho and Tim Salimans. Classifier-free diffusion guidance. arXiv preprint arXiv:2207.12598 , 2022.

Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in Neural Information Processing Systems , 33:6840-6851, 2020.

Emiel Hoogeboom, Didrik Nielsen, Priyank Jaini, Patrick Forré, and Max Welling. Argmax flows and multinomial diffusion: Learning categorical distributions. Advances in Neural Information Processing Systems , 34:12454-12465, 2021.

Sara Hooker, Dumitru Erhan, Pieter-Jan Kindermans, and Been Kim. Evaluating feature importance estimates. web preprint https://research.google/pubs/pub47088/ , 2018.

Yuge Ji, Mohammad Lotfollahi, F Alexander Wolf, and Fabian J Theis. Machine learning for perturbational single-cell omics. Cell Systems , 12(6):522-537, 2021.

Simon Kelow, Bulat Faezov, Qifang Xu, Mitchell I Parker, Jared Adolf-Bryfogle, and Roland L Dunbrack Jr. A penultimate classification of canonical antibody cdr conformations. bioRxiv , pages 2022-10, 2022.

H Benjamin Larman, George Jing Xu, Natalya N Pavlova, and Stephen J Elledge. Construction of a rationally designed antibody platform for sequencing-assisted selection. Proceedings of the National Academy of Sciences , 109(45):18523-18528, 2012.

Jin Sub Lee and Philip M Kim. Proteinsgm: Score-based generative modeling for de novo protein design. bioRxiv , 2022.

Sergey Levine. Reinforcement learning and control as probabilistic inference: Tutorial and review. arXiv preprint arXiv:1805.00909 , 2018.

Xiang Lisa Li, John Thickstun, Ishaan Gulrajani, Percy Liang, and Tatsunori B Hashimoto. Diffusion-lm improves controllable text generation. arXiv preprint arXiv:2205.14217 , 2022.

Hao Liu, Matei Zaharia, and Pieter Abbeel. Ring attention with blockwise transformers for near-infinite context. arXiv preprint arXiv:2310.01889 , 2023.

Shengchao Liu, Yutao Zhu, Jiarui Lu, Zhao Xu, Weili Nie, Anthony Gitter, Chaowei Xiao, Jian Tang, Hongyu Guo, and Anima Anandkumar. A text-guided protein design framework. arXiv preprint arXiv:2302.04611 , 2023.

Andreas Lugmayr, Martin Danelljan, Andres Romero, Fisher Yu, Radu Timofte, and Luc Van Gool. Repaint: Inpainting using denoising diffusion probabilistic models. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition , pages 11461-11471, 2022.

Shitong Luo, Yufeng Su, Xingang Peng, Sheng Wang, Jian Peng, and Jianzhu Ma. Antigenspecific antibody design and optimization with diffusion-based generative models. bioRxiv , 2022.

Pascale Mathonet and Christopher G Ullman. The application of next generation sequencing to the understanding of antibody repertoires. Frontiers in immunology , 4:265, 2013.

Igor Melnyk, Payel Das, Vijil Chenthamarakshan, and Aurelie Lozano. Benchmarking deep generative models for diverse antibody sequence design. arXiv preprint arXiv:2111.06801 , 2021.

Anh Nguyen, Jeff Clune, Yoshua Bengio, Alexey Dosovitskiy, and Jason Yosinski. Plug & play generative networks: Conditional iterative generation of images in latent space. In Proceedings of the IEEE conference on computer vision and pattern recognition , pages 4467-4477, 2017.

Alexander Quinn Nichol and Prafulla Dhariwal. Improved denoising diffusion probabilistic models. In International Conference on Machine Learning , pages 8162-8171. PMLR, 2021.

Tobias H Olsen, Fergus Boyles, and Charlotte M Deane. Observed antibody space: A diverse database of cleaned, annotated, and translated unpaired and paired antibody sequences. Protein Science , 31(1):141-146, 2022.

Antonio Orvieto, Samuel L Smith, Albert Gu, Anushan Fernando, Caglar Gulcehre, Razvan Pascanu, and Soham De. Resurrecting recurrent neural networks for long sequences. arXiv preprint arXiv:2303.06349 , 2023.

Vishakh Padmakumar, Richard Yuanzhe Pang, He He, and Ankur P Parikh. Extrapolative controlled sequence generation via iterative refinement. arXiv preprint arXiv:2303.04562 , 2023.

Michael Poli, Stefano Massaroli, Eric Nguyen, Daniel Y Fu, Tri Dao, Stephen Baccus, Yoshua Bengio, Stefano Ermon, and Christopher Ré. Hyena hierarchy: Towards larger convolutional language models. arXiv preprint arXiv:2302.10866 , 2023.

Machel Reid, Vincent J Hellendoorn, and Graham Neubig. Diffuser: Discrete diffusion via edit-based reconstruction. arXiv preprint arXiv:2210.16886 , 2022.

Zhizhou Ren, Jiahan Li, Fan Ding, Yuan Zhou, Jianzhu Ma, and Jian Peng. Proximal exploration for model-guided protein sequence design. In International Conference on Machine Learning , pages 18520-18536. PMLR, 2022.

Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Björn Ommer. Highresolution image synthesis with latent diffusion models. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition , pages 10684-10695, 2022.

Philip A Romero and Frances H Arnold. Exploring protein fitness landscapes by directed evolution. Nature reviews Molecular cell biology , 10(12):866-876, 2009.

Jeffrey A Ruffolo and Jeffrey J Gray. Fast, accurate antibody structure prediction from deep learning on massive set of natural antibodies. Biophysical Journal , 121(3):155a-156a, 2022.

Nikolay Savinov, Junyoung Chung, Mikolaj Binkowski, Erich Elsen, and Aaron van den Oord. Step-unrolled denoising autoencoders for text generation, 2021.

Harshay Shah, Prateek Jain, and Praneeth Netrapalli. Do input gradients highlight discriminative features? Advances in Neural Information Processing Systems , 34:2046-2059, 2021.

Amir Shanehsazzadeh, Sharrol Bachas, Matt McPartlon, George Kasun, John M Sutton, Andrea K Steiger, Richard Shuai, Christa Kohnert, Goran Rakocevic, Jahir M Gutierrez, et al. Unlocking de novo antibody design with generative artificial intelligence. bioRxiv , pages 2023-01, 2023.

Richard W Shuai, Jeffrey A Ruffolo, and Jeffrey J Gray. Generative language modeling for antibody design. bioRxiv , 2021.

Karen Simonyan, Andrea Vedaldi, and Andrew Zisserman. Deep inside convolutional networks: Visualising image classification models and saliency maps. arXiv preprint arXiv:1312.6034 , 2013.

Sam Sinai, Richard Wang, Alexander Whatley, Stewart Slocum, Elina Locane, and Eric D Kelsic. Adalead: A simple and robust adaptive greedy search algorithm for sequence design. arXiv preprint arXiv:2010.02141 , 2020.

Raghav Singhal, Mark Goldstein, and Rajesh Ranganath. Where to diffuse, how to diffuse, and how to get back: Automated learning for multivariate diffusions. arXiv preprint arXiv:2302.07261 , 2023.

Suraj Srinivas, Kyle Matoba, Himabindu Lakkaraju, and François Fleuret. Efficient training of low-curvature neural networks. Advances in Neural Information Processing Systems , 35: 25951-25964, 2022.

Samuel Stanton, Wesley Maddox, Nate Gruver, Phillip Maffettone, Emily Delaney, Peyton Greenside, and Andrew Gordon Wilson. Accelerating bayesian optimization for biological sequence design with denoising autoencoders. arXiv preprint arXiv:2203.12742 , 2022.

Robin Strudel, Corentin Tallec, Florent Altché, Yilun Du, Yaroslav Ganin, Arthur Mensch, Will Grathwohl, Nikolay Savinov, Sander Dieleman, Laurent Sifre, et al. Self-conditioned embedding diffusion for text generation. arXiv preprint arXiv:2211.04236 , 2022.

Baris E Suzek, Hongzhan Huang, Peter McGarvey, Raja Mazumder, and Cathy H Wu. Uniref: comprehensive and non-redundant uniprot reference clusters. Bioinformatics , 23(10):12821288, 2007.

Christian Szegedy, Vincent Vanhoucke, Sergey Ioffe, Jon Shlens, and Zbigniew Wojna. Rethinking the inception architecture for computer vision. In Proceedings of the IEEE conference on computer vision and pattern recognition , pages 2818-2826, 2016.

Brian L Trippe, Jason Yim, Doug Tischer, Tamara Broderick, David Baker, Regina Barzilay, and Tommi Jaakkola. Diffusion probabilistic modeling of protein backbones in 3d for the motif-scaffolding problem. arXiv preprint arXiv:2206.04119 , 2022.

Robert Verkuil, Ori Kabeli, Yilun Du, Basile IM Wicky, Lukas F Milles, Justas Dauparas, David Baker, Sergey Ovchinnikov, Tom Sercu, and Alexander Rives. Language models generalize beyond natural proteins. bioRxiv , 2022.

Clement Vignac, Igor Krawczuk, Antoine Siraudin, Bohan Wang, Volkan Cevher, and Pascal Frossard. Digress: Discrete denoising diffusion for graph generation. arXiv preprint arXiv:2209.14734 , 2022.

Joseph L Watson, David Juergens, Nathaniel R Bennett, Brian L Trippe, Jason Yim, Helen E Eisenach, Woody Ahern, Andrew J Borst, Robert J Ragotte, Lukas F Milles, et al. Broadly applicable and accurate protein design by integrating structure prediction networks and diffusion generative models. bioRxiv , pages 2022-12, 2022.

Andrew G Wilson and Pavel Izmailov. Bayesian deep learning and a probabilistic perspective of generalization. Advances in neural information processing systems , 33:4697-4708, 2020.

James T Wilson, Riccardo Moriconi, Frank Hutter, and Marc Peter Deisenroth. The reparameterization trick for acquisition functions. arXiv preprint arXiv:1712.00424 , 2017.

Kevin Yang and Dan Klein. Fudge: Controlled text generation with future discriminators. arXiv preprint arXiv:2104.05218 , 2021.

Appendix

Table of Contents

A Extended Background

In this section we provide full descriptions of the diffusion processes introduced in Sec. 3.

A.1 Continuous noise diffusion

The forward process is defined by noise variances β . We use the cosine variance schedule from Nichol and Dhariwal [55]. For convenience we further define

The forward process is defined by the conditional distributions

This is a sample string

This is a [MASK] string

This [MASK] a [MASK] [MASK] [MASK] [MASK] [MASK] [MASK] [MASK] [MASK] [MASK]

where U θ is an embedding matrix. The reverse process is defined by

A.2 Categorical noise diffusion

Following Austin et al. [4] we define the MLM style categorical diffusion using transition matrices

and ¯ Q t ' Q 1 Q 2 ...Q t for noise schedule α t P r 0 , 1 s (see Figure 8 for an illustration). These transition matrices correspond to categorical conditional distributions

The reverse process is defined by

B Methodological Details

B.1 Infilling algorithm

We sample infills using the procedure in Algorithm 1. The infill mask P is constructed by setting the index of conserved residue equal to 1, in this case at every residue that is not included in set of CDR regions being infilled. We use the same algorithm to perform the guided infilling in Subsec. 5.2, where it is extended with a guidance Langevin sampling step.

Algorithm 1 Infilling with categorical denoising diffusion model

Inputs: Denoiser p θ p ˆ w | x t , t q , corruption process p p x t | x 0 q , infilling mask P , and seed sequence s Returns: Sample from ˜ p p w q ' p θ p w | P, s q exp p f p w qq

B.2 Hidden State Langevin Sampling

Design of molecules or images with generative models is often posed as the problem of sampling from a posterior distribution p p x | a q given the unconditional distribution p p x q and attribute model p p a | x q . Indeed, reinforcement learning, the design of good actions in an environment, can also be framed as posterior sampling where p p a | x q is the probability that a given state or state-action pair is optimal [46]. Methods that employ posterior sampling of this form are often call 'plug-and-play' because p p a | x q and p p x q need not share parameters and therefore users can mix and match different instantiations [54, 17, 34, 26]

The most common way to sample from the posterior p p x | a q 9 p p a | x q p p x q is through Langevin sampling on the unnormalized joint density ˜ p p a, x q ' p p a | x q p p x q , with sampling steps

When we work with generative models over continuous random variables that permit a likelihood (e.g. normalizing flows), score function (e.g. diffusions), or energy (e.g. EBMs) ∇ log p p x q has a natural interpretation and sampling can be performed with essentially vanilla Langevin sampling. In other cases where only a denoising function over continuous variables is available, authors have proposed approximate samplers using an approximation of the score function [54].

When we instead hope to sample from a posterior over discrete random variables constructing an analogy to the score function ∇ log p p x q is challenging, and prior work adopts a different approach of regularizing the conditional sampling distribution p p w | a q with unconditional sampling p p w q in order to maintain high likelihood [17]. In autoregressive models, p p w q is broken down using the chain rule, p p w t | w ă t q and thus the appropriate regularization is

In our case, the distribution p p w q is factorized by the transition distributions p p w t | w t ´ 1 q (or their continuous analogies in token embedding space), and we hope to sample from the perturbed transition

The correct regularization term in our case is thus

To put the pieces together, we first recognize that the denoising model p θ p w 0 | w t q can be broken down into an language model head, H θ , and trunk, T θ , with

We can then perform Langevin sampling on the hidden representations, initializing with h t , as shown in Algorithm 2. In the experiments above we set λ 3 ' 0 , as we saw no noticable benefit from adding additional stochasticity. Importantly, sampling from p p w t ´ 1 | w t q already introduces randomness into the reverse process.

Algorithm 2 Guided diffusion sampler

C Infilling / NOS Guidance

All of our diffusion models are train on all paired heavy and light chain sequences from OAS [56] (pOAS) combined with all sequences from SAbDab [25], aligned with ANARCI [24].

C.1 Infilling experiment

For our trained diffusion models, we use Algorithm 1 without guidance, generating P based on the indicated CDRs, using chothia numbering for consistency with DiffAb. For the baselines, we constructed wrapper scripts to convert the chosen CDR ids into each method's native format.

C.2 MCMCcomparison

Following Verkuil et al. [78], we construct a Markov chain using uniform random mutations to map a sequence w to a mutated sequence w 1 , using the following Metropolis-Hastings correction:

where T ą 0 is a temperature hyperparameter. While this method has appealing theoretical properties, obtaining good samples from this Markov chain in practice requires hundreds of thousands of steps of burn-in.

In our experiment (Figure 9), we define the energy, E , by combining sequence level probabilities assigned by IgLM with a beta sheets objective function trained on IgLM's representations. We construct the energy as

Wetune λ to generate sequences with approximately 40% beta sheets. We also tune the NOS λ parameter (Eq. 4) to produce approximately 40% beta sheets.

C.3 PPLM details

In order to generate full (heavy and light chain) optimized antibodies with PPLM and IgLM, we train two separate value function models on IgLM's aggregated hidden representations, one for heavy chain sequences and one for light chain sequences. IgLM uses special tokens for both the chain identity and the species identity of each sequences, and we pass in appropriate corresponding tokens when calculating the hidden representations for each model. To determine the correct species token for each sequence, we use the predicted species returned by ANARCI [24]. Our value function is a

0

5000

10000

15000

20000

Sampling Steps

3.0

3.2

3.4

3.6

Energy

simple one-layer feed-forward neural network trained on top of the mean-aggregated representations for the corresponding chain identity.

To sample using PPLM, we overwrite the forward pass of the huggingface decoder used by IgLM to include a Langevin sampling step over the current hidden representations. We perform K gradient steps to update the current hidden representation h 1 by descending on the objective

where h is the original hidden representation output by the model's encoder, and η and λ are the step size and regularization strength respectively. We ran optimization with both vanilla gradient descent and AdaGrad [23] and found AdaGrad to be more robust to poor specifications of the step size. For the results in Sec. 5, we draw samples and present results for all of the hyperparameter settings in Table 1

One critical difference between controllable autoregressive models and controllable diffusions is the ability to resample previously sampled values. Procedures that allow for resampling are often called 'iterative refinement' procedures because they can produce increasingly plausible generations by refining the model's previous output at each step in an iterative procedure. Because there are many potential differences between our NOS models and PPLM, including but not limited to the nature of iterative refinement, we performed an additional experiment to assess the impact of adapting a discrete diffusion to perform autoregressive sampling. Autoregressive models can themselves be thought of as diffusions with an idiosyncratic corruption process that masks out all tokens to the right of the last sampled token. As in our discrete corruption process, the prior is also a sequence of all mask tokens. Using this insight, we can run our trained discrete diffusions in autoregressive mode by contriving the sampling noise schedule to be autoregressive and recover an approximation of the timestep post-hoc from the percentage of masks at each step in autoregressive sampling.

Figure 10 shows the difference in objective values and likelihood for samples obtained by running the model in typical diffusion mode (iterative refinement) or in contrived autoregressive mode. We can see that on the beta sheets objective, iterative refinement has a noticeable positive impact on the objective values of the sample. This effect is also present in the SASA objective, but to a much more limited extent. We speculate that the iterative refinement facet of NOS is helpful for outperforming other methods but not completely sufficient.

C.4 Model Architecture and Training

The gaussian and categorical diffusions are trained with the bert-small transformer backbone introduced by Bhargava et al. [8]. We use a cosine noise schedule for both diffusions and train for 100 epochs with a batch size of 64, optimizing with AdamW using an initial learning rate of 5e-3 with a linear warmup. The value function is a feed-forward neural network with one hidden layer. The value function is trained jointly with the denoiser by alternating optimization steps, with 5 steps on the generative objective for each step on the discriminative objective. We train the models for 100 epochs in total.

C.5 Hyperparameter settings

For each guided sampling experiment with NOS, we sample using many different hyperparameter combinations in order to generate both conservative and aggressive optimization of the value function.

-500

-400

0.15

0.20

0.25

% Beta Sheets

-500

-450

-400

-350

11000

12000

SASA

Autoregressive

Iterative Refinement

The full hyperparameter settings for both objectives (beta sheets and SASA) and both corruption types (NOS-D and NOS-C) are shown in Table 2. In Table 2, there is an additional hyperparameter, 'guidance layer', which we did not discuss at length in the main text of the paper. This parameter dictates whether we perform guidance in the first layer of the neural network (the token embeddings), as is standard in continuous diffusion models for discrete sequences, or the final layer of the neural network (the layer before the final linear head). In either case, we can use the same gradient descent objective and corruption process in each case and need only change the variable we propagate gradient updates to. Table 2 shows the hyperparameters used in the just Figure 5.

To aid intuition for the effects of each hyperparameter, we show the sample densities that result from each combination of λ and η in Table 2 when guiding in the first (Figure 11) and last (Figure 12) layer of the NOS-D and NOS-C models. We see that the most important parameter is λ , which controls how far samples tend to move from the seeds. We can also observe that guiding in the first hidden state tends to perform better when sampling with NOS-C, while guiding in the final hidden state tends to perform better with NOS-D.

DiGress comparison DiGress [79] is built on top of a model with one-hot encodings and discrete corruptions. The guided sampling procedure can be described as follows (using the notation from our submission): At each denoising step t , we use the one-hot encodings as a continuous variable and construct a perturbation distribution from a learned discriminative model ˆ v ' v θ p w t q ,

We then sample the next value from the base diffusion transition p θ p w t ´ 1 | w t q perturbed with p θ p ˆ v | w t ´ 1 q ,

The key details for guided sampling can be found in the DiGress code repo, where we see that the guided distribution is the normalized product of the original denoising distribution and the softmax of the gradients scaled with λ .

On a theoretical level, this guidance has noticeably different properties from NOS. For large λ , the perturbation p p w t | ˆ v q collapses to a one-hot on the token index with the largest gradient value. For small values of λ , p p w t | ˆ v q becomes a uniform distribution. Therefore λ interpolates p p w t ´ 1 | w t q between the original unguided distribution and a one-hot in the max gradient direction. NOS also reduces to unguided infilling when λ ' 0 , but λ ą 0 only modulates the direction of the gradient update. The distance between the guided and unguided distribution is controlled by the number of langevin steps and the step size hyperparameter η . Digress amounts to a single update step applied directly to the output token probabilities using a continuous relaxation of the one-hot encoded input, whereas NOS performs a sequence of local updates to hidden states that are actually continuous.

In our comparison, the embeddings and corruptions of each model are chosen to be:

NOS-C [Gaussian corruptions + learned embeddings]

11000

12000

13000

η

=

0.1

λ

=

10.0

λ

=

1.0

λ

=

0.1

λ

=

0.01

λ

=

0.001

11000

12000

13000

η

=

0.5

SASA

-600

-400

11000

12000

13000

η

=

1.0

-600

-400

-600

-400

-600

-400

-600

-400

NOS-D (Ours)

NOS-C (Ours)

NOS-D [Discrete (mask) corruptions + learned embeddings]

DiGress [Discrete (mask) corruptions + fixed one-hot encodings]

All models use the same backbone transformer and regression heads, facilitating an apples-toapples comparison. For DiGress, we perform sampling for large range of scaling values λ P t 1 e 5 , 3 e 4 , 1 e 4 , 3 e 3 , 1 e 3 , 3 e 2 , 1 e 2 , 3 e 1 , 1 e 1 , 1 e 0 , 1 e ´ 1 , 1 e ´ 2 , 1 e ´ 3 u . For each model, λ modulates the degree to which the model prefers greedy sampling from the value function gradient.

C.6 Density plots

Because pareto fronts present only a partial view of sampling outcomes (focusing on the best case outcomes along each axis), we also include sample density plots to confirm that our methods consistently yield samples with better trade-off between likelihood and fitness. Figure 13 shows

10000

11000

12000

η

=

0.1

λ

=

10.0

λ

=

1.0

λ

=

0.1

λ

=

0.01

λ

=

0.001

10000

11000

12000

η

=

0.5

SASA

-600

-400

10000

11000

12000

η

=

1.0

-600

-400

-600

-400

-600

-400

-600

-400

NOS-D (Ours)

NOS-C (Ours)

-600

-400

0.2

0.3

% Beta Sheets

RFDiffusion

-600

-400

DiffAb

-600

-400

IgLM

-600

-400

NOS-D (Ours)

-600

-400

NOS-C (Ours)

-600

-400

11000

12000

SASA

-600

-400

-600

-400

-600

-400

-600

-400

RFDiffusion

DiffAb

PPLM

NOS-D (Ours)

NOS-C (Ours)

Seeds

density plots for NOS and baselines when optimizing each of the two objectives (percentage of beta sheets and SASA). We find that DiffAb and IgLM samples tend to cluster around the starting seeds, while RFDiffusion samples tend to generate more diverse samples under the objective, but often with much lower likelihood than the seed sequences. By contrast, both NOS methods consistently improve values of the objective without sacrificing likelihoods.

D LaMBO-2

D.1 Intro to Multi-Objective Bayesian Optimization

When there are multiple objectives of interest, a single best (i.e. strictly dominant) sequence x ˚ may not exist. Suppose there are k objectives, f : X Ñ R k . The goal of multi-objective optimization

Algorithm 3 LaMBO-2: one guided discrete diffusion step

Inputs: Seed sequence w 0 , edit budget projection P , diffusion timestep t , corruption function c p w,t q , constraint function u p w q , encoder g θ p w q , value function v θ p h q , decoder d θ p h q , regularization strength λ , SGLD step-size η and temperature τ . 1

(MOO) is to identify the set of Pareto-optimal (i.e. non-dominated) solutions such that improving one objective within the set leads to worsening another. We say that x dominates x 1 , or f p x q ą f p x 1 q , if f j p x q ě f j p x 1 q for all j P t 1 , . . . , m u and f j p x q ą f j p x 1 q for some j . The set of non-dominated solutions X ˚ is defined in terms of the Pareto frontier (PF) P ˚ ,

MOO algorithms typically aim to identify a finite approximation to X ‹ (which may be infinitely large), within a reasonable number of iterations. One way to measure the quality of an approximate PF P is to compute the hypervolume HV p P | r ref q of the polytope bounded by P Yt r ref u , where r ref P R m is a user-specified reference point .

where P 1 ' P Y t ˆ f p x qu [27, 28, 18]. To decide where to query f next, we search for argmax x E r u EHVI p x , f, D qs , where the expectation is w.r.t. p p f | D q .

D.2 Discrete EHVI

Although expression yield and binding affinity are both continuous measurements, we chose to discretize them and model them as classification with a softmax likelihood (See Appendix D.4). As a result we needed an extension of EHVI for discrete outcomes. Informally, EHVI is simply computing the HVI for different realizations of f and marginalizing f using p p f | D q . Instead of taking f to be the latent function of some regression y ' f p w q ` ε . ε ' N p 0 , σ 2 q , we instead take f to be the logits of a categorical distribution, p p y ' i | w, D q ' ş softmax i p f p w qq p p f | D q df .

Let y ' r y 1 ¨ ¨ ¨ y k s J . Given a set of baseline points B Ă A L we define P (Eq. 8) using the posterior mean ˆ y p w q ' E r y | w, D s , w P B . We model y 1 , . . . , y k as conditionally independent given some shared hidden state h ' g d p w q , so p p y | h, D q factorizes nicely. Finally we define P 1 ' P Yt y u and take the expectation of Eq. 8 w.r.t. p p y | h, D q . Since p p y | h, D q is discrete and factorizes, we can marginalize in closed form when K 1 ˆ¨¨¨ ˆ K k is not too large, where K i is the number of classes corresponding to the discretization of the original continuous f i .

D.3 Architecture and Hyperparameters

The inputs of the LaMBO-2 model for antibody design are the variable heavy (VH) and variable light (VL) regions of the antibody sequence as determined by Aho alignment with ANARCI, as well as the (unaligned) antigen sequence. Note that the concatenation of the antigen to the input makes the samples from the generative head conditional on the antigen as well as the unmasked portion of the antibody sequence. The LaMBO-2 model jointly predicts antigen-conditional categorical

token distributions for corrupted positions and discriminative distributions over protein properties. Discriminative predictions that should not depend on the antigen are made invariant through data augmentation with random antigen sequences. See Algorithm 3 for an overview of a single guided diffusion step with LaMBO-2.

Model Architecture: our architecture for this experiment is inspired by the one proposed by Stanton et al. [73]. In particular we jointly a train an encoder shared between a generative discrete diffusion head and discriminative heads which predict expression and affinity. Rather than use a deep kernel GP, we simply ensemble 10 heads for each discriminative task to obtain uncertainty estimates. Like Stanton et al. [73] for this experiment we use 1D CNN residual blocks (kernel width 9 ), with layer normalization and sinusoidal position embeddings. The shared encoder was comprised of 4 residual blocks, and each task head was comprised of 2 residual blocks followed by a linear layer, with the exception of the generative head which was just a linear layer on top of the shared embeddings. Note that in future work self-attention layers could be used instead of CNN layers, as was the case for the pOAS experiments in Sec. 5. We set the embedding dimension to 32, and the latent channel dimension to 256.

Training Hyperparameters: The LaMBO-2 model is both a jointly trained generative and discriminative model, as well as a true multi-task model, which is necessary since measurements for various protein properties are often missing from a substantial fraction of rows in real-world datasets. We trained for 500K gradient updates using the Adam optimizer with η ' 1e-3 , β 0 ' 0 . 99 , β 1 ' 0 . 999 . At each gradient step we randomly sampled a task head and task minibatch (batch-size 121) and updated the corresponding weights (including shared weights). We used a linear learning rate warmup over 10K gradient updates, and decayed the learning rate to 1e-6 with a cosine schedule. We did not regularize with weight decay or dropout.

Generation Hyperparameters: to generate the designs in Figure 7, we sampled 1K designs from a pool of seed antibody sequences hand-selected by domain experts. For each seed we set the total edit budget shared between chains to B ' 16 . In this experiment each infilling method took 16 diffusion steps, using an inverse linear noise schedule α t ' 1 {p 1 ` t q . Although the models were trained with a standard cosine noise schedule, we found the inverse linear schedule gave better results in terms of sample acquisition value at generation time. Within each diffusion step we took 64 Langevin steps, with noise scale τ ' 1e-2. For guided infills with uniformly distributed edit positions we set τ ' 1e6. For guided infills with saliency-informed edit position selection we set τ ' 0 . 1 . We set λ ' 0 . 5 to balance the tradeoff of sequence likelihood and value during guidance.

Generation Constraints: in addition to the edit budget locality constraint, our LaMBO-2 designs were also constrained to meet certain sequence liabilities constraints:

Canonical Cysteine Conservation: there are specific conserved cysteine residues in antibody sequences which play a crucial role in the formation of disulfide bridges. Disulfide bridges are covalent bonds formed between two cysteine residues through oxidation of their sulfur atoms. These bridges contribute to the overall structural stability and integrity of antibodies.

No Unpaired Cysteines: odd numbers of cysteines within individual chains (i.e. unpaired cysteines) are generally undesirable since they can lead to non-native disulfide bonds between different antibody molecules, which may disrupt assembly, folding, or function.

No Glycosylation Motifs: A glycosylation motif is a specific amino acid sequence within a protein that serves as a recognition site for the attachment of sugar molecules. The presence of a glycosylation motif in a protein can affect its stability, solubility, activity, and function. The addition of sugar molecules can alter the protein's conformation, change its interactions with other proteins or molecules, and affect its trafficking and localization within the cell.

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

Training Data: the expression task heads were trained on a dataset of 10K linear transfection expression measurements, which was subsequently augmented to 160K rows by pairing the same measurements with different random antigens to teach the model to ignore the antigen sequence when predicting expression. The binding task heads were trained on a dataset of 10K SPR affinity measurements for various antigens, which was then augmented to 12K rows by pairing binders with different random antigens and imputing a non-binding label. This augmentation is important for

0

2

4

6

8

10

measured pKD

0.0

0.2

0.4

0.6

0.8

Density

0

1

2 3 4 5

6

7

8

9

10

quantized pKD

0

1

2

3

4

5

6

7

8

9

10

quantized pKD

0

200

400

600

800

1000

count

training a pan-target affinity model, since experimental measurements of affinity to off-target antigens are uncommon. Note that the expression and affinity data only partially overlapped, necessitating the multi-task architecture described in Appendix D.3. The generative diffusion head was trained only on binding antibody-antigen pairs in the SPR binding data.

We did not pretrain our LaMBO-2 models. It is likely that performance could be improved with the right pretraining corpus, however it is unclear if datasets like pOAS are particularly useful for pretraining antibody design models since most do not report antigen sequences and may not have the right level of variability. In any case, it is very encouraging to see positive real-world results before scaling in earnest.

Label Discretization. As noted above, biological data tends to be very imbalanced, and historical experimental data even more so since there are strong selection effects imposed by the scientists collecting the data. We chose to discretize continuous properties like expression yield and binding affinity, making it easier to correct for class imbalance by upsampling minority classes. In Figure 14 we illustrate our discretization scheme. Any antibody-antigen pair with ´ log p KD q (pKD) less than 4 was assigned to the non-binding class 0. Then binders were assigned to classes 1 - 10 based on which pKD decile (computed from binders only) they resided in. One consequence of this scheme is increasing any objective value by one unit corresponds to moving up one decile in the empirical label distribution.

Training Discriminators on Noisy Inputs: the benefits of discretization are not limited to addressing class imbalance. Working with discretized labels also allowed a simple approach to training the discriminator on corrupted inputs inspired by label smoothing [76]. We train the discriminators with the same noise schedule as the diffusion model and the usual cross-entropy loss, using modified labels

where y is the one-hot encoded label and K is the number of classes. Informally, as α t Ñ 0 the discriminator reverts to a uniform prior since the inputs are not distinguishable. Training on corrupted inputs avoids evaluating the value gradient on out-of-distribution inputs during generation, and causes the strength of the value gradient to grow as the diffusion progresses and the samples become more defined.

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

Structure-Based Diversification We have shown that we can effectively optimize antibodies for predicted yield and affinity, and our method performs well compared to unguided sequence-based infilling methods. We expand our evaluation for this task to include unguided infilling with DiffAb and RFDiffusion of CDRs H2 and H3 of hu4D5 (i.e. the seed), a publicly released therapeutic antibody that is ideally suited for structure-based methods since we have a ground truth crystal structure of hu4D5 docked with its target ERBB2. While it is not feasible to validate the resulting designs in vitro during the author response period, we can compare the AntiBERTy naturalness scores and the acquisition value (log expected hypervolume improvement or log-EHVI) of the designs relative to our guided infills (Fig. 15). To summarize, unguided structure-based infilling produces

high likelihood samples, but even when conditioned on the antigen the distribution shift toward better predicted function is very slight.

hu4D5

DiffAb

LaMBO-2

RFDiffusion

0.40

0.45

0.50

0.55

0.60

Naturalness

0

20

40

Density

-50

-40

-30

-20

-10

0

Acquisition Value

0.0

0.2

0.4

0.6

Density

-50

-40

-30

-20

-10

0

Acquisition Value

0.000

0.025

0.050

0.075

0.100

Density

RFDiff w/o antigen

RFDiff w/ antigen

Sequence Diversification This in silico evaluation compares two variants of LaMBO-2 (one using NOS-C, the other NOS-D) against a competing method, walk-jump sampling (WJS), an unguided smoothed discrete sampling algorithm proposed by Frey et al. [30]. Each method generated 1K designs from the same set of seeds, and all methods were restricted to B ' 8 edits. LaMBO-2 chose all edit positions automatically along the entire antibody sequence, whereas WJS was given manually selected edit positions restricted to CDRs only. In the left two panels of Figure 7 we compare the predicted expression yield, predicted binding affinity, and naturalness of the antibody designs, using the metric proposed by . Comparing the Pareto frontiers obtained from each set of designs, we see that while WJS excels at generating 'natural' antibodies, it struggles to generate designs at the higher end of the objective range. Conversely LaMBO-2 designs (particularly those generated with NOS-C) have high predicted objective value but also lower naturalness scores. LaMBO-2 designs generated with NOS-D strike a balance between the two extremes.

Seeds

NOS-C (Ours)

NOS-D (Ours)

WJS

0.40

0.45

0.50

0.55

Naturalness

0.0

2.5

5.0

7.5

10.0

Pred. Expression

0.45

0.50

0.55

Naturalness

4

6

8

10

Pred. Affinity

D.6 Are Saliency Maps Reliable?

There is substantial controversy regarding the reliability of input-gradient-based feature attribution methods, specifically related to their ability to consistently highlight ground truth task-discriminative features and ignore irrelevant features. For example, Hooker et al. [41] claim that random attribution is competitive with input-gradient methods, and Casper et al. [11] claim that gradient-free attribution outperforms input-gradient competitors. On the other hand, many papers claim that specific types

Input Masking

EVQLVES -GGGLVQPGGSLRLSCAASG- FN I KD- - - - - TY I HWVRQAPGKGLEWVAR I YPT - - - NGYTRYADSVKGRFT I SADTSKNTAYLQMNSLRAEDTAVYYCSRWGGDG- - - - - - - - - - - - - - - - - - - FYAMDYWGQGTLVTVSS

0

2

4

6

DI QMTQSPSSLSASVGDRVT I TCRAS - -QDVN- - - - - - TAVAWYQQKPGKAPKLL I YS - - - - - - - - ASF LYSGVPSRFSGSRSG- - TDFTLT I SSLQPEDFATYYCQQHYT - - - - - - - - - - - - - - - - - - - - - - - TPPTFGQGTKVE I K

0

5

10

of regularization can improve the performance of input-gradient attribution, including adversarial training [66], mask denoising [6], and model curvature penalties [72].

A thorough investigation of these claims is beyond the scope of this work, however we have found that saliency maps produced by independent models trained with different corruption processes seem to consistently highlight specific regions of the antibody sequence (Figure 17). It is also worth noting that most of the related literature evaluates feature attribution in the offline setting. In LaMBO-2 feature attributions are used online to intervene on the data collection process (specifically where to introduce changes in the antibody sequences). If LaMBO-2 changes a position that does not affect function it is reasonable to conjecture that input-gradient attributions would adjust accordingly after the model is retrained for the next round. Further investigation into feature attribution in decision-making contexts (as opposed to post hoc interpretability) is an exciting direction for future work.

D.7 Wetlab Validation

1.0

2.0

3.0

4.0

5.0

6.0

7.0

8.0

9.0

10.0

11.0

Edit Distance from Seed

0.05

0.10

0.15

0.20

0.25

0.30

yield (mg)

1.0

2.0

3.0

4.0

5.0

6.0

7.0

8.0

9.0

10.0

11.0

Edit Distance from Seed

-0.2

-0.1

0.0

0.1

Δ yield (mg)

1.0

2.0

3.0

4.0

5.0

6.0

7.0

Edit Distance from Seed

6.0

6.5

7.0

7.5

8.0

8.5

-log10(KD) (M)

1.0

2.0

3.0

4.0

5.0

6.0

7.0

Edit Distance from Seed

-2.5

-2.0

-1.5

-1.0

-0.5

0.0

0.5

-Δlog10(KD) (M)

In this section we briefly summarize the experimental procedures used to validate LaMBO-2 designs in vitro . Designed antibody sequences from LaMBO-2 were expressed and purified, and surface plasmon resonance (SPR) measurements were used to determine binding affinity. See Figure 18 for a plot of design binding affinity vs. edit distance from seed antibody.

Plasmid Construction and Antibody Production: synthesized DNA of antibody variable domains (Twist Biosciences) were cloned into mammalian expression vectors using Gibson assembly. The whole vector was amplified using PrimeStar Max polymerase (Takeda). PCR products were transfected transiently in 1mL Expi293 cell culture. Expression lasted 7 days before harvest. Antibodies were affinity purified over a MAb Select SuRe resin (Cytiva), and their concentration was measured by optical density at 280nM.

Binding Affinity Measurements: affinity of the antibodies towards their target antigen was measured by surface plasmon resonance (SPR) at 37 °C on a Biacore 8K instrument (Cytiva) in HBS-EP+ buffer (10 mM Hepes, pH 7.4, 150 mM NaCl, 0.3mM EDTA and 0.05% vol/vol Surfactant P20). Antibodies were captured on a Protein A chip and their target antigen were injected for 5 minutes and allowed to dissociate for 10 minutes at 30ul/min. The surface was regenerated between cycles with 10 mM glycine pH 1.5. Affinity constants were obtained using Biacore Insight (Cytiva) using a 1:1 binding kinetics model.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper is focused on the task of multi-objective protein sequence design using guidance in diffusion models. Since protein sequences are discrete objects, classifier guidance tools from continuous diffusion are not directly applicable. To bridge this gap, the paper proposes a new method NOS, where the classifier guidance is applied to the latent representation of the sequence rather than the discrete sequence. Empirical evaluations on multiple real world sequence design tasks showcase the improved performance of their proposed methods.\n\nSTRENGTHS:\n1. The paper is well written and easy to follow. The technical choices made across the paper are well motivated.\n\n2. The biggest strength of the paper are the empirical evaluations on real world sequence design tasks. In particular, the antibody lead optimization experiment has both an in-silico component and an in-vitro component, with experimental validation of predicted model designs.\n\nWEAKNESSES:\n1. The technical contributions of the paper are largely incremental. Upon first reading the paper abstract, it appeared as though the paper proposes a new guidance method to work directly with discrete protein sequences. However, the guidance was applied to the continuous latent representation of the sequence. This is a natural and elegant direction to investigate, but diminishes the technical novelty.\n\n2. Some questions pertaining to the experimental evaluation:\n    * In the antibody lead optimization task, could the authors add other baselines such as DiffAb / RFDiffusion?\n\nQUESTIONS:\nPlease refer to the weaknesses section above. \n\nEdit: I have read the authors rebuttal. Post discussion with the authors, some aspects associated with novelty and technical contributions have been made clearer. I raised my score by 1 point, but not more as I am not an expert in Bayesian Optimization and some of the contributions of the work are associated with this field.\n\nLIMITATIONS:\nThe authors have addresses limitations and proposed directions for future work.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 2, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-05"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposes a discrete diffusion model guided by protein properties. They add a property predictor to the diffusion model, sharing the hidden states. The method is largely inspired by previous work, while experimental results are positive.\n\nSTRENGTHS:\n- The paper is well written.\n- The experimental results are convincing and good.\n- The function-guided protein design task is important.\n\nWEAKNESSES:\n- Many baselines addressing a similar problem are overlooked [1-4]. \n- The experimental evaluation is limited to a single attribute. The authors could extend the method to multi-objective optimization.\n\n\n\n[1] Kirjner, Andrew, et al. \"Optimizing protein fitness using Gibbs sampling with Graph-based Smoothing.\" arXiv preprint arXiv:2307.00494 (2023).\n\n[2] Sinai, Sam, et al. \"AdaLead: A simple and robust adaptive greedy search algorithm for sequence design.\" arXiv preprint arXiv:2010.02141 (2020).\n\n[3] Ren, Zhizhou, et al. \"Proximal exploration for model-guided protein sequence design.\" International Conference on Machine Learning. PMLR, 2022.\n\n[4] Brookes, David, Hahnbeom Park, and Jennifer Listgarten. \"Conditioning by adaptive sampling for robust design.\" International conference on machine learning. PMLR, 2019.\n\nQUESTIONS:\n- Q1: Have you performed in vitro experiments or just evaluated the designed sequence by another neural model?\n\n- Q2: Could you compare your method with more baselines [1-4]?\n\n- Q3: Could you extend the method to multi-objective optimization?\n\n[1] Kirjner, Andrew, et al. \"Optimizing protein fitness using Gibbs sampling with Graph-based Smoothing.\" arXiv preprint arXiv:2307.00494 (2023).\n\n[2] Sinai, Sam, et al. \"AdaLead: A simple and robust adaptive greedy search algorithm for sequence design.\" arXiv preprint arXiv:2010.02141 (2020).\n\n[3] Ren, Zhizhou, et al. \"Proximal exploration for model-guided protein sequence design.\" International Conference on Machine Learning. PMLR, 2022.\n\n[4] Brookes, David, Hahnbeom Park, and Jennifer Listgarten. \"Conditioning by adaptive sampling for robust design.\" International conference on machine learning. PMLR, 2019.\n\nLIMITATIONS:\nN/A", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-06"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposes new methods for discrete diffusion guidance, making it possible to optimize protein sequences for local and global properties while retaining high sequence likelihood.\n\nSTRENGTHS:\n1. The paper addresses a problem closely related to the NeurIPS community.\n2. The methods sound reasonable.\n3. I think the introduction is well-written, especially the second paragraph, which effectively highlights the significance of protein sequence design.\n4. The experimental results are also significant.\n\nWEAKNESSES:\n1. Overall, I find the presentation of the paper quite comfortable, but I feel that the reading experience is slightly compromised due to some details not being included in the main text.\n2. Lack of the comparison with the guidance for discrete diffusion in [1].\n\nQUESTIONS:\n1. I looked at the supplementary materials, and it seems that the selection of edit positions is fixed after the initial determination. What would happen if the positions were reselected at each denoising step?\n2. There is another method for discrete guidance [1]. I would like the authors to compare the results of this method with the proposed approach.\n3. This seems to resemble an optimization problem where a seed sequence is given, and during inference, only a few positions are edited for amino acid types. I'm curious about the training process. Was the training performed on the entire sequence?\n4. The notation $\\mathcal{D}$ in Line 168 is not specified or defined.\n5. The notation $d$ in Eq.5 is not specified or defined. I guess it may represent the dimension of embeddings.\n6. $C_{300}^{8}$ instead of $C_{8}^{300}$ in Line 260?\n\nI would be happy to revise my score if the authors can provide satisfactory responses.\n\n**Minor Concern:**\n1. Equation 1 has an extra comma at the end.\n2. The notation $L$ in the article is ambiguous, as it is used in two different contexts. In one instance, it represents the length of the sequence, while in another instance, it represents the depth of the encoder network.\n\nLIMITATIONS:\nThe paper lacks a discussion of the limitations of the work. I hope the authors will address the limitations and include this discussion in the revised version of the paper during the rebuttal period.\n\n**Reference:**\n1. Vignac, Clement, et al. \"DiGress: Discrete Denoising diffusion for graph generation.\" The Eleventh International Conference on Learning Representations. 2022.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 5, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-07"}

comments: SUMMARY:
The submission introduces a new diffusion for protein design. It focuses on overcoming the challenge of discrete nature of protein sequences. The model succeeded in local and global properties optimization of antibody sequences making them more applicable for therapeutic purposes.

STRENGTHS:
Originality: Protein diffusion models aren't new (there are many related contribution into this field), however, the discrete nature of protein sequences challenges the usage of these models directly in sequence space, Therefore, the submission is important, novel, and useful. The claims are well supported.
Quality: Overall quality of the paper is high, but it's better to add and extend the comparison to and discussion of other protein diffusion models.
Clarity: The paper is clearly written and contains all the necessary citations. Methods, experiments, and equations are clearly described. 
Significance: Diffusion is a promising approach in protein design. The submission provides a valuable impact to optimize protein sequences for local and global properties.

WEAKNESSES:
The submission lack of comparison to currently available deep learning protein design approaches such as RFdiffusion.

QUESTIONS:
1. Please provide the comparison to RFdiffusion and other available deep learning protein design tools.

LIMITATIONS:
The paper doesn't describe limitations and future work.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 4
DATE: 2023-07-07

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper presents a novel approach to applying diffusion models to optimize protein sequences, specifically for antibody design. Diffusion models have been successful in handling continuous data like images, but their direct application to protein structures faces challenges due to limited high-quality structural data. The authors propose a new method for discrete diffusion guidance, enabling the optimization of protein sequences for desired properties while maintaining high sequence likelihood. They further apply this method to a real-world protein design task, focusing on optimizing antibodies for increased expression yield and binding affinity to a therapeutic target while adhering to locality and liability constraints.\n\nSTRENGTHS:\nNovel Ideas and Comprehensive Solution: The method proposed in the paper introduces several novel ideas that address a critical challenge in protein design.\n\nA notable strength of the method is its ability to optimize protein sequences without relying on structural information.Protein design approaches often face limitations due to the scarcity of high-quality structural data and the need to convert structures into sequences for synthesis. By bypassing the requirement for structural information, the proposed method reduces complexity and provides a more straightforward and practical solution for designing proteins.\n\nThe absence of a structural component in the optimization process simplifies the method's implementation and reduces computational complexity. \n\nThe authors demonstrate the efficacy of their approach through real-world evaluations. By applying the method to optimize antibodies for increased expression yield and binding affinity to a therapeutic target, the paper offers tangible and meaningful results that resonate with practical protein design needs.\n\nOne of the noteworthy strengths of the method is its ability to optimize protein sequences while maintaining high sequence likelihood and by extension expression efficiency.\n\nWEAKNESSES:\nUnclear Introduction and Insufficient Motivation: The paper's introduction lacks clarity and fails to provide a meaningful digest of key topics in the field of protein design. As a result, the motivation for the research is poorly established, which hampers the reader's understanding and engagement throughout the rest of the paper, making it challenging to appreciate the significance of the presented work.\n\nLack of Reproducibility: The absence of all test data and failure to provide the code used in the experiments significantly impairs the reproducibility of the results. Reproducibility is a fundamental aspect of scientific research, and the unavailability of essential resources restricts other researchers' ability to validate and build upon the findings.\n\nMissing Computational Cost and Time Information: The paper lacks essential details concerning the computational cost and time required for both training and inference of the proposed method. Such information is crucial for assessing the feasibility and practicality of the approach, and its absence leaves readers with an incomplete understanding of the method's resource requirements.\n\nInadequate Description of Model Architecture and Training Procedure: Section 4.2 requires more comprehensive information on the model architecture, feature extraction layers, and the encoder used in the proposed method. Similarly, the training procedure and hyperparameter tuning should be thoroughly explained to ensure clarity and reproducibility.\n\nInsufficient Information on Test Data: The paper does not provide sufficient details about the test data used, such as the number of samples and the criteria for their selection. Without this information, it remains unclear if the reported results generalize well to diverse protein sequences.\n\nLack of Relevant Information in Figure Captions: Figure captions throughout the paper lack relevant information that could aid in the readers' interpretation of the presented results. Clear and informative captions are essential for facilitating a comprehensive understanding of the visual data.\n\nMissing Quality Assessment of IgFold Structures: In section 5.1, the authors fail to report on the quality of the IgFold structures used in structure-based methods. Since the performance of these methods can be highly dependent on the starting structure's quality, this information is crucial for the paper's credibility.\n\nInadequate Explanation of Comparison Choice: In section 5.3, the rationale behind comparing salient selection with uniform random selection from the entire sequence remains unclear. The authors should provide a more insightful justification for this choice, considering the low probability of randomly selecting critical regions like CDRs.\n\nLack of Ablation Analysis for Objective Functions: The authors do not thoroughly evaluate the importance of the chosen objective functions through ablations. Specifically, in section 5.2, the effect of 40% beta on ProtG3T Log Likelihood should be explained and analyzed in more detail.\n\nUnexplained Statement on Sensitive Data: The statement regarding non-disclosure of specific drug targets in the experiments lacks a clear explanation. The authors should clarify the reasons behind this decision and provide reassurance regarding the privacy and ethical handling of sensitive data.\n\nAvoidance of Emotional Statements: The use of emotional statements, such as \"eager to\" or \"encouraging to see,\" should be avoided in scientific writing. Such statements can introduce bias and subjectivity, undermining the paper's objectivity and professionalism.\n\nMinor issues\n* Line 46: “there is no analogous method for proteins, which do not encode the same general-purpose semantics. As a result many […]” It is not clear what the authors are saying here. Explain how the encoded semantics in protein space differs from NLP and why other sampling methods have been used to address this issue.\n* Figure 1 caption: “[...] a powerful method for protein design.” Please stick to informative text about the figure in the caption.\n* Line 84: “A protein is only useful if it can be expressed in living cells [...]” This is incorrect. A protein could be expressed for instance in a cell-free translation system and still be useful.\n* Figure 7 caption: Last sentence doesn’t belong in the caption.\n* Equation (1), remove Z and replace equation with proportionality or explain what Z is.\n* The OAS dataset needs to be briefly described somewhere in the text.\n\nQUESTIONS:\n1. Could you provide a more concise and clear introduction that offers a comprehensive overview of the key topics in the field of protein design? Additionally, could you elaborate on the motivation for the research to better establish the significance of your work throughout the paper?\n2. Can you provide detailed information on the computational cost and time required for both training and inference of your proposed method? This will help readers assess the feasibility and practicality of implementing your approach.\n3. In section 4.2, could you offer a more thorough explanation of the model architecture, feature extraction layers, and the encoder used in your proposed method? Similarly, can you provide more information about the training procedure and hyperparameter tuning to ensure clarity and reproducibility?\n4. Please provide additional details about the test data used, including the number of samples and the criteria used for their selection. This information will help readers understand the generalizability of the reported results to diverse protein sequences.\n5. Could you include more relevant information in the figure captions to aid readers in interpreting the presented visual data?\n6. Can you report on the quality assessment of the IgFold structures used in structure-based methods in section 5.1? Understanding the starting structure's quality is crucial for evaluating the performance of these methods.\n7. Could you clarify the statement regarding the non-disclosure of specific drug targets in the experiments? Please provide an explanation for this decision and assure readers of the ethical handling and privacy of sensitive data.\n8. Line 37, what does this mean: “Searching directly in sequence space eliminates the need to recover sequence from structure.”\n9. Line 38, any examples (references) for: “ Protein sequence models are also comparatively fast, especially during inference, and can leverage sequence datasets that are often several orders of magnitude larger than their structural equivalents.”\n\nLIMITATIONS:\nThe presentation lacks clarity and fails on several points (see weaknesses).\n\nThe unavailability of all test data and the absence of provided code significantly impair reproducibility, limiting the validation and extension of the research by other researchers.\n\nThe paper does not present detailed information on several topics (see weaknesses).\n\nAblation studies are lacking (see weaknesses).", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-26"}