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

TL;DR

UniPROT introduces a theoretically grounded, scalable method for selecting uniformly weighted prototypes using partial optimal transport, improving minority class representation in imbalanced classification tasks.

Contribution

It reformulates optimal transport constraints into a submodular objective, enabling efficient greedy algorithms with guarantees, and demonstrates improved minority class performance.

Findings

Enforces uniform prototype weights improves minority-class representation.

Achieves robust performance gains in language models under domain imbalance.

Provides a scalable, theoretically justified prototype selection method.

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 $\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 \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 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 GitHub\footnote{Code: https://github.com/efficiency-learning/UniPROT}