Reliable Clustering and Quantization via Distortion-Constrained Optimal Transport

Clustering and vector quantization are core primitives for representation learning and large-scale retrieval, yet widely used methods often produce clusters or codewords of highly uneven quality, especially in noisy or heterogeneous data. In this paper, we propose a distortion-constrained optimal transport (DCOT) framework that explicitly enforces per-cluster bounds on the average assignment distortion, defined as the expected distance between data points and their assigned representatives, ensuring more consistent and reliable representations. Our formulation couples an optimal transport assignment objective with cluster-level distortion constraints to control cluster quality. To solve this constrained problem, we develop an efficient alternating optimization algorithm that iteratively updates transport plans, representatives (codewords), and dual variables. We further provide theoretical guarantees including entropic consistency, coordinate descent, and sublinear convergence. Importantly, DCOT serves as a plug-in assignment module that can be seamlessly integrated into a wide range of vector quantization methods. Replacing the assignment step with DCOT consistently improves reconstruction quality across diverse architectures and training settings, demonstrating strong robustness and general applicability.