Protein Design with Guided Discrete Diffusion
Authors: Nate Gruver
Samuel Don Stanton
Nathan C. Frey
Tim G. J. Rudner
Isidro Hotzel
Julien Lafrance-Vanasse
Arvind Rajpal
Kyunghyun Cho
Andrew Gordon Wilson
Venue: NeurIPS 2023 spotlight (2023)
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
Nate Gruver*! Samuel Stanton *2
Nathan Frey?
Tim G. J. Rudner'
Isidro Hotzel
Julien Lafrance- Vanasse
Arvind Rujpal Kyunghyun Cho Andrew Gordon Wilson'
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 fitmess. Given its broad success in conditional sampling, classifier-guidee diffusion modeling is a promising foundation for protein design, leading many develop guided diffusion models for structure with inverse folding to recover sequences. In this work, we propose diffussoN 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-basee methods, including scarce data and challenging inverse design. Moreover, we use NOS to generalize LaMBO, 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 real-
through novel application of saliency maps. We apply LaMBO-2 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.
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 the patient than radiotherapy 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 small enriched library of candidates to find viable option in as few measurements as possible [44]. In practice these 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].
enriched libraries are usually obtained through massive
There many approaches 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
* "quual contribution. NNew York University, ("Prescient Design, 'Antibody Engincering) Genentech.
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

diffusioN Optimized Sampling (NOS)
INDAP ENOUGGRAAE/*
25806 SPn No. juturr varasspvvo Rsa0munns JAVTSSIN NL..
Tx
A44
X
(ww.r)eeSSNN..O0e KUPomSeenTL (LI0M)) TSCADYY00 /OOOLALT (AVINOEO
x
Unguided Sampling
1kolihood
toodt
Figure I: We propose diffusioN Optimized Sampling (NOS), method for gradient-guided sampling from discrete diffusion models. NOS uses I sampling steps of denoising diffusion, where each step consists of applying a corruption, gradient steps to optimize a value function, f. and sampling of the next discrete sequence, or corresponding latent state. NOS generates samples that optimize an arbitrary objective while maintaining high likelihood with respect to a reference distribution of sequences. We combine NOS with LaMBO, a strong Bayesian optimization method for sequence design [73], to make LaMBO-2, our improved method for protein design.
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 result, many sequence search methods resort to gradient-fres free sampling methods like Metropolis-Hasting 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 (Latent Multi-Objectivy Bayesian Optimization), 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, 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 simultancously enforcing an edit budget and choosing optimal edit locations based on that budget.
MCMC [78, 37], which are flexible but computationally
To address the first issue we propose NOS (diffusioN Optimized Sampling), a new method for controllable categorical diffusion (Figure 1). Diffusion models capture complex distributional dependencies by making iterative denoising steps, but there 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 -fate-of-the-art structure and sequence- -aaed baselines on both unguided and guided infilling tasks.2 To address feature
the second problem (choosing optimal edit locations) we propose using umbedding-gradiee attributions (i.c. 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.
hhupss//gitubb.co///ggwv/e////

---

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 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 [3 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, 51J, which also circumvent inverse folding. Because these models still rely on structure information 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 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] 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& 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 enginecring a latent space, Dathathri et al. [17] and Yang and Klein [83] propose gradient-guided autoregressive models by using the decoder's s activations as a gradient- friendly latent space. These methods alternate between sampling from logits and ascending the likelihood of 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 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 E A with alphabet A and fixed length L,² which maximize single objective f(w) (e.g.. binding affinity) or multiple objectives fi(w), ...
fu fu(w) (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 E B (1,..., L 1} to a seed sequence 8 AL, 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 E E c A", where E is the set of all expressible proteins. Since naturally occurring sequences must express in order to be observed, w(w), the likelihood of protein with respect to empirical distribution of natural protein sequences, is often taken as proxy for the tendency of a protein to express. In protein design, these proxies are
Length change enabled by the use of protein sequence alignments, which introduce a gap token "'",
3

---

Continuoux Process
Mamiiirro (1- oa)))
Cutegorical Process
Cat( ..)
Uow
ZT
...
tO
WT
WA
27 1)
[AAK]
go 8o(00,-//o0) (1s))
Figure 2: Two approaches to diffusion generative modeling for categorical variables. (Left) Cate- gorical data is embedded into continuous variables with an accompanying continuous noise process. (Right) Categorical noise is applied directly to sequences, and corrupted sequences denoised using standard language modeling methods.
typically called metrics of naturalness. Since 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
p(w) exp(-E(w))/Z, E(w) - -(w))v(w),
(1)
where Fiww) log ) is scalar energy function, and the value function v: :A - R expresses the utility of a sequence with respect 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 al. [78] can require hundreds of thousands of steps 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, Zg. to samples from a prior =(x) (Figure 2). The forward process (No - .T) is composed of conditional distributions ((-1)) (i.e., the noise process) that admit closed forms for the conditional distributions pl (f (zo) and 1 Fo) (e.g., independent Gaussian corruption). The reverse process (ET ro) converts samples from the prior into samples from the learned data distribution Pa (zo) by repeatedly predicting the denoised variable fo from noisy values at and using the conditional distribution p( (x,, 20) to derive transition distribution, pa(x,-11 ( (xx). The specific choice 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 diffusion generative model, depending on the nature of the noise process. We include brief descriptions below and more detailed account in Appendix A.
learn a distribution p(w), one strategy is to first embed w to continuous Continuous noise
variable with embedding matrix Up and apply Gaussian noise [21]. The prior is taken to be (=)) = (0, 7) while the forward process is p(z; (zo) (xx; Vārza, (1 - āz) for â₂ [O, 1]. The values of âr are determined by user-specified corruption schedule. For the reverse process, we learn a function, pa(u Ze, t), to predict the sequence from noised points x, by minimizing the following loss with respect to 0:
L(0)= Ew..t [- log po(w)(=))], 2, p(æe|=0 Uowa).
Using po pr, 1) we can construct a distribution for the reverse process
po(#1-117)) = Σ (F Ugr) pl p(1)x,,)) !),
(2)
d
where p(2;-1/=4, 2y) also a Gaussian distribution. At inference time, we can use the learned reverse process to convert samples from π(х) into samples from the learned distribution po(zo) by repeatedly sampling Po (#+-1(1;)) followed by sampling w ~ po (mllaoo. 0).
Categorical noise Alternatively, Austin et al. [4] proposed a forward process which operates directly on w, by introducing an absorbing state for each token () = [MASK]. The forward process (wo wy) is defined by discrete transition matrix, describing the probability of mutating token into a [MASK], and the corresponding prior is simply point mass on the sequence of all [MASK] tokens. To learn the parameters of the denoiser po (0 wy, t) we maximize the likelihood of the denoising process on ground truth sequences
(()) Ewn. [- log po powwollue (wo/):)]
We ((w)(wo)

---

Then, as above, we can use the denoiser to construct the reverse process
po(we-1lun Σ p w [w,, whp ( (t
where p(w wy, wo) is also categorical distribution derived using Bayes' rule. To sample, the transition distribution is applied for [T,....0] 0].
4 Methods
Now present practical methods for efficiently sampling from p(w) of p(w) exp(e(w)) (Eq. 1) by modifying the learned transition distribution with a learned value function Up (w). We then show how this sampling method can be used 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 general form of gradient guidance (NOS) for discrete diffusions with categorical denoising models, i.c. 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, any denoising network, e.g. po(ii vt; l), the discrete sequence w, has many corresponding continuous representations in the form of hidden states of the model h, ga(we) Upw-- Notably, go(w)). corruption
for d E {O..... D), where D is the depth of the encoder network and go(ure) for the Gaussian diffusion models in Sec. 3, we can equivalently have and sampling are performed on the learned token embeddings. In the case of the categorical noise
diffusion po( w0o (ay) po tool/t)) and thus for the purpose of guidance, we can consider general pe(1i (w) h2) for both forms of corruption.
To sample from (wa) oc po (wx) exp(2g(wx)), we construct a modified denoising model, po( h4)) oc po he) exp(ve(he)),
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 wt In Appendix D.4 we propose simple procedure for corrupted discriminative training inspired by label smoothing [76]. Using this modified denoising model we can construct modified transition distributions using Eq. or Eq. 3. There is one key difference between these transition distributions: in the continuous (Eq. 2), smooth steps are taken in the token embedding space, while in the discrete case (Eq. 3) the transition leads large jumps from one token embedding to another. In either case, it is possible to sample a discrete sequence to at any point in the chain using the logits of the denoiser po (úi /t)) When using Eq. 2 to derive 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 F > O, with the update step,
M A4 - 1) Nn [XKL(pp (M))) MP (t)))) va(M()] + 2 &~N(0,11,, (4)
where 7) is the step size and À is the regularization strength, followed by sampling po no(w:-1/h) IMa) or (b)))-1())) While the gradient 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 craddeet-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 particularly important for protein design, because function can be determined by complex interactions between distant parts of the sequence. As we see in See. 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
5

---

ENGLNEG SOA.SSOOLOOOOOOOOOOOOOOOOO
Figure An example of binding affinity saliency map produced by LaMBO-2 with NOS-D. For simplicity, only the variable heavy (VH) region of the hu4D5 antibody is shown. Positions corresponding to complementarity defining regions (CDRs) are enclosed in green boxes. Converting this saliency map to an edit position distribution will concentrate computational resources on editing CDRH3, which is often manually selected by experts. Some resources are also allocated to the framework and other CDRs since these positions may also affect binding.
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-base uncertainty quantification.
Architecture and value function In order to apply the methods discussed in Subsec. 4.1 we require generative diffusion model pe(w) and discriminator fe(w) which share hidden layers up depth d. The discriminator is trained to predict the objective function f. Like LaMBO-1 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 n (w) E[u( w, f. D.)]] where D, is the data already available and the expectation is taken with respect to posterior pe (D, 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 different choice of utility function [82]. To estimate the expectation we draw samples from Po( with an approach call partial deep ensembles, where the discriminative layers of the model above the shared encoder are replicated & 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 ho 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 ho we define s. (ho)) the saliency with respect to e of position i {L...., L} as
sc(ho) = max {
Σ (Vhe(thol)a)))
1/7 eJ.
PPodit wn = Σ, sh(u))
s;(ho)
(5)
where 7 >> O is a temperature hyperparameter and 0 < << 1. AS +00, Pledit w) 1/L for all i. For each sequence we draw B edit positions without replacement according Eq. 5. We conserve parts of the input we cannot change (e.g. the antigen sequence) by setting the the saliency to O 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 highly conserved position, then the predicted logits will be low entropy and the guidance will incur large KL penalty for changes (Eq. 4).
Experiments
We evaluate our methods on three increasingly complex antibody design tasks. First we compare trained diffusion models unguided infilling tasks, showing that sequence diffusion methods

---

RFDiffusion
DiffAb
IgLM
NOS-C (Ouns)
NOS-D (Ours)
Seq. Recovery
1.0
0.5
0.0
HCDR1
HCDR2
HCDRR++CDR2++CCRS
LCDR3
Infill Region
Figure 4: We infill antibody CDRs with discrete diffusion models (ours) and compare against structure-based diffusion models (DiffAb [51] and and RFDiffusion [80]) and an autoregressive antibody language model (lgLM [68]). We diffusion on sequences alone-without structural priors-reliable leads to high sequence recovery. For structure based methods, we first fold seed sequences with IgFold [64] and then run joint sampling of sequence and structure for the CDR. We sample 10 infills for each of the 10 antibody seed sequences selected randomly from paired OAS [56].
consistently outperform structure-based methods when only predicted structures are available*. 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 silcco and in vitro experiments.
5.1 Unguided antibody CDR infilling
We focus immunoglobulin G (IgG) format antibodies, which are comprised of heavy (H) chain and light (L.) chain. Each chain has three complimentarity determining regions (CDRs), which tend to have strong effects on binding affinity 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], GPT2 language model trained on OAS, and two structure- based methods: DiffAb [51]. joint eequence-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 IgIM is trained with ti--in-the-midlee augmentations [7]. it does not natively support infilling multiple non-contiguous regions, and we so by replacing regions that are not yet sampled with [UNK] tokens. For the structure- trueture-based methods, we provide starting structures generated with IgFold [64].
target antigen but limited effects other structural
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-structen underperforms sequence-only diffusion, most likely because our sequence- baseeneeebbedd approaches have access to 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 he explained by the relative scarcity of antibody structures in the PDB compared to SAbDab, which has an antibody every structure. The poor performance of structural methods on CDR infilling could also be result of poor sequence recovery from structure during inverse folding, a problem that could amplified for relatively unstructured loop regions like CDRs.
co-design out of the box, often
5.2 Optimizing antibodies for in silico objectives
To test guided sampling from our model, experiments simple single-objective tasks:
The percentage of beta sheets, measured on primary sequence [15]
The solvent accessible surface area (SASA) of the protein's 's predicted structure [64]
'In practical protein design campaigns it is infeasible to get ground truth structural measurements for proposed designs, and predicted structures are the only alternative available.
7

---

Seeds
NOS-D (Ours)
NOS-C (Ours)
RFDiffusion
DiGress
Beta Sheets
1e4
% Beta Sheets
1e4
0.3
SASA
1.2
SASA
1.2
0.2
0.2
XX
1.1
x
1.1
-500
-400
-500
400
-500
-400
-500
400
ProtGPT Log Likelihood
ProtGPT Log Likelihood
Figure 5: Comparing samples from NOS (ours) with alternative guided generation methods and structure-based models. NOS exhibits higher likelihood for similar or dramatically improved values of the objective. (left) Sequence diversification (resampling and selecting improved points) with DiffAb [51] or RFDiffusion [80]. DiffAb generates sequences and structures simultaneously. while sequences for RFDiffusion are obtained using ProteinMPNN [19]. Compared with NOS, these methods do not effectively optimize the objective and yield low- wwwliikl sequences. (right) Guided generation using PPLM [17]. guidance method for autoregressive language models (in this case IgLM [68]) and DiGress, a competing guidance method for discrete diffusion models [79]. NOS, PPLM, and DiGress are sampled for many settings of guidance strength (e.g. 17 and X (Eq. 4)) to demonstrate the full range of trade-offss between the objective and likelihood. We provide details about hyperparameter settings in Appendix C.5 and additional density plots in Appendix C.6.
Since we want plausibly natural antibodies with high objective value we examine the Pareto front for samples optimized for each objective, with log- likeliidddo assigned by ProtGPT [29] (trained on Uniref50 [75]) plotted against the value of the objective. As autoregressive guided baseline, run PPLM, using lgLM as the base generative model (details in Appendix C.3). We use DiGress [79] as a guided diffusion baseline. DiGress uses gradients 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.
one-hot representations to performing guidance
5.3 Antibody lead optimization: in silico evaluation
Having established the performance of NOS on simpler benchmarks, we now turn 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, therapeutic antibody targeting the HER2 antigen"
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 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 simple incquality 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.
$ HHER2 is important target for certain types of breast and stomach cancer [36].
8

---

8 B
=2
0.5
Bust Aco Volue
600
Agg, Volue
2000
400
0.3
Acq.
2000
3000
30
20 Stmpping Steps
30
10
20 Sampping Steps
Diffusion
Diffusion Stop
Uniform Edits Unitorm Gundance
Sollent Edits No Guisance
AdaLead
LaMBO-2 PEX
WJ5
Sallent Edits
Figure 6: (left) Naturalness constraints present challenge for genetic methods, which rapidly decline in naturalness even as their objective value improves. The grey dashed line is an empirical lower bound on naturalness above which in vitro expression is reliable. Although AdaLead and PEX both improve the acquisition value, they quickly leave well-supported areas the search space (drop below the dashed line), shown by the faded section of each curve. By contrast, the naturalness of LaMBO-2 samples degrades much more slowly while consistently improving the acquisition value. (right) Ablating the effects of guidance and edit position selection. We start with the hu4D5 HER2 antibody and vary the edit budget B E {8, 32), optimizing for expression yield and binding affinity. For all choices of edit budget, we find that the effect size of edit position selection is much larger than that of guidance, making salient unguided edits a surprisingly strong baseline.
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 pructure-based unguided infilling (i.c. 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 relatively small edit budget B (B <= 10% of mutable positions). To isolate the effects of salient edits we baseline against edit positions chosen uniformly random, and to isolate the effects of guidance set the step size y (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 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%) is dwarfed by the combinations of possible edit positions (C200). 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 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 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

---

AbSci (2023)
LaMBO-2 (Ours)
Exprepton Bate
1.00
0.75
improved
0.6
Binding Mate
0.6
Binding Improved
0.4
0.3
0.50
9.2
0.2
0.2
9.1
0.00
0.0
0.0
0.0
Round
Round 2
Round
Round
Sound
Rourd
Round
Raund 2 Round
Round
Round
Bound
Figure 7: We use LaMBO-2 to optimize 20 seed antibodies for 4 different target antigens three experimental rounds, retraining the model after each round. Some design choices and hyperparameters varied from round round, with substantial impact the results. In the last round tested 56 antibodies and attained a 99% expression rate and 40% binding rate on average across targets. On average 43% of the expressing designs had higher yield and 21 % of binding designs had higher binding affinity than the corresponding seed. These results are very encouraging when placed in context with a related experiment designing HER2 antíbody libraries [67]. Our results provide evidence that enriched antibody libraries can be created in silico without the assistance of high- throughput in vitro screening.
hu4DS/HER2', 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 targets with a median 5 edits to the seed. See Appendix D.7 for individual yield and affinity measurements and experimental details.
The improvements 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 Round 2 can mainly be attributed to framework residue deletions that arose when (the KI. 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 Shanchsazzadeh 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 true upples-to-apple comparison. In the latter experiment 1M designs were generated for the HER2 target and screened with high-throe thghohhhoughhut assay. After screeming 421 designs were validated with high-fidelity surface plasmon resonance (SPR) assay. addition to wetlab screening, their experiment also restricted edits to specific antibody CDRs. We optimized antibodies for 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 sílico 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 a fruitful direction, 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 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 gene perturbations [42], which can exceed 20K tokens in length and may necessitate the use of recent ad 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 difficult express black-box functions constraints [49, 58].
'Due the sensitive nature of the data, do not disclose the other seeds drug targets.
10

---

References
[1] Rebecca Alford, Andrew Leaver-Fay, Jeliazko R Jeliazkov, Matthew O' MeMara, Frank 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, (3(6):3031-3044 2017.
[2] Namrata Anand and Tudor Achim. Protein structure and sequence generation with equivariant denoising diffusion probabilistic models. arXiv preprint arXiv:2205.15019, 2022.
[3] Charles Audet. A survey on direct search methods for blackbox optimization and their applica- lions. Springer, 2014.
[4] 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:1 19911-19993, 2021.
[5] David Baehrens, Timon Schroeter, Stefan Harmeling, Motoaki Kawanabe, Katja Hansen, and Klber-Robert Müller. How to explain individual classification decisions. The Journal of Machine Learning Research, 11:1003-181,, 2010.
[6] Jasmijn Bastings, Sebastian Ebert, Polina Zablotskaia, Anders Sandholm, and Katja Filippova. " will you find these shorteuts?" a protocol for evaluating the faithfulness of input salience methods for text classification. arXiv preprini arXiv:2111.0738 2021.
[7] 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. 24255 2022.
[8] Prajjwal Bhargava, Aleksandr Drozd, and Anna Rogers. Generalization in nli: Ways (not) to go beyond simple heuristics, 2021.
[9] Carl Ivar Branden and John Tooze Introduction protein structure. Garland Science, 2012.
[10] 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. Profein crystallography: methods and protocols, pages 627-641, 2017.
[11] Stephen Casper, Yuxiao Li, Jiawei Li, Tong Bu, Kevin Zhang, and Dylan Hadeield-Menell. Benchmarking interpretability tools for deep neural networks. arXiv preprint arXiv:2302.10894 2023.
[12] Huiwen Chang, Han Zhang, Jarred Barber, AJ Maschinot, Jose Lezama, Lu Jiang, Ming- Hsuan Yang, Kevin Murphy, William T Freeman, Michael Rubinstein, al. Muse: Text-to-image generation via masked generative transformers. arXiv preprint 2XXiv23301.00000 2023.
[13] Rewon Child, Scott Gray, Alec Radford, and Ilya Sutskever. Generating long sequences with sparse transformers. arXiv preprint arXiv: 10509, 2019.
[14] Hyungjin Chung, Byeongsu Sim, Dohoon Ryu, and Jong Chul Ye. Improving diffusion models for inverse problems using manifold constraints. arXiv preprint XXiv:22066009411 2022.
[15] 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-1444 1423, 2009.
[16] Tri Dao, Dan Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. Flashattention: Fast and memory-efficient exact attention with io-awareness. Advances Neural Information Processing Systems, 35: 35:16344-16359, 2022.
[17] Sumanth Dathathri, Andrea Madotto, Janice Lan, Jane Hung, Eric Frank, Piero Molino, Jason Yosinski, and Rosanne Liu. Plug and play language models: simple approach to controlled text generation. arXiv preprint arXiv:1912.02 2019.

---

[18] Samuel Daulton, Maximilian Balandat, and Eytan Bakshy. Differentiable Expected Hypervol- ume Improvement for Parallel Multi- Objei-O Bayesian Optimization. Advances in Neural Information Processing Systems, 33, 2020.
[19] Justas Dauparas, Ivan Anishchenko, Nathaniel Bennett, Hua Bai, Robert Ragotte, Lukas Milles, Basile IM Wieky, Alexis Courbet, Rob J de Haas, Neville Bethel, et al. Robust deep learning-based protein sequence design using proteinmpnn. Science, 378(6615):44 5):49566, 2022.
[20] Prafulla Dhariwal and Alexander Nichol. Diffusion models beat gans on image synthesis. Advances in Neural Information Processing Systems, 34:8780-8794 2021.
[21] Sander Dieleman, Laurent Sartran, Arman Roshannai, Nikolay Savinov, Yaroslav Ganin, Pierre H Richemond, Arnaud Doucet, Robin Strudel, Chris Dyer, Conor Durkan, et al. Continu- ous diffusion for categorical data. arXiv preprint arXiv:2211.15089, 2022.
[22] Benjamin Doerr and Frank Neumann. Theory of evolutionary computation: Recent develop- ments in discrete optimization. 2019.
[23] John Duchi, Elad Hazan, and Yoram Singer. Adaptive subgradient methods for online learning and stochastic optimization. Journal of machine learning research, 12(7), 2011.
[24] James Dunbar and Charlotte M Deane. Anarci: antigen receptor numbering and receptor classification. Bioinformatics, 32(2):298-300, 2016.
[25] James Dunbar, Konrad Krawezyk, Jinwoo Leem, Terry Baker, Angelika Fuchs, Guy Georges, Jiye Shi, and Charlotte M Deane. Sabdab: the structural antibody database. Nucleic acids research, 42(D1):D1 140-):14410 2014.
[26] Patrick Emami, Aidan Perreault, Jeffrey Law, David Biagioni, and Peter St John. Plug & play directed evolution of proteins with gradient-based discrete meme. Machine Learning: Science and Technology, 4(2):025014, 2023.
[27] Michael Emmerich. Single-and multi mujti-objecctiv evolutionary design optimization assisted by gaussian random field metamodels. dissertation, Universität Dortmund, 2005.
[28] Michael TM Emmerich, André H Deutz, and Jan Willem Klinkenberg. Hypervolume- based expected improvement: Monotonicity properties and exact computation. In 20/ IEEE Congress
of Evolutionary Computation (CEC), pages 2147-2154. IEEE, 2011.
[29] 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.
[30] 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 Spotlight presentation.
[31] Marjan Ghazvininejad, Omer Levy, Yinhan Liu, and Luke Zettlemoyer. Mask-predin Parallel decoding of conditional masked language models. arXiv preprint 1xXiv:1904.09322 2019.
[32] Vladimir Gligorijevic, Daniel Berenberg, Stephen Ra, Andrew Watkins, Simon Kelow, Kyunghyun Cho, and Richard Bonneau. Function-guis guided protein design by deep manifold sampling. bioRxiv, 2021.
[33] Rafael Gómez-Bombardee Jennifer N Wei, David Duvenaud, José Miguel Hernández-Lobate Benjamín Sánchez-Lengeling Dennis Sheberla, Jorge Ryan 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.
Timothy D Hirzel,
[34] Alexandros Graikos, Nikolay Malkin, Nebojsa Jojic, and Dimitris Samaras. Diffusion models as plug-and- -nd-aly priors. arXiv preprint arXiv: 2022.
12

---

[35] Albert Gu, Karan Goel, Ankit Gupta, and Christopher Ré. On the parameterization and initialization of diagonal state space models. Advances Neural Information Processing Systems, 15:35971-35988. 2022.
[36] Carolina Gutierrez and Rachel Schiff. Her2: biology, detection, and clinical implications. Archives of pathology & laboratory medicine, 135(1):55 62, 2011.
[37] Brian Hie, Salvatore Candido, Zeming Lin, Ori Kabeli, Roshan Rao, Nikita Smetanin, Tom Scrcu, and Alexander Rives. A high-level programming language for generative protein design. bioRxiv, pages 2022-12, 2022.
[38] Jonathan Ho and Tim Salimans. Classifier-free diffusion guidance. arXiv:2207.12590 2022.
arXiv preprint
[39] Jonathan Ho, Ajay Jain, and Picter Abbeel. Denoising diffusion probabilistic models. Advances in Neural Information Processing Systems, 33:6840-6883 6851, 2020.
[40] Emiel Hoogehoom, 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-124655 2021.
[41] Sara Hooker, Dumitru Erhan, Pieter-Jan Kindermans, and Been Kim. Evaluating feature importance estimates. web preprint
2018.
[42] Yuge Ji, Mohammad Lotfollahi, F Alexander Wolf, and Fabian Theis. Machine learning for perturbational single-cell omics. Cell Systems, 12(6):522-537, 2021.
[43] Simon Kelow, Bulat Faezov, Qifang Xu, Mitchell Parker, Jared Adolf-Bryfoge and Roland L. Dunbrack Jr. penultimate classification of canonical antibody edr conformations. bioRxiv, pages 2022-10, 10, 2022.
[44] Benjamin Larman, George Jing Xu, Natalya Pavlova, and Stephen J Elledge. Construction of rationally designed antibody platform for sequencing-asisted selection. Proceedings of the National Academy of Sciences, 109(45): 2012.
[45] Jin Sub Lee and Philip M Kim. Proteinsgm: Score-based generative modeling for de novo protein design. bioRxiv, 2022.
[46] Sergey Levine. Reinforcement learning and control as probabilistic inference: Tutorial and review. arXiv preprint arXiv:7805.0090 000,00 2018.
[47] Xiang Lisa Li, John Thickstun, Ishaan Guirajani, Percy Liang, and Tatsunori B Hashimoto. Diffusion-In improves controllable text generation. arXiv preprint 2rXi::2205.14217 2022.
[48] Hao Liu, Matei Zaharia, and Pieter Abbeel. Ring attention with blockwise transformers for near-infinite context. arXiv preprint arXiv:2310.01889, 2023.
[49] 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 23Xx:23022.0671 2023.
[50] 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.
[51] Shitong Luo, Yufeng Su, Xingang Peng, Sheng Wang, Jian Peng. and Jianzhu Ma. Antigen- specific antibody design and optimization with diffusion-based generative models. bioRxiv, 2022.
[52] Pascale Mathonet and Christopher G Ullman. The application next generation sequencing to the understanding of antibody repertoires. Frontiers in immunology, 4:265, 2013.
13

---

[53] Igor Melnyk, Payel Das. Vijil Chenthamarakshan, and Aurelie Lozano. Benchmarking deep generative models for diverse antibody sequence design. arXiv preprint arXiv:2111.0680/ 2021.
[54] 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 computer vision and pattern recognition, pages 4467-4477, 2017.
[55] Alexander Quinn Nichol and Prafulla Dhariwal. Improved denoising diffusion probabilistic models. In International Conference on Machine Learning, pages 8162-8171. PMLR, 2021.
[56] 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.
[57] 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.0633 2023.
[58] Vishakh Padmakumar, Richard Yuanzhe Pang. He He, and Ankur P Parikh. Extrapolative controlled sequence generation via iterative refinement. arXiv preprint arXiv:2303.04562, 2023.
[59] 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.10880 2023.
[60] Machel Reid, Vincent J Hellendoorn, and Graham Neubig. Diffuser: Discrete diffusion via edit-based reconstruction. arXiv preprint arXiv:2210. 16886, 2022.
[61] 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 18820-18836 PMLR, 2022.
[62] Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Björn Ommer. High- resolution image synthesis with latent diffusion models. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 10684-10695, 2022.
[63] Philip A Romero and Frances H Arnold Exploring protein fitness landscapes by directed evolution. Nature reviews Molecular cell biology, 10(12):866-87) 2009.
[64] Jeffrey A Ruffolo and Jeffrey Gray. Fast, accurate antibody structure prediction from deep learning massive set of natural antibodies. Biophysical Journal, 121(3):155a-156) 2022.
[65] Nikolay Savinov, Junyoung Chung, Mikolaj Binkowski, Erich Elsen, and Aaron van den Oord. Step-unrolled denoising autoencoders for text generation, 2021.
[66] Harshay Shah, Prateek Jain, and Praneeth Netrapalli. Do input gradients highlight discriminative features? Advances in Neural Information Processing Systems, 34:2046 2e-29, 2021.
[67] Amir Shanehsazzadeh, Sharrol Bachas, Matt McPartlon, George Kasun, John M Sutton, An- drea 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.
[68] Richard W Shuai, Jeffrey A Ruffolo, and Jeffrey J Gray. Generative language modeling for antibody design. bioRxiv, 2021.
[69] Karen Simonyan, Andrea Vedaldi, and Andrew Zisserman. Deep inside convolutional networks: Visualising image classification models and saliency maps. arXiv preprint arXiv:1312.6034, 2013.
[70] 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 2rX0vv2010002241, 2020.
14

---

[71] Raghav Singhal, Mark Goldstein, and Rajesh Ranganath. Where to diffuse, how to dif- fuse, and how get back: Automated learning for multivariate diffusions. arXiv preprint arXiv:2302.07261, 2023.
[72] 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.
[73] 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.
[74] 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:221..02236, 2022.
[75] Baris E Suzek, Hongzhan Huang, Peter McGarvey, Raja Mazumder, and Cathy H Wu. Uniref: comprehensive and non-redundant uniprot reference clusters. Bioinformatics, 23(10): 1288, 2007.
[76] Christian Szegedy, Vincent Vanhoucke, Sergey loffe, Jon Shlens, and Zbigniew Wojna. Re- thinking the inception architecture for computer vision. In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 281 188-2826, 2016.
[77] 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 motff mofif-scaffolding problem. arXiv preprint arXiv:2206. 04119, 2022.
[78] 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.
[79] Clement Vignac, Igor Krawezuk, Antoine Siraudin, Bohan Wang, Volkan Cevher, and Pas- cal Frossard. Digress: Discrete denoising diffusion for graph generation. arXiv preprint arXiv:2209.1444 14734, 2022.
[80] Joseph L Watson, David Juergens, Nathaniel R Bennett, Brian Trippe, Jason Yim, Helen E Eisenach, Woody Ahem, Andrew 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 122 2022.
[81] Andrew G Wilson and Pavel Izmailov. Bayesian deep learning and probabilistic perspective of generalization. Advances neural information processing systems, 47:48,9-47008 2020.
[82] James T Wilson, Riccardo Moriconi, Frank Hutter, and Marc Peter Deisenroth. The reparame- terization trick for acquisition functions. arXiv preprint arXiv:1712.00424 2017.
[83] Kevin Yang and Dan Klein. Fudge: Controlled text generation with future discriminators. arXiv preprint 2rXiv:2 18, 2021.
15

---

Appendix
Table of Contents
Extended Background
A.1 Continuous noise diffusion
16
A.2 Categorical noise diffusion
17
B Methodological Details
17
B.I Infilling algorithm
17
B.2 Hidden State Langevin Sampling
18
C
Infilling/ / NOS Guidance
19
C.I Infilling experiment
19
C.2 MCMC comparison
19
C,3 PPLM details
19
C.4 Model Architecture and Training
20
Hyperparameter settings
................................
20
C.6 Density plots
22
LaMBO-2
23
D.1 Intro to Multi-Ohjective Bayesian Optimization
23
D.2 Diserete EHVI
24
D.3 Architecture and Hyperparameters
24
D.4 Training Data, Class Imbalance, and Label Smoothing
25
D.5 Baselining LaMBO-2 Against Unguided Sequence and Structure Based Diversifi- cation:
26
D.6 Are Saliency Maps Reliable?
27
D.7 Wetlab Validation
28
Extended Background
In this section we provide full descriptions the diffusion processes introduced Sec. 3.
A.1 Continuous noise diffusion
The forward process is defined by noise variances 3. We use the cosine variance schedule from Nichol and Dhariwal [55]. For convenience we further define
0₂2 1- -B,, Oe
[ ac
The forward process is defined by the conditional distributions
p(z+(x=11) =
Nizx 11- 1111-118111
p( [x)) = Niza vārzo, (1
p(x.ww) = Nizi; Vo, T)
16

---

This is sample string
Forward process
Reverse process
This is [MASK] string
This [MASK] a [MASK] [MASK]
[MASK] [MASK] [MASK] [MASK] [MASK] [MASK] [MASK]
Figure 8: Illustration of string gradually corrupted by [MASK] tokens.
where Us an embedding matrix.
The reverse process is defined by
(②)) (),
p p(1111 1e, 20) N (31-1i Me, 07
v Vot-1 1Bt En
v vas(1 ā
14
Ft
- αe
Oit-1 Be
of
1 Ōz
po (w|xx (x)
Softmax (do())o)
Po ( 1/))
Σ p(x 12,, to (000))o((ū(zz))
A.2
Categorical noise diffusion
Following Austin et al. [4] we define the MLM style categorical diffusion using transition matrices
if i= m if j j= i m [Q = a2 - aq if i= # m
and Q+ QiQ2--Qu for noise schedule 04 [0, (see Figure 8 for illustration). These transition
matrices correspond to categorical conditional distributions
((w,/wr-1 Cat(wa; wt-1Qi) p(welwo) Cat(wr; w)Q
The reverse process is defined by
x (w) = 1[W [MASK]
w.Q.T O wa QT
p( We, wa)
Cat W wo wow
po (Wol we)
Softmax( (oftmax((o((x:)
Σ
p( (. (00-1/w0e wn Do (100 (w))
B Methodological Details
B.1 Infilling algorithm
We sample infills using the procedure in Algorithm 1. The infill mask constructed by setting the index of conserved residue equal to in this case at every residue that is not included in set of CDR regions being infilled. We use the same algorithm perform the guided infilling Subsec. 5.2, where it is extended with guidance Langevin sampling step.
17

---

Algorithm/ Infilling with categorical denoising diffusion model
Inputs: Denoiser pe (1 Iz, t), corruption process p(zz) infilling mask P. and seed sequence
Returns: Sample from ((w)= = po(w] P. ex)(f(w))
Jr p(rr)
ST ~ p(sr(s)
JT - (I PT Plat + PT 8T
for t= T, ..... do
p(=;-1/04) - 2 pl Fe pe,, w (
2t ~ p ((t-1((tt)
p(((s)
It - (I- (I-PT)),; + PT..
end
w po(w) return
B.2 Hidden State Langevin Sampling
Design of molecules or images with generative models often posed as the problem of sampling from posterior distribution p(x|a) given the unconditional distribution p(z) and attribute model p(a|z). Indeed, reinforcement learning, the design of good actions in an environment, can also be framed as posterior sampling where p(a[z) is the probability that given state state-action pair is optimal [46]. Methods that employ posterior sampling of this form are often call "plug-and- and-play" because p(a[2=) and p(x:) need not share parameters and therefore users mix and match different instantiations [54, 17, 34, 26]
The most common way to sample from the posterior p(x|a) x p(a(())(()) is through Langevin sampling on the unnormalized joint density p(a, x)) = p(a))) with sampling steps
x 1/ log + 2m22 z' 1o
x 7j V logp(ala) + V Vlog(()) 2np2, ~'~N(0,1)
When we work with generative models over continuous random variables that permit likelihood (e.g. normalizing flows), score function (e.g. diffusions), or energy (e.g. EBMs) V p(r) has natural interpretation and sampling can be performed with essentially vanilla Langevin sampling. In other cases where only 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 the score function V log p(x:) is challenging, and prior work adopts different approach of regularizing the conditional sampling distribution p( w a)) with unconditional sampling pl w)) in order maintain high likelihood [17]. In autoregressive models, p(w) is broken down using the chain rule, (x) and thus the appropriate regularization is
(L (Wee)) || p(w//wer a))
In our case, the distribution p(w) is factorized by the transition distributions (w)] (or their continuous analogies in token embedding space), and we hope to sample from the perturbed transition
pl ((w-i/ww = po (ww exp(vo(we))
The correct regularization term in our case is thus
(L(p(w,-llul || p(we-rlwe,,))
To put the pieces together, we first recognize that the denoising model pe( (ww (e) can be broken down into an language model head, Hg, and trunk, To. with
h4 To(we)
pe wol w.) = Ho(wplhh))
We can then perform Langevin sampling on the hidden representations, initializing with ht, as shown
in Algorithm 2. In the experiments above set Xa = O, as we saw no noticable benefit from adding additional stochasticity. Importantly, sampling from pl (ww)l/ww) already introduces randomness into the reverse process.
18

---

Algorithm 2 Guided diffusion sampler
Inputs: Denoiser po( Ir, t) [Te, Mo], value function va, and weights \,, 22, A3 Sample from ((w)= = p(w) exp(f(w)
Returns:
= [MASK]
for 1 T ...... do
pl m(w)- 1/(a)) - Ee p(u w ti (w)lw))
hD - To(w:)
for O,..., K 1-1do do
2' N(0, 1)
Ph - Ls p( (0t-1/00 w Hg
hit1 - ht' d1 Vi (h²) 22 V KV - ( + 232
end
Wt-1 Ho(hk)
end
return w0
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 without guidance, generating 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 MCMC comparison
Following Verkuil et al. [78], we construct Markov chain using uniform random mutations to map sequence to mutated sequence w', using the following Hatropolis-Hastinge correction:
3.8
pl placcept w'lu w'(w) min (1, (xpx(((((()/))) exp(-E(w)/TT)
3.4
where > 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.
0
5000 10000 16000 20000 Sampling Steps
Figure 9: (left) Comparing convergence in sampling us- ing a Metropolis Hastings- adjusted MCMC [78] against NOS models. Diffusion mod- els (ours) accelerate sampling by two orders of magnitude while converging similar en- ergy values.
In our experiment (Figure 9), define the energy, E, by combining sequence level probabilities assigned by IgLM with a beta sheets objective function trained IgLM's representations. We construct the energy as
E(w) (w) Xug(w),
We tune À to generate sequences with approximately 40% beta sheets. 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 '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
19

---

simple one-layer feed-forward neural network trained on top of the mean- -gan-agggeeated representations for the corresponding chain identity.
To sample using PPLM, we overwrite the forward pass of the huggingface decoder used by IgLM to include Langevin sampling step the current hidden representations, We perform K gradient steps to update the current hidden representation h' by descending on the objective
AKL(p(((h)) p( - v(h')
where the original hidden representation output by the model's encoder, and 17 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
0.001, 0.01, 0.1, 1.0
0.5, 0.8, 1.1, 1.4, 1.7.2.
K
5, 10 SGD, AdaGrad
optimizer
Table 1: Hyperparameter settings used for PPLM. À controls the strength of the regularization. Large values prevent sampling values that differ significantly from the unguided model. 1) controls the size of steps taken in the latent space. Larger step sizes, when not too large, can increase the distance traveled in the latent space and the extent to which sampling can yield samples with high values of the objective.
One critical difference between controllable autoregressive models and controllable diffusions the ability to resample previously sampled values. Procedures that allow for resampling often called "iterative refinement" procedures because they can produce increasingly plausible generations by refining the model's previous output each step in iterative procedure. Because there many potential differences between our NOS models and PPLM, including but not limited to the nature of iterative refinement, performed additional experiment assess the impact of adapting discrete diffusion to perform autoregressive sampling. Autoregressive models can themselves be thought of diffusions with an idiosyncratic corruption process that masks out all tokens 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, by contriving the sampling noise schedule to be autoregressive and recover an approximation of the timestep post-hoe from the percentage of masks at each step in autoregressive sampling.
trained discrete diffusions in autoregressive mode
Figure 10 shows the difference in objective values and likelihood for samples obtained by running the model in typical diffusion mode (iterative refinement) in contrived autoregressive mode. We can see that on the beta sheets objective, iterative refinement has noticeable positive impact on the objective values of the sample. This effect is also present in the SASA objective, but to 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 intro- duced by Bhargava al. [8]. We use cosine noise schedule for both diffusions and train for 100 epochs with a batch size of 64, optimizing with Adam using an initial learning rate of 5e-3 with 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.
20

---

% Beta Sheets
12000
0.25
SASA
0.20
11000
0.15
-500
-400
-500
-450
-400
-350
ProtGPT Log Likelihood
ProtGPT Log Likelihood
Autoregressive
Iterative Refinement
Figure 10: We compare samples from running our guided discrete diffusion (NOS-D) with diffusion style sampling versus autoregressive style sampling. We find that using an iterative refnement procedure does lead consistent improvements in the objective value, though not to an extent that would suggest iterative refinement is sufficient for strong sampling performance.
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 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, 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 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 A, 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 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 continuous variable and construct perturbation distribution from learned discriminative model è vo(w,),
Po po(0|uz (()) a exp(- V ar, vo(w:), wx-1>)
We then sample the next value from the base diffusion transition Po (ww)) perturbed with 1o)
W (p / (0
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 X, the perturbation p( we (6) collapses to a one-hot on the token index with the largest gradient value. For small values of X, p[w,]e (6) becomes uniform distribution. Therefore À interpolates pl p(w,11 [wy) between the original unguided distribution and a one-hot in the max gradient direction. NOS also reduces 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 7/. Digress amounts single update step applied directly to the output token probabilities using a continuous relaxation the one-hot encoded input, whereas NOS performs sequence of local updates hidden states that actually continuous.
In our comparison, the embeddings and corruptions of each model are chosen to be:
1. NOS-C [Gaussian corruptions Iearned embeddings]
21

---

Guidinnss Layer frat
13000
12000
0
18000
SASA
12000
13000
n 10
12000
11000
-600
-400
-500
-400
-600
-400
-600
-400
-600
-400
ProtCPT Likelhood
NOS-D (Ours)
NOS-C (Ours)
Figure 11: Density plots for every combination of the regularization (A) and step-size (n) parameter, when performing guidance in the first layer (token embeddings) of the neural network denoiser. We observe that lambda has the strongest effect on trading off fitness under the objective with likelihood or closeness to the seed sequences.
2. NOS-D [Discrete (mask) corruptions learned embeddings]
3. DiGress [Discrete (mask) corruptions + fixed one-hot encodings]
All models use the same backbone transformer and regression heads, facilitating an apples-to- apples comparison. For DiGress, we perform sampling for large range of scaling values À {le5, 3e4, le4, 3e3, 3e2, le2, 3e1, Iel, le0, le- le- le Le-3). For each model, À modulates the degree to which the model prefers greedy sampling from the value function gradient.
0.001, 0.01, 0.1, 1.0, 10.0
0.1, 0.5, 1.0
5, 10
guidance layer optimizer
first, last
SGD, AdaGrad
Table 2: NOS guided sampling hyperparameter settings. À controls the regularization strength, con- straining the plausibility of samples. 7). when chosen effectively, can effect the degree of optimization that takes place on the hidden states. The guidance layer is the layer in the neural network over which guidance is applied, the first being the token embeddings and the last being the final representations before the linear head. The same values used for both NOS-D and NOS-C.
À
0, 0.001, 0.01, 0.1, 1.0, 10.0
7
1.0
K
10
optimizer
AdaGrad
Table 3: Hyperparameter settings used in Sec. The guidance layer for NOS-D is final, and the guidance layer for NOS-C is last.
C.6 Density plots
Because pareto fronts present only 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
22

---

Guidenca Layer + lest
12000
0=
10000
12000
10000
12000
11000
10000
-500
-400
-600
-400
-500
-400
-600
-400
-630
-400
ProtGPT Log Likelhood
NOS-D (Ours)
NOS-C (Ours)
Figure 12: Density plots for every combination of the regularization (A) and step-size (n) parameter, when performing guidance in the last layer (pre-logits layer) of the neural network denoiser. NOS-C and NOS-D) exhibit quite different performance as function of guiding the first or final hidden representation.
RFDiffusion
DiftAb
IgLM
NOS-D (Ours)
NOS-C (Ours)
% Beta Sheels
0.3
0.2
800
SASA
12000
11000
-600
-400
-600
-400
-600
-400
-600
-400
-600
-400
ProIGPT Log Likelihood
RFDIffusion
DifIAb
PPLM
NOS-D (Ours)
NOS-C (Ours)
X
Seeds
Figure 13: We compare sample densities for the methods presenting in Sec. 5, in order to augment the limitations simply showing pareto fronts. We see that NOS-C and NOS-D can both consistently generate samples with favorable trade-offs while other methods tend to radically decrease likelihood with little benefit to the value function or be relatively limited to the neighborhood around the seed sequences.
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 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 & objectives, f X - k The goal of multi-objective optimization
23

---

Algorithm 3 LaMBO-2: one guided discrete diffusion step
Inputs: Seed sequence wi, edit budget projection diffusion timestep corruption function c(w, 8),
constraint function u( (()), encoder gol ww), value function Vy (h). decoder dg (h), regularization strength
X, SGLD step-size 1 and temperature т.
Returns: Best feasible sample from SGLD chain with distribution p' (x)x p(x) exp(f g(x))
w* v* Wo. 0g g(wp)
(initialize optimal solution)
wo c(wo: 1)
(apply diffusion noise)
ln go(w/)
(initialize hidden state)
for i 1,.... do
loss AKT (dg (lld
(1g-1) 1/4 11-1- P + 22TT)), ~ 1),
(projected SGLD step) (decode hidden state)
wi (h()
if u* < uo (wi) & u(wx) then
wº w/
0* - Ug o go(ws)
end
end
return wº 0
(MOO) to identify the set Pareto-optimal (i.e. non-dominated) solutions such that improving one objective within the set leads worsening another. We say that dominates x', or f(x) f(x'),
(x) > fx')) for all j E (1,.... m) and fz(x) > [x') for some j. The set non-dominated
solutions 2 is defined in terms of the Pareto frontier (PF) P".
x (x: f(x) E )*), where P = (f(x) x X, 7x s... f(x') > [(x))
(7)
MOO algorithms typically aim identify finite approximation to 2n (which may be infinitely large), within a reasonable number iterations. One way to measure the quality of an approximate PF P is to compute the hypervolume HV(P|rzzt) of the polytope bounded by P u (Fret), where
Fref E I"' a user-specified reference point.
MEHVI((, J,D) HVT(P', P(rret [HV(P') (raor) HV((|rrer))++
(8)
where P = P (f(x) [27, 28, 18). To decide where to query / next, we search for argmax, A[uEnvI(X, J)D)], where the expectation is w.r.t. p(J[D).
D.2 Discrete EHVI
Although expression yield and binding affinity are both continuous measurements, we chose to discretize them and model them classification with softmax likelihood (See Appendix D.4). As result we needed an extension of EHVI for discrete outcomes. Informally, EHVI is simply computing the HVI for different realizations of f and marginalizing using f( [D). Instead of taking f to be the latent function of some regression y = fww) + E. E ~ N00, o ) we instead take f to be the logits of categorical distribution, p(y ilw, D) softmax, (f ((f(Dd
Let [y yil Given set of baseline points B = A define P (Eq, 8) using the posterior mean (w)) = E|y w. D], w E B. We model yy...., 2/ as conditionally independent given some shared hidden state h ga(w), so p(y] h, D) factorizes nicely. Finally we define p' P o {y) and take the expectation of Eq. w.r.t p(y[h, D). Since p(y h, D) discrete and factorizes, we can x K2 is not too large, where K; is the number of classes
marginalize in closed form when K1 corresponding to the discretization of the original continuous ft.
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 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 single guided diffusion step with LaMBO-2.
Model Architecture: our architecture for this experiment is inspired by the one proposed by Stanton al. [73]. In particular we jointly a train an encoder shared between 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 ID CNN residual blocks (kernel width 9), with layer normalization and sinusoidal position embeddings. The shared encoder was comprised of 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 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 discrimi- native model, as well as true multi-task model, which is necessary since measurements for various protein properties are often missing from a substantial fraction of rows in world datasets. We trained for 500K gradient updates using the Adam optimizer with At each gradient step we randomly sampled 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.
1e-3, Bu = 0.99, 31 = 0.999,
Generation Hyperparameters: to generate the designs in Figure 7, we sampled IK 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 B 16. In this experiment each infilling method took 16 diffusion steps, using an inverse linear noise schedule ã. = 1/(1 + t). 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 r 1e-2. For guided infills with uniformly distributed edit positions we set r 1e6. For guided infills with saliency-informe edit position selection we set r = 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 meet certain sequence liabilities constraints:
Canonical Cysteine Conservation: there are specific conserved cysteine residues in anti- body sequences which play 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 protein that serves as a recognition site for the attachment of sugar molecules. The presence of a glycosylation motif in protein can affect its stability, solubility, activity, and function. The addition of sugar molecules can alter the protein's 's conformation, change its interactions with other proteins 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 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 non- non-binding label. This augmentation is important for
25

---

quantized pkD
1 2.3
3455 6
7.
10
1000
0.6
Dersity
600
0.2
measured PRD
quantized pRD
Figure 14: An illustration of using quantization to address heavily imbalanced data. On the right we show the original marginal label distribution in green, and the discretization boundaries as dotted lines. The boundaries are defined by a minimal level of affinity to be considered a binder (pKD 4), and pKD deciles computed from the remaining measurements.
training pan-target affinity model, since experimental measurements of affinity to off tafgeta 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 very encouraging to see positive real-world results before scaling in earnest.
Label Discretization. As noted above, biological data tends 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 aanbbog pair with log (KD) (pKD) less than 4 was assigned to the non-binding class O. Then binders were assigned to classes 1 - 10 based on this scheme is
which pKD decile (computed from binders only) they resided in. One consequence increasing any objective value by one unit corresponds to moving up one decile in the empirical label distribution.
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 the same noise schedule as the diffusion model and the usual cross-entropy loss, using modified labels
Training Discriminators on Noisy Inputs:
[76]. We train the discriminators with
y, = a * y + (1 a,,/K
K 1,
where y is the one-hot encoded label and K is the number of classes. Informally, as at - the discriminator reverts to a uniform prior since the inputs are not distinguishable. Training on corrupted inputs avoids evaluating the value gradient on distribbut 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
26

---

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.6
0.100
RFD wio antigen
Density
40
Density
0.4
Density
000075
antigen
0.050
20
0.025
0.0
0.000
0.40
0.45 Naturalness
0.50
0.55
0.60
-50
-40
-30
-20 -10
-50
-40
-30
-20 -10
Acquisition Value
Acquisition Value
Figure 15: (Ieft) we find that structure-based infills, particularly from DiffAb, tend to score consis- tently well on naturalness. Guided infilling produces a much wider range of scores, but the mode is very close to that of RFDiffusion. (middle) as assessed by the same model used to guide towards higher yield and binding affinity. The guided infills have very high acquisition value, since they were explicitly optimized for that outcome. Given 1024 samples each, DiffAb failed produce any sequences of higher expected value than the seed, and RFDiffusion produced only 7 marginally improved designs. We also took the opportunity to assess the sensitivity of RFDiffusion the antigen by comparing infills generated using the antibody stucture only (right). While the effect is not large, antigen information does produce a small shift the distribution of acquisition values to the right.
Sequence Diversífication This in silico evaluation compares two variants 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 WIS was given manually selected edit positions restricted to CDRs only. In the left two panels of Figure 7 compare the predicted expression yield. predicted binding affinity, and naturalness the antibody designs, using the metric proposed by. Comparing the Pareto frontiers obtained from each set of designs, that while WJS excels at generating "natural" antibodics, 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 balance between the two extremes.
(Ours)
(Ours)
Pred. Expression
10.0
Pred. Affinity
7.5
5.0
x
x
X
2.5
x
0.0
0.40
0.45
0.50
0.55
0.45
0.50
0.55
Naturalness
Naturalness
Figure 16: We evaluate LaMBO-2 the context of real-world antibody lead optimization. LaMBO-2 can use either NOS-C or NOS-D to generate design libraries with higher predicted objective value than the unguided sampling baseline WJS [30], however intensive optimization comes at the cost of reduced naturalness (panels left and center).
D.6 Are Salieney Maps Reliable?
There is substantial controversy regarding the reliability of coput-gradient-beee based feature attribution methods, specifically related to their ability to consistently highlight ground truth task mineiiscriieeeevee features and ignore irrelevant features. For example, Hooker et al. [41] claim that random attribution competitive with input-gradient methods, and Casper et al. [11] claim that gradient-fre attribution outperforms input-gradient competitors. On the other hand, many papers claim that specific types
27

---

Gaussian Input Noise
Input Masking
2
10
0 DIONTORFEOL SORPOOOOR TOOS
Figure 17: Binding affinity feature attributions for hu4D5 produced by independent models trained with different input corruptions. While the attributions do not match exactly, there is substantial agreement on the importance of CDRH3 (top panel) and CDRLI. Some importance is also assigned to various framework regions, which could be related to the fitness of different antihody germlines. We emphasize that these models were trained solely on aligned sequences, with no additional positional information.
of regularization can improve the performance of -rautigt 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 the offline setting. LaMBO-2 feature attributions are used online to intervene on the data collection process (specifically where to introduce changes the antibody sequences). LaMBO-2 changes 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 exciting direction for future work.
D.7 Wetlab Validation
0.25
ye46 (1900
yeedd (mg)
0.70
0.0
0.15
100
11.0
2.0
7.0
Edit Distance from Seed
Edit Distance from Seed
8.5
3eg100K01N0 (M)
8.0
7.8
-0.5
2.0
-100
6.0
-2.0
Edit Distance from Seed
Edit Distance from Seed
Figure 18: Here we show the experimentally validated yield (top) for all expressing designs and affinity (bottom) for all binding designs as a function of edit-distance from the original seed. In the right column we show the absolute measurement, and the left column shows the change relative to seed measurement in the same batch.
28

---

In this section we briefty 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 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 trans- fected transiently in ImL 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 ffinity Measurements: affinity of the antibodies towards their target antigen was measured by surface plasmon resonance (SPR) 37 eC Biacore 8K instrument (Cytiva) in HBS-EP+ buffer (10 mM Hepes, pH 7.4, 150 mM NaCI, 0.3mM EDTA and 0.05% vol/vol Surfactant P20). Antibodies were captured on a Protein chip and their target antigen were injected for minutes and allowed to dissociate for 10 minutes at 30ul/min. The surface was regenerated between cycles with 10 mM glycine pH I.5. Affinity constants obtained using Biacore Insight (Cytiva) using a 1:1 binding kinetics model.
29


Reviews:
[{"review_id": "3fid8AX1Yn", "reviewer": "Reviewer_UJds", "summary": "The 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.", "questions": "1. 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", "limitations": "The 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).\n", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 1, "contribution": 3, "strengths": "Novel 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.", "weaknesses": "Unclear 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", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the detailed feedback. \n\nFirst of all, we would like to push back on the assertion that our paper has inadequate description of models or training procedures or lack of reproducibility. We encourage the reviewer to take a look at the paper appendices and follow the anonymous repo link (https://anonymous.4open.science/r/protein-seq-diffusion-0431/README.md) included on page 2 of the submission.\n\nAnswers to questions:\n1. Currently protein design lacks methods that satisfy the following criteria simultaneously (a) reliable optimization of a goal objective function (b) generated sequences that are likely under the distribution of real proteins (naturalness) (c) fast sampling. Notably, hill-climbing methods often struggle to produce natural sequences, while popular sequence-based methods like MCMC can be slow to converge and computationally expensive. This gap in functionality is the fundamental motivation for our work. While we find this writing feedback helpful, and we have added additional descriptions of our motivation, as you can see in the top-level rebuttal comment, we also think a conference paper can never provide a truly comprehensive introduction to a subject, and this is a task better left to review paper. For protein design, there are many great existing review papers [2,3,4], which discuss the tradeoffs between sequence-based and structural methods and provide extensive context on the challenges of creating design methods that sample from a desired conditional distribution while maintaining speed and trading off exploration-exploitation, as our method is designed to.\n2. The training process for our method converges in about 100 epochs, which takes about 7 hours on one RTX8000 GPU. Sampling from this model takes a few minutes or less on one GPU for both corruption processes. This level of compute is accessible to most academic and industry labs, and therefore there’s no fundamental barrier to others using or building on our method.\n3. Please refer to Appendix E.3 in the original submission, which contains all the requested details about model architecture, as well all necessary details of training and sampling, including an example algorithm and the hyperparameters used.\n4. The precise test seeds were released in the submission, and can be viewed at the following link (https://anonymous.4open.science/r/protein-seq-diffusion-0431/poas_seeds.txt). pOAS is a dataset of around 120,000 paired antibody sequences. For a more detailed description you can also see the following link: https://opig.stats.ox.ac.uk/webapps/oas/documentation_paired. Our test seeds were chosen randomly from pOAS and were removed from the training and validation sets. \n5. We appreciate your detailed writing feedback, but we see stated criticisms as fairly subjective, and not related to a fundamental lack of clarity or informativeness in the figure captions. If there is missing information that a reader would need for a full understanding of the presented number (for example the meaning of error bars), please point those out in your response, and we will be sure to address them so that there is no room for misinterpretation of the results.\n6. IgFold reports predicted RMSD (pRMSD) values for its folded structures. For our 10 seed pairs of sequence and predicted structure, the overall average per-residue pRMSD was 0.34. We also include the average and maximum pRMSD values for each individual seed in the table below. The framework typically has a very low pRMSD while the maximum value is located in CDR H3. These values are around the average or well below average compared to values in the IgFold paper [1], and therefore we don’t have any concerns about using these predicted structures as starting points. \n| Seed | 1 | 2 | 3 | 4 | 5 | 6 | 7 | 8 | 9 | 10 | \n| ----------- | ----------- | ----------- | ----------- | ----------- | ----------- | ----------- | ----------- | ----------- | ----------- | ----------- |\n| Avg | 0.3 | 0.43 | 0.27 | 0.46 | 0.27 | 0.29 | 0.29 | 0.29 | 0.26 | 0.56 | \n| Max | 2.05 | 2.95 | 2.27 | 4.32 | 2.30 | 1.67 | 1.68 | 2.37 | 1.32 | 3.99 | \n\n7. The non-disclosure is the result of intellectual property concerns from our industry collaborator, not ethical or privacy concerns.\n8. By “searching directly in sequence space” we mean that we never have to produce a structure in the process of designing a protein sequence. The only time we use structure representations is when we fold sequences into structures as part of the evaluation step in our optimization experiments. \n9. One example is SabDab (a database of structures and sequences) versus pOAS (a database of sequences alone). SabDab has about 6,000 paired antibodies, while pOAS has about 120,000 paired sequences. If we expand the comparison to unpaired structures and sequences it is about 7,500 structures vs 3,549,291,485 sequences. Beyond antibodies, another example is the size of the PBD, which contains about 200,000 structure entries, versus the size of UniProt, which contains 248,000,000 entries. \n\nAs your concerns centered primarily on clarity, rather than methodological or experimental soundness of the paper, we hope that you will consider revising your score in response to our detailed answers.\n\n[1] Ruffolo, Jeffrey A., and Jeffrey J. Gray. \"Fast, accurate antibody structure prediction from deep learning on massive set of natural antibodies.\" Biophysical Journal 121.3 (2022): 155a-156a.\n[2] Hummer, Alissa M., Brennan Abanades, and Charlotte M. Deane. \"Advances in computational structure-based antibody design.\" Current Opinion in Structural Biology 74 (2022): 102379.\n[3] Tian, Xiaochen, et al. \"SEQUENCE VS. STRUCTURE: DELVING DEEP INTO DATA DRIVEN PROTEIN FUNCTION PREDICTION.\" bioRxiv (2023): 2023-04\n[4] Hie, Brian L., and Kevin K. Yang. \"Adaptive machine learning for protein engineering.\" Current opinion in structural biology 72 (2022): 145-152.\n\n"}, {"review_id": "MDE8LGud4Z", "reviewer": "Reviewer_j6TU", "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. ", "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.", "rating": 7, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 2, "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.\nQuality: 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.\nClarity: The paper is clearly written and contains all the necessary citations. Methods, experiments, and equations are clearly described. \nSignificance: 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. ", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your supportive review and for emphasizing the impact and novelty of our submission! We have included the requested comparison with RFDiffusion in our rebuttal PDF. We also include DiffAb and several other local sequence search methods (AdaLead, PEX) for the protein design task in our extended experiments. The experiments show that our method is dominant for practical applications of generative design, and we hope these results inspire your continued support for our submission."}, {"review_id": "azG17ueXyq", "reviewer": "Reviewer_N1y4", "summary": "This 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.", "questions": "1. 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.", "limitations": "The 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.", "rating": 6, "confidence": 5, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. 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.", "weaknesses": "1. 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].", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your detailed response and suggestions. \n\n**DiGress Comparison**\n\nWhile DiGress looks like a closely related and compelling method, its application appears to be limited to graph representations of small molecules, with a maximum of 9 nodes. As our proteins have 1000 nodes, adapting DiGress to our setting would be highly non-trivial and is an effort better to future work, and potentially its own entire paper. On the point of extending baseline comparisons, however, we would like to point out the extended experiments included in our rebuttal PDFs. In those experiments, we compare our method with AdaLead and PEX, which are two popular methods targeted specifically at protein design. We also provide more comparisons with competing diffusion methods (DiffAb and RFDiffusion) that were missing from the earlier protein design experiments. \n\nTo answer your itemized questions\n1. Integrating selection of the edit locations into the diffusion process is a very interesting idea. The challenge with selecting positions at every iteration would be maintaining the constraint that N residues must be conserved (remain unchanged). If new locations could be chosen at every point in the diffusion, the union of all the chosen locations could be much larger than the budget after several iterations. In an unconstrained protein design setting, this effect might be desirable, the design process could iterate over the entire protein throughout the diffusion process, but this freedom is not desirable in many practical applications. \n2. DiGress does not appear to be directly applicable to proteins, which tend not to be represented as graphs, and which are much larger than small molecules.\n3. The training process is performed on the entire sequence. Both the generative model and discriminative model are trained to take a corrupted version of the entire sequence as input. During sampling, corrupts are applied along the entire sequence as well, even when infilling only a small region. See Algorithm 1 and Figure 8 in the Appendix for a more detailed walkthrough of the procedure and a visualization. Applying noise along the entire length of the sequence ensures there is no distribution shift between the training procedure and any possible infilling size, but global consistency is not hurt by this procedure, as later steps in diffusion have very small corruptions, and the majority of sequence is passed to the model uncorrupted, informing the new amino acid values that are sampled in the infilling region. As an alternative to training with corruptions on the entire sequence, one can also sample masks, and corrupt only the masked region during training (e.g. [1]). The perk of our method as compared to the limited masking approach is that our models can be used with any masking size or location, as opposed to being trained for a particular masking distribution. \n4. Thank you for catching this detail. D refers to the dataset and is a commonly used notation in Bayesian inference, where we frequently define posterior distributions over functions conditioned on a given dataset. We will clarify this in a camera-ready.\n5. You are correct that d is the dimension of the hidden representation $h_0$. This will also be clarified in a camera-ready. \n6. Both $C^n_k$ and $C^k_n$ are used and there is no universal standard. We will change the “choose” notation to binomial coefficient notation in order to remove any possible ambiguity.\n\nWe also thank you for the additional small suggestions listed under “Minor Concerns”.\n\n**Limitations**\n\nAs it is not possible to upload a revised draft of the paper, we include a new limitations section below, which will be included in a camera-ready version of the paper:\n\nLimitations: Methodologically, one of the most important limitations of the paper is that it cannot easily leverage pre-training on general protein families before fine-tuning on any particular protein family. This limitation is the result of our use of alignments to create fixed-length sequences. These alignments are applicable to families of proteins with a large number of sequences (hundreds of thousands to millions of sequences), but do extend easily across families. The solution to this limitation would be to extend our method to diffusions that generate variable-length sequences, for example by applying insertions and deletions during the corruption process. Experimentally, the most notable limitation of the current work is the fact that we only evaluate the design of antibodies. Our method is applicable to other common protein families, and other biological sequences, and we hope to demonstrate that generality in future work.\n\nWe hope that you will consider raising your score in response to answers and additional experiments. If you have any followup questions about our responses, let us know and we will provide further clarification. \n\n[1] Saharia, Chitwan, et al. \"Palette: Image-to-image diffusion models.\" ACM SIGGRAPH 2022 Conference Proceedings. 2022."}, {"review_id": "FxsqEDQSUM", "reviewer": "Reviewer_gzgm", "summary": "This 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.", "questions": "- 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.", "limitations": "N/A", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- The paper is well written.\n- The experimental results are convincing and good.\n- The function-guided protein design task is important.", "weaknesses": "- 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your review and suggestions for improved comparisons. \n\nADDITIONAL EXPERIMENTS\n\nWe have performed experiments with two of the papers you reference, and the results are quite positive in favor of our method. The detailed results are contained in the rebuttal PDF, but the findings can be summarized as follows. When given the same number of optimization steps, LaMBO-2 can achieve values of the objective function that are uniformly better than PEX, but slightly worse than AdaLead. If we examine the types of solutions these three algorithms find, however, we see the LaMBO-2’s samples are much more natural, and therefore likely to be useful in practice. When calibrated against this consideration, LaMBO-2’s samples achieve much higher values than samples from PEX or AdaLead with comparable naturalness. We also find that LaMBO-2 adds a minor computational overhead for these significant gains in sample quality. We encourage you to examine the detailed explanation in the attached PDF. \n\nIN-VITRO AND MULTI-OBJECTIVE RESULTS\n\nWe would like to re-emphasize that our work already includes both in vitro experiments (Figure 7 right) and multi-objective evaluations (Figure 7 left), in which we optimize for both expression yield and binding affinity. Please refer to Section 5.4 and Appendix E5, where we have detailed explanations of the multi-objective evaluation and in vitro experiments. Indeed, our method is specifically targeted at both multi-objective use cases and practical considerations for in vitro protein design, as described in Section 4.2. These considerations are also described at length in Appendix E1, where we provide an introduction to multi-objective Bayesian optimization, and Appendix E4, which includes a detailed description of the training data and methods designed specifically for the in vitro trials.\n\nADDITIONAL BASELINES AND CONCURRENT WORK\n\nWhile Kirjner et al also looks like a relevant baseline for this task, it was made public after the NeurIPS submission deadline. For the purposes of this conference, we don’t think it’s reasonable to expect a comparison, and we leave it for future work. We also leave a comparison with Brookes et al. (CbAS) for future work but note that evaluations in the PEX suggest it is a much stronger baseline than CbAS on all tasks considered in the PEX paper.\n\nADDENDUM: EXPERIMENT DETAILS FOR FIGURE 2\n\nHere are some additional experimental details as promised in the PDF\nFor LaMBO-2 we took a sampling step to be one diffusion step, each of which consisted of 16 gradient updates.\nPEX and AdaLead perform local search in sequence space on a learned approximation of the goal objective, akin to genetic optimization with a learned value function.\nFor AdaLead a sampling step is one survival and recombination event, followed by 32 x 16 = 512 model evaluations on random mutants of the recombinants.\nLikewise a sampling step for PEX is a contraction of the population to the proximal frontier, followed by 512 evaluations of random mutants along the frontier.\nFor both genetic methods we used the authors' recommended hyperparameters. \n"}, {"review_id": "8uOSOzbpR1", "reviewer": "Reviewer_EUPz", "summary": "This 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.", "questions": "Please 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.", "limitations": "The authors have addresses limitations and proposed directions for future work.", "rating": 5, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. 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.", "weaknesses": "1. 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?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your review. We have added DiffAb and RFDiffusion comparisons to the antibody lead optimization task in the rebuttal PDF. There we show that our method generates samples with much higher values of the multi-objective acquisition function than these competing diffusion methods. This result mirrors our finding on the percentage of beta sheets and solvent-accessible surface area (SASA) optimization tasks. \n\n**Technical Novelty**\n\nRegarding the technical novelty of our guidance method, we think there might be a small misunderstanding that should be clarified. In both the discrete and continuous diffusion, there is a neural network parameterizing a categorical distribution over tokens, $p(w | h_t)$. We bias the logits of this categorical distribution by taking gradient steps in the hidden units of the network. We are therefore always working with categorical distributions and guiding categorical distributions, though the corruptions are sometimes continuous and sometimes discrete. Performing guidance in this way is distinct from performing guidance directly over a continuous input space and converting this guided variable into discrete tokens. Guiding the categorical distributions directly avoids many of the common pitfalls of converting between continuous and discrete variables, for example the need for rounding continuous vectors to the nearest embedding value, as in [1]. \n\nAlthough we think there is plenty of novelty in our method of guiding categorical distributions, there is also plenty of additional novelty in the many other tools we introduce in our paper, including our method of using saliency maps for selecting edit positions, multi-objective Bayesian optimization as guided generative modeling, and partial deep ensembles for multi-objective surrogate models. All of these insights were necessary to achieve our strong in vitro results, in which we created real proteins with improved binding affinity.\n\nIn general we find the notion of novelty implied by the reviewer to be overly narrow in scope. To the best of our knowledge, we are first to perform in vitro experiments with a guided diffusion model operating directly on protein sequences. While discrete diffusion models have not seen major adoption in language modeling, they are poised to have a large impact in protein engineering, and our paper outlines precisely how to accomplish this end, with key details at every level of abstraction, from core machine learning methodology to task specific considerations. \n\n**Additional Experiments**\n\nAs the lack of DiffAb and RFDiffusion baselines was one of your major concerns about the submission, we hope you’ll consider raising your score as a response to our extended experiments. As you noted in your review, out empirical evaluations are quite extensive (and have only become more comprehensive with the extended results), and there is a mismatch between this description and the score rubric description, “reasons to reject, e.g., limited evaluation, outweigh reasons to accept”.\n\n[1] Li, Xiang, et al. \"Diffusion-lm improves controllable text generation.\" Advances in Neural Information Processing Systems 35 (2022): 4328-4343."}]

Author rebuttal:
We are glad to share many exciting new results inspired by reviewer comments in the attached PDF, including a new comparison with RFDiffusion and DiffAb on antibody lead optimization (Figure 1) and two additional genetic algorithm baselines mentioned by gzgm (Figure 2). In both of these experiments, we show that our proposed methods (NOS/LaMBO-2) generate better samples than competing methods by optimizing a goal objective without falling off the manifold of plausible sequences.

Our new results further solidify what is already a significant and timely contribution to the field of ML-guided protein design. Two of the baselines suggested by reviewers were ArXived after our submission date and should be considered concurrent work [4,5]. Structural diffusion methods have become a popular approach to generative protein design [1] but come with several drawbacks. Our novel method for guided sequence diffusion provides an alternative to structural diffusion that conserves many of its benefits (e.g. the potential for gradient-guided search) while avoiding several downsides (e.g.  limited structural data or the use of inverse folding models). Our method is also a compelling alternative to sequence hill-climbing [2], which can rapidly fall off-manifold, or MCMC methods [3], which can suffer from slow convergence and high computational overhead. Lastly, it is also noteworthy that among the many recent methods for ML-assisted protein design, we are one of only a handful that provide both open-source in silico evaluations and in vitro experiments.  

We would like to note that the majority of reviewers thought our methods were sound and sufficiently novel but wanted to see additional baselines or wanted clarifications on the paper’s text. We have directly addressed many of these concerns through our rebuttal answers and our additional experiments, and we hope that reviewers will be open to revisiting their scores in light of these updates.

REFERENCES

[1] Watson, Joseph L., et al. "Broadly applicable and accurate protein design by integrating structure prediction networks and diffusion generative models." BioRxiv (2022): 2022-12.

[2] Sinai, Sam, et al. "AdaLead: A simple and robust adaptive greedy search algorithm for sequence design." arXiv preprint arXiv:2010.02141 (2020).

[3] Hie, Brian L., et al. "Efficient evolution of human antibodies from general protein language models." Nature Biotechnology (2023).

[4] Kirjner, Andrew, et al. "Optimizing protein fitness using Gibbs sampling with Graph-based Smoothing." arXiv preprint arXiv:2307.00494 (2023).

[5] Vignac, Clement, et al. "DiGress: Discrete Denoising diffusion for graph generation." The Eleventh International Conference on Learning Representations. 2022.


Meta-review: The paper makes significant methodological and empirical contributions. The proposed approach of controllable discrete diffusion allows for a principled use of sequence data as structural data is scarce. The AC and reviewers feel that all key concerns have been satisfactorily addressed and this will strengthen the submission. We therefore urge the authors to incorporate the clarifying points and additional experiments (s.a. comparisons with DiGress) made during the discussion phase.
Decision: Accept (spotlight)