熱門資訊> 正文
2025-10-28 09:56
剛剛,不發論文、愛發博客的 Thinking Machines Lab (以下簡稱 TML)再次更新,發佈了一篇題為《在策略蒸餾》的博客。
在策略蒸餾(on-policy distillation)是一種將強化學習 (RL) 的糾錯相關性與 SFT 的獎勵密度相結合的訓練方法。在將其用於數學推理和內部聊天助手時,TML 發現在策略蒸餾可以極低的成本超越其他方法。
該公司 CEO Mira Murati 表示,這種方法可用於小模型,使其具備強大的領域性能和持續學習能力。
值得注意的是,在這篇新博客中,TML 明確表示這項新成果受到了 Qwen 團隊研究的啓發,並且其實驗過程中也大量用到了 Qwen3 系列模型。事實上,在原英文博客中,「Qwen」這個關鍵詞一共出現了 38 次之多!比小米 17 系列發佈會雷總提到「蘋果」的 37 次還多一次。
作為一家明星創業公司,TML 的更新也吸引了廣泛關注。有人總結其優勢:
更是有網友盛讚,TML 纔是真 Open AI。
博客地址:https://thinkingmachines.ai/blog/on-policy-distillation/
這篇博客的主要作者是 Thinking Machines Lab 研究者 Kevin Lu。他之前曾在 OpenAI 工作,領導了 4o-mini 的發佈,並參與過 GPT-5 series、GPT-oss、o3 & o4-mini、4.1-nano & 4.1-mini、o1-mini、o3-mini 等模型的研發工作。
下面我們就來詳細看看這篇博客的內容。
大型語言模型(LLM)能夠在特定領域展現出專家級的水平。這是幾種能力共同作用的結果,包括:對輸入的感知、知識檢索、規劃選擇和可靠執行。
要實現這一點,需要一系列的訓練方法。我們可以將其大致分為三個階段:
預訓練(Pre-training):教授通用能力,例如語言使用、寬泛的推理和世界知識。
中訓練(Mid-training):傳授領域知識,例如代碼、醫療數據庫或公司內部文件。
后訓練(Post-training):引導出目標行為,例如遵循指令、解決數學問題或聊天。
在特定專業領域,經過強化訓練的小型模型,其表現往往優於那些大型的通用模型。使用小型模型有很多好處:
出於隱私或安全考慮,它們可以進行本地部署。
它們可以更輕松地持續訓練和更新。
它們還能節省推理成本。
想要利用這些優勢,就需要為訓練的后續階段選擇正確的方法。
后訓練「學生」模型的方法可以分為兩種:
在策略(On-policy)訓練:從學生模型自身採樣軌跡(rollouts),併爲這些軌跡分配某種獎勵。
離策略(Off-policy)訓練:依賴於某個外部來源的目標輸出,學生模型需要學習模仿這些輸出。
例如,我們可能希望訓練一個緊湊模型來解決如下的數學問題:
我們可以通過強化學習(RL)來進行在策略訓練。具體做法是根據學生模型的每個軌跡是否解決了問題來為其評分。這個評分可以由人工完成,也可以由一個能可靠給出正確答案的「教師」模型來完成。
在策略訓練的優勢在於,學生通過在自己的樣本上訓練,能更直接地學會避免錯誤。
但 RL 有一個主要缺點:它提供的反饋非常稀疏(sparse feedback)。無論使用多少 token,它在每個訓練回合(episode)中教授的比特數是固定的。
在我們上面的例子中,學生只知道「21」是錯誤答案,並更新模型以避免產生這個軌跡。但它並沒有學到究竟錯在哪里 —— 是搞錯了運算順序,還是算術本身出了錯。這種反饋的稀疏性使得 RL 在許多應用中效率低下。
離策略訓練通常通過監督微調(SFT)來完成,即在一組精心策劃的、針對特定任務的有標註示例上進行訓練。這些有標註示例的來源可以是一個在當前任務上表現出色的教師模型。
我們可以使用一種稱為蒸餾(distillation)的機制:訓練學生模型來匹配教師模型的輸出分佈。我們在教師的軌跡上進行訓練,這些軌跡是生成的 token 的完整序列,包括中間的思考步驟。
在每一步,我們既可以使用教師完整的「下一個 token 分佈」(常被稱為 「logit 蒸餾」),也可以只採樣給定的序列。實踐證明,採樣序列提供了對教師分佈的無偏估計,並能達到相同的目標。學生模型會根據自己生成該 token 的概率有多低,來相應地更新對序列中每個 token 的學習(在下例中用深色表示):
事實證明,蒸餾大型教師模型,在訓練小型模型方面非常有效,使其能夠:
遵循指令
進行數學和科學推理
從醫療筆記中提取臨牀信息
以及參與多輪聊天對話
用於這些應用和其他應用的蒸餾數據集通常是開源和公開發布的。
離策略訓練的缺點是,學生是在教師經常遇到的上下文中學習,而不是在學生自己將來會經常遇到的上下文中學習。
這可能會導致複合錯誤(compounding error):如果學生早期犯了一個教師從未犯過的錯誤,它會發現自己越來越偏離在訓練中觀察到的狀態。
當我們關心學生在長序列上的表現時,這個問題變得尤為突出。爲了避免這種偏離,學生必須學會從自己的錯誤中恢復。
離策略蒸餾觀察到的另一個問題是,學生可以學會模仿教師的風格和自信,但不一定能學會其事實的準確性。
打個比方:如果你在學習國際象棋,在策略 RL 就好比在沒有教練指導的情況下自己下棋。贏棋或輸棋的反饋與你自己的下法直接相關,但每局只收到一次反饋,而且不會告訴你哪些棋步對結果貢獻最大。離策略蒸餾則類似於觀看一位特級大師下棋 —— 你觀察到的是非常高超的棋步,但這些棋步是在新手玩家很少會遇到的棋局狀態下走出的。
我們希望能將 RL 的在策略相關性與蒸餾的密集獎勵信號結合起來。
對於學習國際象棋來説,這就好比有一位老師來為你自己的每一步棋打分,從「大錯特錯」到「妙不可言」。對於 LLM 的后訓練來説,這就是在策略蒸餾(on-policy distillation)。
在策略蒸餾的核心思想是:從學生模型中採樣軌跡,並使用一個高性能的教師模型來為每個軌跡的每一個 token 評分。
回到我們上面的數學例子,在策略蒸餾會給解題的每一步打分,懲罰那些導致學生得出錯誤答案的錯誤步驟,同時強化那些執行正確的步驟。
在這篇文章中,我們探討了在策略蒸餾在以下任務上的應用:
1. 訓練模型進行數學推理。
2. 訓練一個兼具領域知識和指令遵循能力的助手模型。
我們在已經具備預訓練和中訓練基礎能力的模型上應用在策略蒸餾。我們發現,這是一種廉價而強大的后訓練方法,它成功將在策略訓練的優勢和密集獎勵信號結合到了一起。
我們的在策略蒸餾工作借鑑了 DAGGER(Ross et al, 2010),這是一種迭代式的 SFT 算法,它包含了教師對學生訪問過的狀態的評估。
它也類似於過程獎勵建模(Lightman et al, 2023),這是一種 RL 方法,會對學生模型思維鏈中的每一步都進行評分。
我們擴展了 Agarwal et al.(2023)和 Qwen3 團隊(2025)之前的在策略蒸餾工作。使用 Tinker 訓練 API,我們復刻了 Qwen3 的成果,即通過在策略蒸餾在推理基準上實現了同等性能,而成本僅為 RL 的一小部分。
你可以在這個 Tinker cookbook 中跟着學習實現的每一步:
https://github.com/thinking-machines-lab/tinker-cookbook/tree/main/tinker_cookbook/recipes/distillation
損失函數:反向 KL
在策略蒸餾可以使用多種損失函數來為學生的軌跡評分。為簡單起見,我們選擇逐 token 的反向 KL(reverse KL)—— 即在給定相同先前軌跡的條件下,學生(π_θ)和教師(π_teacher)在每個 token 上的分佈之間的散度:
我們的獎勵函數會最小化反向 KL,這會促使學生在自己所處的每種狀態下都去近似教師的行為。當學生的行為與教師完全一致時,反向 KL 為零。為簡單起見,我們使用的折扣因子為零:在任何給定的時間步,學生只優化眼前的下一個 token,不考慮未來的 token。
反向 KL 與 RL 有着天然的協同作用,RL 通常優化由獎勵模型引導的某種序列級反向 KL。然而,與實踐中的大多數獎勵模型不同,反向 KL 是「不可破解的」(unhackable),因為從教師模型的角度來看,低 KL 總是對應着高概率的期望行為。反向 KL 的另一個有用特性是它是「尋找眾數(mode seeking)」的 —— 它學習一種特定行為(教師的行為),而不是將其分佈分散在幾個次優選項上。
這種方法可節省大量計算資源。因為它不需要等待一個軌跡完成採樣才能計算獎勵,所以我們可以使用更短或部分的軌跡進行訓練。查詢教師的對數概率也只需要大型模型進行一次前向傳播,而軌跡則是由更小、更廉價的學生模型生成的。
我們也不需要單獨的獎勵或標註模型。將基於蒸餾的逐 token 獎勵與序列級的環境獎勵結合起來可能會有好處;這是未來一個有趣的潛在研究領域。
圖解
下面我們來看一個真實的例子,這是一個錯誤的學生軌跡,由教師模型進行評分。這個例子來自 SimpleBench,它要求模型做出一個關鍵觀察:問題的前提很重要。正確答案是 「B. 0」,因為冰塊在煎鍋里會融化。而學生模型(Qwen3-4B-Instruct-2507)錯誤地將其視為一個純粹的數學問題,沒有考慮物理背景。
顏色越深,代表該 token 受到教師模型(Qwen3-235B-A22B-Instruct-2507)的懲罰越高(教師模型正確解決了這個問題)。
我們看到,它懲罰了那些引導學生誤入歧途的短語的起始 token,這直觀上對應了引導推理的重要「分叉 token」(forking tokens)。最終答案(雖然是錯的)並沒有受到懲罰 —— 因為在給定前面所有序列的條件下,這個答案是完全可預測的。
偽代碼
我們在 Tinker 的 RL 腳本之上實現了在策略蒸餾,該腳本已經實現了採樣、獎勵計算和策略梯度式的訓練。
地址:https://thinkingmachines.ai/blog/on-policy-distillation/(https://github.com/thinking-machines-lab/tinker-cookbook/blob/main/tinker_cookbook/rl/train.py)
初始化教師客户端:Tinker API 可以輕松地為不同模型創建不同的客户端。我們使用採樣客户端,因為我們不需要通過教師模型傳播對數概率。
採樣軌跡:我們像在 RL 中一樣從學生模型中採樣軌跡。在採樣期間,RL 已經計算了學生的對數概率 log π_θ(x),用作重要性採樣損失的一部分。
計算獎勵:我們用 compute_logprobs 函數在採樣出的軌跡上查詢教師客户端,它會返回教師在學生採樣的 token x 上的對數概率 log π_teacher (x)。然后我們用這個來計算反向 KL。
使用 RL 進行訓練:我們將逐 token 的優勢(advantage)設置為負的反向 KL,並調用 RL 的重要性採樣損失函數來對學生模型執行訓練更新。
偽代碼如下:
在下面的實驗中,我們通常將在策略蒸餾應用於已經過特定領域知識中訓練的模型。這種訓練提高了學生生成教師分佈範圍內的 token 的概率,儘管這通常遠不足以復刻教師的性能。通常,正如我們將在個性化示例中看到的,生成相關 token 的概率開始時為零,因為學生缺乏任何相關的領域知識。
我們將使用在策略蒸餾進行后訓練,並將其與訓練專家模型的其他最后關鍵階段的方法進行比較。
我們使用蒸餾來訓練 Qwen3-8B-Base 模型的數學推理能力,並使用 Qwen3-32B 作為教師模型。教師(Qwen3-32B)和學生(Qwen3-8B-Base)都是目前 Tinker 上支持的模型,因此你可以使用 Tinker cookbook 復現我們的實驗。
離策略蒸餾
如前所述,我們所有的實驗都以離策略蒸餾(即在教師生成的示例數據集上進行監督微調)的形式作為中訓練的起點。用於數學推理的數據集是 OpenThoughts-3,這是一個由 QwQ-32B(一個類似於 Qwen3-32B 的推理模型)生成的推理提示和響應的集合。
在 40 萬個提示上對學生模型(Qwen3-8B-Base)進行全參數微調(full fine-tuning),在 AIME'24(一個數學問題基準測試)上獲得了 60% 的分數。我們也可以使用 LoRA 進行訓練,但在高容量數據集上訓練時,它落后於全參數微調。在所有情況下,我們都看到性能呈對數線性增長 —— 最初的性能提升很廉價,但后期的提升成本高昂。
我們可以將在 40 萬個提示上微調過的模型視為一個檢查點,然后嘗試各種后訓練方法,將其在 AIME’24 基準上的分數從 60% 提高到 70%。
默認方法是在更多提示上進行微調,即繼續離策略蒸餾的過程。根據對數線性趨勢推斷,我們估計模型在大約 200 萬個提示上能達到 70% 的分數。這個推斷需要 scaling law 持續有效而不停滯,這並不簡單。
強化學習
Qwen3 技術報告稱,在類似的 SFT 初始化基礎上,通過 17,920 個 GPU 小時的 RL,在基準測試上達到了 67.6% 的性能。這很難與蒸餾的成本直接比較,但基於對 SFT 訓練堆棧的一些合理假設,這與在 200 萬個離策略蒸餾提示上訓練的成本相似。
Qwen 團隊還報告稱,使用在策略蒸餾,能以 RL 成本的十分之一,在 AIME’24 上達到了 74.4% 的更高分數。這也啓發了我們的工作。
在策略蒸餾
作為替代方案,我們運行了在策略蒸餾。從 40 萬 SFT 檢查點開始,在策略蒸餾在大約 150 個步驟內就達到了 AIME’24 70% 的成績。
跨方法比較計算成本並非易事。下面,我們用 FLOPs(浮點運算次數)來計算成本。
我們發現,當 SFT 數據集是現成的(如我們的 OpenThoughts-3 示例)或在多次訓練中被攤銷時,基線成本降低了 9 倍。
然而,我們經常希望在一個沒有現成離策略蒸餾數據集的新任務上訓練一個小模型。如果我們將教師模型的全部成本(即包括從教師模型採樣的額外成本)計算在離策略蒸餾中,那麼總成本可降低約 30 倍。
除了將小型模型訓練到在通用任務上表現出色之外,蒸餾的另一個用例是個性化。例子包括在對話中遵循特定的語氣和輸出格式,或者像工具使用和成本預算這樣的能力。我們經常希望在傳授新領域知識的同時訓練這種行為。
同時訓練這兩者通常很困難,輕量級微調(如 LoRA)往往不足以實現這一目標,因此需要更大規模的中訓練。在掌握新知識的基礎上學習后訓練行為,需要一個複雜的后訓練堆棧,通常由專有數據和獎勵模型組成。雖然前沿實驗室可以做到這一點,但其他從業者要復刻可能很困難或成本高昂。
在本節中,我們展示了在策略蒸餾可以有效地用於后訓練專業化行為。這種方法也適用於持續學習或「測試時訓練」:即在模型部署后更新它們,而不會導致基礎性能下降。我們使用一個在公司內部文檔上進行中訓練的模型作為應用示例。
訓練一個內部助手
定製模型的一個常見目標是充當助手:在某個領域擁有專家知識,並且具有可靠的助手式行為。我們可能需要對這兩者進行單獨訓練,尤其是當專業領域知識無法僅從預訓練數據中學到,或者學習它會干擾行為時。
我們的例子是一個公司內部助手,我們有兩個期望:
1. 模型對該領域(公司文檔)知識淵博。預訓練模型沒有見過任何公司內部文檔,因此無論模型規模多大,都只能猜測。我們將使用內部知識召回評估(「內部 QA」)來衡量這一點。
2. 模型表現出強大的后訓練行為,即遵循指令。我們將使用常用的 IF-eval 來衡量這一點。
訓練新知識會降低已學到的行為
我們將從 Qwen3-8B(而不是基礎模型)開始。Qwen3-8B 已經通過 RL 進行了后訓練,掌握了作為助手有用的技能,如指令遵循和推理。先前的研究表明,這種強化學習只訓練了原始模型的一小部分子網絡,因此當網絡在大量數據上進一步訓練時,可能會變得很脆弱。我們研究了這種情況發生的程度,以及如何恢復所需的行為。
爲了減少這種災難性遺忘(catastrophic forgetting),中訓練中一種常見的做法是混入來自模型原始預訓練分佈的「背景數據」。
在這種情況下,我們無法訪問 Qwen3 的預訓練分佈。因此,我們考慮一個更強、成本更高的基線:我們獲取 Tulu3 提示(一個廣泛的聊天和指令遵循數據集),並使用 Qwen3-8B 重新採樣它們,以充當聊天背景數據。
然后,我們在內部文檔和聊天數據的不同混合比例上微調 Qwen3-8B。提高文檔數據比例會直接提升模型的知識水平。然而,儘管混入至少 30% 的聊天數據有助於保留大部分指令遵循能力,但沒有任何一種權重配比能維持在 IF-eval 上的原始性能。
對於任何給定的混合比例,我們都觀察到 IF-eval 性能在微調過程中下降。
另一種常用的方法是使用 LoRA 來約束參數更新,從而減少災難性遺忘的可能性。然而,這種方法仍然不足以保留 IF-eval,而且 LoRA 學到的知識也更少。
在策略蒸餾恢復后訓練行為
接下來,我們試圖在對內部文檔進行微調后,恢復指令遵循行為。這種行為最初是用 RL 訓練的,成本高昂,而且正如我們所見,很脆弱。
取而代之的是,我們在 Tulu3 提示上,使用模型的早期版本 Qwen3-8B 作為教師,來進行在策略蒸餾。請注意,這個訓練階段與內部文檔數據無關,其唯一目的是恢復指令遵循能力。
使用模型的早期版本作為教師來「重新喚起」在微調過程中丟失的能力,這使得在策略蒸餾在持續學習(continuous learning)方面非常有前景。我們可以交替進行「在新數據上微調」和「蒸餾以恢復行為」這兩個階段,使我們的模型能夠隨着時間的推移學習並保持知識的最新狀態。
在 70-30 混合的內部文檔數據和聊天數據上微調后,在策略蒸餾幾乎完全恢復了在 IF-eval 上的性能,且沒有損失任何知識;我們還觀察到聊天能力和模型在內部 QA 評估中的「知識」性能之間存在一些正向遷移。
從本質上講,我們將語言模型本身視為一個獎勵模型,高概率的行為會受到獎勵。這與逆向 RL(inverse RL)有關:高概率的行為對應於假定的潛在偏好模型中的有利獎勵。任何經過指令調優的開源權重模型都可以在這個意義上用作獎勵模型;我們只需要能訪問 compute_logprobs 函數。
密集監督可極大提高計算效率
強化學習和在策略蒸餾都通過反向 KL 進行學習,修剪基礎策略中存在的動作空間。區別在於獎勵的密度。
在 LoRA Without Regret 中,我們提出了信息論的觀點,即強化學習每個回合只教授 O (1) 的比特。相比之下,蒸餾每個回合教授 O (N) 比特,其中 N 是 token 的數量。通過更密集的獎勵,我們到底能獲得多少訓練效率的提升?
我們做了一個實驗來直接比較兩者:
從 Qwen3-8B-Base(沒有額外的 SFT)開始。
在 DeepMath 上運行 RL。我們使用 128 的 LoRA rank。生成的模型是蒸餾的教師。
從 RL 訓練的模型(2)在策略蒸餾回基礎模型(1)。
從相同的初始化開始,在策略蒸餾學習 RL 訓練的策略所需的梯度步數大約少 7-10 倍,這對應於 50-100 倍的計算效率提升。
我們看到,蒸餾達到教師性能水平的速度比 RL 快了大約 7-10 倍。反向 KL 下降到接近零,AIME 分數在 10 個梯度步內就得以恢復,而 RL 則需要 70 步才能達到該水平。
蒸餾可以有效地重用訓練數據以提高數據效率
對於從業者來説,收集大量的訓練提示可能既困難又耗時。因此,我們希望能夠在訓練中多次重用提示。
使用 RL 時,在同一個提示上訓練多個輪次(epochs)常常導致對最終答案的簡單記憶。
相比之下,在策略蒸餾學習的是近似教師的完整分佈,而不是記憶單個答案。這使我們能夠用同一個提示訓練多個樣本。
我們重複了上述在數學上訓練 Qwen3-8B-Base 的實驗,但這次只使用數據集中隨機選擇的一個提示。我們在這個提示上連續訓練了 20 步。儘管我們只在一個提示上訓練,但我們確實達到了與教師模型相當的性能。
RL 在語義策略空間中搜索
我們已經看到,在策略蒸餾可以用少得多的訓練步驟來複刻 RL 提供的學習效果。
一種解釋是,與預訓練不同,RL 並未在梯度步驟本身上花費大量計算。我們應該認為 RL 將其大部分計算用在了搜索(search)上 —— 即推出一個策略並分配功勞 —— 而不是進行更新。
預訓練通過隨機梯度下降探索高維參數空間。預訓練需要海量信息,並且非常難以蒸餾。
相比之下,我們應該認為 RL 是在探索語義策略(semantic strategies)的空間。RL 不是在參數空間中探索,而是靠運氣「偶然」發現新策略 —— 它只是從它已有的權重集合中隨機抽樣。
一旦找到了一個好的策略,蒸餾就成了學習它的捷徑:在策略蒸餾不需要對 RL 課程中的中間策略進行建模,而只需要學習最終的策略。
打個比方:在科學研究中,我們花費大量時間和資源來尋找答案和探索新思想。一旦發現一個結果,用自然語言將其教給別人就簡單得多。相比之下,像運動這樣的直覺性身體技能,就很難教給別人,因為這些知識存在於一種天生的語言中(例如,肌肉記憶),只有我們自己才能輕易理解。運動只能通過反覆練習來學習。
作為持續學習工具的在策略學習
在關於個性化的部分,我們探討了在策略蒸餾將專業訓練行為重新引入模型的能力。這可以推廣到更廣泛的持續學習(continual learning)任務,這些任務要求在不降低先前能力的情況下獲取新知識。
先前的工作發現,在策略學習(RL)比較少地遺忘。然而,RL 只能塑造行為 —— 它不能很好地教授新知識,因此不足以用於持續學習。
在上一節中,我們看到 SFT(包括離策略蒸餾)在支持持續學習方面是失敗的,因為它會降低行為。
我們更深入地研究了這個問題。當我們在模型自己的樣本數據集上運行 SFT 時會發生什麼?我們看到,任何大於零的實用學習率都會導致指令遵循評估的性能下降!
一個可能的解釋是,雖然 KL 散度在期望上為 0,但每個有限的批次(batch)在實踐中都會表現出略微不同的分佈。在這些有限批次上訓練會導致非零的梯度更新,這會使更新后的模型策略偏離其原始狀態。
在策略蒸餾始終保持在在策略狀態,並且由於教師保持不變,學生會收斂於教師的期望行為,而不會像 SFT 那樣在自蒸餾設置中出現性能衰退。這使得在策略蒸餾成為一種非常有前景的持續學習工具。
我們探討了在策略蒸餾在訓練小型模型進行數學推理或持續學習助手等方面的應用。我們將在策略蒸餾與其他兩種后訓練方法進行了比較:離策略蒸餾和在策略 RL。
我們發現,在策略蒸餾結合了兩者的優點:在策略訓練的可靠性能以及密集獎勵信號帶來的成本效益。
后訓練是達到前沿模型能力的關鍵部分。通過利用來自學生的在策略採樣和來自教師的密集監督,在策略蒸餾方案能夠以前沿高計算量 RL 運行成本的一小部分,達到這些能力。
我們的實現可以在 Tinker cookbook 中找到。我們的工作探索了在策略蒸餾的簡單直接的實例化,以清晰地展示其優勢。我們希望繼續研究蒸餾的新應用、改進教師監督的新方法,以及提高數據效率和持續學習的方法。
在 Thinking Machines,我們的使命是為人們提供兼具前沿性能、適應性和個性化的 AI 模型。在策略蒸餾是實現這一目標的有力工具。
參考鏈接
https://x.com/thinkymachines/status/1982856272023302322
https://x.com/miramurati/status/1982856564970254772
本文來自微信公眾號「機器之心」,36氪經授權發佈。