TR2-D2: Tree Search Guided Trajectory-Aware Fine-Tuning for Discrete Diffusion

Introduction

• TR2-D2 uses tree search to generate reward-guided trajectories for off-policy reinforcement learning in discrete diffusion fine-tuning, where gradients are undefined.Off-policy RL effectiveness is limited by the quality of trajectories generated from the policy.Source: §1 · pp. 1–2

Contributions

• TR2-D2 achieves state-of-the-art discrete diffusion fine-tuning for enhancer-activity regulatory DNA design and multiobjective therapeutic peptide design.These results are reported for Secs. 6.1 and 6.2, respectively.Source: §1 · p. 2
• The framework uses Monte Carlo Tree Search to curate a replay buffer of optimal trajectories for off-policy control-based reinforcement learning.The paper presents this as an efficient discrete diffusion fine-tuning strategy.Source: §1 · p. 2
• TR2-D2 introduces multiobjective discrete diffusion fine-tuning by generating Pareto-optimal replay buffers.The paper identifies this as the first method for multiobjective fine-tuning of discrete diffusion models.Source: §1 · p. 2

Continuous-Time Markov Chains

• A CTMC evolves over a discrete state space, and its generator specifies the transition rate from each state to another.The process is defined over time from 0 to T.Source: §2 · p. 2

Discrete Diffusion Models

• Discrete diffusion learns a CTMC generator that transforms an easy-to-sample prior into the data distribution in finite time, using time-reversed noise-injection and denoising processes.Source: §2 · p. 2
• Masked discrete diffusion (MDM) uses a fully masked prior, independently masks data tokens in the forward process, and iteratively decodes masks conditioned on already unmasked tokens.Its optimal generator has a special decomposition (Ou et al., 2024), which motivates a neural network that predicts conditional data-token probabilities; MDMs are trained with denoising cross-entropy (Ou et al., 2024; Sahoo et al., 2024; Shi et al., 2024).Source: §2 · p. 2

Reinforcement Learning for Discrete Diffusion Models

• RL aligns a pre-trained discrete diffusion model with a reward-tilted target distribution by optimizing an entropy-regularized reward objective (Uehara et al., 2024a).The regularization strength α controls tolerance for deviation from the pre-trained model; the resulting fine-tuning connects to CTMC stochastic optimal control by matching the policy path measure to the optimal path measure (Wang et al., 2025; Zhu et al., 2025c).Source: §2 · p. 3

• The discrete sequence space contains D^L possible sequences, while high-reward sequences may lie in low-density regions rarely sampled by the pre-trained model.This makes the search space intractably large even for modest D and L.Source: §3 · pp. 3–4
• A replay buffer lets off-policy RL repeatedly reinforce search-discovered high-reward sequences, amortizing curation cost and improving later buffer generation.The framework updates the policy over multiple epochs using buffer samples, then resets the buffer.Source: §3 · p. 4
• Although supervised fine-tuning can warm up RL with curated favorable sequences, specialized labeled datasets are sparse for tasks such as biological sequence optimization.Source: §3 · p. 4
• MCTS-based reward-guided inference can find high-reward sequences that serve as a pseudo-warm-up dataset for the discrete diffusion policy.The framework searches for optimal samples, then explores similar samples to build its training buffer.Source: §3 · p. 4
• The search and fine-tuning steps are decoupled, allowing the framework to pair any search algorithm with any off-policy RL algorithm.The paper presents this as a general framework, with a later section implementing it for masked discrete diffusion models.Source: §3 · p. 4

4 TR2-D2: TREE SEARCH GUIDED TRAJECTORY-AWARE FINE-TUNING

• TR2-D2 couples off-policy RL fine-tuning with MCTS-generated replay buffers to optimize reward-guided discrete diffusion trajectories.MCTS constructs the buffers, which are then used for fine-tuning under an off-policy RL objective.Source: §4 · p. 4

4.1 OFF-POLICY RL FOR MASKED DISCRETE DIFFUSION MODELS

• Weighted denoising cross-entropy (WDCE) makes fine-tuning tractable when i.i.d. samples from the reward-tilted target distribution are unavailable.WDCE uses importance sampling over trajectories to simulate the target distribution; its reference policy need not track updates to the fine-tuned policy.Source: §4.1 · p. 5
• Rollouts are stored as clean-sequence and log-importance-weight pairs, and the reference policy is periodically aligned to control weight variance.The implementation uses a stop-gradient copy of the current policy as the reference policy and approximates normalization with a batch softmax over importance weights.Source: §4.1 · p. 5

4.2 STRUCTURED TREE SEARCH FOR BUFFER GENERATION

• MCTS balances exploration and exploitation through selection, expansion, rollout, and backpropagation over unmasking trajectories.Selection samples among top-k children by softmaxing their selection rewards; expansion perturbs the current policy's predictions with i.i.d. Gumbel noise to diversify child sequences.Source: §4.2 · p. 6
• Each rollout completes an expanded sequence, adds its terminal reward to the trajectory log-importance weight, and retains the top-B sequences in the buffer.Backpropagation updates predecessor nodes with terminal rewards from generated sequences.Source: §4.2 · p. 6
• Tree search exploits sampled high-reward trajectories while reusing precomputed, unchanged pre-trained-model log-probabilities during selection.Source: §4.2 · p. 7

Pareto Optimization

• Each tree-search iteration moves the buffer toward the Pareto front, maximizing the hypervolume generated by its reward vectors.Buffer sequences are retained when their rewards are non-dominated by those already in the buffer.Source: §5 · p. 7

Multi-Objective Selection

• MCTS balances objective tradeoffs by exploiting Pareto-optimal paths with vector rewards rather than scalarizing them, while exploring Pareto-optimal nodes.During selection, it computes a K-objective reward vector for each candidate step to estimate future reward.Source: §5 · p. 7

6.1 REGULATORY DNA SEQUENCE DESIGN

• The experiments fine-tuned a pre-trained HepG2 DNA enhancer model and a peptide MDM to optimize enhancer activity and five therapeutic properties, respectively.The DNA study compared guidance and fine-tuning baselines using predicted activity, ATAC accessibility, 3-mer correlation, and pre-trained-model log-likelihood; the peptide study used classifier-defined rewards for binding affinity, solubility, non-hemolysis, non-fouling, and membrane permeability, comparing against unconditional pre-trained sampling and PepTune guidance (Tang et al., 2025).Source: §6.1 · p. 8, §6.2 · pp. 8–9
• 9.78 Pred-Activity (vs. 7.64 for the closest benchmark) and 100% ATAC-Acc made TR2-D2 (α = 0.001) best on activity and accessibility.The method also maintained relatively high 3-mer correlation and log-likelihood. With α = 0.1, it achieved the highest 3-mer correlation to the top 0.1% HepG2 sequences while retaining higher predicted activity and chromatin accessibility than DRAKES.Source: §6.1 · p. 8

6.2 MULTI-OBJECTIVE PEPTIDE SEQUENCE DESIGN

• Tree search during buffer generation significantly improved multiple rewards over fine-tuning iterations compared with scalarized-reward off-policy fine-tuning.TR2-D2 also outperformed PepTune after 200 fine-tuning epochs while minimizing the diversity trade-off observed with increasing iterations.Source: Figure 2 · p. 9, §6.2 · p. 9
• TR2-D2 achieved higher scores across nearly all properties for each protein target than PepTune (Tang et al., 2025), using one diffusion pass.Table 2 reports averages over 100 generated peptides; its caption defines TR2-D2 results after 1000 fine-tuning epochs and PepTune results after 100 inference-time Monte Carlo Tree Guidance iterations.Source: Table 2 · p. 9, §6.2 · p. 9

7 CONCLUSION

• TR2-D2 curates MCTS-optimized sequence buffers for off-policy reinforcement learning and succeeds in single- and multiobjective discrete-diffusion fine-tuning.The authors describe it as a general framework for improving the efficiency and reliability of reinforcement learning with structured search.Source: §7 · p. 10
• The framework could extend to full-length proteins and mRNA with structural and functional constraints, and integrate experimentally validated feedback from high-throughput wet-lab pipelines (Wang et al., 2024; Peng et al., 2025; Vincoff et al., 2025; Patel et al., 2025; Zhao et al., 2024; 2025a;b).Wet-lab feedback could be incorporated into the replay buffer for closed-loop sequence discovery.Source: §7 · p. 10
• The authors identify discrete stochastic optimal control, trajectory-weighting variance reduction, and Pareto-efficient multi-objective optimization as theoretical directions, framing TR2-D2 as a platform for therapeutic design.These directions concern how reinforcement learning and search interact with discrete diffusion processes.Source: §7 · p. 10

OVERVIEW OF APPENDIX

• Appendices A–G cover related work, theory and proofs, experiment details, ablations, hyperparameters, and algorithm pseudocode; the appendix also defines sequence, path-measure, and unmasking-tree notation.The notation includes partially masked sequences as tree nodes and M unmasking steps from a node.Source: OVERVIEW OF APPENDIX · p. 21

A RELATED WORKS

• The related-work review situates the paper among reinforcement-learning diffusion fine-tuning, masked discrete diffusion, inference-time steering and search, and multi-objective generative-model optimization.It covers applications to text-to-image and biomolecular design, discrete guidance methods, and approaches to sampling or optimizing across Pareto tradeoffs.Source: A RELATED WORKS · pp. 21–22

B EXTENDED THEORETICAL BACKGROUND

• The theoretical background connects reinforcement-learning fine-tuning of discrete diffusion models to stochastic optimal control of continuous-time Markov chains, restating earlier results for self-contained exposition.The cited prior results are from Uehara et al. (2024b), Wang et al. (2025), and Zhu et al. (2025c).Source: B EXTENDED THEORETICAL BACKGROUND · p. 22

B.1 CONTINUOUS-TIME MARKOV CHAINS (CTMCS)

• The CTMC preliminaries present the forward equation and derive the Radon–Nikodym derivative between two CTMC path measures, with a corresponding KL-divergence result.The derivation uses the generators and initial distributions of the compared processes.Source: §B.1 · pp. 22–24

B.2 ENTROPY-REGULARIZED DIFFUSION FINE-TUNING

• Entropy-regularized fine-tuning maximizes expected terminal reward while penalizing path-measure KL divergence from the pretrained model; α controls the strength of that constraint.Smaller α permits greater divergence from the pretrained model, while larger α keeps the policy closer to it.Source: §B.2 · p. 25

B.3 FINE-TUNING WITH STOCHASTIC OPTIMAL CONTROL

• The stochastic optimal-control treatment derives the optimal generator and reward-tilted path measure from a value function, and states the value function’s Hamilton–Jacobi–Bellman equation.It also derives the Radon–Nikodym derivative between the optimal and reference path measures.Source: §B.3 · pp. 25–27

C.1 OFF-POLICY LEARNING FOR MASKED DISCRETE DIFFUSION FINE-TUNING

• The off-policy section derives the weighted denoising cross-entropy (WDCE) objective for masked discrete diffusion to match the optimal path measure, building on Zhu et al. (2025c).It describes computing denoising loss over multiple potential trajectories by remasking each clean sample and uses a stop-gradient policy for the weights.Source: §C.1 · pp. 28–29

C.2 JUSTIFICATION FOR THE DECOUPLING OF TREE SEARCH AND FINE-TUNING

• Off-policy RL trains on arbitrary-reference trajectories in a frozen replay buffer and fits its distribution with theoretical guarantees, amortizing search by transferring high-quality samples to the current policy.This decoupling enables integration of any pair of tree-search and off-policy RL algorithms.Source: §C.2 · p. 30

C.3 ACHIEVING PARETO-OPTIMALITY WITH MULTI-OBJECTIVE FINE-TUNING

• Adding candidate rewards to the buffer cannot decrease the hypervolume of its non-dominated rewards.The lemma states that the hypervolume for ND(B ∪ S) is at least that of the original buffer B.Source: §C.3 · p. 30
• With each search iteration, the buffer approaches the Pareto front, where its rewards maximize feasible hypervolume.The proposition considers the Pareto frontier of the feasible search space; its proof does not limit buffer size.Source: §C.3 · pp. 30–31
• The convergence argument assumes sufficient exploration, bounded rewards over the feasible space, and positive probability of an ε-sized hypervolume increase each iteration.The stated assumptions require a discovery probability p > 0 for an increase Δ ≥ ε.Source: §C.3 · p. 31
• The proof defines P⋆ as the Pareto frontier of feasible solutions with maximum feasible hypervolume.For this proof, the buffer set B is not size-limited.Source: §C.3 · p. 31
• Any search algorithm satisfying those assumptions has the stated convergence guarantee, and MCTS satisfies them through exploration and its trained reward oracle.The passage says MCTS gives every path non-zero sampling probability and exploits estimated future rewards to favor paths likely to improve hypervolume.Source: §C.3 · p. 31

D REGULATORY DNA EXPERIMENT DETAILS

• The regulatory DNA experiments largely follow Wang et al. (2025) for their experimental setup and evaluation metrics to ensure fair benchmarking.Source: D REGULATORY DNA EXPERIMENT DETAILS · p. 31

Pre-trained Model

• The pre-trained MDLM-based model uses a CNN backbone and linear noise schedule, trained on 700k 200-base-pair Gosai DNA enhancer sequences.The model follows Wang et al. (2025), MDLM (Sahoo et al., 2024), and the schedule in Stark et al. (2024).Source: §D.1 · p. 31

Enhancer Activity Predictor

• Separate Enformer oracles trained on disjoint chromosome-based halves of the Gosai dataset support fine-tuning and evaluation, with held-out Pearson correlations above 0.85.The evaluation oracle computes Pred-Activity; the split follows Lal et al. (2024), and both oracles follow Wang et al. (2025).Source: §D.1 · p. 31

Fine-Tuning Setup

• Fine-tuning uses frozen weights for log-RND computation, unfrozen weights for optimization, buffer size 128, 128 diffusion steps, and AdamW at η = 3 × 10−4.DNA experiments ran on an NVIDIA H100; evaluation covers 640 sequences across three random seeds and reports mean and standard deviation, consistent with Wang et al. (2025) and Zekri & Boullé (2025). Ablations vary α, MCTS use, Nresample, and Niter.Source: §D.1 · p. 31

D.2 ENHANCER EVALUATION METRICS

• Pred-Activity uses disjoint-split fine-tuning and evaluation oracles, while ATAC-Acc tests predicted enhancer activity with a classifier not directly optimized during fine-tuning.Pred-Activity reports median activity from the evaluation oracle; ATAC-Acc predicts HepG2 chromatin accessibility, where positive accessibility indicates increased enhancer activity.Source: §D.2 · p. 32
• 3-mer Corr measures similarity to high-activity training sequences, while App-Log-Lik measures generated-sequence likelihood under the pre-trained model.3-mer Corr compares generated sequences with the top 0.1% of Gosai HepG2 enhancer sequences; App-Log-Lik uses the ELBO, with larger values indicating higher likelihood and helping identify out-of-distribution generation.Source: §D.2 · p. 32

E.1 EXPERIMENT SETUP

• The peptide generator is a bond-dependent masked discrete diffusion model trained on 11 million peptide SMILES and tokenized with SPE, which has 581 SMILES tokens plus 5 special tokens.Its MDLM masking schedule promotes early unmasking of peptide bond tokens; its RoFormer backbone has 8 Transformer layers and 8 attention heads.Source: §E.1 · pp. 32–33
• Fine-tuning uses frozen pre-trained weights for trajectory log-RND and unfrozen weights for optimization, with five classifiers jointly scoring peptide properties.Each target and hyperparameter set is trained for 1000 epochs; evaluation generates 100 sequences in one pass with 128 diffusion steps and reports mean and standard deviation of predicted rewards. Experiments use an NVIDIA A6000, learning rate η = 10^-4, AdamW, and gradient clipping.Source: §E.1 · p. 33
• TR2-D2 is evaluated for peptide binding to six therapeutically relevant targets: TfR, GLP-1R, GLAST, GFAP, AMHR2, and NCAM1.The passage relates these targets to drug delivery, type-2 diabetes and obesity, neurological disorders, Alexander disease, PCOS therapy, and neuronal migration and synaptogenesis, respectively.Source: §E.1 · p. 33

E.2 THERAPEUTIC PROPERTY CLASSIFIERS

• The five pre-trained classifiers from Tang et al. (2025) provide multi-objective rewards, including a target-binding predictor trained on PepLand peptide–protein affinity data.The binding model embeds protein sequences with ESM-2-650M and peptide SMILES with PeptideCLM, then uses cross multi-head attention; its Spearman correlations are 0.869 on training data and 0.633 on held-out validation data. Binding scores are categorized as weak (< 6.0), medium (6.0–7.5), or high (≥7.5).Source: §E.2 · p. 33
• XGBoost logistic classifiers predict solubility, non-hemolysis, and non-fouling from binary peptide data, with positive-score thresholds of 0.500, 0.800, and 0.450, respectively.Positive non-hemolysis and non-fouling indicate lower red-cell destruction and off-target binding; positive solubility indicates greater water solubility and potential drug loading.Source: §E.2 · p. 33
• An XGBoost regression classifier predicts membrane permeability from 34,853 experimentally validated peptide SMILES, with less negative PAMPA lipophilicity scores indicating stronger permeability.The labeled data come from ChEMBL and CycPeptMPDB.Source: §E.2 · p. 33

E.3 BASELINES AND EVALUATION

• Evaluation generates 100 unconditional pre-trained sequences and compares them with PepTune guidance using 100 MCTG iterations, 128 denoising steps, and M = 50.For in-silico binding-affinity confirmation and visualization, the study uses AutoDock VINA, MGITools, RDKit ETKDG, and PyMol.Source: §E.3 · p. 34

F HYPERPARAMETER DISCUSSION AND ABLATIONS

• The hyperparameter discussion covers epoch, enhancer DNA, and multi-objective peptide ablations, along with MCTS and other hyperparameters, and recommends tuning for new modalities and tasks.The cited analyses include Table 4, Table 6, Table 7, and Figures 3–6.Source: F HYPERPARAMETER DISCUSSION AND… · p. 34

F.1 ABLATION ON MCTS SEARCH

• Removing MCTS from buffer generation worsens enhancer DNA metrics and lowers rewards across all five objectives for multi-objective therapeutic peptide design.The ablation instead fills the buffer using independent forward diffusion passes through the current policy, without gradient tracking, while keeping other hyperparameters unchanged.Source: §F.1 · p. 34

F.2 ABLATION ON NUMBER OF FINE-TUNING EPOCHS

• With 200 and 1000 fine-tuning epochs, TR2-D2 outperforms PepTune across almost all objectives for each protein target; mean rewards plateau at 1000 epochs.The reported settings use N_resample = 20, N_iter = 20, and M = 50.Source: §F.2 · p. 34
• At 200 epochs, rewards increase across most objectives while sequence diversity is maintained; at 1000 epochs, reward optimality comes with lower diversity.The authors identify 200–1000 epochs as a suitable range for multi-objective peptide generation and report that N_resample = 10 instead of 20 significantly increases diversity even after 1000 epochs.Source: §F.2 · p. 34

F.3 HYPERPARAMETER DISCUSSION

• The exploration constant c = 0.1 optimally balanced exploration and exploitation of optimal trajectories.For DNA-enhancer search, setting k = M lets every child be explored, and resetting the tree before each buffer-generation step performed best because fine-tuning inherits sequence-generation ability from the prior tree.Source: §F.3 · pp. 35–36
• α = 0.1 improved correlation to the top 0.1% DNA-enhancer rewards, while α = 0.01 improved benchmark reward optimization; peptide runs used α = 0.1 to retain validity while optimizing multi-objective rewards.The regularization scale controls how far the fine-tuned distribution can diverge from the pre-trained model.Source: §F.3 · p. 35
• Even Niter = 5 improved fine-tuning across multiple rewards, with performance rising steadily as iterations increased.Figure 4 shows mean buffer rewards increasing with more MCTS iterations.Source: §F.3 · p. 35, Figure 4 · p. 38
• Exploring more child sequences increased rewards stored in the buffer and improved performance across multiple objectives.Figure 5 shows steady increases in mean buffer rewards as the number of explored children rises.Source: §F.3 · p. 34, Figure 5 · p. 39
• Nresample = 20 enhanced multi-objective rewards but reduced sequence diversity, whereas Nresample = 10 preserved diversity while optimizing all objectives.The ablation also reports that rewards rise as Nresample increases, as the model gets more training iterations on each buffer.Source: §F.3 · p. 36, Figure 6 · p. 40, Table 7 · p. 41

G ALGORITHMS

• The appendix gives pseudocode for reverse unmasking with log-RND tracking, replay-buffer remasking for WDCE loss, and MCTS buffer generation with trajectory selection and Pareto-front updates.The selection and Pareto-front routines support the single- and multi-reward cases, respectively.Source: G ALGORITHMS · pp. 41–43