ICML（International Conference on Machine Learning，國際機(jī)器學(xué)習(xí)大會(huì)）是機(jī)器學(xué)習(xí)領(lǐng)域最具影響力和學(xué)術(shù)聲譽(yù)的國際頂級(jí)會(huì)議之一，也是人工智能研究領(lǐng)域的重要學(xué)術(shù)交流平臺(tái)。會(huì)議長(zhǎng)期聚焦機(jī)器學(xué)習(xí)的基礎(chǔ)理論、核心算法、系統(tǒng)架構(gòu)與實(shí)際應(yīng)用，涵蓋深度學(xué)習(xí)、強(qiáng)化學(xué)習(xí)、機(jī)器學(xué)習(xí)系統(tǒng)以及跨學(xué)科應(yīng)用等前沿方向。PKU-DAIR實(shí)驗(yàn)室的論文《DARTS: Distribution-Aware Active Rollout Trajectory Shaping for Accelerating LLM Reinforcement Learning》、《EchoAttention: Exploiting Token-Pair Redundancy and Frame-Block Similarity for Efficient Long Video Generation》和《SALE:Low-bit Estimation for Efficient Sparse Attention in Long-context LLM Prefilling》 被 ICML 2026 錄用。

DARTS: Distribution-Aware Active Rollout Trajectory Shaping for Accelerating LLM Reinforcement Learning

作者：Yujie Wang, Siwei Chen, Longzan Luo, Xinyi Liu, Xupeng Miao, Fangcheng Fu, Bin Cui

強(qiáng)化學(xué)習(xí)（Reinforcement Learining）已經(jīng)成為訓(xùn)練高性能大語言模型的重要范式。與傳統(tǒng)監(jiān)督學(xué)習(xí)不同，LLM RL 通常需要模型針對(duì)每個(gè) prompt 生成多條響應(yīng)軌跡，并基于獎(jiǎng)勵(lì)信號(hào)更新策略。整個(gè)流程主要包括 rollout 與 training 兩個(gè)階段，其中 rollout 階段需要大量采樣，往往成為訓(xùn)練系統(tǒng)的主要瓶頸。

LLM rollout 的響應(yīng)長(zhǎng)度呈現(xiàn)嚴(yán)重長(zhǎng)尾分布：少數(shù) prompt 會(huì)生成遠(yuǎn)長(zhǎng)于平均水平的軌跡，使得同步 RL 系統(tǒng)必須等待最慢樣本完成，導(dǎo)致批次阻塞和 GPU 利用率下降。現(xiàn)有方法多采用 prompt 級(jí)別的長(zhǎng)尾調(diào)度，例如將未完成的長(zhǎng)軌跡延后處理或截?cái)嗬m(xù)跑。這類方法能夠緩解等待問題，但本質(zhì)上仍是在“繞開長(zhǎng)尾”，并未改變模型自身的生成長(zhǎng)度分布。

長(zhǎng)尾不僅存在于不同 prompt 之間，也廣泛存在于同一個(gè) prompt 的多條 rollout 內(nèi)部，即 intra-prompt long-tail distribution。這說明長(zhǎng)尾并不只是 prompt 難度差異帶來的現(xiàn)象，而是模型生成分布本身的固有問題。更重要的是，其中相當(dāng)一部分長(zhǎng)軌跡并不帶來更好的獎(jiǎng)勵(lì)或推理質(zhì)量，而是冗長(zhǎng)、低效甚至錯(cuò)誤的生成過程。

圖1：大語言模型強(qiáng)化學(xué)習(xí)Rollout階段軌跡長(zhǎng)度的長(zhǎng)尾分布

因此，如何從源頭上減少無效長(zhǎng)尾，同時(shí)保留必要的深度推理軌跡，成為高效 LLM RL 訓(xùn)練亟需解決的問題。

圖2：DARTS 方法示意圖：針對(duì)不同 prompt 的自適應(yīng)采樣策略與冗余分配

我們提出 DARTS，一個(gè)面向 LLM 強(qiáng)化學(xué)習(xí)訓(xùn)練的高效 rollout 框架。其核心思想是 主動(dòng)分布塑形（active distribution shaping）：不再僅僅等待或調(diào)度長(zhǎng)尾軌跡，而是通過分布感知的采樣與資源分配機(jī)制，將模型 rollout 分布逐步塑造成更加簡(jiǎn)潔、穩(wěn)定且高效的形式。DARTS 主要包含三個(gè)關(guān)鍵組件。

1. 分布感知軌跡采樣

DARTS 首先為每個(gè) prompt 構(gòu)造更大的 intra-prompt 候選軌跡池，即通過冗余 rollout 生成多于訓(xùn)練所需數(shù)量的響應(yīng)。隨后，DARTS 采用 dual-end length sampling：一方面選擇最短的一部分軌跡，以鼓勵(lì)模型學(xué)習(xí)簡(jiǎn)潔、直接的解答；另一方面保留少量最長(zhǎng)但有效的軌跡，以避免過度壓縮推理過程，保留必要的深度思考能力。

這種設(shè)計(jì)對(duì)應(yīng)兩類典型情形：對(duì)于“冗長(zhǎng)且無效”的長(zhǎng)尾，DARTS 會(huì)增強(qiáng)其訓(xùn)練影響，引導(dǎo)模型更簡(jiǎn)潔；對(duì)于“復(fù)雜且必要”的長(zhǎng)推理，DARTS 仍會(huì)保留有效長(zhǎng)軌跡，使模型不會(huì)因?yàn)樽非蠖梯敵龆鴵p害能力。

2. 自適應(yīng)冗余分配

并不是所有 prompt 都需要同等規(guī)模的冗余采樣。對(duì)于響應(yīng)長(zhǎng)度分布穩(wěn)定、模型較有把握的問題，過多采樣會(huì)浪費(fèi)計(jì)算；而對(duì)于高方差、長(zhǎng)尾嚴(yán)重的問題，更大的候選空間有助于更有效地塑形分布。

因此，DARTS 提出 variance-based adaptive redundancy allocation，利用歷史響應(yīng)長(zhǎng)度方差來衡量 prompt 的長(zhǎng)尾程度與模型不確定性，并將更多采樣預(yù)算分配給高方差 prompt。這樣，系統(tǒng)能夠在有限計(jì)算預(yù)算下，把資源投入到最需要分布塑形的樣本上，從而同時(shí)提升算法效果與系統(tǒng)效率。

圖3：軌跡響應(yīng)長(zhǎng)度方差越高，長(zhǎng)尾特征越明顯

3. 系統(tǒng)級(jí)優(yōu)化

在系統(tǒng)實(shí)現(xiàn)上，DARTS 進(jìn)一步引入了兩類優(yōu)化。第一是 variance-guided tail pruning：當(dāng)某些 prompt 被識(shí)別為極端長(zhǎng)尾時(shí)，系統(tǒng)切換為 shortest-only sampling，并結(jié)合主動(dòng) early stopping，在收集到足夠有效軌跡后提前終止剩余超長(zhǎng)生成，避免極端尾部拖慢整體訓(xùn)練。

第二是 token-level streaming：不同于傳統(tǒng) sample-level streaming 必須等待整條軌跡生成完畢后再送入訓(xùn)練，DARTS 將生成結(jié)果按 token chunk 細(xì)粒度流式傳輸?shù)接?xùn)練端，使 rollout 與訓(xùn)練能夠更充分重疊，進(jìn)一步減少長(zhǎng)尾場(chǎng)景下的等待時(shí)間。

#3 實(shí) 驗(yàn)

我們?cè)?4張 NVIDIA H20 96GB集群上進(jìn)行實(shí)驗(yàn)。我們使用Qwen2.5-3B、Qwen2.5-Math-7B、Qwen2.5-14B、Qwen2.5-32B 以及 Qwen3-30B-A3B，并在 DAPO-MATH 與 MATH-lighteval 數(shù)據(jù)集上進(jìn)行了廣泛的實(shí)驗(yàn)。結(jié)果表明，DARTS 在不同模型規(guī)模下均穩(wěn)定優(yōu)于現(xiàn)有系統(tǒng)。相較于 VeRL，DARTS 實(shí)現(xiàn)了 1.29×–1.77× 的吞吐量提升；相較于代表性 prompt-level tail scheduling 方法 Tail Batching，最高也可達(dá)到 1.43× 的加速。實(shí)驗(yàn)說明，DARTS 不僅能緩解系統(tǒng)等待，更能通過主動(dòng)塑形 rollout 分布，從源頭上降低長(zhǎng)尾開銷。

圖4：不同模型規(guī)模下的端到端吞吐量加速對(duì)比

Case study表明，DARTS 能顯著壓縮rollout 長(zhǎng)度分布，使其集中到更短、更緊湊的區(qū)域；對(duì)于冗長(zhǎng)無效型 prompt，DARTS 可以有效消除 verbose tails；對(duì)于需要復(fù)雜推理的 prompt，DARTS 仍能保留正確長(zhǎng)軌跡所需的推理深度。

圖5：DARTS通過對(duì)rollout軌跡響應(yīng)長(zhǎng)度的主動(dòng)塑造實(shí)現(xiàn)強(qiáng)化學(xué)習(xí)的有效加速

#4 總 結(jié)

本文提出了 DARTS，一個(gè)通過 Distribution-Aware Active Rollout Trajectory Shaping 加速大語言模型強(qiáng)化學(xué)習(xí)訓(xùn)練的新框架。不同于已有方法主要從 prompt 級(jí)別調(diào)度長(zhǎng)尾，DARTS 從 rollout 分布本身出發(fā)，識(shí)別并優(yōu)化 intra-prompt 長(zhǎng)尾問題，通過分布感知軌跡采樣、自適應(yīng)冗余分配以及系統(tǒng)級(jí)流式優(yōu)化，引導(dǎo)模型生成更加簡(jiǎn)潔、穩(wěn)定且高效的響應(yīng)軌跡。

實(shí)驗(yàn)結(jié)果表明，DARTS 在多種模型規(guī)模和數(shù)據(jù)集上均取得顯著加速，最高實(shí)現(xiàn) 1.77× 的端到端吞吐量提升，同時(shí)保持模型收斂與下游任務(wù)表現(xiàn)不下降。DARTS 為大語言模型強(qiáng)化學(xué)習(xí)訓(xùn)練提供了一種新的系統(tǒng)優(yōu)化思路：不僅要調(diào)度長(zhǎng)尾，更要主動(dòng)塑造產(chǎn)生長(zhǎng)尾的分布本身。

EchoAttention: Exploiting Token-Pair Redundancy and Frame-Block Similarity for Efficient Long Video Generation

作者：Yifei Xia, Fangcheng Fu, Hao Yuan, Suhan Ling, Xupeng Miao, Huixia Li, Yuxi Ren, Xin Xia, Xuefeng Xiao, Bin Cui

近年來，基于擴(kuò)散 Transformer（DiT）的視頻生成模型（如 Wan2.1、CogVideoX 等）發(fā)展迅猛，但長(zhǎng)視頻推理的計(jì)算瓶頸始終制約其實(shí)用化：3D 全注意力（Full Attention）的二次方計(jì)算開銷 O(n²?²d) 可占據(jù)整體推理算力的 60%～82%?，F(xiàn)有稀疏注意力方法通過剪枝 Query-Key 對(duì)來緩解這一瓶頸，但其效果受制于大量"非稀疏 Head"——這些 Head 的注意力質(zhì)量無法被有效稀疏化，導(dǎo)致純稀疏方法存在難以逾越的速度—質(zhì)量天花板。

圖1：Token 對(duì)冗余性（稀疏注意力）與幀塊相似性（Echo 算子）的對(duì)比示意

為突破上述瓶頸，本文發(fā)現(xiàn)了視頻 DiT 中一種此前被忽視的規(guī)律——幀塊相似性（Frame-Block Similarity）：注意力權(quán)重矩陣中同一塊行（或塊列）內(nèi)的各幀塊分布高度相似，且差異可由輕量對(duì)角線性標(biāo)定精確刻畫（圖2）。如圖1右側(cè)所示，只需為每行計(jì)算一個(gè)原型塊，其余塊通過對(duì)角矩陣標(biāo)定恢復(fù)，即可將主導(dǎo)計(jì)算復(fù)雜度從 O(n²?²) 降至 O(n?²)，為非稀疏 Head 的加速提供了全新手段。

圖2：Wan2.1-1.3B 中幀塊相似性的可視化。同一塊行/列內(nèi)各幀塊分布高度一致，行方向 Pearson 相關(guān)系數(shù)接近 1、rMSE 普遍低于 0.1，說明輕量線性標(biāo)定即可精確重建

圖3：EchoAttention 整體架構(gòu)與三階段蒸餾流程。（a）雙算子架構(gòu)及 Echo-Row 三步算法；（b）軟混合 → 直通估計(jì) → 硬路由的三階段訓(xùn)練策略

基于上述觀察，本文提出 EchoAttention（圖3），一個(gè)雙算子注意力加速框架。它包含 Sparse 算子（處理高稀疏 Head）與 Echo 算子（Echo-Row / Echo-Col，處理幀塊相似 Head），并通過兩級(jí)可學(xué)習(xí)門控路由器，在去噪時(shí)間步—層—Head 三個(gè)維度上自動(dòng)選擇最優(yōu)算子。為使路由穩(wěn)定收斂，本文采用"軟混合→直通估計(jì)→硬路由"三階段蒸餾策略，配合自定義 Triton 算子實(shí)現(xiàn)真正的端到端推理加速。

本文在 Wan2.1-1.3B 和 CogVideoX1.5-5B 兩個(gè)代表性視頻 DiT 上與 VSA、SLA、SpargeAttn 等最優(yōu)基線進(jìn)行全面對(duì)比，EchoAttention 在 Wan2.1-1.3B 上實(shí)現(xiàn) 1.97× 端到端加速，在 CogVideoX1.5-5B 上實(shí)現(xiàn) 2.42× 加速，均為所有方法中最高，且 VBench 分?jǐn)?shù)與全注意力持平、PSNR/SSIM 指標(biāo)顯著優(yōu)于所有稀疏基線，驗(yàn)證了其在突破速度—質(zhì)量天花板方面的有效性。

SALE:Low-bit Estimation for Efficient Sparse Attention in Long-context LLM Prefilling

作者：Xiaodong Ji，Hailin Zhang，F(xiàn)angcheng Fu，Bin CUI

論文鏈接：https://arxiv.org/abs/2505.24179

#1 背景與挑戰(zhàn)

隨著大語言模型（LLMs）在長(zhǎng)文檔問答、長(zhǎng)篇摘要和倉庫級(jí)代碼補(bǔ)全等復(fù)雜任務(wù)中的廣泛應(yīng)用，推理時(shí)上下文窗口正在被不斷拉長(zhǎng)。然而，LLM 的自注意力模塊計(jì)算復(fù)雜度與序列長(zhǎng)度呈二次關(guān)系，在預(yù)填充階段（prefilling）成為了長(zhǎng)上下文推理的主要時(shí)延瓶頸。

已有研究發(fā)現(xiàn)，LLM 的注意力圖（Attention map）天然稀疏，只有少量區(qū)域?qū)敵霎a(chǎn)生重要貢獻(xiàn)?；谶@一觀察，學(xué)術(shù)界提出了一系列稀疏注意力（Sparse attention）方法，通過構(gòu)造稀疏掩碼跳過不重要區(qū)域的計(jì)算來加速prefilling。然而，如圖1中(a) (b)兩個(gè)實(shí)例所示，現(xiàn)有動(dòng)態(tài)稀疏注意力方法在"精度—效率權(quán)衡"上仍不理想：它們對(duì)注意力圖的檢查方式要么過于粗糙，要么不夠全面，導(dǎo)致無法準(zhǔn)確識(shí)別重要的注意力區(qū)域。

圖1：不同稀疏注意力方法對(duì)注意力圖的檢查方式對(duì)比。

本質(zhì)上，稀疏注意力方法在注意力圖的"檢查精度"與"額外算力開銷"之間存在固有的trade-off：越細(xì)粒度、越全覆蓋的檢查越能精準(zhǔn)識(shí)別出真正重要的區(qū)域，但需要付出更高的額外算力；而為了壓低額外開銷，現(xiàn)有方法不得不退化為采樣或粗粒度檢查這類近似手段，代價(jià)是漏檢或誤判重要元素。如何在不顯著增加開銷的前提下實(shí)現(xiàn)細(xì)粒度、全覆蓋的注意力圖檢查，是提升稀疏注意力精度—效率權(quán)衡的關(guān)鍵。

#2 方法

為解決這一問題，我們提出了 SALE（Sparse Attention via Low-bit Estimation），一種基于低比特估計(jì)的塊稀疏注意力方法，能夠在幾乎不損失模型精度的前提下，顯著加速 LLM 長(zhǎng)上下文預(yù)填充。SALE的核心設(shè)計(jì)是：用極低比特量化的 query-key 乘積快速近似整張注意力圖，再基于這一近似圖做細(xì)粒度的重要性篩選。圖2展示了SALE的整體工作流。

圖2：SALE 的整體工作流

在 SALE 的注意力圖檢查階段，我們將 Q、K 量化至 4-bit，并在此基礎(chǔ)上計(jì)算注意力圖上每個(gè)位置的近似注意力權(quán)重。該步驟利用了現(xiàn)代 GPU 上低比特 Tensor Core 指令的高吞吐特性，最高能達(dá)到FP16 Tensor Core指令吞吐量的8倍。

得到注意力權(quán)重之后，SALE 提出使用 Relative Attention Score（相對(duì)注意力分?jǐn)?shù)）來評(píng)估每個(gè) query-key 對(duì)的重要性。該指標(biāo)源自于被大量研究證實(shí)的現(xiàn)象（如圖3所示）：LLM每行注意力權(quán)重在 sink（序列開頭）和 local（序列末尾）區(qū)域普遍偏高?；诖耍琒ALE以當(dāng)前位置權(quán)重相對(duì)于 sink-local 區(qū)域權(quán)重的比值作為重要性判據(jù)。相較于直接使用原始注意力分?jǐn)?shù)進(jìn)行篩選，該指標(biāo)無需將整張注意力圖存儲(chǔ)到GPU顯存，計(jì)算/訪存開銷都更小。篩選結(jié)果以塊稀疏掩碼的形式傳遞給隨后的Attention計(jì)算階段。

圖3：注意力權(quán)重分布示意圖

考慮到不同注意力頭之間的稀疏率存在顯著差異，SALE 針對(duì)關(guān)鍵超參數(shù)進(jìn)行了逐頭離線校準(zhǔn)，使每個(gè)頭都能匹配到合適的稀疏度，所得校準(zhǔn)結(jié)果可在推理階段直接復(fù)用；此外，SALE 為注意力圖檢查階段開發(fā)了高性能的 GPU kernel，將該階段的開銷壓縮至全量注意力計(jì)算的11%左右。

#3 實(shí)驗(yàn)

圖4：SALE 與現(xiàn)有稀疏注意力方法的精度—加速比權(quán)衡對(duì)比

圖5：SALE 與現(xiàn)有稀疏注意力方法的精度—加速比權(quán)衡對(duì)比。橫軸為延遲，縱軸為精度得分

如圖4所示，SALE 在處理 64K 以上長(zhǎng)序列時(shí)相對(duì)全量注意力可獲得至少 3.36× 的加速，同時(shí)精度損失可忽略不計(jì)；并且如圖5所示，SALE在"精度—效率"權(quán)衡維度上全面優(yōu)于現(xiàn)有方法（例如MInference，F(xiàn)lexPrefill和SpargeAttn），證明了低比特估計(jì) + 細(xì)粒度檢查這一設(shè)計(jì)范式的有效性。

#4 總結(jié)

本文提出一種新穎的動(dòng)態(tài)稀疏注意力方法SALE，用于解決長(zhǎng)上下文大語言模型預(yù)填充階段稀疏注意力的精度-效率權(quán)衡問題。該方法將查詢（Q）和鍵（K）量化至4比特以快速近似注意力圖，利用低比特計(jì)算的高吞吐特性降低開銷；采用相對(duì)注意力分?jǐn)?shù)進(jìn)行細(xì)粒度重要區(qū)域篩選，結(jié)合逐頭離線校準(zhǔn)與高性能GPU內(nèi)核優(yōu)化。實(shí)驗(yàn)表明，該方法在幾乎不損失模型精度的前提下，實(shí)現(xiàn)了長(zhǎng)序列預(yù)填充推理的顯著加速。

實(shí)驗(yàn)室簡(jiǎn)介

北京大學(xué)數(shù)據(jù)與智能實(shí)驗(yàn)室（Data And Intelligence Research Lab at Peking Univeristy，PKU-DAIR實(shí)驗(yàn)室）由北京大學(xué)計(jì)算機(jī)學(xué)院崔斌教授領(lǐng)導(dǎo)，長(zhǎng)期從事數(shù)據(jù)庫系統(tǒng)、大數(shù)據(jù)管理與分析、人工智能等領(lǐng)域的前沿研究，在理論和技術(shù)創(chuàng)新以及系統(tǒng)研發(fā)上取得多項(xiàng)成果，已在國際頂級(jí)學(xué)術(shù)會(huì)議和期刊發(fā)表學(xué)術(shù)論文200余篇，發(fā)布多個(gè)開源項(xiàng)目。課題組同學(xué)曾數(shù)十次獲得包括CCF優(yōu)博、ACM中國優(yōu)博、北大優(yōu)博、微軟學(xué)者、蘋果獎(jiǎng)學(xué)金、谷歌獎(jiǎng)學(xué)金等榮譽(yù)。PKU-DAIR實(shí)驗(yàn)室持續(xù)與工業(yè)界展開卓有成效的合作，與騰訊、阿里巴巴、蘋果、微軟、百度、快手、中興通訊等多家知名企業(yè)開展項(xiàng)目合作和前沿探索，解決實(shí)際問題，進(jìn)行科研成果的轉(zhuǎn)化落地。