AE Studio最近在Modal平臺開展了一項引人注目的研究,旨在透過強化學習訓練語言模型進行數學定理證明。他們系統比較了兩種方法:當前主流的組相對策略最佳化(GRPO)和一種受自然選擇啟發的進化策略(ES)。研究使用了Lean形式化驗證語言作為獎勵訊號——模型生成的Lean程式碼由編譯器驗證正確性,從而提供明確的獎勵訊號。
為了高效執行實驗,AE Studio選擇了Modal平臺,該平臺提供了並行GPU計算、沙盒隔離和持久化卷儲存等關鍵功能。具體工作流程分為三個部分:證明生成使用vLLM在GPU上執行,證明驗證使用Lean編譯器在CPU上透過Modal沙盒隔離執行,協調程序管理整個訓練迴圈。ES方法透過建立一群略有不同的模型副本,評估它們並更新原始模型,而GRPO則基於組內相對效能進行梯度更新。
在實現細節上,AE Studio充分利用了Modal的特色功能。他們為每個計算角色使用獨立的映象,而不是將所有依賴打包到一個臃腫的環境中。GPU工作節點專注於推理,Lean沙盒專注於驗證,協調器保持輕量。並行GPU扇出透過Modal的.map()實現,每個擾動評估只需要當前檢查點、定理批次、一個擾動種子和生成引數。由於ES的權重擾動完全由種子決定,更新步驟只需種子和獎勵即可聚合所有個體,避免了GPU間昂貴的權重傳輸。驗證環節使用Modal沙盒進行隔離,每個驗證批次啟動一個獨立的Lean伺服器,處理完即關閉,防止錯誤的證明掛起或崩潰整個系統。這種設計使得每輪迭代可以建立3840個證明嘗試,並按批次並行驗證。
ES的一個優雅特性是,整個模型狀態可以由基模型加上一組(種子,獎勵)對來描述。協調器維護一個種子/獎勵歷史列表,並傳遞給每個工作節點。GPU工作節點從卷儲存載入原始基模型,然後重放完整歷史以重構當前權重,再應用自己的擾動。這個重放過程遍歷過去每次迭代的種子和獎勵,在GPU上使用確定性種子重新生成同樣的噪聲,並就地應用加權更新。整個模型狀態作為純Python列表傳遞,小到可以作為函式引數傳給每次遠端呼叫。
效能方面,Modal的實現僅需約250行平臺設定程式碼,不到其他平臺典型600行的一半。複雜度降低使得整個專案從啟動到成功訓練執行僅用了不到兩天時間,比替代平臺快60%。執行時效率也顯著提升:生成步驟(需GPU)平均每輪147秒,但完整迴圈因驗證步驟(僅CPU)而耗時538秒。Modal的彈性計費意味著只需為GPU實際使用時間付費,將浪費的GPU時間減少了約3.7倍。基於這些時間資料,估算整個執行成本在Modal上僅為122美元,而同等硬體上的替代方案成本在180至480美元之間。
早期結果令人鼓舞。在多次定理證明執行中,ES在每輪驗證的證明數量上匹配甚至超越了GRPO基線,並且當訓練資料有限時,ES通常展現出更高的樣本效率。但在其他設定中增益較小,且尚未分離超引數選擇、資料集差異或種群規模等因素帶來的變化。ES在這一領域的縮放行為仍未被充分理解,需要更多實驗來確定其一致優勢以及在大規模訓練中能否與GRPO競爭或超越。
研究問題仍然開放,但Modal上的基礎設施設定將繼續使用,以便執行、檢查和迭代這個基於驗證器的訓練迴圈,而無需投入時間構建定製基礎設施。接下來的實驗包括方差分析(對sigma、種群規模和定理選擇進行系統掃描)和更大模型測試(使用開源的7B蒸餾Kimina模型)。
值得注意的是,這一基礎設施設定並非定理證明所獨有。它在許多機器學習系統中普遍適用:在GPU上生成輸出,對外部驗證器或編譯器執行結果,將結果轉化為訓練訊號,並在大量獨立候選方案上重複。關鍵的技術要點是並行化GPU推斷、隔離驗證以及保持訓練狀態的可移植性——全部在一個工作流中完成。.map()分佈了證明生成,沙盒控制了驗證器故障,卷儲存使檢查點在迭代之間可用,而無需移回本地機器。
透過Modal,團隊能夠在三種截然不同的執行時中執行相同的稀疏獎勵強化學習工作流,而無需每次重新搭建。他們從本地機器啟動執行,將基模型保持在遠端,記錄到MLflow,並根據實驗進展更改種群大小、定理批次大小和驗證並行度。由於Modal平臺的特性和易用性,專案啟動數天內就能看到實驗結果。其他基礎設施平臺可能需要數週測試和迭代才能獲得任何訊號。一旦搭建完畢,團隊可以完全專注於實驗設計,無需擔心繫統可靠性。
完整的實驗程式碼已在GitHub上公開(github.com/agencyenterprise/modal-rl-theorem-case-study),供深入研究或適配到自己的工作流。