Research map / Learning theory and algorithms / Optimal transport

Supervised Distributional Reduction via Optimal Transport and Dependence Maximization

NeurIPS 2026 · Learning theory and algorithms › Optimal transport · approach: Neural optimal transport · 0 citations (OpenAlex)

Read the paperCheck your idea against 108,967 papers

AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.

Proposes Supervised Distributional Reduction (SDR), which extends Fused Gromov-Wasserstein distributional reduction with a CKA dependence term and an RKHS out-of-sample map, used as a non-stationary kernel feature extractor for GPs.

Problem

Distributional reduction (joint clustering and dimensionality reduction) is unsupervised, so learned representations can lose task-relevant signal and lack an out-of-sample mapping for prediction.

Approach

Combine supervised FGW alignment (with prototype targets eliminated in closed form for Bregman losses) with a direct representation-level CKA dependence term on embeddings, plus a kernel ridge out-of-sample map so the embedding can feed a GP with a data-dependent non-stationary kernel.

Key results

• On the GP kernel benchmarks, SDR-GP attains the best MLL on Energy (-0.37) and Concrete (-0.40) and the best MSE on Energy (0.05) and Concrete (0.13), with the best MNIST-2k accuracy (0.95 vs 0.91 for NCA-GP).
• On COIL-20 (m=200), SDR raises CKA to 0.56 from 0.31 for FGW-only; on SNAREseq CKA rises to 0.91 from 0.38.
• Pure DR with KRR on Swiss Roll test R^2 reaches 0.96 for SDR vs 0.53 for Supervised UMAP, and on Friedman test R^2 0.91 vs 0.55 (k-NN, S-UMAP).

Surprising findings

• FGW-only embeddings have lower CKA with targets than the original data (0.31 vs 0.48 on COIL-20), indicating the coupling-mediated supervision is diluted.
• Supervised UMAP gives high train R^2 but poor test R^2 on Swiss Roll (0.47 k-NN), showing overfitting relative to SDR.

Limitations

• Pure DR scales as O(n^3), limiting scalability to larger datasets.
• CKA term is non-differentiable through the transport solver, so T-step ignores supervision and alternating scheme is inexact.
• Sensitive to the eta hyperparameter; trade-off between OOS map approximation and prototype predictive signal for DKL.

Experimental setup

Datasets: COIL-20 Fashion-MNIST SNAREseq Boston Housing Energy Efficiency Concrete Compressive Strength MNIST S-curve Swiss Roll Friedman

Benchmarks: COIL-20 Fashion-MNIST SNAREseq Boston Housing Energy Efficiency Concrete Compressive Strength MNIST-2k S-curve Swiss Roll Friedman

Metrics: homogeneity k-means NMI silhouette score CKA score MLL MSE MLP accuracy R^2 MACE

Baselines: DistR Cluster-then-DR DR-then-Cluster NCA-GP KSPCA-GP UMAP-GP DGP DKL Supervised UMAP KS-PCA FGW-only

Builds on: Distributional Reduction (DistR) Fused Gromov-Wasserstein Centered Kernel Alignment Deep Kernel Learning Supervised UMAP

Closest papers by meaning

• Learning to Embed Distributions via Maximum Kernel Entropy (NeurIPS 2024 · similarity 0.86)
• Distributional Autoencoders Know the Score (NeurIPS 2025 · similarity 0.85)
• On Statistical Learning Theory for Distributional Inputs (ICML 2024 · similarity 0.85)
• Conditional Distribution Compression via the Kernel Conditional Mean Embedding (NeurIPS 2025 · similarity 0.85)
• Supervised Kernel Thinning (NeurIPS 2024 · similarity 0.84)
• Bayesian Optimization with Fisher Information Geometry: Gradient Bounds and Trust-Region Methods (NeurIPS 2026 · similarity 0.84)