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

Prateek Chanda, Prayas Agrawal, Karthik S. Gurumoorthy, Ganesh Ramakrishnan, Bamdev Mishra, Pratik Jawanpuria

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}}

Code Dataset Promise: Yes

Code Dataset Url: https://github.com/efficiency-learning/UniPROT

Format Confirmation: I agree that I have read and followed the formatting instructions for the camera ready version.

Submission Number: 2293