ICML（International Conference on Machine Learning，國際機器學習大會）是機器學習領域最具影響力和學術聲譽的國際頂級會議之一，也是人工智能研究領域的重要學術交流平臺。會議長期聚焦機器學習的基礎理論、核心算法、系統(tǒng)架構與實際應用，涵蓋深度學習、強化學習、機器學習系統(tǒng)以及跨學科應用等前沿方向。PKU-DAIR實驗室的論文《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

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

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

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

圖1：大語言模型強化學習Rollout階段軌跡長度的長尾分布

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

圖2：DARTS 方法示意圖：針對不同 prompt 的自適應采樣策略與冗余分配

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

1. 分布感知軌跡采樣

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

這種設計對應兩類典型情形：對于“冗長且無效”的長尾，DARTS 會增強其訓練影響，引導模型更簡潔；對于“復雜且必要”的長推理，DARTS 仍會保留有效長軌跡，使模型不會因為追求短輸出而損害能力。

2. 自適應冗余分配

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

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

圖3：軌跡響應長度方差越高，長尾特征越明顯

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

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

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

#3 實 驗

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

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

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

圖5：DARTS通過對rollout軌跡響應長度的主動塑造實現(xiàn)強化學習的有效加速

#4 總 結

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

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

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

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

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

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

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

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

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

本文在 Wan2.1-1.3B 和 CogVideoX1.5-5B 兩個代表性視頻 DiT 上與 VSA、SLA、SpargeAttn 等最優(yōu)基線進行全面對比，EchoAttention 在 Wan2.1-1.3B 上實現(xiàn) 1.97× 端到端加速，在 CogVideoX1.5-5B 上實現(xiàn) 2.42× 加速，均為所有方法中最高，且 VBench 分數(shù)與全注意力持平、PSNR/SSIM 指標顯著優(yōu)于所有稀疏基線，驗證了其在突破速度—質(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）在長文檔問答、長篇摘要和倉庫級代碼補全等復雜任務中的廣泛應用，推理時上下文窗口正在被不斷拉長。然而，LLM 的自注意力模塊計算復雜度與序列長度呈二次關系，在預填充階段（prefilling）成為了長上下文推理的主要時延瓶頸。

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

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

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

#2 方法

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

圖2：SALE 的整體工作流

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

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

圖3：注意力權重分布示意圖

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

#3 實驗

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

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

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

#4 總結

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

實驗室簡介

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