背景:自回归图像生成的崛起与推理瓶颈

大语言模型的成功让 "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 的并行算力也用不起来 —— 高分辨率和实时场景下,基本没法实际部署。

现有加速方案

为了突破这一瓶颈,研究者们已经探索了多种加速策略,但都面临不同的局限:

这就引出了一个关键的开放性问题:能否在不从头训练、不改变原始预测目标的前提下,将已有的预训练 AR 模型改造成高度并行的生成器,同时继承其强大的生成能力?

来自浙江大学和阿德莱德大学的研究团队提出了FlashAR—— 一个轻量级的后训练加速框架。不需要从头训练,在 Emu3.5-Image-34B 模型上,仅用原始训练数据的 0.05%(约 8 万张图片),就能将预训练好的自回归模型改造成高度并行的生成器 Emu3.5-34B-Flash,实现最高22.9 倍的端到端加速。





核心思路:从 "逐个生成" 到 "对角线并行"

传统的自回归图像生成模型遵循严格的光栅扫描顺序 —— 从左到右、从上到下,每一步只预测水平方向的下一个 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 如此高效?

我们总结了 FlashAR 的核心优势

FlashAR 证明了一个重要观点:通过精心设计的后训练适配,可以在几乎不改变原始模型训练目标的前提下,将自回归模型改造成高度并行的生成器,完整继承预训练模型的强大能力。