Technical Summary: Guided End-to-End AutoRegression for Image Synthesis (GEAR)

1. Problem Statement

Modern visual generative models, including both autoregressive (AR) transformers and diffusion models, typically follow a two-stage training paradigm:

1. Tokenizer Training: A vector-quantized (VQ) tokenizer (e.g., VQ-VAE) is trained to reconstruct images and then frozen.
2. Generator Training: An AR generator is trained on the discrete indices produced by the frozen tokenizer.

This decoupling creates a fundamental misalignment: the tokenizer is optimized purely for reconstruction fidelity, oblivious to whether the resulting discrete sequence is easy for the downstream generator to model causally. While recent work in diffusion models (e.g., REPA-E) has successfully explored end-to-end joint training by aligning latent representations with external semantic encoders (like DINOv2), applying this to discrete AR generation is significantly harder.

The core obstacle is non-differentiability. The mapping from a continuous latent to a discrete VQ index involves an arg max operation. Consequently, gradients from the AR generator cannot flow back to the tokenizer. Attempts to bridge this gap using a Straight-Through Estimator (STE) have proven unstable, often causing the codebook to collapse (where the tokenizer emits only a few dominant codes to maximize predictability at the cost of reconstruction fidelity), leading to catastrophic failure (gFID $\approx$ 105).

2. Methodology: GEAR

The authors propose GEAR (Guided End-to-end AutoRegression), a framework that jointly trains a VQ tokenizer and an AR generator end-to-end without collapsing the codebook.

Core Mechanism: Dual Read-Out

GEAR resolves the non-differentiability issue by introducing a dual read-out of the codebook assignment for every position:

1. Hard Branch (One-Hot):
   • Uses the standard discrete index ($q_i = \text{arg max}_k A_{ik}$).
   • Feeds discrete tokens into the AR model for Next-Token Prediction (NTP) and a hard representation alignment loss.
   • Crucially: This branch is non-differentiable with respect to the tokenizer. Gradients from NTP and the hard alignment loss update only the AR model, not the tokenizer. This prevents the NTP loss from forcing the tokenizer to collapse into low-entropy codes.
2. Soft Branch (Differentiable):
   • Uses a temperature-weighted softmax interpolation over the codebook entries ($\pi_i = \text{softmax}(A_i / \tau)$).
   • Feeds a soft, continuous embedding into the AR model.
   • Crucially: This branch is fully differentiable. It carries a representation alignment loss (e.g., against DINOv2 features) that flows back through the AR backbone to update only the tokenizer.

Optimization Strategy

The two modules are updated with disjoint objectives:

• Tokenizer ($\theta_{tok}$): Updated by the standard VQ reconstruction loss ($L_{VQ}$) plus the soft alignment loss ($L^s_{align}$). The tokenizer is guided to reshape its index distribution to be more predictable and semantically aligned, but not by the prediction loss itself.
• AR Generator ($\theta_{AR}$): Updated by the NTP loss ($L_{NTP}$) plus the hard alignment loss ($L^h_{align}$).

This design ensures the tokenizer learns to emit tokens that the AR model can easily predict and align with semantic features, while maintaining high reconstruction fidelity.

3. Key Contributions

• Guided End-to-End Training: GEAR introduces a differentiable soft-assignment bridge that allows the AR model's representation alignment objective to guide the tokenizer, overcoming the non-differentiable index barrier that defeats standard STE approaches.
• Shifted Alignment Burden: Unlike diffusion-side methods (e.g., REPA-E, VA-VAE) that make the latent itself more semantic (DINOv2-like), GEAR shifts the alignment burden to the AR generator.
  • The tokenizer's features become less DINOv2-like, re-organizing its index distribution toward a more predictable, lower-entropy usage.
  • The AR generator's hidden states become more DINOv2-like, particularly at the patch level, capturing rich local spatial structure.
• Generality: The mechanism is agnostic to the underlying quantizer, demonstrating improvements across VQVAE, LFQ, and IBQ. It also transfers effectively from class-conditional generation to text-to-image synthesis.

4. Experimental Results

Experiments were conducted on ImageNet-1K (256x256) and the GPIC text-to-image benchmark.

• Convergence Speed: GEAR accelerates ImageNet gFID convergence by up to 10$\times$ compared to the strong LlamaGen-REPA baseline. A naive end-to-end variant using STE diverges completely.
• Generation Quality:
  • ImageNet: At 300 epochs, GEAR (775M params) achieves a gFID of 2.52 (with CFG), outperforming LlamaGen-REPA (2.68) and approaching the performance of strong latent diffusion models (e.g., REPA-E at 1.12, though REPA-E uses a continuous VAE with higher reconstruction fidelity).
  • Text-to-Image: On GPIC, GEAR consistently outperforms LlamaGen-REPA across all training steps (50k to 390k), achieving lower Fréchet distances in DINOv2 feature space (FDD).
• Reconstruction: Unlike methods that sacrifice reconstruction for semantics, GEAR improves or maintains reconstruction quality (rFID) while training end-to-end.
• Ablations:
  • STE Failure: Replacing the soft bridge with STE causes codebook collapse (gFID $\approx$ 105).
  • Temperature: A temperature of $\tau=0.1$ balances generation quality and reconstruction.
  • Model Size: GEAR improves performance across Base, Large, and XLarge scales, with the co-trained tokenizer showing improved reconstruction as the AR model size increases.

5. Significance and Claims

The paper claims that GEAR represents a significant step toward unified, end-to-end visual generation for discrete token pipelines.

• Overcoming the Discrete Bottleneck: GEAR demonstrates that end-to-end training is possible for discrete AR models without the instability of STE, provided the prediction loss is decoupled from the tokenizer update.
• New Paradigm for Alignment: The paper highlights a counter-intuitive finding: for discrete AR models, the optimal strategy is not to make the tokenizer semantic. Instead, the tokenizer should be optimized for predictability (lower entropy, structured indices), while the semantic alignment emerges within the AR model's hidden states. This "local, patch-level structure" is what makes next-token prediction efficient.
• Foundation for Unified Systems: The authors posit that discrete next-token formulations are uniquely suited for unified understanding and generation systems (e.g., handling long contexts and multimodal inputs) because they bound per-step error and inherit mature alignment techniques from Large Language Models (RLHF, DPO). GEAR provides a robust substrate for scaling these systems.

The authors acknowledge that while GEAR narrows the gap between AR and diffusion models, the reconstruction ceiling remains a limitation due to the discrete nature of VQ tokenizers compared to continuous VAEs used in top-tier diffusion models. Future work is suggested to explore decoupling compression from sequence length (e.g., via patchified AR) to further improve fidelity.