arXiv:2605.06148cs.CVcs.AI2026-05

让图像分块模型在训练时就学会匹配生成模型的分布,提升自回归生成效果。

Learning Discrete Autoregressive Priors with Wasserstein Gradient Flow

论文配图:Learning Discrete Autoregressive Priors with Wasserstein Gradient Flow
图 1 · 摘自论文原文
  • 训练分块模型时加入分布匹配信号,用Wasserstein梯度流优化
  • 在CIFAR-10和ImageNet上降低自回归损失并提升生成质量(FID)
  • 无需反向传播,仅需两个自回归模型前向计算,适合高效部署

离散图像分块器通常分两阶段训练:先重建图像,再用固定分块序列拟合先验模型。这种分离使分块器不了解后续生成模型的需求,导致分块虽保留图像信息,却难被自回归先验从左到右预测。我们通过三重变分一致性(TVC)分析此不匹配问题,发现两阶段训练虽保持重建性能,但未包含先验一致性目标——分块分布提前固定,自回归先验无法参与优化。为此,我们在分块器训练中引入分布级先验匹配信号,并使用Wasserstein梯度流进行优化。对于硬分类别分块,该更新简化为一个辅助自回归模型与目标自回归先验之间的逐标记对比,仅需两次前向传播,无需反向传播。所提wAR-Tok分块器在保持重建质量的同时,在CIFAR-10和ImageNet上显著降低自回归损失并提升生成FID。

原文摘要 · Abstract (English)

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.

自回归生成分块模型扩散模型分布对齐

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。