Low-Bit, High-Fidelity: Optimal Transport Quantization for Flow Matching

📄 arXiv: 2511.11418v1 📥 PDF

作者: Dara Varam, Diaa A. Abuhani, Imran Zualkernan, Raghad AlDamani, Lujain Khalil

分类: cs.LG, cs.CV

发布日期: 2025-11-14

备注: 12 pages, 8 figures


💡 一句话要点

提出基于最优传输量化的Flow Matching模型压缩方法,实现低比特高保真生成。

🎯 匹配领域: 支柱二:RL算法与架构 (RL & Architecture)

关键词: Flow Matching 最优传输 模型量化 生成模型 模型压缩

📋 核心要点

  1. Flow Matching模型训练高效,但高精度参数需求阻碍了其在资源受限设备上的部署。
  2. 论文提出基于最优传输的量化方法,通过最小化权重分布差异,在低比特下保持模型性能。
  3. 实验证明,该方法在极低比特(2-3 bits)下,仍能保持生成质量和潜在空间稳定性。

📝 摘要(中文)

Flow Matching (FM)生成模型提供高效的无模拟训练和确定性采样,但其高精度参数要求限制了实际部署。本文将基于最优传输(OT)的后训练量化方法应用于FM模型,最小化量化权重和原始权重之间的2-Wasserstein距离,并系统地将其有效性与均匀、分段和对数量化方案进行比较。理论分析提供了量化下生成退化的上界,在五个不同复杂度的基准数据集上的实验结果表明,基于OT的量化在每个参数2-3比特的情况下,既能保持视觉生成质量,又能保持潜在空间的稳定性,而其他方法则失败。这确立了基于OT的量化作为一种有原则的、有效的方法来压缩FM生成模型,以用于边缘和嵌入式AI应用。

🔬 方法详解

问题定义:Flow Matching (FM) 生成模型虽然具有训练高效和确定性采样的优点,但其对参数精度要求高,导致模型体积大,难以部署在边缘设备或嵌入式系统中。现有量化方法在压缩FM模型时,容易造成严重的性能下降,无法在低比特下保持生成质量和潜在空间的一致性。

核心思路:论文的核心思路是利用最优传输理论,寻找量化后的权重分布与原始权重分布之间的最佳匹配,从而最小化量化带来的信息损失。通过最小化2-Wasserstein距离,确保量化后的模型尽可能地保留原始模型的生成能力。这种方法旨在找到一种更智能的量化方案,而不是简单地进行均匀或对数量化。

技术框架:该方法属于后训练量化,即先训练好一个高精度的FM模型,然后对其权重进行量化。整体流程包括:1) 训练一个高精度的Flow Matching模型;2) 使用最优传输算法,计算量化后的权重值,目标是最小化量化前后权重分布的2-Wasserstein距离;3) 使用量化后的模型进行推理或生成。该框架的关键在于最优传输算法的选择和参数调整。

关键创新:该论文的关键创新在于将最优传输理论引入到Flow Matching模型的量化中。与传统的均匀量化、分段量化或对数量化方法相比,最优传输量化能够更好地保留原始模型的权重分布信息,从而在低比特下实现更高的生成质量和潜在空间稳定性。本质区别在于,传统方法忽略了权重分布的结构信息,而最优传输量化则显式地考虑了这种结构信息。

关键设计:关键设计包括:1) 使用2-Wasserstein距离作为量化损失函数,衡量量化前后权重分布的差异;2) 选择合适的最优传输算法,例如Sinkhorn算法,用于计算量化后的权重值;3) 对不同的Flow Matching模型层进行不同的量化比特数设置,以平衡压缩率和性能;4) 理论分析量化误差对生成模型性能的影响,并给出量化误差上界。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明,基于最优传输的量化方法在2-3比特下,能够显著优于均匀、分段和对数量化方法,保持Flow Matching模型的生成质量和潜在空间稳定性。在多个基准数据集上,该方法在极低比特下仍然能够生成高质量的图像,并且能够保持潜在空间的平滑性,这对于下游任务至关重要。相比于其他量化方法,该方法能够实现更高的压缩率和更小的性能损失。

🎯 应用场景

该研究成果可应用于边缘计算、嵌入式系统和移动设备等资源受限的场景,实现Flow Matching生成模型的高效部署。例如,可以将压缩后的FM模型用于图像生成、音频合成、视频编辑等任务,在本地设备上提供高质量的生成服务,而无需依赖云端计算资源。此外,该方法还可以应用于模型压缩和加速的其他领域。

📄 摘要(原文)

Flow Matching (FM) generative models offer efficient simulation-free training and deterministic sampling, but their practical deployment is challenged by high-precision parameter requirements. We adapt optimal transport (OT)-based post-training quantization to FM models, minimizing the 2-Wasserstein distance between quantized and original weights, and systematically compare its effectiveness against uniform, piecewise, and logarithmic quantization schemes. Our theoretical analysis provides upper bounds on generative degradation under quantization, and empirical results across five benchmark datasets of varying complexity show that OT-based quantization preserves both visual generation quality and latent space stability down to 2-3 bits per parameter, where alternative methods fail. This establishes OT-based quantization as a principled, effective approach to compress FM generative models for edge and embedded AI applications.