SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
u _ {\mathrm {E H V I}} (\mathbf {x}, f, \mathcal {D}) = \operatorname {H V I} \left(\mathcal {P} ^ {\prime}, \mathcal {P} | \mathbf {r} _ {\text {r e f}}\right) = \left[ \operatorname {H V} \left(\mathcal {P} ^ {\prime} | \mathbf {r} _ {\text {r e f}}\right) - \operatorname {H V} \left(\mathcal {P} | \mathbf {r} _ {\text {r e f}}\right) \right] _ {+}, \tag {8}
$$
where  $\mathcal{P}' = \mathcal{P} \cup \{\hat{f}(\mathbf{x})\}$  [27, 28, 18]. To decide where to query  $f$  next, we search for  $\operatorname{argmax}_{\mathbf{x}} \mathbb{E}[u_{\mathrm{EHVI}}(\mathbf{x}, f, \mathcal{D})]$ , where the expectation is w.r.t.  $p(f|\mathcal{D})$ .
D.2 Discrete EHVI
Although expression yield and binding affinity are both continuous measurements, we chose to discretize them and model them as classification with a softmax likelihood (See Appendix D.4). As a result we needed an extension of EHVI for discrete outcomes. Informally, EHVI is simply computing the HVI for different realizations of  $f$  and marginalizing  $f$  using  $p(f|\mathcal{D})$ . Instead of taking  $f$  to be the latent function of some regression  $y = f(w) + \varepsilon$ ,  $\varepsilon \sim \mathcal{N}(0,\sigma^2)$ , we instead take  $f$  to be the logits of a categorical distribution,  $p(y = i|w,\mathcal{D}) = \int \text{softmax}_i(f(w))p(f|\mathcal{D})df$ .
Let  $\mathbf{y} = [y_1 \cdots y_k]^\top$ . Given a set of baseline points  $\mathcal{B} \subset \mathcal{A}^L$  we define  $\mathcal{P}$  (Eq. 8) using the posterior mean  $\hat{\mathbf{y}}(w) = \mathbb{E}[\mathbf{y}|w, \mathcal{D}]$ ,  $w \in \mathcal{B}$ . We model  $y_1, \ldots, y_k$  as conditionally independent given some shared hidden state  $h = g_d(w)$ , so  $p(\mathbf{y}|h, \mathcal{D})$  factorizes nicely. Finally we define  $\mathcal{P}' = \mathcal{P} \cup \{\mathbf{y}\}$  and take the expectation of Eq. 8 w.r.t.  $p(\mathbf{y}|h, \mathcal{D})$ . Since  $p(\mathbf{y}|h, \mathcal{D})$  is discrete and factorizes, we can marginalize in closed form when  $K_1 \times \cdots \times K_k$  is not too large, where  $K_i$  is the number of classes corresponding to the discretization of the original continuous  $f_i$ .
D.3 Architecture and Hyperparameters
The inputs of the LaMBO-2 model for antibody design are the variable heavy (VH) and variable light (VL) regions of the antibody sequence as determined by Aho alignment with ANARCI, as well as the (unaligned) antigen sequence. Note that the concatenation of the antigen to the input makes the samples from the generative head conditional on the antigen as well as the unmasked portion of the antibody sequence. The LaMBO-2 model jointly predicts antigen-conditional categorical
token distributions for corrupted positions and discriminative distributions over protein properties. Discriminative predictions that should not depend on the antigen are made invariant through data augmentation with random antigen sequences. See Algorithm 3 for an overview of a single guided diffusion step with LaMBO-2.
Model Architecture: our architecture for this experiment is inspired by the one proposed by Stanton et al. [73]. In particular we jointly a train an encoder shared between a generative discrete diffusion head and discriminative heads which predict expression and affinity. Rather than use a deep kernel GP, we simply ensemble 10 heads for each discriminative task to obtain uncertainty estimates. Like Stanton et al. [73] for this experiment we use 1D CNN residual blocks (kernel width 9), with layer normalization and sinusoidal position embeddings. The shared encoder was comprised of 4 residual blocks, and each task head was comprised of 2 residual blocks followed by a linear layer, with the exception of the generative head which was just a linear layer on top of the shared embeddings. Note that in future work self-attention layers could be used instead of CNN layers, as was the case for the pOAS experiments in Sec. 5. We set the embedding dimension to 32, and the latent channel dimension to 256.
Training Hyperparameters: The LaMBO-2 model is both a jointly trained generative and discriminative model, as well as a true multi-task model, which is necessary since measurements for various protein properties are often missing from a substantial fraction of rows in real-world datasets. We trained for 500K gradient updates using the Adam optimizer with  $\eta = 1\mathrm{e} - 3$ ,  $\beta_0 = 0.99$ ,  $\beta_{1} = 0.999$ . At each gradient step we randomly sampled a task head and task minibatch (batch-size 121) and updated the corresponding weights (including shared weights). We used a linear learning rate warmup over 10K gradient updates, and decayed the learning rate to 1e-6 with a cosine schedule. We did not regularize with weight decay or dropout.
Generation Hyperparameters: to generate the designs in Figure 7, we sampled 1K designs from a pool of seed antibody sequences hand-selected by domain experts. For each seed we set the total edit budget shared between chains to  $B = 16$ . In this experiment each infilling method took 16 diffusion steps, using an inverse linear noise schedule  $\overline{\alpha}_t = 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  $\tau = 1\mathrm{e} - 2$ . For guided infills with uniformly distributed edit positions we set  $\tau = 1\mathrm{e}6$ . For guided infills with saliency-informed edit position selection we set  $\tau = 0.1$ . We set  $\lambda = 0.5$  to balance the tradeoff of sequence likelihood and value during guidance.
Generation Constraints: in addition to the edit budget locality constraint, our LaMBO-2 designs were also constrained to meet certain sequence liabilities constraints:
- Canonical Cysteine Conservation: there are specific conserved cysteine residues in antibody sequences which play a crucial role in the formation of disulfide bridges. Disulfide bridges are covalent bonds formed between two cysteine residues through oxidation of their sulfur atoms. These bridges contribute to the overall structural stability and integrity of antibodies.
- No Unpaired Cysteines: odd numbers of cysteines within individual chains (i.e. unpaired cysteines) are generally undesirable since they can lead to non-native disulfide bonds between different antibody molecules, which may disrupt assembly, folding, or function.
- No Glycosylation Motifs: A glycosylation motif is a specific amino acid sequence within a protein that serves as a recognition site for the attachment of sugar molecules. The presence of a glycosylation motif in a protein can affect its stability, solubility, activity, and function. The addition of sugar molecules can alter the protein's conformation, change its interactions with other proteins or molecules, and affect its trafficking and localization within the cell.
D.4 Training Data, Class Imbalance, and Label Smoothing
Training Data: the expression task heads were trained on a dataset of 10K linear transfection expression measurements, which was subsequently augmented to 160K rows by pairing the same measurements with different random antigens to teach the model to ignore the antigen sequence when predicting expression. The binding task heads were trained on a dataset of 10K SPR affinity measurements for various antigens, which was then augmented to 12K rows by pairing binders with different random antigens and imputing a non-binding label. This augmentation is important for
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  $(\mathrm{pKD} = 4)$ , and pKD deciles computed from the remaining measurements.
training a pan-target affinity model, since experimental measurements of affinity to off-target antigens are uncommon. Note that the expression and affinity data only partially overlapped, necessitating the multi-task architecture described in Appendix D.3. The generative diffusion head was trained only on binding antibody-antigen pairs in the SPR binding data.
We did not pretrain our LaMBO-2 models. It is likely that performance could be improved with the right pretraining corpus, however it is unclear if datasets like pOAS are particularly useful for pretraining antibody design models since most do not report antigen sequences and may not have the right level of variability. In any case, it is very encouraging to see positive real-world results before scaling in earnest.
Label Discretization. As noted above, biological data tends to be very imbalanced, and historical experimental data even more so since there are strong selection effects imposed by the scientists collecting the data. We chose to discretize continuous properties like expression yield and binding affinity, making it easier to correct for class imbalance by upsampling minority classes. In Figure 14 we illustrate our discretization scheme. Any antibody-antigen pair with  $-\log (\mathrm{KD})$  (pKD) less than 4 was assigned to the non-binding class 0. Then binders were assigned to classes 1 - 10 based on which pKD decile (computed from binders only) they resided in. One consequence of this scheme is increasing any objective value by one unit corresponds to moving up one decile in the empirical label distribution.
Training Discriminators on Noisy Inputs: the benefits of discretization are not limited to addressing class imbalance. Working with discretized labels also allowed a simple approach to training the discriminator on corrupted inputs inspired by label smoothing [76]. We train the discriminators with the same noise schedule as the diffusion model and the usual cross-entropy loss, using modified labels
$$
\overline {{\mathbf {y}}} _ {t} = \overline {{\alpha}} _ {t} * \mathbf {y} + (1 - \overline {{\alpha}} _ {t}) / K * \mathbf {1},
$$
where  $\mathbf{y}$  is the one-hot encoded label and  $K$  is the number of classes. Informally, as  $\overline{\alpha}_t\rightarrow 0$  the discriminator reverts to a uniform prior since the inputs are not distinguishable. Training on corrupted inputs avoids evaluating the value gradient on out-of-distribution inputs during generation, and causes the strength of the value gradient to grow as the diffusion progresses and the samples become more defined.
D.5 Baselining LaMBO-2 Against Unguided Sequence and Structure-Based Diversification:
Structure-Based Diversification We have shown that we can effectively optimize antibodies for predicted yield and affinity, and our method performs well compared to unguided sequence-based infilling methods. We expand our evaluation for this task to include unguided infilling with DiffAb and RFDiffusion of CDRs H2 and H3 of hu4D5 (i.e. the seed), a publicly released therapeutic antibody that is ideally suited for structure-based methods since we have a ground truth crystal structure of hu4D5 docked with its target ERBB2. While it is not feasible to validate the resulting designs in vitro during the author response period, we can compare the AntiBERTy naturalness scores and the acquisition value (log expected hypervolume improvement or log-EHVI) of the designs relative to our guided infills (Fig. 15). To summarize, unguided structure-based infilling produces
high likelihood samples, but even when conditioned on the antigen the distribution shift toward better predicted function is very slight.
Figure 15: (left) we find that structure-based infills, particularly from DiffAb, tend to score consistently 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 to 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 to 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 in the distribution of acquisition values to the right.
Sequence Diversification This in silico evaluation compares two variants of LaMBO-2 (one using NOS-C, the other NOS-D) against a competing method, walk-jump sampling (WJS), an unguided smoothed discrete sampling algorithm proposed by Frey et al. [30]. Each method generated 1K designs from the same set of seeds, and all methods were restricted to  $B = 8$  edits. LaMBO-2 chose all edit positions automatically along the entire antibody sequence, whereas WJS was given manually selected edit positions restricted to CDRs only. In the left two panels of Figure 7 we compare the predicted expression yield, predicted binding affinity, and naturalness of the antibody designs, using the metric proposed by. Comparing the Pareto frontiers obtained from each set of designs, we see that while WJS excels at generating "natural" antibodies, it struggles to generate designs at the higher end of the objective range. Conversely LaMBO-2 designs (particularly those generated with NOS-C) have high predicted objective value but also lower naturalness scores. LaMBO-2 designs generated with NOS-D strike a balance between the two extremes.
Figure 16: We evaluate LaMBO-2 in 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 Saliency Maps Reliable?
There is substantial controversy regarding the reliability of input-gradient-based feature attribution methods, specifically related to their ability to consistently highlight ground truth task-discriminative features and ignore irrelevant features. For example, Hooker et al. [41] claim that random attribution is competitive with input-gradient methods, and Casper et al. [11] claim that gradient-free attribution outperforms input-gradient competitors. On the other hand, many papers claim that specific types
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 CDRL1. Some importance is also assigned to various framework regions, which could be related to the fitness of different antibody germlines. We emphasize that these models were trained solely on aligned sequences, with no additional positional information.
of regularization can improve the performance of input-gradient attribution, including adversarial training [66], mask denoising [6], and model curvature penalties [72].
A thorough investigation of these claims is beyond the scope of this work, however we have found that saliency maps produced by independent models trained with different corruption processes seem to consistently highlight specific regions of the antibody sequence (Figure 17). It is also worth noting that most of the related literature evaluates feature attribution in the offline setting. In LaMBO-2 feature attributions are used online to intervene on the data collection process (specifically where to introduce changes in the antibody sequences). If LaMBO-2 changes a position that does not affect function it is reasonable to conjecture that input-gradient attributions would adjust accordingly after the model is retrained for the next round. Further investigation into feature attribution in decision-making contexts (as opposed to post hoc interpretability) is an exciting direction for future work.
D.7 Wetlab Validation
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 a seed measurement in the same batch.
In this section we briefly summarize the experimental procedures used to validate LaMBO-2 designs in vitro. Designed antibody sequences from LaMBO-2 were expressed and purified, and surface plasmon resonance (SPR) measurements were used to determine binding affinity. See Figure 18 for a plot of design binding affinity vs. edit distance from seed antibody.
Plasmid Construction and Antibody Production: synthesized DNA of antibody variable domains (Twist Biosciences) were cloned into mammalian expression vectors using Gibson assembly. The whole vector was amplified using PrimeStar Max polymerase (Takeda). PCR products were transfected transiently in  $1\mathrm{mL}$  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  $280\mathrm{nM}$ .
Binding Affinity Measurements: affinity of the antibodies towards their target antigen was measured by surface plasmon resonance (SPR) at  $37^{\circ}\mathrm{C}$  on a Biacore 8K instrument (Cytiva) in HBS-EP+ buffer (10 mM Hepes, pH 7.4, 150 mM NaCl, 0.3mM EDTA and 0.05% vol/vol Surfactant P20). Antibodies were captured on a Protein A chip and their target antigen were injected for 5 minutes and allowed to dissociate for 10 minutes at 30ul/min. The surface was regenerated between cycles with 10 mM glycine pH 1.5. Affinity constants were obtained using Biacore Insight (Cytiva) using a 1:1 binding kinetics model.