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

Prateek Chanda, Prayas Agrawal, Karthik S. Gurumoorthy, Ganesh Ramakrishnan, Bamdev Mishra, Pratik Jawanpuria
Proceedings of The 29th International Conference on Artificial Intelligence and Statistics, PMLR 300:5059-5067, 2026.

Abstract

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 \texttt{UniPROT}, 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 \emph{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 \texttt{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, \texttt{UniPROT} enforces uniform source contributions, yielding robust performance gains. Our results establish \texttt{UniPROT} as a scalable, theoretically grounded solution for uniform-weighted prototype selection. Our code is publicly available at GitHub\footnote{Code: \url{https://github.com/efficiency-learning/UniPROT}}

Cite this Paper

BibTeX
@InProceedings{pmlr-v300-chanda26a, title = { UniPROT: Uniform Prototype Selection via Partial Optimal Transport with Submodular Guarantees }, author = {Chanda, Prateek and Agrawal, Prayas and Gurumoorthy, Karthik S. and Ramakrishnan, Ganesh and Mishra, Bamdev and Jawanpuria, Pratik}, booktitle = {Proceedings of The 29th International Conference on Artificial Intelligence and Statistics}, pages = {5059--5067}, year = {2026}, editor = {Khan, Emtiyaz and Li, Yingzhen and Solin, Arno and Ramdas, Aaditya}, volume = {300}, series = {Proceedings of Machine Learning Research}, month = {02--05 May}, publisher = {PMLR}, pdf = {https://raw.githubusercontent.com/mlresearch/v300/main/assets/chanda26a/chanda26a.pdf}, url = {https://proceedings.mlr.press/v300/chanda26a.html}, abstract = { 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 \texttt{UniPROT}, 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 \emph{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 \texttt{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, \texttt{UniPROT} enforces uniform source contributions, yielding robust performance gains. Our results establish \texttt{UniPROT} as a scalable, theoretically grounded solution for uniform-weighted prototype selection. Our code is publicly available at GitHub\footnote{Code: \url{https://github.com/efficiency-learning/UniPROT}} } }
Endnote
%0 Conference Paper %T UniPROT: Uniform Prototype Selection via Partial Optimal Transport with Submodular Guarantees %A Prateek Chanda %A Prayas Agrawal %A Karthik S. Gurumoorthy %A Ganesh Ramakrishnan %A Bamdev Mishra %A Pratik Jawanpuria %B Proceedings of The 29th International Conference on Artificial Intelligence and Statistics %C Proceedings of Machine Learning Research %D 2026 %E Emtiyaz Khan %E Yingzhen Li %E Arno Solin %E Aaditya Ramdas %F pmlr-v300-chanda26a %I PMLR %P 5059--5067 %U https://proceedings.mlr.press/v300/chanda26a.html %V 300 %X 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 \texttt{UniPROT}, 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 \emph{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 \texttt{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, \texttt{UniPROT} enforces uniform source contributions, yielding robust performance gains. Our results establish \texttt{UniPROT} as a scalable, theoretically grounded solution for uniform-weighted prototype selection. Our code is publicly available at GitHub\footnote{Code: \url{https://github.com/efficiency-learning/UniPROT}}
APA
Chanda, P., Agrawal, P., Gurumoorthy, K.S., Ramakrishnan, G., Mishra, B. & Jawanpuria, P.. (2026). UniPROT: Uniform Prototype Selection via Partial Optimal Transport with Submodular Guarantees . Proceedings of The 29th International Conference on Artificial Intelligence and Statistics, in Proceedings of Machine Learning Research 300:5059-5067 Available from https://proceedings.mlr.press/v300/chanda26a.html.

Related Material