ProtInvTree: Deliberate Protein Inverse Folding with Reward-guided Tree Search

Xiaoxue Cheng (Renmin University of China) · Hong Chang (Institute of Computing Technology, Chinese Academy of Sciences) · Shiguang Shan (Chinese Academy of Sciences) · Xilin Chen (Institute of Computing Technology, Chinese Academy of Sciences) · Cheng Tan (Shanghai AI Laboratory) · Zhangyang Gao (Westlake University, China) · Mengdi Liu (Institute of Computing Technology, Chinese Academy of Sciences)
decision-making processempirical evaluationfocus-and-grounding action mechanismgenerative modeljumpy denoising strategyposition selectionpretrained protein language modelsprotein inverse foldingresidue generationreward-guided tree-searchsearch breadthsearch depthsequence generationstructural consistencytest-time scaling

Designing protein sequences that fold into a target 3D structure—known as protein inverse folding—is a fundamental challenge in protein engineering. While recent deep learning methods have achieved impressive performance by recovering native sequences, they often overlook the one-to-many nature of the problem: multiple diverse sequences can fold into the same structure. This motivates the need for a generative model capable of designing diverse sequences while preserving structural consistency. To address this trade-off, we introduce ProtInvTree, the first reward-guided tree-search framework for protein inverse folding. ProtInvTree reformulates sequence generation as a deliberate, step-wise decision-making process, enabling the exploration of multiple design paths and exploitation of promising candidates through self-evaluation, lookahead, and backtracking. We propose a two-stage focus-and-grounding action mechanism that decouples position selection and residue generation. To efficiently evaluate intermediate states, we introduce a jumpy denoising strategy that avoids full rollouts. Built upon pretrained protein language models, ProtInvTree supports flexible test-time scaling by adjusting the search depth and breadth without retraining. Empirically, ProtInvTree outperforms state-of-the-art baselines across multiple benchmarks, generating structurally consistent yet diverse sequences, including those far from the native ground truth. The code is available at https://github.com/A4Bio/ProteinInvBench/.