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

TL;DR

This work addresses the challenge of deploying Flow Matching FM generative models on resource-constrained devices by introducing optimal-transport OT post-training quantization to minimize the $W_2$ distance between original and quantized weights. The authors derive theoretical bounds showing both uniform and OT quantization scale with $2^{-2b}$ in the $FID$ metric, but with different front constants, and demonstrate that equal-mass OT quantization yields tighter bounds (ratio $\rho(b) \approx 0.25$–$0.4$) under realistic weight tail assumptions. Empirically, OT quantization preserves both generation quality (PSNR/SSIM) and latent-space stability across five benchmarks (MNIST, FashionMNIST, CIFAR-10, ImageNet, CelebA), maintaining fidelity down to 2–3 bits per parameter where other schemes fail. The results support OT quantization as a principled, effective approach for edge and embedded AI deployment of FM models, with avenues for hardware-aware optimizations and extension to other architectures.

Abstract

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.

Figures (8)

• Figure 1: System diagram illustrating the quantization process for Flow Matching (FM) generative models. Per-layer weights are extracted from a trained FM model and flattened into a one-dimensional distribution. These weights are quantized using optimal transport (equal-mass bin) quantization method. Quantized models with different bit-widths are then used for sample generation, and the resulting outputs are evaluated for generative fidelity while the latent space is used to evaluate disentanglement.
• Figure 2: Comparison of OT-based quantization to other post-training quantization methods at various bit-widths, demonstrating the robustness and sampling quality for a model trained on the CelebA dataset.
• Figure 3: Quantitative evaluation of generative fidelity under quantization. (A) SSIM and (B) PSNR scores for each quantization scheme and bitwidth, evaluated across all benchmark datasets. OT quantization consistently outperforms alternatives, particularly at extreme low bit-widths.
• Figure 4: Latent variance standard deviation versus bitwidth for each quantization method and dataset. OT quantization maintains stable latent structure across all bitwidths, while alternative methods become increasingly unstable at low bits, especially for complex datasets.
• Figure 5: Generated samples for MNIST.

Theorems & Definitions (23)

• Lemma 1
• Lemma 2
• Theorem 3
• Lemma 4
• Lemma 5
• Theorem 6
• Definition 1
• Definition 2
• proof
• Lemma 7