AR-GRPO: Training Autoregressive Image Generation Models via Reinforcement Learning
Paper: arXiv:2508.06924 Code: Kwai-Klear/AR-GRPO Code reference:
main@59b8d54(2025-11-26)
1. Motivation(研究动机)
这篇论文的问题意识很直接:自回归图像生成模型已经能把图像离散化成 token 序列,并借助 Transformer / Llama 结构做 next-token prediction,但主流训练仍然主要依赖最大似然估计(MLE)。MLE 的目标是让模型在 teacher-forcing 条件下预测训练图像的下一个 token,它并不直接优化人类真正关心的图像质量、偏好、真实性、文本对齐或类别语义。作者指出这种错位带来三个具体问题。第一,likelihood 高不等于视觉上更好,因为 token-level CE 并不等价于感知质量或人类偏好。第二,MLE 可能让结果“平均地合理”,但不会主动奖励高频细节、真实感和美学质量。第三,AR 模型训练时看到 ground-truth prefix,推理时却依赖自己已经采样出的 token,exposure bias 会把早期错误累积成伪影、结构异常或语义漂移。
论文把这个问题类比到 LLM 后训练:LLM 已经证明 RLHF / RLVF 可以在 likelihood 之外把模型推向任务目标,那么图像 AR 模型也可以把完整生成图像当作一个 response,用奖励模型评价后再做 policy optimization。这里的关键不是简单把 diffusion RL 的 recipe 搬过来,而是说明 AR 图像生成的动作空间是离散视觉 token,完整图像质量只能在整条序列采完后评价,因此 reward 是 sequence-level,梯度则要回传到 token-level log probability。AR-GRPO 的核心动机就是把这种“完整图像的黑盒奖励”转成每个生成 token 的 policy gradient 信号。
作者选择 LlamaGen 作为基础模型还有一个现实原因:LlamaGen 已经提供 class-to-image(C2I)和 text-to-image(T2I)两条代表性 AR 图像生成路径,前者以 ImageNet class label 为条件,后者用 FLAN-T5 XL text embedding 作为前缀条件。这样可以同时验证 RL 对类别语义、视觉质量、文本对齐和人类偏好指标的作用。论文贡献可以概括为三点:第一,首次系统地把 GRPO 用于 autoregressive image generation;第二,设计了覆盖条件一致性、图像质量和真实感的多维 reward;第三,在 C2I 与 T2I 两种任务上展示 RL 能提升质量和偏好,但也会带来 entropy 下降、diversity/recall 损失等 trade-off。
Figure 1 解读:teaser 展示的不是一个新 tokenizer 或新 backbone,而是把 LlamaGen 的采样结果接入 reward system 后做 online RL。图中最重要的信息是训练闭环:condition 输入 AR policy,policy 采样多个候选图像,reward model 对完整图像打分,GRPO 用组内相对优势更新 policy。它强调 AR-GRPO 的 reward 不是逐 token 的人工标签,而是完整图像级别的自动评价。
2. Idea(核心思想)
AR-GRPO 的核心思想是:把自回归图像生成重新表述为一个 MDP,然后使用无 critic 的 group-relative policy optimization,把同一个条件下多个候选图像之间的相对好坏变成优势函数。给定条件 (类别或文本)和视觉 token 序列 ,AR 模型本来学习 MLE 训练最小化 ,而 AR-GRPO 关心的是采样完成后图像 的 reward 。因此作者把状态定义为 ,动作定义为从 codebook vocabulary 中选择下一个 visual token ,policy 就是 AR model 的 next-token distribution,reward 则只在完整序列 decode 成图像后计算。这个建模有一个重要含义:单个 token 很难判断质量,真正可评价的是完整图像;但是 GRPO 的 ratio 和 KL 是按 token logprob 计算,所以 sequence-level reward 会通过 normalized advantage 作用到所有生成 token。
为什么选 GRPO 而不是 PPO?论文的隐含理由是图像 token 序列的 state/action 空间很大,训练额外 value model 会增加成本和不稳定性。GRPO 对每个 condition 采样 个输出 ,用 reward system 得到 ,再用组内均值和方差标准化: 直觉上,绝对 reward 的尺度可能在不同 prompt / class 之间不可比,但同一个条件下的多个候选可以比较好坏。这样高于组均值的图像会提高其 token 概率,低于组均值的图像会被压低。目标函数使用 PPO-style clipping: 其中 ,论文设置 。KL penalty 的作用是防止 policy 过快偏离 reference model;实验也证明没有 KL 时 IS/Precision 会更高,但 recall 和 FID 明显变差,出现 mode collapse 和 distribution drift。
这篇论文的新意不在于提出全新的 RL 算法,而在于把 GRPO 的组内相对优势和多维视觉 reward 接到 AR 图像生成的离散 token 序列上,并系统展示它的利弊:C2I 中,RL 让模型更确定、更偏向高 reward 类别语义,IS 和 Precision 上升,但 Recall 下降;T2I 中,文本对齐、图像质量和偏好指标同时改善,说明当 reward 目标更明确时,AR 模型也能像 LLM 一样做 preference/alignment-style post-training。
3. Method(方法)
3.1 AR 图像生成基础
论文基于 LlamaGen。图像先经过 encoder-quantizer-decoder 结构离散化为 visual tokens;quantizer codebook 写作 ,实验使用默认 、。给定 tokenized visual sequence,LlamaGen 使用 Llama-style Transformer:RMSNorm、SwiGLU、RoPE。论文验证了三个模型规模:LlamaGen-B(111M, 12 layers, hidden 768, 12 heads)、LlamaGen-L(343M, 24 layers, hidden 1024, 16 heads)、LlamaGen-XL(775M, 36 layers, hidden 1280, 20 heads)。这些模型不是从零训练,而是在 pretrained LlamaGen checkpoint 上继续 RL fine-tuning。
C2I 条件 是 ImageNet 类别标签;T2I 条件 是文本 prompt,经 FLAN-T5 XL 编码,再通过额外 MLP 投影成 AR 模型的 prefilling token embeddings。两者之后都以 next-token prediction 方式生成视觉 token,直到达到序列最大长度,再由 decoder 还原成图像。AR-GRPO 只改后训练过程,不改变 tokenizer、decoder 和主干架构。
3.2 多维 reward 设计
生成完整 token 序列 后,模型先 decode 得到图像 ,然后计算三类 reward:
- Conditional Reward:衡量图像是否符合条件。论文写作
C2I 中 CLIP/HPSv2 需要把类别条件转成文本语义;T2I 中直接比较图文一致性。这个 reward 主要约束类别、对象存在、数量和文本语义。
- Image Quality Reward:使用 MANIQA 这类 no-reference image quality assessment,写作
它关注清晰度、局部细节和视觉吸引力,弥补 CLIP/HPSv2 不擅长判断低层质量的问题。
- Realism Reward:使用 Qwen2.5-VL-3B-Instruct 作为 VLM judge,写作
Appendix 中说明 Qwen 评估三类维度:Fake Identification(是否明显 AI 生成)、Label Recognition(是否满足类别或文本条件)、Weird Detection(是否有奇怪结构)。完整 prompt 还要求对对象匹配、是否缺失部件、是否真实合理、是否清晰高质量、是否无噪声/缺陷/伪影分别给 0–1 分,并输出 JSON 中的 score/explanation。
最终 reward 为
论文主实验设置三个权重都为 1。Appendix 进一步说明,为了平衡不同 reward 尺度,CLIP 与 HPSv2 乘以 5,MANIQA 乘以 2,Qwen reward 乘以 0.25;为了缓解 CLIP/HPSv2 reward hacking,还加入它们的 quantized 版本,将连续值映射为 0.5、1、1.5 三档,并与原始连续 reward 一起训练。released code 中 C2I reward list 是 clip_distance clip_quantize hpsv2 hpsv2_quantize maniqa qwenvl_3b qwenvl_fake qwenvl_weird;T2I reward list 是 clip_text clip_text_quantize hpsv2_text hpsv2_text_quantize maniqa qwenvl_3b_text qwenvl_fake qwenvl_weird,对应论文的条件、质量、真实感三类目标。
3.3 Training loop 与代码路径
论文层面的训练循环是三阶段交替:Generation Phase 对每个 condition 采样 个输出;Evaluation Phase 对所有输出计算多维 reward;Optimization Phase 用 GRPO objective 更新 policy。released repo 的实现与此基本对应:img_gen_grpo_train.py 负责解析参数、构造 LlamaGen 模型、根据 reward_funcs registry 实例化 reward、构造 dataset,并把 generation kwargs 注入 trainer;modified_grpo_trainer.py 扩展 TRL 的 GRPOTrainer,负责生成、打分、计算 per-token logprob、归一化 advantage 和 GRPO loss;img_gen_grpo_rewards.py 实现 CLIP/HPSv2/MANIQA/Qwen 等 reward 类。
# Paper + released-code level pseudocode, anchored to main@59b8d54
def ar_grpo_train(policy, ref_policy, dataloader, reward_fns, group_size=8):
for condition_batch in dataloader:
# 1. Generation: same class/text condition, multiple on-policy samples
samples = policy.generate(
condition_batch,
num_generations=group_size,
top_k=0,
top_p=1.0,
temperature=1.0,
cfg_scale=train_cfg,
)
# 2. Evaluation: decode image tokens and compute multi-faceted rewards
decoded_images = tokenizer_decoder(samples.visual_tokens)
reward_matrix = []
for reward_fn in reward_fns:
reward_matrix.append(reward_fn(condition_batch, decoded_images))
rewards = weighted_sum(reward_matrix, reward_weights)
# 3. Group-relative advantage: normalize within each condition group
rewards_grouped = rewards.reshape(batch_size, group_size)
advantages = (rewards_grouped - rewards_grouped.mean(dim=1, keepdim=True))
advantages = advantages / (rewards_grouped.std(dim=1, keepdim=True) + 1e-4)
# 4. Token-level GRPO loss with ratio clipping and KL to reference
logp_new = policy.logprob(samples.tokens, condition_batch)
logp_old = samples.old_logprob
ratio = exp(logp_new - logp_old)
clipped_ratio = clamp(ratio, 1 - epsilon, 1 + epsilon)
policy_loss = -minimum(ratio * advantages, clipped_ratio * advantages)
kl_loss = beta * kl(policy, ref_policy, samples.tokens)
loss = masked_mean(policy_loss + kl_loss, samples.completion_mask)
loss.backward(); optimizer.step(); optimizer.zero_grad()Paper Concept → Source File → Key Class/Function
| Paper Concept | Source File | Key Class/Function |
|---|---|---|
| Training entry and reward registry | img_gen_grpo_train.py | reward_funcs_registry, ImgGenGRPOScriptArguments, build_model, main |
| C2I exact launch config | run_c2i_train.sh | reward_funcs, reward_weights, beta, num_generations, gen_cfg_cfg_scale |
| T2I exact launch config | run_t2i_train.sh | dataset_name, gpt_ckpt, reward_funcs, beta, force_model_bf16 |
| GRPO trainer | modified_grpo_trainer.py | GRPOTrainer_M, _generate_and_score_completions, _calculate_rewards, _get_per_token_logps, compute_loss |
| Image reward functions | img_gen_grpo_rewards.py | CLIPDistanceReward, CLIPTextReward, HPSV21DistanceReward, HPSV21TextReward, MANIQAReward, QwenTextReward, QwenFakeDiscrimReward, QwenWeirdDiscrimReward |
| LlamaGen model backbone | autoregressive/models/gpt.py, autoregressive/models/gpt_hf.py | GPT/LlamaGen model definitions used by build_model |
| Sampling/evaluation utilities | autoregressive/sample/*.py, test_t2i.py, test_t2i_rewards.py | C2I/T2I sampling and reward evaluation scripts |
Training config evidence from actual launch scripts, not base defaults:run_c2i_train.sh sets gen_cfg_cfg_scale=2.0, num_generations=8, top_p=1.0, top_k=0, C2I rewards above, all reward weights 1, beta=0.1, learning_rate=1e-5, per_device_train_batch_size=8, bf16=true, num_train_epochs=20, exit_step=${steps} and gpt_ckpt=./pretrained_models/c2i_XL_384.pt for the released C2I run.
run_t2i_train.sh sets dataset_name=./dataset/coco_captions_30000.json, vq_ckpt=./pretrained_models/vq_ds16_t2i.pt, gpt_ckpt=./pretrained_models/t2i_XL_stage1_256.pt, num_generations=8, top_p=1.0, top_k=0, T2I rewards above, reward weights 1, beta=1.0, learning_rate=1e-5, per_device_train_batch_size=16, bf16=true, force_model_bf16=True, num_train_epochs=20。
注意:论文正文称 T2I 训练和推理 CFG 都为 2.0,但当前 released script 的 run_t2i_train.sh 写的是 cfg=4.0;因此复现时应以代码 commit 和实际启动脚本为准,并记录与论文配置的差异。
4. Experiments(实验与结果)
4.1 C2I:ImageNet class-conditional generation
C2I 训练基于 ImageNet 和 LlamaGen pretrained checkpoint,论文默认图像分辨率 256×256;采样时 top_k=0(all)、top_p=1.0、temperature=1.0;训练 CFG 默认 2.0;每个输入采样 8 个候选;启用 KL,;所有 dropout 设为 0;总 batch size 64;AdamW learning rate ,、、weight decay=0.05。表 2 在 ImageNet benchmark 上比较 LlamaGen-B/L/XL 和 RL 后模型。
关键数值如下:LlamaGen-B 256 的 IS 从 193.61 提升到 211.67,Precision 从 0.84 到 0.86,但 FID 从 5.46 变差到 5.91,Recall 从 0.46 降到 0.40。LlamaGen-B 384 的 IS 从 157.17 到 195.45,Precision 从 0.81 到 0.88,但 FID 从 6.44 到 7.64,Recall 从 0.46 到 0.34。LlamaGen-L 384 的 IS 从 256.07 到 264.64,Precision 从 0.83 到 0.87,sFID 从 6.09 小幅改善到 5.91,但 FID 和 Recall 分别从 3.08/0.52 变为 4.24/0.44。LlamaGen-XL 384 的 IS 从 286.88 到 293.07,Precision 从 0.84 到 0.86,但 FID 从 2.78 到 3.63,Recall 从 0.54 到 0.48。
这些结果说明 RL 的收益主要体现在“更像高 reward 类别样本”:IS 和 Precision 普遍提高,视觉语义更贴近类别;代价是 distribution coverage 下降,FID/Recall 往往变差。作者把它解释为 GRPO 鼓励同组高 reward 样本,policy entropy 降低,采样更确定。换句话说,AR-GRPO 不是免费提升所有指标,而是把生成分布向 reward 偏好的高质量模式移动。
Figure 2 解读:C2I scaling 图展示了 B/L/XL 与 256/384 下的视觉样本。它支持论文的一个实用结论:同一套 RL 后训练不只适用于单一模型规模,在更大模型和更高分辨率上也能工作;但表格显示,规模越大,baseline 本身越强,RL 的边际收益更偏向 Precision/IS,而不是 FID/Recall 全面提升。
Figure 3 解读:C2I entropy 曲线验证了“质量-多样性”权衡。RL 后 policy entropy 下降,表示模型在同一条件下更少探索低 reward token 路径;这解释了 Precision/IS 上升,也解释了 Recall 降低。对实际使用而言,如果目标是类别准确、少伪影、偏好分数高,entropy 下降可接受;如果目标是覆盖 ImageNet 分布的多样性,则需要更强 KL、更温和 CFG 或 diversity reward。
4.2 T2I:text-conditional generation
T2I 训练基于 LlamaGen T2I stage-I model,使用 FLAN-T5 XL 提取文本 embedding,训练数据为 LAION-COCO 子集,分辨率 256×256。论文正文称训练和推理 CFG 都为 2.0,其余配置与 C2I 类似。评估包括 GenEval 和 DrawBench;GenEval 测对象位置、颜色、属性绑定、计数、单物体、双物体等 compositional 能力;DrawBench 额外报告 CLIP、HPSv2、Aesthetic、DeQA、ImageReward、PickScore、UnifiedReward。
GenEval 上,LlamaGen-XL overall 从 0.306 提升到 0.324。细项中 color 从 0.550 到 0.593,counting 从 0.197 到 0.228,single object 从 0.750 到 0.791;position 从 0.042 到 0.040、attribute binding 从 0.032 到 0.030,two object 维持 0.263。这个结果说明 reward 对对象存在、颜色和计数更有效,但对更困难的空间位置和属性绑定没有明显改善。作者强调他们没有直接使用 GenEval reward 训练,这意味着提升来自通用 alignment/quality/realism reward 的泛化,而不是 benchmark-specific optimization。
DrawBench 上所有报告指标都改善:CLIP 0.245→0.274,HPSv2 0.153→0.208,Aesthetic 4.701→4.808,DeQA 0.546→0.551,ImageReward -1.784→-0.712,PickScore 0.737→0.769,UnifiedReward 0.271→0.312。相比 C2I 的 FID/Recall trade-off,T2I 的结果更像典型 alignment 后训练:因为 reward 明确覆盖文本对齐、质量和偏好,RL 能同时推高多种指标。
Figure 4 解读:T2I visual comparison 展示 baseline 与 RL model 的定性差异。结合指标看,RL 后模型更倾向于生成语义更明确、视觉上更完整的样本;但这类对比图不能单独证明泛化,真正有说服力的是 GenEval/DrawBench 的多指标提升以及 reward ablation。
Figure 5 解读:T2I entropy 在训练过程中同样下降,说明 GRPO 对 AR policy 的影响具有一致性:不管条件来自类别还是文本,reward 优化都会让采样分布更集中。T2I 中这种集中带来更多 alignment/quality 收益;C2I 中则更明显地牺牲 Recall。
4.3 Reward、KL 与 CFG 的消融
KL penalty 消融使用 LlamaGen-B 256。无 KL 时 IS 220.57、Precision 0.88,比有 KL 的 IS 211.67、Precision 0.86 更高;但 FID 7.83、sFID 11.72、Recall 0.32 明显差于有 KL 的 FID 5.91、sFID 8.58、Recall 0.40。论文的解释是无 KL 会让 policy 过度追逐 reward,造成 mode collapse 和 distribution drift。这个消融是全篇最重要的安全阀证据:GRPO 能优化 reward,但必须用 KL 保持与 reference model 的距离。
Reward design 消融显示三类 reward 都有贡献。C2I 中,去掉 conditional reward 的 RL(A) IS 215.49/FID 6.00/Recall 0.39;去掉 image quality 的 RL(B) IS 205.41/FID 5.77/Recall 0.40;去掉 realism 的 RL(C) IS 230.95/FID 6.55/Recall 0.37;三者全开 RL(D) IS 211.67/FID 5.91/Recall 0.40。这里不能只看 IS:RL(C) 的 IS 最高,但 FID/Recall 更差,说明 realism reward 对抑制“看起来高分但分布漂移”的样本有作用。T2I GenEval 中,去掉 conditional reward 的 overall 只有 0.225,去掉 image quality 为 0.254,去掉 realism 为 0.260,全开为 0.324;这说明文本条件任务中 conditional reward 是最关键的组件之一,因为它直接约束对象存在、数量和文本语义。
Figure 6 解读:overall reward 随训练 step 上升,说明模型确实学会了提高 reward,而不是仅仅评估噪声波动。Appendix 还显示 HPSv2、MANIQA、Qwen Weird reward 增长更明显,而 CLIP、Qwen Fake、Qwen Label 相对稳定;这意味着训练主要改善美学/质量/异常结构相关维度,而基础语义匹配可能在 baseline 中已经较高或被 reward quantization 限制。
CFG 消融只在 C2I LlamaGen-B 256 上做。训练 CFG 变化实验固定 inference CFG ,发现中等训练 CFG 最好,论文称 在 B-size 模型上取得较好平衡;推理 CFG 变化实验用 训练的模型测试不同 inference CFG,发现更高 inference CFG 能把 IS 推到约 390(),但 FID 也恶化到约 18。结论是训练 CFG 和推理 CFG 不必相同;实践中可以在训练 CFG 2–8 范围内寻找平衡,推理时根据更重视质量/分布对齐还是更重视多样性/高 IS 调整 CFG。
Figure 7 解读:论文源码中 cfg-train.pdf 与 cfg-test.pdf 属于同一 figure,这里按 LaTeX source 的左右子图组合为一张图。左图对应训练 CFG,右图对应推理 CFG。它提醒复现者不要把 CFG 当作固定常数:CFG 与 GRPO reward 会共同改变采样分布,CFG 过高可能让 reward/IS 上升但 FID 变坏。
4.4 附加可视化
Appendix T2I 可视化显示 RL-trained model 在多种 prompt 下的样本。它适合用于理解 qualitative behavior,但不应替代指标判断。尤其是 AR-GRPO 会降低 entropy,因此 cherry-picked visual results 可能高估平均多样性;需要结合 GenEval/DrawBench、Recall 和 entropy 曲线一起看。
5. Limitations(局限、复现注意点与启发)
第一,AR-GRPO 的 reward 是自动 reward,不是人工偏好标注。CLIP/HPSv2/MANIQA/Qwen2.5-VL 都有偏差:CLIP/HPSv2 容易被语义捷径或风格偏好影响;MANIQA 关注 no-reference quality,不保证语义正确;VLM judge 可能对 prompt、解析格式和自身幻觉敏感。论文用 reward scaling、quantization 和多 reward 组合缓解 reward hacking,但不能从根本上证明 reward 与真实人类偏好完全一致。因此它更像“自动多目标后训练”而不是严格 RLHF。
第二,质量提升伴随多样性下降。C2I 中 Recall 普遍下降,entropy 曲线也证明 policy 更确定。这个结果不是小缺陷,而是 GRPO 目标的自然后果:如果 reward 偏好少数更安全、更清晰、更符合类别的模式,policy 会收缩到这些模式。实际部署时需要根据目标决定是否加入 diversity reward、提高 KL、降低 CFG、调高 temperature,或限制 RL steps。
第三,论文的部分指标存在 trade-off,不能只汇报单一高分。例如无 KL 的 IS/Precision 更高,但 Recall 和 FID 明显恶化;C2I RL(C) 去掉 realism reward 后 IS 最高,但分布指标更差;推理 CFG 过高也能推高 IS,却让 FID 变差。读这篇论文时应把 IS、Precision、FID、Recall、entropy 放在一起看,而不是把“RL 提升图像质量”理解成全指标同时变好。
第四,代码复现需要锁定 commit 和脚本。论文正文 C2I 默认 ,T2I 正文称 CFG=2.0;但当前 released repo main@59b8d54 的 T2I launch script 使用 cfg=4.0 和 beta=1.0。这不一定代表论文错误,可能是 release 后脚本调整或另一套默认实验;但复现记录必须以 run_c2i_train.sh / run_t2i_train.sh 的实际值为准。代码中 reward registry 还包含 continuous 与 quantized 的 CLIP/HPSv2,以及 Qwen fake/weird/label/text variants,实际 reward 数量比论文的三类抽象更多。
第五,AR-GRPO 对图像生成 RL 的启发是:离散 token AR 模型可以直接借用 LLM post-training 工具链,但视觉任务的 reward 需要更小心设计。LLM 中一个 answer 的 correctness 可能更容易验证,而图像 reward 往往是多目标、噪声大、易被 hack 的。论文最有价值的经验不是“GRPO 一定优于 MLE”,而是展示了一个可复用 recipe:group sampling + multi-faceted reward + KL regularization + entropy/Recall 监控。后续工作可以把这个 recipe 扩展到更强 AR 视觉模型、更多分辨率、人工偏好 reward、diversity-aware objective,或者把 reward 分解成可解释的局部区域/对象级反馈。