標準 MTP 訓練中的梯度干擾問題

多詞元預測（MTP）已成為語言模型預訓練的常見技術，旨在提升下游任務表現並透過投機解碼加速推論。然而，研究指出標準 MTP 訓練存在缺陷：MTP 梯度中的殘差成分會干擾下一個詞元預測（NTP）與多詞元學習，且這種干擾隨模型規模擴大而加劇。

具體而言，MTP 的交叉熵梯度可精確分解為兩部分：一是將 MTP 頭部拉向模型自身 NTP 分布的 KL 自蒸餾項，這有助於改善投機解碼；二是信噪比極低的殘差項，它會掩蓋有用的 KL 信號。

Soft-MTP 的梯度恢復技巧與實作

研究團隊提出「KL 梯度恢復技巧」：當兩個分布共享線性輸出頭時，可透過直通梯度路由（straight-through gradient routing）組合它們的交叉熵梯度，精確恢復其 KL 散度的梯度。與直接最小化 KL 散度不同，此技巧不需要詞彙表大小的 logits 矩陣，且相容於標準的融合交叉熵核心。

基於此技巧，Soft-MTP 被定義為對 MTP 的一行程式碼修改。它移除了雜訊殘差，僅透過重用現有的 NTP 和 MTP 交叉熵計算來反向傳播 KL 自蒸餾信號。該方法支援 DeepSeek 與 Meta 風格的 MTP，不增加參數，且無可測量的額外開銷。

不同參數規模下的效能提升驗證

在 167M 至 1.2B 參數的實驗範圍內，Soft-MTP 在所有測試規模下均優於標準 MTP。具體表現包括降低 NTP 困惑度、加速投機解碼以及提高零樣本準確率。

從 689M 參數開始，Soft-MTP 還能降低 MTP 困惑度。在 1.2B 參數進行微調後，該方法展現出更強的數學與程式碼推理能力。此外，Soft-MTP 在多數指標上表現等同甚至超越直接最小化 KL 散度的方法。