Unmasking On-Policy Distillation: Where It Helps, Where It Hurts, and Why

📌 Unmasking On-Policy Distillation：診斷框架揭示何時有幫助、何時有害

TL;DR：提出免訓練的 per‑token 診斷方法，可量測蒸餾梯度與理想更新的對齊度，發現錯誤推理路徑的對齊度較高，而正確路徑易受噪聲影響。

🎣 開場鉤子
當蒸餾用來訓練推理模型時，直覺會認為「師傅越好，學徒越快進步」——但實際上，師傅的指導在學徒已經答對時可能變得反而雜亂無章。

🤔 背景或問題
On‑policy 蒸餾能為每個 token 提供密集監督，有助於訓練推理模型。然而，目前缺乏低成本的方式來判斷：

• 哪種教師模型（或自蒸餾的哪段上下文）最適合作為監督訊號？
• 這個最佳選擇是否會隨著 token 的不同而變化？
  傳統做法必須靠昂貴的完整訓練跑來觀察總體效能，因而難以觀察單一 token 層面的動態。

🧩 方法或架構
論文提出一個 training‑free 診斷框架，操作 granularity 達到「per token、per question、per teacher」：

1. 理想 per‑node gradient – 定義為能最大程度提升學生模型成功機率的引數更新。
2. 可擴充套件的 targeted‑rollout 演算法 – 用以高效估計上述理想梯度，即使在長鏈中間思考（chain‑of‑thought）情境下也適用。
3. 梯度對齊分數（gradient alignment score） – 計算理想梯度與任意給定蒸餾梯度的餘弦相似度，用來量測特定配置逼近理想訊號的程度。

此框架不需要實際進行完整訓練，即可在不同教師模型或不同監督上下文的效果。

📊 資料或結果
作者在多種自蒸餾設定與外部教師模型上應用該診斷框架，觀察到：

• 蒸餾引導在 錯誤推理路徑（incorrect rollouts） 上與理想梯度的對齊度 顯著高於 正確路徑（correct rollouts），因為在學生已表現良好時，教師訊號易變得雜亂。
• 最佳的蒸餾上下文（即作為監督訊號的特定 token 範圍）同時取決於 學生模型的容量 與 目標任務的特性，並非一成不變。

（摘要未提供具體資料表格或基線比較，故僅敘述上述定性觀察。）

💡 深入分析
透過 per‑token 的梯度對齊分數，研究者得以分辨哪些 token 受益於蒸餾、哪些 token 被噪聲幹擾。這種細粒度檢視解釋為何在某些情境下自蒸餾會提升推理能力，而在其他情境則可能導致效能停滯甚至倒退。框架的訓練免除特性也意味著實驗成本大幅下降，使得快速迭代不同教師模型或監督策略成為可能。

⚠️ 限制

• 框架依賴於能夠計算「理想 per‑node gradient」的近似；若目標任務的成功機率難以定義，估計可能不準。
• 目前的實驗聚焦於特定的推理基準與模型規模，未涵蓋所有可能的師生組合或更長的思考鏈。
• 作者尚未在大規模實際訓練中驗證該診斷分數與最終模型效能的直接相關性。

🎯 實務啟示
對於從事推理模型蒸餾的工程師而言，可先利用此訓練免除的診斷工具：

• 快速篩選不同教師模型或自蒸餾的上下文視窗，選擇在錯誤路徑上對齊度最高的配置。
• 在資源受限的環境中，先以低成本的 per‑token 對齊分數做決策，再投入較少的完整訓練驗證最佳選項。
  如此能減少盲目試誤的訓練成本，並更有針對性地提升模型在錯誤推理路徑上的表現。

🔗 來源

• 標題：Unmasking On-Policy Distillation: Where It Helps, Where It Hurts, and Why
• 作者／機構：Mohammadreza Armandpour*, Fatih Ilhan*, David Harrison, Ajay Jaiswal, Duc N.M Hoang, Fartash Faghri, Yizhe Zhang, Minsik Cho, Mehrdad Farajtabar @ Apple ML
• 連結：https://machinelearning.apple.com/research/unmasking-on-policy-distillation

#AppleML #OnPolicyDistillation #GradientAlignment #TrainingFreeDiagnostic #ReasoningModels #SelfDistillation #TeacherStudent #NLP #MachineLearning #ResearchPaper