金天，麻省理工學(xué)院（MIT）計(jì)算機(jī)科學(xué)與人工智能實(shí)驗(yàn)室（CSAIL）博士五年級(jí)學(xué)生，師從 Michael Carbin 和 Jonathan Ragan-Kelley。他主要研究機(jī)器學(xué)習(xí)與編程系統(tǒng)的結(jié)合。此前曾在 IBM Research 主導(dǎo)實(shí)現(xiàn)深度神經(jīng)網(wǎng)絡(luò)在 IBM 主機(jī)上的推理部署。本科畢業(yè)于 Haverford College，獲計(jì)算機(jī)科學(xué)與數(shù)學(xué)雙學(xué)位。

鄭鈺熹，麻省理工學(xué)院 CSAIL 博士三年級(jí)學(xué)生，師從 Michael Carbin。她的研究方向?yàn)榫幊陶Z(yǔ)言與機(jī)器學(xué)習(xí)的交叉領(lǐng)域。

大語(yǔ)言模型（LLM）的生成范式正在從傳統(tǒng)的「單人書(shū)寫(xiě)」向「分身協(xié)作」轉(zhuǎn)變。傳統(tǒng)自回歸解碼按順序生成內(nèi)容，而新興的異步生成范式通過(guò)識(shí)別語(yǔ)義獨(dú)立的內(nèi)容塊，實(shí)現(xiàn)并行生成。

如圖所示，傳統(tǒng)方法（下）按順序生成所有內(nèi)容，而異步生成（上）同時(shí)處理多個(gè)互不依賴(lài)的內(nèi)容塊。對(duì)比順序生成，異步生成在 AlpacaEval 長(zhǎng)度控制評(píng)測(cè)中實(shí)現(xiàn) 1.21-1.93× 的幾何平均提速，對(duì)應(yīng)生成質(zhì)量變化（勝率）為 +2.2% 至 -7.1%。

MIT 與谷歌研究團(tuán)隊(duì)在最新研究 PASTA（PArallel STructure Annotation）中首次從策略學(xué)習(xí)（policy learning）角度探索異步生成范式的可能。

• 論文標(biāo)題：Learning to Keep a Promise: Scaling Language Model Decoding Parallelism with Learned Asynchronous Decoding
• 論文地址：https://arxiv.org/abs/2502.11517

研究團(tuán)隊(duì)不依賴(lài)人工設(shè)計(jì)規(guī)則來(lái)識(shí)別異步生成機(jī)會(huì)，而通過(guò)策略學(xué)習(xí)讓模型自主發(fā)現(xiàn)并標(biāo)注這些機(jī)會(huì)，系統(tǒng)地優(yōu)化質(zhì)量與速度的平衡。這種方法使 LLM 根據(jù)內(nèi)容特點(diǎn)自適應(yīng)地確定最佳異步生成策略，為生成效率優(yōu)化開(kāi)創(chuàng)學(xué)習(xí)驅(qū)動(dòng)的全新路徑。

PASTA-LANG：劃分獨(dú)立內(nèi)容的標(biāo)記語(yǔ)言

研究人員首先開(kāi)發(fā)了一種新的標(biāo)記語(yǔ)言 PASTA-LANG，專(zhuān)為異步生成而設(shè)計(jì)。大模型使用它在生成過(guò)程中標(biāo)記語(yǔ)義獨(dú)立塊，指示并行生成機(jī)會(huì)。這種語(yǔ)言包含三種核心標(biāo)記：

• <promise topic="..."/>：標(biāo)記語(yǔ)義獨(dú)立的內(nèi)容塊，通過(guò) topic 屬性總結(jié)內(nèi)容主題，大模型用它表明「這部分將會(huì)由一個(gè)獨(dú)立子線(xiàn)程異步生成」。
• <async>...</async>：在 <promise> 后標(biāo)識(shí)對(duì)應(yīng)的異步生成的內(nèi)容，表示這一部分由獨(dú)立子線(xiàn)程負(fù)責(zé)生成。
• <sync/>：在主線(xiàn)程標(biāo)記同步點(diǎn)，表明后續(xù)內(nèi)容生成將會(huì)依賴(lài)于之前的異步生成，主線(xiàn)程需等待所有異步線(xiàn)程完成后才能繼續(xù)。

這些標(biāo)記共同構(gòu)成了一種「承諾-履行」的生成模式：大模型首先通過(guò) <promise/> 標(biāo)記「承諾」生成某些內(nèi)容，推理系統(tǒng)再創(chuàng)建異步線(xiàn)程來(lái)「履行」這些承諾，最后在 <sync/> 處將異步內(nèi)容組合。

如圖所示的線(xiàn)段長(zhǎng)度計(jì)算案例展示了這一機(jī)制：面對(duì)計(jì)算任務(wù)（A），大模型首先識(shí)別出「坐標(biāo)提取」和「長(zhǎng)度公式」兩個(gè)可并行的獨(dú)立部分，生成相應(yīng)的 <promise/> 標(biāo)記（B），隨后用 <sync/> 標(biāo)記（E）表明需要等待這些內(nèi)容完成。圖中紅色和綠色區(qū)域（C、D）顯示了兩個(gè)異步線(xiàn)程并行生成的內(nèi)容，最終在（F）處組合成完整解答。

這個(gè)新的標(biāo)記語(yǔ)言簡(jiǎn)單，可擴(kuò)展性強(qiáng)，開(kāi)啟了新的未來(lái)研究范式。

PASTA 訓(xùn)練：從標(biāo)注到優(yōu)化的雙階段學(xué)習(xí)

如圖所示，PASTA 系統(tǒng)采用雙階段訓(xùn)練流程，使大模型自主學(xué)習(xí)使用上述標(biāo)記語(yǔ)言，完成異步生成。

第一階段：監(jiān)督微調(diào)。研究團(tuán)隊(duì)首先選取 SlimOrca 指令跟隨數(shù)據(jù)集，用 Gemini 1.5 Flash 為 100K 條樣本添加 PASTA-LANG 標(biāo)記，在樣本回答中插入 <promise/>、<async> 和 <sync/> 標(biāo)記，創(chuàng)建 PASTA 微調(diào)數(shù)據(jù)集。團(tuán)隊(duì)隨后對(duì) Gemma 7B 進(jìn)行監(jiān)督微調(diào)，得到能插入 PASTA-LANG 標(biāo)記的 PASTA-SFT 模型。

第二階段：偏好優(yōu)化。為優(yōu)化標(biāo)注策略，團(tuán)隊(duì)設(shè)計(jì)了策略學(xué)習(xí)方案。團(tuán)隊(duì)對(duì)每個(gè)樣本從 PASTA-SFT 模型采樣多種標(biāo)注方案，然后基于兩項(xiàng)指標(biāo)評(píng)估這些方案：理論加速比和內(nèi)容質(zhì)量（由 Gemini 1.5 Pro 評(píng)估）。根據(jù)評(píng)估結(jié)果，團(tuán)隊(duì)構(gòu)建「拒絕采樣數(shù)據(jù)集」，該數(shù)據(jù)集包含每個(gè)輸入的最佳和最差標(biāo)注方案。最后，團(tuán)隊(duì)用 BoNBoN 算法對(duì) PASTA-SFT 模型進(jìn)行偏好優(yōu)化，得到最終的 PASTA 模型。

PASTA 推理系統(tǒng)：并行生成與緩存管理

推理系統(tǒng)設(shè)計(jì)難點(diǎn)。異步并行生成的主要挑戰(zhàn)在于如何協(xié)調(diào)多個(gè)線(xiàn)程高效協(xié)作。傳統(tǒng)方法通常需要為每個(gè)線(xiàn)程創(chuàng)建獨(dú)立的 KV 緩存池——?jiǎng)?chuàng)建新線(xiàn)程時(shí)必須復(fù)制主線(xiàn)程的前綴內(nèi)容到子線(xiàn)程緩存池，完成后再?gòu)?fù)制結(jié)果回主線(xiàn)程。這兩次大規(guī)模矩陣復(fù)制操作嚴(yán)重限制了系統(tǒng)性能，使理論加速難以轉(zhuǎn)化為實(shí)際收益。

KV 緩存的存儲(chǔ)布局。PASTA 設(shè)計(jì)了交錯(cuò)式 KV 緩存布局，所有線(xiàn)程共享單一連續(xù)內(nèi)存池。系統(tǒng)初始以連續(xù)方式存儲(chǔ)用戶(hù)輸入，在推理過(guò)程中動(dòng)態(tài)將不同線(xiàn)程在同一時(shí)間點(diǎn)生成的 token 交錯(cuò)存儲(chǔ)在相鄰位置。

注意力控制與位置編碼。PASTA 通過(guò)兩個(gè)機(jī)制確保大模型正確理解多線(xiàn)程交錯(cuò)存儲(chǔ)的 KV 緩存：

• 注意力掩碼控制：限制子線(xiàn)程只能訪(fǎng)問(wèn)與自己相關(guān)的內(nèi)容，在<sync/>后通過(guò)移除掩碼使主線(xiàn)程能訪(fǎng)問(wèn)所有子線(xiàn)程生成的內(nèi)容。
• 位置編碼調(diào)整：每個(gè)線(xiàn)程都使用獨(dú)立且連續(xù)的位置編碼，使線(xiàn)程處理自己的內(nèi)容時(shí)，將交錯(cuò)存儲(chǔ)的內(nèi)容視為邏輯上連續(xù)的序列，確保模型能正確理解上下文。

這些設(shè)計(jì)共同確保 PASTA 能在提高速度的同時(shí)保持輸出質(zhì)量。

實(shí)驗(yàn)結(jié)果：Pareto 最優(yōu)與可擴(kuò)展性

PASTA 在性能與質(zhì)量的平衡上取得了突破性成果，實(shí)驗(yàn)結(jié)果表明它不僅實(shí)現(xiàn)了顯著加速，還在某些情況下提高了輸出質(zhì)量。研究團(tuán)隊(duì)在 AlpacaEval 基準(zhǔn)上進(jìn)行了全面評(píng)估，該基準(zhǔn)包含 805 個(gè)具有代表性的指令跟隨任務(wù)。

質(zhì)量-速度平衡的 Pareto 前沿。如圖所示，PASTA 通過(guò)調(diào)節(jié)質(zhì)量權(quán)重參數(shù)生成了一系列的模型。在不同的生成質(zhì)量的情況下，PASTA 均能提供非常可觀(guān)的加速。結(jié)果顯示，即使最注重質(zhì)量的 PASTA 模型也能提供顯著加速，而最快的模型則以一定的質(zhì)量犧牲換取接近 2 倍的速度提升。與基于手動(dòng)設(shè)計(jì)的異步生成方案（Skeleton-of-Thought, APAR）相比，PASTA 模型展現(xiàn)出全面優(yōu)勢(shì)。

可擴(kuò)展性。研究結(jié)果展示了 PASTA 方法出色的可擴(kuò)展性，如圖所示。隨著偏好優(yōu)化不斷推進(jìn)，PASTA 模型的性能持續(xù)提升。圖中清晰展示了從第一輪開(kāi)始到第一輪結(jié)束，再到第二輪開(kāi)始和第二輪后半程的整個(gè)優(yōu)化過(guò)程，質(zhì)量-速度的 Pareto 前沿大體持續(xù)向右上方推進(jìn)。

這種穩(wěn)定的改進(jìn)趨勢(shì)表明，PASTA 方法具有良好的可擴(kuò)展性——隨著投入更多計(jì)算資源，仍未飽和。與傳統(tǒng)依賴(lài)固定規(guī)則的異步解碼方法不同，PASTA 通過(guò)策略學(xué)習(xí)驅(qū)動(dòng)的訓(xùn)練算法提供了可持續(xù)的優(yōu)化路徑，能夠有效地將額外計(jì)算資源轉(zhuǎn)化為更高的推理效率。

總結(jié)與展望

PASTA 首次證明，通過(guò)策略學(xué)習(xí)讓 LLM 自主優(yōu)化生成策略，能夠突破傳統(tǒng)自回歸和基于規(guī)則的異步生成的效率極限。這一工作不僅為實(shí)時(shí)大模型應(yīng)用提供了實(shí)用加速方案，更印證了未來(lái) LLM 可能具備推理時(shí)自我優(yōu)化能力的發(fā)展方向。