[Paper Review] Protein Discovery with Discrete Walk-Jump Sampling

Nathan C. Frey, Daniel Berenberg+11 more|arXiv (Cornell University)|Jun 8, 2023

Single-cell and spatial transcriptomics4 citations

TL;DR

This paper introduces Discrete Walk-Jump Sampling (dWJS), a novel non-autoregressive generative modeling framework that combines smoothed energy-based models and score-based denoising to enable efficient, high-quality protein sequence generation. By training on a single noise level and using Langevin MCMC on a smoothed data manifold, dWJS achieves fast, diverse sampling and generates 97–100% expressible, purified antibodies with 35% showing improved binding affinity in first-round lab experiments.

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 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.

Motivation & Objective

• Address the challenges of training and sampling from discrete energy-based models (EBMs) in protein sequence generation.
• Overcome the brittleness, instability, and inefficiency of existing autoregressive and diffusion-based models for antibody design.
• Develop a simplified, generalizable framework that requires only a single noise level and avoids complex training tricks like replay buffers or annealing.
• Enable ab initio discovery of novel, functional antibodies with high biophysical validity and experimental success rates.
• Introduce a new evaluation metric, the Distributional Conformity Score (DCS), to benchmark generative models in protein design.

Proposed method

• Propose Smoothed Discrete Sampling (SDS), a new paradigm that learns an energy function on a smoothed version of the discrete data manifold.
• Use Langevin MCMC to sample from the smoothed energy-based model, enabling efficient exploration of the data distribution.
• Apply one-step denoising to project samples back onto the true discrete data manifold, preserving sequence validity.
• Decouple training into separate score-based and energy-based models: train a score model on noisy data and an EBM on smoothed data.
• Leverage the neural empirical Bayes (NEB) formalism to unify kernel density estimation and empirical Bayes for discrete data.
• Design a novel architecture for discrete EBMs that improves training stability and sample quality without requiring common EBM regularization techniques.

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.

Experimental results

Research questions

• RQ1 Can a single-noise-level, non-autoregressive generative model outperform autoregressive and diffusion-based baselines in generating functional antibody sequences?
• RQ2 Can smoothed discrete sampling with walk-jump sampling achieve fast mixing and diverse exploration of multiple antibody germline classes in a single MCMC chain?
• RQ3 Does the proposed method eliminate the need for complex training tricks like replay buffers, ℓ₂ penalties, or simulated annealing in EBM training?
• RQ4 Can the Distributional Conformity Score (DCS) serve as a reliable, scalar benchmark for evaluating protein generative models?
• RQ5 To what extent can dWJS generate novel, biophysically valid, and functionally improved antibodies in a single round of laboratory testing?

Key findings

• 97–100% of generated antibody sequences were successfully expressed and purified in a single round of in vitro experiments.
• 35% of the generated functional designs exhibited equal or improved binding affinity compared to known functional antibodies.
• The method achieved long-run fast-mixing MCMC chains, visiting diverse antibody protein classes within a single Markov chain.
• The approach eliminated the need for common EBM training stabilizers such as replay buffers, ℓ₂ penalties, and rejection sampling.
• dWJS outperformed autoregressive large language models, diffusion models, and score-based baselines in both sample quality and experimental success.
• The Distributional Conformity Score (DCS) effectively benchmarked generative models and correlated with experimental functionality.

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) pa