金天,麻省理工學院(MIT)計算機科學與人工智能實驗室(CSAIL)博士五年級學生,師從 Michael Carbin 和 Jonathan Ragan-Kelley。他主要研究機器學習與編程系統的結合。此前曾在 IBM Research 主導實現深度神經網絡在 IBM 主機上的推理部署。本科畢業于 Haverford College,獲計算機科學與數學雙學位。
鄭鈺熹,麻省理工學院 CSAIL 博士三年級學生,師從 Michael Carbin。她的研究方向為編程語言與機器學習的交叉領域。
大語言模型(LLM)的生成范式正在從傳統的「單人書寫」向「分身協作」轉變。傳統自回歸解碼按順序生成內容,而新興的異步生成范式通過識別語義獨立的內容塊,實現并行生成。
如圖所示,傳統方法(下)按順序生成所有內容,而異步生成(上)同時處理多個互不依賴的內容塊。對比順序生成,異步生成在 AlpacaEval 長度控制評測中實現1.21-1.93× 的幾何平均提速,對應生成質量變化(勝率)為 +2.2% 至 -7.1%。
MIT 與谷歌研究團隊在最新研究 PASTA(PArallel STructure Annotation)中首次從策略學習(policy learning)角度探索異步生成范式的可能。
- 論文標題:Learning to Keep a Promise: Scaling Language Model Decoding Parallelism with Learned Asynchronous Decoding
- 論文地址:https://arxiv.org/abs/2502.11517
研究團隊不依賴人工設計規則來識別異步生成機會,而通過策略學習讓模型自主發現并標注這些機會,系統地優化質量與速度的平衡。這種方法使 LLM 根據內容特點自適應地確定最佳異步生成策略,為生成效率優化開創學習驅動的全新路徑。
PASTA-LANG:劃分獨立內容的標記語言
研究人員首先開發了一種新的標記語言 PASTA-LANG,專為異步生成而設計。大模型使用它在生成過程中標記語義獨立塊,指示并行生成機會。這種語言包含三種核心標記:
- :標記語義獨立的內容塊,通過 topic 屬性總結內容主題,大模型用它表明「這部分將會由一個獨立子線程異步生成」。
- :在
- 后標識對應的異步生成的內容,表示這一部分由獨立子線程負責生成。
- :在主線程標記同步點,表明后續內容生成將會依賴于之前的異步生成,主線程需等待所有異步線程完成后才能繼續。
這些標記共同構成了一種「承諾-履行」的生成模式:大模型首先通過
標記「承諾」生成某些內容,推理系統再創建異步線程來「履行」這些承諾,最后在
處將異步內容組合。
如圖所示的線段長度計算案例展示了這一機制:面對計算任務(A),大模型首先識別出「坐標提取」和「長度公式」兩個可并行的獨立部分,生成相應的
標記(B),隨后用
標記(E)表明需要等待這些內容完成。圖中紅色和綠色區域(C、D)顯示了兩個異步線程并行生成的內容,最終在(F)處組合成完整解答。
這個新的標記語言簡單,可擴展性強,開啟了新的未來研究范式。
PASTA 訓練:從標注到優化的雙階段學習
如圖所示,PASTA 系統采用雙階段訓練流程,使大模型自主學習使用上述標記語言,完成異步生成。
第一階段:監督微調。研究團隊首先選取 SlimOrca 指令跟隨數據集,用 Gemini 1.5 Flash 為 100K 條樣本添加 PASTA-LANG 標記,在樣本回答中插入
標記,創建 PASTA 微調數據集。團隊隨后對 Gemma 7B 進行監督微調,得到能插入 PASTA-LANG 標記的 PASTA-SFT 模型。
第二階段:偏好優化。為優化標注策略,團隊設計了策略學習方案。團隊對每個樣本從 PASTA-SFT 模型采樣多種標注方案,然后基于兩項指標評估這些方案:理論加速比和內容質量(由 Gemini 1.5 Pro 評估)。根據評估結果,團隊構建「拒絕采樣數據集」,該數據集包含每個輸入的最佳和最差標注方案。最后,團隊用 BoNBoN 算法對 PASTA-SFT 模型進行偏好優化,得到最終的 PASTA 模型。
PASTA 推理系統:并行生成與緩存管理
推理系統設計難點。異步并行生成的主要挑戰在于如何協調多個線程高效協作。傳統方法通常需要為每個線程創建獨立的 KV 緩存池——創建新線程時必須復制主線程的前綴內容到子線程緩存池,完成后再復制結果回主線程。這兩次大規模矩陣復制操作嚴重限制了系統性能,使理論加速難以轉化為實際收益。
KV 緩存的存儲布局。PASTA 設計了交錯式 KV 緩存布局,所有線程共享單一連續內存池。系統初始以連續方式存儲用戶輸入,在推理過程中動態將不同線程在同一時間點生成的 token 交錯存儲在相鄰位置。
注意力控制與位置編碼。PASTA 通過兩個機制確保大模型正確理解多線程交錯存儲的 KV 緩存:
- 注意力掩碼控制:限制子線程只能訪問與自己相關的內容,在
- 后通過移除掩碼使主線程能訪問所有子線程生成的內容。
- 位置編碼調整:每個線程都使用獨立且連續的位置編碼,使線程處理自己的內容時,將交錯存儲的內容視為邏輯上連續的序列,確保模型能正確理解上下文。
這些設計共同確保 PASTA 能在提高速度的同時保持輸出質量。
實驗結果:Pareto 最優與可擴展性
PASTA 在性能與質量的平衡上取得了突破性成果,實驗結果表明它不僅實現了顯著加速,還在某些情況下提高了輸出質量。研究團隊在 AlpacaEval 基準上進行了全面評估,該基準包含 805 個具有代表性的指令跟隨任務。
質量-速度平衡的 Pareto 前沿。如圖所示,PASTA 通過調節質量權重參數生成了一系列的模型。在不同的生成質量的情況下,PASTA 均能提供非常可觀的加速。結果顯示,即使最注重質量的 PASTA 模型也能提供顯著加速,而最快的模型則以一定的質量犧牲換取接近 2 倍的速度提升。與基于手動設計的異步生成方案(Skeleton-of-Thought, APAR)相比,PASTA 模型展現出全面優勢。
可擴展性。研究結果展示了 PASTA 方法出色的可擴展性,如圖所示。隨著偏好優化不斷推進,PASTA 模型的性能持續提升。圖中清晰展示了從第一輪開始到第一輪結束,再到第二輪開始和第二輪后半程的整個優化過程,質量-速度的 Pareto 前沿大體持續向右上方推進。
這種穩定的改進趨勢表明,PASTA 方法具有良好的可擴展性——隨著投入更多計算資源,仍未飽和。與傳統依賴固定規則的異步解碼方法不同,PASTA 通過策略學習驅動的訓練算法提供了可持續的優化路徑,能夠有效地將額外計算資源轉化為更高的推理效率。
總結與展望
PASTA 首次證明,通過策略學習讓 LLM 自主優化生成策略,能夠突破傳統自回歸和基于規則的異步生成的效率極限。這一工作不僅為實時大模型應用提供了實用加速方案,更印證了未來 LLM 可能具備推理時自我優化能力的發展方向。
特別聲明:以上內容(如有圖片或視頻亦包括在內)為自媒體平臺“網易號”用戶上傳并發布,本平臺僅提供信息存儲服務。
Notice: The content above (including the pictures and videos if any) is uploaded and posted by a user of NetEase Hao, which is a social media platform and only provides information storage services.