CMU清華教LLM練成數學高手，LeanSTaR訓練模型邊思考邊證明，登頂新SOTA
如果想訓練LLM證明定理的能力，你會怎么做？
既然模型可以通過海量語料學會生成文本，那如果我們能喂給它足夠數量的形式證明數據，定理證明能力自然水到渠成？
然而，我們看到的事實是，無論用符號形式還是自然語言，GPT等大模型的推理能力都不如人意。
兩句話，讓LLM邏輯推理瞬間崩潰！最新「愛麗絲夢游仙境」曝出GPT、Claude等重大缺陷
就像GPT-4o自信表示13.11比13.8大一樣，AI再聰明卻依舊會在簡單的算術上犯蠢。
然而，LLM的數學能力弱，不代表自動化的定理證明器對數學沒用。
前段時間剛剛被破解的「忙碌海貍」問題中，4萬行Coq代碼功不可沒。
陶哲軒也曾在采訪中強調，使用Lean等自動化工具可以徹底顛覆數學家們的工作方式。這是一股不可小覷的力量。
最近，CMU和清華的一項研究就致力于讓LLM的「自然語言思維鏈」和Lean的形式化證明結合在一起。
論文地址：https://arxiv.org/abs/2407.10040
論文提出，Lean、Coq、Isabelle等基于形式語言（代碼）的自動化證明方法，忽略了大量可能對推理過程有用的「非形式化信息」。
比如，每個證明步驟之前的潛在思維過程是必不可少的，但卻不會形式化地體現在最終的公式和代碼中。
比如，圖1中右側的推理思路，在左側的證明步驟中完全「無處安放」。
因此，作者提出了Lean-STaR訓練框架，讓語言模型既學會逐步推理的思維，也學會形式化的證明方式。
這意味著，需要將自然語言和形式語言交織在一起，也將「思考」和「證明」的過程交織在一起。
方法：Lean-STaR
顧名思義，Lean-STaR這個方法同時結合了之前的兩項成果——Lean和STaR。
Lean是一種函數式編程語言，可以用作交互式定理證明器（Interactive Theorem Prover）。
項目主頁：https://lean-lang.org/
這是由Leonardo de Moura在微軟研究院期間發起的開源項目，目前已經更新到Lean 4。
比如，要想形式化證明，能從n≤m推斷出n+k≤m+k，就可以用Lean寫為如下形式（圖6）：
首先給出一種高級的「策略」（tactic，圖中所示為歸納策略k），將當前要證明的目標狀態簡化為多個子目標（下圖中的case 0和case ih）。
這些子目標又會形成新的「狀態」（state）。當所有子目標都得到證明時，我們就給出了定理的完整證明。
STaR則是來源于斯坦福和谷歌研究院在2022年發表的一篇論文，全稱是「自學推理器」（Self-Taught Reasoner）。
論文地址：https://arxiv.org/abs/2203.14465
其基本思想就是用到了「自舉法」（bootstrapping）。
首先根據訓練數據中的問題和答案，提示語言模型，生成能解釋答案的「原理」（rationale）。
之后，再用這個包含了問題、答案和原理的混合數據集對LM進行微調，提升模型的推理能力（圖1）。
Lean-STaR模型的微調也是采用了「漸進優化」的思路，逐步將以上兩個相關工作的成果融合在一起，完善底層的策略預測模型。模型構建的流水線如圖4所示。
直接策略預測（Direct Tactic Prediction）
首先，將定理證明問題簡單地定義為馬爾科夫決策過程（MDP）
。
從這個角度來看，證明過程是狀態si、策略ai和獎勵ri∈R等3個變量組成的軌跡(s1，a1，r1) (s2，a2，r2)?其中，ITP（比如Lean）用于提供每個新狀態si+1。
在這種經典設置中，證明定理的過程包括向LM提供狀態s，讓模型M生成策略?????(??|??) 。
因此，可以使用僅包含成功證明軌跡的基本數據集
對基本模型進行監督微調，得到SFT模型。
思維增強策略預測（Thought-augmented Tactic Prediction）
結合之前所述的研究動機，我們假設「潛在想法」可以提高模型的策略預測能力，因此引入一個表示「思維」的隱變量ti，然后將模型擴展為：
此時，根據狀態預測下一個策略的分布可以表示為：
如果用這種方式預測，我們就需要一個全新的數據集
用于訓練模型M，然而Lean給出的證明步驟只包含si和ai。
論文的解決方法是：借助一個強大的語言模型G（如GPT-4）作為「預言家」，讓它在給定當前狀態si和真實策略ai的情況下生成ti，從而創建出新的數據集DT（即圖4中的CoT Dataset）。
作者將這種方法稱為「回顧性原理生成」（retrospective rationale generation）。
將SFT模型在DT數據集上再進行一次微調后，就得到了第一個思維增強的策略預測模型Lean-CoT。
自舉思維增強定理證明（Bootstrapping Thought-augmented Theorem Proving）
在Lean-CoT模型的基礎上，作者提出，可以應用「專家迭代」（expert iteration）方法進一步提升性能。
具體來說，從初始的Lean-CoT模型M0以及初始數據集D開始，讓M0對每個問題進行K次采樣，每次采樣都會產生一個證明軌跡 [(s0,t0,a0),(s1,t1,a1),?,(sn,tn,an)]，之后過濾出成功的證明軌跡并去重，得到新數據集D1。
接下來，在數據集DT∪D1上進一步微調M0模型以得到Lean-STaR模型M1。
將上述過程進行多次迭代，即可不斷更新Lean-STaR模型。
評估實驗
為了測試Lean-STaR的具體性能，研究使用了可用的最佳開放語言模型Lean語料庫 (InternLM2-Math-base-7b) 上進行預訓練，并遵循Lean的Mathlib作為底層訓練集的標準實踐。
首先以LeanDojo Benchmark 4 v9作為監督微調(SFT)數據集，包含超過23.1萬個示例，進行1輪微調以獲得SFT模型。
之后從數據集中隨機選擇17256個不同的成功證明軌跡，并使用GPT-4-0125模型注釋出52438個想法，并且執行兩次專家迭代。
實驗在MiniF2F基準上評估Lean-STaR，使用了與之前的實驗工作類似的評估設置，但主要使用的是采樣方法（sampling）而不是最佳優先搜索（best-first search）來進行評估。
實驗結果表明，回顧性原理生成和專家迭代都顯著提高了模型的定理證明能力。
實驗結果
實驗的主要結果如下表所示，Lean-STaR比之前基于Lean的SOTA模型有了顯著的改進。
例如，在類似的推理預算下，同樣使用best-first search，Lean-STaR從InternLM2的30.3%提升至34.8%，也同樣高于使用GPT-4的COPRA（30.7%）。
隨著計算預算的增加，Lean-STAR的性能進一步提升至36.1%。
思維增強改進定理證明
Lean-STaR的第一階段在思維增強的合成數據集上進行微調，訓練模型來交替生成思維和策略。
此階段的微調模型（在表1中表示為Lean-CoT）達到了32.8%的通過率，高于此階段之前的模型（表示為 SFT，29.5%）。
可以證明，第一階段的思維增強能提高語言模型的定理證明能力，即使對于已經專門用于生成Lean策略的語言模型（例如SFT）也依舊成立。
自舉法（Bootstrapping）進一步改進
Lean-STaR的第二階段包括使用當前語言模型生成新的思維和策略，保存正確結果，并結合初始數據集進行訓練。
從表1結果來看，每次迭代都會提高模型的定理證明性能，從32.8%（初始模型）到34%（迭代1次后的L-STR）再到34.8%（迭代2次后的L-STR）。
此外，我們發現該模型可以通過額外采樣進一步改進，將采樣的K值加倍后，分數能進一步提升至36.1%。
無CoT的專家迭代實驗
表5顯示了沒有CoT的專家迭代結果（即僅使用狀態和策略，沒有思維增強），對比Lean-CoT和Lean-STaR的表現。
僅用專家迭代時，準確率就達到了43.0%，低于Lean-STaR (45.5%)。
這表明Lean-STaR的性能提升不僅僅來自于專家迭代的使用，思維增強也有不可忽略的效果。
問題類型與難度
MiniF2F-test中的問題有多個來源，包括AIME、AMC、IMO等數學競賽以及MATH數據集，并進行了手動形式化處理。
這些問題可能有不同的難度和類型。表2展示了成功證明的問題數量，按類型和難度劃分。
Lean-CoT提高了解決所有類別難題的表現，尤其是數學競賽中的難題。
除了這些改進之外，Lean-STAR的改進主要集中在數論方面。
搜索和抽樣預算
表4說明了問題通過率與搜索規模或抽樣預算S×K的關系。
實驗發現，Lean-STAR性能與K值的大小成正比，特別是當K值相對較大時。
對比前兩列和Lean-STaR可以發現，附帶思維的額外采樣能提高定理證明性能，而沒有思維的額外采樣可能會飽和。
作者猜測，可能是因為「思維」增加了輸出的多樣性，并有助于對定理證明空間進行探索。
因此，Lean-STaR更具可擴展性（就推理階段算力而言），并且可以通過額外的專家迭代進一步改進。
更強基礎模型和更多數據實驗
實驗還使用了更強的語言模型InternLM2-Math-plus-7b訓練LeanSTaR，來測試不同語言模型性能的影響。
不僅基座模型更強，為數據集注釋「思維」的模型也從GPT-4升級到GPT-4o，生成了1.4萬條想法。
實驗只執行一次專家迭代，收集了大約6萬條（證明狀態、思維、下一步策略）正確的數據，命名為「STaR 數據集」。
在STaR數據集上進一步微調得到Lean-STAR模型，其測評結果如表3所示，可以看到Lean-STaR仍然比基線有了顯著的改進。
結論和局限性
研究團隊提出了Lean-STaR，這是一種新穎的方法，通過將思維鏈 (CoT) 原理集成到每個證明步驟中，顯著增強了語言模型用形式化數學語言進行定理證明的能力。
方法首先根據ground truth回顧性地為證明步驟生成「原理」，然后微調語言模型，訓練模型學會生成「原理」并預測后續策略，從而得到Lean-CoT模型。
然后使用專家迭代進一步改進該模型，根據被證明為正確的采樣結果進行微調，并使用Lean solver進行驗證。
研究的貢獻包括引入第一個思維增強的定理證明數據集，并證明專家迭代可以進一步提高性能。得到的模型在miniF2F測試上取得最新SOTA，將通過率從30.3%提高到36.1%。
這些進步不僅提高了自動化定理證明的準確性，而且還提供了一個可擴展且高效的框架來促進對數學的理解，這可能會對教育、科學發現和程序驗證產生重大影響。
方法的主要限制在于，其性能可能受限于計算可擴展性，實驗中用于微調Lean-CoT和Lean-STaR模型的數據集都不是很大。
需要注意的是，專家迭代的速度也存在嚴重瓶頸，會受限于Lean ITP的緩慢進程。
此外，使用GPT-4生成合成數據成本較大，并可能引入偏差。
