SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Architecture and value function In order to apply the methods discussed in Subsec. 4.1 we require a generative diffusion model  $p_{\theta}(w)$  and a discriminator  $\hat{f}_{\theta}(w)$  which share hidden layers up to depth  $d$ . The discriminator is trained to predict the objective function  $f$ . Like LaMBO-1 our architecture consists of many task-specific feature extraction layers that share a bidirectional encoder. Bayesian optimization is an iterative cycle of inference and data acquisition. During the data acquisition phase of any iteration  $i$  we need to find sequences with maximal acquisition value  $v_{i}(w) = \mathbb{E}[u(w, f, \mathcal{D}_{i})]$ , where  $\mathcal{D}_i$  is the data already available and the expectation is taken with respect to a posterior  $p_{\theta}(f|\mathcal{D}_i)$  and  $u$  is some utility function. For multi-objective tasks  $u$  is the hypervolume improvement utility function [18], however we note that single-objective tasks are easily accommodated by a different choice of utility function [82]. To estimate the expectation we draw samples from  $p_{\theta}(f|\mathcal{D})$  with an approach we call partial deep ensembles, where the discriminative layers of the model above the shared encoder are replicated  $k$  times and randomly initialized [81]. We provide further details about partial deep ensembles and our learned discriminators in Appendix D.2 and D.3.
Choosing edit positions When  $B \ll L$  encoder-only architectures allow very precise control of edit positions since we will only change positions we corrupt. However, this feature introduces the need for some procedure to choose those positions, ideally where edits will most improve our objective value. We automatically select edit positions by computing the gradient of the value function with respect to  $h_0$  to determine which positions affect the value estimate the most (see Figure 3 for an illustration). This method is related to the use of saliency maps to explain the decisions of classifiers [5, 69]. We use input saliency to induce a distribution over edit positions. Specifically, given an embedded sequence  $h_0$  we define  $s_i(h_0)$ , the saliency with respect to  $v$  of position  $i \in \{1, \dots, L\}$  as
$$
s _ {i} \left(h _ {0}\right) := \max  \left\{\left(\sum_ {j = 1} \left| \left(\nabla_ {h} v \left(h _ {0}\right)\right) _ {i j} \right|\right) ^ {1 / \tau}, \varepsilon \right\}, \quad \mathbb {P} [ \text {e d i t} w _ {0} ^ {(i)} ] = \frac {s _ {i} \left(h _ {0}\right)}{\sum_ {j} s _ {j} \left(h _ {0}\right)}, \tag {5}
$$
where  $\tau > 0$  is a temperature hyperparameter and  $0 < \varepsilon \ll 1$ . As  $\tau \rightarrow +\infty$ ,  $\mathbb{P}[\text{edit } w_0^{(i)}] = 1 / L$  for all  $i$ . For each sequence we draw  $B$  edit positions without replacement according to Eq. 5. We conserve parts of the input we cannot change (e.g. the antigen sequence) by setting the saliency to 0 before computing the edit position distribution. Importantly, the diffusion sampling process can also preserve the original values of selected positions when appropriate. If we select a highly conserved position, then the predicted logits will be low entropy and the guidance will incur a large KL penalty for changes (Eq. 4).
We evaluate our methods on three increasingly complex antibody design tasks. First we compare our trained diffusion models on unguided infilling tasks, showing that sequence diffusion methods
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 (IgLM [68]). We see 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<sup>4</sup>. We then evaluate NOS by optimizing two objectives that can be simulated effectively in silico. Lastly, we evaluate LaMBO-2 on antibody lead optimization, with both in silico and in vitro experiments.
5.1 Unguided antibody CDR infilling
We focus on immunoglobulin G (IgG) format antibodies, which are comprised of a heavy (H) chain and a light (L) chain. Each chain has three complimentarity determining regions (CDRs), which tend to have strong effects on binding affinity to a target antigen but limited effects on other structural properties of the protein. Antibody design methods traditionally focus on proposing mutations to CDRs while leaving the rest of the protein fixed, which can be viewed as an infilling task. We select 10 seeds at random from paired OAS [56] and infill each CDR individually as well as in combination. To evaluate the performance of each model, we measure the sequence recovery rate, which is simply the accuracy of the infilling predictions given the ground truth sequence. As baselines, we include IgLM [68], a GPT2 language model trained on OAS, and two structure-based methods: DiffAb [51], a joint sequence-structure diffusion model trained on SAbDab, and RFDiffusion [80], a structural diffusion model trained on the PDB [10] that uses inverse folding to derive sequences. Although IgLM is trained with fill-in-the-middle augmentations [7], it does not natively support infilling multiple non-contiguous regions, and we do so by replacing regions that are not yet sampled with [UNK] tokens. For the structure-based methods, we provide starting structures generated with IgFold [64].
In Figure 4, we find that diffusion models often generate infills that are on-par or better than that those returned by IgLM by default, especially when multiple regions must be filled simultaneously. We also see that DiffAb, while being capable of sequence-structure co-design out of the box, often underperforms sequence-only diffusion, most likely because our sequence-based approaches have access to a larger training dataset, while paired datasets with sequences and structures are much more limited. Lastly RFDiffusion tends to generate relatively low likelihood CDR infills. The gap between DiffAb and RFDiffusion may be explained by the relative scarcity of antibody structures in the PDB compared to SAbDab, which has an antibody in every structure. The poor performance of structural methods on CDR infilling could also be a result of poor sequence recovery from structure during inverse folding, a problem that could be amplified for relatively unstructured loop regions like CDRs.
5.2 Optimizing antibodies for in silico objectives
To test guided sampling from our model, we run experiments on two simple single-objective tasks:
- The percentage of beta sheets, measured on primary sequence [15]
- The solvent accessible surface area (SASA) of the protein's predicted structure [64]
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.
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-likelihood sequences. (right) Guided generation using PPLM [17], a 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.  $\eta$  and  $\lambda$  (Eq. 4)) to demonstrate the full range of trade-offs 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-likelihood assigned by ProtGPT [29] (trained on Uniref50 [75]) plotted against the value of the objective. As an autoregressive guided baseline, we run PPLM, using IgLM as the base generative model (details in Appendix C.3). We use DiGress [79] as a guided diffusion baseline. DiGress uses gradients on one-hot representations to performing guidance in the embedding layer and is thus closely related to our approach. We discuss differences between the methods and the details of our DiGress experiments in Appendix C.5. For PPLM, DiGress, and NOS, we generate samples for many different settings of the control hyperparameters (Section 4.1), which yields samples across the spectrum from aggressively optimizing the objective to conservatively maintaining high likelihood. We also include DiffAb and RFDiffusion without guidance as baselines, as examples of popular "diversification" procedures, in which new samples are generated for later ranking and filtering. In Figure 5, we see that for both continuous and discrete corruptions NOS offers better trade-offs between optimizing the objective and maintaining high likelihood, while also generating high values of the objective at the extreme.
5.3 Antibody lead optimization: in silico evaluation
Having established the performance of NOS on simpler benchmarks, we now turn to real-world antibody design with LaMBO-2. From this point forward in all experiments we jointly condition on the heavy chain, light chain, and antigen sequence, and we jointly optimize the heavy and light chains only for improved expression yield and binding affinity to the antigen. As we discussed in Subsec. 4.2, one of the critical subproblems in Bayesian optimization is the identification of high value additions to the existing dataset. In this section we show that LaMBO-2 effectively applies NOS to this subproblem in the antibody design setting by finding high acquisition value sequences while preserving naturalness (which we quantify with the metric proposed by Shanehsazzadeh et al. [67]). We focus on optimizing hu4D5, a therapeutic antibody targeting the HER2 antigen<sup>5</sup>
Comparison with genetic algorithms We first compare LaMBO-2 to two discrete optimization baselines, AdaLead and PEX. We generated 32 designs from the hu4D5 seed with each method, optimizing the same acquisition function derived from the same model. To ensure a fair comparison we limited all methods to a total of 512 model evaluations and a maximum of 2 edits per sampling iteration. We evaluated both the sample acquisition value and naturalness after each iteration. We identified an empirical naturalness threshold below which expression became unreliable and treated this threshold as a simple inequality constraint. Note that this experiment evaluates how each method balances the tradeoff between acquisition value and naturalness as sampling progresses, and does not involve evaluations of the actual black-box objective.
<sup>5</sup>HER2 is an important target for certain types of breast and stomach cancer [36].
Figure 6: (left) Naturalness constraints present a 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 of 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 \in \{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 structure-based unguided infilling (i.e. random hit diversification) has the opposite behavior, producing samples with reasonable naturalness but low acquisition value.
Effect of salient edits To separate and independently study the effects of guidance (NOS) and salient position selection, we present an ablation in Figure 6 for optimization with a relatively small edit budget  $B$  ( $B <= 10\%$  of mutable positions). To isolate the effects of salient edits we baseline against edit positions chosen uniformly at random, and to isolate the effects of guidance we set the step size  $\eta$  (Eq. 4) to 0. Small edit distance constraints are common in antibody engineering because the goal is typically to increase binding affinity without altering the binding location on the antigen (i.e. the engineered antibody should bind to the same epitope) [43]. One heuristic way to constrain the design to the same epitope is to set  $B \approx 8$ , (about  $2.7\%$  of the antibody sequence length) [43], precisely the range we consider in Figure 6.
In the few edit regime we find that while both interventions improve sample objective value, selecting positions using saliency has a much larger effect than guidance. Although gradient guidance is a reliable and generally applicable tool for improved sampling, the scale of the edit position search dominates the scale of the search over token replacements that guidance affects. With a vocabulary of 21 tokens the number of possible token combinations  $(21^{8})$  is dwarfed by the combinations of possible edit positions  $(C_{8}^{300})$ . Salient selection of edit positions is, therefore, key to any practical application of NOS in budget-constrained design. Interestingly, this facet of protein design differs significantly from guided sampling of images, where generation is typically limited to fixed locations [50, 14], not a fixed edit budget spread over any location that will optimize the objective. These additional degrees of freedom pose an extra challenge.
5.4 Antibody lead optimization: in vitro evaluation
As our final evaluation in Figure 7 we present results using LaMBO-2 (specifically with NOS-D) to optimize 20 seed antibodies distributed across 4 different therapeutic target antigens, including
Figure 7: We use LaMBO-2 to optimize 20 seed antibodies for 4 different target antigens over three experimental rounds, retraining the model after each round. Some design choices and hyperparameters varied from round to round, with substantial impact on the results. In the last round we 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 antibody libraries [67]. Our results provide evidence that enriched antibody libraries can be created in silico without the assistance of high-throughput in vitro screening.
hu4D5/HER2 $^6$ . We tested 374 designs in total over three rounds, retraining the model after each round and varying a range of design choices and hyperparameters. While expression and binding performance varied from round to round across seeds and targets, by the final round we were able to generate multiple submicromolar binders for all 4 targets with a median of 5 edits to the seed. See Appendix D.7 for individual yield and affinity measurements and experimental details.
The improvements to yield and affinity over time can be attributed both to the data added to the training corpus and methodological insights gleaned after each round. For example, the sharp drop in expression in Round 2 can mainly be attributed to framework residue deletions that arose when  $\lambda$  (the KL penalty coefficient) was set too small. In the following round we tried a range of larger  $\lambda$  values and fixed the sequence lengths and expression immediately recovered.
Figure 7 also reports binding affinity results of a related experiment from Shanehsazzadeh et al. [67] for context, though we emphasize that there are substantial differences between our wetlab validation and that of Shanehsazzadeh et al. [67] which prevent a true apples-to-apples comparison. In the latter experiment 1M designs were generated for the HER2 target and screened with a high-throughput assay. After screening 421 designs were validated with a high-fidelity surface plasmon resonance (SPR) assay. In addition to wetlab screening, their experiment also restricted edits to specific antibody CDRs. We optimized antibodies for a range of targets including HER2 and relied exclusively on in silico screening before validating with SPR, while placing no explicit restrictions on the edit locations. Despite these differences, our results provide initial evidence that it is possible to generate enriched libraries of antibody designs exclusively with in silico methods operating only on primary sequence. While the experimental validation provided is preliminary, we are actively pursuing more rigorous experimental testing in the form of up-scaled and repeated expression and binding experiments and specificity assessment.
There are many exciting directions for future work. The original LaMBO algorithm was used to optimize small molecules in addition to proteins, and applying LaMBO-2 to small molecule design is a fruitful direction, as LaMBO-2's improvements are not protein-specific. While sequence alignments are a convenient solution to the length change problem in protein design, padding methods [47] or diffusion with a variable-length corruption process (e.g. [60]) will be needed for applications like small molecules which do not admit alignments. We are also eager to consider optimizing much longer sequences, such as gene perturbations [42], which can exceed 20K tokens in length and may necessitate the use of recent advancements such as implicit convolutions [35, 57, 59] or clever modifications of self-attention [16, 13, 48]. More general notions of guidance such as classifier-free guidance [38] for text or class-conditional generation [62, 12] are another intriguing direction, since some goals are difficult to express as black-box functions or constraints [49, 58].
Due to the sensitive nature of the data, we do not disclose the other seeds or drug targets.
[71] Raghav Singhal, Mark Goldstein, and Rajesh Ranganath. Where to diffuse, how to diffuse, and how to get back: Automated learning for multivariate diffusions. arXiv preprint arXiv:2302.07261, 2023.
[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 Alché, Yilun Du, Yaroslav Ganin, Arthur Mensch, Will Grathwohl, Nikolay Savinov, Sander Dieleman, Laurent Sifre, et al. Self-conditioned embedding diffusion for text generation. arXiv preprint arXiv:2211.04236, 2022.
[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):1282-1288, 2007.
[76] Christian Szegedy, Vincent Vanhoucke, Sergey Ioffe, Jon Shlens, and Zbigniew Wojna. Rethinking the inception architecture for computer vision. In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 2818-2826, 2016.
[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 motif-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 Krawczuk, Antoine Siraudin, Bohan Wang, Volkan Cevher, and Pascal Frossard. Digress: Discrete denoising diffusion for graph generation. arXiv preprint arXiv:2209.14734, 2022.
[80] Joseph L Watson, David Juergens, Nathaniel R Bennett, Brian L Trippe, Jason Yim, Helen E Eisenach, Woody Ahern, Andrew J Borst, Robert J Ragotte, Lukas F Milles, et al. Broadly applicable and accurate protein design by integrating structure prediction networks and diffusion generative models. bioRxiv, pages 2022–12, 2022.
[81] Andrew G Wilson and Pavel Izmailov. Bayesian deep learning and a probabilistic perspective of generalization. Advances in neural information processing systems, 33:4697-4708, 2020.
[82] James T Wilson, Riccardo Moriconi, Frank Hutter, and Marc Peter Deisenroth. The reparameterization 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 arXiv:2104.05218, 2021.
A Extended Background 16
A.1 Continuous noise diffusion 16
A.2 Categorical noise diffusion 17
B Methodological Details 17
B.1 Infilling algorithm 17
B.2 Hidden State Langevin Sampling 18
C Infilling / NOS Guidance 19
C.1 Infilling experiment 19
C.2 MCMC comparison 19
C.3 PPLM details 19
C.4 Model Architecture and Training 20
C.5 Hyperparameter settings 20
C.6 Density plots 22
D LaMBO-2 23
D.1 Intro to Multi-Objective Bayesian Optimization 23
D.2 Discrete 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 Diversification: 26
D.6 Are Saliency Maps Reliable? 27
D.7 Wetlab Validation 28
A Extended Background
In this section we provide full descriptions of the diffusion processes introduced in Sec. 3.
A.1 Continuous noise diffusion
The forward process is defined by noise variances  $\beta$ . We use the cosine variance schedule from Nichol and Dhariwal [55]. For convenience we further define
$$
\alpha_ {t} = 1 - \beta_ {t}, \bar {\alpha} _ {t} = \prod_ {i} ^ {t} \alpha_ {i}
$$
The forward process is defined by the conditional distributions
$$
\begin{array}{l} p \left(x _ {t} \mid x _ {t - 1}\right) = \mathcal {N} \left(x _ {t}; \sqrt {1 - \beta_ {t}} x _ {t - 1}, \beta_ {t} I\right) \\ p \left(x _ {t} \mid x _ {0}\right) = \mathcal {N} \left(x _ {t}; \sqrt {\bar {\alpha} _ {t}} x _ {0}, (1 - \bar {\alpha} _ {t}) I\right) \\ p \left(x _ {t} \mid w\right) = \mathcal {N} \left(x _ {t}; \sqrt {\bar {\alpha} _ {t}} U _ {\theta} w, (1 - \bar {\alpha} _ {t}) I\right) \\ \end{array}
$$