揭示顯式CoT訓練機制:思維鏈如何增強推理泛化能力
基于逐步生成解決方案的大語言模型(LLMs)訓練范式在人工智能領域獲得了廣泛關注,并已發展成為行業內的主流方法之一。
例如,OpenAI 在其「12 Days of OpenAI」直播系列的第二日推出了針對 O1 模型的強化微調(Reinforcement Fine-Tuning,RFT),進一步推動了 AI 定制化的發展[1]。RFT/ReFT[2] 的一個關鍵組成部分是使用思維鏈(Chain-of-Thought,CoT)注釋[3] 進行監督微調(Supervised Fine-Tuning,SFT)。在 DeepSeek-R1 模型[4] 中,引入了少量長 CoT 冷啟動數據,以調整模型作為初始強化學習的代理。
然而,為了全面理解采用 CoT 訓練的策略,需要解決兩個關鍵問題:
- Q1:與無 CoT 訓練相比,采用 CoT 訓練有哪些優勢?
- Q2:如果存在優勢,顯式 CoT 訓練的潛在機制是什么?
由于實際訓練過程中涉及眾多因素,分析顯式 CoT 訓練的優勢及其潛在機制面臨顯著挑戰。為此,我們利用清晰且可控的數據分布進行了詳細分析,并揭示了以下有趣現象:
- CoT 訓練的優勢
(i)與無 CoT 訓練相比,CoT 訓練顯著增強了推理泛化能力,將其從僅適用于分布內(in-distribution, ID)場景擴展到 ID 和分布外(out-of-distribution, OOD)場景(表明系統性泛化),同時加速了收斂速度(圖 1)。
圖表 1: 模型在優化過程中對訓練和測試兩跳推理事實的準確率。
(ii)即使 CoT 訓練中包含一定范圍的錯誤推理步驟,它仍能使模型學習推理模式,從而實現系統性泛化(圖 4 和圖 5)。這表明數據質量比方法本身更為重要。訓練的主要瓶頸在于收集復雜的長 CoT 解決方案,而推理步驟中存在少量的錯誤是可以接受的。
- CoT 訓練的內部機制
(i)數據分布的關鍵因素(如比例 λ 和模式 pattern)在形成模型的系統性泛化中起著決定性作用。換句話說,在 CoT 訓練中僅接觸過兩跳數據的模型無法直接泛化到三跳情況,它需要接觸過相關模式。
(ii)通過 logit lens 和 causal tracing 實驗,我們發現 CoT 訓練(基于兩跳事實)將推理步驟內化到模型中,形成一個兩階段的泛化電路。推理電路的階段數量與訓練過程中顯式推理步驟的數量相匹配。
我們進一步將分析擴展到推理過程中存在錯誤的訓練數據分布,并驗證了這些見解在現實數據上對更復雜架構仍然有效。
據我們所知,我們的研究首次在可控制的實驗中探索了 CoT 訓練的優勢,并提供了基于電路的 CoT 訓練機制解釋。這些發現為 CoT 以及 LLMs 實現穩健泛化的調優策略提供了寶貴的見解。
- 論文標題:Unveiling the Mechanisms of Explicit CoT Training: How Chain-of-Thought Enhances Reasoning Generalization
- 論文鏈接:https://arxiv.org/abs/2502.04667
一、預備知識與定義
本部分介紹研究使用的符號定義,具體如下:
原子與多跳事實:研究使用三元組來表示原子(一跳)事實,并基于原子事實和連接規則來表示兩跳事實以及多跳事實。
訓練數據:研究使用的訓練數據包括所有的原子(一跳)事實(即),以及分布內(ID)的兩跳事實(即
)。其中記 | 兩跳事實 |:| 原子事實 |= λ。
訓練方式:對于原子(一跳)事實,模型的訓練和評估通過預測最終尾實體來完成。對于兩跳事實,考慮是否使用 CoT 注釋進行訓練。
(1) Training without CoT:模型輸入 ,預測目標只有最終尾實體
;
(2) Training with CoT:模型輸入 ,預測橋接實體
和最終尾實體
。
評估:為更好地評估模型的泛化能力,我們從分布內(ID)和分布外(OOD)兩個維度進行性能評估。
(1)分布內泛化旨在通過評估模型完成未見過的兩跳事實的能力,判斷模型是否正確學習了潛在模式。
(2)分布外泛化則用于評估模型獲得的系統性能力,即模型將學習到的模式應用于不同分布知識的能力,這是通過在事實上測試模型來實現的。若模型在分布內數據上表現良好,可能僅表明其記憶或學習了訓練數據中的模式。然而,在分布外數據上的優異表現則表明模型確實掌握了潛在模式,因為訓練集僅包含原子事實
,而不包含
。
二、系統性組合泛化
本研究聚焦于模型的組合能力,即模型需要將不同事實片段「串聯」起來的能力。盡管顯式的推理步驟表述(如思維鏈推理)能夠提升任務表現 [4-8],但這些方法在大規模(預)訓練階段并不可行,而該階段正是模型核心能力形成的關鍵時期 [9-10]。已有研究對基于 Transformer 的語言模型是否能夠執行隱式組合進行了廣泛探討,但均得出了否定結論 [11-12]。
具體而言,存在顯著的「組合性鴻溝」[11],即模型雖然掌握了所有基礎事實卻無法進行有效組合的情況,這種現象在不同大語言模型中普遍存在,且不會隨模型規模擴大而減弱。
更準確地說,Wang 等人 [13] 的研究表明,Transformer 模型能夠在同分布泛化中學習隱式推理,但在跨分布泛化中則表現欠佳(如圖 1 左所示)。
這自然引出一個問題:如果在訓練過程中使用顯式推理步驟,模型的泛化能力將受到何種影響?(即回答 Q1:與無思維鏈訓練相比,基于思維鏈的訓練具有哪些優勢?)
思維鏈訓練顯著提升推理泛化能力
如圖 1 所示,我們展示了模型在訓練和測試兩跳事實上的準確率隨優化過程的變化,其中 λ = 7.2。
(1)Training without CoT(圖 1 左)。我們觀察到了與 Wang 等人 [13] 相同的現象(稱為頓悟現象 [14]),即模型能夠較好地泛化到分布內測試樣本,但高性能只有在經過大量訓練后才能實現,遠超過過擬合點。此外,即使經過數百萬次優化步驟的訓練,仍未觀察到分布外泛化(
)的跡象,這表明這是一種缺乏系統性的延遲泛化現象。模型可能只是記憶或學習了訓練數據中的模式。
(2)Training with CoT(圖 1 右)。使用思維鏈標注后,模型在訓練集上的收斂速度加快,且在訓練過程中更早地實現了較高的測試性能,特別是在分布內測試樣本上。模型在大約 4,000 次優化步驟后,在同分布測試集上的準確率就達到了接近完美的水平,表明與無思維鏈訓練相比,泛化能力得到了顯著提升。分布外泛化(
)也顯示出明顯改善,這突出表明思維鏈提示訓練不僅在分布內泛化方面,而且在分布外泛化方面都發揮著關鍵作用,盡管效果程度有所不同。
關鍵影響因素探究
研究進一步開展了消融實驗,以評估不同因素在思維鏈訓練中的影響。
圖表 2: 分布外測試集上的推理泛化速度。
適當的 λ 值能夠加速模型收斂。圖 2(左)展示了不同 λ 值下的分布外測試準確率。可以看出,λ 值與泛化速度存在強相關性。更有趣的是,較小的 λ 值能夠加速由思維鏈訓練帶來的分布外泛化能力提升,從而減少對長時間訓練的需求。然而,λ 值并非越小越好,因為過小的 λ 值可能導致模型無法學習相關規則。
不同模型規模 / 層數和訓練集大小的影響。我們在模型層數∈{2,4,8} 和 λ∈{3.6,7.2,12.6} 的條件下進行實驗。總體而言,可以觀察到擴大模型規模并不會從根本上改變其泛化行為,主要趨勢是較大的模型能夠在更少的優化步驟中收斂。關于訓練集大小(|E|)的影響,我們的結果與 [13] 一致:當固定 λ 值時,訓練集大小不會對模型的泛化能力產生本質影響。
兩跳到多跳分析
在本部分中,研究將重點轉向多跳場景:在思維鏈訓練階段僅接觸過兩跳事實的模型,能否泛化到三跳事實?
在思維鏈訓練中,我們僅使用單跳 / 兩跳事實,并測試模型是否能夠泛化到三跳事實的推理(這里研究使用來表示三跳事實)。
結果:在思維鏈訓練中僅接觸過兩跳數據的模型無法直接泛化到三跳場景。然而,當訓練集中加入一定量的三跳數據后,模型能夠快速實現泛化(前提是模型需要接觸過相關模式)。另一方面,當我們人為地將一個三跳事實拆分為兩個兩跳事實進行測試時,模型也能夠有效泛化。換句話說,我們分別測試 預測
和
預測
,當兩者都正確時,我們認為
預測
是正確的。這些發現與 [15] 結果一致:思維鏈與重現訓練集中出現的推理模式有關。
總結:至此,我們已經證明在受控實驗中引入顯式思維鏈訓練能夠顯著提升推理泛化能力,使其從僅限分布內泛化擴展到同時涵蓋分布內和分布外泛化。數據分布的關鍵因素(如比例和模式)在形成模型的系統性泛化能力中起著重要作用。然而,驅動這些改進的內部機制仍不明確,我們將進一步探討(回答 Q2:如果存在優勢,顯式思維鏈訓練的潛在機制是什么?)。
圖表 3: 兩跳事實訓練對應的兩階段泛化電路(模型層數:8)。
三、兩階段泛化電路
研究通過兩種主流方法分析模型在泛化過程中的內部工作機制:logit lens [16] 和 causal tracing [17],本部分研究使用表示兩跳推理。
圖 3 展示了發現的泛化電路,該電路代表了 8 層模型在實現兩跳分布外(OOD)泛化后的因果計算路徑。具體而言,我們識別出一個高度可解釋的因果圖,該圖由第 0 層、第 l 層和第 8 層的狀態組成,其中弱節點和連接已被剪枝(If perturbing a node does not alter the target state (top-1 token through the logit lens), we prune the node)。
(1)在第一跳階段,第 l 層將電路分為上下兩部分:下部從輸入中檢索第一跳事實,并將橋接實體
存儲在狀態
中;上部通過殘差連接將的信息傳遞到輸出狀態(其中
表示對應位置的激活)。由于數據分布可控,l 層可以精確定位(對于 ID 為第 3 層,對于 OOD 為第 5 層)。
(2)在第二跳階段,自回歸模型使用第一跳階段生成的。該階段省略了
,并從輸入
處理第二跳,將尾實體
存儲到輸出狀態
中。
系統性泛化解釋
(1)兩階段泛化電路表明,使用思維鏈訓練可以將推理步驟內化到模型中。這也解釋了為什么模型在思維鏈訓練下能夠在跨分布測試數據上表現出良好的泛化能力。
(2)該電路由兩個階段組成,與訓練期間模型中的顯式推理步驟相一致。因此,模型在思維鏈訓練期間僅接觸兩跳數據時無法在測試階段直接泛化到三跳場景。
四、更普適的分析
總體而言,我們目前的研究為通過受控數據分布上的思維鏈訓練來深入理解和增強 Transformer 的泛化能力鋪平了道路。然而,現實世界中的訓練數據分布往往更為復雜。在本部分中,我們將分析擴展到推理過程中存在錯誤的分布,并展示思維鏈訓練能提高模型的泛化能力的結論在更復雜的場景中仍然成立。
數據分布帶噪
方法:我們旨在分析通過思維鏈訓練獲得的系統性泛化能力在噪聲訓練數據下的魯棒性。我們通過隨機選擇一個有效實體向引入噪聲(真實訓練目標為
):
(1)僅第二跳有噪聲,即;
(2)兩跳均有噪聲,即。
需要注意的是,噪聲比例用 ξ 表示,我們將探討不同 ξ 值的影響。
圖表 4: 僅第二跳噪聲對分布內和分布外的影響。
圖表 5: 模型在不同噪聲比例(兩跳均有噪聲)下對訓練和測試兩跳推理事實的準確率。
結果:我們針對兩種情況分析了不同的 ξ(噪聲比例)候選集:僅第二跳有噪聲時為 {0.05, 0.2, 0.4, 0.6, 0.8},兩跳均有噪聲時為 {0.05, 0.1, 0.2, 0.4}。比較結果如下:
(1)圖 4 清晰地展示了僅第二跳噪聲對分布內和分布外泛化的影響。總體而言,在思維鏈訓練條件下,模型仍能夠從噪聲訓練數據中實現系統性泛化,但其泛化能力隨著噪聲比例的增加而降低。
更具體地說,隨著訓練的進行,分布外泛化最初保持不變,然后增加,而分布內泛化先增加后減少。分布內泛化的減少與分布外泛化的增加相對應。
然而,隨著噪聲比例的增加,分布內和分布外泛化的最終性能都會下降。特別是當噪聲比例(ξ < 0.2)相對較小時,模型幾乎不受影響,這展示了思維鏈訓練的魯棒性。
此外,我們同樣檢查了泛化電路。由于我們僅在第二跳添加噪聲,第一跳階段的電路學習得相對較好,而第二跳階段的電路受噪聲影響更大。
(2)圖 5 展示了在兩跳噪聲 ξ 值為 0.05、0.1、0.2 和 0.4 時的結果比較。與僅在第二跳添加噪聲相比,在兩跳都添加噪聲對模型泛化的抑制效果要強得多。大于 0.2 的噪聲比例足以幾乎消除分布內和分布外泛化能力。
總而言之,即使在訓練數據存在噪聲的情況下,當噪聲在一定范圍內時,思維鏈訓練仍能使模型實現系統性泛化。特別是當噪聲比例較小時,這些噪聲數據仍能幫助模型學習泛化電路。
五、討論
總結
本文通過在受控和可解釋的環境中展示系統性組合泛化如何通過顯式思維鏈(CoT)訓練在 Transformer 中產生,揭示了思維鏈訓練的核心機制。具體而言:
(1)與無思維鏈訓練相比,思維鏈訓練顯著增強了推理泛化能力,使其從僅限分布內(ID)泛化擴展到同時涵蓋分布內和分布外(OOD)場景。
(2)通過 logit lens 和 causal tracing 實驗,我們發現思維鏈訓練(使用兩跳事實)將推理步驟內化到 Transformer 中,形成了一個兩階段泛化電路。然而,模型的推理能力受訓練數據復雜性的限制,因為它難以從兩跳情況泛化到三跳情況。這表明思維鏈推理主要是重現了訓練集中存在的推理模式。
(3)我們進一步將分析擴展到推理過程中存在錯誤的訓練數據分布,證明當噪聲保持在一定范圍內時,思維鏈訓練仍能使模型實現系統性泛化,此類噪聲數據的結構或許有助于泛化電路的形成。
有趣的是,我們的工作還突出了思維鏈訓練的瓶頸:訓練數據分布(比例 λ 和模式)在引導模型實現泛化電路方面起著關鍵作用。模型需要在訓練過程中接觸過相關模式(特別是思維鏈步驟的數量)。
這可能解釋了為什么 DeepSeek-R1 [4] 在冷啟動階段構建和收集少量長思維鏈數據來微調模型。我們的發現為調整大語言模型(LLMs)以實現穩健泛化的策略提供了關鍵見解。
不足與未來展望
(1)盡管我們的自下而上的研究為實際應用提供了寶貴的見解,但我們工作的一個關鍵局限是實驗和分析基于合成數據,這可能無法完全捕捉現實世界數據集和任務的復雜性。雖然我們的一些結論也在 Llama2-7B [18] 等模型中得到了驗證,但有必要在更廣泛的模型上進行進一步驗證,以彌合我們的理論理解與實際應用之間的差距。
(2)我們的分析目前僅限于使用自然語言。未來,我們旨在探索大型語言模型在無限制潛在空間中的推理潛力,特別是通過訓練大型語言模型在連續潛在空間中進行推理 [19] 等方法。
(3)最近的一種方法,「backward lens」[20],將語言模型的梯度投影到詞匯空間,以捕捉反向信息流。這為我們完善思維鏈訓練的潛在機制分析提供了一個新的視角。
作者介紹
劉勇,中國人民大學,長聘副教授,博士生導師,國家級高層次青年人才。長期從事機器學習基礎理論研究,共發表論文 100 余篇,其中以第一作者 / 通訊作者發表頂級期刊和會議論文近 50 篇,涵蓋機器學習領域頂級期刊 JMLR、IEEE TPAMI、Artificial Intelligence 和頂級會議 ICML、NeurIPS 等。獲中國人民大學「杰出學者」、中國科學院「青年創新促進會」成員、中國科學院信息工程研究所「引進優青」等稱號。主持國家自然科學面上 / 基金青年、北京市面上項目、中科院基礎前沿科學研究計劃、騰訊犀牛鳥基金、CCF - 華為胡楊林基金等項目。
姚鑫浩,中國人民大學高瓴人工智能學院博士研究生,本科畢業于中國人民大學高瓴人工智能學院。當前主要研究方向包括大模型推理與機器學習理論。