PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Siyong Jian, Siyuan Li, Luyuan Zhang, Zedong Wang, Xin Jin, Ying Li, Cheng Tan, Huan Wang

arXiv:2605.21195v1cs.CV

TL;DR

Discrete AR text-to-image post-training typically updates the policy while freezing the decoder, creating a latent distribution mismatch that can improve alignment while degrading fidelity. RankE co-evolves the policy and decoder through alternating optimization, improving both fidelity and alignment across two backbones and multiple evaluation axes.

  • Problem

    Frozen-decoder policy tuning faces Latent Covariate Shift because sampled policy tokens diverge from the ground-truth codes used to train the decoder.

  • Method

    RankE alternates policy and decoder optimization, using ranking-based alignment objectives and parameter-space stability anchors to co-evolve both components.

  • Results

    RankE consistently improves fidelity and alignment across two backbones, with FID 15.21 and CLIP 33.76 on LlamaGen-XL’s MS-COCO 30K evaluation.

  • Takeaways & Limitations

    Decoder co-evolution overcomes the fidelity–alignment trade-off observed in frozen-decoder baselines for discrete AR text-to-image models.

  • Takeaways & Limitations

    RankE’s gains depend on alignment between the SFT corpus and the backbone’s pre-training distribution, limiting available post-training headroom when alignment is poor.

Abstract

from arXiv · show

Discrete autoregressive (AR) text-to-image (T2I) models pair a VQ tokenizer with an AR policy, and current post-training pipelines optimize only the policy while keeping the VQ decoder frozen. Recent diffusion T2I work, exemplified by REPA-E, has shown that the VAE itself constitutes a key alignment bottleneck, yet no analogous investigation exists for discrete AR models. We show that policy-only optimization induces Latent Covariate Shift: as the policy evolves, the resulting token distribution diverges from the ground-truth distribution on which the decoder was trained, such that reward scores improve while decoded image quality degrades. To address this mismatch, we propose RankE, the first end-to-end post-training framework for discrete T2I generation. Rather than optimizing the policy against a fixed decoder, RankE co-evolves both components through alternating optimization: each module maximizes a ranking-based alignment objective while being regularized by a stability-preserving anchor suited to its parameter space. This co-evolution breaks the fidelity--alignment trade-off that plagues frozen-decoder approaches: on LlamaGen-XL (775M), standard RL improves CLIP but degrades FID, whereas RankE improves both simultaneously (FID 15.21, CLIP 33.76 on MS-COCO 30K). Consistent gains on Janus-Pro (1B) confirm that decoder co-evolution reliably converts reward optimization into pixel-space quality improvements.

1 Introduction

Discrete AR text-to-image post-training typically optimizes the policy while freezing the VQ decoder, causing Latent Covariate Shift as generated token distributions depart from decoder-training codes. RankE jointly co-evolves policy and decoder through alternating ranking-based optimization, improving fidelity and alignment across evaluated models and metrics.

  • Background: Discrete AR T2I models use a VQ tokenizer to map images into discrete codes and an autoregressive policy to model token sequences.This factorization supports unified multimodal architectures and inherits large-language-model scaling infrastructure.
  • Problem: Latent Covariate Shift arises because the decoder is trained on deterministic ground-truth codes but inference feeds it policy-sampled tokens whose distribution diverges under reward optimization.Standard RL widens the token-distribution gap relative to SFT and exacerbates this decoder-side mismatch.
  • Problem: Policy-only GRPO on LlamaGen-XL improves CLIP while degrading FID, exposing a fidelity–alignment trade-off that frozen-decoder tuning cannot resolve.The widening KL divergence against ground-truth token statistics confirms the distributional shift under standard RL.
  • Method: RankE jointly evolves the AR policy and VQ decoder through alternating token-level and pixel-level ranking objectives without differentiating through categorical sampling or VQ quantization.The policy uses group-relative advantages in GRPO, while the decoder uses a reward-weighted adversarial loss.
  • Results: 15.21 FID and 33.76 CLIP are achieved by RankE on LlamaGen-XL (775M) over MS-COCO 30K, improving both fidelity and alignment simultaneously.RankE also consistently improves alignment and zero-shot GenEval on Janus-Pro-1B under the HPSv2 reward.

2 Related Work

Diffusion T2I post-training has advanced through reinforcement learning, preference optimization, and reward fine-tuning, enabled by differentiable denoising chains. REPA-E further extends this paradigm by jointly optimizing the VAE.

  • Post-Training and Alignment in T2I: Diffusion T2I post-training uses online RL, offline preference optimization, and direct reward fine-tuning.These approaches have contributed to the rapid maturation of post-training in the diffusion family.
  • Post-Training and Alignment in T2I: Differentiable denoising chains allow reward gradients to flow from pixel space back to the generator via ∇θR(x).This structural property enables end-to-end optimization in diffusion models but is unavailable in discrete AR models.
  • Post-Training and Alignment in T2I: REPA-E advances diffusion post-training by unlocking the VAE for joint optimization.The passage presents this as a step beyond prior diffusion post-training methods.

3 Method

RankE co-evolves the autoregressive policy and VQ decoder through alternating, ranking-based updates because categorical sampling severs direct gradient flow between them. Policy alignment uses GRPO with KL anchoring, while decoder adaptation uses reward-weighted pixel-space alignment and manifold-preserving regularization, with alternation motivated by a Generalized EM view.

  • Alternating co-evolution: RankE alternates policy and decoder updates, allowing reward information to cross the discrete sampling bottleneck through signals suited to each parameter space.The unified formulation connects this alternation to a Generalized EM procedure.
  • Unified objective: Each module maximizes ranking-based alignment toward reward-favored samples while a parameter-space anchor tethers it to a trusted prior.The method ranks G same-prompt rollouts rather than relying on absolute reward magnitudes.
  • Stage 1: token-level ranking: Stage 1 updates the policy with GRPO advantages from G decoded rollouts and stabilizes it using KL regularization against an EMA reference policy.The group-normalized advantage records both whether a rollout beats its peers and by how much.
  • Stage 2: pixel-level ranking: Stage 2 freezes the policy and adapts the decoder toward reward-preferred decodings using pixel-level ranking, reconstruction, and consistency losses.The ranking signal uses direct reward back-propagation when available and Rank-GAN for black-box rewards.
  • Decoder stability: Alignment-only decoder training risks abandoning the ground-truth manifold, so reconstruction on ground-truth codes preserves fidelity against drift and catastrophic forgetting.Rank-GAN concentrates updates on policy-preferred samples, while the reconstruction anchor retains the deterministic training distribution.

4 Experiments

Experiments evaluate RankE against frozen-decoder post-training across two discrete AR T2I backbones, complementary metrics, and controlled tests. RankE improves alignment while preserving fidelity and generalization by co-evolving the decoder with the policy, absorbing Latent Covariate Shift, and benefiting from coordinated optimization.

  • Experimental setup: Experiments compare RankE with Base, SFT, and frozen-decoder Std. RL on LlamaGen-XL (775M) and Janus-Pro-1B (1B).The evaluation uses fixed reward, data, and compute to test whether decoder co-evolution yields measurable gains.
  • Experimental setup: Evaluation spans FID and CLIP Score on MS-COCO 30K, HPSv2 preference subsets, and zero-shot GenEval compositional reasoning.GenEval covers Two-Object, Counting, and Color binding without task-specific supervision.
  • Controlled comparison: 16.58→17.76 FID under LlamaGen-XL Std. RL shows fidelity degradation, whereas RankE reaches FID 15.21 and CLIP 33.76 under CLIP-based optimization.On Janus-Pro-1B, RankE achieves the best CLIP and zero-shot GenEval, although all post-training variants regress FID versus Base.
  • Qualitative verification: RankE produces faithful prompt attributes and high perceptual quality without the artifact bands observed for standard RL.Standard RL improves adherence but introduces visible artifacts, while the base model often misses color, count, and spatial-relational attributes.
  • Mechanism verification: +24% KL divergence under standard RL contrasts with RankE remaining near SFT initialization, while RankE preserves entropy near the real-image level of 13.87 bits.These diagnostics provide evidence that co-evolution absorbs Latent Covariate Shift rather than merely regularizing its downstream effects.
  • Ablations: Full RankE is best on both FID and GenEval and matches the best CLIP, outperforming policy-only and decoder-only adaptation in the component ablation.Decoder-only adaptation raises CLIP to 33.41 but degrades FID to 18.68, whereas full RankE achieves FID 15.21.

5 Conclusion … A.3 The Gradient Barrier in Discrete AR

RankE addresses Latent Covariate Shift in discrete autoregressive text-to-image post-training by co-evolving the AR policy and VQ decoder rather than differentiating through discrete bottlenecks. The appendix grounds this alternating design theoretically, diagnostically, and reproducibly, while situating it among tokenizer and AR-generator research.

  • 5 Conclusion: RankE co-evolves the AR policy and VQ decoder through a Generalized EM procedure to absorb Latent Covariate Shift during post-training.The shift arises because decoders train on deterministically quantized ground-truth codes, while inference samples tokens from a stochastic reward-driven policy.
  • 5 Conclusion: RankE’s resource overhead mainly comes from holding the discriminator and EMA decoder, motivating adaptive scheduling or less frequent decoder updates.The conclusion also identifies sensitivity to SFT corpus alignment, joint encoder training, pre-training co-evolution, online human feedback, and multi-objective rewards as future directions.
  • Roadmap: The appendix progresses from theoretical grounding through empirical validation to implementation details.Its three parts cover foundations, mechanism diagnostics and robustness, and reproducibility.
  • Part I | Theoretical Foundations: Part I formalizes RankE’s GEM convergence guarantee, positioning alternating co-evolution as principled rather than heuristic.This section also explains where RankE sits in the literature and why its alternation is principled rather than ad hoc.
  • A Extended Related Work: RankE’s alternating design is motivated by discrete visual tokenizers, autoregressive generator factorizations, and the gradient barrier between policy and pixels.The related-work appendix expands the paper’s background along these three axes.
  • A.1 Discrete Visual Tokenizers: Discrete visual tokenizers compress images into latent-code grids, with vector quantization learning finite codebooks as in VQ-VAE and VQGAN.Recent work targets quantization fidelity, regularization, and representation structure.
  • A.2 AR-based Image Generators: Autoregressive image generators represent images as token sequences under varied factorizations, including raster-order, next-scale, masked-parallel, and randomized decoding.These approaches include DALL-E, CogView, Parti, VAR, MaskGIT, MUSE, and RandAR.
  • A.3 The Gradient Barrier in Discrete AR: Discrete AR reward back-propagation is blocked by categorical sampling and VQ argmax, while STE is biased and unstable and Gumbel-Softmax degrades at large vocabularies.Unlike diffusion’s smooth reparameterization, RankE alternates around the discrete bottleneck, using discrete-friendly RL for the policy and continuous gradients for the decoder.

B Generalized EM Formal Derivation · Part II | Mechanism Diagnostics and Robustness · C Latent Covariate Shift: Measurement Protocol

RankE admits a generalized-EM interpretation in which policy and decoder updates stochastically ascend one MAP-augmented ELBO, while latent covariate shift is measured from token statistics and worsens monotonically during policy training.

  • B Generalized EM Formal Derivation: RankE alternates θ E-steps and ϕ M-steps as stochastic ascent on the same MAP-augmented evidence lower bound.The decoder prior anchors parameters to the pre-training manifold through reconstruction and consistency regularization.
  • B.1 Setup and the Optimality Variable: log Z(y; ϕ) is the marginal optimality likelihood whose maximization produces high-reward outputs under prior πref.Here Z(y; ϕ) = Ez∼πref[exp(r/β)] and remains finite because r is upper-bounded.
  • B.3 E-step: GRPO as ELBO Ascent on θ: Advantage normalization preserves the expected ascent direction while reducing variance, and PPO clipping constrains importance-ratio updates to a trust region.For ϵ=0.2, the expected ascent direction matches the unclipped gradient despite finite-step bias of order ϵ.
  • B.3 E-step: GRPO as ELBO Ascent on θ: The KL penalty βDKL(πθ∥πref) supplies ELBO regularization, with EMA reference πref serving as a slowly varying latent prior.Thus GRPO is a stochastic, variance-reduced, trust-region-clipped ascent step on ∇θL.
  • B.4 M-step: Decoder Update as MAP Ascent on ϕ: For black-box rewards, Rank-GAN recovers an importance-weighted surrogate using w(ˆzi) ∝exp(ri/τ), where τ softens reward weighting.The reward-tilted posterior is p∗(z | y) ∝ πref(z | y) exp(r/β), and πθ ≈p∗ at optimum.
  • B.5 Convergence: Under Robbins–Monro learning-rate conditions and bounded gradient variance, alternating updates converge to a stationary point of L.Each stage need only improve L given the other, yielding the generalized-EM guarantee.
  • C Latent Covariate Shift: Measurement Protocol: DKL increases monotonically over training, quantifying worsening latent covariate shift and motivating decoder adaptation in RankE.Policy sequences are evaluated at checkpoints t ∈{500, 1000, . . . , 6000} with fixed sampling settings.

D Training Dynamics and Convergence Behavior · D.1 Per-Step Dynamics under CLIP and HPSv2 · D.2 Convergence across Co-Evolution Rounds

RankE exhibits stable, reward-agnostic per-step dynamics under CLIP and HPSv2, while co-evolution rounds show monotonic reward gains, transient E-step FID degradation, and M-step recovery. These dynamics support decoder recalibration to the policy’s evolving distribution while keeping KL bounded.

  • D.1 Per-Step Dynamics under CLIP and HPSv2: Fig. D.1 tracks six diagnostics over 3,000 steps under CLIP and HPSv2 with identical hyperparameters.The two reward settings exhibit qualitatively identical convergence behavior across all panels.
  • D.1 Per-Step Dynamics under CLIP and HPSv2: Policy-sampled reconstruction loss stays near zero under HPSv2 and briefly spikes around steps 1,000–2,000 before recovering under CLIP.Ground-truth reconstruction remains 0.3–0.4, with HPSv2 marginally higher, indicating a mild trade-off rather than destabilization.
  • D.1 Per-Step Dynamics under CLIP and HPSv2: The discriminator reaches ∼0.5 within 500 steps, declining to ∼0.35 under CLIP.Consistency loss stays below 0.015, consistent with EMA-anchored consistency distillation preventing abrupt decoder changes.
  • D.2 Convergence across Co-Evolution Rounds: Across co-evolution rounds, rewards rise monotonically.This pattern is reported alongside FID recovery after E-step degradation and bounded KL under GRPO regularization.
  • D.2 Convergence across Co-Evolution Rounds: FID transiently worsens during each E-step because of Latent Covariate Shift, then recovers during the subsequent M-step.The recovery is visible in Fig. D.1(b) and Fig. 4(b) of the main paper.
  • D.2 Convergence across Co-Evolution Rounds: KL stays bounded under GRPO regularization across co-evolution rounds.This boundedness accompanies monotonically rising rewards and alternating FID degradation and recovery.
  • D Training Dynamics and Convergence Behavior: Each M-step recalibrates the decoder to the policy’s evolving distribution, sustaining a virtuous improvement cycle.The reported round-level patterns confirm this core mechanism.

E Extended Ablations and Sensitivity Studies · E.1 Consistency Distillation Weight λc

RankE is broadly robust across reasonable hyperparameter ranges, but excessive consistency weight and sequential K=1 training are sensitive regimes. For λc, the default λc=1 provides the most stable training, while λc=50 causes decoder collapse despite λc=10 achieving mild drift and the best FID/GenEval.

  • E Extended Ablations and Sensitivity Studies: RankE remains robust across reasonable ranges of consistency weight, IS temperature, and EMA decay, with sensitivity concentrated in sequential K=1 training and excessive λc.The sweep uses LlamaGen-XL with CLIP reward on MS-COCO 30K.
  • E.1 Consistency Distillation Weight λc: λc is evaluated at λc=10 (Run A), λc=1 (Run B, default), and λc=50 (Run C), with dynamics and final metrics reported separately.The ablation compares three consistency-distillation weights.
  • E.1 Consistency Distillation Weight λc: Run A with λc=10 introduces mild drift but achieves the best FID/GenEval.This setting is less stable than Run B but leads the reported FID/GenEval comparison.
  • E.1 Consistency Distillation Weight λc: Training dynamics track GAN reconstruction losses, discriminator and distillation consistency losses, policy KL divergence, and reward curves under CLIP and HPSv2.The figure distinguishes generated-image and ground-truth reconstruction losses and uses dual y-axes for reward curves.
  • E.1 Consistency Distillation Weight λc: Run B with λc=1 yields the most stable training and is adopted for overall stability.The selected setting is the default Run B.
  • E.1 Consistency Distillation Weight λc: Run C with λc=50 causes decoder collapse around step 1,500, as GT reconstruction diverges and the discriminator drops to near zero.The collapse indicates that excessive consistency overwhelms the adversarial signal.
  • E.1 Consistency Distillation Weight λc: Table E.1 reports the final metrics for the λc ablation.The table is associated with the three runs λc=10, λc=1, and λc=50.

E.2 Importance Sampling Temperature τ · E.3 EMA Decay Rate α · Part III | Reproducibility

RankE is most effective with moderate importance-sampling temperature and a slowly updated EMA teacher, balancing reward signal, diversity, stability, and adaptability. The reproducibility materials provide the hyperparameters, pseudocode, and data pipeline needed to reproduce the end-to-end method.

  • E.2 Importance Sampling Temperature τ: τ=0.1 achieves the best performance across all three metrics by balancing reward concentration and sampling diversity.High τ approaches uniform sampling and dilutes the reward signal, whereas overly low τ causes hard selection and degrades FID to 16.12.
  • E.2 Importance Sampling Temperature τ: Table E.2 evaluates RankE’s sensitivity to the importance-sampling temperature τ.The evaluation compares the temperature choices underlying the reported reward-signal and diversity trade-off.
  • E.3 EMA Decay Rate α: α=0.999 achieves optimal results across all three metrics by striking the best stability–adaptability trade-off.With α=0.900, the teacher tracks the student too closely, yielding FID 15.75 and CLIP 33.48.
  • E.3 EMA Decay Rate α: λc=50 causes decoder collapse, while moderate λc values remain stable in the training-dynamics ablation.Run A uses λc=10, Run B uses λc=1 as the default, and Run C uses λc=50.
  • E.3 EMA Decay Rate α: Table E.3 evaluates how the EMA decay rate α affects RankE’s performance.The ablation identifies the decay-rate choice that governs teacher updating and the stability–adaptability trade-off.
  • Part III | Reproducibility: The reproducibility section provides everything needed to reproduce RankE end-to-end: hyperparameters, pseudocode, and the data pipeline.These materials cover the method’s implementation and training workflow.

F Implementation Details and Hyperparameters

RankE alternates policy alignment and decoder adaptation using specified optimizers, reward models, and VQ tokenizer components. The method adds modest memory and runtime requirements while avoiding extra kernel-level or infrastructure cost.

  • Alternating optimization: Stage 1 trains the AR policy with GRPO at 1×10−5, while Stage 2 updates the decoder with AdamW at 5×10−5.Stage 2 uses a PatchGAN discriminator and an EMA teacher with decay 0.999; Stage 1 schedules β against an EMA reference policy.
  • VQ tokenizer architecture: The shared VQ-VAE uses 16× spatial downsampling, a 16,384-entry codebook with embedding dimension 8, and produces 256 tokens from a 256×256 input.Its convolutional ResNet uses channel multipliers [1] [2] [4] and self-attention at 16×16.
  • Reward models: CLIP Score and HPSv2 are frozen reward models whose gradients pass through the decoder reward loss but stop before non-differentiable sampling.CLIP uses ViT-g/14 cosine similarity scaled by 100, while HPSv2 is trained on large-scale pairwise human judgments.
  • Compute footprint: 56 GB peak VRAM and approximately 20 hours for 6k steps characterize RankE, versus 33 GB and 19 hours for the frozen-decoder GRPO baseline.The alternating update scheme adds no extra kernel-level or infrastructure cost.

G Training Algorithm

RankE trains the discrete text-to-image system end to end by alternating optimization of the autoregressive policy and decoder refinement under the multi-objective loss in Eq. 4.

  • G Training Algorithm: RankE alternates policy optimization for the AR generator with decoder refinement under the multi-objective loss in Eq. 4.Algorithm 1 summarizes this end-to-end training procedure.

H Training Data and Curation

RankE constructs a 15K training set from BLIP3o-60k and curates it for reward-based post-training in two stages. Long synthetic captions are compressed into concise visual tags to avoid CLIP context truncation while preserving reward-relevant semantics.

  • Training data: RankE builds a 15K training set from BLIP3o-60k image–text pairs curated across diverse scenes, objects, and human gestures.The source pairs were curated with GPT-4o.
  • Training data: Long synthetic captions can exceed CLIP’s 77-token context window, causing reward truncation and unstable gradients.The affected captions come from synthetic sources including DALL-E 3 and JourneyDB.
  • Caption summarization pipeline: Qwen2.5-Instruct compresses captions into visual tags of at most 50 words while preserving subjects, styles, lighting, colors, and composition.The pipeline removes narrative prose, subjective descriptions, and filler words.
  • Caption summarization pipeline: Caption compression preserves reward-relevant semantic content while keeping the full caption within CLIP’s effective context window.This addresses the context-length issue introduced by long synthetic captions.
Loading 2605.21195v1…