一、關于stable diffusion
很多人都知道stable diffusion,stable diffusion的出現改變了機器生成領域,讓AI技術第一次無比的接近正常人。大語言模型,AIGC概念于是興起。基于stable diffusion 大家開發了lora, hyperwork等微調技術,使得我們可以精確的控制生成圖片,于是我們可以生成如下一些圖片,以假亂真:
???????
神奇的stable diffusion模型的架構如下:
stable diffusion模型架構簡單理解:首先使用一個編碼器,將圖片編碼到隱式空間的潛向量,然后基于該潛向量實現diffusion model的增噪和去噪過程,生成新的潛向量,然后再使用解碼器,解碼潛向量,生成新圖片。關于stable diffusion的介紹,全網到處都是,感興趣的可以仔細研究研究。
那么問題來了,在分子生成領域,在藥物設計領域,是不是可以有一個stable diffusion 模型?實現真正的可控分子生成?讓藥物設計、分子設計領域的大語言模型真正的興起,實現Drug-AIGC?要實現這一點的難點在于:藥物分子是具有空間幾何特征的(即旋轉不變性和空間平移不變性),且不連續稀疏的。而圖片更簡單,是平面的且連續的。因此,基機器視覺開發的stable diffusion是不能直接使用于分子設計、藥物設計等領域的,需要專門開發一個能處理空間幾何與稀疏不連續的stable diffusion模型架構。
新發表的GEOLDM 模型實現了這一點。也許從此刻開始,變天了。Drug-AIGC真的要來了。分子設計領域有了stable diffusion模型,那么分子設計領域的lora還遠嗎?
二、分子生成領域的stable diffusion - GeoLDM
GEOLDM 模型是發表在2023年 ICML(全稱:International Conference on Machine Learning)會議上的會刊 <Geometric Latent Diffusion Models for 3D Molecule Generation> 文章。第一作者是斯坦福大學計算機系的Minkai Xu。文章鏈接為:https://arxiv.org/pdf/2305.01140.pdf
GEOLDM模型可以生成3D的分子,在有條件或者無條件下都可以生成。
GEOLDM的全稱是Geometric Latent Diffusion Models,幾何隱式擴散模型。GEOLDM 是分子幾何領域的第一個隱式擴散模型(簡稱:DM),由將結構編碼為連續隱式向量的自動編碼器和在隱式空間中運行的擴散模型組成。 文章的關鍵創新在于,對 3D 分子幾何進行建模,通過構建具有不變標量和等變張量的點結構隱式空間來捕獲其關鍵的旋轉平移等變約束。
在GEOLDM模型中,分子中原子坐標和特征,即x,h會,被編碼器編碼成隱式向量,即z_x和z_h。然后在隱式空間內,對z_x和z_h進行T步的加噪過程,使用深度學習網絡預測噪音,反向去噪,生成隱式向量。解碼器根據去噪后的隱式向量生成分子。
注意,GEOLDM的擴散過程是操作在隱式向量空間的。這一點與之前的分子生成模型,EDM,Difflinker,Targetdiff等不同。之前的分子生成模型的擴散過程是直接操作在原子空間的,即添加噪音在坐標x和原子特征h上。而GEOLDM是將擴散過程操作在坐標x和原子特征h,生成的嵌入隱空間z上。這一點完全與機器視覺中的stable diffusion對齊了。
訓練GEOLDM模型的偽代碼:
首先,需要按照常規的變分自動編碼器一樣,在沒有加噪和去噪的擴散過程下,先訓練變分自動編碼器,即First stage。然后,固定自動變分編碼器,加入擴散的加噪和去噪過程,訓練模型學習去噪。所以,可以簡單的理解GEOLDM模型由編碼器,擴散去噪,解碼器,三個AI網絡組成。
使用GEOLDM模型生成分子的偽代碼如下:
三、GeoLDM模型的性能
作者在QM9和DRUG數據上進行了測試,考察了分子生成的質量(有效性,唯一性)。作者對每個模型都進行了三次分子生成,每次生成1w個分子。測試結果如下:
與之前的方法相比,同時具有不變變量和等變變量的潛在空間使GEOLDM在大分子生成的有效性方面實現了高達7%的改進。
stable diffusion模型厲害之處,主要是其非常強大的可改造能力,最簡單的,就是條件控制生成。這里作者做了嘗試。作者嘗試讓GEOLDM模型生成HOMO等更低的分子,結果發現,GEOLDM確實要好于我們之前常用的EDM模型。如下圖:
另一方面,在條件控制生成中,GEOLDM,可以對生成的條件的程度進行控制。作者以生成極化的分子為例,使用了不同的極化控制參數,如下圖:
四、總結
整體感覺GEOLDM這個模型非常好模擬了stable diffusion模型的結構,將擴散過程的加噪去噪過程操作在隱式向量空間內。大大降低了,在分子生成領域的計算成本。同時,因為在潛向量空間內,也打開了分子生成領域的優化可能,我們可以CLIP, LORA, Controlnet 等技術,微調模型,實現精確控制分子生成。
看來Drug-AIGC的日子不遠了,真的不遠了。距離作者正式發表GEOLDM有幾個月了,距離作者開始在GitHub上傳代碼有10個月了。我們將很快看到很多的相關文章出現。老了,腦子跟不上了。