Protein Discovery with Discrete Walk-Jump Sampling
Authors: Nathan C. Frey
Dan Berenberg
Karina Zadorozhny
Joseph Kleinhenz
Julien Lafrance-Vanasse
Isidro Hotzel
Yan Wu
Stephen Ra
Richard Bonneau
Kyunghyun Cho
Andreas Loukas
Vladimir Gligorijevic
Saeed Saremi
Venue: ICLR 2024 oral (2024)
Keywords: generative modeling
langevin mcmc
energy-based models
score-based models
protein design
protein discovery
Primary area: applications to physical sciences (physics, chemistry, biology, etc.)
TL;DR: We resolve difficulties in training and sampling from a discrete generative model by learning a smoothed energy function, sampling from the smoothed data manifold, and projecting back to the true data manifold with one-step denoising.

Abstract: We resolve difficulties in training and sampling from a discrete generative model by learning a smoothed energy function, sampling from the smoothed data manifold with Langevin Markov chain Monte Carlo (MCMC), and projecting back to the true data manifold with one-step denoising. Our $\textit{Discrete Walk-Jump Sampling}$ formalism combines the contrastive divergence training of an energy-based model and improved sample quality of a score-based model, while simplifying training and sampling by requiring only a single noise level. We evaluate the robustness of our approach on generative modeling of antibody proteins and introduce the $\textit{distributional conformity score}$ to benchmark protein generative models. By optimizing and sampling from our models for the proposed distributional conformity score, 97-100\% of generated samples are successfully expressed and purified and 70\% of functional designs show equal or improved binding affinity compared to known functional antibodies on the first attempt in a single round of laboratory experiments. We also report the first demonstration of long-run fast-mixing MCMC chains where diverse antibody protein classes are visited in a single MCMC chain.

Published as a conference paper at ICLR 2024
PROTEIN DISCOVERY WITH DISCRETE WALK-JUMP SAMPLING
Nathan C. Frey Daniel Berenberg 33, Karina Zadorozhny'
Joseph Kleinhenz'
Julien Lafrance- V Lafrance-Vanasee
Isidro Hötzel
Yan Wu²
Stephen Ra
Richard Bonneau' Kyunghyun Cho 1.31
Andreas Loukas' Vladimir Gligorijevié' Saeed Saremi
Prescient Design, Genentech AAtibbddy Engineering, Genentech
Department of Computer Science, New York University
Center for Data Science, New York University
ABSTRACT
We resolve difficulties in training and sampling from discrete generative model by learning smoothed energy function, sampling from the smoothed data manifold with Langevin Markov chain Monte Carlo (MCMC), and projecting back to the true data manifold with one-step denoising. Our Discrete Walk-Jump Sampling formalism combines the contrastive divergence training of an energy-based model and improved sample quality of a score-based model, while simplifying training and sampling by requiring only a single noise level. We evaluate the robustness of our approach on generative modeling of antibody proteins and introduce the distributional conformity score to benchmark protein generative models. By opti- mizing and sampling from our models for the proposed distributional conformity score, 97-100% of generated samples are successfully expressed and purified and 70% of functional designs show equal or improved binding affinity compared known functional antibodies on the first attempt in single round of laboratory experiments. We also report the first demonstration of long-run fast fast-mixing MCMC chains where diverse antibody protein classes are visited in a single MCMC chain.
INTRODUCTION
Discrete sequence generation poses number of challenges to gradient-based generative models. Generative models must be expressive enough to faithfully capture the underlying data distribution, while also having controllable outputs that are novel, unique, diverse, and respect the constraints of the problem space. Energy-based models (EBMs) (Hinton & Sejnowski, 1986 LeCun et al. (006) fit energy function that specifies probability distribution over data analogous to the Boltzmann distribution from statistical physics. Giving access to an easily computable energy is an advantage of EBMs, but on the flip-side they can be difficult to train and sample from. Denoising objectives based on score matching (Hyvärinen) 2005 Vincent 2011) and the related advancements in diffusion models (Sohl-Dickstein al.) 2015 Ho al.] 2020) overcome these issues, but these either model the energy gradient or only provide access to an empirical lower-bound of the likelihood.
Protein design an instance of the discrete sequence generation problem, wherein the challenge is to find useful proteins in the large, discrete, and sparsely functional space (Romero & Arnold 2009) of dimension 20 for proteins of length Here, we consider the specific problem of generative modeling of antibodies, class of proteins with highly conserved structure that are of immense interest for therapeutics. addition to the qualities mentioned above, generative models for antibodies must be sample-efficient because of the relatively small size of datasets with therapeutic antibodies (Kim) ct al., 2023). Antibodics consist of well-conserved domains and high-entropy variable regions, so leveraging evolutionary information from pre-trained protein language models is not immediate solution. We distinguish ab initio protein discovery and design (producing novel, functional proteins
"Equal contribution
1 Cooresponnence to: (frey. Etrry.nathan www.aan.hannnn. nf1, saremi saeed) @gene.com com

---

Published as a conference paper at ICLR 2024
t=
t=T
((yt)
Selected samples from single Markov chain Monte Carlo sampling run of discrete Walk- Jump sampling (our method). Protein color corresponds different antibody germlines (classes). Samples are folded with EquiFold (Lee et al. (202)) for visualization purposes. Discrete walk-jump sampling exhibits fast mixing and explores diverse modes of the distribution single chain.
Figure 1:
given some training samples), which is the focus of this work, from de novo design, which we define as the generation of novel proteins without starting material. Existing autoregressive protein design methods (Jin et al., 2021) are inefficient and can suffer from accumulation of errors and high inference latency, while current non-autoregressive diffusion models are similarly inefficient and poorly optimized for real discovery and design tasks (Kong al. 2022). Our goal here is invent an efficient, -onaautorerressvee generative modeling paradigm for discrete data that produces high quality, novel samples.
To this end, we introduce Smoothed Discrete Sampling (SDS), a new formalism for training and sampling from discrete generative models. We propose novel algorithm, discrete Walk- Walk-Jump Sampling (dWJS), a method building on the neural empirical Bayes (NEB) (Saremi & Hyvärinen 2019) formalism. that addresses the brittleness of discrete EBMs and diffusion models and in doing so, provides a robust and general framework for protein discovery and design ! We also design metric called the Distributional Conformity Score (DCS), which a simple scalar score for protein sample quality. Our results rescue EBMs for discrete distribution modeling and question the need for diffusion models with multiple noise scales in protein discovery.
Our contributions are as follows:
We introduce a new paradigm for modeling discrete data distributions, Smoothed Discrete Sampling (SDS), building on the neural empirical Bayes framework. We propose the discrete Walk-Jump sampling algorithm, which uses uncoupled, separately trained score- and energy- based models learn noisy data distributions and sample discrete data, dWJS enables fast, non-autoregressive sampling with variable length discrete outputs. We also design novel architecture for discrete EBMs.
Our method simplifies score-based model training for discrete data by requiring only single noise level and no noise schedule, which alleviates the brittleness, training instabilities, and slow sampling of diffusion models. Our method also resolves difficulties in training EBMs, obviating the need for many common EBM training tricks (replay buffer, lz norm penalty, rejection sampling, etc.), while preserving good sample quality and fast sampling.
We demonstrate the utility of our approach in the context of ab initio protein discovery and design generating novel, biophysically-valid protein sequences from models trained on repertoires of functional molecules. We validate method with in vitro experiments. Our method outperforms autoregressive and masked protein language models, large language models, discrete sequence and ttructur-ssqqeece diffusion, and based baselines.
2

---

Published as a conference paper at ICLR 2024
2 BACKGROUND
2.1 ENERGY-BASED MODELS
EBMs are a class of models that learn an energy function fo X -> mapping inputs x (in IR to scalar "energy" value. The data distribution is approximated by the Boltzmann distribution
pe x fo(x))
EBMs typically trained via contrastive divergence (Hinton 2002), , and new samples drawn from Po(2) (x) by Markoy-Chain Monte Carlo (MCMC). Details of the loss function used in this work given in Section E] In Langevin MCMC, samples are initialized from known data point or random noise and refined with (discretized) Langevin diffusion
Tk+1 Fi 8 V [o(xx) + 2 Ek, ex N(0, Ta);
(1)
where V denotes the gradient of the energy function with respect to inputs, & is the sampling step, d is the (discretization) step size, and the noise E is drawn from the normal distribution at each step.
2.2 NEURAL EMPIRICAL BAYES
In NEB, the random variable X is transformed with additive Gaussian noise Y = X N(O. o² Ta).
The least-squares estimator of X given Y = y is given by (Robbins) 1956 Miyasawa 1961)
{ + 0² log)(v),
(2)
where p(y) p(y)ar)((a)d is the probability distribution function of the smoothed density This estimator is often expressed directly in terms of g(y) = V logp(y) known as the score func-
tion (Hyvärinen) (2005) which is parameterized with neural network denoted by go Rd R2 The least-squares estimator then takes the following parametric form:
(v) = y + o² (v)
(3)
Putting this all together leads to the following learning objective
( En--de). 2
which is optimized with stochastic gradient descent. Notably, MCMC sampling is required
during learning. In short, the objective "learning to denoise" with an empirical Bayes formulation (discussed further in Appendix B.3)
ANTIBODY DISCOVERY AND DESIGN
3.1 DISCRETE JALK-JUMP SAMPLING
Following training of the denoising network, 9a. one can sample noisy data using the learned score function 96 (y) with Langevin MCMC (replace - Vf with in [q. 1p. For any such draws D/s clean samples from the true data manifold M obtained by "jumping" back to M with the least-squares estimator (yx) = y/ + o² go (yx). This is the walk- jump sampling (WJS) scheme. A key property of WJS is the fact that the least-squares estimation (jump) is decoupled from the Langevin MCMC (walk).
Here, we take advantage of this decoupling train EBM with maximum likelihood estimation the smoothed distribution of noisy sequences, generate noisy samples with Langevin MCMC, and denoise samples with a separately trained neural network, the least-squares estimator. The algorithm for discrete walk-jump sampling is given in Algo. [ Our algorithm is general and applies to any discrete sequence inputs of fixed vocabulary. In Fig. [ we show samples generated from a single chain of MCMC. Unlike a diffusion model, every sample along the chain collectively forms valid set of samples from the underlying distribution, because of the decoupled walk (sampling) and jump (denoising) steps. dwJS also produces fast- bast-mixing chains, such that many diverse modes (protein
We follow the convention p(x) == px(x).p()) = py (V), etc.
3

---

Published as a conference paper at ICLR 2024
o
b
+N
x
y
dEBM
fa(y)
y e- fo(
ByteNet
V Togp(y)
d
dC0M
ML-I
ME-11
Walk
Linrar
Z
h
Pid
MLP
CNN
MLP
faly)
B Cencatccate
Figure 2: Discrete walk-jump sampling. The noising and denoising process is applied to antibody proteins. Discrete inputs are smoothed with isotropic Gaussian noise, ε ~ N(O, I2)), noisy inputs, z + E. A discrete energy- based model (dEBM) parameterizes the energy function fo( of noisy data. Noisy data is sampled with the energy function, and denoised with separate denoising ByteNer network to clean samples, (v). The "walk" sampling steps the noisy data manifold with Langevin MCMC are totally decoupled from the "jump" steps to clean samples. d The dEBM takes noisy inputs y, concatenates them with 1d positional encoding, Pid- passes through MLP and a 3 layer CNN, and concatenates the outputs with an embedding 2, of the inputs into a hidden state, h. is passed through an MLP and returns the energy fo (y))
classes) are sampled in single chain. Samples folded with EquiFold (Lee al. 2022 for visualization and confirmation of structural validity.
The EBM is trained by maximizing the log-likelihood of noisy data under the model:
arg max Eyppy lo arg max (E,, ~p)(s) (f((y))) 1 By Ioy ())(
(5)
8
where y* are noisy training data and y are noisy data sampled from the model.
With this objective, the model aims decrease the energy of noisy training data ("positive" samples y+t) while increasing the energy of noisy data sampled from the model ("negative" samples y in expectation. The following identity is behind the positive/negative phases in the EBM training:
Ve log (B) = (of - Vo log Z(0) fol())) dy
Vefo(v) V Vefo(y)e folu)
Z(0)
(6)
-Vefo())) Vafalv): po(u)dy
-Voffo + Eyrpa( (e [VFoov)))
where Z(0)- =
f - dy is the partition function.
Algorithm 1: Discrete Walk-Jump Sampling
Input: Denoiser, go (v), energy-based model, fo(y)
Output: Noisy samples y p(y), denoised samples i(y)
Vo Uni 1]d) N(0, TT))
for f. = 0,...,
do
De+1 + - ā V ( + v 20s, ~ N(O,
end
fr + yr + ozg
return arg max ār // to recover one-hot encoding
4

---

Published as a conference paper at ICLR 2024
Variable length protein sequence generation. We represent antibody protein molecules as I = (I1....,Tz)) where zı E 20). corresponds to the amino acid (AA) type at position l. Sequences from the Observed Antibody Space (OAS) database OOssen et al. 2022) are aligned ac- cording to the AHo numbering scheme (Honegger & PluEckthun 2001 using the ANARCI (Dunbar & Deane 2016) package and one-hot encoded. Aligning sequences in this way is a practical solution to handling insertions and deletions, which are otherwise troublesome for models that require fixed length inputs and outputs; alignment introduces a "gap" token that can be introduced or removed during sampling effectively change the length of sequences. This allows the model to capture the distribution of lengths present in natural antibodies. The alignment step maps heavy and light chain sequences of varying lengths standard, gapped input size of 149 and 148 respectively with 21 possible discrete tokens including the gap. Thus, the input dimension for every sequence becomes d= (149 148)= 21. Without loss of generality, any set of proteins can be aligned with a multiple tokens sequence alignment (Rao et al.) (02T). For other classes of discrete data, aseudo-alignmed can be used and randomly inserted into the inputs, or simple BOS and EOS tokens can be used and sampled. An EBM is trained via contrastive divergence on the manifold of smoothed, noisy one-hot encodings, y, given by y = x s, ~ N(O, I2), where x E {0, 13d. A separate denoising model trained with the objective in Eq. 4 New antibody sequences are generated (Fig. ( by sampling noisy
samples with Langevin MCMC following gradients from the EBM, denoising with the least-squares estimator, and taking argmax è to recover one-hot encoding. Further details related to training and network architecture are given in Appendix A
Protein design vs discovery. Protein discovery is the task of generating novel, unique, and valid samples. Protein design refers to taking some starting sequence and making edits to improve function. With dWJS we achieve discovery through unconditional sampling, while design is performed via constrained sampling and scoring. That is, we impose the following constraint in the form of a binary projection matrix
Id argmax ((y,)) t) = P
for P e {0, 1] where e is the number of conserved tokens in the sequence, y/ is the noisy sequence time step t of Langevin MCMC, 2(y, is the denoised sample at time t, and is the starting sequence. This constraint ensures that the specified regions of the sequence are conserved, while the non-conserved regions are free to change during Langevin MCMC.
3.2 DERIVATION OF OPTIMAL NOISE LEVEL FOR DISCRETE, SEQUENCE DATA
Throughout the experiments in Section we must choose what noise level, a, to use for training. Empirically, we find that in the protein discovery setting, o 2 0.5 is sufficient for getting good sample quality. Here, we provide some intuition for choosing a good ar, based on a geometric picture of the concentration of the measure (Saremi & Hyvärinen 2019). We define the matrix X with entries
||X, Xx
Xie
2vd
where d is the dimension of the data and the 2vd scaling comes from the concentration of isotropic Gaussians in high dimensions. The critical noise level, oe, is defined as
X
such that for >> de. all noisy data have some degree of overlap. For our antibody sequence data, the statistics of the X matrix are given in Table [ and the histogram of Xie' values is shown in Appendix [A.6] We find that 02 = 0.5, which agrees with our empirical hyperparameter optimization. Estimating de in this way serves to motivate the empirical success of the used in our experiments, and provides helpful guidance on the scale of a to use for discrete data. Here we take to be the length of the input vector (d = T. 297 for aligned antibody sequences); for the flattened sparse one-hot matrices with vocabulary size 21, = 6237. This scales ae by 0.22, which still gives a useful scale for but is not optimal because of the sparsity of the one-hot matrices.
3.3 DISTRIBUTIONAL CONFORMITY SCORE
The Fréchet inception distance (FID) Heusel et al. 2017) (a metric for image generation quality) and the BL.U (BiL.nggal Evaluation Understudy) score (Papineni et al., 2002) (for evaluating the quality
5

---

Published as a conference paper at ICLR 2024
Table 1: Statistics of the distance matrix, X, for discrete antibody sequence data.
min
median
mean
max RJ de
0.17
0.42
0.41
0.51
In slliro designs
[ikelihood of cack soquese
Conformicy measure
Compute metrics
returilecs netier hytiophathittty
joint density estimation
Chagg jundes
Likelihood okahnkk--t retrrsseset
Figure 3: in silico designs sampled with dWJS are compared to reference set of validation samples. Distributions are characterized with a set of sample quality metrics. Joint density estimation is used compute the likelihood of designs versus the validation and likelihoods condensed into distributional conformity score that characterizes the similarity of generated samples to the reference set.
of machine translation) greatly simplify the evaluation of proposed methods; protein generation lacks such metrics, which motivates us to introduce the "distributional conformity score" (DCS) (Fig. 3). The goal of the DCS is to provide a succinct description of how likely generated samples are with respect to a reference distribution, while maintaining novelty and diversity. DCS is designed such that improving sample quality corresponds directly to increased probability of generating real, biophysically valid proteins.
We evaluate the probability that generated sequences conform to reference distribution using the conformal transducer system (Shafer & Vovk 2008; Vovk et al. 2016). Let X E IRd, y E R. and Z X y A conformity measure is measurable function that maps sequence (Zi,..., 2n) E Z" to set of real numbers (a1, an) and is equivariant under permutations. Here, we define be the likelihood under the joint density over various properties, including biophysical properties and statistical properties, such as a log- tog-probabiling under a protein language model (the properties considered, further details, and comparisons to other measures of protein-likeness are given in Appendix FF. To avoid overfitting the estimator, we split the reference set into fitting set and validation set (Algo. 2). In our context, the reference distribution comprises all antibodies and label represents the property of interest such expression or binding. Empirically, we find that DCS is a useful in silico evaluation metric for developing generative methods and hyperparameter optimization, and that methods with DCS 0.3 yield nearly 100% expressing proteins in the wet lab.
Algorithm 2: Distributional conformity scores for evaluating generated designs
Input: Reference distribution DD test example æ e X. conformity measure A, label y
Output: p-value yd (the fraction of validation examples that are less similar Dalp than x:)
Sample 2a)...,,,n))) ~ D2lp and held-out validation set (ā,,..., ēx-1), 1), 24 ~ Dzly
Set on +- (x,
for i to ls do
a4 t A A(21..... Zn, 2n)
end
nd + & Dimla: < ax)
return pd
II the discussion of distributional conformity score, refers to sample features; elsewhere in the paper refers clean data. Here, y refers to lahels; elsewhere in the paper refers noisy data.

---

Published as a conference paper at ICLR 2024
Table 2: Ablikeness metries, uniqueness, diversity, and distributional conformity scores.
Model
W property
Unique f
Edist f
IntDiv t
DCS 1
dWJS (energy-based)
0.056
1.0
58.4
55.3
0.38
dWJS (score-based)
0.065
0.97
62.7
65.1
0.49
Seq VDM
0.062
1.0
60.0
57.4
0.40
DEEN
0.087
0.99
50.9
42.7
0.41
GPT 3.5
0.14
0.66
55.4
46.1
0.23
IgLM
0.08
1.0
48.6
34.6
0.533
ESM2
0.15
1.0
70.99*
77.56*
0.061
EXPERIMENTS
We evaluate our method, discrete Walk-jump sampling (dWJS) 2 on three antibody generation tasks: 1) distribution learning on paired observed antibody space (Olsen et al.) 2022): 2) the vitro expression and purification of novel antibodies; and 3) most importantly. functional therapeutic antibody design (Mason et af., 202T). Crucially, we compare methods using our distributional contormity score which is sample-to-distributio metric to assess sample quality (analogous to an FID score), rather than the sequence recovery metries used in previous antibody design work (Kong et al. [022]] Jin et al., (202)). Sequence recovery is poor objective for our goal, which is the discovery of novel (large edit distance from known examples), functional antibodies. Details related to model architectures, training, baseline methods, and sequence sampling in Appendix A
4.1 DWJS GENERATES NATURAL, NOVEL, DIVERSE ANTIBODIES in silico
We measure generative model performance with suite of "antibody likeness" (ab-likeness) metrics including labels derived from the AA sequence with Biopython (Cock et al. 2009). Sequence property metrics are condensed into single scalar metric by computing the distributional conformity score and the normalized average Wasserstein distance W Wrooperty between the property distributions of samples and a validation set. The average total edit distance Ediat summarizes the novelty and diversity of samples compared to the validation set, while internal diversity (IntDiv) represents the average total edit distance between samples. Our method achieves strong ablikeness results (Table simply by increasing a to 2 0.5.
dWIS with dEBM sampling achieves the best agree- ment with the validation set property distribution and highest percentage of unique samples, while dWJS with score-based sampling has high distributional conformity score, novelty, and diversity. We compare to latent sequence diffusion method VDcq) (a discrete generalization of variational diffusion; Kingma al. 2021). ccore-based model with ergy parameterization (DEEN), a transformer boased language model trained specifically for antibody de- sign (IgLM): Shuai et al. (202T), ESM2[Lin et al.] 2023). and a pre-trained large language model (LLM) (GPT 3.5). Our dWJS methods have faster sampling time and lower memory footprint than diffusion. autoregressive, and score- based baselines (Table while also having better sample quality. edit distance, and internal diversity metries,
Table 3: Measured protein synthesis.
Model
total f
dWJS (score-based)
1.0
dWIS (energy-based)
0.97
EBM
0.42
Our methods outperform IgLM on stibbod-llikenese whereas lgL.M has the best DC score and the lowest (worst) IntDiv. This indicates that IgLM samples are extremely close to the reference set and samples are similar to one another. As our acore-based dWJS has DC score of 0.49 and 100% expression rate in the lab, it is likely that all IgLM samples would successfully express in the lab as well. IgLM does produce 100% unique samples, outperforming score-baseds dWJS and GPT 3.5. Due IgLM being autoregressive sampler, method (energy-basee dWJS) has 43x faster sampling speed, which is useful when generating and ranking large numbers of designs. To mimic the ab initio generation task for ESM2, increase the
7

---

Published as a conference paper at ICLR 2024
masking percentage as high set sequences to generate new samples. As expected, ESM2 performs extremely poorly in generating ab initio samples with good untibody-likeness, because it is not trained for generation. It generates highly repetitive sequences that are very dissimilar to antibodies (hence the high, but meaningless, Edist and Div scores). As masked language model, ESM2 infilling is the fastest sampler, but at the cost of poor sample quality. Details on the baseline methods, and IgLM and GPT 3.5 prompts are given in Appendices A and D
(40%) while still generating valid samples and infill validation
4.2 DWIS GENERATES NATURAL, NOVEL, DIVERSE ANTIBODIES in vitro
Out of more than 277 designed antibody sequences tested in the laboratory, 270 were successfully expressed and purified (Table We achieved the 97.47% in virro success rate by developing dWJS to capture the antibody distribution in silico as measured by W property and distributional conformity score. For comparison. sequences from an EBM (trained on clean data with samples drawn using traditional Langevin MCMC) achieved a 42% expression rate. An antibody sequence comprised of random vocabulary tokens would be expected to have a 0% expression rate, and in laboratory experiments we have confirmed that small number of edits (< 4) can destroy expression if the proposal distribution (generative model) is poorly optimized.
4.3 DWJS GENERATES FUNCTIONAL ANTIBODY VARIANTS in vitro
To further show the ro- bustness of our method, we consider the task of
Table 4: Predicted and measured antibody binding affinity.
training generative mod- els on hu415 antibody mutant dataset (Mason et al. 2021) and compare to baseline models. The dataset consists of 9k bind- ing and 25k non-binding hu4D5 CDR H3 mutants with up to 10 mutations (after de-duplication and removing samples that are labeled both binding and non-binding). This yields a 10 dimensional search space. The mutants were measured in lab exper- iments to determine their binding to HER2 antigen. The goal of this benchmark task to produce unique samples that also bind to HER2. We trained dWJS models (score-based and energy- on only the binder set a noise level of a = 0.5, while also training a ID-CNN binary classifier to classify binders and non- non-binders. The classifier achieves 86% accuracy
Model
Poind
totalbind f
dWJS (energy-based) (Ours)
0.96
0.70
dWJS (score-based) (Ours)
0.95
N/A
LaMBO-2 (Gruver et al. 2023
N/A
0.25
AbDiffuser (Martinkus et al. 2023)
0.94
0.22 (0.57)
Seq VDM
0.75
N/A
GPT 4
0.74
N/A
Transformer
0.60
N/A
EGNN
0.58
N/A
IID validation split. Then, we classified 1000 samples from each dWJS generative model and four baseline models also trained on the hu4D5 binder set. We compare in silico to three diffusion models: 1) sequence transformer based on BERT (Devlin et al., 2018) that generates sequences, 2) an E(n) Equivariant Graph Neural Network (EGNN) (Satorras et al. 2021) l that codesigns (sequence, structure), and 3) latent sequence diffusion model, Seq and a pre-trained LLM, GPT 4. The specific prompt used for GPT is given in Appendix D The probability of binding for unique designs from each model is reported in Table 4 Our methods produce the highest percentage of unique predicted binders.
We also report in vitro wetlab validation results for the dWJS energy- designs. dWJS produces the highest percentage of functional antibodies that bind to target (totalpind in Table 4. We achieved a 70% binding rate the first round of experimental wet lab validation on HER2, by redesigning the CDR H3 loop of trastuzumab. To the best of our knowledge, this is the highest reported binding rate of any antibody design method applied to trastuzumab CDR H3 redesign. Concurrent with the first appearance of our work, two diffusion models reported in vitro wet-lab results for this task: AbDiffuser (Martinkus et al. 2023), , a (sequence, structure) antibody codesign method, and LaMBO-2 (Gruver et al. (023). a guided discrete diffusion method. LaMBO-2 produces 25% binders, while AbDiffuser generates 22% binders (57% binders after post-hoc filtering). Our method, dWIS, produces the highest percentage of unique binders (70%) and requires post-hoc filtering. While the
8

---

Published as a conference paper at ICLR 2024
diffusion methods include other important capabilities including structure generation (AbDiffuser) and guided sampling (LaMBO-2), our method outperforms diffusion for discrete distribution learning, which is our goal. For this experiment, our model is trained only on the publicly available (Mason et al. 2021) dataset. This training dataset, our code and and our experimental results are open and we encourage other researchers to compare results in the same in vitro setting. Further details wetlab experiments are presented in Appendix
5 RELATED WORK
Energy-based models (EBMs) LeCunn et al. 2006) are a class of physics-inspired models that learn an energy function defining probability distribution over data with a rich history that goes back to Boltzmann machines (Hinton & Sejnowski, 1986). Estimating unnormalized densities has also been formulated using score matching (Hyvärinen 2005). This formulation led probabilistic models for denoising autoencoders (Vincent) 2011 Alain & Bengio 2014, Saremi et al., 2018), but also has an empirical Bayes interpretation that most related to this work. particular, the neural empirical Bayes (NEB) (Saremi & Hyvärinen (209)) formalism unifies kernel density estimation (Pazen 1962) and empirical Bayes (Robbins] 1956) to transform the unsupervised learning problem into more tractable form where a neural network energy function is parameterized to capture smoothed data distribution. Our work is the first study of the NEB formalism for discrete data. Discrete diffusion models such (Austin et al. 2023) learn iterative denoising process over many different noise levels by prescribing noise process over discrete data that converges to a known categorical distribution.
Approaches borrowing from traditional MI generative modeling have been used to model antibodies (Shuai et al., 2021] Gligorijevié et al. 2021] Ferruz Höcker 2022; Tagasovska et al. 2022). but typical natural-languagee
methods struggle to capture the data distribution of antibodies, for which there is limited training data (~IK- 1M high-quality sequences depending on the distribution of interest) and additional challenges due to the high-entropy variable regions of the sequence. Here, we address the above challenges with training and sampling discrete sequences using a novel formulation of decoupled energy- and acore-based modeling.
6 CONCLUSIONS
We proposed Smoothed Discrete Sampling (SDS), new paradigm for modeling discrete distributions that uses Langevin Markov-Chain Monte Carlo to sample from smoothed data distributions. We introduce the discrete Walk-Jump Sampling (dWJS) algorithm and evaluate it on the antibody discovery and design problems, showing the capability of our method to generate novel, diverse, and functional antibodies measured by synthetic biophysical property distributions, similarity metrics, and in vitro experiments. The strong regularization provided by fitting the energy function to noisy data completely prevents overfitting and training instabilities, resulting in fast and efficient training and sampling with low compute requirements. dWJS discards many of the commonly used techniques for improving EBM training with Langevin MCMC (replay buffers, & norm penalty, simulated annealing, rejection sampling, etc.) and reduces the engineering complexity of training EBMs and diffusion-basee based models to single hyperparameter choice: the noise level, a.. Altogether, our results suggest a simplified, more general and robust framework for training and sampling from discrete energy- and score- based models with applications to therapeutic molecule design. Future work will probe the generality of our results to other classes of molecules and even other data modalities (e.g., images), as well as theoretical investigation into the results presented here.

---

Published as a conference paper at ICLR 2024
ACKNOWLEDGEMENTS
The authors acknowledge the entire Prescient Design team and the Antibody Engincering department at Genentech for providing helpful discussions and input that contributed the research results reported within this paper. The authors would like to especially acknowledge Simon Kelow, Franziska Seeger, and Andrew Watkins for helpful discussions related to antibody discovery, and Allen Good- man for consulting on large language model benchmarks. We would additionally like to acknowledge the reviewers for helpful comments and questions that improved the paper.
REFERENCES
Guillaume Alain and Yoshua Bengio. What regularized auto-encoders learn from the data-generating distribution. Journal of Machine Learning Research, 15()):3563-3593 2014.
Jacob Austin. Daniel D. Johnson, Jonathan Ho, Daniel Tarlow, and Rianne van den Berg. Structured denoising diffusion models in discrete state-spaces, 2023.
Peter JA Cock, Tiago Antao, Jeffrey T Chang, Brad A Chapman, Cymon J Cox, Andrew Dalke, Iddo Friedberg, Thomas Hamelryck, Frank Kauff, Bartek Wilczynski, al. Biopython: freely available python tools for computational molecular biology and bioinformatics. Bioinformatics, 25(11) (1):: 1422-1423, 2009.
Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint 1rXiv:1810.04000 2018.
Yilun Du and Igor Mordatch. Implicit generation and modeling with energy based models. Advances in Neural Information Processing Systems, 32, 2019.
James Dunbar and Charlotte Deane. ANARCI: antigen receptor numbering and receptor classifica- tion. Bioinformatics, 32(2):298-300 2016.
Noelia Ferruz and Birte Höcker. Controllable protein design with language models. Nature Machine Intelligence, 4(6):521-532, 2022.
Vladimir Gligorijević. Daniel Berenberg, Stephen Ra, Andrew Watkins, Simon Kelow, Kyunghyun Cho. and Richard Bonneau. Funetion- guided protein design by deep manifold sampling. bioRxiv, 2021.
Nate Gruver, Samuel Stanton, Nathan C Frey, Tim GJ Rudner, Isidro Hotzel, Julien Lafrance-Vanasse, Arvind Rajpal, Kyunghyun Cho, and Andrew Gordon Wilson. Protein design with guided discrete diffusion. arXiv preprint 23Xx:23355.2009 2023.
Martin Heusel, Hubert Ramsauer, Thomas Unterthiner, Bernhard Nessler, and Sepp Hochreiter Gans trained by a two time-scale update rule converge to a local nash equilibrium. Advances in neural information processing systems, 30, 2017.
Geoffrey E Hinton. Training products of experts by minimizing contrastive divergence. Neural computation, 14(8):177--1800, 2002.
Geoffrey E Hinton and Terrence Sejnowski. Learning and relearning in Boltzmann machines.
Parallel distributed processing: Explorations in the microstructure of cognition, 1(282-317):2 1986.
Jonathan Ho. Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in Neural Information Processing Systems, 33:6840-6851, 2020.
Annemarie Honegger and Andreas Pluf:ckthun. Yet another numbering scheme for immunoglobulin variable domains: automatic modeling and analysis tool. Journal of molecular biology, 309(3): 657-670, 2001.
Yi-Chun Hsiao, Ying-Jiun J Chen, Leonard D Goldstein, Jia Wu, Zhonghua Lin, Kellen Schneider, Subhra Chaudhuri, Aju Antony, Kanika Bajaj Pahuja, Zora Modrusan, et al. Restricted epitope specificity determined by variable region germline segment pairing in rodent antibody repertoires. In MAbs, volume 12, pp. 1722541. Taylor & Francis, 2020.
10

---

Published as a conference paper at ICLR 2024
Aapo Hyvärinen. Estimation of non-normalized statistical models by score matching. Journal of Machine Learning Research, 6(Apr):695-70) 7099,709000 2005.
Wengong Jin, Jeremy Wohlwend, Regina Barzilay, and Tommi Jaakkola. Iterative refinement graph neural network for antibody sequence-structure co-design. arXiv preprint arXtv:2110.04624, 2021.
Nal Kalchbrenner, Lasse Espeholt, Karen Simonyan, Aaron van den Oord, Alex Graves, and Koray Kavukcuoglu. Neural machine translation linear time. arXiv preprint 1rXiv:1600.10099 2016.
Jisun Kim, Matthew McFee, Qiao Fang, Osama Abdin, and Philip M Kim. Computational and artificial intelligence-based methods for antibody development. Trends in Pharmacological Sciences, 2023.
Diederik P Kingma, Tim Salimans, Ben Poole, and Jonathan Ho. Variational diffusion models. In A. Beygelzimer, Y. Dauphin, P. Liang, and J. Wortman Vaughan (eds.), Advances in Neural Information Processing Systems, 2021.
Xiangzhe Kong, Wenbing Huang, and Yang Liu, Conditional antibody design as 3d equivariant graph translation. arXiv preprint arXiv:2208 06073, 2022.
Yann LeCun, Sumit Chopra, Raia Hadsell, M Ranzato, and Fujie Huang. A tutorial on energy-based learning. Predicting structured data, 100), 2006.
Jae Hyeon Lee, Payman Yadollahpour, Andrew Watkins, Nathan C Frey, Andrew Leaver-Fay, Stephen Ra, Kyunghyun Cho, Vladimir Gligorijevic, Aviv Regev, and Richard Bonneau. Equifold: protein structure prediction with novel boarse-grained structure representation. bioRxiv, pp. 2022-10, 2022.
Zeming Lin, Halil Akin, Roshan Rao, Brian Hie, Zhongkai Zhu, Wenting Lu, Nikita Smetanin, Robert Verkuil, Ori Kabeli, Yaniv Shmueli, al. Evolutionary-scale prediction of atomic-level protein structure with a language model. Science, 11996667):1123-111 2023.
Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. arXiv preprint arXiv:1711.05101, 2017.
Karolis Martinkus, Jan Ludwiczak, Kyunghyun Cho, Wei-Ching Lian, Julien Lafrance- Lafrance-Vanasse, Isidro Hotzel, Arvind Rajpal, Yan Wu, Richard Bonneau, Vladimir Gligorijevic, et al. Abdiffuser: Full-atom generation of in-vitro functioning antibodies. arXiv preprint arXiv:2308.C 05027, 2023.
Derek M Mason, Simon Friedensohn, Cédric R Weber, Christian Jordi, Bastian Wagner, Simon M Meng, Roy A Ehling, Lucia Bonati, Jan Dahinden, Pablo Gainza, et al. Optimization of therapeutic antibodies by predicting antigen specificity from antibody sequence via deep learning. Nature Biomedical Engineering, 5(6))600-61) 612, 2021.
Chenlin Meng, Kristy Choi, Jiaming Song, and Stefano Ermon. Concrete score matching: Generalized score matching for discrete data, 2023.
Koichi Miyasawa. An empirical Bayes estimator of the mean of normal population. Bulletin of the International Statistical Institute, 38(4):181-18)) 1961.
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.
Kishore Papineni, Salim Roukos, Todd Ward, and Wei-Jing Zhu. Bleu: a method for automatic evaluation of machine translation. In Proceedings of the 40th annual meeting of the Association for Computational Linguistics, pp. 311-318, 2002.
Emanuel Parzen. On estimation of a probability density function and mode. The annals of mathemat- ical statistics, 1065):766, 1962.
11

---

Published as a conference paper at ICLR 2024
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Kopf, Edward Yang. Zachary Vetito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala. Pytorch: An imperative style, high-performance deep learning library. In Advances in Neural Information Processing Systems 32, pp. 8024-8035. Curran Associates, Inc., 2019.
Roshan M Rao, Jason Liu, Robert Verkuil, Joshua Meier, John Canny. Pieter Abbeel, Tom Sercu, and Alexander Rives. Msa transformer. In International Conference on Machine Learning, pp- 1844-8856, PMLR, 2021.
Matthew IJ Raybould, Claire Marks, Konrad Krawczyk, Bruck Taddese, Jaroslaw Nowak, Alan P Lewis, Alexander Bujotzek, Jiye Shi, and Charlotte M Deane. Five computational developability guidelines for therapeutic antibody profiling. Proceedings of the National Academy of Sciences, 116(10):402544033 2019.
Herbert Robbins. An empirical Bayes approach to statistics. In Proc. Third Berkeley Symp., volume 1, pp. 163-163, 1956.
Philip A Romero and Frances H Arnold. Exploring protein fitness landscapes by directed evolution. Nature reviews Molecular cell biology, (0(12):866-876 2009.
Matthias Sachs, Benedict Leimkuhler, and Vincent Danos. Langevin dynamics with variable coeffi- cients and nonconservative forces: from stationary states to numerical methods. Entropy, 19(12): 647, 2017.
Saced Saremi and Aapo Hyvärinen. Neural empirical Bayes. Journal of Machine Learning Research, 20:1-23, 2019.
Saeed Saremi and Rupesh Kumar Srivastava. Multimeasurement generative models. In International Conference on Learning Representations, 2022.
Saeed Saremi, Arash Mehrjou, Bernhard Schölkopf. and Aapo Hyvärinen. Deep energy estimator networks. arXiv preprini arXiv: /00 08306, 2018
Victor Garcia Satorras, Emiel Hoogeboom, and Max Welling. E(n) equivariant graph neural networks. In International conference on machine learning, pp. 9323-9332. PMLR, 2021.
Glenn Shafer and Vladimir Vovk. A tutorial on conformal prediction. Journal of Machine Learning Re- search, 9(1)):711-22 2008 URL http::/tttmmil http:///ww....ro//pppp////////aa
html
Richard W Shuai, Jeffrey A Ruffolo, and Jeffrey Gray. Generative language modeling for antibody design. bioRxiv, 2021.
Jascha Sohl-Dickstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsupervised learning using nonequilibrium thermodynamics. In International Conference on Machine Learning, pp. 2256-2265. PMLR, 2015.
Nataša Tagasovska, Nathan C Frey, Andreas Loukas, Isidro Hötzel, Julien Lafrance- Vanaase--Vanasee Ryan Lewis Kelly, Yan Wu, Arvind Rajpal, Richard Bonneau, Kyunghyun Cho, et al. A pareto- optimal compositional anergy-based model for sampling and optimization of protein sequences. arXiv preprint arXiv:2210.10833 2022.
Pascal Vincent. A connection between score matching and denoising autoencoders. Neural compula- fion, 2367)11611-1674, 2011.
Vladimir Vovk, Valentina Fedorova, Ilia Nouretdinov, and Alexander Gammerman. Criteria of efficiency for conformal prediction. CoRR, 10011600304416. 2016. URL httt:::/arxiv..o////
bbs//603..m 04416
Kevin K. Yang, Alex X. Lu, and Nicolo Fusi. Convolutions are competitive with transformers for
protein sequence pretraining. bioRxiv, 2022. doi: 10. 19.110022022.00.999997777 URL https?
12

---

Published as a conference paper at ICLR 2024
Appendix
Table of Contents
Network architectures and training details
A.1
Discrete Julk-Jump Samplers
13
A.2
dWJS stabilizes and simplifies training
14
A.3
Diflusion baselines
14
A.4
Language model basclines
14
A.5
Effect of choice of a)
15
A.6
Estimation of d. hyperparameter
15
B
Additional algorithms
15
B.1
Gradient Nlw enables local minima finding
15
B.2 Langevin MCMC Update
16
B.3 Neural Empirical Bayes
16
C
Performance profiling
17
D Few-shot, in-context learning prompts
17
GPT 4 prompt
D.2 GPT 3.5 prompl
Ig|. prompt
E in vitro validation
18
Experimental results
18
E.2 Experimental details
18
F Further discussion distributional conformity score
19
G
Further discussion of related work
19
A NETWORK ARCHITECTURES AND TRAINING DETAILS
A.1 DISCRETE WALK-JUMP SAMPLERS
For all experiments we use an identical architecture for the EBM consisting of three Conv1D layers with kernel sizes 15, 5, and 3 and padding 1, ReL.U non-linearities and output linear layer of size 128. The denoising model is 35-layer ByteNet (Kalchbrenner et al., 2016) architecture with hidden dimension 128, trained from scratch. The Bytenet architecture has been shown to perform competitively with transformers for protein sequence pretraining tasks (Yang et al. (022).. All models were trained with the AdamW (Loshchilov & Hutter (207)) optimizer Py Torch (Paszke al. 2019). We used a batch size of 256, an initial learning rate of 1 x 10-4, and trained with early stopping.
Transformer implementation of dWJS denoising model In addition to the ByteNet implementa- tion, we implemented transformer-based architecture for the denoising network, with 12 hidden layers, 8 attention heads, dimension of 2048 for the feed-forward layers, 256 features in the en- coder/decoder inputs, and SiLU activations. The score-basee based dWJS with a transformer architecture exhibits comparable performance to the ByteNet architecture, indicating that the performance of our method is not reliant on a particular architecture, and any sufficiently expressive architecture will work.
13

---

Published as a conference paper at ICLR 2024
Taxonomy of Smoothed Discrete Sampling Because of the decoupled walk and jump steps. there are many natural implementations of Smoothed Discrete Sampling. Empirically, we find that Algo. [] takes advantage of both energy- and ccsre-based modeling to produce the highest quality, novel, unique, diverse samples. Four natural choices for performing sampling, which arise from different combinations of energy- and score-based parameterizations,
summarized in Table Discrete Walk-Jump Sampling refers to walking with the EBM, fa (v), and jumping with the denoising network, go(y). Score-based dWIS uses (2 for both walking and jump steps. The Deep Energy Estimator Network (DEEN) (Saremi et al. 20188) approach uses denoiser that is trained by taking the derivative of an energy and using the same learning objective as Eq. 4 DEEN can be thought of as an energy parameterization of a score-based generative model. Finally, dWJJSEBBM uses an EBM for sampling and the gradient of the energy, [yy),, for denoising. Empirically, we find that the most performative method in terms of efficiency, sample quality, and diversity is the EBM walk and denoiser jump, which we refer to as Discrete Jalk-Jump Sampling.
Table 5: Smoothed Discrete Sampling implementations.
Model
Walk (sampling)
Jump (denoising)
dWJS (energy-based)
FBM
Denoiser
dWJS (score-based)
Denoiser
Denoiser
Deep Energy Estimator Network
Denoiser (energy)
Denoiser (energy)
dWJS-EBM
EBM
EBM
A.2 DWJS STABILIZES AND SIMPLIFIES TRAINING
We observe that the dWJS algorithm prevents instabilities during maximum likelihood training. EBMs commonly exhibit issues with training stability and divergences in the energy, due the energy landscape becoming too complicated to sample. Noising the data provides strong regularization that prevents overfitting and instabilities. This is seen over a range of noise levels [0.5, 4.0] for EBMs trained over 3,000 steps. Training instabilities recur for < 0.5. We investigate the effects of discarding many of the techniques for improved EBM training that, while introduced to ameliorate challenges with EBMs, also introduce complexities that make EBMs brittle, inflexible, and difficult to optimize. In particular, we discard the replay buffer, the ES norm penalty loss term to regularize the energies, Metropolis rejection sampling. and time step annealing. We use the Langevin MCMC algorithm (Algo. 4 from (Sachs et al.) 2017) and eliminate the need for careful hyperparameter finetuning; is the only free hyperparameter in dWJS.
A.3 DIFFUSION BASELINES
In our comparison study we use the Sequence- Variational Diffusion Model (Seq VDM) proposed by Kingma et al. Kingma et al. (2021). adapted for protein sequence data. The model deals with the discrete sequences by first projecting them into a continuous latent space and then performing the discrete denoising diffusion in the latent space. The VDM learns the data distribution by modeling the reverse of a diffusion process latent space. In all our experiments we used T 1000 steps with the fixed noise schedule Ymin = -13 and umar 5. The encoder, decoder and score network model are parameterized with 3 blocks of residual MLP layers applied on flattened 1-hot encoding representations of sequences. The MLP layers project the initial sequence representation down to a a 512 dimensional latent space. The model is simultaneously trained to optimize the diffusion loss (i.e., the score-matching loss) and the sequence reconstruction loss. SeqVDM is trained on paired OAS with the Adam optimizer and the initial learning rate of 2 x 10 for 50 epochs. The sampling is done by starting from latent vector initialized with Gaussian noise.
A.4 LANGUAGE MODEL BASELINES
We generate samples from IgLM using the prompt given in Appendix D IgLM is GPT2-style model trained to conditionally generate antibody heavy and light chains. It is therefore a strong autoregressive baseline for antibody design.
14

---

Published as a conference paper at ICLR 2024
To mimic the ab initio generation task presented in Table 2 increase the masking percentage of the masked language model ESM2 Lin et al. (202)) as high as possible (40%) and infill validation set sequences to generate containing non-amino acid characters. As a masked language model, ESM2 capable of limited infilling tasks, but it is not designed to perform true ab initio protein discovery. Indeed, it does not generate antibody-like sequences, and the high Edist and IntDiv scores are therefore meaningless. We include purely as familiar and powerful general protein language model baseline to show the gap in performance between general, pre-trained protein MLM and our methods.
samples. Beyond 40% masking, the model produces invalid sequences
A.5 EFFECT OF CHOICE OF
In Table we show sample quality results for score-based dWJS as a function of the noise level, a. We that = 0.5 produces the best quality samples, while maintaining uniqueness and diversity. Setting the noise too low (a = 0.1) leads to samples that do not capture the biophysical property distribution of the training training set (Eaist > 120). This is because at low a. the density is not smoothed and sampling does not perform well. At extremely high noise level (o 3). the samples more closely match the training distribution, but they start to show signs of "mode collapse" (lower average edit distance to the training set and lower internal diversity), because the density is over-smoothed.
(high W property)) and have extremely high edit distances from the
Table 6: a, ablikeness metrics, uniqueness, diversity for score-based dWJS.
W property +
Unique f
Edist f
IntDiv 1
0.1
0.378
1.0
120.6*
60.0
0.5
(From Table 2
0.065
1.0
62.7
65.1
3.0
0.130
0.995
44.2
30.0
A.6 ESTIMATION OF ae HYPERPARAMETER
600
500
400
Count
300
200
100
0
0.2
0.3
0.4
0.5
Xa
Figure 4: Histogram of Xw values for random samples from the paired observed antibody space
Olsen et 2022 dataset.
B ADDITIONAL ALGORIHTM
B. GRADIENT FLOW ENABLES LOCAL MINIMA FINDING
We define the gradient flow (t) log f(y(t), where sampling is performed by following the flow of the gradient of the probability density function in a deterministic dynamics, rather than stochastic Langevin dynamics. We initialize sampling from noise l 0, yo En, Eo N(0, Ta).
15

---

Published as a conference paper at ICLR 2024
and sample noisy samples (t) following the gradient flow. In this way, discover local "attractors" on the data manifold that correspond to local minima of the learned energy function. The algorithm for discrete gradient flow is given in Algo.
Algorithm 3:
Discrete gradient flow
Input:
Denoiser, (v), energy-basel model, f3( Output: Denoised samples syy. t)
yo N(0, for f = T -1 1 do
Ut+1 f- yt - 5 V. (v)(
end
fT < y/ + o² you())
return arg max äT
B.2 LANGEVIN MCMC UPDATE
Algorithm 4: Walk-jump sampling (Saremi Hyvärinen, 2019) using the discretization of Langevin diffusion by (aahss et al. (2017). Lines 6-13 correspond to walk step and line 14 is the jump step.
1: Input 5 (step size), (inverse mass), Y (friction), K (steps taken)
2: Input Learned score function ge (y) Rs log p(y) and noise level a
3: Output R â
4: Yo ~ N(0, Ia) + La((,,1)
5: to 0
6: for k == =0....; K 1 do
7:
yk++ & y/k + Uuk
8:
Jk+1 < (0(4)+1)
9:
uk+1 t Uk + T gk+1
ui
g~ s~N(0,1a))
11:
uk+1 & exp( xxp(-76)0s++ + 28
T + Vu exp(-270))
12:
yk+1 + yk+1 + 2 lk+1
end for
14: ZK t yK + (ukx)
B.3 NEURAL EMPIRICAL BAYES
Here, include additional discussion and motivation for Neural Empirical Bayes and generating discrete samples with decoupled walk and jump steps. Discrete data in this work is viewed as taking continuous values (embedded in Fuclidean space) in the NEB formalism, and by choosing large noise levels we can smooth out the original distribution greatly which makes it easier to sample from. In addition, we can use single-step denoising back discrete values. In short, the walk-jump sampling is especially well suited for discrete data.
followed the initialization scheme in (Saremi & Srivastava 2022 since In line of Algorithm ||
here the discrete data is viewed as being embedded in Euclidean space. Typically having the normal distribution as part of the initialization more important for high noise levels, and we chose the uniform distribution for simplicity.
The walk and jump steps are completely decoupled. The models used for denoising and sampling can have the same parameterization, but they can also be trained separately. Because the walk and jump steps do not depend on each other in any way, we consider them to decoupled. The energy and score are related by derivative with respect to the inputs, and any model that gives access to score can be used for sampling via Langevin MCMC. However, EBMs and score-based models are trained in completely different ways with different loss functions and architectures, the empirical performance can differ greatly. DEEN models are not trained with contrastive divergence, so the EBM formulation is completely distinct in terms of parameterization, training, and sampling.
16

---

Published as a conference paper at ICLR 2024
C
PERFORMANCE PROFILING
Table 7: Model size, sampling time and memory consumption.
Model
Parameters (M)
GPU time / /ample (ms) +
GPU memory (MB)
dWJS (energy-based
9.87
64.89
4734.76
dWJS (score-based)
1.03
327.60 12.31
16.7
VDM
940.40
246.39
DEEN
1.03
976.99
16.72
IgLM
12.89
2800
ESM2
7.8
5.25
D FEW-SHOT, IN-CONTEXT LEARNING AND PROMPTS
D.1 GPT 4 PROMPT
GPT 4 accessed on April 27, 2023. Prompt was repeated 20 times generate additional samples for statistics.
Prompt: You H3 variants of trastuzumab that were reported binders to the HER2 antigen in the paper "Optimization of therapeutic antibodies by predicting antigen specificity from antibody quence via deep learning". You will generate 100 new CDR H3 variants that you predict will also bind HERR2 Output the 100 samples python list. Here are the examples: TWWHINFYYFH'' FODHGMYQHV', YYLAFGYVEL'' WWLYHSYLFN" 'YNRYG- FYVED', 'WRKSGFYTFD' 'WANRFFYAND' "WPSCGMFALL', 'WSNYGMFVES', 'WS- MGGFYVFV', "WGQLGFFAYA", 'WPILGLYVFI', 'WHRNGMYAFD', "WPLYSMYVYK' 'WGL.GLLAYQ'.]
expert antibody engineer. I am going give you examples of CDR
Response: As an AI language model, I can't 't run simulations or generate sequences based on actual molecular interactions. However, I can use patterns found in your provided list to generate sequences that may resemble the CDR H3 variants you' ve shared. Keep in mind that these new sequences have not been validated experimentally and may not bind to HER2.
D.2
GPT 3.5 PROMPT
GPT 3.5 accessed on May 2, 2023, Prompt was repeated 5 times to generate additional samples for statistics.
expert antibody engineer. am going give you examples of antibody heavy chain variable regions from the paired observed antibody space database. You will generate 10 new antibody heavy chains that not in the database. Output the 10 samples python list. Here are the examples: CQL.OOEEEEGGAMPTTTTTSSCCCAAAAAAT
Prompt: You
NDVIROPOKKLLEWITIISSSTTT WTWAATTTAALLLLLLTSSSSSISSSESEEEE TAATTARRTVVDETCRRCCETTTTTTTTTTTT SOC.CCCCTTTTTTTTTTTTOOOOOOOOOOAO "QVQLVESSCGGVVQ" GRSLRI C.B.AA/AEAASSSSSSSSSIVVVVVPEEEES 'ADSVK- 'ADSVV
V A SAT.SSSSSSSSSSSOTTTTOOAAAASSSSSS
GLEWVSSISSSSSYIVY:
EEYCSCGGCCSONGGOCCTIYVVS:::: SWA.A//TTTTTTTTOOLEEEEOOOOOOOOOT
SKANEEESEDDDAIITTCCYCCAAAAAAAAAS 'AONFIGRVIITADKSTT
17

---

Published as a conference paper at ICLR 2024
Response: Sure, I can help you generate 10 new antibody heavy chains. Here's 's a list of 10 new antibody heavy chains that not in the given database:
These heavy chains were generated using the antibody engineering tools and techniques to create novel heavy chains. The generated sequences have similar length and structure as the provided examples, but they have different amino acid sequences, and are unique.
D.3
IGLM PROMPT
IgLM samples were generated using the prompt in the IgLM README, adjusted produce 2000 samples for comparison with other methods.
igim_ generate --prompt_sequence sequence EVQ \
--chain token [HEAVY] \\
- -sspeciessstteeee token [HUMAN] -e-nuumseeeee 2000
E in vitro VALIDATION
E.1 EXPERIMENTAL RESULTS
0.10 0.25
0.20
0,25
7.5
8.0
8.5
Exp cosrossion wiald (W- 67)
Affinity pKD (W 47 accepted)
Figure 5: Expression yield (mg) and binding affinity (pKD) of sequence designs from our method targeting the ERBB2 antigen.
E.2 EXPERIMENTAL DETAILS
in vitro validation of generated antibody proteins was performed following Hsiao et al. (202)). Discrete Walk-June Julk-Jump Sampling (dWJS) was used to generate antibody sequences, which were then expressed and purified in the laboratory. Surface plasmon resonance (SPR) measurements were used determine binding affinity.
Plasmid Construction and Antibody Production: synthesized DNA (provided by Twist Biosciences) of antibody variable domains were cloned into mammalian expression vectors using Gibson assembly. We amplified the whole vector using PrimeStar Max polymerase (Takeda). We transfected PCR products transiently ml Expi293 cell culture. Expression lasted 7 days before harvest. We affinity purified antibodies over MAb Select SuRe resin (Cytiva), and measured their concentration by optical density at 280nM.
Binding Affinity Measurements: we measured affinity of the antibodies towards their target antigen by surface plasmon resonance (SPR) at 37 "C on Biacore 8K instrument (Cytiva) HBS-EP+ buffer (10 mM Hepes, pH 7.4, 150 mM NaCI, 0.3mM EDTA and 0.05% vol/vol Surfactant P20). We captured antibodies on a Protein A chip and injected their target antigens for 5 minutes and allowed them to dissociate for 10 minutes at 30 ul/min. We regenerated the surface between cycles with 10
18

---

Published as a conference paper at ICLR 2024
mM glycine pH 1.5. We obtained affinity constants using Biacore Insight (Cytiva) using 1:1 binding kinetics model.
F FURTHER DISCUSSION OF DISTRIBUTIONAL CONFORMITY SCORE
Given a new example 2, we use the conformal transducer, A, to measure how similar 2 is to (Z1,,.,,22)) The conformal transducer is then defined as system of p-values where for each label E Y, reference sequence (21....,22) 2) E Z', and a test example & E X, we have: pd = y² ("1,,...., R1; (I,, प. Ei Ong where (Ay) .... Oyy (g 1 == A(z1, ,..,, 21, (2, v)). Intuitively, pd is the fraction of examples that have a smaller degree of conformity to the reference distribution than (æ, y)-
The difference between DCS and property alignment (W property) provides valuable insights into the p nature of the DCS statistic. DCS, being a measure of joint distribution alignment, might prioritize capturing relationships among properties as opposed to alignment of individual properties. Addi- tionally, given that DCS uses KDE, might be more influenced by extreme data points within the distribution. These factors suggest that DCS offers a distinct perspective the overall quality of generative model performance.
We considered sequence-basee properties (calculated with BioPython (Cock et al. 2009) of average hydprophilicity, molecular weight, grand average of hydropathy, as well as two structure-based properties (calculated with the Therapeutic Antibody Profiler (Raybould et al.) (209))): surface hydrophobicity patches around the CDR region and the symmetry of structural variable chain charges.
No multiple sequenee alignment pre- pre-processing of the sequences is required. For convenience and because we have small numbers of examples and low dimensions, we use kernel density estimation (KDE) to compute the joint density. However, DCS is completely general and can be combined with any density estimator.
Kernel density estimation was performed using Gaussian kernels as implemented in the open-source library awkde available at https:/gihhu..co//mmmtttoo///// com/mmnntthor/wwke We estimated the global bandwidth of the kernel using Silverman's method, set the adaptive local kernel bandwidth to 0.15, and employed a diagonal covariance matrix.
G FURTHER DISCUSSION OF RELATED WORK
training using Gibbs sampling was proposed to estimate the gradient of the log partition function, wherein input data is usually discrete and MCMC chains are initialized from training data, leading to long mixing times in high dimensions. Using continuous inputs and Langevin MCMC initialized from uniform noise with a replay buffer of past samples, efficient training was achieved (Du & Mordatch] 2019). The Langevin MCMC approach to sampling and maximum likelihood training yield advantages simplicity (only one network trained), flexibility (no constraints imposed by prior distribution), and compositionality (energy functions can be summed). Whereas our approach relies on smoothing discrete data and learning energies and scores over the smooth distribution, |Meg et al. (2023) formulates discrete score matching by constructing a faithful approximation of continuous score matching via an inductive prior on the local topology of the data space.
Contrastive divergence (Hinton 2002)
Although generative modeling is widely adopted in image and natural language generation, successful applications of generative modeling in the sciences few and far between, due the over- representation of image and text datasets, challenges in evaluation, and the need for generating samples that are novel and diverse while respecting the underlying symmetries and structure of particular domain. We consider the application of designing new molecules, focusing on therapeutic antibodies. Antibodies are proteins consisting of a heavy and light chain that can represented as discrete sequences of amino acids (AAs), which comprise standard vocabulary of 20 characters.
19

---

Published as a conference paper at ICLR 2024
20


Reviews:
[{"review_id": "5B2FVxxxKv", "reviewer": "Reviewer_aiU9", "summary": "This paper introduces a new method called discrete Walk-Jump Sampling (dWJS) for generative modeling and sampling of discrete protein sequences. The key ideas are as follows. The discrete data distribution is smoothed by adding Gaussian noise, which makes it easier to model and sample from. A discrete energy-based model (dEBM) is used to learn the distribution of noisy protein sequences and is trained via contrastive divergence. A denoising model implemented as a ByteNet is separately trained to map the noisy sequences back to the original discrete space. Sampling is performed by first using Langevin MCMC to sample noisy sequences from the dEBM. Then the denoising model is used to map these noisy samples back to valid discrete protein sequences. The walk (MCMC sampling) and jump (denoising) steps are decoupled, which provides flexibility. The authors argue this provides benefits over autoregressive models, diffusion models, and traditional EBM training. The proposed method is shown to be effective on antibody protein sequence modeling and design tasks using both in silico and in vitro experiments, \n\nOverall, I believe this is a novel contribution which demonstrates promising results on an important and challenging problem. The proposed dWJS approach is simple yet effective, providing a robust alternative to existing generative models of proteins. With additional analysis and experimental validation, this could become a leading technique for antibody design and beyond.", "questions": "Your proposed approach operates purely at the sequence level, focusing specifically on antibody sequences. What are your thoughts on the relative pros and cons of sequence-only approaches compared to structure-aware approaches that leverage protein 3D structure information (e.g. inverse folding methods)? Could structure information, either from experiments or structure prediction, be incorporated to potentially improve the performance of your model? For example, might recent powerful inverse folding techniques, where they can be viewed as structure-conditional sequence generative models, like ProteinMPNN [1], ESM-IF [2], PiFold [3], or LM-Design [4] be combined with dWJS to create even more advanced antibody design frameworks? I'm curious to hear your perspective on the value of adding structural awareness and how feasible it would be to integrate with your dWJS approach.\n\n---\n[1] Dauparas, et al. Robust deep learning–based protein sequence design using ProteinMPNN. Science 2022.  \n\n[2] Hsu, et al. Learning inverse folding from millions of predicted structures. In ICML 2022\n\n[3] Gao, et al. PiFold: Toward effective and efficient protein inverse folding. ICLR 2023.\n\n[4] Zheng, et al. Structure-informed Language Models Are Protein Designers. ICML 2023", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 4, "strengths": "(+) The proposed method is intuitive and technically sound. Decoupling the sampling and denoising steps is an elegant idea.\n\n(+)Thorough in silico evaluation using antibody-specific metrics, uniqueness, diversity, etc. In particular, a distributional conformity score  was introduced to evaluate the quality of generated protein sequences compared to a reference distribution. Extensive comparison to strong baselines.\n\n(+) Impressive wet lab validation demonstrating high expression yields and binding rates.", "weaknesses": "(-) Only a single task (antibody design) is evaluated. Testing on other protein classes or discrete domains would be useful.\n\n(-) The distributional conformity score for evaluation is introduced late with little detail. More motivation and analysis would improve clarity.\nCertain details are unclear, like how sequences are aligned and handled.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "1GxLzdwnSR", "reviewer": "Reviewer_RGZR", "summary": "The authors present a new approach for modeling discrete data distributions called Smoothed Discrete Sampling (SDS), founded on the neural empirical Bayes framework. They detail the discrete Walk-Jump sampling algorithm (dWJS), which facilitates rapid, non-autoregressive sampling and can handle variable-length discrete outputs, built upon a distinct architecture for discrete EBMs. This approach facilitates the training of score-based models for discrete data, requires just a single noise level and eliminates the need for a noise schedule. As a result, the pitfalls of diffusion models, such as brittleness, training instabilities, and sluggish sampling, are sidestepped. Furthermore, the introduced method simplifies the training of EBMs by bypassing several commonly used EBM training techniques, ensuring both swift sampling and high-quality samples. The authors evaluate their method for ab initio protein discovery and design and compares against several diffusion and LLM-based models.", "questions": "- Is the same value of σ=0.5 used for all reported experiments? It would be nice to see what effect lowering or increasing this value will have on some of the reported metrics for a better assessment of the sensitivity of the approach to this parameter.\n- How is d obtained in this sentence? \"... for the flattened sparse one-hot matrices with vocabulary size 21, d = 6237 ...\" Not sure why d is so high if d=L.\n- Please define the FID score (Fréchet inception distance (FID)) as well as the BLEU and add references.\n\n\nI thank the authors for addressing my comments and running additional experiments to demonstrate the effect of sigma. The rebuttal addresses my concern about the DCS score. I increased my score one notch.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "- The current approach builds upon foundational principles of EBMs and NEB but introduces unique methodologies for handling discrete data, formulating decoupled modeling, and tackling challenges in antibody sequence generation. In particular the study of NEB (Neural Empirical Bayes) for discrete data seems to be unique. \n\n- The current paper distinguishes its approach from discrete diffusion models like those by Austin et al. (2023), who learn an iterative denoising process over different noise levels. Although generative modeling has been previously applied to antibodies (cited works include Shuai et al., Gligorijević et al., Ferruz & Höcker, and Tagasovska et al.), the present work highlights the unique challenges faced due to limited training data and the complexities of antibody sequences. The paper suggests that typical natural-language-based methods might struggle in this domain, pointing towards a differentiation from those methods. The current work introduces a novel formulation of decoupled energy- and score-based modeling to address the challenges in training and sampling discrete sequences, which seems distinct from previously mentioned methods.\n\n- A new distributional conformity score (DCS) is proposed, which is useful for evaluating generated samples in comparison to a reference set. DCS provides a measure of joint distribution alignment, enabling it to capture relationships among properties rather than just aligning individual properties. This could mean that DCS provides a more comprehensive assessment of how well a generative model captures the nuances and interconnectedness of the properties within the data.", "weaknesses": "- While the paper touts the simplification to a single hyperparameter choice (noise level, σ) as a strength, it might also be seen as a limitation since the entire model's performance could be sensitive to this single parameter. \n\n- Set of properties includes both continuous, binary, and discrete values and estimating the distribution by a kernel density approach may not be very effective. Also DCS may be overly influenced by outliers and extreme data points in the dataset.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "ez6LPPZPOL", "reviewer": "Reviewer_BFxy", "summary": "The paper introduces an interesting, novel approach to sampling new, diverse antibody sequences and validates the method's performance using real wet-lab experiments with real antibodies. They explore a couple variations on energy-based models with some interesting tricks-of-the-trade such as the pretrained denoiser to help guide the sampling. The experimental results are strong.\n\n\n==After author's response==\nI appreciate all of the details in response to my question. I continue to advocate for acceptance of this high-quality paper. My only final feedback is that the discussion of mixing should mention that mixing is not a concern for autoregressive models, since for them it's trivial to draw independent samples.", "questions": "The idea of relaxing the discrete problem to a continuous modeling problem on one-hots is interesting. I was surprised, however, that you treated them as one-hots instead of elements of a probability simplex. Have you explored this? A key advantage is that you could sample from the probabilities rather than taking the argmax in the final step.\n\nOne of my concerns for using your method vs. LMs is that there are a number of difficult-to-set hyperparameters. Can you discuss the sensitivity to these choices? How did you validate them?\n\nThe DCS seems like a good way to filter samples. Is there a way that you could evaluate the impact of such filtering in-silico?\n\nThe DCS is a one-vs-rest comparison, using a single query vs. a reference set. How could you extend this to generate a *diverse* set of samples?\n\nThe paper claims informally that the sampling mixes quickly. Is there a way to demonstrate this quantitatively?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "The paper does real wet-lab experiments! By that, I don't mean just that experiments were performed in a wet lab in a toy setting, but instead that the authors did real-world antibody design. \n\nThe paper provides an interesting/refreshing take on how to sample antibodies that is quite different from most papers' use of language models. It's drawing on a whole line of work on energy-based models with details that are both old-school and new-school. There will be many ML+Bio researchers that will be unfamiliar with this background material and will benefit from seeing this paper. \n\nI thought it was interesting that they included GPT3.5 as a baseline. Sometimes the results are surprisingly reasonable! As generalist foundation models get stronger, such a baseline seems important.\n\nThe authors have release code for their modeling, which is key because it is more niche than LM approaches for which there are standard packages.", "weaknesses": "Some of the comparisons to autoregressive language models (LMs) were a bit under-developed. For example, the paper mentions fast mixing sampling, but this is trivial for LMs. Second, LMs provide a natural way to condition on metadata; how would the current model do that? Finally, claims that LM sampling is expensive in terms of compute aren't that compelling, since for these antibody design applications the wet-lab experiments are orders of magnitude more expensive than the compute to generate the library.\n\n\nThe 'distributional conformity scores' (DCS) section was interesting and cool, but not central to the paper and felt like a way to add extra math into the paper. Please clarify: was this used to filter any of the samples for the wet-lab experiments or was it only used to provide the rightmost column in Table 2? Given that your proposed does not have strong DCS scores compared to some other baselines, how should I be interpreting these results? Overall, I would recommend removing the DCS content and allocating more space in the main paper to details/derivations of the model fitting/sampling.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (oral)