SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Nate Gruver*1 Samuel Stanton*2 Nathan Frey*2 Tim G. J. Rudner*1 Isidro Hotzel*3
Julien Lafrance-Vanasse<sup>3</sup> Arvind Rajpal<sup>3</sup> Kyunghyun Cho<sup>1,2</sup> Andrew Gordon Wilson<sup>1</sup>
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.
Optimizing protein sequences for improved function has the potential for widespread impact [63]. Among many potential applications in engineering and medicine, engineered antibodies can be used to create cancer therapeutics that are much less harmful to the patient than radiotherapy or chemotherapy. Because the space of possible proteins is vast and discrete, and experimental validation is slow and expensive, all practical methods for protein design must restrict themselves to a small enriched library of candidates to find a viable option in as few measurements as possible [44]. In practice these enriched libraries are usually obtained through massive high-throughput in vitro screening [67], or in the case of antibodies by injecting a live animal with the target antigen and sequencing the animal's immune response [52]. Generative protein models offer the tantalizing prospect of enriched libraries produced nearly instantly at a fraction of the cost. Success in real-world applications, however, has proven elusive, in part because naive generative models produce outputs that are similar to their training data and are therefore unlikely to improve target qualities [53].
There are many approaches to guided generation of proteins, but one broad and important distinction is between methods that search in sequence space and those that search in structure space. A basic tenet of molecular biology is "sequence determines structure, structure determines function" [9]. Hence when optimizing a protein for a desired function, it may seem more direct to design the protein in structure space, where gradient-based sampling methods can be used in tandem with carefully
*Equal contribution.  ${}^{1}$  New York University,  ${\left\{  {}^{2}\text{Prescient Design,}{}^{3}\text{Antibody Engineering}\right\}  }$  Genentech.
Figure 1: We propose diffusion-based Optimized Sampling (NOS), a method for gradient-guided sampling from discrete diffusion models. NOS uses  $T$  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 a result, many sequence search methods resort to gradient-free sampling methods like Metropolis-Hastings MCMC [78, 37], which are flexible but computationally expensive, eroding a key advantage over structure search. Several methods have been proposed that maintain gradient-based search by performing guidance in a continuous latent space, with a learned decoder to sample discrete sequences [33, 32]. Notably, Stanton et al. [73] proposed LaMBO (Latent Multi-Objective 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, it has two key limitations. First, LaMBO is built on top of MLMs which, while strong representation learners are not strong generative models. In particular, MLMs lag behind other methods in producing high likelihood samples or infills. Second, despite being designed to improve known sequences instead of designing them completely from scratch, LaMBO and related methods have no principled framework for simultaneously enforcing an edit budget and choosing optimal edit locations based on that budget.
To address the first issue we propose NOS (diffusionOptimized Sampling), a new method for controllable categorical diffusion (Figure 1). Diffusion models capture complex distributional dependencies by making iterative denoising steps, but there is relatively little previous work on how to control these processes. NOS generates sequences with both high likelihood and desirable qualities by taking many alternating steps between corruption, guidance, and denoising in the continuous latent space of the model. Our in silico validation shows that NOS outperforms many state-of-the-art structure and sequence-based baselines on both unguided and guided infilling tasks.2 To address the second problem (choosing optimal edit locations) we propose using embedding-gradient feature attributions (i.e. saliency maps) to determine which positions on the sequence are most important to edit to improve the guidance objective value. We combine NOS with saliency-selected edits to create LaMBO-2, a more powerful variant of the original LaMBO algorithm. Exploratory in vitro experimental validation of our designs provides evidence that LaMBO-2 can be used to create enriched antibody libraries without the aid of additional in vitro high-throughput screening.
$^{2}$ https://github.com/ngruver/NOS
Discrete diffusions Austin et al. [4] and Hoogeboom et al. [40] constructed diffusion models for discrete categorical data using a categorical noise process. Recently, categorical diffusion has shown promise as a competitor to autoregressive models in text generation for machine translation and summarization. The approaches can be roughly grouped into methods that apply categorical noise distributions directly to sequences (CMLM [31], SUNDAE [65]), and those that apply Gaussian corruptions to token-vector embeddings (SED [74], CDCD [21]). In this work we show that NOS can guide both types of categorical diffusions. Within the space of protein design, our method is also closely related to joint diffusion models over both sequence and structure [2, 51], which also circumvent inverse folding. Because these models still rely on structure information at training time, they can be limited by data availability in the same manner as pure structure models.
Discrete generative guidance Gradient guidance typically augments sampling from a generative model with gradient steps to increase the occurrence of a desired attribute [54]. Gradient guidance is natural within the framework of continuous diffusion models [20], and Li et al. [47] use this connection to perform gradient-guided sampling from a diffusion language model. To obtain a continuous space, they perform Gaussian diffusion [39] on word embeddings, decoding out to tokens using a linear head. The original method required many careful engineering interventions, e.g. clamping latent representations to the nearest word embedding, that have been improved by recent methods, such as CDCD [21], but gradient guidance has not been discussed for these more recent formulations.
To achieve a similar form of gradient guidance without carefully engineering a latent space, Dathathri et al. [17] and Yang and Klein [83] propose gradient-guided autoregressive models by using the decoder's activations as a gradient-friendly latent space. These methods alternate between sampling from logits and ascending the likelihood of a separately trained classifier model. Surprisingly, despite work on gradient guidance for continuous noise diffusions and autoregressive language models, there has been little work on gradient guidance for general categorical diffusions that predict denoised categorical distributions (e.g. CMLM, SUNDAE, CDCD), which is a topic we explore in detail. One closely related method proposed in the context of generative models of small molecules is DiGress [79], which performs gradient guidance on one-hot token embeddings to bias the categorical sampling distribution of a denoising model. In our setting we show that categorical and Gaussian discrete diffusions guided with NOS outperform PPLM and DiGress (Subsec. 5.2).
Genetic algorithms Evolutionary algorithms are a popular solution for black-box optimization in discrete spaces [3, 22]. These methods are often evaluated on their ability to optimize in silico proxy estimates of actual in vitro fitness, e.g. deep learning models trained on experimental datasets. In Subsec. 5.3 we baseline NOS against two genetic optimizers from protein design literature, AdaLead [70] and Proximal Exploration (PEX) [61]. We show these baselines rapidly degrade sequence likelihood as the proxy fitness is improved, limiting the effective number of edits that can be made before checking the actual fitness of the samples, ultimately limiting their sample efficiency and rate of convergence to optimal solutions. By contrast, NOS consistently improves proxy fitness while maintaining sequence likelihood.
We pose protein design as the problem of finding sequences,  $w \in \mathcal{A}^L$  with alphabet  $\mathcal{A}$  and fixed length  $L$ , which maximize a single objective  $f(w)$  (e.g., binding affinity) or multiple objectives  $f_1(w), \ldots, f_k(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 \in \{1, \ldots, L-1\}$  to a seed sequence  $s \in \mathcal{A}^L$ . A protein is only useful if it can be synthesized (i.e. expressed), and the objective value of non-expressing proteins is undefined since their properties cannot be measured. Therefore we must introduce the constraint  $w \in \mathcal{E} \subset \mathcal{A}^L$ , where  $\mathcal{E}$  is the set of all expressible proteins. Since naturally occurring sequences must express in order to be observed,  $p(w)$ , the likelihood of a protein with respect to an empirical distribution of natural protein sequences, is often taken as a proxy for the tendency of a protein to express. In protein design, these proxies are
3Length change is enabled by the use of protein sequence alignments, which introduce a gap token " -".
Figure 2: Two approaches to diffusion generative modeling for categorical variables. (Left) Categorical data is embedded into continuous variables with an accompanying continuous noise process. (Right) Categorical noise is applied directly to sequences, and corrupted sequences are denoised using standard language modeling methods.
typically called metrics of naturalness. Since we are looking for sequences that by definition have not yet been identified in nature, naturalness and our other objectives are often in tension.
We can trade off naturalness and objective value by drawing samples from the unnormalized density
$$
\tilde {p} (w) = \exp (- \tilde {E} (w)) / Z, \quad \tilde {E} (w) = E (w) - v (w),, \tag {1}
$$
where  $E(w) = -\log p(w)$  is a scalar energy function, and the value function  $v: \mathcal{A}^L \to \mathbb{R}$  expresses the utility of a sequence with respect to our objectives. When designing proteins from primary sequence, sampling efficiently from the resulting energy function can be challenging. Simple approaches, such as the MCMC sampler used by Verkuil et al. [78] can require hundreds of thousands of steps to converge (Appendix C.2). Guided diffusion models are an appealing alternative because they construct a fixed-length Markov chain that quickly generates low-energy samples.
Diffusion models Denoising diffusion models construct samples by reversing a diffusion process that maps clean data points,  $x_0$ , to samples from a prior  $\pi(x)$  (Figure 2). The forward process  $(x_0 \to x_T)$  is composed of conditional distributions  $p(x_t | x_{t-1})$  (i.e., the noise process) that admit closed forms for the conditional distributions  $p(x_t | x_0)$  and  $p(x_{t-1} | x_t, x_0)$  (e.g., independent Gaussian corruption). The reverse process  $(x_T \to x_0)$  converts samples from the prior into samples from the learned data distribution  $p_\theta(x_0)$  by repeatedly predicting the denoised variable  $\hat{x}_0$  from noisy values  $x_t$  and using the conditional distribution  $p(x_{t-1} | x_t, \hat{x}_0)$  to derive a transition distribution,  $p_\theta(x_{t-1} | x_t)$ . The specific choice of noise process has been shown to significantly affect the likelihood and quality of image samples [71]. For categorical data there are two common approaches to constructing a diffusion generative model, depending on the nature of the noise process. We include brief descriptions below and a more detailed account in Appendix A.
Continuous noise To learn a distribution  $p(w)$ , one strategy is to first embed  $w$  to a continuous variable  $x_0$  with embedding matrix  $U_{\theta}$  and apply Gaussian noise [21]. The prior is taken to be  $\pi(x) = \mathcal{N}(0, I)$  while the forward process is  $p(x_t | x_0) = \mathcal{N}(x_t; \sqrt{\bar{\alpha}_t} x_0, (1 - \bar{\alpha}_t)I)$  for  $\bar{\alpha}_t \in [0, 1]$ . The values of  $\bar{\alpha}_t$  are determined by a user-specified corruption schedule. For the reverse process, we learn a function,  $p_{\theta}(\hat{w} | x_t, t)$ , to predict the sequence from noised points  $x_t$  by minimizing the following loss with respect to  $\theta$ :
$$
L (\theta) = \mathbb {E} _ {w _ {0}, t} \left[ - \log p _ {\theta} (w _ {0} | x _ {t}) \right], x _ {t} \sim p (x _ {t} | x _ {0} = U _ {\theta} w _ {0}).
$$
Using  $p_{\theta}(\hat{w} |x_t,t)$  we can construct a distribution for the reverse process
$$
p _ {\theta} \left(x _ {t - 1} \mid x _ {t}\right) = \sum_ {\hat {w}} p \left(x _ {t - 1} \mid x _ {t}, \hat {x} _ {0} = U _ {\theta} \hat {w}\right) p (\hat {w} \mid x _ {t}, t), \tag {2}
$$
where  $p(x_{t - 1}|x_t,x_0)$  is also a Gaussian distribution. At inference time, we can use the learned reverse process to convert samples from  $\pi (x)$  into samples from the learned distribution  $p_{\theta}(x_0)$  by repeatedly sampling  $p_{\theta}(x_{t - 1}|x_t)$ , followed by sampling  $w\sim p_{\theta}(\hat{w} |x_0,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  $w^{(i)} = [\mathrm{MASK}]$ . The forward process  $(w_0 \to w_T)$  is defined by a discrete transition matrix, describing the probability of mutating a token into a [MASK], and the corresponding prior is simply a point mass on the sequence of all [MASK] tokens. To learn the parameters of the denoiser  $p_{\theta}(\hat{w}_0 | w_t, t)$  we maximize the likelihood of the denoising process on ground truth sequences
$$
L (\theta) = \mathbb {E} _ {w _ {0}, t} \left[ - \log p _ {\theta} (w _ {0} | w _ {t}) \right], w _ {t} \sim p (w _ {t} | w _ {0})
$$
Then, as above, we can use the denoiser to construct the reverse process
$$
p _ {\theta} \left(w _ {t - 1} \mid w _ {t}\right) = \sum_ {\hat {w} _ {0}} p \left(w _ {t - 1} \mid w _ {t}, \hat {w} _ {0}\right) p _ {\theta} \left(\hat {w} _ {0} \mid w _ {t}, t\right) \tag {3}
$$
where  $p(w_{t-1}|w_t, w_0)$  is also a categorical distribution derived using Bayes' rule. To sample, the transition distribution is applied for  $t = [T, \dots, 0]$ .
Now we present practical methods for efficiently sampling from  $\tilde{p}(w) \propto p(w) \exp(v(w))$  (Eq. 1) by modifying the learned transition distribution with a learned value function  $v_{\theta}(w)$ . We then show how this sampling method can be used to perform protein design through guided infilling in sequence space. As before, we provide the most salient information below and the full details in Appendix B.
4.1 NOS: diffusion Optimization Sampling
We introduce a general form of gradient guidance (NOS) for discrete diffusions with categorical denoising models, i.e. diffusion models that predict logits over the ground truth tokens (e.g. [21, 4]). The key challenge in applying gradient guidance to categorical data is simply the lack of a continuous representation. Fortunately, in any denoising network, e.g.  $p_{\theta}(\hat{w} |x_t,t)$ , the discrete sequence  $w_{t}$  has many corresponding continuous representations in the form of hidden states of the model  $h_t = g_d(w_t)$  for  $d\in \{0,\dots ,D\}$ , where  $D$  is the depth of the encoder network and  $g_0(w_t) = U_\theta w_t$ . Notably, for the Gaussian diffusion models in Sec. 3, we can equivalently have  $x_{t} = g_{0}(w_{t})$ , as corruption and sampling are performed on the learned token embeddings. In the case of the categorical noise diffusion  $p_{\theta}(\hat{w}_0|w_t) = p_{\theta}(\hat{w}_0|h_t)$ , and thus for the purpose of guidance, we can consider a general  $p_{\theta}(\hat{w}|h_t)$  for both forms of corruption.
To sample from  $\tilde{p}_{\theta}(w_t) \propto p_{\theta}(w_t) \exp(v_{\theta}(w_t))$ , we construct a modified denoising model,
$$
\tilde {p} _ {\theta} (\hat {w} | h _ {t}) \propto p _ {\theta} (\hat {w} | h _ {t}) \exp (v _ {\theta} (h _ {t})).
$$
This formulation requires that the denoising model and the value function share hidden states up to depth  $d$ , and that the value function also be trained on corrupted inputs  $w_{t}$ . In Appendix D.4 we propose a simple procedure for corrupted discriminative training inspired by label smoothing [76]. Using this modified denoising model we can construct modified transition distributions using Eq. 2 or Eq. 3. There is one key difference between these transition distributions: in the continuous case (Eq. 2), smooth steps are taken in the token embedding space, while in the discrete case (Eq. 3) the transition leads to large jumps from one token embedding to another. In either case, it is possible to sample a discrete sequence  $w$  at any point in the chain using the logits of the denoiser  $p_{\theta}(\hat{w}|h_t)$ . When using Eq. 2 to derive a continuous transition distribution, we call the method NOS-C, and when using Eq. 3 for discrete transitions, we call the method NOS-D.
To sample from the modified transition distribution at each diffusion step, we use Langevin dynamics with temperature  $\tau > 0$ , with the update step,
$$
h _ {t} ^ {\prime} \leftarrow h _ {t} ^ {\prime} - \eta \nabla_ {h _ {t} ^ {\prime}} \left[ \lambda \mathrm {K L} \left(p _ {\theta} (\hat {w} | h _ {t} ^ {\prime}) | | p _ {\theta} (\hat {w} | h _ {t})\right) - v _ {\theta} \left(h _ {t} ^ {\prime}\right) \right] + \sqrt {2 \eta \tau} \varepsilon , \quad \varepsilon \sim \mathcal {N} (0, I), \tag {4}
$$
where  $\eta$  is the step size and  $\lambda$  is the regularization strength, followed by sampling  $p_{\theta}(w_{t - 1}|h_t^{\prime})$  or  $p_{\theta}(h_{t - 1}|h_t^{\prime})$ . While the gradient  $\nabla_h v_\theta$  guides towards high values of the objective, the KL term ensures the resulting transition distribution still maximizes the likelihood of the original prediction.
NOS is related to the popular method plug-and-play language model (PPLM), which can be used for gradient-guidance of autoregressive language models [17]. PPLM guides sampling by taking gradient steps similar to Eq. 4 for each autoregressive decoding step (details in Appendix B). Unlike PPLM, NOS is a form of iterative refinement, meaning that tokens across the entire sequence can be modified at each optimization step. This distinction is particularly important for protein design, because function can be determined by complex interactions between distant parts of the sequence. As we see in Sec. 5, NOS leads to better trade-offs between likelihood and objective value.
4.2 LaMBO-2: function-guided protein design
Many unique challenges arise when applying guided diffusion to real-world protein design tasks. Our approach builds on the LaMBO-1 algorithm proposed by Stanton et al. [73], which explicitly
Figure 3: An example of a 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-based uncertainty quantification.