NeurIPS 2024 | REBASE,比MCTS更高效的Inference Scaling算法
?還在為大模型推理速度慢、成本高而煩惱嗎?這篇來自 CMU 和清華大學的論文給你帶來了一個顛覆性的解決方案! 他們發現,在推理大型語言模型時,“大力出奇跡”不再是真理! 通過深入研究推理過程中的計算與性能關系,他們提出了“推理縮放定律”,并革命性地推出了一種名為 REBASE 的全新算法。
REBASE 就像一位聰明的向導,能夠巧妙地利用獎勵信號,指引模型在推理的迷宮中高效探索,避免了傳統方法中耗時費力的盲目搜索。 實驗結果令人驚嘆:在數學推理任務上,REBASE 搭配小模型,不僅速度更快,而且精度更高,輕松擊敗了更大的模型!
這篇論文的意義不僅僅在于提出了一種新算法,更在于它顛覆了我們對大模型推理的傳統認知,為我們指明了一條通往更高效、更經濟的 AI 之路。 想要了解如何讓你的大模型推理效率飆升?快來閱讀這篇論文,一起探索 LLM 推理的奧秘吧!
一、概述
?Title:INFERENCE SCALING LAWS: AN EMPIRICAL ANALYSIS OF COMPUTE-OPTIMAL INFERENCE FOR LLM PROBLEM-SOLVING
?URL:?? https://arxiv.org/abs/2408.00724??
?Authors:Yangzhen Wu, Zhiqing Sun, Shanda Li, Sean Welleck, Yiming Yang
?Code:?? https://thu-wyz.github.io/inference-scaling/??
1.Motivation
? 目前關于大型語言模型 (LLM) 的研究主要集中在訓練階段的縮放定律上,而對于推理階段的最佳配置研究較少。
? 現有的推理技術通常需要在推理時增加額外的計算,以最大化性能,但這些技術的計算成本必須考慮在內,以實現計算最優推理。
? 傳統的采樣算法會有性能上限(論文Figure 1),需要更好的推理算法來取得更好的結果。
2.Methods
文章評估了不同模型大小和推理算法在數學推理基準測試上的性能,設計了一個新的樹搜索算法REBASE,該算法在計算效率上優于常用的采樣和MCTS方法。
2.1 省流版總結:
核心思想: 研究在給定計算預算下,如何選擇最優的語言模型大小和有效的推理策略,以最大化問題解決的準確性。
創新點: 提出了一個新的樹搜索算法REBASE(Reward Balanced Search),該算法與加權投票結合使用,能夠在準確性和推理計算量之間實現帕累托最優的權衡。
帕累托最優的權衡意味著:
? REBASE 算法在給定的計算預算下,能夠獲得最佳的性能(例如,最高的數學問題解決準確率)。
? 如果你想進一步提高準確率,你就必須增加計算預算(例如,生成更多的 tokens,進行更多的搜索)。
? 如果你想減少計算預算,你就必須接受較低的準確率。
2.2 REBASE詳細方法和步驟:
目標: 在給定固定計算預算的情況下,如何選擇最優的模型大小和推理策略,以最大化模型在特定任務上的性能。
2.2.1 REBASE 算法是干什么的的?
為了解決 MCTS 的計算成本高問題,論文提出了新的樹搜索算法 REBASE。REBASE 通過獎勵模型來控制節點的擴展,從而避免了 MCTS 中顯式的 rollout 步驟。
REBASE 算法避免 MCTS 中顯式 rollout 步驟的關鍵在于它巧妙地利用了獎勵模型 (Reward Model) 和 平衡擴展 (Balanced Expansion) 策略。
2.2.2 MCTS 的 Rollout 步驟有什么問題?
在傳統的 MCTS 中,為了評估一個節點的價值,需要進行 rollout(模擬)。Rollout 是指從當前節點開始,根據某種策略(例如隨機策略或基于模型的策略)模擬生成多個完整的解決方案,然后根據這些解決方案的結果來估計該節點的價值。
Rollout 的問題在于:
1)計算成本高:Rollout 需要生成多個完整的解決方案,這在計算上非常昂貴,特別是當解決)方案很長時。
2)估計偏差:Rollout 使用的策略可能與實際的策略不同,這會導致價值估計存在偏差。
2.2.3 REBASE 如何利用獎勵模型控制節點擴展?
REBASE 摒棄了顯式的 rollout 步驟,而是直接使用過程獎勵模型 (PRM) 來評估節點的質量。具體來說,REBASE 在每次迭代中執行以下步驟:
1)獎勵分配 (Reward Assignment):
- 對于當前深度?
?i?
?? 的每個節點??nj?
??,REBASE 使用 PRM 來計算該節點的獎勵??R(nj)?
??。這個獎勵反映了從根節點到節點??nj?
? 的部分解決方案的質量。注意,這里直接利用了 PRM 的輸出,而不需要進行 rollout。
2)平衡擴展 (Balanced Expansion):
Wj = Round ( Bi * exp(R(nj) / Tb) / Σk exp(R(nk) / Tb) )
其中:
直觀理解: 這個公式的含義是,每個節點的擴展寬度與該節點的獎勵的指數成正比。獎勵越高的節點,其指數值越大,因此獲得的擴展寬度也越大。通過這種方式,REBASE 將更多的計算資源分配給更有希望的節點。
???Bi?
?? 是深度??i?
? 的可用采樣預算(即還可以生成的總節點數)。
???R(nj)?
?? 是節點??nj?
? 的獎勵。
???Tb?
?? 是一個平衡溫度參數,用于控制擴展的“平滑”程度。較高的??Tb?
?? 會導致更均勻的擴展,而較低的??Tb?
? 會導致更集中于高獎勵節點的擴展。
???Round()?
? 函數將結果四舍五入到最接近的整數。
? 分母??Σk exp(R(nk) / Tb)?
?? 對所有節點??nk?
?? 的獎勵進行 softmax 歸一化,確保所有節點的擴展寬度之和等于??Bi?
?。
? REBASE 的目標是根據節點的獎勵來“平衡”地擴展節點。這意味著獎勵越高的節點,將被分配越多的計算資源(即擴展出更多的子節點)。
? 具體來說,REBASE 計算每個節點??nj?
?? 的擴展寬度??Wj?
?,公式如下:
3)節點擴展 (Node Expansion):
? 根據計算出的擴展寬度??Wj?
??,REBASE 從策略模型??πθ?
?? 中采樣??Wj?
?? 個子節點,作為節點??nj?
? 的子節點。
2.2.4 總結一下REBASE方法的改進點
REBASE 通過以下方式避免了 MCTS 的 rollout 步驟:
?直接使用獎勵模型評估節點質量:無需進行耗時的 rollout 模擬。
?基于獎勵模型進行平衡擴展:根據節點的獎勵動態分配計算資源,將更多資源分配給更有希望的節點,從而更有效地利用計算預算。
通過這種方式,REBASE 在保持樹搜索的“探索”和“利用”特性的同時,顯著降低了計算成本,使其在準確性和計算效率之間取得了更好的平衡,從而實現了帕累托最優的權衡。
3 Conclusion
?較小的模型(例如 Llemma-7B)在相同的計算預算下可以勝過較大的模型,并且較小的模型與先進的推理算法相結合可以產生帕累托最優的成本性能權衡。
? 論文提出的REBASE 算法在所有測試的計算預算中都實現了最佳性能,優于常用的加權多數投票(Majority Voting)和 MCTS 方法。
? 通過采樣更多的樣本,基于采樣的多數投票策略不可避免地會飽和到一個取決于底層生成策略的分布。因此,通過設計替代推理策略來改變采樣分布是有意義的。
4 Limitation
? 該分析主要針對數學問題解決,推理縮放規律和計算優化推理在其他任務中是否適用還有待進一步研究。
? 在GSM8K和MATH500數據集上的評估可能不足以泛化到其他任務和數據集。
? 模型大小和推理策略的探索仍然存在局限性。
二、詳細情況
1 實驗設置
其他推理策略對比: 論文主要考慮了以下幾種推理策略:
? 貪婪搜索(Greedy Search):每次選擇概率最高的token。
? Best-of-n:生成多個候選答案,選擇得分最高的答案。
? 多數投票(Majority Voting):生成多個候選答案,選擇出現頻率最高的答案。
? 加權多數投票(Weighted Majority Voting):根據獎勵模型給出的分數對候選答案進行加權,然后進行投票。
? 樹搜索算法(Tree Search):包括 MCTS 和作者提出的 REBASE。
模型和數據集:
? 使用了 Pythia 和 Llemma 系列的多種大小的模型,并在數學推理基準測試集 GSM8K 和 MATH500 上進行了實驗。
? 使用 MetaMath 數據集對模型進行微調。
理論分析:
? 對基于采樣的投票方法進行了理論分析,表明隨著采樣次數的增加,準確率會收斂到一個極限,這個極限取決于語言模型和獎勵模型的分布。
- 實驗評估:
? 通過實驗,論文比較了不同模型大小和推理策略在不同計算預算下的性能。結果表明,較小的模型與高級的推理算法相結合可以實現帕累托最優的成本效益。
2 最優錯誤率和模型大小以及預算的關系
1.1 錯誤率和計算量以及模型大小的關系(左)
?隨著計算量的增加,錯誤率下降:所有的曲線都顯示,隨著推理計算量的增加 (即橫軸向右移動),錯誤率逐漸下降,說明模型性能在提升。這是因為增加計算量通常意味著模型可以進行更多的采樣,從而找到更好的解。
?模型性能最終趨于飽和:每條曲線最終都趨于平緩,說明模型性能有一個極限。增加更多的計算量后,錯誤率的下降幅度會變得越來越小,最終趨于飽和。
?不同大小的模型在不同階段表現不同:在計算量較少的時候,較小的模型 (例如 410M, 1.4B) 下降速度可能更快。隨著計算量增加,較大模型 (例如 6.9B, 12B) 會逐漸展現出更好的性能,但其飽和也需要更大的計算量。
1.2 最佳模型大小和預算的關系(右)
?最佳模型大小隨計算預算變化:圖中顯示,在不同的推理計算量預算下,最佳的模型大小是不同的。
例如,在較小的計算量預算下 (241 和 244
隨著計算預算的增加 (247
?小模型在較低計算量下表現出色:如圖中所示,在 241 和 244 FLOPs 的計算量下,較小的模型是計算最優的。這意味著,如果計算資源有限,與其使用大模型,不如使用小模型搭配適當的推理策略可能效果更好。
3 REBASE 算法在計算效率上有明顯優勢(7B模型+REBASE算法可以超過MCTS等方法)
方法:在 MATH 數據集上,對比不同推理策略 (weighted majority voting 和 best-of-n) 和不同模型大小對模型性能的影響 。
結論:
?在 MATH 數據集上,REBASE 算法在不同的推理策略(加權多數投票和最佳 n 選 1)下,均優于其他方法(包括傳統的采樣和 MCTS)。
?REBASE 算法的優越性不僅僅局限于特定的模型大小,其在使用 7B 模型時通常能夠達到最佳的性能。
?MCTS 方法在數學推理任務中,性能不如其他方法。
4 REBASE 算法可以在更低的計算預算下實現更高的準確率
?列:包括 # SAMPLES (采樣數量), FLOPS (推理計算量) 和 MATH500 (在 MATH500 數據集上的準確率) 三列。
?行:比較了使用采樣方法和 REBASE 算法在不同模型 (Mistral-7B, Llemma-7B, Llemma-34B) 下的性能。
?總結:
REBASE 算法可以在更低的計算預算下實現更高的準確率:對于所有模型,REBASE 算法在使用更少的樣本和更低的計算量 (FLOPS) 的情況下,都比采樣方法實現了更高的 MATH500 準確率。
驗證了 REBASE 的高效性:這個表格用具體數值驗證了 REBASE 的高效性,即能夠在降低計算成本的同時,實現更高的性能。
5 消融實驗
實驗說明:
?Sampling W.M.:使用采樣方法和加權多數投票策略。
?Sampling BoN:使用采樣方法和 best-of-n 策略。
?REBASE W.M.:使用 REBASE 算法和加權多數投票策略。
?REBASE BoN:使用 REBASE 算法和 best-of-n 策略。
實驗結論:
?REBASE 算法性能普遍優于采樣:在所有模型 (Llemma-7B, Llemma-34B, Mistral-7B) 和所有計算預算下,使用 REBASE 算法 (無論是加權多數投票還是 best-of-n) 的曲線都低于使用采樣方法的曲線,說明 REBASE 算法性能更佳。
?REBASE + 加權多數投票 策略表現最佳:REBASE 算法配合加權多數投票策略通常可以達到最優的性能。
?不同模型之間在相同策略下,性能隨計算預算增加有差異:例如,在低預算下,Llemma-7B 在 REBASE 策略下表現可能更好;在高預算下,不同模型之間的差距可能縮小。
三、總結
結論1: 計算最優的推理策略依賴于模型大小和計算預算。 論文表明,在推理時使用更小的模型并進行多次采樣,在計算成本方面可能比使用更大的模型更有效。此外,論文提出的REBASE算法在多個數據集上優于MCTS方法和傳統的采樣方法。
結論2: 基于采樣的投票策略具有性能上限,需要更復雜的推理算法。 論文的理論分析表明,基于采樣的投票方法(如多數投票和加權多數投票)在計算預算增加時,性能會飽和,且達到一個基于底層生成策略的極限,這表明了需要更復雜的推理策略。
結論3: REBASE算法在計算效率和準確率方面都有優勢。 REBASE算法通過獎勵模型來控制節點的擴展,從而避免了顯式的rollout,從而在計算效率上優于MCTS。同時,其性能在所有測試的計算預算下優于采樣方法。
本文轉載自??NLP PaperWeekly??,作者: NLP PaperWeekly ????
