AIxiv專欄是機器之心釋出學術、技術内容的欄目。過去數年，機器之心AIxiv專欄接收報道了2000多篇内容，覆寫全球各大高校與企業的頂級實驗室，有效促進了學術交流與傳播。如果您有優秀的工作想要分享，歡迎投稿或者聯系報道。投稿郵箱：[email protected]；[email protected]
本論文作者來自華為蒙特利爾諾亞方舟實驗室的康計堃，李信擇，陳熙， Amirreza Kazemi，陳博興。
人工智能（AI）在過去十年裡取得了長足進步，特别是在自然語言處理和計算機視覺領域。然而，如何提升 AI 的認知能力和推理能力，仍然是一個巨大的挑戰。
近期，一篇題為《MindStar: Enhancing Math Reasoning in Pre-trained LLMs at Inference Time》的論文提出了基于樹搜尋的推理時間能力提升方法 MindStar [1]，該方法在開源模型 Llama-13-B 與 Mistral-7B 上達到了近似閉源大模型 GPT-3.5 與 Grok-1 在數學問題上的推理能力。
論文标題：MindStar: Enhancing Math Reasoning in Pre-trained LLMs at Inference Time
論文位址：https://arxiv.org/abs/2405.16265v2
MindStar 在數學問題上的應用效果：
圖 1 ：不同大型語言模型的數學準确率。LLaMA-2-13B 在數學性能上與 GPT-3.5 (4-shot) 類似，但節省了大約 200 倍的計算資源。
1. 引言
随着模型規模的快速增長，基于 Transformer 的大型語言模型（LLMs）在指令遵循 [1,2]、編碼輔助 [3,4] 和創意寫作 [5] 等領域展示了令人印象深刻的成果。然而，解鎖 LLMs 解決複雜推理任務的能力仍然是一大挑戰。最近的一些研究 [6,7] 嘗試通過監督微調（Supervised Fine-Tuning, SFT）來解決，通過将新的推理資料樣本與原始資料集混合，使 LLMs 學習這些樣本的底層分布，并嘗試模仿所學邏輯來解決未見過的推理任務。盡管這種方法有性能提升，但它嚴重依賴于大量的訓練和額外的資料準備 [8,9]。
Llama-3 報告 [10] 強調了一個重要的觀察：當面對一個具有挑戰性的推理問題時，模型有時會生成正确的推理軌迹。這表明模型知道如何産生正确答案，但在選擇上存在困難。基于這一發現，我們提出了一個簡單的問題：我們能否通過幫助 LLMs 選擇正确的輸出來增強它們的推理能力？為探索這一點，我們進行了一項實驗，利用不同的獎勵模型進行 LLMs 輸出選擇。實驗結果表明，步驟級選擇顯著優于傳統的 CoT 方法。
2. MindStar 方法
圖 2 MindStar 的算法架構圖
我們引入了一種新的推理搜尋架構 ——MindStar（M*），通過将推理任務視為搜尋問題，并利用過程監督的獎勵模型（Process-supervised Reward Model, PRM），M * 在推理樹空間中有效導航，識别近似最優路徑。結合束搜尋（Beam Search, BS）和 Levin 樹搜尋（Levin Tree Search, LevinTS）的思想，進一步增強了搜尋效率，并保證在有限計算複雜度内找到最佳推理路徑。
2.1 過程監督獎勵模型
2.2 推理路徑擴充
2.3 推理路徑選擇
在擴充推理樹後，我們使用預訓練的過程監督獎勵模型（PRM）來評估每個新生成的步驟。正如前面提到的，PRM 采用路徑和步驟 ，并傳回相應的獎勵值。在評估之後，我們需要一種樹搜尋算法來選擇下一個要擴充的節點。我們的架構不依賴于特定的搜尋算法，在這項工作中，我們執行個體化了兩種最佳優先搜尋方法，即 Beam Search 和 Levin Tree Search。
3. 結果與讨論
在 GSM8K 和 MATH 資料集上的廣泛評估顯示，M * 顯著提升了開源模型（如 LLaMA-2）的推理能力，其表現可與更大規模的閉源模型（如 GPT-3.5 和 Grok-1）媲美，同時大幅減少了模型規模和計算成本。這些發現突顯了将計算資源從微調轉移到推理時間搜尋的潛力，為未來高效推理增強技術的研究開辟了新途徑。
表 1 展示了各種方案在 GSM8K 和 MATH 推理基準上的對比結果。每個條目的數字表示問題解決的百分比。符号 SC@32 表示在 32 個候選結果中的自一緻性，而 n-shot 表示少樣本例子的結果。CoT-SC@16 指的是在 16 個思維鍊（CoT）候選結果中的自一緻性。BS@16 代表束搜尋方法，即在每個步驟級别涉及 16 個候選結果，而 LevinTS@16 詳細說明了使用相同數量候選結果的 Levin 樹搜尋方法。值得注意的是，MATH 資料集上 GPT-4 的最新結果為 GPT-4-turbo-0409，我們特别強調這一點，因為它代表了 GPT-4 家族中的最佳性能。
圖 3 我們研究了 M * 性能如何随着步驟級别候選數量的變化而變化。我們選擇 Llama-2-13B 作為基礎模型，并分别選擇束搜尋（BS）作為搜尋算法。
圖 4 Llama-2 和 Llama-3 模型家族在 MATH 資料集上的尺度定律。所有結果均來自它們的原始資源。我們使用 Scipy 工具和對數函數來計算拟合曲線。
表 2 不同方法在回答問題時的平均 token 生産數量
4. 結論
本文介紹了 MindStar（M*），一種新穎的基于搜尋的推理架構，用于增強預訓練大型語言模型的推理能力。通過将推理任務視為搜尋問題并利用過程監督的獎勵模型，M* 在推理樹空間中有效導航，識别近似最優路徑。結合束搜尋和 Levin 樹搜尋的思想，進一步增強了搜尋效率，并保證在有限計算複雜度内找到最佳推理路徑。廣泛的實驗結果表明，M* 顯著提升了開源模型的推理能力，其表現可與更大規模的閉源模型媲美，同時大幅減少了模型規模和計算成本。
這些研究成果表明，将計算資源從微調轉移到推理時間搜尋具有巨大的潛力，為未來高效推理增強技術的研究開辟了新途徑。
參考文獻：
[1] Nisan Stiennon, Long Ouyang, Jeffrey Wu, Daniel Ziegler, Ryan Lowe, Chelsea Voss, Alec Radford, Dario Amodei, and Paul F Christiano. Learning to summarize with human feedback. Advances in Neural Information Processing Systems, 33:3008–3021, 2020.
[2] Long Ouyang, Jeffrey Wu, Xu Jiang, Diogo Almeida, Carroll Wainwright, Pamela Mishkin, Chong Zhang, Sandhini Agarwal, Katarina Slama, Alex Ray, et al. Training language models to follow instructions with human feedback. Advances in neural information processing systems, 35:27730–27744, 2022.
[3] Ziyang Luo, Can Xu, Pu Zhao, Qingfeng Sun, Xiubo Geng, Wenxiang Hu, Chongyang Tao, Jing Ma, Qingwei Lin, and Daxin Jiang. Wizardcoder: Empowering code large language models with evol-instruct. arXiv preprint arXiv:2306.08568, 2023.
[4] Mark Chen, Jerry Tworek, Heewoo Jun, Qiming Yuan, Henrique Ponde de Oliveira Pinto, Jared Kaplan, Harri Edwards, Yuri Burda, Nicholas Joseph, Greg Brockman, et al. Evaluating large language models trained on code. arXiv preprint arXiv:2107.03374, 2021.
[5] Carlos Gómez-Rodríguez and Paul Williams. A confederacy of models: A comprehensive evaluation of llms on creative writing. arXiv preprint arXiv:2310.08433, 2023.
[6] Longhui Yu, Weisen Jiang, Han Shi, Jincheng Yu, Zhengying Liu, Yu Zhang, James T Kwok, Zhenguo Li, Adrian Weller, and Weiyang Liu. Metamath: Bootstrap your own mathematical questions for large language models. arXiv preprint arXiv:2309.12284, 2023.
[7] Zhihong Shao, Peiyi Wang, Qihao Zhu, Runxin Xu, Junxiao Song, Mingchuan Zhang, YK Li, Y Wu, and Daya Guo. Deepseekmath: Pushing the limits of mathematical reasoning in open language models. arXiv preprint arXiv:2402.03300, 2024.
[8] Keiran Paster, Marco Dos Santos, Zhangir Azerbayev, and Jimmy Ba. Openwebmath: An open dataset of high-quality mathematical web text. arXiv preprint arXiv:2310.06786, 2023.
[9] Peiyi Wang, Lei Li, Zhihong Shao, RX Xu, Damai Dai, Yifei Li, Deli Chen, Y Wu, and Zhifang Sui. Math-shepherd: Verify and reinforce llms step-by-step without human annotations. CoRR, abs/2312.08935, 2023.
[10] Meta AI. Introducing meta llama 3: The most capable openly available llm to date, April 2024. URL https://ai.meta.com/blog/meta-llama-3/. Accessed: 2024-04-30.