Uncertainty-Aware Discrete Diffusion Improves Protein Design

Abstract

Protein inverse folding involves generating amino acid sequences that adopt a specified 3D structure--a key challenge in structural biology and molecular engineering. While discrete diffusion models have demonstrated strong performance, existing methods often apply uniform denoising across residues, overlooking position-specific uncertainty. We propose an uncertainty-aware discrete denoising diffusion model that employs a prior-posterior signaling mechanism to dynamically guide the denoising process. Our approach further integrates learned priors from a pretrained protein large language model and a structure encoder within a modular framework, jointly optimized through multi-objective training. Across multiple benchmarks, our method achieves substantial improvements over state-ofthe-art baselines, offering a principled framework for structure-conditioned sequence generation in proteins and beyond.