约束解耦潜在扩散模型：蛋白质反向映射的新突破

《Journal of Chemical Theory and Computation》：Constraint Decoupled Latent Diffusion for Protein Backmapping

本文介绍了一种名为CODLAD（约束解耦潜在扩散）的创新框架，用于解决从粗粒度（CG）结构重建全原子（AA）蛋白质构象的挑战。该研究通过两阶段方法——首先将原子结构压缩为离散潜在表征以嵌入结构约束，随后在潜在空间进行高效去噪扩散——实现了原子精度、构象多样性和计算效率的显著提升。实验表明，CODLAD在多个蛋白质数据集上均达到最先进性能，尤其在未见过的动态系统（如DES数据集）上表现出卓越的泛化能力，为蛋白质结构分析提供了高效可靠的解决方案。

　　

1. 引言

粗粒度（CG）分子动力学模拟通过将原子基团表示为统一珠子来简化系统复杂性，平滑能量景观，从而能够探索蛋白质折叠和构象转变等长时间尺度现象。然而，CG表示固有地牺牲了原子级细节，而这些细节对于分子识别、蛋白质-配体对接和蛋白质-蛋白质相互作用等关键任务至关重要。反向映射（Backmapping）任务旨在从CG表示中重建全原子结构，恢复详细结构分析所需的分辨率。

有效的反向映射方法需要具备原子精度以重建物理真实的原子细节，构象多样性以捕捉广泛的生物相关结构集合，计算效率以及跨不同分子系统的强大泛化能力。传统方法通常基于启发式规则算法生成初始原子结构，随后进行几何优化或能量最小化等细化步骤。然而，这些细化计算成本高，经常因引入非物理伪影（如原子冲突）而损害结构有效性，并且由于其确定性性质难以捕捉多样的构象状态。

尽管生成模型如VAE和GAN在反向映射中提供计算优势，但它们往往难以捕捉分子系统的完整热力学多样性，导致有限的化学可转移性和模式崩溃。最近的研究探索了扩散模型在反向映射中的应用，利用其通过随机采样生成多样化和详细结构的潜力。尽管取得了这些进展，但当前扩散模型在反向映射中的应用仍面临挑战，包括由于残基级去噪导致的计算开销增加和全局几何形状受损，或跨不同构象系统的有限泛化。

为了应对这些挑战，本研究提出了CODLAD（约束解耦潜在扩散），一个新颖的两阶段框架。该框架首先将原子结构压缩成离散潜在表征，明确嵌入结构约束，从而将约束处理与生成解耦。随后，它在此潜在空间中进行高效去噪扩散，以产生结构有效且多样化的全原子构象。

2. 方法

2.1. 预备知识

蛋白质反向映射从其CG表示重建全原子（AA）结构。全原子结构表示为AA = {(x_i, a_i)}_i=1ⁿ，其中x_i表示原子坐标，a_i表示原子类型。类似地，CG结构表示为CG = {(X_i, A_i)}_i=1^N，其中X_i表示珠子坐标，A_i表示氨基酸类型。给定CG结构，目标是生成原子坐标x，其中原子类型a由氨基酸序列决定。该任务可以表述为学习从条件分布p(x|X, A)中采样。

内部坐标用于表示全原子蛋白质结构，以键长d_i、键角θ_i和二面角τ_i编码分子几何形状，而不是绝对笛卡尔位置。这种表示为T = {(d_i, θ_i, τ_i)}_i=1^K，其中K是三元组的总数（通常每个残基最多13个重原子，N个残基为N×13）。内部坐标捕获原子间的相对空间关系，保留了基本结构特征，同时允许跨残基的统一处理。

去噪扩散概率模型（DDPM）通过两个阶段操作：将数据转换为噪声的前向扩散，以及重建原始数据的反向过程。在前向扩散期间，初始数据x₀通过马尔可夫转移逐渐被破坏为高斯噪声x_T。反向过程通过迭代去噪状态{x_t}_t=1^T来重建x₀。神经网络ε_θ(x_t, t)预测每个时间步的噪声。

2.2. CODLAD流程

CODLAD整合了分层编码器-解码器和潜在扩散模型，以实现CG到AA的反向映射。该框架在两个耦合阶段运行：阶段（a）通过分层编码和重构学习全原子结构的残基级潜在表征，满足几何约束；而阶段（b）在此潜在空间中学习条件扩散过程，以生成多样化且物理合理的全原子构象，条件于CG输入。

在训练中，阶段（a）使用SE(3)-等变分层GNN将AA结构编码到紧凑潜在空间并重建它们，同时强制执行几何一致性；原子和残基图之间的跨级消息传递。向量量化规范化潜在流形并产生离散代码，这有助于减轻构象空间中的模式崩溃。在阶段（b）中，扩散模型在此潜在空间中训练，以将噪声潜在映射到干净的、约束一致的潜在，条件于CG图。阶段（a）建立潜在流形和解码器，而阶段（b）学习在此流形上的CG条件采样。

在推理中，给定新的CG输入，CODLAD首先运行阶段（b）以从随机潜在去噪到干净的潜在（条件于CG图），然后应用阶段（a）的解码器获得内部坐标，这些坐标被确定性地转换为笛卡尔坐标以产生最终的AA结构。在推理时不使用编码器。

2.3. 通过压缩和重构解耦约束

给定粗粒度蛋白质结构CG = {(X_i, A_i)}_i=1^N，蛋白质反向映射重建全原子结构x ∈ R^n×3，其中n = ∑_i=1^Nn_i表示原子总数。一个关键挑战在于以高精度捕捉多样构象状态，同时保持效率和对未见蛋白质系统的泛化能力。在全原子空间中直接学习结构细节和执行结构约束增加了复杂性并可能降低重建质量，而后生成细化则存在结构偏差风险。

该方法侧重于获得有效的蛋白质结构表征，通过由键长、键角和扭转角组成的内部坐标来保持结构有效性并捕捉局部变化。该策略将结构约束与生成过程解耦，允许它们在相对简单和轻量级的自动编码器阶段内处理。为了获得蛋白质的潜在表征，首先压缩全原子结构。

然而，直接获得的蛋白质表征显示出低重建精度。与图像表征不同，蛋白质结构需要原子之间的3D几何关系，以方向向量u_ij= (x_j– x_i) /∥x_j– x_i∥和空间距离d_ij= ∥ x_i– x_j∥为特征。这意味着仅编码全局成对距离会导致不足的一维表征。

为了在压缩过程中更好地捕捉3D几何特征，使用双表征表示全原子结构x ∈ R^n×3：低维图X ∈ R^N×3和潜在嵌入h ∈ R^N×d。这里，h编码相对于X中节点的相对原子位置，而图保留残基之间的几何关系。这种双表征在降低维数的同时保留了必要的3D信息。此外，结构约束通过应用于编码器和解码器阶段的几何损失自然解耦。

使用具有分层消息传递的SE(3)-等变图神经网络实现这种双表征。该架构通过联合操作于原子和C_α级图来捕捉局部和全局相互作用，从而反映了我们的两级设计。

蛋白质表示为两个耦合图G_atom= (V_a, E_a)和G_res= (V_r, E_r)，由SE(3)-等变分层GNN联合处理。在原子级，每个节点i ∈ V_a代表一个原子类型为a_i的原子，具有初始嵌入h_i⁽⁰⁾= Embed_a(a_i)。边(i,j) ∈ E_a连接局部截断半径r_c^atom= 9 ?内的原子，捕捉键合和近邻相互作用。每个边特征编码径向和角向几何。

在残基级，每个节点k ∈ V_r代表一个残基（C_α原子），具有残基类型索引A_k和初始嵌入h_k^α(0)= Embed_r(A_k)。残基边(k, l) ∈ E_r连接截断半径r_c^res= 21 ?内的残基，并由几何特征注释。

跨级边(i, k) ∈ E_a?r通过映射i → k和截断r_c^cross= 21 ?将每个原子i连接到其父残基A。

每个编码器层在保持SE(3)-等变性的同时，交替进行级内和级间（原子-残基）消息传递。对于原子和残基特征h_i^(l)和h_k^α(l)，更新定义如公式7所示。这里，N_i^a和N_k^r分别表示原子级和残基级邻域，而N_k^a是属于残基k的原子集合。级间模块φ_res→atom和φ_atom→res实现双向信息交换：残基到原子广播和原子到残基聚合。每个φ(·)作为张量积卷积实现，确保在E(3)对称性下的旋转和平移等变性。

获得潜在表征后，进一步离散化以实现高效稳定的生成。具体使用向量量化将连续特征映射到离散代码，这有助于减轻在连续VAE中观察到的模式崩溃。线性投影首先降低特征维度：h′ = Linear(h) ∈ R^d′(d′ < d)，提高计算效率并平滑潜在空间。然后计算离散表征h_q，如公式8所示，其中E = {e₁, ..., e_K}是大小为K的可学习码本。这种离散编码在保留基本蛋白质构象特征的同时降低了生成复杂性。

在解码阶段，解码器采用SE(3)-不变消息传递网络将残基级潜在表征h映射到每个残基的内部几何约束T = {l_ij, θ_ijk, φ_ijkl}。通过操作于不变成对关系（如残基间距离），解码器捕捉几何形状而不受全局旋转或平移的影响，同时残基类型嵌入为局部结构提供氨基酸特异性先验。预测的内部坐标然后被确定性地转换为笛卡尔坐标，确保重建符合有效的分子几何形状。

为确保几何保真度和化学合理性，训练目标包括以下损失项：重建损失（包括键长损失、键角损失和扭转角损失）、笛卡尔损失、冲突损失和图损失。总损失如公式12所示，其中L_recon指上述三项之和，L_vq是向量量化承诺损失。

2.4. 压缩空间中的潜在扩散

传统的构象集成扩散模型通常需要复杂的约束处理（例如，键长强制执行或残基级去噪）以确保生成过程中的几何有效性。为了简化这一点，扩散在学习的潜在空间中进行，该空间固有地编码结构先验，确保生成过程遵循条件分布p(x|X,A)。

具体地，VAE编码器E将全原子结构x映射到潜在特征h₀= E(x) ∈ H_struct，这是一个保持结构有效性的流形。扩散然后直接在此潜在空间中进行。在训练期间，前向过程q(h_t|h_t–1)逐渐添加高斯噪声，而反向过程p_θ(h_t–1|h_t, X, A)学习去噪，条件于CG图。因为潜在流形H_struct已经通过阶段（a）编码了几何有效性，反向过程不需要额外的坐标级约束，确保生成的结构在物理上与输入CG配置一致。

模型使用标准噪声预测目标进行训练，如公式13所示。条件分布p(x|X, A)通过在潜在空间中的去噪扩散过程学习，如公式14所示。因此，通过学习的反向扩散采样h并解码产生遵循所需条件分布的结构，同时严格遵守解码器施加的几何约束。

去噪网络ε_θ基于ProteinMPNN，接受四个输入：噪声潜在表征h_t、CG坐标X、蛋白质序列A和扩散时间步t。为了指导h_t的一致去噪，X和A作为图架构中的节点和边特征，而t通过自适应层归一化（adaLN）处理，用于归一化层中的时间相关调制。

该过程直接更新潜在特征h_t，而不是索引，允许在细化期间动态调整结构元素（例如，响应原子间相互作用的侧链）。此外，它以非自回归方式运行，在条件于全局CG结构