跳转至

AdaNAT: Exploring Adaptive Policy for Token-Based Image Generation

会议: ECCV 2024
arXiv: 2409.00342
代码: https://github.com/LeapLabTHU/AdaNAT
领域: 图像生成
关键词: 非自回归Transformer, 强化学习, 自适应生成策略, 对抗奖励, token生成

一句话总结

提出AdaNAT，将非自回归Transformer（NAT）的生成策略配置建模为MDP，通过轻量策略网络+PPO强化学习+对抗奖励模型自动为每个样本定制生成策略（重掩码比例、采样温度、CFG权重等），在ImageNet-256上仅用8步达到FID 2.86，相比手工策略实现约40%的相对提升。

背景与动机

基于离散token的图像生成是当前与扩散模型并行的重要范式。其中非自回归Transformer（NAT，如MaskGIT、Muse、MAGE）通过并行解码机制在少量步数内生成图像，展现了效率与质量兼顾的潜力。然而NAT的生成过程需要配置复杂的策略——每一步都要设定重掩码比例\(m^{(t)}\)、采样温度\(\tau_1^{(t)}\)、重掩码温度\(\tau_2^{(t)}\)以及CFG引导强度\(w^{(t)}\)，总共\(4 \times T\)个超参数。现有方法全部依赖人工设计的调度函数（cosine schedule等），这不仅费时费力，而且效果次优——人工调度是全局共享的"一刀切"策略，无法适应不同样本的复杂度差异。

核心问题

如何自动地、自适应地为NAT的每个生成样本配置最优策略？具体挑战有三：（1）\(4T\)个超参数的联合搜索空间巨大，人工调优不现实；（2）不同复杂度的样本理应使用不同的生成策略（简单图案早收敛少调整，复杂结构需要更多精细化步骤），但现有方法对所有样本用同一策略；（3）离散token生成过程不可微，无法直接端到端优化策略网络。

方法详解

整体框架

AdaNAT在一个预训练好的、冻结的NAT模型之上，训练一个轻量级策略网络。输入是当前步的生成状态（时间步\(t\)和当前token序列\(\mathbf{v}^{(t)}\)），输出是当前步的策略配置（\(m^{(t)}, \tau_1^{(t)}, \tau_2^{(t)}, w^{(t)}\)）。整个流程不修改底层NAT模型的参数，仅学习"如何更好地使用它"。

关键设计

  1. MDP建模: 将NAT的\(T\)步生成过程建模为马尔可夫决策过程。状态\(s_t = (t, \mathbf{v}^{(t)})\)是当前token序列，动作\(a_t\)是四个策略参数，状态转移由冻结NAT模型决定（给定动作后执行一步并行解码+重掩码），奖励仅在最后一步给出（评估最终生成图像的质量）。这个建模巧妙地绕过了离散token不可微的问题。

  2. 轻量策略网络: 策略网络复用NAT模型已有的输出特征\(f_\theta(\mathbf{v}^{(t)})\)作为输入（不需要额外编码器），仅包含一个深度卷积层+逐点卷积层+MLP，附加AdaLN注入时间步信息。整体推理开销仅占NAT总推理开销的0.03%，几乎可忽略。策略以高斯分布形式\(\pi_\phi(a_t|s_t) = \mathcal{N}(\eta_\phi(s_t), \sigma I)\)输出，训练时做随机探索，推理时取均值。

  3. 对抗奖励模型: 这是论文最核心的贡献。作者系统对比了三种奖励设计：

  4. FID奖励：FID是统计指标无法提供样本级信号，实测策略网络训练失败或产生视觉质量差的图像（FID低但图像模糊/畸变）——说明FID可被"hack"
  5. 预训练奖励模型（ImageReward）：可提供样本级信号，但策略网络倾向生成风格趋同的图像，多样性严重不足——"过拟合"静态奖励
  6. 对抗奖励（本文方案）：引入一个类GAN判别器作为奖励模型\(r_\psi\)，与策略网络构成minimax博弈——策略网络最大化奖励，奖励模型同时更新以区分真假图像。由于奖励是动态变化的，策略网络无法过拟合，最终在保真度和多样性之间取得平衡。

损失函数 / 训练策略

  • 策略网络用PPO算法优化，采用clipped surrogate objective + 价值函数损失
  • 对抗奖励模型用标准GAN判别器损失训练（判断图像真假的二分类交叉熵）
  • 两者交替更新，每轮各做5次梯度更新以稳定minimax游戏
  • 整体仅需1000轮迭代即可收敛，batch size 4096，探索参数\(\sigma\)从0.6在500轮后衰减到0.3
  • 关键：整个过程中NAT模型参数完全冻结，不需要反向传播梯度穿过NAT

实验关键数据

数据集 模型 步数 TFLOPs FID-50K 对比SOTA
ImageNet-256 AdaNAT-S 4 0.2 4.54 MaskGIT(8步): 6.18
ImageNet-256 AdaNAT-S 8 0.3 3.71 MaskGIT-RS(8步): 4.02
ImageNet-256 AdaNAT-L 4 0.5 3.63 U-ViT-H†(4步): 8.45
ImageNet-256 AdaNAT-L 8 0.9 2.86 DiT-XL†(8步): 5.18
ImageNet-512 AdaNAT-L 8 1.2 3.66 ADM-G: 7.72
MS-COCO AdaNAT-S 8 0.3 5.75 U-ViT†(8步): 6.37
CC3M AdaNAT-Muse 8 2.8 6.83 Muse(8步): 7.67

消融实验要点

  • 可学习性贡献: 手工策略→可学习（非自适应）策略：FID从7.65降到5.40（-30%）；再加自适应：5.40→4.54（再降16%）。总相对提升约40%
  • 对抗奖励 vs 替代方案: FID奖励导致自适应策略训练崩溃（FID 55.4）；预训练奖励模型的图像多样性差；对抗奖励在质量和多样性间取得最佳平衡
  • 策略网络开销: 仅占总推理开销的0.03%，完全可忽略

亮点

  • "优化使用者而非生成器"的思路：不改动预训练NAT模型，只学习如何更好地使用它——这是一种高效的后处理式优化范式，可推广到任何需要复杂策略配置的生成模型
  • RL+对抗奖励的组合：通过动态变化的奖励信号解决策略过拟合问题，这个见解对所有基于RL的生成优化（如RLHF）都有启发
  • 自适应策略的可视化：论文清晰展示了策略网络对不同复杂度样本的差异化行为——简单图像早停精调，复杂图像持续大幅调整——验证了自适应的合理性
  • 极致轻量：策略网络复用NAT特征，额外开销0.03%，等于"免费"增强

局限性 / 可改进方向

  • 作者指出尚未在超大规模数据集（如LAION-5B）和超大模型（>1B参数）上验证可扩展性
  • 对抗奖励模型的训练稳定性可能在不同数据集/模型规模上有所不同，缺乏深入分析
  • 仅验证了class-conditional和text-to-image生成，未探索其他生成任务（如图像编辑、视频生成）
  • 对抗奖励模型采用StyleGAN-T的判别器架构，未探索更多判别器设计的影响
  • 未与同期的AutoNAT（CVPR 2024）做充分的公平对比——两者思路相似但奖励设计不同

与相关工作的对比

  • vs MaskGIT/MAGE: 这些工作使用手工设计的cosine schedule等调度函数，AdaNAT证明这些设计严重次优（FID相对差40%），且无法做样本级自适应。AdaNAT作为后处理可直接增强这些模型
  • vs AutoNAT (CVPR 2024): 同样优化NAT策略，但AutoNAT使用FID作为优化目标。AdaNAT发现FID目标会导致视觉质量差（尽管FID数值低），提出对抗奖励作为更优替代。定量上AdaNAT-FID (2.56) < AutoNAT (2.68)，但AdaNAT-Adv (2.86)在视觉质量和多样性上更优
  • vs 扩散模型RL优化（DPOK、DDPO等）: 这些工作直接微调扩散模型以对齐人类偏好；AdaNAT不改动生成器，只优化策略配置——思路正交，可用于不同生成范式

启发与关联

  • 与EVATok的对比启发: EVATok采用"先离线估最优→再训路由器模仿"的两步范式为视频token分配自适应长度，而AdaNAT用RL端到端学习最优策略——两种自适应策略学习的范式各有优劣，EVATok更稳定但需要离线搜索，AdaNAT更灵活但依赖奖励设计
  • RL优化生成策略的通用框架: AdaNAT的核心框架（冻结生成器+RL策略网络+对抗奖励）可能推广到其他需要复杂推理时配置的生成模型——如扩散模型的自适应步数/CFG/采样器选择
  • 与 RL驱动的自适应成像参数优化 idea共享"RL学习最优策略参数"的核心范式

评分

  • 新颖性: ⭐⭐⭐⭐ 将NAT策略配置建模为MDP并用RL优化不算全新，但对抗奖励设计的洞察和系统对比非常有价值
  • 实验充分度: ⭐⭐⭐⭐ 四个数据集覆盖class-conditional和text-to-image，消融和分析详实，但缺少大规模验证
  • 写作质量: ⭐⭐⭐⭐⭐ 问题motivate清晰，三种奖励设计的递进分析像讲故事一样引人入胜
  • 价值: ⭐⭐⭐⭐ "不改模型只改策略"的后处理优化范式和对抗奖励设计都有较好的可迁移性