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

Prateek ChandaPrayas AgrawalKarthik S. GurumoorthyGanesh RamakrishnanBamdev MishraPratik Jawanpuria

article2026arXiv2 citations

Introduces UniPROT, a subset selection framework that reformulates optimal transport into a submodular objective with a (1 - 1/e) greedy approximation guarantee, effectively preserving minority-class representation in imbalanced classification and language model training.

Listen

Modern artificial intelligence workflows increasingly rely on subset selection to identify small, representative groups of data points—known as prototypes—to summarize massive datasets and accelerate model training. Existing subset selection methods implicitly assign uneven importance scores to chosen examples. Under real-world conditions where data is imbalanced, this behavior disproportionately favors majority categories, leaving minority classes and rare sub-domains severely underrepresented and poorly modeled.

The article introduces and evaluates UniPROT, a framework designed to select an equally weighted, highly representative prototypical subset from source data to match a target data distribution. Its primary objective is to demonstrate that enforcing uniform prototype importance improves minority representation and training efficiency across imbalanced computer vision and natural language processing tasks.

To overcome the computational intractability of enforcing equal prototype weights, the authors reformulate the underlying optimal transport problem using a partial optimal transport approach with submodular properties. This reformulation mathematically proves that a fast, greedy selection algorithm can achieve a guaranteed theoretical approximation ratio while maintaining computational costs comparable to standard clustering techniques. The authors validate UniPROT across multiple settings, including nearest-neighbor image classification on skewed benchmarks such as MNIST and CIFAR-10, parameter-efficient fine-tuning of open-source language models like Phi-2, Phi-3, and Zephyr-3B on the highly imbalanced MathInstruct dataset, and pretraining LLaMA architectures on web text.

The experimental findings show clear advantages across several dimensions. First, UniPROT consistently boosts minority-class classification accuracy in long-tailed vision benchmarks without compromising majority-class performance. Second, in language model instruction fine-tuning at a 50% batch reduction budget, UniPROT outperforms full-batch fine-tuning and state-of-the-art data selection competitors across both in-domain and out-of-domain mathematical reasoning evaluations. Third, as data pruning becomes more aggressive—reducing mini-batch budgets to 25% or 12.5%—UniPROT maintains stable validation loss, whereas competing selection baselines degrade significantly. Finally, in language model pretraining, UniPROT achieves lower validation perplexity than standard full-batch training while processing fewer tokens.

These results demonstrate that enforcing equal exemplar weighting is an effective strategy for preventing data selection algorithms from neglecting critical minority groups. In practical deployments, this enables organizations to substantially cut compute expenses, memory overhead, and training timelines without sacrificing model quality or safety on rare edge cases. The framework also improves model interpretability by ensuring all selected exemplars contribute equally.

Organizations training models on heterogeneous, long-tailed data should consider adopting uniform prototype selection to construct high-quality training batches and coresets. When deploying the framework, engineering teams should configure the internal regularization parameter carefully, as empirical results confirm that tighter regularization yields superior transport fidelity and downstream model accuracy.

While the mathematical guarantees and empirical results are robust across vision and language benchmarks, the framework requires computing gradient or feature representations to establish pairwise similarities, which introduces a modest initial computational cost. Confidence in the reported performance gains is high for fine-tuning and classification tasks, though broader application to multi-modal architectures and distributed trillion-token training regimes remains an area for further validation.

Indian Institute of Technology Bombay Microsoft Walmart Labs classification classification models genetic selection Uniform Prototype Selection Partial Optimal Transport Submodular Guarantees

arXiv: 2604.10952 efficiency-learning/UniPROT

Paper: Optimal Transport for Domain Adaptation, Nicolas Courty et al. (2014). Establishes the foundational framework of aligning source and target distributions via optimal transport, providing key mathematical context for transport-based data alignment.
Paper: Active Learning for Convolutional Neural Networks: A Core-Set Approach, Ozan Sener et al. (2018). Formulates representative data subset selection as a geometric core-set problem, providing essential background for discrete optimization and greedy subset guarantees.
Paper: Prototypical Networks for Few-shot Learning, Jake Snell et al. (2017). Introduces the foundational concept of representing target data distributions and categories using prototypical exemplars in an embedding space.

No sufficiently relevant recommendations were found.

Cover for 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 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: this https URL}

Topics

classification classification models genetic selection language model large language models marginals robust performance protein database submodularity

Keywords

Uniform Prototype Selection Partial Optimal Transport Submodular Guarantees coreset selection submodular reformulation entropic regularization greedy algorithm mini-batch selection

1 Introduction
2 Preliminaries
3 Proposed Approach
3.1 Uniform Prototype Selection via Optimal Transport
3.2 Submodular Reformulation of ()
3.3 Computationally efficient approximate greedy algorithm for ()
4 Algorithm Details
5 Experimental Evaluation
5.1 Long Tailed Image Classification
5.2 High quality Mini-Batch Selection for LLM training
5.3 Evaluation Results and Discussion
6 Conclusion
7 Acknowledgements
References
A Organization of Appendix
B Theoretical Results
C Implementation Details
C.1 Hardware and License
C.2 Algorithm Implementation
C.3 Finetuning experiments
C.4 Details of baselines
C.5 Calculation of gradient features
C.6 Pretraining Experiments
D Experimental Setup Details
D.1 Model Details
D.2 Datasets
D.3 Training Details
D.4 Evaluation Datasets and Metrics
D.5 Evaluation Setup
E Additional Experimental Results
E.1 Additional Ablations on Entropic Regularization
E.2 Additional Results on UniPROT-PB
Experiments on batch size=256 for UniPROT-PB
Experiments on full-batch prototype selection
E.3 Additional Results on Zephyr-3B
E.4 Additional Results on Phi-2
F Additional Related Works
G Broader Impact
H Code

Knowls

Knowl 1 — Uniform Prototype Selection via Optimal Transport Formulation
model/method
Let $\mathcal{S} = \{x_i\}_{i=1}^m$ be a discrete source set and $\mathcal{T} = \{y_j\}_{j=1}^n$ be a target set, endowed with a ground cost matrix $C \in \mathbb{R}_+^{m \times n}$ where $C_{ij} = c(x_i, y_j)$ represents the cost of transporting unit mass from $x_i$ to $y_j$ . Let $S \in \mathbb{R}_+^{m \times n}$ be a non-negative similarity matrix defined by $S_{ij} = \beta - C_{ij}$ for a constant $\beta > \max_{ij} C_{ij}$ . For any subset of prototypes $P \subseteq \mathcal{S}$ , let $\mathbf{1}_P \in \{0, 1\}^m$ denote the indicator vector of $P$ (where $(\mathbf{1}_P)_i = 1$ if $x_i \in P$ and $0$ otherwise).

To enforce equal importance across all selected prototypes, the prototypical empirical distribution is restricted to uniform marginals $\mu_P = \frac{1}{|P|}\mathbf{1}_P$ . Assuming a uniform target distribution $\nu = \frac{1}{n}\mathbf{1}_n$ , the uniform prototype selection problem under a cardinality constraint $k$ seeks a subset $P \subseteq \mathcal{S}$ with $|P| \le k$ that maximizes the optimal transport similarity:

$\max_{P \subseteq \mathcal{S}, |P| \le k} g(P) \quad \text{where} \quad g(P) = \max_{\gamma \in \Gamma(\frac{1}{|P|}\mathbf{1}_P, \frac{1}{n}\mathbf{1}_n)} \langle S, \gamma \rangle$

Here, $\Gamma(\mu, \nu) = \{\gamma \in \mathbb{R}_+^{m \times n} \mid \gamma \mathbf{1}_n = \mu, \gamma^\top \mathbf{1}_m = \nu\}$ denotes the set of admissible transport couplings. Maximizing $g(P)$ is equivalent to maximizing the scaled proxy objective:

$\max_{P \subseteq \mathcal{S}, |P| \le k} h(P) \quad \text{where} \quad h(P) := |P|g(P) = \max_{\gamma \in \Gamma(\mathbf{1}_P, \frac{|P|}{n}\mathbf{1}_n)} \langle S, \gamma \rangle$

Unlike traditional $k$ -medoids or submodular facility location objectives that implicitly learn non-uniform prototype weights $\gamma \mathbf{1}_n$ , this formulation ensures every selected prototype contributes an identical share of total mass.
Knowl 2 — Submodular Partial Optimal Transport Reformulation and Approximation Guarantee
theoretical result
The scaled uniform prototype selection objective $h(P) = \max_{\gamma \in \Gamma(\mathbf{1}_P, \frac{|P|}{n}\mathbf{1}_n)} \langle S, \gamma \rangle$ is non-negative, monotone, and super-additive over disjoint sets ( $h(P_1 \cup P_2) \ge h(P_1) + h(P_2)$ for $P_1 \cap P_2 = \emptyset$ ), making direct greedy approximation intractable. To resolve this, a semi-relaxed partial optimal transport (POT) surrogate objective $f: 2^\mathcal{S} \to \mathbb{R}_+$ is formulated as:

$f(P) := \text{POT}\left(\mu_P = \mathbf{1}_P, \nu = \frac{k}{n}\mathbf{1}_n\right) = \max_{\gamma \in \Gamma_\le(\mathbf{1}_P, \frac{k}{n}\mathbf{1}_n)} \langle S, \gamma \rangle$

where $\Gamma_\le(\mathbf{1}_P, \frac{k}{n}\mathbf{1}_n) = \{\gamma \in \mathbb{R}_+^{m \times n} \mid \gamma \mathbf{1}_n = \mathbf{1}_P, \gamma^\top \mathbf{1}_m \le \frac{k}{n}\mathbf{1}_n\}$ .

This reformulation satisfies three key theoretical properties:
1. Submodularity and Monotonicity: The set function $f(P)$ is non-negative, monotone, and submodular over subsets $P \subseteq \mathcal{S}$ subject to cardinality constraint $k$ .
2. Exact Equivalence at Cardinality $k$ : For any set $P$ with $|P| = k$ , $f(P) = h(P)$ . Consequently, any optimal subset $P^*$ of cardinality $k$ for $\max_{|P|\le k} h(P)$ is also an optimal solution for $\max_{|P|\le k} f(P)$ , and vice-versa.
3. Approximation Bound: The classical greedy algorithm applied to $f(P)$ returns a set $\hat{P}$ of size $k$ satisfying:
$h(\hat{P}) = f(\hat{P}) \ge \left(1 - \frac{1}{e}\right)\text{OPT}$

where $\text{OPT} = h(P^*)$ is the optimal value of the original super-additive problem.
Knowl 3 — Computationally Efficient Approximate Marginal Gain Estimator and Weak Submodularity Guarantee
theoretical result
Evaluating the exact marginal gain $f(x_j \mid P) = f(P \cup \{x_j\}) - f(P)$ for all candidates $x_j \in \mathcal{S} \setminus P$ requires solving $(m - |P|)$ partial optimal transport (POT) problems per greedy step. To reduce this computational burden, an approximate marginal gain estimator $\hat{f}(x_j \mid P)$ is defined by fixing the transport plan on $P$ to $\gamma_P = \arg\max_{\gamma \in \Gamma_\le(\mathbf{1}_P, \frac{k}{n}\mathbf{1}_n)} \langle S, \gamma \rangle$ and optimizing only over the allocation vector $v \in \mathbb{R}_+^n$ for the candidate point $x_j$ :

$\hat{f}(x_j \mid P) = \max_{v \in \mathbb{R}_+^n, v^\top \mathbf{1}_n = 1, v \le \frac{k}{n}\mathbf{1}_n - \gamma_P^\top \mathbf{1}_m} \langle S(j, :), v^\top \rangle$

For a given $\gamma_P$ , this optimization admits an exact closed-form solution by greedily saturating the upper bounds $(\frac{k}{n}\mathbf{1}_n - \gamma_P^\top \mathbf{1}_m)$ on indices corresponding to the largest entries of the similarity row $S(j, :)$ , requiring only $O(n \log n)$ time via sorting.

Let $\alpha_{j,\min}$ and $\alpha_{j,\max}$ denote $\frac{1}{\lfloor n/k \rfloor}$ times the sum of the $\lfloor n/k \rfloor$ smallest and largest entries of $S(j, :)$ , respectively, and define the ratio:

$\alpha = \min_{j \in [m]} \frac{\alpha_{j,\min}}{\alpha_{j,\max}}$

If $\hat{P}$ is the subset of size $k$ selected by the greedy algorithm using the approximate marginal gain estimator $\hat{f}(x_j \mid P_i)$ at each step, then:

$f(\hat{P}) = h(\hat{P}) \ge \left(1 - e^{-\alpha}\right)\text{OPT}$

where $\text{OPT} = f(P^*) = h(P^*)$ is the optimal objective value for cardinality $k$ .
Knowl 4 — UniPROT Algorithm for Uniform Prototype Selection
algorithm
UniPROT selects a subset $P_k \subseteq \mathcal{S}$ of $k$ uniformly weighted prototypes from source set $\mathcal{S}$ ( $|\mathcal{S}| = m$ ) to represent target set $\mathcal{T}$ ( $|\mathcal{T}| = n$ ) using similarity matrix $S \in \mathbb{R}_+^{m \times n}$ .
Input: Similarity matrix $S$ between $\mathcal{S}$ and $\mathcal{T}$ , number of prototypes $k$ , entropic regularization parameter $\lambda$
Output: Uniformly weighted prototypical set $P_k \subseteq \mathcal{S}$
$P_0 \leftarrow \emptyset$
for $i = 1$ to $k$ do
     $\gamma^*_{P_{i-1}} \leftarrow \arg\max_{\gamma \in \Gamma_\le(\mathbf{1}_{P_{i-1}}, \frac{k}{n}\mathbf{1}_n)} \langle S, \gamma \rangle - \lambda \langle \gamma, \ln \gamma \rangle$
     $x^* \leftarrow \arg\max_{x \in \mathcal{S} \setminus P_{i-1}} \hat{f}(x \mid P_{i-1})$
     $P_i \leftarrow P_{i-1} \cup \{x^*\}$
end for
return $P_k$
In step 3, the regularized partial optimal transport problem is solved using Bregman-Dykstra projections or Sinkhorn iterations in $O(i \cdot n)$ time. In step 4, the approximate marginal gain $\hat{f}(x_j \mid P_{i-1}) = \max_{v \in \mathbb{R}_+^n, v^\top \mathbf{1}_n = 1, v \le \frac{k}{n}\mathbf{1}_n - (\gamma^*_{P_{i-1}})^\top \mathbf{1}_m} \langle S(j, :), v^\top \rangle$ is computed for candidate elements in $O((m - i) n \log n)$ time by sorting the similarity rows $S(j, :)$ .

For large source sets $\mathcal{S}$ , candidate evaluation can be restricted to a stochastic random sample $R \subseteq \mathcal{S} \setminus P_{i-1}$ of size $n k^{-1} \log(1/\epsilon)$ . By pre-sorting each row of $S$ in an initial $O(m n \log n)$ step using $O(mn)$ additional memory, the overall selection complexity of UniPROT is reduced to $O(kmn)$ .
Knowl 5 — Super-Additivity and Monotonicity of the Uniform Prototype Selection Proxy Function
theoretical result
Let $\mathcal{S}$ be a discrete source set and $S \in \mathbb{R}_+^{m \times n}$ be a non-negative similarity matrix. The set function $h: 2^\mathcal{S} \to \mathbb{R}_+$ defined by:

$h(P) = \text{OT}\left(\mu_P = \mathbf{1}_P, \nu = \frac{|P|}{n}\mathbf{1}_n\right) = \max_{\gamma \in \Gamma(\mathbf{1}_P, \frac{|P|}{n}\mathbf{1}_n)} \langle S, \gamma \rangle$

satisfies the following fundamental properties:
1. Non-negativity: $h(P) \ge 0$ for all $P \subseteq \mathcal{S}$ , since $S \ge 0$ and $\gamma \ge 0$ .
2. Monotonicity: $h(P_1) \le h(P_2)$ for all $P_1 \subseteq P_2 \subseteq \mathcal{S}$ . For an optimal coupling $\gamma_{P_1}$ of $P_1$ and $P_2 = P_1 \cup \{x_i\}$ , constructing a feasible coupling $\hat{\gamma}$ on $P_2$ with $\hat{\gamma}(i, :) = \frac{1}{n}\mathbf{1}_n^\top$ yields $h(P_2) \ge h(P_1) + \langle S(i, :), \frac{1}{n}\mathbf{1}_n \rangle \ge h(P_1)$ .
3. Super-additivity over Disjoint Sets: For any disjoint sets $P_1, P_2 \subseteq \mathcal{S}$ with $P_1 \cap P_2 = \emptyset$ :
$h(P_1 \cup P_2) \ge h(P_1) + h(P_2)$

This holds because concatenating the optimal couplings $\gamma_{P_1}$ and $\gamma_{P_2}$ forms a feasible transport coupling for $P_1 \cup P_2$ with marginals $\mathbf{1}_{P_1 \cup P_2}$ and $\frac{|P_1| + |P_2|}{n}\mathbf{1}_n$ .
Knowl 6 — Gradient-Based Similarity Construction for Mini-Batch Subset Selection in LLM Training
model/method
In mini-batch prototype selection for language model training, a subset of size $k$ is selected from a candidate batch $B_t$ of size $|B_t|$ to approximate the full batch gradient update. For each sample $z_i \in B_t$ , the gradient of the loss with respect to the LoRA-adapted value projection tensor of the topmost transformer block, $W_{V,\text{LoRA}}^{(L)}$ , is computed via backpropagation and flattened to $g_{i,t}^{\text{vp}} = \nabla_v \ell(\theta_t; z_i) \in \mathbb{R}^{d_{\text{vp}}}$ .

To align sample similarities with the Adam optimizer update direction, per-example gradients are normalized coordinate-wise using the running second-moment accumulator $v_t = \beta_2 v_{t-1} + (1 - \beta_2)\bar{g}_t^2$ :

$\phi_{i,t} = \frac{g_{i,t}^{\text{vp}}}{\epsilon + \sqrt{v_t}} \in \mathbb{R}^{d_{\text{vp}}}$

The pairwise similarity matrix $S \in \mathbb{R}_+^{|B_t| \times |B_t|}$ is computed using cosine affinities between the normalized vectors $\phi_{i,t}$ . Subset selection is then applied under two settings:
- Per-Source Selection (UniPROT-PS): When data source labels are available for $Q$ constituent domains, candidate batch $B_t$ is partitioned into $B_t^q = B_t \cap V_q$ , and UniPROT is solved independently on each source $B_t^q$ with local budget $k_q$ (such that $\sum_{q=1}^Q k_q = k$ ) and similarity matrix $S_q \in \mathbb{R}^{|B_t^q| \times |B_t^q|}$ .
- Per-Batch Selection (UniPROT-PB): In source-agnostic regimes, UniPROT is solved across the entire candidate batch $B_t$ directly using $S \in \mathbb{R}^{|B_t| \times |B_t|}$ .

Knowl 7 — Mathematical Reasoning Instruction Tuning Performance of Phi-3 on MathInstruct

data/table

Performance of Phi-3 (3.8B) fine-tuned on MathInstruct (260K instruction pairs across 14 imbalanced data sources) for 2048 steps with candidate batch size $|B| = 128$ and prototype budget $k = 64$ (effective batch size 64, 50% ratio). Results compare source-wise ("PS", left of "/") and batch-wise ("PB", right of "/") selection strategies against Full-Finetuning (FT) and other mini-batch selection baselines across in-domain and out-of-domain mathematical reasoning benchmarks:

Method	In-domain				Out-of-domain			Avg-All
	GSM8K	MATH	NumGLUE	Avg	SVAMP	Mathematics	SimulEq	Avg
FT	76.72	36.54	62.57	58.61	85.10	33.30	62.78	60.39	59.50
MaxLoss	70.64 / 69.44	32.05 / 30.02	57.80 / 58.90	53.50 / 52.79	80.60 / 79.20	31.45 / 30.06	57.19 / 55.82	56.41 / 55.03	54.96 / 53.91
GradNorm	76.04 / 75.40	36.10 / 35.03	64.01 / 64.10	58.72 / 58.18	85.30 / 84.17	38.00 / 36.50	61.84 / 65.70	61.71 / 62.12	60.22 / 60.15
SBERT	73.20 / 72.80	35.54 / 35.06	60.26 / 57.60	56.33 / 55.15	79.31 / 77.90	34.05 / 34.00	61.70 / 58.90	58.35 / 56.93	57.34 / 56.04
COLM	76.80 / 76.36	37.28 / 36.42	64.11 / 64.10	59.40 / 58.96	85.10 / 85.30	38.00 / 37.40	62.25 / 63.60	61.78 / 62.10	60.59 / 60.53
GREATS	76.72 / 77.80	37.84 / 37.28	67.46 / 64.40	60.67 / 59.83	86.10 / 85.00	35.60 / 38.19	62.06 / 61.92	61.25 / 61.64	60.96 / 60.73
UniPROT (Ours)	79.07 / 78.16	38.40 / 37.76	68.80 / 66.02	62.09 / 60.65	86.20 / 85.70	36.90 / 37.20	66.73 / 68.28	63.28 / 63.73	62.68 / 62.19

UniPROT consistently outperforms full fine-tuning (59.50% overall average) and all subset selection baselines (highest baseline GREATS at 60.96% PS / 60.73% PB) in both source-wise (62.68%) and batch-wise (62.19%) configurations.

Knowl 8 — Source-Agnostic LLM Finetuning Performance on SuperGLUE

data/table

When fine-tuning datasets lack domain or source annotations, per-source partitioning is not possible, requiring prototype selection to operate directly across the entire batch (batch-wise mode). The table below reports classification accuracy for Phi-3 fine-tuned on SuperGLUE classification tasks (SST-2, MultiRC, CB) for 512 steps with batch size $|B| = 32$ and a 25% prototype budget ( $k = 8$ ), compared against full-batch training (FT):

Method	SST2	MultiRC	CB	Avg
FT	93.91	86.05	92.72	90.89
GradNorm	87.94	57.54	69.10	71.53
SBERT	90.10	82.11	87.27	86.49
CoLM	94.72	82.99	93.05	90.25
GREATS	94.81	88.42	93.12	92.12
UniPROT-PB (Ours)	94.65	88.03	94.54	92.41

UniPROT-PB achieves the highest average accuracy (92.41%), outperforming full fine-tuning (90.89%) and competing subset selection methods (GREATS at 92.12%, CoLM at 90.25%), demonstrating that uniform optimal transport selection maintains high representational quality even in the absence of source metadata.

Knowl 9 — Minority Class Representation in Long-Tailed Image Classification
empirical result
Under unsupervised prototype selection from a balanced source set $\mathcal{S}$ to represent an imbalanced target set $\mathcal{T}$ with label shifts, evaluating a 1-nearest neighbour (1-NN) classifier parameterized by the selected prototype set $P$ demonstrates that UniPROT substantially improves classification accuracy on minority classes relative to $k$ -medoids and random selection.

On CIFAR10-LT and synthetically skewed MNIST (where two minority classes are restricted to 3% or 5% of target samples while majority classes occupy 90–94%), $k$ -medoids assigns low transport weights to prototypes from minority classes due to its unconstrained marginal formulation, leading to under-representation. UniPROT enforces uniform prototype weights in its objective, resulting in higher prototype allocations to minority classes and consistent accuracy gains on minority categories across prototype counts ( $k \in [10, 200]$ ) without degrading overall majority-class accuracy.
Knowl 10 — Robustness of UniPROT to Prototype Pruning Budgets and Entropic Regularization Parameter
empirical result
Empirical evaluations on language model training demonstrate that UniPROT is robust to aggressive subset pruning ratios and sensitive to the entropic regularization coefficient $\lambda$ :
1. Sensitivity to Prototype Ratios: When fine-tuning Zephyr-3B on MathInstruct for 2048 steps across prototype selection ratios $r \in \{50\%, 25\%, 12.5\%\}$ from candidate batches of size 128, UniPROT maintains stable validation log-perplexity across all ratios. In contrast, CoLM exhibits severe perplexity degradation as $r$ decreases, and GREATS displays a noticeable performance loss.
2. Entropic Regularization Strength: Evaluating downstream performance of Phi-3 fine-tuned on MathInstruct across $\lambda \in \{0.1, 0.05, 0.01\}$ shows that decreasing $\lambda$ from $0.1$ to $0.01$ consistently improves accuracy across downstream tasks (e.g., GSM8K accuracy increases from 47.76% at $\lambda = 0.1$ to 49.40% at $\lambda = 0.01$ , and SVAMP increases from 52.9% to 54.5%). Smaller entropic regularization improves the fidelity of the optimal transport matching to the true target geometry.
3. Pretraining Performance: Pretraining LLaMA-3 (60M and 500M) on OpenWebText for 20k steps with 50% subset selection demonstrates that UniPROT-PB achieves lower validation log-perplexity trajectories than full-batch pretraining, CoLM, and GREATS throughout training.

Coverage note — No substantial contributed material was omitted; the theoretical properties, submodular reformulation, approximate marginal gain estimator, algorithm details, gradient modeling methods, and empirical evaluations across image classification, instruction tuning, and pretraining benchmarks are all represented.

References

1.Adam, K. D. B. J. et al. (2014). A method for stochastic optimization. arXiv preprint arXiv:1412.6980, 1412(6).
2.Agueh, M. and Carlier, G. (2011). Barycenters in the wasserstein space. SIAM Journal on Mathematical Analysis, 43(2):904–924.
3.Benamou, J.-D., Carlier, G., Cuturi, M., Nenna, L., and Peyre, G. (2015). Iterative bregman projections for regularized transportation problems. SIAM Journal on Scientific Computing, 37(2):A1111–A1138.
4.Bien, J. and Tibshirani, R. (2011). Prototype selection for interpretable classification.
5.Bradbury, J., Frostig, R., Hawkins, P., Johnson, M. J., Leary, C., Maclaurin, D., Necula, G., Paszke, A., VanderPlas, J., Wanderman-Milne, S., and Zhang, Q. (2018). JAX: composable transformations of Python+NumPy programs.
6.Chapel, L., Alaya, M. Z., and Gasso, G. (2020). Partial optimal tranport with applications on positive-unlabeled learning. In NeurIPS.
7.Chen, C., Li, O., Tao, D., Barnett, A., Rudin, C., and Su, J. K. (2019). This looks like that: deep learning for interpretable image recognition. Advances in neural information processing systems, 32.
8.Cobbe, K., Kosaraju, V., Bavarian, M., Chen, M., Jun, H., Kaiser, L., Plappert, M., Tworek, J., Hilton, J., Nakano, R., et al. (2021). Training verifiers to solve math word problems. arXiv preprint arXiv:2110.14168.
9.Cuturi, M. (2013). Lightspeed computation of optimal transportation distances. Advances in Neural Information Processing Systems, 26(2):2292–2300.
10.Cuturi, M. and Doucet, A. (2014). Fast computation of wasserstein barycenters. In International conference on machine learning, pages 685–693. PMLR.
11.Das, A. and Kempe, D. (2018a). Approximate submodularity and its applications: Subset selection, sparse approximation and dictionary selection. Journal of Machine Learning Research, 19(3):1–34.
12.Das, A. and Kempe, D. (2018b). Approximate submodularity and its applications: Subset selection, sparse approximation and dictionary selection. Journal of Machine Learning Research, 19(3):1–34.
13.Davies, A., Velickovi ˇ c, P., Buesing, L., Blackwell, S., Zheng, D., Tomasev, N., Tanburn, R., Battaglia, P., Blundell, C., Juhasz, A., Lackenby, M., Williamson, G., Hassabis, D., and Kohli, P. (2021). Advancing mathematics by guiding human intuition with AI. Nature, 600(7887):70–74.
14.Dhurandhar, A. and Gurumoorthy, K. S. (2020). Classifier invariant approach to learn from positive-unlabeled data. In IEEE ICDM.
15.Elenberg, E. R., Khanna, R., Dimakis, A. G., and Negahban, S. (2018). Restricted strong convexity implies weak submodularity. Annals of Statistics, 46(6B):3539–3568.
16.Gurumoorthy, K. S., Dhurandhar, A., Cecchi, G., and Aggarwal, C. (2019). Efficient data representation by selecting prototypes with importance weights. In IEEE ICDM.
17.Gurumoorthy, K. S., Jawanpuria, P., and Mishra, B. (2021). Spot: A framework for selection of prototypes using optimal transport. In ECML PKDD.
18.Hendrycks, D., Burns, C., Kadavath, S., Arora, A., Basart, S., Tang, E., Song, D., and Steinhardt, J. (2021). Measuring mathematical problem solving with the math dataset. arXiv preprint arXiv:2103.03874.
19.Ho, N., Nguyen, X., Yurochkin, M., Bui, H. H., Huynh, V., and Phung, D. (2017). Multilevel clustering via wasserstein means. In International conference on machine learning, pages 1501–1509. PMLR.
20.Hong, F., Lyu, Y., Yao, J., Zhang, Y., Tsang, I., and Wang, Y. (2024). Diversified batch selection for training acceleration. In International Conference on Machine Learning, pages 18648–18667. PMLR.
21.Hu, E. J., Shen, Y., Wallis, P., Allen-Zhu, Z., Li, Y., Wang, S., Wang, L., Chen, W., et al. (2022). Lora: Low-rank adaptation of large language models. ICLR, 1(2):3.
22.Javaheripi, M., Bubeck, S., Abdin, M., Aneja, J., Bubeck, S., Mendes, C. C. T., Chen, W., Del Giorno, A., Eldan, R., Gopi, S., et al. (2023). Phi-2: The surprising power of small language models. Microsoft Research Blog.
23.Kantorovich, L. (1942). On the transfer of masses (in russian). Doklady Akademii Nauk, 37(2):227–229.
24.Katharopoulos, A. and Fleuret, F. (2018). Not all samples are created equal: Deep learning with importance sampling. In International conference on machine learning, pages 2525–2534. PMLR.
25.Kawano, K., Koide, S., and Otaki, K. (2022). Partial Wasserstein covering. In AAAI.
26.Killamsetty, K., Durga, S., Ramakrishnan, G., De, A., and Iyer, R. (2021a). Grad-match: Gradient matching based data subset selection for efficient deep model training. In International Conference on Machine Learning, pages 5464–5474. PMLR.
27.Killamsetty, K., Sivasubramanian, D., Ramakrishnan, G., and Iyer, R. (2021b). Glister: Generalization based data subset selection for efficient and robust learning. In Proceedings of the AAAI conference on artificial intelligence, volume 35, pages 8110–8118.
28.Kim, B., Khanna, R., and Koyejo, O. O. (2016). Examples are not enough, learn to criticize! criticism for interpretability. In NeurIPS.
29.Koncel-Kedziorski, R., Roy, S., Amini, A., Kushman, N., and Hajishirzi, H. (2016). Mawps: A math word problem repository. In Proceedings of the 2016 conference of the north american chapter of the association for computational linguistics: human language technologies, pages 1152–1157.
30.Kothawade, S., Kaushal, V., Ramakrishnan, G., Bilmes, J. A., and Iyer, R. K. (2021). Submodular mutual information for targeted data subset selection. CoRR, abs/2105.00043.
31.Krause, A. and Golovin, D. (2014). Submodular function maximization. Tractability, 3(71-104):3.
32.Krizhevsky, A., Hinton, G., et al. (2009). Learning multiple layers of features from tiny images.
33.Li, Y., Bubeck, S., Eldan, R., Del Giorno, A., Gunasekar, S., and Lee, Y. T. (2023). Textbooks are all you need ii: phi-1.5 technical report. arXiv preprint arXiv:2309.05463.
34.Lin, H. and Bilmes, J. (2011). A class of submodular functions for document summarization. In Proceedings of the 49th Annual Meeting of the Association for Computational Linguistics: Human Language Technologies.
35.Liu, H., Li, Y., Xing, T., Wang, P., Dalal, V., Li, L., He, J., and Wang, H. (2025). Dataset distillation via the wasserstein metric. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pages 1205–1215.
36.Liu, Z., Karbasi, A., and Rekatsinas, T. (2024). Tsds: Data selection for task-specific model finetuning. Advances in Neural Information Processing Systems, 37:10117–10147.
37.Menon, A. K., Rawat, A. S., Reddi, S., Kim, S., and Kumar, S. (2021). A statistical perspective on distillation. In International Conference on Machine Learning, pages 7632–7642. PMLR.
38.Minoux, M. (1978). Accelerated greedy algorithms for maximizing submodular set functions. In Optimization Techniques.
39.Mirzasoleiman, B., Badanidiyuru, A., Karbasi, A., Vondrak, J., and Krause, A. (2015). Lazier than lazy greedy. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 29.
40.Mirzasoleiman, B., Bilmes, J., and Leskovec, J. (2020a). Coresets for data-efficient training of machine learning models. In International Conference on Machine Learning.
41.Mirzasoleiman, B., Cao, K., and Leskovec, J. (2020b). Coresets for robust training of deep neural networks against noisy labels. Advances in Neural Information Processing Systems, 33:11465–11477.
42.Mishra, S., Mitra, A., Varshney, N., Sachdeva, B., Clark, P., Baral, C., and Kalyan, A. (2022). Numglue: A suite of fundamental yet challenging mathematical reasoning tasks. arXiv preprint arXiv:2204.05660.
43.Nemhauser, G. L., Wolsey, L. A., and Fisher, M. L. (1978). An analysis of approximations for maximizing submodular set functions—i. Mathematical programming, 14:265–294.
44.Nguyen, A. D., Nguyen, T. D., Nguyen, Q. M., Nguyen, H. H., Nguyen, L. M., and Toh, K.-C. (2024). On partial optimal transport: Revising the infeasibility of sinkhorn and efficient gradient methods. In AAAI.
45.Nguyen, D., Yang, W., Anand, R., Yang, Y., and Mirzasoleiman, B. (2025). Mini-batch coresets for memory-efficient language model training on data mixtures. In International Conference on Learning Representations.
46.Nguyen, T., Novak, R., Xiao, L., and Lee, J. (2021). Dataset distillation with infinitely wide convolutional networks. Advances in Neural Information Processing Systems, 34:5186–5198.
47.Patel, A., Bhattamishra, S., and Goyal, N. (2021). Are nlp models really able to solve simple math word problems? arXiv preprint arXiv:2103.07191.
48.Peyre, G., Cuturi, M., et al. (2019). Computational optimal transport: With applications to data science. Foundations and Trends® in Machine Learning, 11(5-6):355–607.
49.Radford, A., Wu, J., Child, R., Luan, D., Amodei, D., Sutskever, I., et al. (2019). Language models are unsupervised multitask learners. OpenAI blog, 1(8):9.
50.Reimers, N. and Gurevych, I. (2019). Sentence-bert: Sentence embeddings using siamese bert-networks. arXiv preprint arXiv:1908.10084.
51.Riaz, B., Karahan, Y., and Brockmeier, A. J. (2023). Partial optimal transport for support subset selection. Transactions on Machine Learning Research.
52.Schlegel, M., Pan, Y., Chen, J., and White, M. (2017). Adapting kernel representations online using submodular maximization. In Proceedings of the 34th International Conference on Machine Learning.
53.Shalev-Shwartz, S. and Wexler, Y. (2016). Minimizing the maximal loss: How and why. In International Conference on Machine Learning, pages 793–801. PMLR.
54.Solso, R. L., MacLin, O. H., and MacLin, M. K. (2017). Cognitive Psychology. Pearson Education.
55.Tan, H., Wu, S., Huang, W., Zhao, S., and QI, X. (2025). Data pruning by information maximization. In The Thirteenth International Conference on Learning Representations.
56.Tunstall, L., Beeching, E., Lambert, N., Rajani, N., Rasul, K., Belkada, Y., Huang, S., Von Werra, L., Fourrier, C., Habib, N., et al. (2023). Zephyr: Direct distillation of lm alignment. arXiv preprint arXiv:2310.16944.
57.Wang, A., Pruksachatkun, Y., Nangia, N., Singh, A., Michael, J., Hill, F., Levy, O., and Bowman, S. (2019). Superglue: A stickier benchmark for general-purpose language understanding systems. Advances in neural information processing systems, 32.
58.Wang, J., Dai, T., Zhang, B., Yu, S., Lim, E. G., and Xiao, J. (2025). Pot: Prototypical optimal transport for weakly supervised semantic segmentation. In Proceedings of the Computer Vision and Pattern Recognition Conference, pages 15055–15064.
59.Wang, J. T., Wu, T., Song, D., Mittal, P., and Jia, R. (2024). Greats: Online selection of high-quality data for llm training in every iteration. Advances in Neural Information Processing Systems, 37:131197–131223.
60.Yang, Y., Kang, H., and Mirzasoleiman, B. (2023). Towards sustainable learning: coresets for data-efficient deep learning. In International Conference on Machine Learning.
61.Yue, X., Qu, X., Zhang, G., Fu, Y., Huang, W., Sun, H., Su, Y., and Chen, W. (2023). Mammoth: Building math generalist models through hybrid instruction tuning. arXiv preprint arXiv:2309.05653.
62.Zhang, G., Zhang, H., Wang, Y., Li, R., Tan, H., and Liang, J. (2024). Hyperspherical multi-prototype with optimal transport for event argument extraction. In Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pages 9271–9284.
63.Zhao, B., Mopuri, K. R., and Bilen, H. (2021). Dataset condensation with gradient matching. In ICLR.
64.Zheng, H., Liu, R., Lai, F., and Prakash, A. (2023). Coverage-centric coreset selection for high pruning rates. In The Eleventh International Conference on Learning Representations.

Citation

MLA

Chanda, P., et al. “UniPROT: Uniform Prototype Selection via Partial Optimal Transport with Submodular Guarantees”. arXiv, 2026, http://arxiv.org/abs/2604.10952v1.

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. arXiv. http://arxiv.org/abs/2604.10952v1

Chicago

Chanda, P., P. Agrawal, K. S. Gurumoorthy, G. Ramakrishnan, B. Mishra, and P. Jawanpuria. 2026. “UniPROT: Uniform Prototype Selection via Partial Optimal Transport with Submodular Guarantees”. arXiv. http://arxiv.org/abs/2604.10952v1.

Harvard

Chanda, P. et al. (2026) “UniPROT: Uniform Prototype Selection via Partial Optimal Transport with Submodular Guarantees”, arXiv [Preprint]. Available at: http://arxiv.org/abs/2604.10952v1.

Vancouver

1. Chanda P, Agrawal P, Gurumoorthy KS, Ramakrishnan G, Mishra B, Jawanpuria P (2026) UniPROT: Uniform Prototype Selection via Partial Optimal Transport with Submodular Guarantees. arXiv

BibTeX

@article{chanda2026uniprot,
  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},
  year = {2026},
  journal = {arXiv},
  url = {http://arxiv.org/abs/2604.10952v1},
  eprint = {2604.10952}
}

Metadata:arXiv

Source Code

This paper has an official code repository available. Click below to access the source code.

View Repository

Access the Paper

This paper is available from its original source. Click below to access the PDF.

Open PDF

License: https://creativecommons.org/licenses/by/4.0/

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

Abstract

Topics

Keywords

Table of Contents

Knowls

Knowl 1 — Uniform Prototype Selection via Optimal Transport Formulation

Knowl 2 — Submodular Partial Optimal Transport Reformulation and Approximation Guarantee

Knowl 3 — Computationally Efficient Approximate Marginal Gain Estimator and Weak Submodularity Guarantee

Knowl 4 — UniPROT Algorithm for Uniform Prototype Selection

Knowl 5 — Super-Additivity and Monotonicity of the Uniform Prototype Selection Proxy Function

Knowl 6 — Gradient-Based Similarity Construction for Mini-Batch Subset Selection in LLM Training