3位牛津本科生學霸,4個月復現AlphaFold 3直接開源!

新智元報道

編輯:喬楊 好睏

【新智元導讀】5月發佈的AlphaFold 3有論文、沒代碼,讓許多機構和團隊紛紛開啓了「復現AF3」的工作。率先做出成果的,是一家成立不到一年的初創Ligo,3位創始人全都是牛津大學的本科生。

成就如此卓著、意義如此重大,AF3的成果公開自然讓衆多科學家翹首以盼。

然而,DeepMind團隊卻潑下了一盆巨大的冷水。他們只放出了論文,並沒有公佈任何相關的代碼或模型權重。

如今剛剛過去不到4個月,初創公司Ligo宣佈——他們已經完成了AlphaFold 3開源復現的工作。

項目地址:https://github.com/Ligo-Biosciences/AlphaFold3

這個成果相當激動人心,也得到了Figure創始人Brett Adcock的轉發。

Ligo團隊表示,他們正在使用AlphaFold 3的想法來進行酶設計,於是順便開啓了支線任務——復現AF3。

作爲生物分子的結構預測模型,AF3主要可以用於三類任務:

預測蛋白質結構

預測藥物-蛋白質相互作用結構

預測核酸-蛋白質複合物結構

這是結構建模技術的根本進步,整個生物科技行業理應從中受益。其應用範圍廣泛,包括:

CRISPR基因編輯技術:科學家可以準確看到DNA如何與「剪刀」Cas蛋白相互作用

癌症研究:預測潛在藥物如何與癌症靶標結合,AF3論文的亮點之一就是預測KRAS抑制劑與其靶標的複合物

抗體/納米抗體的靶向預測:AlphaFold3在這一類分子上的準確性比現有的最佳工具提高了兩倍

而此次Ligo發佈的模型是在單鏈蛋白質上訓練的,可以完成上述三項功能中的第一項,即預測蛋白質結構,其他兩個功能將在不久後完成訓練併發布。

GitHub倉庫中目前僅公開了代碼,但團隊表示,一旦訓練和基準測試完成就會發布權重,而且會使用Apache 2.0許可證,實現「真正的開源」!

項目地址:https://github.com/Ligo-Biosciences/AlphaFold3

如何「復刻」AlphaFold 3?

由於DeepMind在論文中發佈了模型的完整架構,以及每個組件的僞代碼,因此任何團隊都有權復現。

但復現這件事,說起來容易,做起來難。

Ligo選擇將其完全翻譯爲PyTorch代碼,其中涉及到不少逆向分析和重構工作,遠多出他們的想象。

復現過程中,他們也發現了原始論文中存在的多個問題,會干擾訓練,恰好是深度學習領域的關注點,因此一併放出,供社區參考。

1. MSE損失縮放公式有誤

如下圖所示,公式中使用的是加號而非乘號;如果使用加法,就無法在高噪聲水平下正確降低權重,且MSE在初始化時不是單位化的。

這與Karras等人2022年發表的論文不同,可能只是一個筆誤。

原文地址:https://arxiv.org/abs/2206.00364

2. 論文中省略了原DiT包含的殘差層

代碼中將其添加了回來,並對兩種情況進行了對比實驗,發現引入殘差層可以改善梯度流和收斂性。

3. 當前形式的MSA模塊中存在無效層,

如果使用論文所述的MSA模塊的通信步驟,最後的配對加權平均和轉換層無法對配對錶示(pair representation)做出貢獻,因此沒有梯度。

代碼使用了AlphaFold2中ExtraMsaStack的順序;另一種解決方案是使用權重共享,但論文中是否有此操作尚不明確

除了復刻模型,Ligo團隊也在探索更加快速高效的實現方式。

比如,重用了OpenFold的三角注意力(triangular attention),還將初代AF提出的MSARowAttentionWithPairBias重用於DiT,這是AF3論文中沒有提及的操作。

下面這個動畫由Ligo復現的模型生成(未使用模板),模型僅使用了8個A100 GPU訓練 10 小時。

分析顯示,儘管進行了優化,但模型超過60%的操作仍然受限於內存。因此除了擴展到更多功能,團隊還在借鑑ScaleFold的理念,致力於實現一個更具擴展性、更高效的方案。

「閉源」惹衆怒,3位本科生率先復現

作爲一種生物分子結構預測模型,AlphaFold 3大大加速了蛋白質結構的繪製,解決了一個70多年來一直無法破解的巨大問題,因而對科學進步具有極其重要的意義。

研究人員可能需要用讀完一整個博士的時間(4~6年)才能建模出一個結構,但AlphaFold 3只需幾分鐘,即可獲得與實驗精度相當的預測結果。

但除了博客和論文外,DeepMind只開放了一個服務器,允許科學家們進行非商業用途的使用,每天調用上限爲20次。

網站地址:https://gdm-alphafold.corp.google.com/welcome

之所以如此一反常態,「捂緊」AF3的成果,很可能是有商業盈利用途的考量。

他們和新成立的子公司Isomorphic Labs正在進軍藥物研發領域,以滿足大型製藥公司的需求。

今年年初,Isomorphic Labs已經與諾華和禮來簽訂了價值30億美元的合同,而最新的AF3的開發也有Isomorphic Labs的貢獻,可被用於加速藥物發現。

雖然爲了商業盈利也無可厚非,但這種做法引起了科學家們的強烈不滿。

文章地址:https://undark.org/2024/06/06/opinion-alphafold-3-open-source/

同時也是在批評Nature——爲什麼在沒有公開代碼時接收了論文。

在反對的聲浪下,DeepMind沒有頂住壓力,隨後改口表示,會在6個月內公佈模型及權重。

Nature也不得不站出來迴應,原原本本地向研究者們交代接收論文的考量和依據。

文章地址:https://www.nature.com/articles/d41586-024-01463-0

然而,很多反對者依舊不買賬,畢竟在爭分奪秒的研究領域,6個月實在太久;而且DeepMind出於合作方Isomophic Lab的盈利目的,很可能只會放出「閹割版」。

因此,除了Ligo,很多實驗室和機構都開始了破解、復現AlphaFold 3的工作。

文章地址:https://www.nature.com/articles/s41586-024-07487-w

Ligo所借鑑的OpenFold團隊就是其中之一,由哥倫比亞大學助理教授Mohammed AlQuraishi領頭。

他們此前就對AlphaFold 2進行過復現和重新訓練,成果在今年5月剛剛被Nature接收,代碼也如數公佈。

原文地址:https://www.nature.com/articles/s41592-024-02272-z

倉庫地址:https://github.com/aqlaboratory/openfold

就在發推宣傳這項成果時,Alquraishi表示,AlphaFold 3項目在進行時了。

此外,GitHub上的開源大佬、舊金山的獨立軟件工程師Phil Wang也組織起了一個衆包開源項目,同樣是用PyTorch復現AlphaFold 3,正進行得如火如荼。

項目地址:https://github.com/lucidrains/alphafold3-pytorch

本科生 VS DeepMind

有如此多的競爭者,能在不到4個月的時間率先復現AF3,Ligo這家初創究竟是什麼來頭?

這家YC系初創成立於去年12月,總部位於倫敦,3位創始人都來自牛津,且有豐富的研究經歷。

CEO Edward Harris曾在普林斯頓大學入讀計算機科學系,之後在2021年轉入牛津醫學院,目前正輟學全職創業。

進入牛津前,Harris曾在墨西哥創辦外賣平臺Abas2Go,目前年營業額超過120萬美元。

CSO Emily Egerton-Warburton目前在牛津分子和細胞生物化學系學習,曾獲得英國皇家化學學會頒佈的化學奧賽金獎。

CTO Arda Goreci是牛津大學生物醫學系的學生,2023年憑藉計算生物學方面研究成果入選Google Cloude Research Innovator計劃,他也是AF3開源項目的主要參與者和領導者。

參考資料:

https://x.com/adcock_brett/status/1832812049421807898

https://github.com/Ligo-Biosciences/AlphaFold3

https://news.ycombinator.com/item?id=41448439

https://x.com/ArdaGoreci/status/1830744265007480934