最近，DeepSeek-R1 和 OpenAI o1/03 等推理大模型在后訓(xùn)練階段探索了長度擴展（length scaling），通過強化學(xué)習(xí)（比如 PPO、GPRO）訓(xùn)練模型生成很長的推理鏈（CoT），并在奧數(shù)等高難度推理任務(wù)上取得了顯著的效果提升。

受此啟發(fā)，研究人員開始探索預(yù)訓(xùn)練階段的長度擴展，已有方法包括在序列中插入文本、插入潛在向量（如 Coconut）、復(fù)用中間層隱藏狀態(tài)（如 CoTFormer）以及將中間隱藏狀態(tài)映射為概念（如 COCOMix）。不過，這些方法普遍存在問題，比如需要更大的 KV 緩存導(dǎo)致推理慢 / 占內(nèi)存多。

本文中，來自 ByteDance Seed 團隊的研究者提出了更簡單的方法：直接重復(fù)輸入 tokens（1/2/3/4 次），不做中間層處理。他們觀察到了訓(xùn)練損失和模型性能隨重復(fù)倍數(shù)擴展的趨勢，如下圖 1a 和 1b 所示。但是，直接重復(fù) tokens 也帶來了新問題，包括 KV 緩存規(guī)模線性增加，內(nèi)存壓力大；預(yù)填充時間超線性增加；解碼延遲變長。這些都是實現(xiàn)預(yù)訓(xùn)練長度擴展需要重點解決的挑戰(zhàn)。

• 論文標(biāo)題：Efficient Pretraining Length Scaling
• arXiv 地址：https://arxiv.org/pdf/2504.14992

研究者提出了一種推理友好的新穎長度擴展方法，核心是 PHD-Transformer（Parallel Hidden Decoding Transformer），它保持了與原始 transformer 相同的 KV 緩存大小，同時實現(xiàn)有效的長度擴展。PHD-Transformer 通過創(chuàng)新的 KV 緩存管理策略實現(xiàn)了這些能力。

具體來講，研究者將第一個 token 表示原始 token，將重復(fù)的 token 表示為解碼 token。同時僅保留從原始 token 生成的 KV 緩存來用于長距離依賴建模，并在隱藏解碼 token 用于下一個 token 預(yù)測之后丟棄它們的 KV 緩存。因此，PHD-Transformer 提供了與原始 transformer 相同的 KV 緩存，同時相較于簡單的 token 重復(fù)實現(xiàn)了顯著的推理加速（如圖 1d 所示）。

另外，為了更好地保留隱藏解碼 token 的 KV 緩存的性能優(yōu)勢，研究者引入了一種滑動窗口注意力 ——PHD-SWA，保持了這些 token 的局部滑動窗口緩存，在實現(xiàn)顯著性能提升的同時，僅需要的額外 KV 緩存內(nèi)存。

研究者還注意到，在 PHD-SWA 中，隱藏解碼 token 的 KV 緩存表現(xiàn)出了順序依賴關(guān)系，這導(dǎo)致預(yù)填充時間呈線性增長。為了解決這個問題，研究者提出了逐塊滑動窗口注意力 —— PHD-CSWA，從而限制了每個塊內(nèi)的順序依賴關(guān)系。

因此，得益于只有最后一個塊的預(yù)填充時間呈線性增長，PHD-CSWA 顯著縮短了預(yù)填充時間（如圖 1c 所示）。

方法概覽

PHD 的架構(gòu)下圖 2 所示，與原始 Transformer 相比，PHD 保留了相同的模型架構(gòu)，僅在輸入序列和注意力矩陣的設(shè)計上有所不同。具體而言，他們僅允許原始 token生成 KV 緩存，并且可以被所有 token 全局關(guān)注；同時隱藏狀態(tài)的 KV 緩存在并行隱藏解碼后會被立即丟棄。注意力矩陣的策略具體如下： 

研究者在推理過程中實現(xiàn)了與原始 Transformer 相同的 KV 緩存大小和內(nèi)存訪問模式。雖然需要 K 次 FLOP，但這些計算可以并行處理，從而在內(nèi)存受限的推理場景中最大限度地降低延遲開銷。該架構(gòu)的核心優(yōu)勢在于原始 token 和隱藏解碼 token 之間的解耦。在預(yù)填充期間，只有原始 token 需要計算。

這種設(shè)計確保預(yù)填充時間與原始 Transformer 相同，并且無論擴展因子 K 如何變化，預(yù)填充時間都保持不變。而對于損失計算，研究者僅使用 token 的最終副本進行下一個 token 的預(yù)測?？傊褂?token 的第一個副本進行 KV 緩存生成，使用 token 的最后一個副本進行下一個 token 的預(yù)測。

內(nèi)核設(shè)計

M^ij_mn 的簡單實現(xiàn)會導(dǎo)致注意力層計算量增加 K^2 倍，F(xiàn)FN 層計算量也增加 K 倍。然而，由于注意力是稀疏計算的，的注意力可以大幅降低。因此，研究者將原始 token 和隱藏解碼 token 分成兩組，并將它們連接在一起。

下圖 3 展示了 K = 3 的示例，可以得到一個包含 t 個原始 token 的序列和一個包含 2t 個隱藏解碼序列的序列。通過重新排列 token 的位置，研究者將掩碼注意力的位置保留在一個連續(xù)塊中，從而優(yōu)化了注意力計算，將注意力計算復(fù)雜度降低到。

PHD-SWA 和 PHD-CSWA

與簡單的 token 重復(fù)相比，PHD-Transformer 在保持原始 KV 緩存大小的同時實現(xiàn)了長度擴展。然而通過經(jīng)驗觀察到，為隱藏解碼 token 保留一些 KV 緩存可以帶來顯著的性能提升。因此，為了在保持效率的同時獲得這些優(yōu)勢，研究者引入了 PHD-SWA，將滑動窗口注意力限制在 W 個先前的隱藏解碼 token 上。

如下圖 4 所示，PHD-SWA 的注意力模式將對原始 token 的全局訪問與對 W 個最近隱藏解碼 token 的局部訪問相結(jié)合。這種改進的注意力機制實現(xiàn)了顯著的性能提升，同時僅需要的額外 KV 緩存內(nèi)存。

雖然 PHD-SWA 滑動窗口方法提升了模型性能，但由于隱藏解碼 token 的 KV 緩存中存在順序依賴關(guān)系，它會產(chǎn)生 K 倍的預(yù)填充開銷。為了解決這個問題，研究者引入了 PHD-CSWA，它可以在獨立的塊內(nèi)處理注意力。 

如下圖 4 所示，PHD-CSWA 將滑動窗口注意力限制在單個塊內(nèi)運行。這種架構(gòu)創(chuàng)新將額外的預(yù)填充開銷減少到最終塊內(nèi)的 K 次重復(fù)，而不是整個序列重復(fù)，這使得額外的計算成本幾乎可以忽略不計，同時保留了局部注意力模式的優(yōu)勢。

實驗結(jié)果

在實驗中，研究者使用 OLMo2 作為代碼庫，并在 ARC、HellaSwag、PIQA、Winogrande、MMLU 和 CommonsenseQA 等公開基準(zhǔn)測試集上進行了評估。

訓(xùn)練細(xì)節(jié)：研究者使用 1.2B 參數(shù)規(guī)模的模型，它是一個 16 層的密集模型。每個 token 的隱藏層維數(shù)設(shè)置為 2048，F(xiàn)FN 層的隱藏層大小設(shè)置為 16384。同時使用組查詢注意力 (Group-Query Attention，GQA)，它包含 32 個查詢頭和 8 個鍵 / 值頭，每個頭的隱藏層維數(shù)設(shè)置為 64。研究者使用 500B 個 token 訓(xùn)練該模型。

對于本文提出的 PHD 系列設(shè)置，研究者預(yù)訓(xùn)練了以下兩種 PHD-CSWA 變體：

• PHD-CSWA-2-16-32，其中訓(xùn)練 token 重復(fù)兩次。保留一個包含 16 個 token 的局部窗口，并將塊大小設(shè)置為 32 個 token。 
• PHD-CSWA-3-16-32，其中訓(xùn)練 token 重復(fù)三次。局部窗口大小和塊大小與 PHD-CSWA-2-16-32 的設(shè)置相同。

PHD-CSWA 在各個基準(zhǔn)測試中均實現(xiàn)了持續(xù)的性能提升。下圖 5 中展示了訓(xùn)練曲線，下表 1 中展示了主要結(jié)果。本文提出的 PHD-CSWA-2-16-32 在這些基準(zhǔn)測試中平均實現(xiàn)了 1.5% 的準(zhǔn)確率提升，訓(xùn)練損失降低了 0.025；而 PHD-CSWA-3-16-32 在這些基準(zhǔn)測試中平均實現(xiàn)了 2.0% 的準(zhǔn)確率提升，訓(xùn)練損失降低了 0.034。

研究者還分析了 PHD 和 PHD-SWA 的擴展性能，以分析擴展解碼計算的性能。 訓(xùn)練細(xì)節(jié)：使用相同的 550M 模型配置，將窗口大小 W 設(shè)置為 16，并在 {2, 3, 5} 范圍內(nèi)改變擴展因子 K。對于局部窗口大小，研究者在所有實驗中都將窗口大小設(shè)置為 16。

PHD-SWA 的性能在增加擴展因子時有效擴展。如下圖 8 所示，使用固定窗口大小時，損失曲線和下游性能會隨著 token 重復(fù)次數(shù)而有效擴展。通過將擴展因子設(shè)置為 5，可以實現(xiàn)接近 0.06 的損失降低，同時顯著提升下游性能。

下表 2 中的定量結(jié)果表明，當(dāng)擴展至 K = 5 時，所有基準(zhǔn)測試的平均準(zhǔn)確率提高了 1.8%，這證實了本文的方法在更激進的擴展方面仍然有效。

更多實驗結(jié)果請參閱原論文。