# Protein Discovery with Discrete Walk-Jump Sampling

Nathan C. Frey1, Daniel Berenberg1,3,∗ Karina Zadorozhny1 Joseph Kleinhenz1
Julien Lafrance-Vanasse2 Isidro Hötzel2 Yan Wu2
Stephen Ra1 Richard Bonneau1 Kyunghyun Cho1,3,4
Andreas Loukas1 Vladimir Gligorijević1 Saeed Saremi1
1Prescient Design Equal contribution  Genentech 2Antibody Engineering
Genentech
3Department of Computer Science  New York University
4Center for Data Science  New York University

###### 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 Discrete Walk-Jump
Sampling formalism combines the maximum likelihood 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 optimizing and sampling from our models for the proposed
distributional conformity score, 97-100% of generated samples are successfully
expressed and purified and 35% 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.

###  1 Introduction

Discrete sequence generation poses a 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., 2006) fit an energy function that specifies a 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 et al., 2015; Ho et
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 is 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^{L}$ for
proteins of length $L$. Here, we consider the specific problem of generative
modeling of antibodies, a class of proteins with highly conserved structure
that are of immense interest for therapeutics. In 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 et al., 2023). Antibodies consist of well-conserved domains and high-
entropy variable regions, so leveraging evolutionary information from pre-
trained protein language models is not an immediate solution. We distinguish
ab initio protein discovery and design (producing novel, functional proteins
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 et
al., 2022). Our goal here is to invent an efficient, non-autoregressive
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 a novel
algorithm, discrete 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.111https://github.com/Genentech/walk-jump We also design a metric
called the Distributional Conformity Score (DCS), which is 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 to learn noisy data
distributions and sample discrete data. dWJS enables fast, non-autoregressive
sampling with variable length discrete outputs. We also design a novel
architecture for discrete EBMs.

  * •

Our method simplifies score-based model training for discrete data by
requiring only a 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, $\ell_{2}$ 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 our
method with in vitro experiments. Our method outperforms autoregressive large
language models, diffusion, and score-based baselines.

Figure 1:  Selected samples from a single Markov chain Monte Carlo sampling
run. Protein color corresponds to different antibody germlines (classes).
Discrete walk-jump sampling exhibits fast mixing and explores diverse modes of
the distribution in a single chain.

###  2 Background

####  2.1 Energy-based models

EBMs are a class of models that learn an energy function
$f_{\theta}:\mathcal{X}\rightarrow\mathbb{R}$ mapping inputs $x$ (in
$\mathbb{R}^{d}$) to a scalar “energy" value. The data distribution is
approximated by the Boltzmann distribution

$\displaystyle p_{\theta}(x)\propto e^{-f_{\theta}(x)}.$

EBMs are typically trained via contrastive divergence (Hinton, 2002), and new
samples are drawn from $p_{\theta}(x)$ by Markov-Chain Monte Carlo (MCMC).
Details of the loss function used in this work are given in Section 3. In
Langevin MCMC, samples are initialized from a known data point or random noise
and refined with (discretized) Langevin diffusion

$\displaystyle x_{k+1}=x_{k}-\delta\,\nabla
f_{\theta}(x_{k})+\sqrt{2\delta}\,\varepsilon_{k},\
\varepsilon_{k}\sim\mathcal{N}(0,I_{d}),$ (1)

where $\nabla$ denotes the gradient of the energy function with respect to
inputs, $k$ is the sampling step, $\delta$ is the (discretization) step size,
and the noise $\varepsilon_{k}$ 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+\mathcal{N}(0,\sigma^{2}I_{d}).$ The least-squares estimator of $X$ given
$Y=y$ is given by (Robbins, 1956; Miyasawa, 1961)

$\displaystyle\hat{x}(y)=y+\sigma^{2}\nabla\log p(y),$ (2)

where $p(y)=\int p(y|x)p(x)dx$ is the probability distribution function of the
smoothed density.222 We follow the convention $p(x)\coloneqq
p_{X}(x),p(y)\coloneqq p_{Y}(y)$, etc.  This estimator is often expressed
directly in terms of $g(y)=\nabla\log p(y)$ known as the score function
(Hyvärinen, 2005) which is parameterized with a neural network denoted by
$g_{\phi}:\mathbb{R}^{d}\rightarrow\mathbb{R}^{d}.$ The least-squares
estimator then takes the following parametric form:

$\displaystyle\hat{x}_{\phi}(y)=y+\sigma^{2}g_{\phi}(y).$ (3)

Putting this all together leads to the following learning objective

$\displaystyle\mathcal{L}(\phi)=\mathbb{E}_{x\sim p(x),y\sim
p(y|x)}\|x-\hat{x}_{\phi}(y)\|^{2},$ (4)

which is optimized with stochastic gradient descent. Notably, no MCMC sampling
is required during learning. In short, the objective is “learning to denoise”
with an empirical Bayes formulation.

###  3 Antibody discovery and design

####  3.1 Discrete walk-jump sampling

Figure 2:  Discrete walk-jump sampling. a The noising and denoising process is
applied to antibody proteins. b Discrete inputs $x$ are smoothed with
isotropic Gaussian noise, $\varepsilon\sim\mathcal{N}(0,\sigma^{2}I_{d})$, to
noisy inputs, $y=x+\varepsilon$. A discrete energy-based model (dEBM)
parameterizes the energy function $f_{\theta}(y)$ of noisy data. Noisy data is
sampled with the energy function, and denoised with a separate denoising
ByteNet network to clean samples, $\hat{x}_{\phi}(y)$. c The “walk" sampling
steps on 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 a 1d positional encoding, $p_{1d}$, passes through an
MLP and a 3 layer CNN, and concatenates the outputs with an embedding $z_{s}$
of the inputs into a hidden state, $h$. $h$ is passed through an MLP and
returns the energy $f_{\theta}(y)$.

Following training of the denoising network, $g_{\phi}$, one can sample noisy
data using the learned score function $g_{\phi}(y)$ with Langevin MCMC
(replace $-\nabla f$ with $g$ in Eq. 1). For any such draws $y_{k}$, clean
samples from the true data manifold $\mathcal{M}$ are obtained by “jumping”
back to $\mathcal{M}$ with the least-squares estimator
$\hat{x}_{\phi}(y_{k})=y_{k}+\sigma^{2}g_{\phi}(y_{k}).$ 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 to train an EBM with maximum
likelihood estimation on 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. 1. Our algorithm is general
and applies to any discrete sequence inputs of a fixed vocabulary. In Fig. 1
we show samples generated from a single chain of MCMC. Unlike a diffusion
model, every sample along the chain collectively forms a valid set of samples
from the underlying distribution, because of the decoupled walk (sampling) and
jump (denoising) steps. dWJS also produces fast-mixing chains, such that many
diverse modes (protein classes) are sampled in a single chain. Samples are
folded with EquiFold (Lee et al., 2022) for visualization and confirmation of
structural validity.

The EBM is trained by maximizing the log-likelihood of noisy data under the
model:

$\displaystyle\operatorname*{arg\,max}_{\theta}\mathbb{E}_{y\sim p_{Y}}[\log
p_{\theta}(y)]=\operatorname*{arg\,max}_{\theta}\left({\mathbb{E}}_{y^{-}\sim
p_{\theta}(y)}[f_{\theta}(y^{-})]-{\mathbb{E}}_{y^{+}\sim
p_{Y}}[f_{\theta}(y^{+})]\right),$ (5)

where $y^{+}$ are noisy training data and $y^{-}$ are noisy data sampled from
the model.

With this objective, the model aims to decrease the energy of noisy training
data (“positive” samples $y^{+}$) 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:

$\begin{split}\nabla_{\theta}\log
p_{\theta}(y)&=-\nabla_{\theta}f_{\theta}(y)-\nabla_{\theta}\log Z(\theta)\\\
&=-\nabla_{\theta}f_{\theta}(y)+\frac{\int\nabla_{\theta}f_{\theta}(y)e^{-f_{\theta}(y)}dy}{Z(\theta)}\\\
&=-\nabla_{\theta}f_{\theta}(y)+\int\nabla_{\theta}f_{\theta}(y)\cdot
p_{\theta}(y)dy\\\ &=-\nabla_{\theta}f_{\theta}(y)+\mathbb{E}_{y\sim
p_{\theta}(y)}[\nabla_{\theta}f_{\theta}(y)],\end{split}$ (6)

where $Z(\theta)=\int e^{-f_{\theta}(y)}dy$ is the partition function.

Input: Denoiser, $g_{\phi}(y)$, energy-based model, $f_{\theta}(y)$

Output: Noisy samples $y\sim p(y)$, denoised samples $\hat{x}(y)$

1

2$y_{0}\sim{\rm Unif}([0,1]^{d})+\mathcal{N}(0,\sigma^{2}I_{d})$

3 for  _$t=0,\dots,T-1$_ do

4 $y_{t+1}\leftarrow
y_{t}-\delta\,\nabla_{y}f_{\theta}(y_{t})+\sqrt{2\delta}\varepsilon_{t}$ ,
$\varepsilon_{t}\sim\mathcal{N}(0,I_{d})$

5

6 end for

7$\hat{x}_{T}\leftarrow y_{T}+\sigma^{2}g_{\phi}(y_{T})$

8 return _$\operatorname*{arg\,max}\hat{x}_{T}$_

Algorithm 1 Discrete Walk-Jump Sampling

##### 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.
1 takes advantage of both energy- and score-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, are summarized in Table 1. Discrete Walk-Jump
Sampling refers to walking with the EBM, $f_{\theta}(y)$, and jumping with the
denoising network, $g_{\phi}(y)$. Score-based dWJS uses $g_{\phi}(y)$ for both
walking and jump steps. The Deep Energy Estimator Network (DEEN) (Saremi et
al., 2018) approach uses a 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, dWJS-EBM uses an EBM for sampling and the gradient of the energy,
$\nabla f_{\theta}(y)$, 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 Walk-Jump
Sampling.

Table 1: Smoothed Discrete Sampling implementations. Model | Walk (sampling) | Jump (denoising)
---|---|---
dWJS (energy-based) | EBM | Denoiser
dWJS (score-based) | Denoiser | Denoiser
Deep Energy Estimator Network | Denoiser (energy) | Denoiser (energy)
dWJS-EBM | EBM | EBM

##### Variable length protein sequence generation.

We represent antibody protein molecules as
$\mathbf{}{x}=(x_{1},\ldots,x_{d})$, where $x_{l}\in\\{1,\ldots,20\\}$
corresponds to the amino acid (AA) type at position $l$. Sequences from the
Observed Antibody Space (OAS) database (Olsen et al., 2022) are aligned
according to the AHo numbering scheme (Honegger & PluÈckthun, 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 to 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 to a 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)\times 21$. An EBM is trained via
contrastive divergence on the manifold of smoothed, noisy one-hot encodings,
$y$, given by $y=x+\varepsilon$,
$\varepsilon\sim\mathcal{N}(0,\sigma^{2}I_{d})$, where $x\in\\{0,1\\}^{d}$. A
separate denoising model is trained with the objective in Eq. 4. New antibody
sequences are generated (Fig. 2) by sampling noisy samples with Langevin MCMC
following gradients from the EBM, denoising with the least-squares estimator,
and taking ${\rm argmax}\>\hat{x}$ to recover a 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

$P^{\top}\hat{x}(y,t)=P^{\top}s$

for $P\in\\{0,1\\}^{L\times c}$, where $c$ is the number of conserved tokens
in the sequence, $y$ is the noisy sequence at time step $t$ of Langevin MCMC,
$\hat{x}(y,t)$ is the denoised sample at time $t$, and $s$ 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 4, we must choose what noise level,
$\sigma$, to use for training. Empirically, we find that in the protein
discovery setting, $\sigma\geq 0.5$ is sufficient for getting good sample
quality. Here, we provide some intuition for choosing a good $\sigma$, based
on a geometric picture of the concentration of the measure (Saremi &
Hyvärinen, 2019). We define the matrix $\chi$ with entries

$\chi_{ii^{\prime}}=\frac{||X_{i}-X_{i^{\prime}}||}{2\sqrt{d}},$ (7)

where $d$ is the dimension of the data and the $\frac{1}{2\sqrt{d}}$ scaling
comes from the concentration of isotropic Gaussians in high dimensions. The
critical noise level, $\sigma_{c}$, is defined as

$\displaystyle\sigma_{c}=\max_{ii^{\prime}}\chi_{ii^{\prime}}$

such that for $\sigma>\sigma_{c}$, all noisy data have some degree of overlap.
For our antibody sequence data, the statistics of the $\chi$ matrix are given
in Table 2 and the histogram of $\chi_{ii^{\prime}}$ values is shown in
Appendix A.4. We find that $\sigma_{c}\approx 0.5$, which agrees with our
empirical hyperparameter optimization. Estimating $\sigma_{c}$ in this way
serves to motivate the empirical success of the $\sigma$ used in our
experiments, and provides helpful guidance on the scale of $\sigma$ to use for
discrete data. Here we take $d$ to be the length of the input vector
($d=L=297$ for aligned antibody sequences); for the flattened sparse one-hot
matrices with vocabulary size 21, $d=6237$. This scales $\sigma_{c}$ by 0.22,
which still gives a useful scale for $\sigma$, but is not optimal because of
the sparsity of the one-hot matrices.

Table 2: Statistics of the distance matrix, $\chi$, for discrete antibody sequence data. | min | median | mean | max $\approx\sigma_{c}$
---|---|---|---|---
$\chi$ | 0.17 | 0.42 | 0.41 | 0.51
Figure 3:  in silico designs sampled with dWJS are compared to a reference set
of validation samples. Distributions are characterized with a set of sample
quality metrics. Joint density estimation is used to compute the likelihood of
designs versus the validation set and likelihoods are condensed into a
distributional conformity score that characterizes the similarity of generated
samples to the reference set.

####  3.3 Distributional conformity score

The FID score in computer vision and metrics like the BLEU score in machine
translation greatly simplify the evaluation of proposed methods; protein
generation lacks such metrics, which motivates us to introduce the
“distributional conformity score” (Fig. 3). The goal of the distributional
conformity score is to provide a succinct description of how likely generated
samples are with respect to a reference distribution, while maintaining
novelty and diversity. Distributional conformity score is designed such that
improving sample quality corresponds directly to increased probability of
generating real, biophysically valid proteins.

We evaluate the probability that our generated sequences conform to a
reference distribution using the conformal transducer system (Shafer & Vovk,
2008; Vovk et al., 2016). Let $\mathcal{X}\in\mathbb{R}^{d}$,
$\mathcal{Y}\in\mathbb{R}$, and
$\mathbf{Z}=\mathcal{X}\times\mathcal{Y}$.333In the discussion of
distributional conformity score, $x$ refers to sample features; elsewhere in
the paper $x$ refers to clean data. Here, $y$ refers to labels; elsewhere in
the paper $y$ refers to noisy data. A conformity measure A is a measurable
function that maps a sequence $(z_{1},\dots,z_{n})\in\mathbf{Z}^{n}$ to a set
of real numbers $(\alpha_{1},\dots,\alpha_{n})$ and is equivariant under
permutations. Given a new example $z$, we use A to measure how similar $z$ is
to $(z_{1},\dots,z_{n})$. The conformal transducer is then defined as a system
of p-values where for each label $y\in\mathcal{Y}$, a reference sequence
$(z_{1},\dots,z_{l})\in\mathbf{Z}^{l}$, and a test example $x\in X$, we have:
$p^{y}:=p^{y}(z_{1},\dots,z_{l},(x,y))=\frac{1}{l+1}\sum_{i=1}^{l+1}[\alpha_{y_{i}}<\alpha_{y_{l+1}}]$
where
$(\alpha_{y_{1}},...,\alpha_{y_{l}},\alpha_{y_{l+1}}):=A(z_{1},...,z_{l},(x,y))$.
Intuitively, $p^{y}$ is the fraction of examples that have a higher degree of
conformity to the reference distribution than $(x,y)$. Here, we define $A$ to
be the likelihood under the join density over various properties, including
biophysical properties and statistical properties, such as a log-probability
under a protein language model. We use kernel density estimation to compute
the joint density. To avoid overfitting the estimator, we split the reference
set into a fitting set and a validation set, with the latter used to compute
the p-values (Algo. 2). In our context, the reference distribution
$\mathcal{D}$ comprises all antibodies and the label $y$ represents the
property of interest such as expression or binding.

Input: Reference distribution $\mathcal{D}$, test example $x\in\mathcal{X}$,
conformity measure $A$, label $y$

1

Output: p-value $p^{y}$ (the fraction of validation examples that more similar
to $D_{z|y}$ than $x$)

2

3Sample $(z_{1},\dots,z_{n}),z_{i}\sim D_{z|y}$ and a held-out validation set
$(\tilde{z}{1},\dots,\tilde{z}{k-1}),\tilde{z}_{i}\sim D_{z|y}$

4Set $\tilde{z}_{k}\leftarrow(x,y)$

5 for  _$i=1$ to $k$_ do

6 $\alpha_{i}\leftarrow A(z_{1},\dots,z_{n-1},\tilde{z}_{i})$

7 end for

8$p^{y}\leftarrow\frac{1}{k}\sum_{i=1}^{k}[\alpha_{i}<\alpha_{k}]$

9

10return _$p^{y}$_

Algorithm 2 Distributional conformity scores for evaluating generated designs

###  4 Experiments

We evaluate our method, discrete Walk-jump sampling (dWJS) (Fig. 2), on three
antibody generation tasks: 1) distribution learning on paired observed
antibody space (Olsen et al., 2022); 2) the in vitro expression and
purification of novel antibodies; and 3) most importantly, functional
therapeutic antibody design (Mason et al., 2021). Crucially, we compare
methods using our distributional conformity score, which is a sample-to-
distribution metric to assess sample quality (analogous to an FID score),
rather than the sequence recovery metrics used in previous antibody design
work (Kong et al., 2022; Jin et al., 2021). Sequence recovery is a poor
objective for our goal, which is the discovery of novel (large edit distance
from known examples), functional antibodies. We also present surprising
empirical results related to EBM parameter optimization with maximum
likelihood training on noisy data that suggest a more general and robust
training procedure for EBMs. Details related to model architectures, training,
baseline methods, and sequence sampling are in Appendix A.

Table 3: Ablikeness metrics, uniqueness, diversity, and distributional conformity scores. Model | $W_{{\rm property}}\downarrow$ | Unique $\uparrow$ | $E_{{\rm dist}}\uparrow$ | ${\rm IntDiv}\uparrow$ | $\rm DCS\uparrow$
---|---|---|---|---|---
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
SeqVDM | 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

####  4.1 dWJS generates natural, novel, diverse antibodies in silico

We measure generative model performance with a 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 a
single scalar metric by computing the distributional conformity score and the
normalized average Wasserstein distance $W_{{\rm property}}$ between the
property distributions of samples and a validation set. The average total edit
distance $E_{{\rm dist}}$ summarizes the novelty and diversity of samples
compared to the validation set, while internal diversity (${\rm IntDiv}$)
represents the average total edit distance between samples. Our method
achieves strong ablikeness results (Table 3), simply by increasing $\sigma$ to
$\geq 0.5$. dWJS with dEBM sampling achieves the best agreement with the
validation set property distribution and highest percentage of unique samples,
while dWJS with score-based sampling has the best distributional conformity
score, novelty, and diversity. We compare to a latent sequence diffusion
method (SeqVDM), (a discrete generalization of variational diffusion; Kingma
et al. 2021), a score-based model with an energy parameterization (DEEN), and
a pre-trained large language model (LLM) (GPT 3.5). Our dWJS methods have
faster sampling time and lower memory footprint than diffusion and score-based
baselines (Table 6), while also having high sample quality. Details on the
baseline methods and GPT 3.5 prompts are given in Appendices A and D.

Table 4: Protein synthesis. Model | $p_{\rm expression}\uparrow$
---|---
dWJS (score-based) | 1.0
dWJS (energy-based) | 0.97
EBM | 0.42

####  4.2 dWJS 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 4). We achieved the 97.47% in
vitro success rate by developing dWJS to capture the antibody distribution in
silico as measured by 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 a small
number of edits (< 4) can destroy expression if the proposal distribution
(generative model) is poorly optimized.

Table 5: Predicted and measured antibody binding affinity. Model | $p_{\rm bind}\uparrow$ | ${\rm total_{bind}}\uparrow$ | ${\rm improved_{bind}}\uparrow$
---|---|---|---
dWJS (energy-based) | 0.96 | 0.34 | 0.35
dWJS (score-based) | 0.95 | N/A | N/A
SeqVDM | 0.75 | 0.19 | 0.0
GPT 4 | 0.74 | N/A | N/A
Transformer | 0.60 | N/A | N/A
EGNN | 0.58 | N/A | N/A

####  4.3 dWJS generates functional antibody variants in vitro

To further show the robustness of our method, we consider the task of training
generative models on a hu4D5 antibody mutant dataset (Mason et al., 2021) and
compare to baseline models. The dataset consists of 9k binding 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^{13}$ dimensional search space. The mutants were measured in lab
experiments to determine their binding to HER2 antigen. The goal of this
benchmark task is to produce samples that are also predicted to bind to HER2.
We trained dWJS models (score-based and energy-based) on only the binder set
at a noise level of $\sigma=0.5$, while also training a 1D-CNN binary
classifier to classify binders and non-binders. The classifier achieves 86%
accuracy on an IID validation split. Then, we classified 1000 samples from
each dWJS generative model and four baseline models trained on the hu4D5
mutant dataset. We compare to three diffusion models: 1) a 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) that
codesigns $(\rm sequence,\rm structure)$, and 3) a latent sequence diffusion
model, SeqVDM; and a pre-trained LLM, GPT 4. The specific prompt used for GPT
4 is given in Appendix D. The probability of binding for unique designs from
each model is reported in Table 5. dWJS models produce the highest percentage
of unique predicted binders.

We also report in vitro wetlab validation results for active drug discovery
projects.444 Due to the sensitive nature of data, we do not disclose specific
drug targets for these experiments.  dWJS produces the highest percentage of
functional antibodies that bind to target (${\rm total_{bind}}$ in Table 5),
and the highest percentage of antibodies that bind with equal or better
affinity compared to previously known binders (${\rm improved_{bind}}$ in
Table 5). Further details on wetlab experiments are presented in Appendix E

###  5 Related Work

Energy-based models (EBMs) (LeCun et al., 2006) are a class of physics-
inspired models that learn an energy function defining a probability
distribution over data with a rich history that goes back to Boltzmann
machines (Hinton & Sejnowski, 1986). Contrastive divergence (Hinton, 2002)
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 in simplicity (only one network is
trained), flexibility (no constraints imposed by a prior distribution), and
compositionality (energy functions can be summed).

Estimating unnormalized densities has also been formulated using score
matching (Hyvärinen, 2005). This formulation led to probabilistic models for
denoising autoencoders (Vincent, 2011; Alain & Bengio, 2014; Saremi et al.,
2018), but also has an empirical Bayes interpretation that is most related to
this work. In particular, the neural empirical Bayes (NEB) (Saremi &
Hyvärinen, 2019) formalism unifies kernel density estimation (Parzen, 1962)
and empirical Bayes (Robbins, 1956) to transform the unsupervised learning
problem into a more tractable form where a neural network energy function is
parameterized to capture a smoothed data distribution. Our work is the first
study of the NEB formalism for discrete data. Whereas our approach relies on
smoothing discrete data and learning energies and scores over the smooth
distribution, (Meng 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. Additionally,
discrete diffusion models such as (Austin et al., 2023) learn an iterative
denoising process over many different noise levels by prescribing a noise
process over discrete data that converges to a known categorical distribution.

Although generative modeling is widely adopted in image and natural language
generation, successful applications of generative modeling in the sciences are
few and far between, due to 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 a 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 be represented as discrete
sequences of amino acids (AAs), which comprise a standard vocabulary of 20
characters. Approaches borrowing from traditional ML 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-language-
based methods struggle to capture the data distribution of antibodies, for
which there is limited training data ($\sim$1K - 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 score-based modeling.

###  6 Conclusions

We proposed Smoothed Discrete Sampling (SDS), a 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 as 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 retraining and sampling with low compute requirements. dWJS
discards many of the commonly used techniques for improving EBM training with
Langevin MCMC (replay buffers, $\ell_{2}$ norm penalty, simulated annealing,
rejection sampling, etc.) and reduces the engineering complexity of training
EBMs and diffusion-based models to a single hyperparameter choice: the noise
level, $\sigma$. 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.

### Acknowledgments and Disclosure of Funding

The authors acknowledge the entire Prescient Design team and the Antibody
Engineering department at Genentech for providing helpful discussions and
input that contributed to 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 Goodman for consulting on large language model benchmarks.

### References

  * Alain & Bengio (2014) Guillaume Alain and Yoshua Bengio.  What regularized auto-encoders learn from the data-generating distribution.  _Journal of Machine Learning Research_ , 15(1):3563–3593, 2014.
  * Austin et al. (2023) Jacob Austin, Daniel D. Johnson, Jonathan Ho, Daniel Tarlow, and Rianne van den Berg.  Structured denoising diffusion models in discrete state-spaces, 2023.
  * Cock et al. (2009) Peter JA Cock, Tiago Antao, Jeffrey T Chang, Brad A Chapman, Cymon J Cox, Andrew Dalke, Iddo Friedberg, Thomas Hamelryck, Frank Kauff, Bartek Wilczynski, et al.  Biopython: freely available python tools for computational molecular biology and bioinformatics.  _Bioinformatics_ , 25(11):1422–1423, 2009.
  * Devlin et al. (2018) Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova.  Bert: Pre-training of deep bidirectional transformers for language understanding.  _arXiv preprint arXiv:1810.04805_ , 2018.
  * Du & Mordatch (2019) Yilun Du and Igor Mordatch.  Implicit generation and modeling with energy based models.  _Advances in Neural Information Processing Systems_ , 32, 2019.
  * Dunbar & Deane (2016) James Dunbar and Charlotte M Deane.  ANARCI: antigen receptor numbering and receptor classification.  _Bioinformatics_ , 32(2):298–300, 2016.
  * Ferruz & Höcker (2022) Noelia Ferruz and Birte Höcker.  Controllable protein design with language models.  _Nature Machine Intelligence_ , 4(6):521–532, 2022.
  * Gligorijević et al. (2021) Vladimir Gligorijević, Daniel Berenberg, Stephen Ra, Andrew Watkins, Simon Kelow, Kyunghyun Cho, and Richard Bonneau.  Function-guided protein design by deep manifold sampling.  _bioRxiv_ , 2021.
  * Hinton (2002) Geoffrey E Hinton.  Training products of experts by minimizing contrastive divergence.  _Neural computation_ , 14(8):1771–1800, 2002\.
  * Hinton & Sejnowski (1986) Geoffrey E Hinton and Terrence J Sejnowski.  Learning and relearning in Boltzmann machines.  _Parallel distributed processing: Explorations in the microstructure of cognition_ , 1(282-317):2, 1986.
  * Ho et al. (2020) Jonathan Ho, Ajay Jain, and Pieter Abbeel.  Denoising diffusion probabilistic models.  _Advances in Neural Information Processing Systems_ , 33:6840–6851, 2020.
  * Honegger & PluÈckthun (2001) Annemarie Honegger and Andreas PluÈckthun.  Yet another numbering scheme for immunoglobulin variable domains: an automatic modeling and analysis tool.  _Journal of molecular biology_ , 309(3):657–670, 2001.
  * Hsiao et al. (2020) 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.
  * Hyvärinen (2005) Aapo Hyvärinen.  Estimation of non-normalized statistical models by score matching.  _Journal of Machine Learning Research_ , 6(Apr):695–709, 2005.
  * Jin et al. (2021) Wengong Jin, Jeremy Wohlwend, Regina Barzilay, and Tommi Jaakkola.  Iterative refinement graph neural network for antibody sequence-structure co-design.  _arXiv preprint arXiv:2110.04624_ , 2021.
  * Kalchbrenner et al. (2016) Nal Kalchbrenner, Lasse Espeholt, Karen Simonyan, Aaron van den Oord, Alex Graves, and Koray Kavukcuoglu.  Neural machine translation in linear time.  _arXiv preprint arXiv:1610.10099_ , 2016.
  * Kim et al. (2023) 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.
  * Kingma et al. (2021) 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.
  * Kong et al. (2022) Xiangzhe Kong, Wenbing Huang, and Yang Liu.  Conditional antibody design as 3d equivariant graph translation.  _arXiv preprint arXiv:2208.06073_ , 2022.
  * LeCun et al. (2006) Yann LeCun, Sumit Chopra, Raia Hadsell, M Ranzato, and Fujie Huang.  A tutorial on energy-based learning.  _Predicting structured data_ , 1(0), 2006.
  * Lee et al. (2022) 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 a novel coarse-grained structure representation.  _bioRxiv_ , pp. 2022–10, 2022.
  * Loshchilov & Hutter (2017) Ilya Loshchilov and Frank Hutter.  Decoupled weight decay regularization.  _arXiv preprint arXiv:1711.05101_ , 2017.
  * Mason et al. (2021) 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–612, 2021.
  * Meng et al. (2023) Chenlin Meng, Kristy Choi, Jiaming Song, and Stefano Ermon.  Concrete score matching: Generalized score matching for discrete data, 2023.
  * Miyasawa (1961) Koichi Miyasawa.  An empirical Bayes estimator of the mean of a normal population.  _Bulletin of the International Statistical Institute_ , 38(4):181–188, 1961.
  * Olsen et al. (2022) 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.
  * Parzen (1962) Emanuel Parzen.  On estimation of a probability density function and mode.  _The annals of mathematical statistics_ , 33(3):1065–1076, 1962.
  * Paszke et al. (2019) 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 DeVito, 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.
  * Robbins (1956) Herbert Robbins.  An empirical Bayes approach to statistics.  In _Proc. Third Berkeley Symp._ , volume 1, pp. 157–163, 1956.
  * Romero & Arnold (2009) Philip A Romero and Frances H Arnold.  Exploring protein fitness landscapes by directed evolution.  _Nature reviews Molecular cell biology_ , 10(12):866–876, 2009.
  * Sachs et al. (2017) Matthias Sachs, Benedict Leimkuhler, and Vincent Danos.  Langevin dynamics with variable coefficients and nonconservative forces: from stationary states to numerical methods.  _Entropy_ , 19(12):647, 2017.
  * Saremi & Hyvärinen (2019) Saeed Saremi and Aapo Hyvärinen.  Neural empirical Bayes.  _Journal of Machine Learning Research_ , 20:1–23, 2019\.
  * Saremi et al. (2018) Saeed Saremi, Arash Mehrjou, Bernhard Schölkopf, and Aapo Hyvärinen.  Deep energy estimator networks.  _arXiv preprint arXiv:1805.08306_ , 2018.
  * Satorras et al. (2021) Vıctor Garcia Satorras, Emiel Hoogeboom, and Max Welling.  E(n) equivariant graph neural networks.  In _International conference on machine learning_ , pp. 9323–9332. PMLR, 2021.
  * Shafer & Vovk (2008) Glenn Shafer and Vladimir Vovk.  A tutorial on conformal prediction.  _Journal of Machine Learning Research_ , 9(12):371–421, 2008.  URL http://jmlr.org/papers/v9/shafer08a.html.
  * Shuai et al. (2021) Richard W Shuai, Jeffrey A Ruffolo, and Jeffrey J Gray.  Generative language modeling for antibody design.  _bioRxiv_ , 2021.
  * Sohl-Dickstein et al. (2015) 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.
  * Tagasovska et al. (2022) Nataša Tagasovska, Nathan C Frey, Andreas Loukas, Isidro Hötzel, Julien Lafrance-Vanasse, Ryan Lewis Kelly, Yan Wu, Arvind Rajpal, Richard Bonneau, Kyunghyun Cho, et al.  A pareto-optimal compositional energy-based model for sampling and optimization of protein sequences.  _arXiv preprint arXiv:2210.10838_ , 2022.
  * Vincent (2011) Pascal Vincent.  A connection between score matching and denoising autoencoders.  _Neural computation_ , 23(7):1661–1674, 2011\.
  * Vovk et al. (2016) Vladimir Vovk, Valentina Fedorova, Ilia Nouretdinov, and Alexander Gammerman.  Criteria of efficiency for conformal prediction.  _CoRR_ , abs/1603.04416, 2016.  URL http://arxiv.org/abs/1603.04416.
  * Yang et al. (2022) Kevin K. Yang, Alex X. Lu, and Nicolo Fusi.  Convolutions are competitive with transformers for protein sequence pretraining.  _bioRxiv_ , 2022.  doi: 10.1101/2022.05.19.492714.  URL https://www.biorxiv.org/content/early/2022/05/25/2022.05.19.492714.

## Appendix

###  Appendix 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, ReLU non-
linearities and an output linear layer of size 128. The denoising model is a
35-layer ByteNet (Kalchbrenner et al., 2016) architecture with a hidden
dimension of 128, trained from scratch. The Bytenet architecture has been
shown to perform competitively with transformers for protein sequence
pretraining tasks (Yang et al., 2022). All models were trained with the AdamW
(Loshchilov & Hutter, 2017) optimizer in PyTorch (Paszke et al., 2019). We
used a batch size of 256, an initial learning rate of $1\times 10^{-4}$, and
trained with early stopping.

####  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 to 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 $\sigma\in[0.5,4.0]$ for EBMs trained over 3,000 steps. Training
instabilities recur for $\sigma<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 $\ell_{2}$ 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; $\sigma$ is the only free
hyperparameter in dWJS.

####  A.3 Diffusion baselines

In our comparison study we use the Sequence-based Variational Diffusion Model
(SeqVDM) 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 in a latent space. In all our
experiments we used $T=1000$ steps with the fixed noise schedule
$\gamma_{min}=-13$ and $\gamma_{max}=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 $d=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 AdamW optimizer and the initial
learning rate of $2\times 10^{-4}$ for 50 epochs. The sampling is done by
starting from a latent vector initialized with Gaussian noise.

####  A.4 Estimation of $\sigma_{c}$ hyperparameter

Figure 4:  Histogram of $\chi_{ii^{\prime}}$ values for random samples from
the paired observed antibody space (Olsen et al., 2022) dataset.

###  Appendix B Additional algorithms

####  B.1 Gradient flow enables local minima finding

We define the gradient flow as $y^{\prime}(t)=-\nabla\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 at $t=0$,
$y_{0}=\varepsilon_{0},\varepsilon_{0}\sim\mathcal{N}(0,I_{d})$, and sample
noisy samples $y^{\prime}(t)$ following the gradient flow. In this way, we
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. 3.

Input: Denoiser, $g_{\phi}(y)$, energy-based model, $f_{\theta}(y)$

Output: Denoised samples $\hat{x}(y,t)$

1

2$y_{0}\sim\mathcal{N}(0,\sigma^{2}I_{d})$

3 for  _$t=0,\dots,T-1$_ do

4 $y_{t+1}\leftarrow y_{t}-\delta\,\nabla_{y}f_{\theta}(y_{t})$

5

6 end for

7$\hat{x}_{T}\leftarrow y_{T}+\sigma^{2}g_{\phi}(y_{T})$

8 return _$\operatorname*{arg\,max}\hat{x}_{T}$_

Algorithm 3 Discrete gradient flow

####  B.2 Langevin MCMC Update

1: Input $\delta$ (step size), $u$ (inverse mass), $\gamma$ (friction), $K$
(steps taken)

2: Input Learned score function $g_{\theta}(y)\approx\nabla\log p(y)$ and
noise level $\sigma$

3: Output $\hat{x}_{K}$

4: $y_{0}\sim\mathcal{N}(0,\sigma^{2}I_{d})+\mathcal{U}_{d}(0,1)$

5: $v_{0}\leftarrow 0$

6: for $k=0,\dots,K-1$ do

7: $y_{k+1}\leftarrow y_{k}+\frac{\delta}{2}v_{k}$

8: $g_{k+1}\leftarrow g_{\theta}(y_{k+1})$

9: $v_{k+1}\leftarrow v_{k}+\frac{u\delta}{2}g_{k+1}$

10: $\varepsilon\sim\mathcal{N}(0,I_{d})$

11:
$v_{k+1}\leftarrow\exp(-\gamma\delta)v_{k+1}+\frac{u\delta}{2}g_{k+1}+\sqrt{u\left(1-\exp(-2\gamma\delta)\right)}\varepsilon$

12: $y_{k+1}\leftarrow y_{k+1}+\frac{\delta}{2}v_{k+1}$

13: end for

14: $\hat{x}_{K}\leftarrow y_{K}+\sigma^{2}g_{\theta}(y_{K})$

Algorithm 4 Walk-jump sampling (Saremi & Hyvärinen, 2019) using the
discretization of Langevin diffusion by Sachs et al. (2017). Lines 6-13
correspond to _walk_ step and line 14 is the jump step.

###  Appendix C Performance profiling

Table 6: Model size, sampling time and memory consumption. Model | Parameters (M) | GPU time / sample (ms) $\downarrow$ | GPU memory (MB) $\downarrow$
---|---|---|---
dWJS (energy-based) | 9.87 | 64.89 | 4734.76
dWJS (score-based) | 1.03 | 327.60 | 16.7
SeqVDM | 12.31 | 940.40 | 246.39
DEEN | 1.03 | 976.99 | 16.72

###  Appendix D Few-shot, in-context learning

####  D.1 GPT 4 prompt

GPT 4 accessed on April 27, 2023. Prompt was repeated 20 times to generate
additional samples for statistics.

Prompt: You are an expert antibody engineer. I am going to give you examples
of CDR 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 sequence via deep learning". You will
generate 100 new CDR H3 variants that you predict will also bind to HER2.
Output the 100 samples as a python list. Here are the examples: [’WHINGFYVFH’,
’FQDHGMYQHV’, ’YLAFGFYVFL’, ’WLNYHSYLFN’, ’YNRYGFYVFD’, ’WRKSGFYTFD’,
’WANRSFYAND’, ’WPSCGMFALL’, ’WSNYGMFVFS’, ’WSMGGFYVFV’, ’WGQLGFYAYA’,
’WPILGLYVFI’, ’WHRNGMYAFD’, ’WPLYSMYVYK’, ’WGLCGLYAYQ’,]

Response: As an AI language model, I can’t run simulations or generate
sequences based on actual molecular interactions. However, I can use patterns
found in your provided list to generate new 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.

Prompt: You are an expert antibody engineer. I am going to 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 are not in
the database. Output the 10 samples as a python list. Here are the examples:
[’QLQLQESGPGLVKPSGTLSLICAVAGSSVTTSNVYWDWIRQPPGKGLEWIGTISYSATT’,
’YYNPSLKSRVTLFVDTSKNQFSLKLNSVTAADTAMYYCARHVVDCTGGRCLWKTGFDYWGQGTLVTVSS’,
’QVQLVESGGGVVQPGRSLRLSCAASGFTFSSYGMHWVRQAPGKGLEWVAVIWYDGSNKYY’,
’ADSVKGRFTISRDNSKNTLYLQMNSLRAEDTAVYYCARDQGLSTYYGYDKDSSYGMDVWGQGTTVTVSS’,
’EVQLVESGGGLVKPGGSLRLSCAASGFTFSSYSMNWVRQAPGKGLEWVSSISSSSSYIYY’,
’ADSVKGRFTISRDNAKNSLYLQMNSLRAEDTAVYYCARDFGYCSGGSCSGWGQGTLVTVSS’,
’QVQLKESGPGLVQPSQTLSLTCTVSGFSLTSYHVHWVRQPPGKGLEWMGVMWSDGDTSYN’,
’SALKSRLSISRDTSKSQVFLKMSSLQTEDTATYYCARSGKAVMDAWGQGASVTVSS’,
’QVQLVQSGAEVKKPGSSVKVSCKASGGPFRGSTINWVRQAPGQGLEWMGRIIPILNAPNY’,
’AQNFHGRVTITADKSTSTAYMELSSLTSDDTALYFCATYSGTYHMAPGAFDIWGHGTMVTVSS’]

Response: Sure, I can help you generate 10 new antibody heavy chains. Here’s a
list of 10 new antibody heavy chains that are 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 a
similar length and structure as the provided examples, but they have different
amino acid sequences, and are unique.

###  Appendix E in vitro validation

####  E.1 Experimental details

in vitro validation of generated antibody proteins was performed following
Hsiao et al. (2020). Discrete Walk-Jump Sampling (dWJS) and SeqVDM were used
to generate antibody sequences, which were then expressed and purified in the
laboratory. Surface plasmon resonance (SPR) measurements were used to
determine binding affinity.