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

TL;DR

This paper introduces an optimal transport-based quantization method for Flow Matching generative models, significantly reducing parameter precision requirements while maintaining high-quality generation, enabling efficient deployment on edge devices.

Contribution

It proposes a novel OT-based post-training quantization technique for FM models, with theoretical bounds and empirical validation showing superior performance over traditional schemes.

Findings

OT quantization preserves visual quality down to 2-3 bits.

Compared to other methods, OT quantization maintains latent space stability.

Empirical results across five datasets demonstrate effectiveness.

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.