Conditional Distribution Quantization in Machine Learning

TL;DR

This work introduces Conditional Competitive Learning Vector Quantization (CCLVQ), a framework to approximate multimodal conditional distributions $\mathcal{L}(Y\mid X)$ via $n$ learnable, input-conditioned points $f_i(X)$. By minimizing the distortion $\Delta_n(f)=\mathbb{E}[\min_i|Y-f_i(X)|^2]$ and leveraging Wasserstein-2 distance, CCLVQ yields a quantified, multimodal representation of $Y$ given $X$ with an accompanying expert-weight classifier for uncertainty. The authors provide a theoretical foundation connecting conditional quantization to optimal Wasserstein approximations, and demonstrate practical gains in uncertainty-aware inpainting, multi-value regression, normalizing flows, and GANs. The approach enhances diversity and coverage of conditional distributions while maintaining or improving output quality, with broad applicability to uncertainty quantification and multimodal data generation.

Abstract

Conditional expectation \mathbb{E}(Y \mid X) often fails to capture the complexity of multimodal conditional distributions \mathcal{L}(Y \mid X). To address this, we propose using n-point conditional quantizations--functional mappings of X that are learnable via gradient descent--to approximate \mathcal{L}(Y \mid X). This approach adapts Competitive Learning Vector Quantization (CLVQ), tailored for conditional distributions. It goes beyond single-valued predictions by providing multiple representative points that better reflect multimodal structures. It enables the approximation of the true conditional law in the Wasserstein distance. The resulting framework is theoretically grounded and useful for uncertainty quantification and multimodal data generation tasks. For example, in computer vision inpainting tasks, multiple plausible reconstructions may exist for the same partially observed input image X. We demonstrate the effectiveness of our approach through experiments on synthetic and real-world datasets.

Figures (12)

• Figure 1: CCLVQ training loss on the MNIST dataset for an inpainting task, showing the effect of adding new experts using the splitting strategy every 400 epochs.
• Figure 2: Visualization of synthetic data showing true data points $Y$ (in grey) and predicted values from every expert (in colors). Each expert captures a specific mode of conditional distribution.
• Figure 3: Weights estimations of experts versus probabilities of mode w.r.t $X$.
• Figure 4: Multi-modal inpainting results on MNIST. Each row shows the ground truth $Y$, masked input $X$, and reconstructions from different experts, followed by the classifier's weights as a measure of uncertainty.
• Figure 5: Multi-modal denoising results on MNIST. Each row shows the ground truth $Y$, noised input $X$, and reconstructions from different experts, followed by the classifier's weights as a measure of uncertainty.

Theorems & Definitions (10)

• Proposition 3.1
• proof
• Proposition 3.2
• proof
• Proposition 3.3: procPages
• Theorem 4.1
• proof
• Theorem 4.2
• Theorem
• proof