AI for Science:重建巴別塔 PROF. MENGDI WANG Professor, Princeton University

摘要

這場演講探討了通用 AI 架構和強化學習如何加速從醫學到物理學等領域的科學研究。內容涵蓋了 AI 代理的演進,從複雜系統到像 Alita 這樣的極簡、自我進化框架,以及生成式 AI 在材料、生物學和晶片技術中創造新穎設計的廣泛影響。最終的願景是將 AI 作為所有科學學科的統一語言。

Highlight

1.

在前 AI 時代,我們有生物學、醫學、化學、物理學,這些是不同的學科。一個人必須在學校和大學裡花費數年時間才能真正了解其中的奧妙。但現在,我們有了一個統一的技術,希望能以可擴展的方式解決科學和工程問題。

2.強化學習

強化學習是關於控制一個隨機系統以優化獎勵。在大型語言模型的背景下,我們有一個可控的系統。在這個系統中,狀態會被初始化為提示(prompt)或前綴(prefix)。然後我們有多層的 Transformer,最終輸出的是 logits。接著有一個隨機抽樣過程來解碼下一個 token。然後,下一個 token 會被附加到原始狀態上,這是一個隨機的狀態轉移。所以我們可以將大型語言模型視為一個特定的控制器或策略。

現在,人們可以應用強化學習來控制這些大型語言模型,以解決一系列問題,從對齊(alignment)開始,就像我們工作坊的標題一樣,透過對齊讓 ChatGPT 變得更加用戶友好。我們也可以進行解碼,甚至可以嘗試控制多個大型語言模型的解碼過程,以提高整體吞吐量和效率。我們還可以透過推理讓模型變得更聰明。

更進一步,我們甚至可以使用這種控制策略來控制更複雜的系統,包括涉及人類、實驗室和製造業的系統。說到底,這一切都是關於找出方法,利用強化學習來優化這個控制策略。

3.

為什麼推理如此困難?推理也是關於控制大型語言模型的解碼過程,但在推理中,我們真正談論的是長期規劃(long-horizon planning),即找出一個非常長的 token 序列來解決一個困難的問題。當我們思考推理和強化學習時,學習最佳策略的複雜性和計算的複雜度會隨著輸出長度呈指數級增長。

一旦我們擁有了能夠推理的大型語言模型,它就可以成為大腦,然後我們可以將這個 AI 大腦與記憶體連接起來,進行規劃並在虛擬世界或現實世界中採取行動。

4.

分而治之的訓練策略

當我們面對一個非常非常複雜的任務時,人類不會試圖獨自解決它。人類會嘗試使用工具,將任務分解到數天內完成,或與協作者合作共同完成。同樣地,我們可以在強化學習中應用分而治之的策略。這在機器人控制中其實是一個非常標準的做法。

這就是分層強化學習(Hierarchical Reinforcement Learning)的思想。

我們將這個想法應用於訓練語言模型來進行數學問題的推理和解決。

高層次的想法是,我們不試圖透過讓模型正確生成每一個 token 來直接解決問題。相反,我們建立了一個包含解決方案模板的庫。每個模板可能是一種技巧,例如柯西-施瓦茨不等式。

當接收到一個新的輸入問題時,模型會在抽象層面上進行直觀推理,並嘗試將一系列模板組合成一個軌跡。然後,模型會根據這個模板軌跡是否被正確編譯,也就是高層次概念和抽象是否正確組合,來獲得分數和獎勵。透過這種方式,我們可以大幅壓縮強化學習訓練的複雜性。

透過這種方式,我們實際上可以顯著降低訓練成本和數據收集成本,並且能夠在各種數學基準測試中,達到與 O1 等級相當的推理能力。

5.

科學領域的數據挑戰與自動化管道

在生命科學等硬科學領域,很難找到一個包含標準化問題和答案的大型測驗庫,能夠反映實際實驗室研究中的細微之處。

因此,我們建立了一個自動化的數據管道,可以將來自 Google 論壇、Slack 訊息等來源的科學討論,轉化、解析、提取,並最終整理成可訓練的結構化數據。這些數據可以用於訓練大型語言模型,也可用於對特定領域的代理進行基準測試。

我們證明,透過將強化學習應用於在整理過的科學討論數據上訓練開源模型,可以將推理能力提高近 20%。

6.

我們合作建立了 CRISPR-GPT。這是世界上第一個經過實驗室驗證的生物醫學代理,它可以進行假設生成、數位研究、實驗方案設計,以及透過與新手人類研究員協作,進行實際的實驗室操作,使他們能夠在第一天就執行先進的實驗室研究。

這個代理系統已經被應用於多種生物醫學場景,從肺癌研究到阿茲海默症治療以及藥物靶點篩選。

逐字稿

AI 模型已經變得如此強大,我們實際上正處於一個可以應用這種相當統一、可推廣的架構來解決一系列任務的時刻,從理解基因組學到證明數學定理。

Anthropic 的創始人 Dario 寫了一篇筆記,我認為大概是去年夏天。在他那篇非常長的筆記中,闡述了他對 AI 的理解以及 AI 將如何推動創新。他說,由 AI 賦能的生物學和醫學將壓縮研究過程。過去可能需要一百年的事情,現在或許十年內就能完成。這將是驚人的,這意味著我們過去可能需要等待許久才能找到某種癌症的特定療法,但現在我們或許可以在有生之年就獲得它。

跨學科的統一技術

這不僅僅適用於生物學和醫學。在前 AI 時代,我們有生物學、醫學、化學、物理學,這些是不同的學科。一個人必須在學校和大學裡花費數年時間才能真正了解其中的奧妙。但現在,我們有了一個統一的技術,希望能以可擴展的方式解決科學和工程問題。

核心方法論:強化學習

讓我們回過頭來談談其中一個核心方法論:強化學習。

強化學習是關於控制一個隨機系統以優化獎勵。在大型語言模型的背景下,我們有一個可控的系統。在這個系統中,狀態會被初始化為提示(prompt)或前綴(prefix)。然後我們有多層的 Transformer,最終輸出的是 logits。接著有一個隨機抽樣過程來解碼下一個 token。然後,下一個 token 會被附加到原始狀態上,這是一個隨機的狀態轉移。所以我們可以將大型語言模型視為一個特定的控制器或策略。

現在,人們可以應用強化學習來控制這些大型語言模型,以解決一系列問題,從對齊(alignment)開始,就像我們工作坊的標題一樣,透過對齊讓 ChatGPT 變得更加用戶友好。我們也可以進行解碼,甚至可以嘗試控制多個大型語言模型的解碼過程,以提高整體吞吐量和效率。我們還可以透過推理讓模型變得更聰明。

更進一步,我們甚至可以使用這種控制策略來控制更複雜的系統,包括涉及人類、實驗室和製造業的系統。說到底,這一切都是關於找出方法,利用強化學習來優化這個控制策略。

克服推理的複雜性

那麼,為什麼推理如此困難?推理也是關於控制大型語言模型的解碼過程,但在推理中,我們真正談論的是長期規劃(long-horizon planning),即找出一個非常長的 token 序列來解決一個困難的問題。當我們思考推理和強化學習時,學習最佳策略的複雜性和計算的複雜度會隨著輸出長度呈指數級增長。所以推理非常非常困難,近乎棘手。

但是,一旦我們擁有了能夠推理的大型語言模型,它就可以成為大腦,然後我們可以將這個 AI 大腦與記憶體連接起來,進行規劃並在虛擬世界或現實世界中採取行動。

讓我們從推理開始。實際上,我們可以相當有效且聰明地教導模型進行推理。這是一個我們去年做的專案,基本上是與 DeepMind 的 AlphaDev 合作的後續工作。我們找到了一種方法來訓練一個 32B 規模的模型,使其達到 O1-mini 的推理能力,而 O1-mini 是 2024 年 1 月時最先進的水平。我的同事們當時都非常懷疑,如何能夠用這麼少的計算資源如此有效地訓練模型進行推理,因為我們都知道,從根本上說,解決某些任務是有其複雜性的。

這個想法其實是一個非常普遍的實踐概念:分而治之。

當我們面對一個非常非常複雜的任務時,人類不會試圖獨自解決它。人類會嘗試使用工具,將任務分解到數天內完成,或與協作者合作共同完成。同樣地,我們可以在強化學習中應用分而治之的策略。這在機器人控制中其實是一個非常標準的做法。

當我們試圖控制一個機器人系統時,這個系統自然地會分解為更高層次的感知和規劃,以及較低層次的物理控制。在較高層次上,是關於感知環境、決定移動到哪裡,以及確定宏觀層面的行動。而在較低層次上,這個較高層次的控制可以應用於人形機器人和狗,但在較低層次上,要實際控制一隻蜘蛛或一隻狗,我們必須從一個預先計算好的基本動作庫中,建立一個物理控制策略,以確定究竟如何控制馬達以及如何精確移動它的手指。這就是分層強化學習(Hierarchical Reinforcement Learning)的思想。

我們將這個想法應用於訓練語言模型來進行數學問題的推理和解決。

高層次的想法是,我們不試圖透過讓模型正確生成每一個 token 來直接解決問題。相反,我們建立了一個包含解決方案模板的庫。每個模板可能是一種技巧,例如柯西-施瓦茨不等式。

當接收到一個新的輸入問題時,模型會在抽象層面上進行直觀推理,並嘗試將一系列模板組合成一個軌跡。然後,模型會根據這個模板軌跡是否被正確編譯,也就是高層次概念和抽象是否正確組合,來獲得分數和獎勵。透過這種方式,我們可以大幅壓縮強化學習訓練的複雜性。

然後,當我們實際使用經過強化學習訓練的推理模型來解決一個問題時,我們將會啟用這種雙層分層結構。同樣地,當我們有一個輸入問題時,我們會配置這個模板軌跡,但這些模板最初是完全空的。接著,在較低層次上,我們不需要額外的訓練,只需調用基礎模型,來實例化並填補每個模板中缺失的步驟,一個模板接一個模板地進行,以提供完整的解決方案。最後,我們將所有已實例化的模板拼接在一起,得到最終的解決方案。

透過這種方式,我們實際上可以顯著降低訓練成本和數據收集成本,並且能夠在各種數學基準測試中,達到與 O1 等級相當的推理能力。

科學領域的數據挑戰與自動化管道

這就是我們在數學領域所做的工作。然而,如果我們想進入更普遍的科學推理領域,會遇到一個巨大的挑戰。在數學領域,很容易找到包含大量問題和答案的數據集。但在生命科學等硬科學領域,很難找到一個包含標準化問題和答案的大型測驗庫,能夠反映實際實驗室研究中的細微之處。

因此,我們建立了一個自動化的數據管道,可以將來自 Google 論壇、Slack 訊息等來源的科學討論,轉化、解析、提取,並最終整理成可訓練的結構化數據。這些數據可以用於訓練大型語言模型,也可用於對特定領域的代理進行基準測試。

我們證明,透過將強化學習應用於在整理過的科學討論數據上訓練開源模型,可以將推理能力提高近 20%。我們還基於這種分層結構,專門訓練了一個路由器模型,它可以為每個問題選擇最佳的後訓練子模型來回答。也就是說,它可以根據問題的內容為每個問題選擇最佳模型。這取得了最先進的性能,超越了市面上最好的商業模型。

換句話說,從日常的技術討論中提取有用的數據是可能的,並且可以利用這類數據來訓練模型,以提高其在專業領域的理解和推理能力。

作為一個副產品,我們還建立了 Genome-Bench,這是世界上第一個基於真實科學討論的推理評估基準。這個基準是關於如何對實驗室中某項技術失敗的原因進行故障排除。我們成功地從討論論壇中提取了超過三千個高品質的問題和答案。

AI 代理的崛起與演進

一旦我們有了一個能夠真正理解情況的推理模型,我們就可以進一步建構代理。

這在剛才的座談會中提到過,我們合作建立了 CRISPR-GPT。這是世界上第一個經過實驗室驗證的生物醫學代理,它可以進行假設生成、數位研究、實驗方案設計,以及透過與新手人類研究員協作,進行實際的實驗室操作,使他們能夠在第一天就執行先進的實驗室研究。

這個代理系統已經被應用於多種生物醫學場景,從肺癌研究到阿茲海默症治療以及藥物靶點篩選。

透過應用這種方法,甚至允許感知、視覺輸入和機器人自動化,我們為量子材料的發現建立了 AI-Agentic Lab。這是與 Princeton 的物理學家和 Princeton 量子研究所合作的項目。在這個專案中,我們試圖將兩層非常薄的石墨烯疊加在一起,然後透過調整石墨烯的配置,將會產生非常新穎的晶體結構。這就是研究人員尋找所謂的超導體的方法。

更進一步,我們可以透過為 AI 代理增加更多的工具和數據來持續擴展其能力。這是一個由 Stanford 和 Genentech 的同仁領導的合作專案。在這個專案中,我們建立了這個我稱之為 Bionet 的,一個全方位的生物醫學 AI 代理,它整合了大量的工具、數據庫和軟體,以自動化生物資訊學和藥物發現研究。

現在我們可能會問,我們該如何建構越來越強大的代理?是透過納入更多工具、建立更多數據庫、擴展機架來進行規模化嗎?這就是正確的方式嗎?

事實證明,有更聰明的方法。

幾個月前,我的一位學生在一個名為 Alita 的新型代理 AI 框架上非常努力地工作。當 Alita 發表時,第一作者 Charles 收到了邀請,在 AI2 和我所說的美國和中國所有前沿實驗室中談論這個框架。

為什麼這件事如此令人驚訝?Alita 是一個超級簡單、極簡主義的代理基礎設施。它不是一個整合了 100 個或 200 個工具的代理。它是一個僅從三個 API 開始的代理。Alita 透過學習自己的工具和重寫自身來解決新的任務。

透過使用這種超級極簡的代理架構,Alita 解決了世界上最困難的代理基準測試 Gaia。那是三個月前。上個月,基於 Alita 框架,我們建立了 Physics Supernova,這是一個用於硬核理論物理的 AI 代理。Physics Supernova 在所有國際物理奧林匹克競賽者中排名第 14。Physics Supernova 的表現實際上與最頂尖的國際物理奧林匹克金牌得主相當。

所以,我認為如果我們現在要重新思考代理的建構,代理並不僅僅是關於不斷地擴展工程規模。代理是關於找出那種極簡的設計,讓 AI 能夠自我進化、自我學習和自我提升。

生成式 AI 的廣泛應用

除了大型語言模型、推理和代理之外,生成式 AI 工具正被應用於幾乎所有研究領域。研究人員應用例如像擴散模型(diffusion models)這樣的生成式 AI 模型來生成機器人控制策略。

我們已經訓練了語言模型和擴散模型來解碼蛋白質、學習基因組序列的嵌入(embeddings)以及生成新穎的 mRNA 疫苗。研究人員也正在訓練生成式 AI 模型來提出晶體結構,而這些結構可能需要人類研究人員花費 800 年才能想出。

最後一個例子,我們也正在應用擴散模型來生成 AI 晶片設計。在這個專案中,我們在 IEEE 的一個會議上獲得了最佳論文獎,並且我們能夠為射頻積體電路(RFIC)的電磁結構進行逆向合成訓練這個擴散模型。我們能夠將合成的時間從幾分鐘壓縮到幾秒鐘。這個專案是美國 Netcast 補助金用於 AI 晶片生成的三個創始獲獎者之一。

結論:重建科學的巴別塔

最後,我們已經討論了 AI 在科學和工程領域的這麼多不同應用。我認為 AI 真正獨特之處在於,AI 以模型和代理的形式,正在將來自不同學科的科學語言,轉譯成相同的 AI 語言——token 序列。我們基本上正在打破不同學科之間的壁壘,並希望這是一個讓人類和 AI 一起重建巴別塔的時代。

就是這樣,非常感謝大家。

💡 對我們的 AI 研究助手感興趣嗎?

使用 AI 技術革新您的研究流程、提升分析效率並發掘更深層次的洞見。

了解更多