# Learning Discrete Autoregressive Priors with Wasserstein Gradient Flow

arXiv: https://arxiv.org/abs/2605.06148 · HTML: https://mlsanity.com/paper/2605.06148/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 30 of its 46 references that are in the map.
- First posted on arXiv 7 May 2026.

## Builds on (most cited first, 30 of 30)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [High-Resolution Image Synthesis with Latent Diffusion Models](https://mlsanity.com/paper/2112.10752/) (2021)
- [The Unreasonable Effectiveness of Deep Features as a Perceptual Metric](https://mlsanity.com/paper/1801.03924/) (2018)
- [Diffusion Models Beat GANs on Image Synthesis](https://mlsanity.com/paper/2105.05233/) (2021)
- [Masked Autoencoders Are Scalable Vision Learners](https://mlsanity.com/paper/2111.06377/) (2021)
- [Neural Discrete Representation Learning](https://mlsanity.com/paper/1711.00937/) (2017)
- [Zero-Shot Text-to-Image Generation](https://mlsanity.com/paper/2102.12092/) (2021)
- [Taming Transformers for High-Resolution Image Synthesis](https://mlsanity.com/paper/2012.09841/) (2020)
- [Categorical Reparameterization with Gumbel-Softmax](https://mlsanity.com/paper/1611.01144/) (2016)
- [Variational Inference with Normalizing Flows](https://mlsanity.com/paper/1505.05770/) (2015)
- [Estimating or Propagating Gradients Through Stochastic Neurons for Conditional Computation](https://mlsanity.com/paper/1308.3432/) (2013)
- [The Concrete Distribution: A Continuous Relaxation of Discrete Random Variables](https://mlsanity.com/paper/1611.00712/) (2016)
- [One-step Diffusion with Distribution Matching Distillation](https://mlsanity.com/paper/2311.18828/) (2023)
- [Conditional Image Generation with PixelCNN Decoders](https://mlsanity.com/paper/1606.05328/) (2016)
- [Autoregressive Model Beats Diffusion: Llama for Scalable Image Generation](https://mlsanity.com/paper/2406.06525/) (2024)
- [SoundStream: An End-to-End Neural Audio Codec](https://mlsanity.com/paper/2107.03312/) (2021)
- [Generating Sentences from a Continuous Space](https://mlsanity.com/paper/1511.06349/) (2015)
- [Stein Variational Gradient Descent: A General Purpose Bayesian Inference Algorithm](https://mlsanity.com/paper/1608.04471/) (2016)
- [VAE with a VampPrior](https://mlsanity.com/paper/1705.07120/) (2017)
- [Diff-Instruct: A Universal Approach for Transferring Knowledge From Pre-trained Diffusion Models](https://mlsanity.com/paper/2305.18455/) (2023)
- [InfoVAE: Information Maximizing Variational Autoencoders](https://mlsanity.com/paper/1706.02262/) (2017)
- [FlexTok: Resampling Images into 1D Token Sequences of Flexible Length](https://mlsanity.com/paper/2502.13967/) (2025)
- [ImageFolder: Autoregressive Image Generation with Folded Tokens](https://mlsanity.com/paper/2410.01756/) (2024)
- [One-D-Piece: Image Tokenizer Meets Quality-Controllable Compression](https://mlsanity.com/paper/2501.10064/) (2025)
- [LARP: Tokenizing Videos with a Learned Autoregressive Generative Prior](https://mlsanity.com/paper/2410.21264/) (2024)
- [Flow to the Mode: Mode-Seeking Diffusion Autoencoders for State-of-the-Art Image Tokenization](https://mlsanity.com/paper/2503.11056/) (2025)
- [Divergence Triangle for Joint Training of Generator Model, Energy-based Model, and Inference Model](https://mlsanity.com/paper/1812.10907/) (2018)
- [Selftok: Discrete Visual Tokens of Autoregression, by Diffusion, and for Reasoning](https://mlsanity.com/paper/2505.07538/) (2025)
- [Complexity Matters: Rethinking the Latent Space for Generative Modeling](https://mlsanity.com/paper/2307.08283/) (2023)
- [When Worse is Better: Navigating the compression-generation tradeoff in visual tokenization](https://mlsanity.com/paper/2412.16326/) (2024)

## Abstract (arXiv)

Discrete image tokenizers are commonly trained in two stages: first for reconstruction, and then with a prior model fitted to the frozen token sequences. This decoupling leaves the tokenizer unaware of the model that will later generate its tokens. As a result, the learned tokens may preserve image information well but still be difficult for an autoregressive (AR) prior to predict from left to right. We analyze this mismatch using Tripartite Variational Consistency (TVC), which decomposes latent-variable learning into three consistency conditions: conditional-likelihood consistency, prior consistency, and posterior consistency. TVC shows that two-stage training preserves the reconstruction side but leaves prior consistency outside the tokenizer objective: the overall token distribution is fixed before the AR prior participates in training. Motivated by this view, we add a distribution-level prior-matching signal during tokenizer training, while keeping the reconstruction objective unchanged. We optimize this signal with a Wasserstein-gradient-flow update. For hard categorical tokens, the update reduces to a token-level contrast between an auxiliary AR model that tracks the tokenizer's current token distribution and the target AR prior. It requires only forward passes through the two AR models and does not backpropagate through either of them. The resulting tokenizer, wAR-Tok, reduces AR loss and improves generation FID on CIFAR-10 and ImageNet at comparable reconstruction quality.