用半离散最优传输提升生成模型训练效率与效果
AlignFlow: Improving Flow-based Generative Models with Semi-Discrete Optimal Transport
- 引入半离散最优传输构建噪声与数据点的精确配对
- 在图像生成任务中显著提升主流生成模型性能
- 可无缝集成到现有模型，计算开销极低
基于流的生成模型（FGMs）能有效将噪声转换为复杂数据分布。将最优传输（OT）用于训练过程中噪声与数据的耦合，已被证明可增强流轨迹的直线性，从而提升推理效率。然而，现有基于OT的方法依赖采样噪声和数据点的（小）批量来估计OT计划，限制了其在大规模高维数据集上的可扩展性。本文提出AlignFlow，利用半离散最优传输（SDOT）实现噪声分布与数据点间的显式、最优对齐，并保证收敛性。SDOT通过将噪声空间划分为拉格朗日单元，每个单元映射至对应数据点。训练时，独立同分布的噪声样本通过SDOT映射与数据点配对。AlignFlow在大规模数据集和复杂模型架构下仍保持良好扩展性，计算开销几乎可忽略。实验表明，AlignFlow可提升多种先进FGM算法的性能，且可作为即插即用组件直接集成。代码已开源：https://github.com/konglk1203/AlignFlow。
原文摘要 · Abstract (English)
Flow-based Generative Models (FGMs) effectively transform noise into complex data distributions. Incorporating Optimal Transport (OT) to couple noise and data during FGM training has been shown to improve the straightness of flow trajectories, enabling more effective inference. However, existing OT-based methods estimate the OT plan using (mini-)batches of sampled noise and data points, which limits their scalability to large and high-dimensional datasets in FGMs. This paper introduces AlignFlow, a novel approach that leverages Semi-Discrete Optimal Transport (SDOT) to enhance the training of FGMs by establishing an explicit, optimal alignment between noise distribution and data points with guaranteed convergence. SDOT computes a transport map by partitioning the noise space into Laguerre cells, each mapped to a corresponding data point. During FGM training, i.i.d. noise samples are paired with data points via the SDOT map. AlignFlow scales well to large datasets and model architectures with negligible computational overhead. Experimental results show that AlignFlow improves the performance of a wide range of state-of-the-art FGM algorithms and can be integrated as a plug-and-play component. Code is available at: https://github.com/konglk1203/AlignFlow.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。