Accelerating Inference of Masked Image Generators via Reinforcement Learning

📄 arXiv: 2512.01094v1 📥 PDF

作者: Pranav Subbaraman, Shufan Li, Siyan Zhao, Aditya Grover

分类: cs.CV

发布日期: 2025-11-30

备注: 15 pages, 9 figures

💡 一句话要点

提出Speed-RL，通过强化学习加速掩码图像生成模型推理，显著减少采样步骤。

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)

关键词: 掩码图像生成模型 强化学习 模型加速 图像生成 推理优化

📋 核心要点

  1. 掩码生成模型生成高质量图像能力强，但推理速度慢，需要大量采样步骤。
  2. Speed-RL将加速问题转化为强化学习问题，结合质量和速度奖励微调模型。
  3. 实验表明，Speed-RL能在保持图像质量的同时，将模型推理速度提升3倍。

📝 摘要（中文）

掩码生成模型(MGM)在生成高质量图像方面表现出强大的能力。然而，它们需要大量的采样步骤才能生成高质量的图像，导致推理速度缓慢。本文提出了一种新的范式Speed-RL，用于加速预训练的MGM，使其能够在更少的步骤中生成高质量的图像。与传统的蒸馏方法将加速问题定义为分布匹配问题不同（即训练一个少步数的学生模型来匹配多步数教师模型生成的分布），我们将这个问题视为一个强化学习问题。由于加速的目标是在更少的步骤中生成高质量的图像，我们可以将质量奖励与速度奖励相结合，并使用强化学习对基础模型进行微调，并将组合奖励作为优化目标。通过大量的实验，我们表明所提出的方法能够在保持相当图像质量的同时，将基础模型加速3倍。

🔬 方法详解

问题定义：论文旨在解决掩码图像生成模型（MGM）推理速度慢的问题。现有的MGM需要大量的采样步骤才能生成高质量的图像，这限制了它们在实际应用中的部署。传统的加速方法，如蒸馏，通常将加速问题视为分布匹配问题，但这种方法可能难以捕捉到MGM的复杂生成过程。

核心思路：论文的核心思路是将MGM的加速问题建模为一个强化学习问题。通过定义一个结合图像质量和推理速度的奖励函数，利用强化学习算法来优化MGM的采样策略，从而在更少的步骤内生成高质量的图像。这种方法直接优化了加速的目标，避免了传统蒸馏方法中分布匹配的间接性。

技术框架：Speed-RL的技术框架主要包含以下几个部分：1) 预训练的MGM作为基础模型；2) 强化学习环境，其中状态是当前生成的图像，动作是MGM的采样步骤；3) 奖励函数，结合图像质量（如FID分数）和推理速度（采样步骤数）；4) 强化学习算法（如PPO），用于优化MGM的采样策略。整个流程是，MGM在强化学习环境中进行采样，根据生成的图像和采样步骤计算奖励，然后使用强化学习算法更新MGM的参数，使其能够生成更高质量的图像，同时减少采样步骤。

关键创新：Speed-RL的关键创新在于将MGM的加速问题建模为一个强化学习问题。与传统的蒸馏方法不同，Speed-RL直接优化了加速的目标，即在更少的步骤内生成高质量的图像。此外，Speed-RL使用了一个结合图像质量和推理速度的奖励函数，这使得模型能够在两者之间进行权衡，从而获得更好的加速效果。

关键设计：奖励函数的设计是Speed-RL的关键。论文中使用的奖励函数通常包含两部分：图像质量奖励和速度奖励。图像质量奖励可以使用FID分数或Inception Score等指标来衡量生成图像的质量。速度奖励则与采样步骤数成反比，鼓励模型在更少的步骤内完成生成。此外，强化学习算法的选择也很重要，论文中通常使用PPO等算法来优化MGM的采样策略。具体的参数设置，如奖励函数的权重、学习率等，需要根据具体的任务进行调整。

📊 实验亮点

实验结果表明，Speed-RL能够在保持相当图像质量的同时，将基础MGM的推理速度提升3倍。具体来说，Speed-RL在多个数据集上都取得了显著的加速效果，并且生成的图像质量与原始MGM相当。这些结果表明，Speed-RL是一种有效的MGM加速方法，具有很强的实用价值。

🎯 应用场景

Speed-RL具有广泛的应用前景，可以应用于各种需要快速图像生成的场景，例如实时图像编辑、视频生成、游戏开发等。通过加速MGM的推理速度，可以降低计算成本，提高用户体验，并促进MGM在实际应用中的部署。此外，该方法也可以推广到其他类型的生成模型，例如文本生成模型和音频生成模型。

📄 摘要（原文）

Masked Generative Models (MGM)s demonstrate strong capabilities in generating high-fidelity images. However, they need many sampling steps to create high-quality generations, resulting in slow inference speed. In this work, we propose Speed-RL, a novel paradigm for accelerating a pretrained MGMs to generate high-quality images in fewer steps. Unlike conventional distillation methods which formulate the acceleration problem as a distribution matching problem, where a few-step student model is trained to match the distribution generated by a many-step teacher model, we consider this problem as a reinforcement learning problem. Since the goal of acceleration is to generate high quality images in fewer steps, we can combine a quality reward with a speed reward and finetune the base model using reinforcement learning with the combined reward as the optimization target. Through extensive experiments, we show that the proposed method was able to accelerate the base model by a factor of 3x while maintaining comparable image quality.