基于 Wasserstein 梯度流学习离散自回归先验
计算机视觉与模式识别
2026-05-08 v1 人工智能
机器学习
摘要
离散图像分词器通常分两阶段训练：首先进行重建，然后拟合一个先验模型到冻结的 token 序列上。这种解耦使得分词器无法感知后续将生成其 token 的模型。因此，学习到的 token 可能很好地保留了图像信息，但自回归（AR）先验仍然难以从左到右进行预测。我们使用三方变分一致性来分析这种不匹配，该方法将潜变量学习分解为三个一致性条件：条件似然一致性、先验一致性和后验一致性。TVC 表明，两阶段训练保留了重建方面，但将先验一致性排除在分词器目标之外：在 AR 先验参与训练之前，整体 token 分布已经固定。受此观点启发，我们在分词器训练期间添加了一个分布级别的先验匹配信号，同时保持重建目标不变。我们使用 Wasserstein 梯度流更新来优化该信号。对于硬分类 token，该更新简化为追踪分词器当前 token 分布的辅助 AR 模型与目标 AR 先验之间的 token 级别对比。它仅需通过两个 AR 模型的前向传递，且不需要通过其中任何一个进行反向传播。由此产生的分词器 wAR-Tok 降低了 AR 损失，并在具有可比重建质量的情况下，改善了 CIFAR-10 和 ImageNet 上的生成 FID。
引用
@article{arxiv.2605.06148,
title = {Learning Discrete Autoregressive Priors with Wasserstein Gradient Flow},
author = {Bowen Zheng and Yihong Luo and Tianyang Hu},
journal= {arXiv preprint arXiv:2605.06148},
year = {2026}
}