YAML Metadata Warning:empty or missing yaml metadata in repo card

Check out the documentation for more information.

TD3B: Transition-Directed Discrete Diffusion for Allosteric Binder Generation

👁 Screenshot 2026-05-10 at 6.52.29 PM

TD3B is a sequence-based generative framework that designs peptide binders with specified agonist or antagonist behavior. It combines a Direction Oracle, a soft binding-affinity gate, and amortized fine-tuning of a pre-trained discrete diffusion model (MDLM).

Installation

conda env create -f env.yml
conda activate td3b
pip install -e .

Demo

An interactive inference demo is provided in notebooks/TD3B_Inference_Demo.ipynb (configured for a Colab T4 GPU).

👁 Open In Colab

To run it: download the notebook from this repository, open Google Colab, upload it via File → Upload notebook, and select a GPU runtime (Runtime → Change runtime type → T4 GPU).

Data and Checkpoints

The checkpoints ship with this repository. The expected layout is:

TD3B/
├── checkpoints/
│ ├── pretrained.ckpt # Pre-trained MDLM weights
│ ├── td3b.ckpt # Fine-tuned TD3B model
│ └── direction_oracle.pt # Direction Oracle weights
├── scoring/functions/classifiers/
│ ├── binding-affinity.pt
│ ├── hemolysis-xgboost.json
│ ├── nonfouling-xgboost.json
│ ├── permeability-xgboost.json
│ └── solubility-xgboost.json
└── tokenizer/
 ├── new_vocab.txt
 └── new_splits.txt

Code Structure

TD3B/
├── inference.py # Generate binders (main inference entry point)
├── finetune_multi_target.py # Multi-target TD3B training
├── launch_multi_target.sh # Training launcher script
├── models/
│ ├── diffusion.py # MDLM backbone (TR2-D2)
│ ├── roformer.py # RoFormer wrapper
│ └── noise_schedule.py # Noise schedules
├── training/
│ ├── finetune_utils.py # Training utilities
│ └── distributed_utils.py # Distributed training helpers
├── mcts/
│ └── peptide_mcts.py # MCTS tree search
├── td3b/
│ ├── direction_oracle.py # Direction Oracle (f_φ)
│ ├── td3b_scoring.py # Gated reward R = g_ψ · σ(d*·(f_φ−0.5)/τ)
│ ├── td3b_losses.py # L_WDCE + λ·L_ctr + β·L_KL
│ ├── td3b_mcts.py # TD3B-extended MCTS
│ ├── td3b_finetune.py # Training loop
│ └── data_utils.py # Data loading utilities
├── scoring/ # Affinity predictor (g_ψ) and property classifiers
├── baselines/ # CG, SMC, TDS, PepTune, Unguided baselines
├── tokenizer/ # SMILES tokenizer (vocab + splits)
├── configs/ # Model and training configs
└── utils/ # Misc utilities

Inference

Generate agonist/antagonist binders for target proteins:

python inference.py \
 --ckpt_path checkpoints/td3b.ckpt \
 --val_csv data/test.csv \
 --save_path results/ \
 --seed 42 \
 --num_pool 32 \
 --val_samples_per_target 8 \
 --resample_alpha 0.1

This generates 32 candidates per (target, direction), scores them with the Direction Oracle and affinity predictor, applies Algorithm 2 weighted resampling, and saves only valid peptide samples.

Output: results/td3b_results_seed42.csv with columns: target, sequence, direction, affinity, gated_reward, direction_oracle, direction_accuracy.

Training

Multi-target TD3B

1. Edit launch_multi_target.sh — set paths to checkpoints, data, and oracle:

BASE_PATH="/path/to/TD3B"
PRETRAINED_CHECKPOINT="${BASE_PATH}/checkpoints/pretrained.ckpt"
TRAIN_CSV="${BASE_PATH}/data/train.csv"
ORACLE_CKPT="${BASE_PATH}/checkpoints/direction_oracle.pt"

2. Launch training:

bash launch_multi_target.sh

Key hyperparameters (in launch_multi_target.sh):

• CONTRASTIVE_WEIGHT=0.1 — λ for L_ctr
• KL_BETA=0.1 — β for L_KL
• SIGMOID_TEMPERATURE=0.1 — τ for gated reward
• NUM_ITER=20 — MCTS iterations per round
• NUM_CHILDREN=16 — Children per MCTS expansion

Baselines

Run baseline methods (CG, SMC, TDS, PepTune, Unguided):

cd baselines/
bash run.sh --baseline cg --device cuda:0
bash run.sh --baseline smc --device cuda:0
bash run.sh --baseline tds --device cuda:0

Citation

@inproceedings{
 cao2026tdb,
 title={{TD3B: Transition-Directed Discrete Diffusion for Allosteric Binder Generation}},
 author={Hanqun Cao and Aastha Pal and Sophia Tang and Yinuo Zhang and Jingjie Zhang and Pheng-Ann Heng and Pranam Chatterjee},
 booktitle={Forty-third International Conference on Machine Learning (Spotlight)},
 year={2026}
}