基于非平衡最优传输的一步生成建模

中文摘要 AI 辅助

漂移模型通过将分布传输摊销到训练中实现一步生成，但这种效率对每次更新时估计的传输场提出了更高的要求。在大规模训练中，该场由生成的样本和真实样本的有限小批量计算得出，这些样本仅对底层分布提供不完美的近似。平衡最优传输在每个小批量内强制执行精确的质量匹配，使得估计的传输场对真实数据批次的特定组成敏感。我们发现生成样本和真实样本应被不对称地对待：让分配给真实样本的质量自适应调整，同时保持每个生成样本完全传输，在受控消融实验中提升了六个特征空间指标，并且对松弛强度的鲁棒性优于同时松弛两个边际，后者在更强的松弛下低于平衡传输。受此观察启发，我们提出了非平衡最优传输梯度流（UOT-GF），该方法保持生成样本边际固定，仅松弛真实数据边际。在ImageNet-256上DiT-B/2的相同设置下，UOT-GF将Fréchet初始距离（FID）从平衡W-Flow基线的1.53提升至1.46；将相同方案扩展到L/2和XL/2分别得到1.34和1.22的FID，这是我们比较的一步模型中最佳的FID。我们进一步推导了诱导的UOT传输力，建立了动力学Vlasov-Fokker-Planck公式，其过阻尼零温度极限恢复了漂移动力学，并刻画了非目标稳态以及收敛的充分条件。

英文摘要

Drifting models enable one-step generation by amortizing distribution transport into training, but this efficiency places greater demands on the transport field estimated at each update. In large-scale training, the field is computed from finite mini-batches of generated and real samples, which provide only imperfect approximations to the underlying distributions. Balanced optimal transport enforces exact mass matching within every mini-batch, making the estimated field sensitive to the particular composition of the real-data batch. We find that generated and real samples should be treated asymmetrically: letting the mass assigned to real samples adapt while keeping every generated sample fully transported improves generation across six feature-space metrics in controlled ablations, and is more robust to the relaxation strength than relaxing both marginals simultaneously, which falls below balanced transport under stronger relaxation. Motivated by this observation, we propose Unbalanced Optimal Transport Gradient Flow (UOT-GF), which keeps the generated-sample marginal fixed and relaxes only the real-data marginal. Under identical settings at DiT-B/2 on ImageNet-256, UOT-GF improves Fréchet Inception Distance (FID) from 1.53 to 1.46 over the balanced W-Flow baseline; scaling the same recipe yields 1.34 and 1.22 FID at L/2 and XL/2, the best FID among the one-step models we compare. We further derive the induced UOT transport force, establish a kinetic Vlasov--Fokker--Planck formulation whose overdamped zero-temperature limit recovers the drifting dynamics, and characterize non-target stationary states together with sufficient conditions for convergence.

发表机构

• Imperial College London（帝国理工学院）
• Ant Group（蚂蚁集团）
• Brunel University of London（伦敦布鲁内尔大学）

机构由 AI 辅助整理，请以论文原文为准。