What changes when clustering learns the geometry itself?

A new preprint shows that quantization can approximate both data distribution and the shape of the space it lives in. The method extends k-means by treating distances as learnable rather than fixed.

Clustering usually means splitting a dataset into groups. The most common version, k-means, picks center points and assigns each data point to the nearest one. That procedure is mathematically equivalent to quantizing a probability measure with respect to the Wasserstein distance. Wasserstein quantization works inside a fixed metric space. The geometry of that space is given and does not change.

Gromov-Wasserstein quantization changes the premise. It compares metric measure spaces by their internal distance structures rather than by pointwise positions. The goal becomes finding a discrete space that approximates both the distribution of mass and the pairwise geometry of the original space. The paper proves that solutions to this problem exist. It also gives a characterization that justifies an iterative algorithm analogous to Lloyd's method for k-means.

The algorithm alternates between two steps

One step fixes the discrete geometry and updates the assignment of mass. The other step fixes the assignment and updates the discrete distances. Each iteration decreases the Gromov-Wasserstein objective. The procedure converges to a stationary point.

The authors implement this scheme and test it on synthetic and real geometries.

Quantization rates reveal the cost of learning geometry

For Euclidean spaces equipped with the standard metric, the paper calculates the asymptotic quantization rate. The rate matches the classical Wasserstein rate up to a constant factor that depends only on dimension. Learning the geometry does not change the polynomial decay exponent. The constant factor quantifies the extra difficulty of simultaneously approximating distances.

Experiments go beyond point clouds

The method is applied to geodesic distances on three-dimensional shapes. It also performs structured pruning of neural networks by clustering weight matrices with their internal correlation geometry. In both settings the algorithm produces discrete approximations whose quality aligns with the theoretically optimal rates.

The paper does not report wall-clock times or memory footprints for these experiments. The existence proof assumes compact metric measure spaces. The rate calculation covers Euclidean geometries with the standard metric. The algorithm is guaranteed to decrease the objective but not to reach a global minimum. The numerical section does not compare against every existing GW solver. The authors do not discuss differentiability of the discrete geometry with respect to the input data.

The sources this article was written from

  • "arXiv:2608.11016v1
    https://arxiv.org/abs/2608.11016v1"

Read next