©PaperWeekly 原创 · 作者 | 张一帆

学校 | 中科院自动化所博士生

研究方向 | 计算机视觉

关于分子构象生成的新的工作，文中对 bi-level optimization 问题通过一个 e2e 的模型进行优化，取得了远超当前 SOTA 的结果。

深度学习在分子建模方面已经取得了不错的进展，包括分子性质预测，分子生成等。在这些应用中，分子通常被表示为以原子为节点，共价化学键为边的图。图虽然是有效的，但是在现实中分子是三维结构（也称为构象 Conformation），其中每个原子的特征是三维笛卡尔坐标。这样的 3D 结构更加本质，蕴含的信息更丰富，决定了许多化学和生物学特性。

但是 3D 结构的预测比起图来更加具有挑战性，如何预测有效的分子构象一直是计算化学中一个非常重要的课题。最近的工作首先预测原子之间的距离，然后通过求解距离几何问题，根据预测的距离生成分子构象。这种基于距离几何的方法有效地考虑了分子构象的旋转和平移不变性，因此取得了很好的性能。但是两阶段的方法通常会遭受误差传递的负面影响，即第一个阶段效果不好，第二个阶段只会雪上加霜，同样这里如果第一阶段预测的距离本身就不能支撑一个合理的三维结构，那么第二个生成阶段会生成非常不合理的分子构象。

本文是第一篇尝试使用 e2e 架构解决这个问题的文章。本文提出了一种基于条件 VAE 的架构 ConfVAE，首先分子图会被映射到特征空间，然后通过求解一个双层优化问题来生成他的 3D 结构。

论文标题：

An End-to-End Framework for Molecular Conformation Generation via Bilevel Programming

论文来源：

ICML 2021

论文链接：

https://arxiv.org/abs/2105.07246

代码链接：

https://github.com/MinkaiXu/ConfVAE-ICML21

前置知识

1.1 Problem Definition

给定一个分子图 ，这里 分别是原子集合与边集合，每个 描述源自特征比如种类，每个边 描述两个原子间的化学键。对构象集合 ，每个原子 的特征是他的三维笛卡尔坐标系，整个原子的集合表示为：

本文使用 表示两个原子 之间的距离,那么所有边的距离向量可以表示为 。

分子构象生成问题是一个条件生成过程，给定分子图 ，其目标是模拟分子构象的条件分布 。

1.2 Bilevel Optimization

二层优化简单描述就是外层优化目标  的参数是另一个（内层）优化目标 的解，相应的参数记作 ，那么二层优化可以正式定义为：

由于内层优化不能得到闭式解，因此直接处理上式是不可取的。通常的方法是将内层优化 替换成一个近似解，这个近似解可以通过迭代优化动力学 （比如 SGD）得到。从初始参数 开始，我们运行 次内部优化就能得到一个近似解 ,。

接下来就让我们看看作者是如何解决这个二层优化问题并设计相应的 e2e 框架的。

Implicit Distance Geometry

由于一个分子可以有多个稳定构象，我们用条件变分自编码器（CVAE）模拟了基于分子图的构象的条件化分布（即 ），其中引入了一个潜在变量 来模拟分子构象生成的不确定性。

CVAE 模块包含一个先验分布 和一个解码器 来捕捉给定  时， 的条件分布。在训练过程中还包括一个编码器 ，优化目标即 ELBO：

ELBO 的第一部分即重构误差，第二部分隐空间正则项。实际上 和 都是对角高斯分布 ，，其中均值和方差都是 GNN 算出来的。接下来就让我们看看作者如何将 two-stage 的方法融入这个 CVAE。

内部优化目标：首先我们要知道，直接生成笛卡尔坐标的构象很大程度上依赖于任意旋转和平移。因此之前的工作都会先生成原子之间的距离 ，有了距离之后，通过解决如下距离几何问题就可以得到最终的构象：

外部优化目标：这里我们需要一个重构函数，对旋转和平移具有不变性，因此我们首先通过对齐函数产生另一个构象 ，其中 是我们生成的， 是 reference，对齐函数通过旋转，平移 reference 使得以下评测指标最小：

这里的  是分子数目，然后重构函数就可以写作三维笛卡尔坐标的重构损失之和：

有了内/外部优化目标，我们将这个过程写为一个双层优化的形式：

这个双层优化同样是很难找到解的，原因有 2 (i) 内部优化找不到闭式解。(ii) 在连续的隐层空间求期望是不可行的。下文将介绍作者是如何通过设计 decoder 的结构和损失函数来处理这个问题的。

接下来就让我们通过以下伪代码和模型图对算法细节加以描述：

首先第一个损失函数 就是常见的一个 KL divergence 的正则化手段。重点是第二个损失函数 ，为什么从真实距离向量 预测 呢？

像上面说的那样，我们的解码器由两个级联组成：距离预测模型 从隐变量 解码回一组距离，还有一个可微的距离几何生成过程从距离恢复构象 。 被实现为一个连续归一化流模型（continuous normailizing flow  CNF）将一个从 采样出来的噪声变量也就是初始距离 转化为最终距离 ，转化的过程是依赖于隐变量以及 的。

因为是 flow model，给定真实距离 ， 可以被很容易的通过下式算出来。

因此 可以定义为：

在实践中， 可以作为基于距离的任意辅助目标来监督训练。这也就是为什么在算法中我们先通过 ground truth 的距离向量 反推出 了。

在此之后我们就可以采样距离 然后用 flow model 预测真实距离 。有了预测距离，我们在 inner lop 中剩下的工作就是生成构象了，与之前使用半定规划的方法不同，这里使用了梯度下降进行迭代优化：

给定合适的条件假设并且 的时候，GD 可以收敛到一个正确的构象。总而言之，inner-loop 的梯度传递如下图所示：

现在 inner-loop 已经解决好了，要 e2e 的优化这个二层优化目标，那么就需要计算我们 outer-loop 的目标函数 对于网络参数的梯度了：

由于后一项 对 的梯度 inner-lop 直接梯度反传就可以，所以只要 对 是可微的，那么这个梯度就能算出来，pytorch 也能自动实现梯度的反传。而文中 就是一个简单的重构损失，显然是可微的：

总结一下其实整个模块分成了三大部分，第一部分 CVAE 进行 encode，然后对隐变量空间进行约束。第二步通过流模型计算分子节点的距离，迭代生成构象，完成 inner-loop 的优化，第三步得到构象之后计算重构损失。这样一个 bi-level 的优化问题就在一个框架中得到解决了。

Sampling

训练完模型，我们如何进行生成呢，首先从 中采样隐变量 。然后从高斯分布采样随机距离 ，然后将它传入距离计算网络 ，以  为条件生成距离集合 ，然后通过 inner-loop 的参数和步骤生成 即可。

Experiments

一张图说明 e2e 模型的强大，在多个数据集上将大多数 baseline 远远抛在脑后了。

再来看看可视化结果，比较的对象是当前性能最优的其他模型。

更多阅读

#投 稿 通 道#

 让你的文字被更多人看到 

如何才能让更多的优质内容以更短路径到达读者群体，缩短读者寻找优质内容的成本呢？答案就是：你不认识的人。

总有一些你不认识的人，知道你想知道的东西。PaperWeekly 或许可以成为一座桥梁，促使不同背景、不同方向的学者和学术灵感相互碰撞，迸发出更多的可能性。

PaperWeekly 鼓励高校实验室或个人，在我们的平台上分享各类优质内容，可以是最新论文解读，也可以是学术热点剖析、科研心得或竞赛经验讲解等。我们的目的只有一个，让知识真正流动起来。

???? 稿件基本要求：

• 文章确系个人原创作品，未曾在公开渠道发表，如为其他平台已发表或待发表的文章，请明确标注

• 稿件建议以 markdown 格式撰写，文中配图以附件形式发送，要求图片清晰，无版权问题

• PaperWeekly 尊重原作者署名权，并将为每篇被采纳的原创首发稿件，提供业内具有竞争力稿酬，具体依据文章阅读量和文章质量阶梯制结算

???? 投稿通道：

• 投稿邮箱：hr@paperweekly.site

• 来稿请备注即时联系方式（微信），以便我们在稿件选用的第一时间联系作者

• 您也可以直接添加小编微信（pwbot02）快速投稿，备注：姓名-投稿

△长按添加PaperWeekly小编

????

现在，在「知乎」也能找到我们了

进入知乎首页搜索「PaperWeekly」

点击「关注」订阅我们的专栏吧

·