教學:microsoft/retrochimera — 前沿 Ensemble 逆合成模型完整指南

第 1 章:專案定位與價值主張

這個專案解決什麼問題?

在 drug discovery(藥物發現)的 lead optimization(先導化合物優化)流程中,medicinal chemist(藥物化學家)設計出新的 analog(類似物)後,最關鍵的問題是:這個分子做得出來嗎? Retrosynthesis(逆合成分析)就是回答這個問題的核心技術——從目標分子反推可行的合成路徑,直到所有起始物料(building block)都能商業購買。

現有的 AI retrosynthesis 模型主要分兩類:

  1. Template-based(基於模板):從已知反應中提取 reaction template(反應模板),預測哪些模板可套用在目標分子上。優點是化學上較可靠,缺點是受限於訓練資料中出現過的模板。
  2. Template-free / de novo(無模板 / 從頭生成):用 seq2seq 模型直接從 product SMILES 生成 reactant SMILES。優點是不受模板限制,缺點是可能產生化學上不合理的反應(hallucination)。

RetroChimera 的核心創新是:把這兩種方法當作互補的 inductive bias(歸納偏置),透過 learned ensemble(學習式集成)結合,同時取得 template-based 的可靠性與 template-free 的創造力。在 Pistachio 資料集上,RetroChimera 的 top-1 accuracy 遠超所有現有模型,且在 blind test 中被工業有機化學家更偏好。

對 lead optimization 工作流的價值

需求RetroChimera 如何滿足
合成步驟不要太多接入 syntheseus 的多步搜尋(retro*、MCTS),可設定最大深度
Building block 要可商業購買syntheseus 支援自訂 building block 庫作為終止條件
預測要可靠Ensemble 結合兩種互補模型,大幅降低 hallucination 風險
批量處理新 analog支援 batch inference,一次丟多個分子
排名有信心度每個預測附帶 probability(概率),方便篩選

誰適合使用?

角色使用情境
Computational chemist(計算化學家)Lead optimization 產出新 analog 後做 retrosynthesis 可行性評估
Medicinal chemist(藥物化學家)評估合成路徑的化學合理性
Bioinformatics analyst(生物資訊分析師)整合進自動化 DMTA cycle 管線
CASP 平台開發者整合到 CASP(Computer-Aided Synthesis Planning)平台

第 2 章:安裝指南

環境需求

  • Python 3.9(固定版本,因 rdchiral_cpp 相容性限制)
  • CUDA 12.1(推論用 GPU 加速,CPU 亦可但較慢)
  • Conda(管理 RDKit + PyTorch + PyG 的複雜依賴)

步驟一:建立 Conda 環境

1# 下載 retrochimera
2git clone https://github.com/microsoft/retrochimera.git
3cd retrochimera
4
5# 建立環境(約 10-15 分鐘,依網速)
6conda env create -f environment.yml
7conda activate retrochimera

environment.yml 會安裝以下核心依賴:

  • RDKit 2023.09.6:化學資訊處理
  • rdchiral_cpp:反應模板應用引擎(C++ 版,比 Python 版快 10-100x)
  • PyTorch 2.2.2 (CUDA 12.1):深度學習框架
  • PyTorch Geometric 2.5.2:圖神經網路框架
  • pytorch-lightning 2.2.2:訓練框架

步驟二:安裝 retrochimera

 1# 基本安裝
 2pip install retrochimera
 3
 4# 或從 local clone 安裝(開發用)
 5pip install -e .
 6
 7# 若需要跑 USPTO-50K checkpoint(使用 Graphium GNN)
 8pip install retrochimera[graphium]
 9
10# 完整安裝(含測試 + graphium)
11pip install retrochimera[all]

步驟三:下載 Checkpoint

1# Pistachio checkpoint(主力,效能最強)
2wget https://figshare.com/ndownloader/files/59468882 -O pistachio_checkpoint.tar.gz
3tar xzf pistachio_checkpoint.tar.gz
4
5# USPTO-50K checkpoint(基準測試用)
6wget https://figshare.com/ndownloader/files/59511926 -O uspto50k_checkpoint.tar.gz
7
8# USPTO-FULL checkpoint(大規模基準測試用)
9wget https://figshare.com/ndownloader/files/59494598 -O ustpo_full_checkpoint.tar.gz

驗證安裝

 1import retrochimera
 2print(retrochimera.__all__)
 3# 預期輸出:
 4# ['BasicTemplateClassificationModel', 'RetroChimeraDeNovoModel',
 5#  'RetroChimeraEditModel', 'RetroChimeraModel']
 6
 7from syntheseus import Molecule
 8mol = Molecule("c1ccccc1")  # benzene
 9print(mol.smiles)
10# 預期輸出:c1ccccc1

第 3 章:核心架構解析

3.1 Ensemble 設計哲學

RetroChimera 的名字來自 chimera(奇美拉;嵌合體),反映了它將兩種根本不同的逆合成模型嵌合的設計。關鍵洞察是:template-based 和 template-free 模型犯的錯誤具有低相關性——當 A 模型預測錯誤時,B 模型往往能給出正確答案。

 1Target Molecule
 2 3     ├──→ RetroChimeraEditModel (Template Localization)
 4     │    ├── GNN Encoder → molecule & atom embeddings
 5     │    ├── Template Classification → 排名 top-K templates
 6     │    ├── Localization Scoring → 精確對齊 template 到 atom
 7     │    └── Rule Application → RDKit/rdchiral 套用模板
 8 9     ├──→ RetroChimeraDeNovoModel (SMILES Transformer)
10     │    ├── Root-Aligned Augmentation → 多起點 SMILES
11     │    ├── Transformer Encoder-Decoder → beam search 生成
12     │    └── Rank Aggregation → 跨 augmentation 排名
1314     └──→ Ensemble Combiner (combine_results)
15          ├── Rank-Based Weight Scoring → 每個 rank 有獨立權重
16          ├── Score Normalization + Softmax → probability
17          └── Final Ranked Output

3.2 子模型一:Template Localization (Edit Model)

這是 RetroChimera 的獨特子模型,不同於傳統 template classification 只預測「用哪個 template」,它額外加了 localization(定位)——預測 template 應該套用在分子的哪些 atom上。

架構細節:

  • Input Encoder:5-layer GNN(hidden_channels=1024),輸出 molecule-level + atom-level representation
  • Rewrite Encoder:5-layer GNN(hidden_channels=192),為每個 reaction template 學習 atom-level embedding
  • Classification Branch:用 molecule-level representation 對所有 template 做 softmax 分類
  • Localization Branch:用 template atom embedding 與 input atom embedding 的矩陣乘法算 localization score

最終分數 = log(classification_prob) + w * localization_score,其中 wlocalization_score_weight(預設 2.25)。

3.3 子模型二:SMILES Transformer (De Novo Model)

基於 Root-Aligned SMILES Transformer 架構:

  • Data Augmentation:對 product SMILES 做 root-aligned augmentation,產生 10 個不同起點的 SMILES 表示
  • Transformer:8-layer encoder-decoder,hidden_dim=512,8 heads,使用 SiLU activation + RMS LayerNorm
  • Beam Search:beam_size=20,每個 augmentation 產生 20 個候選
  • Rank Aggregation:跨 10 × 20 = 200 個候選,用 compute_rank 函式做去重 + 排名

3.4 Ensemble 權重機制

與簡單的「每個模型投一票」不同,RetroChimera 對每個排名位置分配不同權重。例如 template_localization 模型的權重可能是 [10.0, 8.5, 7.2, 6.0, ...],表示 rank-1 預測比 rank-5 預測貴重得多。

權重存在 models.json 中:

1{
2  "template_localization": ["TemplateLocalizationModel", [10.0, 8.5, 7.2, ...]],
3  "smiles_transformer": ["SmilesTransformerModel", [9.0, 7.8, 6.5, ...]]
4}

合併演算法:

  1. 收集所有子模型的預測結果
  2. 對每個 unique reactant set,從各模型取得其 rank
  3. 用對應 rank 的 weight 加總成 ensemble score
  4. 按 ensemble score 排序
  5. 用 softmax(score × temperature / max_possible_score) 轉為 probability

第 4 章:快速開始 — 單步逆合成

4.1 基本使用

 1from retrochimera import RetroChimeraModel
 2from syntheseus import Molecule
 3
 4# 載入 ensemble 模型(首次約 30-60 秒)
 5model = RetroChimeraModel(model_dir="/path/to/pistachio_checkpoint/")
 6
 7# 定義目標分子(SMILES 格式)
 8target = Molecule("Oc1ccc(OCc2ccccc2)c(Br)c1")
 9
10# 執行單步逆合成,取 top 5 預測
11predictions = model([target], num_results=5)
12
13# 印出結果
14for i, pred in enumerate(predictions[0], 1):
15    reactants = ".".join(m.smiles for m in pred.reactants)
16    prob = pred.metadata.get("probability", 0) * 100
17    print(f"  #{i}: {reactants} ({prob:.1f}%)")

4.2 Batch Inference(批量推論)

 1# 準備多個目標分子(例如 lead optimization 產出的 analog)
 2targets = [
 3    Molecule("CC(=O)Oc1ccccc1C(=O)O"),      # aspirin
 4    Molecule("CC(C)Cc1ccc(C(C)C(=O)O)cc1"),  # ibuprofen
 5    Molecule("c1ccc2c(c1)cc1ccc3cccc4ccc2c1c34"),  # pyrene
 6]
 7
 8# Batch inference
 9all_predictions = model(targets, num_results=5)
10
11for mol, preds in zip(targets, all_predictions):
12    print(f"\nTarget: {mol.smiles}")
13    for i, pred in enumerate(preds, 1):
14        reactants = ".".join(m.smiles for m in pred.reactants)
15        prob = pred.metadata.get("probability", 0) * 100
16        print(f"  #{i}: {reactants} ({prob:.1f}%)")

4.3 存取預測 Metadata

每個預測結果附帶豐富的 metadata:

 1pred = predictions[0][0]  # 第一個分子的第一個預測
 2
 3# 基本資訊
 4print(pred.metadata["probability"])       # ensemble 概率
 5print(pred.metadata["score"])             # ensemble 原始分數
 6
 7# 各子模型的獨立排名
 8print(pred.metadata["individual_ranks"])
 9# 例如:{'template_localization': 0, 'smiles_transformer': 2}
10# → template_localization 排第 1,smiles_transformer 排第 3
11
12# 各子模型的詳細 metadata
13print(pred.metadata["individual_metadata"])

第 5 章:與 Syntheseus 整合 — 多步合成搜尋

5.1 為什麼需要多步搜尋?

單步逆合成只告訴你「目標分子可以從哪些前體合成」。但前體本身可能也不是商業可得的 building block,需要繼續往回拆。多步搜尋(multi-step retrosynthetic planning)就是遞迴地拆解,直到所有起始物料都是可購買的 building block。

RetroChimera 透過繼承 syntheseus 的 ExternalBackwardReactionModel 介面,可直接接入 syntheseus 提供的搜尋演算法。

5.2 設定 Building Block 庫

對 lead optimization 工作流而言,building block 可購買性 是最關鍵的約束條件。你需要準備一個 building block SMILES 清單:

1# building_blocks.txt — 每行一個 SMILES
2# 可從以下來源取得:
3# - Enamine REAL Database(~39B)
4# - eMolecules
5# - Sigma-Aldrich / TCI / Alfa Aesar 目錄
6# - 公司內部倉庫 inventory

5.3 使用 syntheseus 的搜尋演算法

1# 透過 RetroChimera 的 CLI 啟動搜尋
2python retrochimera/cli/run_search.py \
3    model_class=RetroChimera \
4    model_dir=/path/to/pistachio_checkpoint/ \
5    search_targets_file=targets.smiles \
6    building_blocks_file=building_blocks.txt \
7    time_limit_s=120 \
8    search.algorithm=retro_star

run_search.py 內部包裝了 syntheseus 的 search.main(),支援的搜尋演算法包括:

  • Retro*:A*-like 搜尋,使用 cost function 引導
  • MCTS(Monte Carlo Tree Search):蒙地卡羅樹搜尋,適合探索未知空間

5.4 控制合成步驟數

透過 syntheseus 的搜尋參數控制:

1# 限制最大搜尋深度(合成步驟數)
2python retrochimera/cli/run_search.py \
3    model_class=RetroChimera \
4    model_dir=/path/to/checkpoint/ \
5    search_targets_file=targets.smiles \
6    building_blocks_file=building_blocks.txt \
7    search.algorithm=retro_star \
8    search.max_depth=5 \
9    time_limit_s=300

第 6 章:使用子模型

RetroChimera 的每個子模型也可以獨立使用:

6.1 Template Localization Model(Edit 模型)

 1from retrochimera import RetroChimeraEditModel
 2from syntheseus import Molecule
 3
 4model = RetroChimeraEditModel(model_dir="/path/to/template_localization/")
 5target = Molecule("Oc1ccc(OCc2ccccc2)c(Br)c1")
 6
 7predictions = model([target], num_results=5)
 8
 9for pred in predictions[0]:
10    print(pred.reactants, pred.metadata["probability"])
11    # 額外可存取:
12    # pred.metadata["template_probability"]  — template classification 概率
13    # pred.metadata["localization_score"]    — localization 分數
14    # pred.metadata["classification_score"]  — classification 分數
15    # pred.metadata["combined_score"]        — 兩者加權組合

6.2 SMILES Transformer Model(De Novo 模型)

 1from retrochimera import RetroChimeraDeNovoModel
 2from syntheseus import Molecule
 3
 4model = RetroChimeraDeNovoModel(model_dir="/path/to/smiles_transformer/")
 5target = Molecule("Oc1ccc(OCc2ccccc2)c(Br)c1")
 6
 7predictions = model([target], num_results=5)
 8
 9for pred in predictions[0]:
10    print(pred.reactants, pred.metadata["probability"])

6.3 Basic Template Classification Model

1from retrochimera import BasicTemplateClassificationModel
2from syntheseus import Molecule
3
4model = BasicTemplateClassificationModel(model_dir="/path/to/template_classification/")
5target = Molecule("Oc1ccc(OCc2ccccc2)c(Br)c1")
6
7predictions = model([target], num_results=5)

第 7 章:訓練自己的模型

7.1 資料準備流程

RetroChimera 的訓練是一條完整 pipeline:

 1原始 reaction SMILES
 2 3 4[1] extract_templates    ← 從反應中提取 SMARTS template
 5 6 7[2] preprocess           ← 將資料轉為模型輸入格式
 8 910[3] build_tokenizer      ← 建立 SMILES tokenizer(僅 transformer 需要)
111213[4] augment_rsmiles      ← Root-aligned SMILES augmentation(僅 transformer 需要)
141516[5] train                ← 訓練模型
171819[6] eval                 ← 評估模型
202122[7] optimize_ensembles   ← 自動搜尋最佳 ensemble 權重

7.2 訓練指令範例

 1# 步驟 1:提取 reaction templates
 2python retrochimera/cli/extract_templates.py \
 3    data_dir=/path/to/reaction_data/
 4
 5# 步驟 2:前處理
 6python retrochimera/cli/preprocess.py \
 7    data_dir=/path/to/reaction_data/ \
 8    model_type=template_localization
 9
10# 步驟 3:訓練 template localization 模型
11python retrochimera/cli/train.py \
12    --config-path=retrochimera/cli/config/pistachio \
13    --config-name=template_localization \
14    data_dir=/path/to/preprocessed_data/
15
16# 步驟 4:訓練 smiles transformer 模型
17python retrochimera/cli/train.py \
18    --config-path=retrochimera/cli/config/pistachio \
19    --config-name=smiles_transformer \
20    data_dir=/path/to/preprocessed_data/

7.3 Ensemble 權重優化

訓練完子模型後,需要用 optimize_ensembles.py 找到最佳 ensemble 權重:

1python retrochimera/cli/optimize_ensembles.py \
2    data_dir=/path/to/data/ \
3    results_dir=/path/to/eval_results/ \
4    output_dir=/path/to/output/ \
5    input_checkpoints_dir=/path/to/checkpoints/ \
6    output_checkpoints_dir=/path/to/ensemble_checkpoint/

此腳本會:

  1. 載入所有子模型的預測結果
  2. 用 differentiable ranking loss 最佳化 rank weights
  3. 嘗試所有子模型組合,找出最佳 ensemble
  4. 將最佳 checkpoint 打包(含 models.json

第 8 章:超參數調校

8.1 Ensemble 層超參數

參數預設值說明
probability_from_score_temperature8.0Softmax temperature;越高→概率分佈越平均

8.2 Template Localization 超參數

參數預設值說明
localization_score_weight2.25Localization 分數相對 classification 的權重
classification_temperature30.0Classification logits 的 temperature
output_temperature1.0最終輸出的 temperature
num_templates_per_result10嘗試的 template 數量倍率
num_processesCPU cores平行 rule application 程序數

8.3 SMILES Transformer 超參數

參數預設值說明
beam_size20Beam search 寬度
augmentation_size10Root-aligned augmentation 數量
max_generated_seq_len512生成 SMILES 最大長度
probability_from_score_temperature3.0Softmax temperature
filter_duplicate_augmentationsTrue過濾重複 augmentation

8.4 推論效能建議

根據 #15 和 #16 的最新最佳化:

  • Template Localization 模型的 localization score computation 已加速
  • 若記憶體允許,可增加 num_processes 加速 rule application
  • num_results 設定 5-10 為推薦值(排名越後 hallucination 風險越高)

第 9 章:對 Lead Optimization 工作流的整合建議

9.1 典型 DMTA Cycle 整合

 1Design → Make → Test → Analyze → (repeat)
 2  │        ▲
 3  │        │
 4  │    RetroChimera:
 5  │    "這個 analog 做得出來嗎?"
 6  │    "合成路徑是什麼?"
 7  │    "Building block 買得到嗎?"
 8 910 新 analog SMILES

9.2 建議的使用守則

  1. 不要相信排名太後面的預測:論文警告「reactions ranked lower in the output list are increasingly likely to be hallucinations」。建議 num_results=5,最多 10。

  2. 一定要讓化學家驗證:論文明確聲明「predictions must be risk-assessed and verified independently by chemistry experts」。RetroChimera 是輔助工具,不是替代化學家判斷的工具。

  3. 注意 out-of-distribution 問題:如果你的 analog 含有不常見的 functional group 或 scaffold,預測可靠性會下降。Pistachio checkpoint 涵蓋範圍最廣,但仍有限。

  4. Batch 處理效率:一次丟一批 analog 比逐一丟效率高,因為 GNN encoder 可以 batch 處理。

  5. Building block 庫更新:building block 可用性會隨時間變化,建議每季度更新一次 building block SMILES 清單。

9.3 與其他工具的搭配

工具角色整合方式
RetroChimera單步逆合成預測model([mol], num_results=5)
Syntheseus多步搜尋 + building block 過濾run_search.py
RDKit分子操作、SMILES 處理RetroChimera 內建使用
Enamine REAL SpaceBuilding block 庫作為 building_blocks.txt 輸入
公司內部 inventoryBuilding block 庫同上

第 10 章:資安掃描報告

掃描結果

風險類別狀態說明
任意程式執行 (eval / exec / os.system)✅ 安全未發現
shell=Truesubprocess 注入✅ 安全subprocess 僅在測試中使用,無外部 input
pickle 反序列化⚠️ 低風險chem/rewrite.py__reduce__ 用於 multiprocessing,屬正常用途
硬編碼 secrets / API keys✅ 安全未發現
外部網路呼叫⚠️ 低風險wandb(可透過 WANDB_MODE=offline 停用)
供應鏈風險⚠️ 注意依賴 syntheseus-root-aligned==0.2.0,需確認來源

建議

  • 推論環境可設定 WANDB_MODE=offlineWANDB_DISABLED=true 防止意外上傳
  • Checkpoint 檔案是 PyTorch .ckpt,載入時使用 torch.load(),確保只從可信來源下載
  • dgllife/ 目錄是從 AWS 的 dgllife v0.3.0 複製的子集,已避免引入完整 dgl 依賴

第 11 章:常見問題與故障排除

Q1:安裝 rdchiral_cpp 失敗

 1# rdchiral_cpp 需要 C++ 編譯器
 2# Linux
 3sudo apt-get install build-essential
 4
 5# macOS
 6xcode-select --install
 7
 8# 若仍失敗,可用純 Python 版
 9pip install rdchiral
10# 但效能會慢 10-100x

Q2:GPU 記憶體不足 (OOM)

1# 減少 batch size
2predictions = model(targets[:10], num_results=5)  # 分批處理
3
4# 或指定 CPU
5model = RetroChimeraModel(model_dir="/path/to/checkpoint/", device="cpu")

Q3:Template Localization 載入很慢

根據 PR #15 和 #16 的最新最佳化,升級到最新版本可顯著加速。若仍覺慢:

1# 減少平行程序數,降低 memory footprint
2model = RetroChimeraEditModel(
3    model_dir="/path/to/template_localization/",
4    num_processes=2
5)

Q4:如何判斷預測是否 hallucination?

  1. 檢查 probability — 低於 5% 的預測要格外警惕
  2. 檢查 individual_ranks — 如果兩個子模型都排前面,可信度較高
  3. 用 RDKit 驗證反應平衡(原子守恆)
  4. 交給化學家做最終確認

Q5:可以用自己的反應資料訓練嗎?

可以,但需要:

  1. 反應資料格式為 reactants>>product SMILES
  2. 至少數千筆反應
  3. GPU(4x A100 或類似等級,Pistachio 訓練配置用 4 GPU)
  4. 完整跑 extract_templates → preprocess → train → eval → optimize_ensembles pipeline

Q6:Pistachio vs USPTO 用哪個?

  • Pistachio:商業資料集,涵蓋範圍最廣、效能最強,推薦用於實際工作
  • USPTO-50K / USPTO-FULL:公開資料集,適合基準測試和學術比較