一文看懂LLM推理,UCL汪軍教授解讀OpenAI ο1的相關方法

機器之心報道

機器之心編輯部

OpenAI 最近發佈的 o1 系列模型堪稱邁向強人工智能的一次飛躍,其強大的推理能力爲我們描繪出了下一代人工智能模型的未來圖景。近日,倫敦大學學院(UCL)人工智能中心汪軍教授撰寫了一份「LLM 推理教程」,深入詳細地介紹了 OpenAI ο1 模型背後的相關方法。

他將在 10 月 12 號本週星期六早上於香港科技大學(廣州)RLChina 2024 大會(http://rlchina.org/rlchina_2024/)上作相關內容的主題報告,併發布其團隊開發的 LLM 推理開源框架以推動 o1 相關模型的發展。

鏈接:https://github.com/openreasoner/openr/blob/main/reports/Tutorial-LLM-Reasoning-Wang.pdf

o1 的訓練使用了強化學習技術,通過顯式地嵌入一個原生「思維鏈」(NCoT)過程,可出色地完成複雜的推理任務。也就是說,o1 在生成響應之前可通過一步步地推理實現「深度思考」。

從 OpenAI 發佈的數據看,相比於之前的 ChatGPT 4o,o1 在數學和編程任務上的表現要強 5 倍。它在競爭性編程中排名第 89 位,在美國著名的數學奧林匹克資格賽中名列前 500 名,並在物理、生物和化學基準測試中超越了人類博士級的準確度。

o1 的一個關鍵創新是它允許在推理過程中花費更多時間進行推理,這標誌着一種範式轉變:從快速、直接的反應轉向緩慢、深思熟慮、多步驟的推理時間計算。見圖 1。

圖 1:推理時間計算。(a) 自迴歸 LLM 是直接基於給定問題生成答案。(b) 思維鏈和逐步思考的概念則涉及到在得到最終答案之前,整合中間推理步驟。這些重複步驟操作允許 1) 不斷重複訪問之前的輸出,2) 逐步推進到後續推理階段,3) 探索多個推理路徑或軌跡。

有趣的是,在指導人類決策和行爲方面,人類認知中存在兩種相關但不同的認知處理模式,其中每種都有各自不同的大腦回路和神經通路,見圖 2。

圖 2:人類認知和 LLM 的類比。(a) 和 (b) 人類有意識或無意識控制的行爲依賴於部分不同的大腦回路。(a) 人類的無意識控制由一些專門的大腦區域維持,例如前腦島和前補充運動區(pre-SMA)。(b) 而自主控制則涉及更大的網絡,激活頂葉和前額葉內的許多區域。無意識控制通常快速而本能,通常由自動過程驅動,而有意識控制往往涉及更審慎、計算和深入的思考,需要仔細的反思和透徹的分析。

系統 1 思維快速、自動且直觀,毫不費力且通常是無意識的。它依賴於能夠快速處理的神經通路,尤其是在需要快速反應或認知資源受限的情況下。

系統 2 思維是深思熟慮的、費力的和有意識的,涉及集中注意力和分析式推理。它處理信息的速度更慢,常用於複雜的問題求解、邏輯推理和決策任務。

o1 的誕生非常激動人心,因爲 大語言模型(LLM) 現在不僅能使用學習到的模式進行快速響應,而且還能通過思維鏈或其它形式的搜索等機制模擬複雜的推理過程。這就類似於人類的更深度的、步步執行的思考方式。

也正因如此,OpenAI ο1 在科學、編程和數學領域都取得了相當卓越的表現。此外,o1 在 AI 安全和對齊方面也取得了進展。該模型的思維鏈推理爲整合人類價值觀和原則提供了新的機會,從而可提高安全評估和越獄測試的性能。

事實上,在 LLM 領域,思維鏈推理和分步驟思考方法並非新技術。之前已有研究表明,如果在輸入中添加「describe your reasoning in steps」或「explain your answer step by step」這樣的指令或提供少樣本示例,就可以讓 LLM 生成中間推理步驟,進而提升其解決問題的能力,尤其是對於數學和編程任務。

但是,這些方法都基於已有的 LLM,並沒有將思維鏈嵌入到模型本身之中。因此,LLM 無法內化這種學習能力,導致許多研究者在探索如何將其直接整合進模型訓練中。之前人們提出的方法包括收集專門的訓練數據、構建獎勵模型和增加解碼的計算複雜度,但目前還沒有一種方法能大規模地在性能上取得重大突破。

汪軍教授表示,我們目前尚不清楚 OpenAI 的 o1 創新是否植根於模型本身,還是依然依賴於外部提示系統。如果它確實涉及在架構中明確嵌入分步推理,那麼這將是一個重大突破。

在大幅提高性能的基礎上,OpenAI o1 還表明,傳統上在訓練期間應用的擴展原則現在也與推理階段相關了。

這樣一來,就需要考慮給推理階段多分配一些算力了。如果能讓 LLM 通過增加測試時間計算來提升輸出,那便是朝着自我改進式智能體(self-improving agent)邁出的重要一步。

這個研究方向被汪軍教授暫且稱爲LLM 原生思維鏈(LLM-Native Chain-of-Thought/NativeCoT),其應當能夠固有地反映人類系統 2 思維所具有的深思熟慮的分析過程。

不過,考慮到 o1 是一個閉源系統,因此它究竟是如何實現如此強大的推理能力的還依然是一個謎。

在本文中,汪軍教授全面回顧了可能的相關文獻,並探討了這一突破背後可能的核心技術和方法。此外,他還提出了基於近期研究成果實現相應開源版本的方法,以加速該領域的研究。

下面首先將介紹典型自迴歸 LLM 常遇到的兩個挑戰,以強調對世界模型和思維鏈機制的需求。然後將給出一個 MDP 公式,用於將原生 CoT 整合進 LLM(進而得到類似 o1 的推理模型);同時還會探索其實現細節。最後會以文獻評論作結,並給出未來的研究方向。

自迴歸 LLM 面臨的挑戰

這裡就略過自迴歸 LLM 的基礎介紹,重點來看汪軍教授提到的兩個挑戰。

第一個挑戰是預測下一 token 的目標。雖然有些人認爲預測下一 token 可能會造就通用智能(AGI),但汪軍教授表示,僅僅專注於預測下一個詞會限制智能的潛力。爲了得到更深層次的智能,可能需要不同的優化目標和學習範式。

傳統自迴歸 LLM 面臨着一個關鍵挑戰:如何使系統超越其訓練數據的界限並開發出新穎的、可能更優的策略?因爲智能體的表現常常受限於其學習的演示的質量,無法超越其訓練數據所體現出的技能水平。

但是,如果使用數據來開發更深度的理解或世界模型,就有可能實現複雜策略的演進,進而超越訓練數據的限制。

世界模型(world model)代表了智能體對環境的理解。學習和優化這個世界模型,再加上模擬潛在結果的能力,有望極大提升 AI 智能體的能力。這些內部世界模型所提供的模擬能力將能實現深度思考(模擬),從而增強智能體的推理和泛化能力。基於模型的策略(如蒙特卡洛樹搜索 (MCTS))是這種方法的經典例證。向系統 2 型推理的過渡(o1 可能就是一個例證)依賴於建立某種類型的世界模型並利用強化學習(獎勵最大化),而不僅僅是最小化預測誤差。這種方法的轉變可能是 OpenAI o1 強大推理能力背後的關鍵過渡技術之一。

通過將 LLM 的預測能力與強化學習和世界建模的策略深度相結合,像 o1 這樣的 AI 系統可以解決更復雜的問題和實現更復雜的決策過程。這種混合方法既可以實現快速模式識別(類似於系統 1 思維),也可以實現深思熟慮的逐步推理(系統 2 思維的特徵)。這也許能解釋 o1 表現出的強大性能。

第二個挑戰則來自計算複雜性角度:LLM 運行時受到二次計算複雜性的約束。當 LLM 遇到多步數學難題時,這種約束會變得尤爲明顯。

但是,思維鏈卻有望減輕這一限制。其可通過一系列「思維」步驟來擴展響應,由此支持一定數量的額外計算資源;它本質上是一個有限的內存,支持寫入但缺乏刪除或覆蓋的能力。儘管該方法頗具潛力,但它仍然不是一個完全動態的內存系統,並且沒有原生地融入解碼階段。這種必要性使得研究社區亟需超越當前 Transformer 解碼器網絡能力的高級計算架構。事實上,存在這樣的需求:在推理和解碼階段實現類似於蒙特卡洛樹搜索 (MCTS)的基於模型的複雜策略。

這種先進的推理時間計算系統將使 AI 模型能夠維護和動態更新問題空間的表徵,從而促進更復雜的推理過程。這種方法與認知科學中的工作記憶(working memory)概念一致;這對於複雜的問題解決和深度思考至關重要。通過整合這些功能,AI 系統可以模擬多個步驟,評估不同的場景,並做出更明智的決策 —— 類似於人類專家推理的深思熟慮過程。

將 LLM 推理看作是馬爾可夫決策過程

爲了建模問答或問題解答等任務中的推理過程,這裡要將推理的結構調整成 Q → {R} → A 序列的形式。

這種結構允許 LLM 生成一系列推理步驟,從邏輯上將問題 Q 與最終答案 A 聯繫起來。

汪軍教授表示,可以將該推理過程定義爲一個馬爾可夫決策過程(MDP)。MDP 能爲建模推理提供一個靈活的框架。它允許模型自迴歸地生成邁向最終答案的順序推理步驟,同時還通過在每個步驟採樣多條路徑來實現樹結構以獲得備選推理軌跡。通過結合順序推理和分支推理這兩種方法,該模型可以探索各種解決方案,從而創建一個多功能且全面的推理過程。

現在可以使用狀態、動作、策略和獎勵來描述這個推理過程了。其中 LLM 的任務是逐步生成與推理步驟和最終答案相對應的連貫 token 序列。

圖 3:在該形式的馬爾可夫決策過程中,LLM 的任務是逐步生成推理步驟和問題的最終答案。其中,LLM 策略的運作方式是生成 token,而,這些 token 可以形成更高級的推理結構。狀態表示迄今爲止的推理步驟序列,動作對應於選取新推理步驟或最終答案。LLM 策略控制動作的選擇,過程獎勵模型(PRM)的作用是提供有關推理步驟和最終答案質量的反饋。通過優化策略以最大化獎勵,LLM 可以在 PRM 的引導下生成準確且有意義的推理過程。

其中狀態表示迄今爲止的推理步驟順序,而動作對應於選擇新推理步驟或最終答案。LLM 策略控制着動作的選擇,而過程獎勵模型 (PRM)則提供有關推理步驟和最終答案質量的反饋。通過優化策略以最大化獎勵,LLM 可以在 PRM 的引導下生成準確且有意義的推理過程。

詳細的論證過程這裡就不多談了。總之,可以基於此得到 LLM 的世界模型的定義:

定義 1:LLM 的世界模型可以定義爲 (, ),其中:

由於轉換是確定性的並且直接遵循策略,因此過程獎勵模型 (s_t, a_t) 封裝了 LLM 與其環境之間的整個交互,可評估每個推理步驟或 token 對所得出的最終答案的貢獻程度。

實際實現

接下來將介紹如何收集中間推理數據,並使用它來訓練過程獎勵模型(PRM),再利用 PRM 來訓練 LLM 策略,並在解碼階段引導推理過程。

自動獲取推理步驟數據

爲了模擬出高級推理能力,就需要推理軌跡數據。最直接的方法當然是人工標註推理步驟,但這種方法缺點也很明顯。

一種無需人類監督,特別有效的收集數據和提升 LLM 推理的方法是 Self-Taught Reasoner(STaR)。

使用 STaR 方法時,模型會自主生成中間推理步驟並使用它們來驗證其內部推理能力。更方法的基礎是 LLM 有能力通過生成中間步驟 {R_1, R_2, . . . , R_n} 從問題 Q 推理到最終答案 A,並使用自己的策略驗證正確性。

也就是說,該方法首先會採用 LLM 的策略 π_LLM,基於初始問題 Q 和最終答案 A 來生成推理步驟 {R}。

生成 {R} 之後,就要驗證其正確性。這裡可以再次使用這個 LLM 策略。

之後,收集到的 {Q, {R}, A} 就可進一步用於訓練策略 π_LLM,提升有效推理步驟的生成過程。

當推理序列較長時,還會用到蒙特卡洛樹搜索(MCTS)。

自我增強式訓練

如圖 4 所示,PRM v (s) 和 LLM 策略 π_LLM 可以相互增強以實現自我提升。

圖 4:將 PRM 的價值函數與 LLM 的策略生成相結合,可確保得到的結果是經過引導且可控的。在訓練過程中,LLM 策略得到的生成結果和 PRM 提供的評估相互增強,從而可讓這兩個組件不斷自我改進和優化。

有了推理數據之後,下一步就是訓練世界模型了,也被稱爲過程獎勵模型(PRM)。也就是說,由於狀態轉變是確定和已知的,因此重點就變成了學習一個之後可用於引導搜索、推理和解碼過程的通用獎勵模型。

該獎勵模型通常被稱爲驗證器,記爲 vPRM (s),可以使用有標註的推理步驟數據集進行訓練。其訓練通常涉及根據推理步驟的正確性優化一個分類損失函數:

還有另一種方法,是將 PRM 視爲一個可以通過價值迭代方法訓練的價值函數,使其能夠預測累積獎勵並通過最佳動作選擇指導推理過程。

假設有一個推理過程,其中狀態 s 表示當前狀態,並且其整合了之前的所有狀態。該價值迭代方法的目標是學習一個由 θ 參數化的價值函數 V_θ (s),其可預測從狀態 s 開始的預期累積獎勵。該價值函數可通過評估不同動作的潛在結果來指導推理過程。r_φ (s) 是獎勵函數,其會根據中間推理步驟或最終答案的正確性爲狀態 s 分配一個標量獎勵。γ 是折扣因子,決定了未來獎勵的相對重要性。該 PRM 的貝爾曼方程爲:

爲了學習該價值函數的 θ,這裡將 TD 損失函數定義成當前值與貝爾曼目標之間的平方誤差:

得到了 PRM 之後,就可以訓練 LLM 策略以提升其推理能力了。這需要超越傳統監督學習框架的方法。通過整合在線強化學習來優化推理任務,PRM 在此過程中發揮着重要作用。

這裡來看看 Group Relative Policy Optimisation(GRPO),即分組相對策略優化。

假設對於每個問題 Q = q,策略都會生成推理步驟 {o_1, o_2, . . . , o_G},每個輸出 o_i 由多個步驟 {a_{i,1}, a_{i,2}, . . . , a_{i,Ki} } 組成,其中 K_i 是輸出 o_i 中的推理步驟(或 token)總數。現在,可以構建通過 PRM 學習 LLM 策略的 GRPO 了,如下所示。

對於每個問題 q,GRPO 從舊策略 π_θ_old 採樣一組輸出 {o_1, o_2, . . . , o_G},目標是通過最大化以下目標來優化策略:

GRPO 不會將 KL 懲罰直接納入獎勵,其規範策略的方式是將當前策略 π_θ 和參考策略 π_θ_ref 之間的 KL 散度直接添加到損失函數中。這可確保更新後的策略在訓練期間不會過度偏離參考策略,從而有助於保持穩定性。

這種 GRPO 形式是通過利用推理步驟和最終步驟中的分組相對獎勵來優化 LLM 策略,專門適用於通過過程獎勵模型的推理任務。歸一化的優勢函數(advantage function)是根據相對性能計算的,鼓勵策略偏向在一組採樣輸出中表現更好的輸出。此外,KL 正則化可確保更新後的策略與參考策略保持接近,從而提高訓練穩定性和效率。該框架提供了一種穩健的方法,可通過基於 PRM 的優化來指導 LLM 推理。

另外,還有 token 級 DPO 等不使用 PRM 的更高效的離線方法,詳見相關論文《Token-level direct preference optimization》。

推理時間計算

訓練完成後,LLM 策略必須在推理過程中高效地生成輸出。LLM 常用的方法是自迴歸,即根據之前的 token 逐一生成新 token。但是,對於推理任務,還必需更復雜的解碼技術。

爲了在效率和效果之間取得平衡,有研究發現,波束搜索等更靈活的方法有利於推理任務。對於更復雜的推理任務,可以使用 MCTS 等向前看的模型。

MCTS 可模擬多條推理路徑,並根據獎勵系統對其進行評估,選擇預期獎勵最高的路徑。這允許模型在推理過程中探索更大範圍的可能性,從而增加其獲得最優解的機會。使用 MDP,可以從形式上定義其推理過程結構。

定義 2:原生思維鏈(NCoT)是指大型語言模型(LLM)固有的推理能力,這讓其無需外部提示詞便可以自動執行逐步式的結構化推理。該能力可以表述爲一個馬爾可夫決策過程(MDP)(S, A, π, R),其中

該模型既可以通過展開 MDP 來遵循順序推理路徑,也可以通過在每個狀態下采樣不同的推理步驟來探索多個軌跡,形成樹狀結構(圖 5)。過程獎勵模型 R 提供了對該空間的引導搜索,其控制推理軌跡的方式是支持能得到更有意義或更正確的推理步驟的動作。

圖 5:藉助 PRM,LLM 可以通過三種方法執行非自迴歸推理:1) 採樣多條推理軌跡,2) 對潛在推理路徑的樹結構進行蒙特卡洛搜索,3) 結合前兩種方法來增強推理的靈活性和穩健性。

相關文獻

下面將回顧並討論相關領域的幾篇關鍵論文,介紹它們的貢獻和侷限性。圖 6 描述了這些研究與更廣泛的研究圖景之間的聯繫。

圖 6:LLM 原生思維鏈研究。

推理時間計算

有幾篇論文重點關注優化 LLM 在推理時間的推理,比如:

驗證器模型

驗證器模型(結果獎勵模型和過程獎勵模型)已成爲提高 LLM 推理可靠性的重要研究領域。

獲取用於推理任務的數據

獲取推理數據一直都是一大研究重心,研究包括:

理解和系統級提升