Unified Biomolecular Trajectory Generation via Pretrained Variational Bridge¶

会议: ICLR 2026
arXiv: 2602.07588
代码: 无
领域: 医学图像
关键词: 分子动力学, 轨迹生成, 变分桥匹配, 预训练, 强化学习微调

一句话总结¶

PVB（Pretrained Variational Bridge）通过编码器-解码器架构结合增强桥匹配，统一了单结构预训练和配对轨迹微调的训练目标，实现了跨领域生物分子轨迹生成，并通过RL微调加速蛋白质-配体holo态探索。

研究背景与动机¶

领域现状：分子动力学（MD）模拟是表征分子行为的基本工具，但计算代价极高（需要飞秒级时间步长）。近年来，深度生成模型开始学习粗化时间步长上的动力学，以高效生成轨迹。

现有痛点：现有方法存在三个关键问题——（1）跨分子系统泛化能力不足；（2）轨迹数据的分子多样性有限，无法充分利用结构信息；（3）单分子模拟居多，蛋白质-配体复合物等多分子系统研究较少。

核心矛盾：最相关的前作 UniSim 虽然通过3D分子预训练实现了跨领域泛化，但预训练（无条件生成单结构 \(x\)）和微调（条件生成轨迹对 \((x_t, x_{t+\tau})\)）之间的训练目标不一致，导致预训练知识传递不充分。

本文目标：（1）如何设计统一的训练框架，使预训练和微调共享相同的生成目标？（2）如何将生成轨迹应用于蛋白质-配体对接的holo态快速探索？

切入角度：引入隐变量 \(\mathbf{Y}_0\)，将生成过程建模为马尔可夫链 \(\mathbf{X}_0 \to \mathbf{Y}_0 \to \mathbf{Y}_1\)，通过变分编码器将初始结构映射到噪声隐空间，再用增强桥匹配解码器传输到目标状态。

核心 idea：通过编码器-解码器+增强桥匹配的统一框架，消除预训练与微调间的目标不一致，并通过基于伴随匹配的RL微调加速holo态探索。

方法详解¶

整体框架¶

PVB 采用编码器-解码器架构。输入为分子构象 \((z, C, x)\)（原子序号、共价键、3D坐标），输出为下一时刻的构象。训练分为三个阶段： 1. 预训练：在大量高分辨率单结构数据上训练，设 \((\mathbf{X}_0, \mathbf{Y}_1) = (x, x)\) 2. 微调：在MD轨迹配对数据上微调，设 \((\mathbf{X}_0, \mathbf{Y}_1) = (x_t, x_{t+\tau})\) 3. RL微调（可选）：用伴随匹配进行强化学习，引导轨迹快速趋向holo态

关键设计¶

变分编码器（Variational Encoder）:
- 功能：将初始状态 \(\mathbf{X}_0\) 映射到隐变量 \(\mathbf{Y}_0\)
- 核心思路：设先验为 \(q_e(d\mathbf{Y}_0|\mathbf{X}_0) \coloneqq \mathcal{N}(x_0, \sigma_e^2 I)\)，其中 \(\sigma_e = \sqrt{0.5}\) Å。通过神经网络 \(\varphi_e\) 学习后验分布 \(p_e\)，最小化KL散度：\(\mathcal{L}_{KL} = -\frac{1}{2}\mathbb{E}[1 + \log \mathbf{V} - 2\log\sigma_e - \frac{\mathbf{V}}{\sigma_e^2}]\)
- 设计动机：引入隐变量的关键目的是防止单结构预训练时条件分布退化为Dirac测度。较大的 \(\sigma_e\) 确保编码过程保留足够的结构信息，同时避免解码退化为平凡情况
增强桥匹配解码器（Augmented Bridge Matching Decoder）:
- 功能：从隐变量 \(\mathbf{Y}_0\) 生成目标状态 \(\mathbf{Y}_1\)，同时保持 \((\mathbf{Y}_0, \mathbf{Y}_1)\) 之间的耦合
- 核心思路：定义布朗桥路径测度，训练向量场 \(\varphi_d\) 最小化 \(\mathcal{L}_{ABM} = \mathbb{E}_{t, (\mathbf{Y}_0, \mathbf{Y}_1)}[\|\varphi_d(t, \mathbf{Y}_0, \mathbf{Y}_t) - \frac{\mathbf{Y}_1 - \mathbf{Y}_t}{1-t}\|^2]\)。推理时通过模拟非马尔可夫SDE \(d\mathbf{Y}_t = \varphi_d^*(t, \mathbf{Y}_0, \mathbf{Y}_t)dt + \sigma d\mathbf{B}_t\) 生成目标
- 设计动机：增强桥匹配确保了端点耦合 \(\Pi_{0,1}\) 在生成过程中被精确保持，这对于忠实重现MD的动力学性质至关重要。由Proposition 1保证了编码器-解码器组合能无偏地估计目标条件分布
基于伴随匹配的RL微调:
- 功能：引入控制向量场 \(u\)，调节生成分布使轨迹快速趋向蛋白质-配体的holo态
- 核心思路：优化KL正则化目标 \(\max_u \mathbb{E}[r(\mathbf{Y}_1) - \frac{\beta}{2}\int_0^1 \|u\|^2 dt]\)。通过Girsanov定理和伴随匹配（adjoint matching），将SOC问题转化为 \(\mathcal{L}_{adj} = \mathbb{E}[\|u(t, \mathbf{Y}_0, \mathbf{Y}_t) + \sigma\tilde{a}(t)\|^2]\)，其中精简伴随状态 \(\tilde{a}\) 通过ODE反向传播
- 设计动机：直接从apo态模拟到holo态需要毫秒级时间尺度，计算上不可行。RL微调通过显式奖励函数引导生成分布，绕过低效的局部探索。使用伴随匹配而非直接梯度累积，实现了内存高效的训练

损失函数 / 训练策略¶

预训练+微调阶段：\(\mathcal{L} = w_{KL} \cdot \mathcal{L}_{KL} + w_{ABM} \cdot \mathcal{L}_{ABM}\)
RL微调阶段：\(\mathcal{L}_{adj}\)，奖励函数为 \(r(\mathbf{X}) = -\text{rmsd}(\mathbf{X}, \mathbf{X}_{ref})\)
控制向量场重参数化：\(u = \frac{1}{\sigma}(\varphi_d^u - \varphi_d^*)\)，无需引入额外网络
预训练数据：PCQM4Mv2 + ANI-1x（小分子）、PDB（蛋白质）、PDBBind2020（蛋白质-配体复合物）

实验关键数据¶

主实验（ATLAS蛋白质轨迹生成）¶

模型	JSD-Rg ↓	JSD-TIC ↓	JSD-MSM ↓	VAL-CA ↑	Decorr-TIC0 ↑
MD (10ns)	0.379	0.684	0.596	0.926	0.000
ITO	0.792	0.400	0.469	0.001	0.714
MDGEN	0.493	0.400	0.463	0.098	0.857
UniSim	0.538	0.372	0.344	0.106	0.786
PVB	0.457	0.371	0.333	0.975	0.929

消融实验 / 蛋白质-配体复合物（MISATO）¶

模型	EMD-ligand ↓	EMD-CoM ↓	RMSE-CONTACT ↓
ITO	0.494	0.479	0.987
UniSim	0.196	0.128	0.049
PVB	0.133	0.089	0.055

关键发现¶

PVB 在 VAL-CA（构象有效性）上远超所有基线（0.975 vs UniSim的0.106），说明生成的构象物理合理性极高
在慢动力学模式（TIC、MSM）上，PVB超过了10ns的MD模拟，且去相关率最高（0.929）
在蛋白质-配体复合物上，PVB比UniSim的配体RMSD误差低32%
RL微调后，在蛋白质-配体对接任务上，PVB的中位数配体RMSD比Vina+PVB（无RL）降低约18%

亮点与洞察¶

统一预训练-微调框架：通过隐变量+增强桥匹配巧妙解决了单结构预训练与配对轨迹微调的目标不一致问题，这一思想可以迁移到其他需要跨任务迁移的生成模型中
VAL-CA的巨大优势：PVB生成的构象几乎不存在键断裂或原子碰撞（97.5%有效），远超ITO的0.1%，这得益于编码器-解码器架构对结构信息的良好保持
伴随匹配的应用：将SOC理论引入分子轨迹生成的RL微调，实现了内存高效的训练，这一范式可推广到其他需要引导生成方向的扩散/流匹配模型

局限与展望¶

目前仅考虑重原子，未建模氢原子和溶剂效应
RL微调需要已知holo态作为奖励信号，限制了在真实盲对接场景中的应用
预训练数据规模仍有限，特别是蛋白质-配体复合物的结构数据
生成轨迹的时间分辨率受限于粗化时间步长 \(\tau\)

评分¶

新颖性: ⭐⭐⭐⭐ 编码器-解码器+增强桥匹配的统一框架设计精巧，伴随匹配RL微调有新意
实验充分度: ⭐⭐⭐⭐⭐ 覆盖了蛋白质单体（ATLAS+mdCATH）和蛋白质-配体复合物（MISATO+PDBBind），评价指标全面
写作质量: ⭐⭐⭐⭐ 数学推导严谨，框架描述清晰
价值: ⭐⭐⭐⭐ 为跨域分子动力学模拟提供了统一高效的方案，构象有效性的巨大提升具有实际应用价值