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

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

• The paper introduces UniPROT, which enforces uniform prototype weighting by reformulating OT-based selection as a tractable submodular maximization problem.
• The paper demonstrates significant empirical gains, with nearly 10-point improvements on minority classes in long-tailed image classification and robust mini-batch selection for LLM training.
• The paper provides theoretical guarantees with a (1-1/e) optimality bound, bridging intractable super-additive maximization to efficient partial OT-based selection.

Uniform Prototype Selection via Partial Optimal Transport with Submodular Guarantees

Problem Formulation and Motivation

The paper addresses the longstanding problem of subset selection—specifically, the selection of a representative subset of $k$ prototypes from a source distribution to summarize a target distribution in high-dimensional settings. While optimal transport (OT) has been previously utilized as a distribution-matching criterion for prototype selection, existing algorithms (e.g., $k$-medoids, facility location, submodular maximization) inherently rely on weighted prototype sets, where implicitly-learned importance weights tend to favor majority classes or high-frequency regions of the data manifold. This results in poor representativity for minority modes—an effect empirically substantiated in the imbalance settings explored in this paper. The lack of interpretability and fairness in weighted selections is particularly problematic in human-in-the-loop and long-tailed domains.

The authors introduce UniPROT, a framework enforcing uniform weights on the selected prototypes. The objective transitions from unconstrained OT-based prototype selection to minimizing OT distance under the uniform source marginal constraint. The enforcement of uniform prototype weighting is shown to be essential for bias mitigation, transparent interpretability, and robust summarization of structured, imbalanced data.

Theoretical Framework and Algorithmic Contributions

The core technical challenge lies in the fact that uniform prototype selection renders the OT-based subset selection a monotone, non-negative, super-additive set function maximization under cardinality constraint (i.e., maximizing a function with increasing returns), which is intractable for classical greedy optimization and offers no a priori approximation bounds.

The main contributions are as follows:

• Formulation: The authors formalize the uniform prototype selection problem as the cardinality-constrained maximization of a super-additive functional (Eq. 6), where the empirical distribution over prototypes is uniform and OT distance to the target is minimized.
• Submodular Surrogate via Partial OT: Recognizing that super-additive maximization is intractable, they derive a tight, monotone, submodular surrogate by relaxing the target marginal constraints, resulting in a partial optimal transport (POT)-based objective (Eq. 7). The partial OT problem admits relaxation of target-side marginal constraints, making the function monotone and submodular.
• Equivalence and Guarantees: Crucially, it is proven that, at cardinality $k$, the surrogate and the original problem are tightly equivalent, meaning the surrogate maximization (a tractable submodular maximization) yields a solution with value at least $(1-1/e)$ times the global optimum for the original super-additive problem (see Lemma 4).
• Algorithmic Realization: UniPROT is instantiated as a greedy algorithm—at each step, computing marginal gains via the POT formulation and selecting elements with the highest approximate gain. The authors introduce an $O(n \log n)$ cost marginal gain approximation (Eq. 8) to accelerate selection, with a theoretical guarantee on the quality of solutions obtained.
• Computational Complexity: The total computational cost is $O(kmn \log n)$, matching the complexities of practical submodular selection baselines. Further acceleration techniques (caching, stochastic greedy) are discussed for very large datasets.

Empirical Results

Long-Tailed Image Classification

Experiments on MNIST and CIFAR-LT demonstrate that UniPROT significantly improves minority class accuracy over $k$-medoids and other subset selection methods, often by a margin of nearly 10 points on minority classes under high skew. The effect is pronounced in scenarios where weighted selection induces strong prototype mass concentration on majority classes, leading to minimal or null coverage of minority modes.

Mini-Batch Selection in LLM Pre-Training and Fine-Tuning

The utility of UniPROT is established in scalable LLM training, using mini-batch selection to simulate large-batch gradients under memory constraints. Methods that do not enforce uniform weighting (e.g., GREATS, CoLM, GradNorm) consistently underperform on minority-source domains and degrade for small selection budgets. In contrast, UniPROT maintains high quality across in-domain (MATHINSTRUCT, GSM8K, MATH, NumGLUE) and out-of-domain (SVAMP, Mathematics, SimulEq) evaluation sets, for multiple model families (PHI-3, ZEHPYR-3B, PHI-2).

In all reported configurations, UniPROT outperforms baselines with improvements up to 2–4 absolute points in average accuracy. Importantly, the method exhibits robustness to mini-batch selection ratio: validation perplexities vary little as selection budgets are reduced (down to 12.5%), unlike baselines that experience significant drops.

Additional Analyses

Ablation studies confirm that the entropic regularization used in POT computations can be tuned for both approximation speed and solution quality, with lower regularization yielding improved downstream accuracy on mathematical reasoning benchmarks.

Theoretical and Practical Implications

Enforcing uniform weighting for prototype selection directly counteracts systematic biases toward majority classes induced by standard OT-based or submodular selection methods. This shifts the landscape for interpretable summarization, memory- and compute-efficient coreset selection, and balanced training in the presence of long tails or rare phenomena.

The super-additive to submodular surrogate reduction, together with the theoretical $(1-1/e)$ optimality guarantee, sets a precedent for tackling intractable combinatorial objectives by careful reformulation through partial relaxation. The approach extends to other structured subset selection regimes, suggesting future avenues in constrained or multi-objective distributional matching.

For LLMs, the demonstrated gains in out-of-domain generalization underscore the importance of subset selection strategies that enforce equitable representation across diverse data sources. This is critical as models are increasingly finetuned or continually pretrained on massive but heterogeneously composed corpora.

Conclusion

The UniPROT framework solves the uniform prototype selection problem by reformulating the intractable OT-based super-additive maximization as a partial OT-based submodular maximization, enabling scalable greedy algorithms with tight theoretical guarantees. Across a comprehensive suite of experiments, UniPROT consistently outperforms weighted selection baselines, offers substantial improvements in minority-class and minority-source performance, and enables robust mini-batch selection for large model training in imbalanced regimes. The approach provides a blueprint for prototype selection applications demanding fairness, interpretability, and computational efficiency.

(2604.10952)