VA-$π$: Variational Policy Alignment for Pixel-Aware Autoregressive Generation

TL;DR

This work tackles the misalignment between pixel-space fidelity and token-level autoregressive training in visual generation. It introduces VA-$\pi$, a post-training, variational-policy framework that couples a pixel-space reconstruction objective with a token-space prior through an ELBO, and optimizes the AR generator with a reconstruction-based reinforcement learning reward via GRPO. VA-$\pi$ achieves large gains in image fidelity and diversity on class-to-image and text-to-image tasks using only about 1% of data and modest compute, without external reward models, and scales to large multimodal models. The approach provides a practical, theoretically grounded method to bridge tokenized representations and pixel-level distribution, enabling efficient, high-quality visual generation across tasks.

Abstract

Autoregressive (AR) visual generation relies on tokenizers to map images to and from discrete sequences. However, tokenizers are trained to reconstruct clean images from ground-truth tokens, while AR generators are optimized only for token likelihood. This misalignment leads to generated token sequences that may decode into low-quality images, without direct supervision from the pixel space. We propose VA-$π$, a lightweight post-training framework that directly optimizes AR models with a principled pixel-space objective. VA-$π$ formulates the generator-tokenizer alignment as a variational optimization, deriving an evidence lower bound (ELBO) that unifies pixel reconstruction and autoregressive modeling. To optimize under the discrete token space, VA-$π$ introduces a reinforcement-based alignment strategy that treats the AR generator as a policy, uses pixel-space reconstruction quality as its intrinsic reward. The reward is measured by how well the predicted token sequences can reconstruct the original image under teacher forcing, giving the model direct pixel-level guidance without expensive free-running sampling. The regularization term of the ELBO serves as a natural regularizer, maintaining distributional consistency of tokens. VA-$π$ enables rapid adaptation of existing AR generators, without neither tokenizer retraining nor external reward models. With only 1% ImageNet-1K data and 25 minutes of tuning, it reduces FID from 14.36 to 7.65 and improves IS from 86.55 to 116.70 on LlamaGen-XXL, while also yielding notable gains in the text-to-image task on GenEval for both visual generation model (LlamaGen: from 0.306 to 0.339) and unified multi-modal model (Janus-Pro: from 0.725 to 0.744). Code is available at https://github.com/Lil-Shake/VA-Pi.

Figures (14)

• Figure 1: Pixel-Aware Alignment via VA-$\boldsymbol{\pi}$. VA-$\pi$ enables efficient post-training via variational policy optimization, aligning the pixel-space distribution of AR generated images with that of ground-truth images.
• Figure 2: Overview of VA-$\boldsymbol{\pi}$. VA-$\boldsymbol{\pi}$ aligns the visual AR model with tokenizer via variational optimization. Given a reference image and its ground-truth tokens, VA-$\pi$ adds context noise and lets the AR model compute logits under teacher forcing and samples target tokens. These sampled tokens are decoded back into an image, and the reconstruction reward is defined against the reference image. This reward is then used for policy updates within an RL framework such as GRPO yuan2025argrpo. Additionally, a likelihood regularization using cross-entropy loss between the logits and ground-truth tokens is retained to preserve the model’s original next-token prediction ability.
• Figure 3: Left: Qualitative comparison of C2I generation among LlamaGen-XL llamagen (top), AR-GRPO yuan2025argrpo (middle) and VA-$\boldsymbol{\pi}$ (bottom) on the ImageNet-1k imagenet classes. Both models use a CFG scale of 2.0. VA-$\pi$ produces clearer object structures (like the car mirror) than LlamaGen-XL (top) and AR-GRPO (middle), demonstrating that pixel-space alignment encourages realistic generations. Right: Qualitative comparison of T2I generation between Janus-Pro 1B januspro and VA-$\boldsymbol{\pi}$ on the GenEval Benchmark geneval. Both models use a CFG scale of 5.0. VA-$\pi$ produces better object combination and counting accuracy, demonstrating stronger capability.
• Figure 4: Ablation on regularization weight (w/o cfg). CE regularization consistently outperforms KL regularization on FID and IS. Moderate CE regularization ($0.1$) provides the best results.
• Figure 6: Qualitative comparison of C2I generation among LlamaGen-XXL llamagen, post-train tokenizer (PT), STE based post-train AR and VA-$\boldsymbol{\pi}$ on the ImageNet-1k imagenet classes. Both models use a CFG scale of 1.0. VA-$\pi$ shows better semantic alignment and image quality, demonstrating that pixel-space alignment encourages realistic generations.