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

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

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

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

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

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

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

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

Implicit Distance Geometry

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

CVAE 模塊包含一個先驗分布 和一個解碼器 來捕捉給定? 時， 的條件分布。在訓練過程中還包括一個編碼器 ，優(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ù) ，為什么從真實距離向量 預測 呢？

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

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

因此 可以定義為：

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

在此之后我們就可以采樣距離 然后用 flow model 預測真實距離 。有了預測距離，我們在 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

訓練完模型，我們?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)容還不錯，歡迎將生活随笔推薦給好友。