Antibody sequence optimization with gradient-guided discrete walk-jump sampling

For enhancing the rapid discovery and delivery of antibody-based drugs, antibody attribute optimization is essential for the real-world application of therapeutic antibody sequence design. Using a generative machine learning model for antibody design is a promising direction for such tasks. However, existing methods struggle in balancing error accumulation, scalability, and targeted attribute optimization. In this work, we propose gradient-guided discrete walk-jump sampling (gg-dWJS), a novel discrete sequence generation method for antibody attribute optimization. Leveraging gradient guidance in the noisy manifold, we sample from the smoothed data manifold by applying discretized Markov chain Monte Carlo (MCMC) using a denoising model with the gradient-guidance from a discriminative model. This is followed by jumping to the discrete data manifold using a conditional one-step denoising. Through evaluation on both discrete image and antibody sequence generation tasks, we show that our method generates high-quality samples that are well-optimized for specific tasks.