Source-linked AI summary
Protein Discovery with Discrete Walk-Jump Sampling
Nathan C. Frey, Daniel Berenberg, Karina Zadorozhny, Joseph Kleinhenz, Julien Lafrance-Vanasse, Isidro Hotzel, Yan Wu, Stephen Ra, Richard Bonneau, Kyunghyun Cho, Andreas Loukas, Vladimir Gligorijevic, Saeed Saremi
TL;DR
Discrete generative models must balance faithful modeling with controllable, novel outputs, and antibody design operates in a large, sparsely functional space with relatively small therapeutic datasets. The paper combines smoothed-distribution learning with discrete Walk-Jump Sampling and evaluates it for antibody discovery. The authors report stronger performance than several protein-generation baselines and high in-vitro expression success.
Problem
Discrete generation requires expressive models that capture the data distribution while producing controllable, novel, diverse outputs; antibody design also faces a large, sparsely functional sequence space and small therapeutic datasets.
Source: §1 · pp. 1–2Method
Smoothed Discrete Sampling and discrete Walk-Jump Sampling use separately trained score- and energy-based models, Langevin sampling on noisy data, and one-step denoising, with a single noise level.
Source: §1 · p. 2, §3.1 · p. 3Results
The method outperforms autoregressive, masked protein-language, large-language, diffusion, and score-based baselines; 97.47% of tested sequences were expressed and purified.
Source: §1 · p. 2, §4.2 · p. 8Takeaways & Limitations
The authors conclude that fitting the energy function to noisy data prevents overfitting and training instabilities while enabling fast, efficient training and sampling with low compute requirements.
Source: §6 · p. 9
Abstract
from arXiv · showhide
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 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 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.
1 INTRODUCTION
Discrete generation requires models that balance fidelity with controllable, novel, diverse outputs, while antibody design adds a large, sparsely functional sequence space and small therapeutic datasets. The paper introduces Smoothed Discrete Sampling and discrete Walk-Jump Sampling, then evaluates the approach for antibody discovery and design.
1 INTRODUCTION
- Antibody design seeks useful proteins in a large, sparsely functional sequence space, and therapeutic-antibody datasets are relatively small.The authors distinguish ab initio discovery and design from de novo design, which starts without material.Source: §1 · pp. 1–2
- Figure 1 shows a single MCMC chain visiting diverse antibody germline classes, with colors identifying the classes.The caption describes the chain as fast-mixing and the folded structures as visualizations.Source: Figure 1 · p. 2
- Discrete Walk-Jump Sampling uses uncoupled, separately trained score- and energy-based models to learn noisy distributions and sample discrete data, including variable-length outputs.The approach builds on Smoothed Discrete Sampling and the neural empirical Bayes framework.Source: §1 · p. 2
- The method requires only one noise level and no noise schedule, while avoiding common EBM training tricks and retaining good sample quality and fast sampling.The paper says this alleviates brittleness, training instabilities, and slow sampling associated with diffusion models.Source: §1 · p. 2
- The authors report that their method outperforms autoregressive, masked protein-language, large-language, diffusion, and score-based baselines for protein discovery and design.They validate it with in vitro experiments.Source: §1 · p. 2
2 BACKGROUND
The background introduces energy-based models as probability models defined by an energy function and sampled using MCMC. It then describes neural empirical Bayes as learning a denoiser from Gaussian-corrupted data without MCMC during training.
2.1 ENERGY-BASED MODELS
- Energy-based models fit an energy function that specifies a probability distribution over data, and draw samples using Markov-chain Monte Carlo.They are typically trained via contrastive divergence.Source: §2.1 · p. 3
2.2 NEURAL EMPIRICAL BAYES
- Neural empirical Bayes adds Gaussian noise to data and learns a score-based least-squares estimator of the original data.Its denoising objective is optimized with stochastic gradient descent and requires no MCMC sampling during learning.Source: §2.2 · p. 3
3 ANTIBODY DISCOVERY AND DESIGN
Discrete Walk-Jump Sampling separates noisy-manifold sampling from denoising, while a critical-noise estimate guides the noise scale for antibody sequences. The section also introduces a reference-based conformity score for evaluating generated proteins.
3.1 DISCRETE WALK-JUMP SAMPLING
- dWJS samples noisy sequences with Langevin MCMC and separately denoises them, decoupling the walk from the jump to produce clean samples.The EBM is trained on smoothed data; the separate denoiser projects samples back to the data manifold.Source: §3.1 · p. 3
- Antibody sequences use gap tokens to accommodate variable lengths, with aligned heavy and light chains mapped to fixed input sizes of 149 and 148 positions.The model smooths one-hot encodings with Gaussian noise, samples using Langevin MCMC, then denoises and takes argmax to recover a one-hot sequence.Source: §3.1 · p. 5
- Constrained sampling preserves specified sequence regions while allowing non-conserved regions to change during antibody design.Unconditional sampling is used for discovery, whereas design uses constrained sampling and scoring.Source: §3.1 · p. 5
3.2 DERIVATION OF OPTIMAL NOISE LEVEL FOR DISCRETE SEQUENCE DATA
- σc ≈ 0.5 for antibody sequence data, matching the empirically sufficient noise scale; above σc, noisy data have some degree of overlap.For flattened sparse one-hot matrices, the stated scaling gives σc a factor of 0.22, but is not optimal because of sparsity.Source: §3.2 · p. 5
3.3 DISTRIBUTIONAL CONFORMITY SCORE
- The Distributional Conformity Score (DCS) condenses likelihoods under joint property densities into a score of generated-sample similarity to a reference distribution.The conformity measure uses biophysical and statistical properties, including protein-language-model log-probability; the reference set is split into fitting and validation sets.Source: §3.3 · p. 6
4 EXPERIMENTS
The experiments evaluate dWJS on antibody generation, measuring in-silico sample properties and testing expression and binding in the laboratory. dWJS produces high-quality samples across these evaluations, including functional HER2 binders.
4.1 DWJS GENERATES NATURAL, NOVEL, DIVERSE ANTIBODIES in silico
- dWJS with dEBM sampling best matched validation-set property distributions and had the most unique samples; score-based dWJS showed high conformity, novelty, and diversity.Strong antibody-likeness results were achieved by increasing σ to ≥0.5.Source: §4.1 · p. 7
- dWJS sampled faster, used less memory, and achieved better sample quality than diffusion, autoregressive, and score-based baselines; energy-based dWJS was 43× faster than IgLM.IgLM had the best DC score but the lowest internal diversity, while dWJS outperformed it on antibody-likeness, edit distance, and internal diversity.Source: §4.1 · p. 7
4.2 DWJS GENERATES NATURAL, NOVEL, DIVERSE ANTIBODIES in vitro
- 270 of more than 277 designed antibodies were expressed and purified (97.47%), versus 42% for clean-data EBM samples drawn with traditional Langevin MCMC.The authors report that fewer than four edits can destroy expression when the proposal distribution is poorly optimized.Source: Table 3 · p. 7, §4.2 · p. 8
4.1 DWJS GENERATES NATURAL, NOVEL, DIVERSE ANTIBODIES in silico
- ESM2 infilling was the fastest sampler but generated repetitive sequences unlike antibodies, making its high edit-distance and internal-diversity scores misleading.ESM2 was not trained for ab-initio generation; masking was increased to 40% while retaining valid samples.Source: §4.1 · p. 8
4.3 DWJS GENERATES FUNCTIONAL ANTIBODY VARIANTS in vitro
- On hu4D5 CDR H3 mutants, dWJS models trained on binders produced the highest percentage of unique predicted HER2 binders among the compared models.The dataset included 9,000 binding and 25,000 non-binding mutants; dWJS was compared with diffusion models and GPT-4.Source: §4.3 · p. 8
- 70% of dWJS energy-based designs bound HER2 in first-round validation, versus 25% for LaMBO-2 and 22% for AbDiffuser (57% after filtering).The dWJS designs required no post-hoc filtering; the comparison is for discrete distribution learning, while the diffusion methods include other capabilities.Source: Table 4 · p. 8, §4.3 · pp. 8–9
5 RELATED WORK
Prior work connects energy-based modeling and score matching to empirical Bayes, while discrete diffusion learns denoising across noise levels. Antibody generation remains challenging with limited training data and high-entropy variable regions.
5 RELATED WORK
- Neural empirical Bayes unifies kernel density estimation and empirical Bayes, parameterizing a neural energy function to capture a smoothed data distribution.The paper presents this as its first study of the NEB formalism for discrete data.Source: §5 · p. 9
- Discrete diffusion models learn iterative denoising across multiple noise levels using a prescribed process that converges to a known categorical distribution.Source: §5 · p. 9
- Antibody modeling methods can struggle with limited training data and high-entropy variable regions, motivating decoupled energy- and score-based modeling for discrete sequences.The cited antibody datasets range from approximately 1K to 1M high-quality sequences, depending on the distribution of interest.Source: §5 · p. 9
6 CONCLUSIONS
The paper presents SDS and dWJS as a simplified, robust approach to discrete energy- and score-based modeling, evaluated through antibody discovery and design. Its results support further testing beyond antibodies and theoretical investigation.
6 CONCLUSIONS
- dWJS generated novel, diverse, and functional antibodies, as assessed through synthetic biophysical distributions, similarity metrics, and in-vitro experiments.Source: §6 · p. 9
- Fitting the energy function to noisy data provided strong regularization that the authors report prevents overfitting and training instabilities.Source: §6 · p. 9
- dWJS uses Langevin MCMC on smoothed data distributions and reduces training and sampling complexity to choosing the noise level σ.It discards common EBM techniques including replay buffers, ℓ2 norm penalties, simulated annealing, and rejection sampling.Source: §6 · p. 9
- The authors identify generality to other molecule classes and data modalities, such as images, and theoretical investigation as future work.Source: §6 · p. 9
A.1 DISCRETE WALK-JUMP SAMPLERS
The section compares Smoothed Discrete Sampling variants and identifies EBM walking followed by denoiser jumping as the strongest combination. Its reported performance is not tied to a single denoiser architecture.
A.1 DISCRETE WALK-JUMP SAMPLERS
- The experiments use a three-layer Conv1D energy model and a 35-layer ByteNet denoiser, trained with AdamW and early stopping.Both models use hidden dimensions of 128; training uses batch size 256 and initial learning rate 1 × 10^-4.Source: §A.1 · p. 13
- A transformer-based dWJS denoiser performs comparably to ByteNet, suggesting the method is not reliant on one architecture.The transformer uses 12 hidden layers and 8 attention heads.Source: §A.1 · p. 13
- EBM walking followed by a denoiser jump is reported as the most effective variant for efficiency, sample quality, and diversity.The taxonomy also includes score-based dWJS, DEEN, and dWJS-EBM.Source: §A.1 · p. 14
- Table 5 summarizes alternative implementations of Smoothed Discrete Sampling.The supplied caption does not report table entries.Source: Table 5 · p. 14
A.2 DWJS STABILIZES AND SIMPLIFIES TRAINING
The appendix describes training stability, baseline and noise-level choices, sampling procedures, and supplementary implementation details. It reports that noise regularizes EBM training, while the best stated sample-quality setting is σ = 0.5.
A.2 DWJS STABILIZES AND SIMPLIFIES TRAINING
- Training remains stable across σ ∈ [0.5, 4.0] for 3,000 steps, but instabilities recur below σ = 0.5.Noising regularizes the energy model; the method discards replay buffers, ℓ2 penalties, Metropolis rejection, and time-step annealing, leaving σ as its only free hyperparameter.Source: §A.2 · p. 14
A.3 DIFFUSION BASELINES
- The diffusion baseline uses 1,000 steps with a fixed noise schedule and jointly trains diffusion and sequence-reconstruction losses.SeqVDM is trained on paired OAS for 50 epochs.Source: §A.3 · p. 14
A.4 LANGUAGE MODEL BASELINES
- The language-model baselines include conditional IgLM generation and ESM2 infilling, which the authors say is not true ab initio protein discovery.ESM2 infilling is limited to 40% masking; higher masking produces invalid sequences, and the reported high edit-distance and internal-diversity scores are described as meaningless because samples are not antibody-like.Source: §A.4 · pp. 14–15
A.5 EFFECT OF CHOICE OF σ
- σ = 0.5 produces the best reported score-based dWJS sample quality while maintaining uniqueness and diversity.At σ = 0.1, samples poorly capture the training set’s biophysical-property distribution and have edit distances above 120; at σ = 3, samples show signs of mode collapse.Source: §A.5 · p. 15, Table 6 · p. 15
B.1 GRADIENT FLOW ENABLES LOCAL MINIMA FINDING
- The appendix describes discrete gradient flow as deterministic motion along the probability-density gradient from an initial noise sample to local attractors on the data manifold.The passage identifies these attractors with local minima of the learned energy function.Source: §B.1 · pp. 15–16
B.2 LANGEVIN MCMC UPDATE
- The Langevin MCMC algorithm’s walk uses iterative updates, followed by a jump step that produces denoised samples.The supplied algorithm passage lists initialization and update inputs, while its caption identifies lines 6–13 as the walk and line 14 as the jump.Source: §B.2 · p. 16
B.3 NEURAL EMPIRICAL BAYES
- The Neural Empirical Bayes view embeds discrete data in Euclidean space, smooths it with noise, and then uses single-step denoising to return to discrete values.Walk and jump are decoupled; the sampling and denoising models may be trained separately, and their empirical performance can differ.Source: §B.3 · p. 16
C PERFORMANCE PROFILING
- The performance-profiling table is described as reporting model size, sampling time, and memory consumption.The supplied caption gives no values.Source: Table 7 · p. 17
E.1 EXPERIMENTAL RESULTS
The figure reports expression yield and binding affinity for sequence designs targeting ERBB2.
E.1 EXPERIMENTAL RESULTS
- Figure 5 plots expression yield (mg) and binding affinity (pKD) for designs targeting ERBB2.The supplied caption gives no numerical outcomes.Source: Figure 5 · p. 18
E.2 EXPERIMENTAL DETAILS
The experiments validate generated antibodies through laboratory expression, purification, and surface plasmon resonance binding measurements.
E.2 EXPERIMENTAL DETAILS
- Generated dWJS antibody sequences were expressed and purified, with binding affinity measured by surface plasmon resonance.The validation followed Hsiao et al. (2020).Source: §E.2 · p. 18
- Antibody variable domains were cloned into mammalian expression vectors, transiently expressed in Expi293 cells for seven days, and affinity-purified.Concentrations were measured by optical density at 280 nm.Source: §E.2 · p. 18
- Binding affinity was measured by SPR at 37 °C using a Biacore 8K and a 1:1 binding-kinetics model.Antibodies were captured on a Protein A chip before target-antigen injection.Source: §E.2 · pp. 18–19
F FURTHER DISCUSSION OF DISTRIBUTIONAL CONFORMITY SCORE
The discussion explains how distributional conformity score (DCS) measures conformity to a reference distribution and how it differs from property alignment. It also situates the approach among density-estimation choices and prior methods for discrete generative modeling.
F FURTHER DISCUSSION OF DISTRIBUTIONAL CONFORMITY SCORE
- A conformal transducer assigns p-values that represent the fraction of examples with lower conformity to the reference distribution than a test example.The text describes the transducer as measuring how similar a new example is to the reference sequence.Source: F FURTHER DISCUSSION OF DISTRIBUTIONAL… · p. 19
- DCS may prioritize relationships among properties over individual-property alignment and may be more influenced by extreme points because it uses KDE.The authors describe this as a distinct perspective on generative-model quality, not a replacement for property alignment.Source: F FURTHER DISCUSSION OF DISTRIBUTIONAL… · p. 19
- The evaluated properties include three sequence-based measures and two structure-based antibody measures.The sequence-based measures are average hydrophilicity, molecular weight, and grand average of hydropathy; the structure-based measures are CDR-region surface hydrophobicity patches and symmetry of variable-chain charges.Source: F FURTHER DISCUSSION OF DISTRIBUTIONAL… · p. 19
- The authors use KDE for the small-sample, low-dimensional setting, while DCS can be combined with any density estimator and requires no sequence alignment or preprocessing.Their KDE uses Gaussian kernels, Silverman’s global bandwidth, adaptive local bandwidth 0.15, and diagonal covariance.Source: F FURTHER DISCUSSION OF DISTRIBUTIONAL… · p. 19
G FURTHER DISCUSSION OF RELATED WORK
- The discussion contrasts discrete-data Gibbs-sampling contrastive divergence with continuous-input Langevin training and with discrete score matching via a local-topology prior.The cited approaches are contrastive divergence, continuous inputs and Langevin MCMC, and discrete score matching.Source: G FURTHER DISCUSSION OF RELATED WORK · p. 19
- The paper frames antibody design as a scientific generative-modeling task involving discrete amino-acid sequences in heavy and light chains.It notes challenges including evaluation and generating novel, diverse samples that respect domain structure and symmetries.Source: G FURTHER DISCUSSION OF RELATED WORK · p. 19