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

TL;DR

ProtInvTree is a novel reward-guided tree search framework for protein inverse folding that generates diverse, structurally consistent sequences by modeling the process as a step-wise decision-making task.

Contribution

It introduces a reward-guided tree search approach with a two-stage action mechanism and jumpy denoising, enabling diverse sequence design without retraining, outperforming existing methods.

Findings

Outperforms state-of-the-art baselines on multiple benchmarks.

Generates diverse sequences with high structural consistency.

Supports flexible scaling at test time without retraining.

Abstract

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