TD3B: Transition-Directed Discrete Diffusion for Allosteric Binder Generation¶
会议: ICML 2026
arXiv: 2605.09810
代码: https://huggingface.co/ChatterjeeLab/TD3B (有)
领域: 医学与药物 / 离散扩散 / 蛋白质生成
关键词: 变构调控、激动剂/拮抗剂、掩码离散扩散、方向 Oracle、门控奖励
一句话总结¶
TD3B 把激动剂/拮抗剂的设计当作「方向性转移算子」生成任务,用一个方向 Oracle + 亲和力门控 + 树搜索摊销微调的掩码离散扩散框架,让预训练肽段生成器学会写出能定向偏移蛋白质活/失活构象转移的多肽序列。
研究背景与动机¶
领域现状:当前主流的结合剂设计方法(RFdiffusion、BindCraft、BoltzGen、RareFoldGPCR 等)都把蛋白质当成一个固定 3D 结构,把任务定义为「稳定某个目标构象/界面」,本质是平衡态结构匹配。
现有痛点:变构调控(尤其 GPCR 的临床药效)取决于结合剂偏移「激活态 ↔ 失活态」转移方向的能力,而不是单纯稳定某一个构象。激动剂、拮抗剂之间的差别是动力学路径上的非对称扰动,纯结构方法没法系统区分它们,只能靠后置过滤或经验偏置,效果有限。
核心矛盾:变构功能本质是动力学/非平衡现象(非可逆、定向),而结构生成模型只编码平衡态先验,二者表征空间根本不匹配——结构-中心的方法压根没法表达「这个结合剂使转移方向偏向活化」这件事。
本文目标:设计一个能 (i) 显式建模激动 vs 拮抗的方向性、(ii) 与亲和力解耦但又只对真结合体生效、(iii) 复用已有强大肽段离散扩散先验的生成框架。
切入角度:作者借用 Markov 状态模型把变构动力学抽象为序列条件转移算子 \(Q(y)=Q_0+\Delta Q(y)\),关键量是有向不对称性 \(\Delta_{ij}(y)=Q(y)(s_i,s_j)-Q(y)(s_j,s_i)\);但实际无法观测其连续值,只能拿到 \(\mathrm{sign}(\Delta(y))\in\{+1,-1\}\) 的离散标签。这给了一个非常诚实的监督口径:不回归动力学速率,只用方向信号。
核心 idea:把方向控制当作摊销目标引导(amortized guidance)叠在预训练 MDLM 上:方向 Oracle 给方向梯度、亲和力模型当软门控,组合成门控奖励再用 TR2-D2 风格的重要性加权去噪做微调。
方法详解¶
整体框架¶
TD3B 三个阶段:(1) 训练目标感知的方向 Oracle \(f_\phi(y,x)\to[-1,1]\),输入靶蛋白 \(x\) + 候选肽 \(y\),输出激动/拮抗倾向;(2) 把预训练的亲和力预测器 \(g_\psi(y,x)\) 当作软门,与方向信号 \(\sigma(d^\star f_\phi(y,x)/\tau)\) 相乘组成门控奖励 \(R(y;d^\star,x)\);(3) 用门控奖励通过重要性加权去噪交叉熵(WDCE)+ 对比损失 + KL 正则微调预训练的掩码离散扩散语言模型 PepTune,并在采样时叠加树搜索探索方向条件下的轨迹空间。整套流程只动序列空间,从不进 3D 结构。
关键设计¶
-
目标感知方向 Oracle \(f_\phi\):
- 功能:给定靶蛋白序列 \(x\) 和候选肽段 \(y\),判断它是激动倾向(\(+1\))还是拮抗倾向(\(-1\))。
- 核心思路:用预训练编码器分别做池化得到 \(h_x,h_y\),再用门控融合 \(z=g\odot h_x+(1-g)\odot h_y\) 其中 \(g=\sigma(W_g[h_x;h_y]+b_g)\),最后过 MLP 出标量分数。监督用带置信度权重的二分类 \(\mathcal{L}_{\text{dir}}=\mathbb{E}[\kappa(y)\log(1+\exp(-d\cdot f_\phi(y,x)))]\),其中 partial agonist 给较低置信度 \(\kappa_{\text{part}}\in(0,1)\)。
- 设计动机:方向信息只有粗粒度的 \(\{+1,-1\}\),强行回归连续动力学速率会把模型带偏;门控融合允许 Oracle 同时利用靶点上下文和结合剂结构,比简单拼接更灵活。
-
亲和力软门控奖励:
- 功能:把方向信号和亲和力捏成单一奖励,避免后置 Pareto trade-off。
- 核心思路:\(R(y;d^\star,x)=g_\psi(y,x)\cdot\sigma(d^\star\cdot f_\phi(y,x)/\tau)\),预训练亲和力模型 \(g_\psi\in[0,1]\) 作为乘性门,方向项做加性偏移。只有「真能结合 + 方向正确」的序列才会拿到高奖励,非结合体被门掉零、方向反的被压低。
- 设计动机:直接把方向作为损失会让模型生成「方向对但根本不结合」的废序列;显式加权方向和亲和力的 Pareto 又难调权重。乘性门把「先要是结合体」当成硬条件,方向只在结合空间里挑。
-
TR2-D2 风格摊销微调 + 方向对比损失:
- 功能:把门控奖励烘进 MDLM 采样分布,并在表示空间显式拉开激动/拮抗。
- 核心思路:训练目标 \(p^\star(y)\propto p_{\theta_0}(y)\exp(S(y)/\alpha)\) 用 WDCE 优化,轨迹级重要性权重 \(w(y_{0:1})\propto\exp(S(y_1)/\alpha)\prod_n p_{\theta_0}/p_{\bar\theta}\) 矫正 proposal 偏差;同时叠加 margin-based 对比损失 \(\mathcal{L}_{\text{ctr}}=\sum_P\|h_\theta(y_i)-h_\theta(y_j)\|^2+\sum_N\max(0,m-\|\cdot\|)^2\) 把同向样本拉近、反向拉远;最后 KL 项把 \(\theta\) 拴在 \(\theta_0\) 附近避免漂走。采样侧用 PepTune 风格的轨迹感知树搜索,由门控奖励指引重要性加权选支。
- 设计动机:纯 RL 在离散扩散上方差大,摊销 + 树搜索结合既有探索(树搜索)又有 amortize 内化(WDCE);对比损失在表示空间额外加压,避免 Oracle 只在分类头层面学到方向差异。
损失函数 / 训练策略¶
总损失 \(\mathcal{L}=\mathcal{L}_{\text{WDCE}}+\lambda_{\text{ctr}}\mathcal{L}_{\text{ctr}}+\lambda_{\text{reg}}\mathcal{L}_{\text{KL}}\)。训练数据 \(\{(x,y,a)\}\) 来自带功能标签(full/partial agonist、antagonist、negative)的肽段-靶点对,negative 不参与方向损失但贡献亲和力门训练。
实验关键数据¶
主实验¶
论文用 GPCR 等临床相关靶点验证 TD3B 能否在「方向选择性」上超过结构基线和推理时引导基线。核心评估维度是同一靶点生成的激动 vs 拮抗序列在功能空间的可分性与亲和力维持。
| 设定 | 评估指标 | TD3B | 结构基线 (RFdiffusion 等) | 关键差异 |
|---|---|---|---|---|
| 激动剂定向生成 | 方向选择性 | 显著正向 | 接近随机 | 结构方法无法编码方向 |
| 拮抗剂定向生成 | 方向选择性 | 显著负向 | 接近随机 | 同上 |
| 亲和力维持 | 预测亲和力 | 与结构基线相当 | 基准 | 门控保证不退化 |
| 推理时引导基线 | 后置过滤后方向 | 不如 TD3B | — | 后置过滤丢吞吐 |
消融实验¶
| 配置 | 现象 |
|---|---|
| Full TD3B | 方向 + 亲和力同时达成 |
| w/o 亲和力门 | 容易生成「方向对但不结合」的废序列 |
| w/o 对比损失 | Oracle 在表示空间方向分离度下降 |
| 把门换成 Pareto 加权 | 权重难调,方向-亲和力此消彼长 |
| 推理时引导而非微调 | 多样性和方向准确率都下降 |
关键发现¶
- 摊销微调比单纯推理时引导更可靠:在离散空间里梯度引导本来就受限,把奖励烘进模型分布是更稳的路径。
- 把亲和力做成软门而非 Pareto 项是关键工程决策;后者会让模型在两个目标间「来回摇摆」。
- 即使方向监督是粗粒度二值标签,对比损失也能在表示空间放大它的可分性。
亮点与洞察¶
- 「方向作为生成目标」:第一个把变构功能的方向性显式作为序列生成的优化方向,而不是后置过滤;这给「功能导向蛋白质设计」开了新接口。
- 门控奖励的设计哲学:把「必要条件(结合)」做成软门、「方向偏好」做成加性项,是一个比 Pareto 加权更干净的多目标融合范式,可迁移到任何「先要 X 再优化 Y」的生成任务。
- 诚实的监督粒度:作者明确不去回归连续动力学速率,只用 \(\mathrm{sign}(\Delta)\),这种「理论框架超前于监督粒度但不强行外推」的处理在生物 ML 里值得借鉴。
局限与展望¶
- 监督只到方向粗粒度,没法直接量化「强度」——临床上 partial agonist 的细分还需更细标签或主动学习。
- 整套方法基于序列空间,不显式建模 3D 接口,复杂构象间通讯路径可能丢失结构特异性。
- Oracle 训练数据集(带功能标签的肽-靶点对)规模受限,GPCR 之外能否泛化未充分验证。
- 树搜索 + WDCE 计算成本不小,相比 inference-only 引导吞吐更慢。
- 亲和力门 \(g_\psi\) 本身是预训练模型,它的偏差会被无缝传染到 TD3B。
相关工作与启发¶
- vs RFdiffusion / BindCraft / BoltzGen: 它们都是结构-中心方法,目标是稳定接触界面;TD3B 把目标改为偏移转移方向,互补而非替代。
- vs PepTune / TR2-D2: 同样基于 MDLM 的引导微调,但任务目标是亲和力或多目标 Pareto;TD3B 引入方向监督把目标拓展到动力学层面。
- vs DRAKES / GLID2E: 都用 RL 风格更新离散扩散策略,TD3B 用更稳定的摊销 + 树搜索路径并把奖励结构化为门控形式。
- vs Classifier Guidance / SMC: 离散域里梯度引导受限,本文用 amortize 化解。
评分¶
- 新颖性: ⭐⭐⭐⭐⭐ 把方向性变构控制做成扩散生成目标,这条线之前几乎是空白
- 实验充分度: ⭐⭐⭐ GPCR 验证是好的起点,但跨家族泛化和真实湿实验验证仍待补
- 写作质量: ⭐⭐⭐⭐ 数学框架(转移算子→方向监督→门控奖励)层层递进,动机非常清晰
- 价值: ⭐⭐⭐⭐ 给 GPCR 这类临床高价值靶点的功能性结合剂设计提供了新范式