discretize_distributions: Efficient Quantization of Gaussian Mixtures with Guarantees in Wasserstein Distance

TL;DR

The paper tackles the challenge of providing scalable, certified discretizations of Gaussian mixtures for uncertainty propagation and verification in cyber-physical systems. It introduces discretize_distributions, a two-stage framework that constructs mode-aware quantization schemes (grid and sigma-point-inspired) and applies them with closed-form or tractable $W_2$ error bounds. The work delivers a modular PyTorch-based package with distribution primitives, automatic scheme generation, and a discretization operation that yields discrete approximations with $W_2$ certificates. This enables accurate, efficient quantization of complex Gaussian mixtures, demonstrated across high-dimensional and degenerate cases with favorable computation times and compression of support points.

Abstract

We present discretize_distributions, a Python package that efficiently constructs discrete approximations of Gaussian mixture distributions and provides guarantees on the approximation error in Wasserstein distance. The package implements state-of-the-art quantization methods for Gaussian mixture models and extends them to improve scalability. It further integrates complementary quantization strategies such as sigma-point methods and provides a modular interface that supports custom schemes and integration into control and verification pipelines for cyber-physical systems. We benchmark the package on various examples, including high-dimensional, large, and degenerate Gaussian mixtures, and demonstrate that discretize_distributions produces accurate approximations at low computational cost.

Figures (3)

• Figure 1: Illustration of the quantization of a continuous probability distribution ${\mathbb{P}}$, yielding its discrete approximation ${\Delta}_{{\bm{\mathcal{R}}},{\mathcal{C}}}\#{\mathbb{P}}$ with a support of size $4$. The probability mass ${{{\pi}}^{(i)}}$ of ${\mathbb{P}}$ in each region ${\mathcal{R}}_i$ is mapped to the location ${{c}}_i$.
• Figure 2: Quantization of a 2D Gaussian distribution (black dots) using a standard Gaussian (blue dots) as a template for grid and cross configurations. In the eigenbasis $T$ of the covariance matrix (orange arrows), the Gaussian’s dimensions are independent, allowing efficient construction of the discrete support. In \ref{['fig:grid']}, the support is obtained by taking the cross-product of one-dimensional quantizations of the univariate marginals, whereas in \ref{['fig:cross']}, symmetric locations are placed along each principal axis (green markers). The resulting quantized distribution in the original space is obtained by mapping these locations through $T^{-1}$. Blue and black lines indicate the Voronoi partitions in the transformed and original spaces, respectively.
• Figure 3: Example 2D Gaussian mixture ($M=3$) used in Section \ref{['sec:tool']} and its quantization (red dots) constructed per mode \ref{['fig:per_mode']} and per component \ref{['fig:per_comp']}. The mode-wise quantization achieves a comparable Wasserstein error ($0.47$ vs. $0.46$) while requiring a much smaller support ($19$ vs. $29$ points).