Adaptive Test-Time Compute Allocation for Reasoning LLMs via Constrained Policy Optimization

Zhiyuan Zhai, Bingcong Li, Bingnan Xiao, Ming Li, Xin Wang

View Original ↗
AI 導讀 technology AI 重要性 4/5

捨棄齊頭式算力分配,改採動態推論預算路由,模型在數學測試的準確率最高提升 12.8%。

  • 提出雙階段框架,將推論預算分配轉化為毫秒級的輕量分類任務。
  • 動態路由的運算開銷極低,佔不到單次語言模型推論成本的 1%。
  • 總算力不變下,動態分配讓模型的數學準確率最高躍升 12.8%。

大型語言模型在推論階段的算力擴展往往採用齊頭式分配,引發龐大資源浪費。復旦大學與蘇黎世聯邦理工學院的最新研究證明,只要將推理預算從簡單任務轉移至複雜問題,在總算力不變下,模型於 MATH 測試的相對準確率最高可躍升 12.8%。團隊提出的雙階段框架,成功將算力最佳化問題轉化為毫秒級的分類任務。

齊頭式分配與 Fixed-b 造成的推論算力浪費

提升大型語言模型(LLM)效能的傳統配方,長期專注於增加預訓練階段的參數與數據量。近期產業焦點逐漸轉移至測試階段的算力擴展(Test-time scaling),藉由在推論時投入額外運算資源來強化輸出品質。開發者廣泛採用 CoT(思維鏈)、自我一致性多數決與樹狀搜尋等技術,讓模型在單一查詢上反覆推演。這些強大的推論期技術支撐了當前最先進的推理系統,卻也帶來呈線性成長的運算成本。

面對 API 成本、GPU 負載極限與吞吐量要求,實務上的推論預算通常設有嚴格上限。多數現有系統預設採用齊頭式的 Fixed-b(固定預算)策略,不論輸入題目的難易度為何,均強制分配相同數量的生成樣本。觀察真實世界的運算負載可以發現,許多簡單問題只需一次前向傳播即可獲得正確解答,過多採樣純屬浪費;反之,極度複雜的難題可能因為算力配額不足,錯失了透過大量推演找出正解的機會。

針對此異質性極高的查詢需求,研究團隊將資源分配定義為一個受約束的最佳化問題。核心挑戰在於,如何在共享固定總預算的限制下,找出能最大化整體準確率的個別題目預算分配方式。由於預算限制將所有輸入題目全數耦合在一起,單純針對單一輸入進行最佳化並無法得出全局最佳解。

拉格朗日鬆弛法與雙階段推論預算最佳化

拆解全局耦合的難題,研究團隊提出了一套名為 Solve-then-Learn 的雙階段運算力分配框架。首個階段為「求解(Solve)」,引入拉格朗日乘數 $\lambda$ 將總預算約束項轉化為懲罰函數。這個數學轉換成功將原本複雜的全局問題,解構為針對個別輸入獨立運作的子問題。在給定的 $\lambda$ 值下,系統可針對每一道題目計算出完美對應成本與準確率的「神諭(Oracle)分配」。

賦予神諭分配明確的經濟學意義,拉格朗日乘數 $\lambda$ 實質上扮演了「算力單位價格」的角色。當 $\lambda=0$ 時代表算力免費,系統會為所有題目分配最高預算;隨著 $\lambda$ 增加,算力變得昂貴,系統便會優先從那些增加預算也難以提升準確率的簡單題目開始削減資源。研究團隊進一步透過數學證明了成本單調性,確保隨著 $\lambda$ 提升,整體成本必然呈現下降或持平的階梯式函數。

建立這項數學特性後,開發者得以透過簡單的二分搜尋法,快速鎖定符合目標平均預算的 $\lambda$ 預測值。由於這項搜尋僅需要在事前計算好的效能矩陣上操作,不涉及實質的大型模型推論,運算複雜度極低。這些在訓練集上求得的神諭標籤,確立了理想狀態下對抗資源限制的最佳預算分配基準。

提取廉價文本特徵與 XGBoost 輕量級推論路由

確立神諭標籤後,框架的第二階段進入「學習(Learn)」,目標是讓系統能在不依賴昂貴試誤推論的前提下,預測出最佳預算分配。由於可選的預算配置為離散集合(例如設定樣本數為 1、2、4、8、16),此階段可直接視為一個標準的分類任務。系統會在耗費鉅資呼叫 LLM 之前,先針對輸入文本提取極度輕量化的特徵向量。

選定特徵的唯一標準,在於其計算成本必須遠低於完整的生成推論,同時又能有效反映題目難度。團隊測試的 16 項文本特徵涵蓋了詞彙統計(如輸入長度、單詞數)、結構指標(如問號數量、數學符號出現頻率),以及單次初步推論產生的標準化熵值。這些低成本訊號被送入一個輕量級分類器中,由模型負責學習輸入特徵與神諭預算標籤之間的非線性映射關係。

相較於線性回歸或傳統神經網絡,採用 GBM(梯度提升機)架構的 XGBoost 展現出壓倒性的優勢。實驗證明,樹狀分類器能精確捕捉分配邊界的非線性特徵,其平均實測預算完美貼合神諭預算的分佈輪廓。在實際部署時,這套由 XGBoost 驅動的預先路由機制,幾乎不會對系統產生可觀的延遲,運算開銷甚至不到一次 LLM 呼叫成本的 1%。

三大模型於數學測試最高提升 12.8% 準確率

驗證框架實效,團隊選用了 DeepSeek-V3、GPT-4o-mini 與 Qwen2.5-7B 三款極具代表性的 LLM 進行嚴格測試。實驗場域包含極具挑戰性的 MATH 競賽級別數據集,以及作為簡單對照組的 GSM8K 國小數學應用題庫。針對每道題目,團隊完整蒐集了 48 次推論結果,並以無重疊視窗的方式準確估算不同預算設定下的自我一致性準確率。

量化數據揭示了自適應預算分配的強大威力。在 MATH 數據集且平均預算上限設定為 3.0 時,這套動態分配策略全面擊潰了齊頭式與隨機分配的基準表現。具體而言,DeepSeek-V3 取得了 5.8 個百分點的絕對提升、GPT-4o-mini 提升 5.2 個百分點,而 Qwen2.5-7B 更是創下了 6.4 個百分點的驚人增長。若換算為相對準確率,部分配置下的提升幅度高達 12.8%

即便在基礎準確率已高達 94.6% 的 GSM8K 測試中,此方法依然能硬生生擠出額外 1.9 個百分點的準確率紅利。更值得注意的是,學習策略的實耗成本往往低於預算上限,例如在預算限制為 3.0 的情況下,DeepSeek-V3 處理 MATH 測試的實際平均算力僅 2.73。這證實了系統精確識別出只需少量樣本即可作答的題目,並將資源轉交給具備高邊際效益的中等難度題目。

帕雷托前緣分析與 91% 以上的神諭模仿表現

深入分析準確率與算力成本的帕雷托前緣(Pareto frontier),可發現自適應策略在任何預算設定下都穩居支配地位。當預算落在中等區間(B=2 至 6)時,系統能發揮最大的題目分流效益,使得此階段的相對增益最為顯著。相對地,當預算超出一定數值後,受限於題目本身的無解性或極度簡單,所有方法的表現皆會進入平緩的飽和期。

建構在嚴謹的數學推導之上,團隊進一步確立了模仿誤差與任務最終表現之間的理論邊界。理論保證指出,分類器每增加一個百分點的模仿準確率,就能按比例縮小與神諭最佳解之間的差距。一旦學習策略的平均成本趨近於神諭水準,任務準確率的落差將完全由分類準確率來決定,成功將約束推論問題降維至監督式分類範疇。

比對最終結果,輕量級的 GBM 分類器在所有模型與數據集組合中,均達成了 91% 到 99% 的極高神諭模仿準確率。即便存在微小的分類錯誤,系統通常也只是將題目錯派至鄰近的預算區間,因此與理想神諭分配之間的準確率差距被完美限縮在 0.4 至 1.4 個百分點內。這項實證不僅呼應了理論推導的穩健性,也為大規模部署推論應用提供了全新典範。

輕量級動態路由成功終結齊頭式的算力浪費,是未來語言模型低成本部署的關鍵基石。

Abstract

Test-time compute scaling, the practice of spending extra computation during inference via repeated sampling, search, or extended reasoning, has become a powerful lever for improving large language model performance. Yet deploying these techniques under finite inference budgets requires a decision that current systems largely ignore: which inputs deserve more compute, and which can be answered cheaply? We formalize this as a constrained optimization problem (maximize expected accuracy subject to an average compute budget) and solve it with a two-stage Solve-then-Learn pipeline. In the solve stage, Lagrangian relaxation decomposes the global constraint into per-instance sub-problems, each admitting a closed-form oracle action that optimally prices accuracy against cost. We prove that the induced cost is monotone in the dual variable, enabling exact budget targeting via binary search. In the learn stage, a lightweight classifier is trained to predict oracle actions from cheap input features, amortizing the allocation rule for real-time deployment. We establish that the task-level regret of the learned policy is bounded by its imitation error times the worst-case per-instance gap, yielding a clean reduction from constrained inference to supervised classification. Experiments on MATH and GSM8K with three LLMs (DeepSeek-V3, GPT-4o-mini, Qwen2.5-7B) show that our method consistently outperforms uniform and heuristic allocation baselines, achieving up to 12.8% relative accuracy improvement on MATH under matched budget constraints, while closely tracking the Lagrangian oracle upper bound with over 91% imitation accuracy.