RankE: End-to-End Post-Training for Discrete Text-to-Image Generation with Decoder Co-Evolution¶

会议: NeurIPS 2026
arXiv: 2605.21195
论文: https://arxiv.org/abs/2605.21195
项目: https://syjmelody.github.io/RankE/
领域: 图像生成
关键词: 离散自回归生成, 文本到图像生成, 后训练对齐, 潜在协变量偏移, 解码器协同演变

一句话总结¶

针对离散自回归文生图后训练中策略分布偏移导致固定VQ解码器画质退化的“潜在协变量偏移”瓶颈，RankE通过交替协同演变机制联合优化自回归策略与VQ解码器，在统一广义EM框架下打破保真度与文本对齐的权衡困境。

研究背景与动机¶

离散自回归（AR）文生图模型通常采用两阶段范式：首先通过向量量化（VQ）分词器将连续图像像素网格映射为离散的codebook索引，随后利用自回归Transformer策略建模文本条件下的token序列。这种范式天然继承了大语言模型成熟的Scaling机制与训练基建，展现出统一多模态架构的巨大潜力。为了提升图像生成质量与人类偏好一致性，现有后训练（Post-Training）方法普遍引入强化学习对齐（如GRPO）。然而，这些方法几乎全部沿袭了“冻结VQ解码器、仅优化自回归策略网络”的固有惯例。

然而，策略单边优化必然引发严重的“潜在协变量偏移”（Latent Covariate Shift）。在分词器预训练阶段，VQ解码器完全拟合于确定性编码得到的基准离散编码分布；而在后训练强化学习阶段，自回归策略为了迎合奖励模型评分，其采样的token分布逐步偏离解码器的训练流形，且codebook有效熵显著坍缩。这种分布脱节导致了系统性的保真度-对齐度张力：策略在奖励模型下得分提高，但渲染出的图像却出现明显的非自然伪影与纹理崩坏，导致FID严重劣化。由于离散token采样和VQ量化两道非微导操作阻断了端到端梯度回传，传统扩散模型中的端到端微调方法无法直接应用。

针对离散分词器阻断像素级奖励反向传播的根本困难，本文的核心思路是不依赖有偏的梯度替代估计，而是通过双阶段交替优化实现策略与解码器的协同演化。核心 idea：将离散文生图后训练形式化为广义EM过程，以token级优势排序驱动自回归策略更新，同时以奖励加权对抗损失（Rank-GAN）结合流形一致性锚点驱动解码器协同适应，在离散瓶颈之外实现端到端奖励对齐。

方法详解¶

整体框架¶

RankE将策略网络 $\pi_\theta(z \mid y)$ 与VQ解码器 $D_\phi(z)$ 的联合优化构建为围绕离散瓶颈的交替演化流程。由于分类采样 $z \sim \pi_\theta$ 与硬量化切断了像素奖励到策略参数的直接梯度路径，RankE并不试图对离散采样进行有偏的连续松弛，而是在每一轮训练中将参数更新拆解为结构对称的两个阶段：Stage 1 冻结解码器，以Prompt级多候选采样的相对优势执行token级排序优化；Stage 2 冻结自回归策略，利用策略采样的离散token驱动解码器执行像素级排序更新与流形锚定。

从变分推断的角度审视，该交替优化等价于控制即推断（Control-as-Inference）框架下的广义期望最大化（Generalized EM）。Stage 1 对应变分E步，在固定渲染流形下更新隐变量后验分布 $\pi_\theta$；Stage 2 对应最大后验M步，更新生成流形参数 $\phi$，利用重构损失和指数移动平均（EMA）自蒸馏作为先验约束，防止解码器发生流形遗忘或局部过度拟合。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["输入文本 Prompt y 与真实图像 x_gt"] --> B["自回归策略采样 G 个候选 Token 序列"]
    B --> C["Token级排序优化 GRPO<br/>组内相对优势归一化 + EMA参考策略KL锚定"]
    C -->|交替轮次驱动隐空间分布| D["像素级双通道奖励对齐<br/>可微奖励反传 + 奖励加权对抗 Rank-GAN"]
    D --> E["双重流形锚定正则化<br/>GT真实编码重构 + EMA教师网络一致性蒸馏"]
    E --> F["输出协同进化的自回归策略与高保真解码器"]

关键设计¶

1. Token级排序优化 GRPO：基于组内相对优势的自回归策略对齐

在策略更新阶段，为了消除由于策略过度追求奖励标量而导致的分布崩溃，RankE采用组相对策略优化（GRPO）。针对每个输入文本prompt $y$，模型自回归采样出 $G$ 个候选离散token序列 $\{z_i\}_{i=1}^G \sim \pi_\theta(\cdot \mid y)$，通过当前解码器渲染后由奖励模型打分得到标量评价 $\{r_i\}_{i=1}^G$。GRPO并不直接依赖绝对奖励值，而是计算每个候选序列在Prompt组内的标准化优势 $A_i = (r_i - \mu_r) / \sigma_r$，将绝对标量转换为样本间的排序信号。

为了维持策略分布的平稳性并构建广义EM的E步变分下界，优化目标结合了PPO裁剪重要性采样与EMA参考策略 $\pi_{\text{ref}}$ 的KL散度约束： $$\mathcal{L}_{\pi}(\theta) = -\mathbb{E}_y \left[ \frac{1}{G}\sum_{i=1}^G \min\left(\rho_i A_i, \text{clip}(\rho_i, 1\pm\epsilon)A_i\right) - \beta \mathbb{D}_{\text{KL}}(\pi_\theta \parallel \pi_{\text{ref}}) \right]$$ 其中重要性权重 $\rho_i = \pi_\theta(z_i \mid y) / \pi_{\theta_{\text{old}}}(z_i \mid y)$。减去组均值在数学上充当无偏基线以降低梯度方差，PPO裁剪构成置信区间投影，而KL惩罚项则严格约束了策略演进速率，使其与后续解码器适应步骤保持动态平衡。

2. 像素级双通道奖励对齐：融合直接梯度反传与Rank-GAN加权判别

当策略更新后，解码器需要快速跟进并适配已漂移的token分布。RankE在像素空间为解码器引入排序驱动的对齐信号 $\mathcal{A}_\phi$。针对具有连续解析梯度的奖励模型（如CLIP Score），解码器提供了从离散token到像素空间的完全可微路径。此时直接冻结输入token，将奖励梯度反向传播至解码器参数 $\phi$： $$\mathcal{L}_{\text{reward}}(\phi) = -\mathbb{E}_{\hat{z} \sim \pi_\theta} \left[ R(D_\phi(\hat{z}), y) \right]$$ 由于输入 $\hat{z}$ 是脱离计算图的确定索引，梯度流被严格限制在像素解码器内部，避免了传统离散求导的不可行性。

而对于不可微黑盒人类偏好模型（如HPSv2），直接梯度通道失效。若退回常规GAN训练，判别器将无差别地逼近所有生成样本，抹平了策略阶段刚刚建立的样本排序偏好。为此，RankE构建了Rank-GAN机制，借助奖励加权回归思想为每个策略生成样本赋予玻尔兹曼权重 $w(\hat{z}_i) \propto \exp(r_i / \tau)$，并在组内归一化至 $\sum_i w(\hat{z}_i) = G$： $$\mathcal{L}_{\text{Rank-GAN}}(\phi) = -\mathbb{E}_{\hat{z} \sim \pi_\theta} \left[ w(\hat{z}) \cdot \text{Disc}(D_\phi(\hat{z})) \right]$$ 通过奖励加权，判别器与解码器的对抗梯度高度聚焦于高奖励的优质解码样本，使黑盒标量评价以非对称梯度更新形式注入解码器。两通道互补融合，既保留了低方差的像素级对齐推力，又具备通用偏好引导能力。

3. 双重流形锚定正则化：真实编码重构与EMA教师一致性蒸馏

解码器若仅在随机策略采样的离散token和对抗奖励下更新，将极易背离预训练建立的高维自然图像流形，发生类似于“奖励作弊”的解码器退化。为此，RankE在解码器更新的M步中设置了两道互补的流形锚点 $\Omega_\phi$。

第一重锚点是真实编码重构约束 $\mathcal{L}_{\text{recon}}(\phi) = \|x_{\text{gt}} - D_\phi(z_{\text{gt}})\|_1 + \mathcal{L}_{\text{GAN}}(D_\phi(z_{\text{gt}}))$，在每个优化步骤中持续将确定性基准特征 $z_{\text{gt}} = \text{Quantize}(E(x_{\text{gt}}))$ 投影回真实像素空间，从全局上防御灾难性流形遗忘。第二重锚点针对离散量化空间不连续引发的高频扰动敏感性，引入缓慢移动的自蒸馏EMA教师解码器 $D_{\phi_{\text{ema}}}$，构建局部一致性损失： $$\mathcal{L}_{\text{consist}}(\phi) = \mathbb{E}_{\hat{z} \sim \pi_\theta} \left[ \mathcal{L}_{\text{LPIPS}}(D_\phi(\hat{z}), \text{sg}[D_{\phi_{\text{ema}}}(\hat{z})]) \right]$$ 该感知一致性约束平滑了极端token扰动下的解码跳跃，过滤单步对抗更新的局部噪声。真实流形重构与策略流形平滑协同工作，牢固约束了解码器的搜索空间。

损失函数 / 训练策略¶

RankE的完整训练流程严格在自回归策略与VQ解码器之间交替迭代。解码器的总优化目标为： $$\mathcal{L}_D(\phi) = \lambda_d \mathcal{L}_{\text{reward}} + \lambda_g \mathcal{L}_{\text{Rank-GAN}} + \lambda_r \mathcal{L}_{\text{recon}} + \lambda_c \mathcal{L}_{\text{consist}}$$ 在实现细节上，模型采用AdamW优化器，自回归策略学习率设为 $1 \times 10^{-5}$，解码器与PatchGAN判别器学习率设为 $5 \times 10^{-5}$。平衡超参数设置为 $\lambda_r = 1.0, \lambda_g = 0.5, \lambda_c = 1.0, \lambda_d = 0.1$，GRPO候选组大小 $G=8$，KL惩罚系数 $\beta=0.05$，重要性采样温度 $\tau=0.1$，EMA衰减系数 $\alpha=0.999$。系统在8张NVIDIA A100 GPU上进行bf16混合精度训练，总共迭代6,000步，耗时约20小时（仅比固定解码器的GRPO基线增加5%时间开销）。

实验关键数据¶

主实验¶

在MS-COCO 30K零样本生成基准与GenEval组合推理基准上，RankE在LlamaGen-XL（775M）与Janus-Pro（1B）上均实现了显著优于基线模型的综合性能。尤其在LlamaGen-XL上，标准RL对齐带来了明显的保真度下降（FID由基线的15.24大幅恶化至17.76），而RankE不仅将CLIP得分由31.54提升至33.76，同时将FID降低至15.21，成功打破了对齐与保真度的互斥困境。

模型骨干	方法 (奖励信号)	解码器状态	MS-COCO CLIP ↑	MS-COCO FID ↓	GenEval 2-Obj ↑	GenEval Count ↑	GenEval Color ↑	GenEval Avg ↑
LlamaGen-XL (775M)	Base (零样本)	Frozen	31.54	15.24	0.331	0.175	0.559	0.309
	SFT	Frozen	31.86	16.58	0.288	0.228	0.723	0.374
	Std. RL (CLIP)	Frozen	32.45	17.76	0.361	0.272	0.779	0.417
	RankE (CLIP, 本文)	Co-Evol	33.76 (+1.31)	15.21 (-2.55)	0.379	0.356	0.763	0.425 (+0.008)
Janus-Pro (1B)	Base (零样本)	Frozen	33.20	18.95	0.810	0.553	0.894	0.740
	SFT	Frozen	33.31	26.73	0.826	0.572	0.889	0.739
	Std. RL (CLIP)	Frozen	33.60	25.59	0.838	0.478	0.891	0.746
	RankE (CLIP, 本文)	Co-Evol	33.86 (+0.26)	25.19 (-0.40)	0.836	0.487	0.892	0.750 (+0.004)

在黑盒偏好奖励（HPSv2）优化实验中，RankE同样展现出一致优势：在LlamaGen-XL上，HPSv2平均偏好分从基线的0.2196及Std. RL的0.2451进一步提升至0.2531，同时GenEval平均分维持在0.423高位。

消融实验¶

在LlamaGen-XL骨干上，作者针对训练范式及解码器损失各组成部分进行了严格的消融对比（CLIP奖励，MS-COCO 30K评估）。

实验组别	配置说明	$\mathcal{L}_{\text{r}}$	$\mathcal{L}_{\text{g}}$	$\mathcal{L}_{\text{c}}$	$\mathcal{L}_{\text{d}}$	CLIP ↑	FID ↓	GenEval Avg ↑
训练模式消融	SFT (无后训练)	-	-	-	-	31.86	16.58	0.374
	仅优化策略 (Policy-Only GRPO)	-	-	-	-	32.45	17.76	0.417
	仅适应解码器 (Decoder-Only)	-	-	-	-	33.41	18.68	0.403
	完整 RankE 协同演化	✓	✓	✓	✓	33.76	15.21	0.425
解码器损失消融	移除真实重构 $\mathcal{L}_{\text{recon}}$	✗	✓	✓	✓	33.26	17.69	0.401
	移除对抗排序 $\mathcal{L}_{\text{Rank-GAN}}$	✓	✗	✓	✓	32.12	18.59	0.397
	移除EMA一致性 $\mathcal{L}_{\text{consist}}$	✓	✓	✗	✓	34.17	19.03	0.408
	移除可微奖励反传 $\mathcal{L}_{\text{reward}}$	✓	✓	✓	✗	32.53	20.16	0.412

关键发现¶

协同优化的非线性增益：单纯优化策略（Policy-Only）导致FID恶化至17.76；单独优化解码器虽然能利用像素先验提升CLIP至33.41，但缺乏策略引导同样使FID恶化至18.68。只有双向协同演变（RankE）能将FID压低至15.21，证明两个模块的交替适应在数学与表征上具有强互补性。
潜在协变量偏移的量化印证：通过跟踪策略token与真实图像token的KL散度 $D_{\text{KL}}(\pi_\theta \parallel p_{\text{real}})$ 发现，标准RL训练过程中散度增加超过24%，且有效码本熵急剧下降；而RankE全程将KL散度稳定在SFT水平附近，并保持了接近真实图像的码本熵（13.87 bits），直接证实了协同演变机制能够实时吸收隐空间分布偏移。
流形锚点的失效模式：去掉EMA一致性约束 $\mathcal{L}_{\text{consist}}$ 会产生病态的对齐虚高（CLIP虚高至34.17但FID崩至19.03），表明解码器对随机单步采样过度拟合；去掉真实图像重构 $\mathcal{L}_{\text{recon}}$ 则引发流形遗忘，导致FID暴跌至17.69。

亮点与洞察¶

跳出离散微分死结的算法视角：面对离散量化操作的不可导屏障，RankE没有陷入Straight-Through Estimator或Gumbel-Softmax等有偏近似的泥潭，而是巧妙借鉴广义EM理念，将联合优化解耦为E步token优势排序与M步像素流形适配，以交替闭环实现了真正的端到端对齐效果。
即插即用的Rank-GAN奖励注入机制：利用重要性加权回归思想将黑盒标量奖励直接转换为判别器梯度权重，解决了无解析导数奖励模型无法指导连续生成流形优化的长期难题。
极佳的算力收益比：整个协同演化框架仅在显存中常驻EMA解码器与判别器（显存峰值由33GB升至56GB），但因交替算子无需额外的反向计算图，端到端耗时仅较常规GRPO增加5%，兼具工程可行性与理论优雅性。

局限与展望¶

显存占用提升与调度开销：在阶段2中维持判别器与EMA解码器参数使得显存峰值有所上涨，未来可探索根据策略奖励增长平台期自适应触发解码器更新的动态调度策略。
对初始SFT数据对齐的依赖：在闭源大模型Janus-Pro上，由于公开SFT语料与模型原始预训练数据分布存在断层，导致SFT阶段基准FID由18.95退化至26.73，限制了后训练阶段的上限突破空间。
VQ编码器仍保持冻结：当前框架为保证真实编码的参照基准而冻结了VQ编码器，探索编码器、自回归策略与解码器的全三元组协同演变将是自回归文生图未来的突破方向。

评分¶

新颖性: ⭐⭐⭐⭐⭐ 首次提出离散自回归文生图的策略-解码器协同后训练范式，GEM理论建模完备。
实验充分度: ⭐⭐⭐⭐⭐ 涵盖两大多模态模型、跨可微与黑盒双重奖励、包含token分布与熵的多重机理诊断。
写作质量: ⭐⭐⭐⭐⭐ 逻辑结构严密，数学推导清晰自洽，深入剖析了潜在协变量偏移的根源。
价值: ⭐⭐⭐⭐⭐ 为离散自回归与统一多模态模型的强化学习后训练开辟了全新的系统级设计空间。