One-sentence Summary

Genentech and New York University propose Discrete Walk-Jump Sampling, which unifies contrastive divergence training of energy-based models with Langevin MCMC on smoothed manifolds and one-step denoising at a single noise level, yielding high-quality antibody protein generation where 97-100% of samples are successfully expressed and 70% of functional designs show equal or improved binding affinity compared to known functional antibodies on the first laboratory attempt.

Key Contributions

• Smoothed Discrete Sampling is introduced for discrete generative modeling by learning a smoothed energy function, sampling from the smoothed data manifold with Langevin Markov chain Monte Carlo, and projecting samples back to the true data manifold with one-step denoising. The approach combines contrastive divergence energy-model training with score-based sampling behavior while requiring only a single noise level.

• Discrete Walk-Jump Sampling is developed for antibody protein sequences and uses the smoothed formulation to remove common energy-model stabilization machinery, reduce engineering complexity to a single noise-level choice, and enable robust and efficient training and sampling on limited, high-entropy biological data.

• A distributional conformity score is introduced for benchmarking protein generative models, and antibody design experiments show that 97% to 100% of generated samples are expressed and purified and 70% of functional designs equal or improve binding affinity on the first attempt, including the first demonstration of long-run fast-mixing Markov chain Monte Carlo chains that visit diverse antibody protein classes within a single chain.

Introduction

Discrete sequence generation is difficult for gradient-based models because they must learn expressive distributions while producing diverse, novel samples that obey application constraints. This matters in protein design, especially antibody discovery, where the sequence space is enormous, therapeutic datasets are small, and sequences contain both conserved structure and highly variable regions. Existing methods fall short: autoregressive protein models suffer from error accumulation and high latency, diffusion or score models use complex noise schedules and can be brittle or slow, and energy-based models are hard to train and sample. The authors leverage the neural empirical Bayes framework to introduce Smoothed Discrete Sampling and discrete Walk-Jump Sampling, which combine separately trained energy and score models with a single noise level, enable fast non-autoregressive sampling, and demonstrate strong biological validity in antibody design.

Method

The authors leverage a generative framework that separates noisy continuous sampling from discrete decoding. The method operates on smoothed representations of discrete sequences rather than on raw categorical tokens. For antibody data, heavy and light chains are aligned to fixed token lengths with a gap token and encoded as sparse one-hot tensors. These discrete representations $x \in \{0,1\}^{d}$x∈{0,1}d are converted into continuous noisy inputs by additive isotropic Gaussian noise,

The smoothed noisy manifold contains sequences that are close to valid discrete encodings but are not themselves categorical. The framework therefore trains a density model on this continuous noisy space and a separate denoiser that projects noisy states back to clean discrete sequences. Refer to the framework diagram:

Within this pipeline, a discrete energy-based model parameterizes an energy function $f_{\theta}(y)$fθ​(y) over noisy sequence representations. The corresponding noisy density is expressed as

The model is trained by maximum likelihood on the smoothed data distribution. The objective can be written in positive and negative sample form as

where $y^{+}$y+ denotes noisy training data and $y^{-}$y− denotes noisy samples generated from the current model. This training objective encourages the model to assign lower energy, equivalently higher density, to realistic noisy data while assigning higher energy to noisy samples produced by the model. The contrastive divergence implementation of this objective leverages the identity

which replaces the intractable partition function gradient with an expectation over model samples.

A separate neural denoiser is trained through neural empirical Bayes. The smoothed marginal density of noisy samples is

and the least-squares estimator of the clean signal $x$x given a noisy observation $y$y is

The score function $g(y) = \nabla \log p(y)$g(y)=∇logp(y) is approximated by a neural network $g_{\phi}: \mathbb{R}^{d} \to \mathbb{R}^{d}$gϕ​:Rd→Rd. This yields a parametric denoising map

The network $g_{\phi}$gϕ​ is optimized with the denoising objective

This objective is a supervised least-squares regression on noised and denoised pairs and does not require MCMC sampling during denoiser training. In practice, the denoiser maps a smoothed noisy input to a continuous estimate of the clean one-hot encoded sequence, and an $\operatorname{argmax}$argmax projection recovers a valid discrete sequence.

The density model used in the pipeline accepts the noisy tensor $y$y, concatenates it with a positional encoding $p_{1d}$p1d​, passes the result through an MLP and a convolutional feature extractor, combines those features with an embedding of the input, and outputs the energy $f_{\theta}(y)$fθ​(y). The architecture is designed to capture dependencies along the fixed aligned sequence representation while operating in the noisy continuous space rather than directly on categorical tokens.

At inference time, generation proceeds in two interleaved roles. First, the energy model performs a “walk” by running Langevin MCMC on the noisy data manifold. If the noisy state is $y_t$yt​, a Langevin update can be viewed as a gradient-informed random walk over the density modeled by $p_{\theta}$pθ​. Second, after each noisy state $y_k$yk​, the denoiser performs a “jump” by returning

Applying an $\operatorname{argmax}$argmax operation to $\hat{x}_{\phi}(y_k)$x^ϕ​(yk​) recovers a discrete sequence. Because the walk and jump components are decoupled, every noisy state encountered during the Markov chain can be decoded into a candidate sequence. This differs from standard diffusion-style sampling, where intermediate states are typically noisy and must be gradually transported through a fixed noise schedule. In the walk-jump formulation, the Markov chain samples the smoothed density, while the denoiser maps any such sample back toward the discrete data manifold.

A characteristic sampling trace from a single Markov chain is shown in the figure below:

The trace illustrates that one chain can move across diverse regions of the antibody distribution, including different germline classes, without requiring multiple independent generation runs. The separation of sampling and denoising also permits constrained sequence editing. For protein discovery, the model can be sampled unconditionally. For protein design, a starting sequence $s$s can be partially preserved by applying a binary projection matrix $P$P after each denoising step, enforcing

Here $P$P selects the positions that should remain conserved. The unconstrained positions remain free to change as the chain evolves on the noisy manifold.

The noise level $\sigma$σ is chosen to make the noisy representations overlap sufficiently for learning and sampling. The authors motivate this choice through pairwise distances between data samples. Define

where the normalization reflects concentration properties of isotropic Gaussians in high dimensions. The critical noise level is

When $\sigma > \sigma_c$σ>σc​, noisy samples tend to overlap in the smoothed space, which supports both stable density estimation and efficient sampling. For the aligned antibody representations, this critical value is approximately $0.5$0.5, matching the range used empirically.

The generated sequences are assessed through a distributional conformity score, which evaluates whether samples resemble a reference distribution while retaining novelty. The workflow computes multiple properties for in silico designs and for a held-out reference set. These properties may include biophysical attributes and statistical measures such as sequence naturalness, alignment-based similarity, hydrophobicity, charge-related properties, and language-model likelihood. Joint density estimation is then used to compare the designs against the reference distribution, and the resulting likelihoods are condensed into a conformity measure. The evaluation pipeline is shown below:

Experiment

Experiments evaluate discrete Walk jump sampling for antibody distribution learning, in vitro expression, and functional therapeutic design, using distributional conformity, antibody likeness, and wet lab validation. In silico results show that the method generates natural, novel, diverse antibodies with strong sample quality and favorable sampling efficiency compared with several baselines. Wet lab studies further confirm that the generated sequences are highly expressible and can function as target binders, supporting the method’s practical utility for designing realistic and therapeutically useful antibodies.

The table summarizes normalized pairwise distances for discrete antibody sequences and provides a geometric basis for choosing a training noise level. Most distances cluster in the midrange, while the largest distance is just above 0.5 and serves as the critical noise level. This supports using a noise level near 0.5, the value adopted for the antibody generation experiments. Normalized antibody sequence distances show a broad but moderate spread rather than extreme variation. The largest observed distance closely matches the critical noise level, justifying a noise setting near 0.5. The selected noise level was used for training antibody generators that produced high-quality functional binder candidates.

The antibody likeness results show that the dWJS methods generate samples with strong agreement to validation sequence properties while preserving uniqueness. The energy based variant best matches the validation property distribution and has the highest uniqueness, while the score based variant emphasizes novelty and internal diversity. Baseline models show tradeoffs, with IgLM giving strong distributional conformity but weak diversity, and GPT 3.5 showing weaker antibody likeness and uniqueness. dWJS energy based sampling gives the best property distribution agreement and highest uniqueness among the compared models. dWJS score based sampling provides strong novelty and internal diversity while keeping distributional conformity competitive. IgLM has the strongest distributional conformity score but much lower internal diversity, indicating reduced sample variety. GPT 3.5 is less antibody like and less unique than the dWJS methods. ESM2 shows high edit distance and internal diversity but poor conformity to antibody sequence properties.

The measured protein synthesis results show that dWJS generated antibody sequences were expressed and purified with very high in vitro success. A score-based dWJS model reached full measured expression among tested designs, while an energy-based dWJS model also remained very high. In comparison, an EBM baseline expressed at a substantially lower rate, indicating the advantage of dWJS for capturing the antibody distribution. Near complete in vitro expression was observed for dWJS generated antibody sequences. Energy based dWJS maintained very high expression and substantially outperformed the EBM baseline. Few sequence edits can disrupt expression when the generative proposal is poorly optimized.

The table compares predicted and experimentally measured antibody binding for CDR H3 redesign methods. The energy based dWJS model shows the strongest predicted binding and the highest experimentally measured fraction of functional binders among the reported results. Compared with concurrent diffusion methods, dWJS reaches a much higher wet lab binding rate without relying on post hoc filtering. energy based dWJS designs achieve the highest measured functional binding rate among the methods with wet lab validation predicted binding probabilities rank dWJS above the other sequence models, while AbDiffuser also receives a high predicted binder probability measured rates for LaMBO-2 and unfiltered AbDiffuser are lower than those achieved by dWJS

The experiments evaluate the antibody generation pipeline using sequence distance calibration, antibody likeness metrics, in vitro expression, and functional binding validation. Normalized pairwise distances show a moderate spread and support using a training noise level near 0.5. dWJS sampling methods generate realistic, unique, and diverse antibody-like sequences, with energy based sampling giving the best property agreement and score based sampling emphasizing novelty. Experimental results further show that dWJS designs, especially energy based models, express highly in vitro and achieve the strongest measured functional binding among the compared methods.