# UniPROT: Uniform Prototype Selection via Partial Optimal Transport with Submodular Guarantees

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

## Key facts

- Cited by 2 papers in the map (data through 14 Oct 2026).
- Builds on 28 of its 68 references that are in the map.
- First posted on arXiv 13 Apr 2026.

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

- [Sentence-BERT: Sentence Embeddings using Siamese BERT-Networks](https://mlsanity.com/paper/1908.10084/) (2019)
- [Measuring Mathematical Problem Solving With the MATH Dataset](https://mlsanity.com/paper/2103.03874/) (2021)
- [SuperGLUE: A Stickier Benchmark for General-Purpose Language Understanding Systems](https://mlsanity.com/paper/1905.00537/) (2019)
- [Are NLP Models really able to Solve Simple Math Word Problems?](https://mlsanity.com/paper/2103.07191/) (2021)
- [This Looks Like That: Deep Learning for Interpretable Image Recognition](https://mlsanity.com/paper/1806.10574/) (2018)
- [Textbooks Are All You Need II: phi-1.5 technical report](https://mlsanity.com/paper/2309.05463/) (2023)
- [Zephyr: Direct Distillation of LM Alignment](https://mlsanity.com/paper/2310.16944/) (2023)
- [MAmmoTH: Building Math Generalist Models through Hybrid Instruction Tuning](https://mlsanity.com/paper/2309.05653/) (2023)
- [Not All Samples Are Created Equal: Deep Learning with Importance Sampling](https://mlsanity.com/paper/1803.00942/) (2018)
- [GRAD-MATCH: Gradient Matching based Data Subset Selection for Efficient Deep Model Training](https://mlsanity.com/paper/2103.00123/) (2021)
- [GLISTER: Generalization based Data Subset Selection for Efficient and Robust Learning](https://mlsanity.com/paper/2012.10630/) (2020)
- [Dataset Distillation with Infinitely Wide Convolutional Networks](https://mlsanity.com/paper/2107.13034/) (2021)
- [Fast Computation of Wasserstein Barycenters](https://mlsanity.com/paper/1310.4375/) (2013)
- [NumGLUE: A Suite of Fundamental yet Challenging Mathematical Reasoning Tasks](https://mlsanity.com/paper/2204.05660/) (2022)
- [Lazier Than Lazy Greedy](https://mlsanity.com/paper/1409.7938/) (2014)
- [Restricted Strong Convexity Implies Weak Submodularity](https://mlsanity.com/paper/1612.00804/) (2016)
- [Multilevel Clustering via Wasserstein Means](https://mlsanity.com/paper/1706.03883/) (2017)
- [Towards Sustainable Learning: Coresets for Data-efficient Deep Learning](https://mlsanity.com/paper/2306.01244/) (2023)
- [Dataset Distillation via the Wasserstein Metric](https://mlsanity.com/paper/2311.18531/) (2023)
- [Minimizing the Maximal Loss: How and Why?](https://mlsanity.com/paper/1602.01690/) (2016)
- [TSDS: Data Selection for Task-Specific Model Finetuning](https://mlsanity.com/paper/2410.11303/) (2024)
- [Data Pruning by Information Maximization](https://mlsanity.com/paper/2506.01701/) (2025)
- [Diversified Batch Selection for Training Acceleration](https://mlsanity.com/paper/2406.04872/) (2024)
- [SPOT: A framework for selection of prototypes using optimal transport](https://mlsanity.com/paper/2103.10159/) (2021)
- [Mini-batch Coresets for Memory-efficient Language Model Training on Data Mixtures](https://mlsanity.com/paper/2407.19580/) (2024)
- [Submodular Mutual Information for Targeted Data Subset Selection](https://mlsanity.com/paper/2105.00043/) (2021)
- [On Partial Optimal Transport: Revising the Infeasibility of Sinkhorn and Efficient Gradient Methods](https://mlsanity.com/paper/2312.13970/) (2023)
- [Partial Wasserstein Covering](https://mlsanity.com/paper/2106.00886/) (2021)

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

- [A Riemannian Approach to Low-Rank Optimal Transport](https://mlsanity.com/paper/2606.12120/) (2026)
- [Minibatch Selection for Language Models via Partition Matroid Constrained Gradient Matching](https://mlsanity.com/paper/2606.07954/) (2026)

## Abstract (arXiv)

Selecting prototypical examples from a source distribution to represent a target data distribution is a fundamental problem in machine learning. Existing subset selection methods often rely on implicit importance scores, which can be skewed towards majority classes and lead to low-quality prototypes for minority classes. We present methodprop, a novel subset selection framework that minimizes the optimal transport (OT) distance between a uniformly weighted prototypical distribution and the target distribution. While intuitive, this formulation leads to a cardinality-constrained maximization of a super-additive objective, which is generally intractable to approximate efficiently. To address this, we propose a principled reformulation of the OT marginal constraints, yielding a partial optimal transport-based submodular objective. We prove that this reformulation enables a greedy algorithm with a (1-1/e) approximation guarantee relative to the original super-additive maximization problem. Empirically, we showcase that enforcing uniform prototype weights in UniPROT consistently improves minority-class representation in imbalanced classification benchmarks without compromising majority-class accuracy. In both finetuning and pretraining regimes for large language models under domain imbalance, UniPROT enforces uniform source contributions, yielding robust performance gains. Our results establish UniPROT as a scalable, theoretically grounded solution for uniform-weighted prototype selection. Our code is publicly available at GitHubfootnoteCode: https://github.com/efficiency-learning/UniPROT