# Measuring Generalization with Optimal Transport

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

## Key facts

- Cited by 25 papers in the map (data through 14 Oct 2026).
- Builds on 41 of its 62 references that are in the map.
- First posted on arXiv 7 Jun 2021.

## Builds on (most cited first, 30 of 41)

- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [A Simple Framework for Contrastive Learning of Visual Representations](https://mlsanity.com/paper/2002.05709/) (2020)
- [Momentum Contrast for Unsupervised Visual Representation Learning](https://mlsanity.com/paper/1911.05722/) (2019)
- [mixup: Beyond Empirical Risk Minimization](https://mlsanity.com/paper/1710.09412/) (2017)
- [Wide Residual Networks](https://mlsanity.com/paper/1605.07146/) (2016)
- [Supervised Contrastive Learning](https://mlsanity.com/paper/2004.11362/) (2020)
- [Understanding deep learning requires rethinking generalization](https://mlsanity.com/paper/1611.03530/) (2016)
- [Concrete Problems in AI Safety](https://mlsanity.com/paper/1606.06565/) (2016)
- [Spectral Normalization for Generative Adversarial Networks](https://mlsanity.com/paper/1802.05957/) (2018)
- [Spectrally-normalized margin bounds for neural networks](https://mlsanity.com/paper/1706.08498/) (2017)
- [Contrastive Learning with Hard Negative Samples](https://mlsanity.com/paper/2010.04592/) (2020)
- [Computing Nonvacuous Generalization Bounds for Deep (Stochastic) Neural Networks with Many More Parameters than Training Data](https://mlsanity.com/paper/1703.11008/) (2017)
- [Fantastic Generalization Measures and Where to Find Them](https://mlsanity.com/paper/1912.02178/) (2019)
- [Lipschitz regularity of deep neural networks: analysis and efficient estimation](https://mlsanity.com/paper/1805.10965/) (2018)
- [Debiased Contrastive Learning](https://mlsanity.com/paper/2007.00224/) (2020)
- [A PAC-Bayesian Approach to Spectrally-Normalized Margin Bounds for Neural Networks](https://mlsanity.com/paper/1707.09564/) (2017)
- [Stronger generalization bounds for deep nets via a compression approach](https://mlsanity.com/paper/1802.05296/) (2018)
- [Norm-Based Capacity Control in Neural Networks](https://mlsanity.com/paper/1503.00036/) (2015)
- [Efficient and Accurate Estimation of Lipschitz Constants for Deep Neural Networks](https://mlsanity.com/paper/1906.04893/) (2019)
- [Size-Independent Sample Complexity of Neural Networks](https://mlsanity.com/paper/1712.06541/) (2017)
- [Uniform convergence may be unable to explain generalization in deep learning](https://mlsanity.com/paper/1902.04742/) (2019)
- [Generalization Bounds of Stochastic Gradient Descent for Wide and Deep Neural Networks](https://mlsanity.com/paper/1905.13210/) (2019)
- [Statistics of Robust Optimization: A Generalized Empirical Likelihood Approach](https://mlsanity.com/paper/1610.03425/) (2016)
- [Deep learning generalizes because the parameter-function map is biased towards simple functions](https://mlsanity.com/paper/1805.08522/) (2018)
- [Towards Understanding the Role of Over-Parametrization in Generalization of Neural Networks](https://mlsanity.com/paper/1805.12076/) (2018)
- [Predicting the Generalization Gap in Deep Networks with Margin Distributions](https://mlsanity.com/paper/1810.00113/) (2018)
- [Large Margin Deep Networks for Classification](https://mlsanity.com/paper/1803.05598/) (2018)
- [Exactly Computing the Local Lipschitz Constant of ReLU Networks](https://mlsanity.com/paper/2003.01219/) (2020)
- [ClusterFit: Improving Generalization of Visual Representations](https://mlsanity.com/paper/1912.03330/) (2019)
- [In Search of Robust Measures of Generalization](https://mlsanity.com/paper/2010.11924/) (2020)

## Built on by (most cited first, 25 of 25)

- [Robust Contrastive Learning against Noisy Views](https://mlsanity.com/paper/2201.04309/) (2022)
- [Scalable Optimal Transport Methods in Machine Learning: A Contemporary Survey](https://mlsanity.com/paper/2305.05080/) (2023)
- [Some Fundamental Aspects about Lipschitz Continuity of Neural Networks](https://mlsanity.com/paper/2302.10886/) (2023)
- [Confidence and Dispersity Speak: Characterising Prediction Matrix for Unsupervised Accuracy Estimation](https://mlsanity.com/paper/2302.01094/) (2023)
- [X-Boundary: Establishing Exact Safety Boundary to Shield LLMs from Multi-Turn Jailbreaks without Compromising Usability](https://mlsanity.com/paper/2502.09990/) (2025)
- [Integral Probability Metrics PAC-Bayes Bounds](https://mlsanity.com/paper/2207.00614/) (2022)
- [Sample Complexity Bounds for Estimating Probability Divergences under Invariances](https://mlsanity.com/paper/2311.02868/) (2023)
- [Towards Bridging Generalization and Expressivity of Graph Neural Networks](https://mlsanity.com/paper/2410.10051/) (2024)
- [On Measuring Excess Capacity in Neural Networks](https://mlsanity.com/paper/2202.08070/) (2022)
- [InfoOT: Information Maximizing Optimal Transport](https://mlsanity.com/paper/2210.03164/) (2022)
- [What Does Softmax Probability Tell Us about Classifiers Ranking Across Diverse Test Conditions?](https://mlsanity.com/paper/2406.09908/) (2024)
- [Input margins can predict generalization too](https://mlsanity.com/paper/2308.15466/) (2023)
- [Slicing Mutual Information Generalization Bounds for Neural Networks](https://mlsanity.com/paper/2406.04047/) (2024)
- [In Search of Probeable Generalization Measures](https://mlsanity.com/paper/2110.12259/) (2021)
- [Generalizing to any diverse distribution: uniformity, gentle finetuning and rebalancing](https://mlsanity.com/paper/2410.05980/) (2024)
- [Beyond External Monitors: Enhancing Transparency of Large Language Models for Easier Monitoring](https://mlsanity.com/paper/2502.05242/) (2025)
- [Data-Agnostic Augmentations for Unknown Variations: Out-of-Distribution Generalisation in MRI Segmentation](https://mlsanity.com/paper/2505.10223/) (2025)
- [Understanding Generalization in Node and Link Prediction](https://mlsanity.com/paper/2507.00927/) (2025)
- [Generalization of Graph Neural Networks through the Lens of Homomorphism](https://mlsanity.com/paper/2403.06079/) (2024)
- [Meet You Halfway: Explaining Deep Learning Mysteries](https://mlsanity.com/paper/2206.04463/) (2022)
- [Generalization within in silico screening](https://mlsanity.com/paper/2307.09379/) (2023)
- [Twin: Tuning Learning Rate and Weight Decay of Deep Homogeneous Classifiers without Validation](https://mlsanity.com/paper/2403.05532/) (2024)
- [Bispectral OT: Dataset Comparison using Symmetry-Aware Optimal Transport](https://mlsanity.com/paper/2509.20678/) (2025)
- [Towards Initialization-dependent and Non-vacuous Generalization Bounds for Overparameterized Shallow Neural Networks](https://mlsanity.com/paper/2604.00505/) (2026)
- [Optimal Transport PAC-Bayes Bounds](https://mlsanity.com/paper/2610.16560/) (2026)

## Abstract (arXiv)

Understanding the generalization of deep neural networks is one of the most important tasks in deep learning. Although much progress has been made, theoretical error bounds still often behave disparately from empirical observations. In this work, we develop margin-based generalization bounds, where the margins are normalized with optimal transport costs between independent random subsets sampled from the training distribution. In particular, the optimal transport cost can be interpreted as a generalization of variance which captures the structural properties of the learned feature space. Our bounds robustly predict the generalization error, given training data and network parameters, on large scale datasets. Theoretically, we demonstrate that the concentration and separation of features play crucial roles in generalization, supporting empirical results in the literature. The code is available at urlhttps://github.com/chingyaoc/kV-Margin.