# Supervised Distributional Reduction via Optimal Transport and Dependence Maximization

arXiv: https://arxiv.org/abs/2605.27619 · HTML: https://mlsanity.com/paper/2605.27619/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 18 of its 42 references that are in the map.
- First posted on arXiv 26 May 2026.

## Builds on (most cited first, 18 of 18)

- [Fashion-MNIST: a Novel Image Dataset for Benchmarking Machine Learning Algorithms](https://mlsanity.com/paper/1708.07747/) (2017)
- [UMAP: Uniform Manifold Approximation and Projection for Dimension Reduction](https://mlsanity.com/paper/1802.03426/) (2018)
- [Similarity of Neural Network Representations Revisited](https://mlsanity.com/paper/1905.00414/) (2019)
- [GPyTorch: Blackbox Matrix-Matrix Gaussian Process Inference with GPU Acceleration](https://mlsanity.com/paper/1809.11165/) (2018)
- [Deep Kernel Learning](https://mlsanity.com/paper/1511.02222/) (2015)
- [Accurate Uncertainties for Deep Learning Using Calibrated Regression](https://mlsanity.com/paper/1807.00263/) (2018)
- [Fast Computation of Wasserstein Barycenters](https://mlsanity.com/paper/1310.4375/) (2013)
- [Optimal Transport for structured data with application on graphs](https://mlsanity.com/paper/1805.09114/) (2018)
- [Doubly Stochastic Variational Inference for Deep Gaussian Processes](https://mlsanity.com/paper/1705.08933/) (2017)
- [The Promises and Pitfalls of Deep Kernel Learning](https://mlsanity.com/paper/2102.12108/) (2021)
- [Linear-Time Gromov Wasserstein Distances using Low Rank Couplings and Costs](https://mlsanity.com/paper/2106.01128/) (2021)
- [Uncertainty Toolbox: an Open-Source Library for Assessing, Visualizing, and Improving Uncertainty Quantification](https://mlsanity.com/paper/2109.10254/) (2021)
- [Dirichlet-based Gaussian Processes for Large-scale Calibrated Classification](https://mlsanity.com/paper/1805.10915/) (2018)
- [Statistical Optimal Transport via Factored Couplings](https://mlsanity.com/paper/1806.07348/) (2018)
- [Large-Scale Kernel Methods for Independence Testing](https://mlsanity.com/paper/1606.07892/) (2016)
- [Semi-relaxed Gromov-Wasserstein divergence with applications on graphs](https://mlsanity.com/paper/2110.02753/) (2021)
- [SNEkhorn: Dimension Reduction with Symmetric Entropic Affinities](https://mlsanity.com/paper/2305.13797/) (2023)
- [Bias-variance decompositions: the exclusive privilege of Bregman divergences](https://mlsanity.com/paper/2501.18581/) (2025)

## Abstract (arXiv)

Learning representations that capture both intrinsic data geometry and target-relevant structure remains a fundamental challenge, particularly in settings where data reduction must balance compression with predictive fidelity. While distributional reduction-encompassing joint clustering and dimensionality reduction-offers a principled way to summarize data, its supervised variants remain relatively under-explored, despite the importance of retaining task-relevant signal for downstream prediction and decision-making. We propose Supervised Distributional Reduction (SDR), an algorithm for learning target-aware representations by combining optimal transport with explicit dependence maximization. SDR builds on the Fused Gromov-Wasserstein (FGW) objective to align the relational structure of the input distribution with a set of representative points, while augmenting it with a direct dependence term that encourages the learned embeddings to capture predictive signal more explicitly. This results in compact representations that reflect both geometric structure and supervision. Beyond representation learning, SDR naturally induces a data-dependent, non-stationary geometry that can be leveraged for settings such as Gaussian Process (GP) modelling. By redefining distances through target-aware distributional alignment, SDR enables the construction of adaptive kernels that respond to local variations in both data geometry and supervision, offering an optimal transport-based perspective on non-stationary kernel design.