FlashAR:僅用0.05%數據,讓預訓練好的自回歸圖像模型飛起來
![]()
背景:自回歸圖像生成的崛起與推理瓶頸
大語言模型的成功讓 "next-token prediction" 這套范式從文本延伸到了圖像領域。把圖像用視覺分詞器編碼成離散 token,用數預訓再一個接一個的練好投票拉票價格預測出來 —— 這就是自回歸(AR)圖像生成的核心思路。從早期的回歸 PixelCNN、iGPT、圖像Parti,模型到近期的飛起 Emu3.5、LlamaGen、僅據讓Lumina-mGPT、用數預訓GLM-Image,練好AR 模型的回歸生成質量已經全面逼近甚至超過了擴散模型。
但 AR 模型有個繞不開的圖像問題:慢。標準的模型 raster-scan 解碼從左到右、從上到下,飛起投票拉票價格一步只出一個 token。僅據讓生成一張 512×512 的圖要走 32×32 = 1024 步串行前向傳播,單卡耗時超過兩分鐘。延遲隨分辨率線性增長,GPU 的并行算力也用不起來 —— 高分辨率和實時場景下,基本沒法實際部署。
現有加速方案
為了突破這一瓶頸,研究者們已經探索了多種加速策略,但都面臨不同的局限:
- 重新設計生成范式:例如如 VAR 的 "下一尺度預測"、NAR 的 "近鄰預測"、PAR 的分組并行解碼,雖然能大幅降低解碼步數,但這些方法需要從頭預訓練,無法復用已有的大規模預訓練 AR 模型,訓練成本高昂。
- 離散擴散適配:例如 Emu3.5 原文采用的 DiDA,他們通過后訓練將 AR 模型改造為支持并行解碼的離散擴散模型。但這種方法改變了原始的預測目標,引入了預訓練和推理之間的不一致性,往往導致生成質量顯著下降,在我們的復現實驗中,相同數據量,Emu3.5 在 block diffusion 后訓練過程中,geneval 分數會在總體會有一個比較大的 drop
- 推測解碼:作為一種無需訓練的加速插件,實際加速效果受限于草稿模型的接受率,提升效果相對比較有限。
這就引出了一個關鍵的開放性問題:能否在不從頭訓練、不改變原始預測目標的前提下,將已有的預訓練 AR 模型改造成高度并行的生成器,同時繼承其強大的生成能力?
來自浙江大學和阿德萊德大學的研究團隊提出了FlashAR—— 一個輕量級的后訓練加速框架。不需要從頭訓練,在 Emu3.5-Image-34B 模型上,僅用原始訓練數據的 0.05%(約 8 萬張圖片),就能將預訓練好的自回歸模型改造成高度并行的生成器 Emu3.5-34B-Flash,實現最高22.9 倍的端到端加速。
![]()
![]()
- 論文標題:FlashAR: Efficient Post-Training Acceleration for Autoregressive Image Generation
- 論文主頁:https://lxazjk.github.io/FlashAR/
- 論文鏈接:https://arxiv.org/abs/2605.09430
- 代碼鏈接:https://github.com/lxazjk/Emu3.5-FlashAR
核心思路:從 "逐個生成" 到 "對角線并行"
傳統的自回歸圖像生成模型遵循嚴格的光柵掃描順序 —— 從左到右、從上到下,每一步只預測水平方向的下一個 token。對于一張由 H×W 個 token 構成的圖像,需要 H×W 步才能完成生成。
FlashAR 的關鍵洞察在于:圖像天然具有 2D 結構,如果我們為模型新增垂直方向的下一個 token 的預測能力,在每個步驟中,水平解碼頭和垂直解碼頭并行工作,解碼步數從 H×W 驟降至 H+W-1。以 512×512 分辨率(16×16 下采樣倍率)為例,解碼步數從 1024 步直接降到 63 步。
但要讓一個已經訓練好的 "水平方向" 預測模型具備 "垂直方向" 預測能力,并不容易。FlashAR 為此設計了三個關鍵組件:
![]()
1. 中間層分支(Intermediate Branching)
FlashAR 沒有把輕量級的 Vertical Head 接在預訓練模型的最終層,而是從中間層分出一條支路,讓它和原有的水平預測頭并行工作。
為什么不直接用最終層?因為經過完整訓練后,最終層特征已經更偏向原本的水平方向光柵預測任務,針對這個目標做了充分適配,但也因此不一定適合再拿來做垂直方向預測。相比之下,中間層往往還保留著更豐富的二維空間信息,用來適配新的預測方向會更自然。
這樣的設計還有一個額外好處:從中間層分支之后,Vertical Head 可以和原有分支并行執行,從而為整體吞吐帶來提升。
我們也在消融實驗里驗證了這一點。具體來說,我們用 linear probing 系統評估了預訓練模型不同層的特征,結果發現,最終層特征并不是最適合做垂直預測的。這也進一步支持了我們從中間層引出 Vertical Head 的設計。
![]()
2. 可學習融合門(Learnable Fusion Gate)
水平和垂直方向的預測分別建模了互補的空間依賴,其貢獻在不同空間位置上并不一致?;谶@一觀察,FlashAR 引入了一個輕量級的 MLP 融合門,在逐位置的粒度上自適應地融合兩個方向的預測結果,以避免簡單平均所導致的預測模糊。
3. 兩階段適配訓練(Two-Stage Adaptation)
具體而言,訓練過程分為兩個階段:
- 在第一階段,凍結骨干網絡,僅優化垂直預測頭,以使其快速學習到有意義的預測能力;
- 在第二階段,進一步聯合微調垂直預測頭和骨干網絡,使模型更好地適配新的解碼范式。這樣的漸進式訓練策略提升了后訓練過程的穩定性,并提高了數據利用效率。
在推理階段,FlashAR 還部署了硬件感知的推理優化管線:利用 FlexAttention 動態編譯稀疏的二維近鄰注意力掩碼,配合批量化 KV 緩存更新,將理論上的并行性切實轉化為真實的加速效果。
實驗結果
Emu3.5-Image-34B 加速
將 FlashAR 擴展到 340 億參數的 Emu3.5 模型上,是對框架能力的嚴格考驗:
![]()
用0.05%的原始訓練數據(80M token,約 8 萬張圖片),FlashAR 將 512×512 圖像生成速度從 130.10 秒壓縮到 5.68 秒,實現22.9 倍加速。更關鍵的是,加速幾乎不損失質量。在 GenEval 基準上,FlashAR 的 GenEval 總分僅下降 0.19 分(80.48→80.29),在顏色(+1.59)和位置(+7.00)兩個子項上甚至超過了原始模型。相比之下,BlockDiffusion 在相同設置下性能大幅下降至 73.83。
ImageNet 類別條件生成
在 ImageNet 256×256 基準上,FlashAR 在四個模型規模(B/L/XL/XXL)上全面超越現有后訓練方法 BlockDiffusion。
![]()
值得注意的是:
- FlashAR-L 的 IS(289.0)甚至超過了從頭訓練的 NAR-L(263.9),而 FlashAR 僅需輕量級后訓練;
- FlashAR-B 達到 447.2 img/s 的吞吐量,超過了 NAR-B(419.7 img/s);
- FlashAR 僅需25 個 epoch的后訓練 —— 只有 BlockDiffusion 訓練量的三分之一。
為什么 FlashAR 如此高效?
我們總結了 FlashAR 的核心優勢
- 無需從頭訓練:直接復用現有預訓練 AR 模型,通過輕量后訓練實現加速;
- 數據極致高效:僅需 0.05% 的原始訓練數據;
- 性能保持優異:生成質量幾乎無損,部分指標甚至提升;
- 框架通用性強:在 LlamaGen(120M~1.4B)和 Emu3.5(34B)上均驗證有效;
- 實際加速顯著:最高 22.9 倍端到端加速。
FlashAR 證明了一個重要觀點:通過精心設計的后訓練適配,可以在幾乎不改變原始模型訓練目標的前提下,將自回歸模型改造成高度并行的生成器,完整繼承預訓練模型的強大能力。