跳转至

MARché: Fast Masked Autoregressive Image Generation with Cache-Aware Attention

会议: ECCV 2026
论文: ECCV 2026 论文页
代码: 待开源
领域: 图像生成
关键词: 掩码自回归生成、KV缓存加速、缓存感知注意力、选择性刷新、训练无关加速

一句话总结

针对掩码自回归图像生成中每步重复计算全局注意力与前馈网络的严重冗余,提出无需重训的 MARché 框架,通过缓存感知解耦注意力、基于注意力的选择性 KV 刷新及周期校准,在几乎无损生成质量的前提下实现高达 1.79 倍推理加速。

研究背景与动机

掩码自回归模型(MAR)通过在连续 Token 空间引入双向注意力和扩散损失函数,成功结合了自回归生成的强建模能力与掩码并行预测的生成效率,成为高分辨率图像生成的前沿范式。然而,MAR 在推理过程中依赖多步解码(例如默认 64 步),并在每一个去掩码步对所有 Token(无论是尚未生成的掩码 Token 还是此前已生成的已知 Token)重复执行完整的双向自注意力与前馈网络(FFN)映射,导致极大的计算开销与延迟,严重制约了模型向高分辨率和更大参数规模的扩展。

深入分析 MAR 解码过程中的内部表征可以发现,在相邻去掩码步之间,大多数 Token 的 Key 和 Value 投影向量展现出极高的时序稳定性,相邻步间的余弦相似度常常超过 0.95。然而,自回归文本大模型中成熟的因果 KV 缓存机制无法直接套用到 MAR 中,因为 MAR 采用全局双向注意力,新生成 Token 会动态改变已有上下文的表征;如果直接冻结并无限期复用历史 KV 投影,会迅速引发严重的语义漂移并导致图像生成质量彻底崩塌。

解决这一瓶颈的关键切入点在于区分表征变化的敏感度,将动态演化与稳定背景解耦处理。核心 idea:利用新生成 Token 的跨步注意力显著性定位受上下文影响的关键 Token,构建以缓存感知双向解耦与选择性刷新为核心的训练无关解码框架,仅重算活跃子集并跳过稳定 Token 的 FFN 计算,以极小开销换取全量双向上下文的高效维持。

方法详解

整体框架

MARché 在不修改底层 Transformer 结构且完全无需重新训练的前提下,重构了 MAR 的单步解码流水线。模型在每一去掩码步将全序列 Token 划分为由“生成 Token(当前步待预测掩码)”、“缓存 Token(前一步刚生成的已知 Token)”和“刷新 Token(受新生成内容强烈影响的上下文 Token)”构成的活跃集(Active Set),以及其余保持稳定的缓存集(Cached Set)。在解码前两层执行全局注意力计算并提取跨步相关性分数以动态确立活跃集预算;从第三层起启用缓存感知双向注意力算子,将活跃 Query 对活跃 KV 与缓存 KV 的注意力独立评估并通过安全在线 Softmax(Safe Online Softmax)无缝拼接,同时跳过缓存 Token 的 FFN 映射;最后配合定期的全量缓存刷新阻断误差累积。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["输入序列<br/>掩码与已知 Token 特征"] --> B["基于生成注意力的选择性 KV 刷新<br/>前两层全局注意聚合生成 Query 权重,动态选出 Top-K 刷新 Token 填入活跃集"]
    B --> C{"是否触发周期性全量刷新?<br/>(步数 mod 3 == 0 或 初始步)"}
    C -->|是| D["周期性全量刷新与漂移校准<br/>全层重新投影并同步全局 KV 缓存"]
    C -->|否| E["双路解耦注意力与在线融合<br/>活跃 Query 分别交互活跃与缓存 KV,安全在线 Softmax 融合并跳过缓存 FFN"]
    D --> F["解码特征输出<br/>送入扩散采样器预测当前步掩码 Token"]
    E --> F

关键设计

1. 基于生成注意力的选择性 KV 刷新:精准捕捉受新生成内容动态影响的上下文子集

传统 KV 缓存压缩或稀疏策略多采用自上而下的丢弃逻辑,而 MARché 针对双向自回归的特性,反向利用注意力分数来定位“谁需要被重新计算”。在图像去掩码过程中,新生成的图像斑块会对其语义关联区域(例如同一物体的边缘、强纹理关联区域)产生剧烈的上下文更新,若不及时刷新这些位置的 KV,图像局部结构将出现错位。MARché 充分权衡层级深浅与表征对齐度(浅层相关性与深层重合度低,过深层重算开销大),选定解码器第 2 层作为决策锚点。在该层执行全局注意力时,提取当前步所有待生成 Token(Generating Tokens)对全局所有 Token 的注意力矩阵,在全部注意力头上求均值得到全局相关性得分,随后据此自适应选出相关度最高的 Top-\(K\) 个 Token 作为刷新 Token(Refreshing Tokens)。结合当前步生成 Token、上一轮新落盘的 Caching Token,共同组装进固定预算(如 64 个 Token)的活跃集 \(A^{(t)}\),在保持计算负载可控的同时保障了高动态上下文的即时更新。

2. 双路解耦注意力与在线融合:消除张量拼接开销并规避冗余前馈计算

针对划分出的活跃集 \(A^{(t)}\) 与缓存集 \(C^{(t)}\),传统实现若将活跃 KV 与缓存 KV 进行显式拼接再调用通用 Attention Kernel,会产生严重的非连续内存搬运与访存延迟。MARché 设计了解耦的双路径注意力计算流:活跃 Token 独立生成查询向量 \(q_i\)、键向量 \(K_A\) 与值向量 \(V_A\),而缓存 Token 仅提供只读的预存向量 \((K_C, V_C)\),完全不生成 Query 也不参与后续 FFN。对任意活跃 Query \(q_i\),注意力计算被显式解耦为对活跃集与缓存集的两路内积:

\[\alpha^{(A)}, \alpha^{(C)}, \ell_i = \text{SafeOnlineSoftmax}(q_i, K_A, K_C)\]
\[z_i = \frac{\alpha^{(A)} V_A + \alpha^{(C)} V_C}{\ell_i}\]

其中 Safe Online Softmax 维持两路局域最大值的动态归一化因子 \(\ell_i\),实现了与标准全局双向自注意力完全等价的输出精度,同时获得极优的 GPU 共享内存局部性。在注意力层之后,只有活跃向量 \(z_i\) 被送入 FFN 映射,其余处于缓存集的 Token 直接跳过 FFN 阶段,显著斩断了占据 Transformer 近三分之二参数量的密集线性变换耗时。

3. 周期性全量刷新与漂移校准:三道防护屏障彻底阻断长步数误差累积

尽管选择性刷新覆盖了强相关 Token,但连续数十步的局部缓存复用不可避免地会在残差连接与归一化层中沉淀微小的数值漂移,导致去噪后期细节边缘产生虚影或伪影。为此,MARché 设立了三道确定性的安全校准机制:第一,在初始步(Step 0),由于缓存尚为空,强制执行贯穿全部层的标准双向自注意力;第二,在任意解码步的第 1 与第 2 层均执行标准全量注意力,确保刷新 Token 的判定始终建立在未失真的新鲜浅层表征之上;第三,引入步长为 3 的周期性全量刷新机制(Periodic Full Refresh),每经历 3 个去掩码步便强制执行一次全量解码与 KV 覆写,以极低的平均计算摊销(仅增加微量计算)重置全局潜在漂移,从而在 64 步全流程中牢牢锁住图像生成质量的基线指标。

实验关键数据

主实验

在 ImageNet \(256 \times 256\) 条件图像生成基准下,对比在相同单张 NVIDIA H100 GPU 环境中的标准 64 步解码延迟、生成质量指标(FID 与 IS)以及模型参数规模。

模型 图像生成延迟 (s/im) ↓ FID ↓ Inception Score ↑ 参数量 加速比 ↑
MaskGIT 0.440 6.18 182.1 227M -
DiT-XL/2 0.787 2.27 278.2 675M -
LlamaGen-XXL 0.897 3.09 253.6 1.4B -
LlamaGen-3B 1.011 3.05 222.3 3.1B -
MAR-B (基线) 0.104 2.35 281.1 208M 1.00×
LazyMAR-B 0.074 5.32 235.1 208M 1.41×
MARché-B (本文) 0.064 2.56 270.3 208M 1.57×
MAR-L (基线) 0.193 1.84 296.3 479M 1.00×
LazyMAR-L 0.132 4.32 246.6 479M 1.46×
MARché-L (本文) 0.115 2.16 278.6 479M 1.68×
MAR-H (基线) 0.336 1.62 298.6 943M 1.00×
LazyMAR-H 0.220 4.00 251.9 943M 1.52×
MARché-H (本文) 0.188 2.02 281.4 943M 1.79×

消融实验

活跃集 Token 构建策略消融 检验在活跃集构造中各类 Token 组分不可或缺的实际影响。

活跃集构建配置 FID ↓ 说明
MARché (完整模型) 2.56 包含生成、缓存与高注意力刷新 Token
MARché 移除缓存 Token (w/o caching tokens) 2.70 仅使用生成 Token 与刷新 Token,略有掉点但无计算节省
MARché 移除生成 Token (w/o generating tokens) 504.82 缺失当期生成预测核心,语义崩塌失效
随机选取 Token 重算 (Random selection) 564.61 脱离注意力引导随机重算,模型彻底无法收敛解码

加速收益构成与计算分流消融 在 MAR-B 模型下解耦缓存感知注意力算子与 FFN 跳过对整体吞吐的边际贡献。

实验方案 实现配置 单图延迟 (s) 相对加速比
原生 MAR 标准注意力 + 全量 FFN 0.104 1.00×
MAR + 选择性 FFN 标准注意力 + FFN 跳过 0.092 1.12×
MAR + 缓存感知注意力 缓存感知注意力 + 全量 FFN 0.067 1.55×
MARché (完整模型) 缓存感知注意力 + FFN 跳过 0.064 1.63×

关键发现

  • 生成 Token 是不可动摇的刚性支柱:消融表明,若在重算集合中剔除当前步的生成 Token,FID 恶化至 504.82,直接导致解码彻底瓦解;而高注意力引导的刷新策略(FID 2.56)显著碾压低注意力策略(FID 3.80)与随机刷新(FID 3.01),证明双向依赖下注意力强度即为信息更新敏感度的直接度量。
  • 解耦注意力算子贡献绝大部分加速壁垒:从加速拆解表可见,单独启用缓存感知注意力即可带来 1.55 倍提速(延迟从 0.104s 骤降至 0.067s),而 FFN 跳过进一步贡献 1.12 倍独立增益,两相结合达成 1.63 倍复合提速;这表明显存局部性优化与规避大矩阵 KV 拼接同样是生成加速的关键瓶颈。
  • 层级决策权衡清晰:在第 2 层选取刷新 Token 达成 66.8% 的深层平均重合度与最优 Pareto 平衡;若在第 1 层决策虽延迟低至 0.154s 但重合率仅 49.3% 且 FID 掉至 2.62,而在第 4 层决策虽提升 FID 至 2.49 但由于前 4 层均需全量算力导致延迟升至 0.1613s。

亮点与洞察

  • 训练无关与零架构侵入性:完全不调整原始 MAR 模型的权重、拓扑或生成顺序,即插即用且完全保留了基于扩散连续空间的建模优势,避免了蒸馏或结构重训高昂的算力负担。
  • 反转注意力过滤思维的妙用:LLM 场景中注意力分数常用于驱逐长文本历史 KV,而 MARché 将其反向用于指引“谁需要被赋予计算权”,以精准的动态依赖感知打破了双向注意力全局全量重算的算力死锁。
  • 在线 Softmax 规避物理拼接:借鉴 FlashAttention 的在线流式归一化思想,将活跃与缓存两路内积得分在片上寄存器层面直接平滑融合成期望特征,消除了大显存搬运带来的吞吐反噬。

局限与展望

  • 作者承认的局限:方法依然需要在前两层执行全局注意力计算,且活跃集预算目前依赖人工预设的超参数经验值(如固定 64),尚未根据图像内容的几何复杂度动态伸缩活跃集大小。
  • 潜在改进方向:当前决策层固定在第 2 层,未来可探索根据不同分辨率或不同生成阶段(前期构图 vs 后期纹理细化)自适应切换刷新决策层级;此外,将此缓存感知注意力框架扩展到视频生成模型与多模态视觉-语言理解的自回归生成中同样具有巨大的工程探索潜力。

相关工作与启发

  • vs LazyMAR (Yan et al., 2025):LazyMAR 同样试图通过缓存特征加速 MAR,但其依赖隐层特征相似度启发式地跳过某些 Token 的生成顺序,直接破坏了预设的自回归轨迹,导致 FID 大幅恶化(Base 模型从 2.35 掉到 5.32);MARché 严格维持官方生成顺序,通过 KV 缓存与在线软合并实现更高加速比(1.57× vs 1.41×)的同时将 FID 牢固控制在 2.56。
  • vs ENAT (Ni et al., NeurIPS 2024):ENAT 需要额外设计定制的时空注意力模块并进行端到端微调重训;MARché 则是纯粹的推理期训练无关算法,兼容任意已发布的预训练权重。
  • vs LLM 经典 KV 压缩 (Keyformer / H2O):大语言模型因果掩码天然保护了单向时间轴的历史 KV;MARché 则成功克服了视觉图像中 2D 双向全局注意力的表征演化难题,为密集视觉自回归系统的缓存设计提供了范式参考。

评分

  • 新颖性: ⭐⭐⭐⭐☆ 巧妙地将双向注意力的局部平稳性转化为两路解耦注意力与动态刷新机制,理念简练而精妙。
  • 实验充分度: ⭐⭐⭐⭐⭐ 涵盖了 B/L/H 三种规模模型与完整基线对比,提供了深入详实的 Token 构建消融与算子耗时分解。
  • 写作质量: ⭐⭐⭐⭐⭐ 逻辑结构严密清晰,动机直观且分析深入。
  • 价值: ⭐⭐⭐⭐⭐ 无需训练且零精度破坏的即插即用特性,对视觉自回归图像生成的工业部署与边缘端提速极具实用价值。