Learning Discrete Autoregressive Priors with Wasserstein Gradient Flow

• URL: http://arxiv.org/abs/2605.06148v1
• Date: Thu, 07 May 2026 12:41:46 GMT
• Title: Learning Discrete Autoregressive Priors with Wasserstein Gradient Flow
• Abstract summary: Discrete image tokenizers are commonly trained in two stages: first for reconstruction, and then with a prior model fitted to the frozen token sequences.<n>We analyze this mismatch using Tripartite Variational Consistency (TVC), which decomposes latent-variable learning into three consistency conditions.<n>We optimize this signal with a Wasserstein-gradient-flow update.
• Score: 17.114307498606724
• License: http://creativecommons.org/licenses/by/4.0/
• Abstract: 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.

Related papers

• Rectify Then Diffuse: Disentangling Concepts Before Denoising Trajectory Unfolds [18.12683458124076]
  Text-to-image diffusion models often omit or merge concepts incorrectly with multiple concepts.<n>We propose Rectify-then-Diffuse (RTD), a training-free framework that rectifies the initial allocation once before standard denoising.<n>RTD achieves state-of-the-art compositional fidelity and robust gains.
  arXiv  Detail & Related papers  (2026-08-04T05:03:17Z)
• GEAR: Guided End-to-End AutoRegression for Image Synthesis [45.98061600562716]
  GEAR trains a vector-quantized (VQ) tokenizer and an autoregressive (AR) generator jointly and end-to-end.<n>GEAR speeds up ImageNet gFID convergence by up to 10x relative to the strong LlamaGen-REPA baseline.
  arXiv  Detail & Related papers  (2026-06-30T17:59:57Z)
• Let Triggers Control: Frequency-Aware Dropout for Effective Token Control [8.72880783870241]
  We propose Frequency-Aware Dropout (FAD) to improve controllability without adding new parameters.<n>FAD consists of two key components: co-occurrence analysis and curriculum-inspired scheduling.<n>Our method provides a simple yet effective dropout strategy that enhances controllability and personalization in text-to-image generation.
  arXiv  Detail & Related papers  (2026-03-28T08:55:54Z)
• CaTok: Taming Mean Flows for One-Dimensional Causal Image Tokenization [122.88484422855934]
  We present CaTok, a 1D causal image tokenizer with a MeanFlow decoder.<n>By selecting tokens over time intervals, CaTok learns causal 1D representations that support both fast one-step generation and high-fidelity multi-step sampling.<n> Experiments demonstrate that CaTok achieves state-of-the-art results on ImageNet reconstruction, reaching 0.75 FID, 22.53 PSNR and 0.674 SSIM.
  arXiv  Detail & Related papers  (2026-03-06T16:39:17Z)
• ResTok: Learning Hierarchical Residuals in 1D Visual Tokenizers for Autoregressive Image Generation [64.84095852784714]
  Residual Tokenizer (ResTok) is a 1D visual tokenizer that builds hierarchical residuals for both image tokens and latent tokens.<n>We show that restoring hierarchical residual priors in visual tokenization significantly improves AR image generation, achieving a gFID of 2.34 on ImageNet-256 with only 9 sampling steps.
  arXiv  Detail & Related papers  (2026-01-07T14:09:18Z)
• REAR: Rethinking Visual Autoregressive Models via Generator-Tokenizer Consistency Regularization [130.46612643194973]
  reAR is a simple training strategy introducing a token-wise regularization objective.<n>On ImageNet, it reduces gFID from 3.02 to 1.86 and improves IS to 316.9 using a standardization-based tokenizer.<n>When applied to advanced tokenizers, it achieves a gFID of 1.42 with only 177M parameters, matching the performance with larger state-of-the-art diffusion models (675M)
  arXiv  Detail & Related papers  (2025-10-06T02:48:13Z)
• Rethinking Discrete Tokens: Treating Them as Conditions for Continuous Autoregressive Image Synthesis [79.98107530577576]
  DisCon is a novel framework that reinterprets discrete tokens as conditional signals rather than generation targets.<n>DisCon achieves a gFID score of 1.38 on ImageNet 256$times $256 generation, outperforming state-of-the-art autoregressive approaches by a clear margin.
  arXiv  Detail & Related papers  (2025-07-02T14:33:52Z)
• Bridging Continuous and Discrete Tokens for Autoregressive Visual Generation [85.82112629564942]
  We propose TokenBridge, which maintains the strong representation capacity of continuous tokens while preserving the modeling simplicity of discrete tokens.<n>We introduce a dimension-wise quantization strategy that independently discretizes each feature dimension, paired with a lightweight autoregressive prediction mechanism.<n>Our approach achieves reconstruction and generation quality on par with continuous methods while using standard categorical prediction.
  arXiv  Detail & Related papers  (2025-03-20T17:59:59Z)
• Frequency Autoregressive Image Generation with Continuous Tokens [31.833852108014312]
  We introduce the frequency progressive autoregressive (textbfFAR) paradigm and instantiate FAR with the continuous tokenizer.<n>We demonstrate the efficacy of FAR through comprehensive experiments on the ImageNet dataset.
  arXiv  Detail & Related papers  (2025-03-07T10:34:04Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.