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

W1: No mention of computational cost at inference time which to me is significant omission. An analysis of the cost of creating, e.g., 10 designs for a given backbone for ProtInvTree vs. ESM-3 and other models would be particularly beneficial.

Reply: We agree that test-time computational efficiency is critical, particularly for methods that involve large pretrained models and structure prediction. We provide a detailed analysis of the compute–performance trade-off. Specifically, for each given backbone, we generate 10 designs and compute the average generation time for each design. We fix the depth and scale the expand number from 2 to 6, as follows:

We also scaled the number of sampling iterations for ESM-3 to assess its performance–efficiency trade-off. The resulting sc-TMScore and average inference time per design are summarized in the table below:

W2: Few ablation studies to support modeling choices.

Reply: We ran an ablation study. Starting from the full ProtInvTree, we conducted an ablation study by disabling each key mechanism in isolation to assess its contribution:

W3 & Q8-1: Hyperparameter Ablation.

Reply: We conducted a controlled sweep over several key planning hyperparameters: the UCT exploration weight $\omega$, the sampling temperature $T$, the reward threshold used for termination and selection number schedules.

W4: No code in the supplementaries.

Reply: We will release the code upon publication to facilitate reproducibility and further research.

W5: No mention of statistical significance, despite the checklist saying so.

Reply: All reported results are the average over three independent runs, and we will update the standard deviation to reflect variability. Below, we present a subset of results with standard deviation values:

W6: Very limited discussion on model limitations.

Reply: In the revised manuscript, we will discuss more imitations: (1) the reliance on large pretrained models such as ESM-3 and ESMFold, which introduces non-negligible inference cost, (2) the assumption that structural consistency alone is sufficient for functional accuracy, which may not hold in all biological contexts.

Q1: L118: Shouldn’t $n$ be $L$?

Reply: Thank you for catching this typo. It should be $L$ to properly denote the sequence length. We will fix this in the revised version.

Q2: L214: What is DDIM?

Reply: DDIM refers to Denoising Diffusion Implicit Models [1], a sampling method used in diffusion models that accelerates the generation by jumpy denoising. In our method, we apply a single-step DDIM sampling to approximate the final denoised sequence, substantially reducing computational cost during MCTS rollout. We will add the citation [1] in the final version.

Q3: equation 10 essentially implies that using some different model for inverse folding is part of the proposed inverse folding model.

Reply: No, it is the same inverse folding model. The jumpy denoising mechanism simply takes advantage of ESM-3’s ability to perform flexible infilling under arbitrary masking ratios.

Q4: In equation 8, how are the tokens chosen? Are they sampled or do you apply argmax to the distribution over amino acids?

Reply: Concretely, we first soften the model’s logits with a temperature T (T = 0.3 in all experiments) and then sample a residue from the resulting categorical distribution.

Q5: L197 and L201: Why is $N$ used? Shouldn’t it be $K_t$ after the top-$K_t$ positions have been selected?

Reply: Yes, it should be $K_t$. We will fix this in the revised version.

Q6 & Q7: L201: Curious that random ends up being optimal.

Reply: We were also initially surprised that the random position selection strategy outperformed some more "informed" heuristics. However, our analysis suggests that random selection introduces useful exploration into the MCTS process and avoids biasing the search toward fixed or overconfident regions of the sequence. We do not claim that random selection is inherently optimal. Instead, we view it as a strong and surprisingly effective baseline. In future work, we plan to explore hybrid strategies that incorporate biologically or structurally informed priors—e.g., conservation profiles, solvent accessibility, or functional annotations—to guide the modification positions more precisely.

Q8-2&3: Did you explore other PLMs or folding models? Why is it not possible to use a different structure-conditioned model than ESM-3?

Reply: ESM-3 is the first unified structure–sequence model that supports arbitrary masking ratios, making it particularly well-suited for our planning-based inverse folding framework. In contrast, models like ProteinMPNN and KWDesign (based on ESM-2) lack this flexibility and typically require retraining to support structure-conditioned generation—conflicting with the core design of ProtInvTree, which relies on: 1. Stepwise sequence construction; 2. Test-time reward-guided planning without finetuning, and 3. Flexible infilling for focus-and-grounding with jumpy denoising. In addition, we use ESMFold for structure prediction due to its favorable trade-off between accuracy and computational efficiency, which is crucial for iterative structure evaluations during MCTS.

Q9: The termination number is set to 50. How many iterations are typically required before the algorithm terminates? Is it due to reaching this number, or is the structural consistency goal often met? An analysis of this would be very interesting.

Reply: The actual number of iterations is dynamically determined by the expand N, depths, and predefined structural consistency threshold. We analyzed the average number of iterations required under different thresholds:

These results show that early termination frequently occurs, especially under lower thresholds or smaller planning budgets. When both the expand number and depth are large (e.g., N=10, depth=10), the planner often reaches the iteration cap (50).

Q10: In the same vein, how many nodes end up being generated? How many structures are predicted to arrive at a given design? This feeds back into my earlier questions regarding the compute cost. Are we talking a handful of structures or hundreds/thousands?

Reply: The table below reports the average number of nodes explored under different configurations:

Each node, except for the root node, requires a structure prediction call (ESMFold). Moreover, a single MCTS tree can yield multiple high-scoring, diverse candidates from the final pool; accordingly, we report the amortized inference time per design by dividing the total runtime by the number of extracted outputs, shown in reply to W1.

Q11: Figure 4: Noverty→Novelty

Reply: Thank you for pointing this out. We will correct the typo in the revised version of the manuscript.

Thank you for your detailed and insightful feedback. We appreciate your recognition of the strengths of our approach, and your constructive suggestions highlight important areas for further improvement. Below, we address each of your comments and questions:

W1: The fairness of the evaluation.

Reply: During search we predict structures with ESMFold and compute TM‑score against the target backbone. For all tables/figures we report sc‑TMScore computed with TM‑align between the backbone and the structure returned by AlphaFold2 given the designed sequence. Thus the evaluator is different from the reward model, eliminating information leakage. We sample 10 sequences per backbone for every baseline, and the results are

W2 & Q7:  Computational efficiency.

Reply: For each given backbone, we generate 10 designs and compute the average generation time for each design. We fix the depth and scale the expand number from 2 to 6, as follows:

W3 & Q9 & Q11: Reproducibility. The details of the jumpy denoiser.

Reply: We clarify that the jumpy denoiser is not separately trained by us. It leverages the pretrained ESM-3 model, which is capable of filling in arbitrary mask ratios through its masked denoising objective. We will release the code upon publication to facilitate reproducibility and further research.

Q1: Related work

Reply: While both ProtInvTree and EvoPlay adopt Monte Carlo Tree Search (MCTS) to guide sequence generation, they differ significantly in motivation, problem formulation, and detailed methodology:

• Motivation and Problem Formulation: ProtInvTree focuses on protein inverse folding that aims at designing diverse sequences while preserving structural consistency. In contrast, EvoPlay focuses on function-oriented protein engineering, where the objective is to optimize protein functionality (e.g., fluorescence, binding affinity)
• Methodology: ProtInvTree is a test-time scaling framework that proposed a focus-and-grounding mechanism and a jumpy denoising strategy. In contrast, EvoPlay relies on AlphaZero-style self-play with training of policy-value networks and uses full simulations to evaluate each move via surrogate scoring functions.

We will revise the Related Work section to explicitly compare ProtInvTree with EvoPlay and other MCTS-based methods in protein design.

Q2, Q3, Q5, Q13-17: Terminology & clarity issues

Reply: Thank you for your careful reading and thoughtful suggestions regarding terminology and clarity. We have carefully revised the manuscript to address all noted issues.

Q4: Figure 3: Why are there not the same number of points in both graphs for ESM-3 and ProtinvTree?

Reply: Thank you for your careful observation. This is because we included an additional run in the left plot to make the performance trend more visually distinguishable. We will revise the right plot to include the same number of data points for both methods to ensure a fair and consistent comparison across the two subplots.

Q6: (Line 269-270) Can you clarify the sentence “Notably, when at comparable sc-TMScore levels, the baselines of AlphaDesign, PiFold, and KW-Design exhibit progressively higher diversity and lower novelty”? Should I check this pattern in Fig. 3?

Reply: Yes, as shown in Figure 3, when fixing the sc-TMScore (i.e., drawing a vertical line at a specific value), you can observe that AlphaDesign exhibits the lowest diversity (left panel) and the highest novelty (right panel) compared to PiFold and KW-Design.

Q8: The example in Fig. 7 would be more informative if it showcased an example in which the sc-TMScore is increased by applying the proposed method when compared to baseline methods (even if it is a cherry-picked example).

Reply: In the revision, we will update the figure to showcase an example where our method achieves a higher sc-TMScore compared to baseline methods. We note that such improvements are consistently observed across the majority of test cases, as reflected in our quantitative results.

Q10: For the reviewer, it is still not clear if the sequence in a state s_t contains all the amino acid identities and these are mutated depending on the positions sampled, or if there are positions masked that are being filled. From Fig. 2(b), it seems that there are masked positions that are being filled, but it is not clear in the writing. Also, this affects the MCTS termination conditions, as I expected the MCTS to be terminated when the full sequence is generated (no need of the denoiser to calculate the reward function).

Reply: In our framework, the sequence in a state s_t is a partially generated sequence, which means there are positions masked that are being filled. And the MCTS terminated when the full sequence was generated and also the reward exceeded the predefined threshold. Importantly, to enable intermediate feedback during MCTS, we must evaluate partially generated sequences before reaching the final state. This is why we propose a jumpy denoising strategy, which completes the partially filled sequence via a fast approximation.

Q12: In the manuscript, the authors used ESM-3 as their inverse folding method, but any other protein inverse folding could be combined with the proposed MCTS and jumpy denoiser framework. Did the authors test other models as the inverse folding method?

Reply: To the best of our knowledge, ESM-3 is the first structure-sequence unified model that supports arbitrary masking ratios for sequence completion, which makes it particularly suitable for our planning-based inverse folding framework. Other models, such as ProteinMPNN and KWDesign (which relies on ESM-2), do not support arbitrary masking ratios and typically require retraining or finetuning to enable structure-conditioned generation. This conflicts with several core design principles of ProtInvTree:

1. Stepwise generation: Our method incrementally constructs sequences, with each decision depending on the prior t−1 steps—requiring the model to flexibly accept partially completed sequences.
2. Test-time reward-guided search: ProtInvTree performs planning by querying the foundation model at test time, without retraining or gradient-based updates.
3. Focus-and-grounding with jumpy denoising design, which requires flexible infilling under arbitrary masking ratios.

Thank you for your thoughtful and detailed feedback. We appreciate your positive assessment of the strengths of our work, as well as your suggestions for areas of improvement. We address each of your comments in detail below:

W1 & Q1 & Q2: The claimed benefits of planning (intermediate feedback, uncertainty tracking, backtracking) are not directly validated. There is no comparison to beam search or other simpler baselines that lack these features. Add Beam Search Baseline.

Reply: We compare our MCTS-based planner (ProtInvTree) against two commonly used test-time strategies: Beam Search and Best-of-N sampling. To ensure a fair comparison in computational cost, all methods are restricted to approximately the same number of model calls per design.

W2 & Q5: Novelty & Collapse “Figure 5 shows diversity decreasing with deeper planning, but novelty is not reported. How does novelty behave and when does the search collapse to native‑like sequences?”

Reply: We will add novelty curves in the revised Fig. 5. The key trend is a smooth trade‑off—better scTM comes with a gradual, not abrupt, drop in novelty. Even at our most aggressive settings the average fraction of mutated residues remains > 49 %. A condensed view of the sweep is below:

W3: “Quantify the test‑time compute vs. accuracy trade‑off.”

Reply: We provide a detailed analysis of the compute–performance trade-off. Specifically, for each given backbone, we generate 10 designs and compute the average generation time for each design. We fix the depth and scale the expand number from 2 to 6, as follows:

W4: There is no experimental or functional validation

Reply: Wet‑lab assays are beyond scope. In protein engineering, function is overwhelmingly constrained by 3‑D structure—hence the maxim “structure dictates function.” A design that reproduces the target fold with higher scTMScore is therefore more likely to preserve (or enable) the desired biochemical activity.

Q1 & Q2 & Q3: Intermediate Feedback & Uncertainty. Ablation of Planning Features.

Reply: Starting from the full ProtInvTree, we conducted an ablation study by disabling each key mechanism in isolation to assess its contribution:

We also evaluated the impact of jumpy denoising on runtime performance. Specifically, for each given backbone, we generate 10 designs and compute the average generation time for each design.

Q4: Uncertainty Quantification

Reply: ProtInvTree follows the classic UCT rule (Eq. 1 in the paper): $\text{UCT}(s)=V(s)+\omega\sqrt{\frac{\ln N(\text{parent}(s))}{N(s)}} .$ We treat this visit‑count–based bonus as an implicit measure of epistemic uncertainty—lower $N(s)$ ≈ higher uncertainty. No additional statistics are stored or tuned. We provice more results by setting $w\!=\!0$ (greedy value search) as follows:

Removing the uncertainty‑driven exploration degrades structural quality and shrinks the explored sequence space, demonstrating that the built‑in visit‑count heuristic already captures useful uncertainty information.

Q6: Failure Mode Characterization

Reply: We attempted to explore higher values (as shown in the following) as further extreme cases.

However, we regret that when both the expansion number and planning depth exceed this level, the memory consumption of the model exceeds the capacity of a single A100 GPU (80 GB). Despite this, our current scaling experiments already reveal a clear saturation point: although novelty decreases slightly as both expansion number and planning depth increase, the decline is gradual and does not indicate a full collapse to native-like sequences.

Q7: Parameter Sensitivity Study

Reply: To systematically characterize the diversity–novelty–consistency landscape, beyond the scaling analysis we have already conducted on planning depth and expansion number, we conducted a controlled sweep over several key planning hyperparameters: the UCT exploration weight $\omega$, the sampling temperature $T$, the reward threshold used for termination, and selection number schedules.

Q8: Recovery-Oriented Trade-off

Reply: We agree that introducing an explicit recovery-based reward would bias the search toward the native target. But we do not perform this ablation because adding a recovery term would require revealing the ground‑truth sequence during generation.

Q9: Structure Complexity Effects.

Reply: We thank the reviewer for this valuable question. As partially reflected in Table 1, our evaluation already includes a breakdown of model performance across different protein lengths (i.e., "Short" proteins with length ≤ 100, "Single-chain" proteins, and "All" proteins).

The fold classes results are shown below:

Thank you for your detailed and constructive feedback on our work. We appreciate your insights into both the strengths and weaknesses of our approach, and we address each of your points below:

W1: Comparisons to other test-time scaling methods are lacking. Specifically, the author should compare MCTS to simpler methods like beam search and best-of-N to demonstrate the need for a more complex MCTS approach.

Reply: We compare our MCTS-based planner (ProtInvTree) against two commonly used test-time strategies: Beam Search and Best-of-N sampling. To ensure a fair comparison in computational cost, all methods are restricted to approximately the same number of model calls per design.

Our experiments show that ProtInvTree consistently outperforms these baselines, which we attribute to its ability to leverage intermediate structural feedback and to perform both lookahead planning and backtracking-based revision.

W2: The explanations for the Jumpy Denoising strategy are unclear.

Reply: We clarify that the Jumpy Denoising strategy is designed to efficiently estimate the final reward in MCTS without simulating full denoising trajectories. Specifically, instead of performing step-by-step denoising from noise to the final sequence, we adopt a single-step DDIM-based reverse sampling method that directly samples a denoised sequence Sample $\tilde{\mathbf{x_{T}}}$ from an intermediate state $\mathbf{x_{t+1}}$, conditioned on $\mathbf{c}$. Notably, the jumpy denoiser is not separately trained by us. It leverages the pretrained ESM-3 model, which is capable of filling in arbitrary mask ratios through its masked denoising objective.

W3: MCTS also brings additional computational costs, which should be noted and discussed.

Reply: Thank you for raising this important concern. We agree that test-time computational efficiency is critical, particularly for methods that involve large pretrained models and structure prediction. We provide a detailed analysis of the compute–performance trade-off. Specifically, for each given backbone, we generate 10 designs and compute the average generation time for each design. We fix the depth and scale the expand number from 2 to 6, as follows:

We also scaled the number of sampling iterations for ESM-3 to assess its performance–efficiency trade-off. The resulting sc-TMScore and average inference time per design are summarized in the table below:

We also evaluate the method of Best of N, which generates 10 designs and finds the best one with 58.19, with an average sc-TMScore of 0.839.

Q1: What is J(.) in Equation 10, and how does it save computation compared to standard decoding?

Reply: J(·) approximates the single-step reverse denoising $p(\mathbf{x_T} \mid \mathbf{x_{t+1}},\mathbf{c})$, This contrasts with standard diffusion decoding $\prod_{s=t+1}^{T-1} p_{\theta}(\mathbf{x}_{s+1} \mid \mathbf{x}_s, \mathbf{c})$, which requires iterative denoising steps. We bypass this iterative process and generate the final sample in one step, significantly reducing the computational cost for each node evaluation in the MCTS procedure.

Q2: What’s DDIM in line 214?

Reply: DDIM refers to Denoising Diffusion Implicit Models [1], a sampling method used in diffusion models that accelerates the generation by jumpy denoising. In our method, we apply a single-step DDIM sampling to approximate the final denoised sequence, substantially reducing computational cost during MCTS rollout. We will add the citation [1] in the final version for clarity.

Q3: Line 135: MPD → MDP

Reply: Thank you for pointing this out. We will correct the typo in the revised version of the manuscript.

Q4: Address the computational cost in the limitations section.

Reply: Yes, we discuss the computational cost in reply of W3, and we will add the limitations in our revised version.

[1] Jiaming Song et al, Denoising Diffusion Implicit Models, ICLR 2021