Monte Carlo Tree Diffusion with Multiple Experts for Protein Design

TL;DR

MCTD-ME, Monte Carlo Tree Diffusion with Multiple Experts, which integrates masked diffusion models with tree search to enable multi-token planning and efficient exploration under the guidance of multiple experts to enable multi-token planning and efficient exploration under the guidance of multiple experts.

Abstract

The goal of protein design is to generate amino acid sequences that fold into functional structures with desired properties. Prior methods combining autoregressive language models with Monte Carlo Tree Search (MCTS) struggle with long-range dependencies and suffer from an impractically large search space. We propose MCTD-ME, Monte Carlo Tree Diffusion with Multiple Experts, which integrates masked diffusion models with tree search to enable multi-token planning and efficient exploration under the guidance of multiple experts. Unlike autoregressive planners, MCTD-ME uses biophysical-fidelity-enhanced diffusion denoising as the rollout engine, jointly revising multiple positions and scaling to large sequence spaces. It further leverages experts of varying capacities to enrich exploration, guided by a pLDDT-based masking schedule that targets low-confidence regions while preserving reliable residues. We propose a novel multi-expert selection rule ( PH-UCT-ME) extends Shannon-entropy-based UCT to expert ensembles with mutual information. MCTD-ME achieves superior performance on the CAMEO and PDB benchmarks, excelling in protein design tasks such as inverse folding, folding, and conditional design challenges like motif scaffolding on lead optimization tasks. Our framework is model-agnostic, plug-and-play, and extensible to denovo protein engineering and beyond.

Figures (5)

• Figure 1: Overview of MCTD-ME in diffusion-based protein sequence generation. Nodes are partially denoised sequences expanded via masked diffusion and guided by multi-expert rollouts.
• Figure 2: Masked discrete diffusion step: low-confidence sites (mask) are resampled while confident tokens remain fixed, yielding a progressively “unmasked” sequence.
• Figure 3: Length-binned improvements for inverse folding: mean $\Delta$AAR, $\Delta$ normalized reward, and $\Delta$scTM (Final$-$Baseline) for single-expert vs. multi-expert planners. Larger values indicate better recovery and structural consistency; the multi-expert planner yields consistent gains across bins, with larger margins on long proteins.
• Figure 4: CAMEO 7dz2_C refinement. Light gray: baseline lead. Colored overlay: MCTD-ME edits where improvement is achieved, shaded by closeness to the native backbone (red = very close; yellow = closer but not as close). Gray segments indicate no change from the lead. The predominance of red/yellow across many regions visualizes how MCTD-ME moves the lead toward the native geometry via pLDDT-aware masking and multi-expert selection.
• Figure 5: Motif scaffolding (PDB 7mrx). Light/white segment: fixed input motif; gray ribbon: baseline scaffold (lead). Colored overlay: MCTD-ME’s improved scaffold residues (red = very close to native; yellow = closer), shown only where the scaffold improves relative to the lead after motif-aligned superposition. The localized reddening near the motif indicates targeted refinement of the interface without over-editing stable regions.