?PaperWeekly 原創(chuàng) ·?作者?|?張一帆

學校?|?中科院自動化所博士生

研究方向?|?計算機視覺

關(guān)于分子構(gòu)象生成的新的工作，文中對 bi-level optimization 問題通過一個 e2e 的模型進行優(yōu)化，取得了遠超當前 SOTA 的結(jié)果。

深度學習在分子建模方面已經(jīng)取得了不錯的進展，包括分子性質(zhì)預(yù)測，分子生成等。在這些應(yīng)用中，分子通常被表示為以原子為節(jié)點，共價化學鍵為邊的圖。圖雖然是有效的，但是在現(xiàn)實中分子是三維結(jié)構(gòu)（也稱為構(gòu)象 Conformation），其中每個原子的特征是三維笛卡爾坐標。這樣的 3D 結(jié)構(gòu)更加本質(zhì)，蘊含的信息更豐富，決定了許多化學和生物學特性。

但是 3D 結(jié)構(gòu)的預(yù)測比起圖來更加具有挑戰(zhàn)性，如何預(yù)測有效的分子構(gòu)象一直是計算化學中一個非常重要的課題。最近的工作首先預(yù)測原子之間的距離，然后通過求解距離幾何問題，根據(jù)預(yù)測的距離生成分子構(gòu)象。這種基于距離幾何的方法有效地考慮了分子構(gòu)象的旋轉(zhuǎn)和平移不變性，因此取得了很好的性能。但是兩階段的方法通常會遭受誤差傳遞的負面影響，即第一個階段效果不好，第二個階段只會雪上加霜，同樣這里如果第一階段預(yù)測的距離本身就不能支撐一個合理的三維結(jié)構(gòu)，那么第二個生成階段會生成非常不合理的分子構(gòu)象。

本文是第一篇嘗試使用 e2e 架構(gòu)解決這個問題的文章。本文提出了一種基于條件 VAE 的架構(gòu) ConfVAE，首先分子圖會被映射到特征空間，然后通過求解一個雙層優(yōu)化問題來生成他的 3D 結(jié)構(gòu)。

論文標題：

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

給定一個分子圖 ，這里 分別是原子集合與邊集合，每個 描述源自特征比如種類，每個邊 描述兩個原子間的化學鍵。對構(gòu)象集合 ，每個原子 的特征是他的三維笛卡爾坐標系，整個原子的集合表示為：

本文使用 表示兩個原子 之間的距離,那么所有邊的距離向量可以表示為 。

分子構(gòu)象生成問題是一個條件生成過程，給定分子圖 ，其目標是模擬分子構(gòu)象的條件分布 。

1.2 Bilevel Optimization

二層優(yōu)化簡單描述就是外層優(yōu)化目標? 的參數(shù)是另一個（內(nèi)層）優(yōu)化目標 的解，相應(yīng)的參數(shù)記作 ，那么二層優(yōu)化可以正式定義為：

由于內(nèi)層優(yōu)化不能得到閉式解，因此直接處理上式是不可取的。通常的方法是將內(nèi)層優(yōu)化 替換成一個近似解，這個近似解可以通過迭代優(yōu)化動力學 （比如 SGD）得到。從初始參數(shù) 開始，我們運行 次內(nèi)部優(yōu)化就能得到一個近似解 ,。

接下來就讓我們看看作者是如何解決這個二層優(yōu)化問題并設(shè)計相應(yīng)的 e2e 框架的。

Implicit Distance Geometry

由于一個分子可以有多個穩(wěn)定構(gòu)象，我們用條件變分自編碼器（CVAE）模擬了基于分子圖的構(gòu)象的條件化分布（即 ），其中引入了一個潛在變量 來模擬分子構(gòu)象生成的不確定性。

CVAE 模塊包含一個先驗分布 和一個解碼器 來捕捉給定? 時， 的條件分布。在訓(xùn)練過程中還包括一個編碼器 ，優(yōu)化目標即 ELBO：

ELBO 的第一部分即重構(gòu)誤差，第二部分隱空間正則項。實際上 和 都是對角高斯分布 ，，其中均值和方差都是 GNN 算出來的。接下來就讓我們看看作者如何將 two-stage 的方法融入這個 CVAE。

內(nèi)部優(yōu)化目標：首先我們要知道，直接生成笛卡爾坐標的構(gòu)象很大程度上依賴于任意旋轉(zhuǎn)和平移。因此之前的工作都會先生成原子之間的距離 ，有了距離之后，通過解決如下距離幾何問題就可以得到最終的構(gòu)象：

外部優(yōu)化目標：這里我們需要一個重構(gòu)函數(shù)，對旋轉(zhuǎn)和平移具有不變性，因此我們首先通過對齊函數(shù)產(chǎn)生另一個構(gòu)象 ，其中 是我們生成的， 是 reference，對齊函數(shù)通過旋轉(zhuǎn)，平移 reference 使得以下評測指標最小：

這里的? 是分子數(shù)目，然后重構(gòu)函數(shù)就可以寫作三維笛卡爾坐標的重構(gòu)損失之和：

有了內(nèi)/外部優(yōu)化目標，我們將這個過程寫為一個雙層優(yōu)化的形式：

這個雙層優(yōu)化同樣是很難找到解的，原因有 2 (i) 內(nèi)部優(yōu)化找不到閉式解。(ii) 在連續(xù)的隱層空間求期望是不可行的。下文將介紹作者是如何通過設(shè)計 decoder 的結(jié)構(gòu)和損失函數(shù)來處理這個問題的。

接下來就讓我們通過以下偽代碼和模型圖對算法細節(jié)加以描述：

首先第一個損失函數(shù) 就是常見的一個 KL divergence 的正則化手段。重點是第二個損失函數(shù) ，為什么從真實距離向量 預(yù)測 呢？

像上面說的那樣，我們的解碼器由兩個級聯(lián)組成：距離預(yù)測模型 從隱變量 解碼回一組距離，還有一個可微的距離幾何生成過程從距離恢復(fù)構(gòu)象 。 被實現(xiàn)為一個連續(xù)歸一化流模型（continuous normailizing flow ?CNF）將一個從 采樣出來的噪聲變量也就是初始距離 轉(zhuǎn)化為最終距離 ，轉(zhuǎn)化的過程是依賴于隱變量以及 的。

因為是 flow model，給定真實距離 ， 可以被很容易的通過下式算出來。

因此 可以定義為：

在實踐中， 可以作為基于距離的任意輔助目標來監(jiān)督訓(xùn)練。這也就是為什么在算法中我們先通過 ground truth 的距離向量 反推出 了。

在此之后我們就可以采樣距離 然后用 flow model 預(yù)測真實距離 。有了預(yù)測距離，我們在 inner lop 中剩下的工作就是生成構(gòu)象了，與之前使用半定規(guī)劃的方法不同，這里使用了梯度下降進行迭代優(yōu)化：

給定合適的條件假設(shè)并且 的時候，GD 可以收斂到一個正確的構(gòu)象。總而言之，inner-loop 的梯度傳遞如下圖所示：

現(xiàn)在 inner-loop 已經(jīng)解決好了，要 e2e 的優(yōu)化這個二層優(yōu)化目標，那么就需要計算我們 outer-loop 的目標函數(shù) 對于網(wǎng)絡(luò)參數(shù)的梯度了：

由于后一項 對 的梯度 inner-lop 直接梯度反傳就可以，所以只要 對 是可微的，那么這個梯度就能算出來，pytorch 也能自動實現(xiàn)梯度的反傳。而文中 就是一個簡單的重構(gòu)損失，顯然是可微的：

總結(jié)一下其實整個模塊分成了三大部分，第一部分 CVAE 進行 encode，然后對隱變量空間進行約束。第二步通過流模型計算分子節(jié)點的距離，迭代生成構(gòu)象，完成 inner-loop 的優(yōu)化，第三步得到構(gòu)象之后計算重構(gòu)損失。這樣一個 bi-level 的優(yōu)化問題就在一個框架中得到解決了。

Sampling

訓(xùn)練完模型，我們?nèi)绾芜M行生成呢，首先從 中采樣隱變量 。然后從高斯分布采樣隨機距離 ，然后將它傳入距離計算網(wǎng)絡(luò) ，以? 為條件生成距離集合 ，然后通過 inner-loop 的參數(shù)和步驟生成 即可。

Experiments

一張圖說明 e2e 模型的強大，在多個數(shù)據(jù)集上將大多數(shù) baseline 遠遠拋在腦后了。

再來看看可視化結(jié)果，比較的對象是當前性能最優(yōu)的其他模型。

更多閱讀

#投 稿?通 道#

?讓你的文字被更多人看到?

如何才能讓更多的優(yōu)質(zhì)內(nèi)容以更短路徑到達讀者群體，縮短讀者尋找優(yōu)質(zhì)內(nèi)容的成本呢？答案就是：你不認識的人。

總有一些你不認識的人，知道你想知道的東西。PaperWeekly 或許可以成為一座橋梁，促使不同背景、不同方向的學者和學術(shù)靈感相互碰撞，迸發(fā)出更多的可能性。?

PaperWeekly 鼓勵高校實驗室或個人，在我們的平臺上分享各類優(yōu)質(zhì)內(nèi)容，可以是最新論文解讀，也可以是學術(shù)熱點剖析、科研心得或競賽經(jīng)驗講解等。我們的目的只有一個，讓知識真正流動起來。

?????稿件基本要求：

? 文章確系個人原創(chuàng)作品，未曾在公開渠道發(fā)表，如為其他平臺已發(fā)表或待發(fā)表的文章，請明確標注?

? 稿件建議以?markdown?格式撰寫，文中配圖以附件形式發(fā)送，要求圖片清晰，無版權(quán)問題

? PaperWeekly 尊重原作者署名權(quán)，并將為每篇被采納的原創(chuàng)首發(fā)稿件，提供業(yè)內(nèi)具有競爭力稿酬，具體依據(jù)文章閱讀量和文章質(zhì)量階梯制結(jié)算

?????投稿通道：

? 投稿郵箱：hr@paperweekly.site?

? 來稿請備注即時聯(lián)系方式（微信），以便我們在稿件選用的第一時間聯(lián)系作者

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

△長按添加PaperWeekly小編

????

現(xiàn)在，在「知乎」也能找到我們了

進入知乎首頁搜索「PaperWeekly」

點擊「關(guān)注」訂閱我們的專欄吧

·

總結(jié)

以上是生活随笔為你收集整理的Bengio等人新作：基于双层规划的端到端分子构象生成框架的全部內(nèi)容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網(wǎng)站內(nèi)容還不錯，歡迎將生活随笔推薦給好友。