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

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

cs.LGApr 13, 2026

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

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

Organizations: Department of Computer Science and Engineering, Indian Institute of Technology Bombay · Microsoft Research India · Walmart Global Tech, India · Centre for Machine Intelligence and Data Science, Indian Institute of Technology Bombay ·

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

Figures & tables

Figure 1:

k

-medoids ( 4 ) consistently learn skewed weights for prototypes on CIFAR10 . The plot shows ranked weights of prototypes, averaged over 5 runs.

Figure 2: Exact marginal gain versus the proposed approximate marginal gain ( 8 ) on MNIST.

Panel of Figure 3 — Figure 3: Minority Class Accuracy Analysis :On CIFAR10-LT and synthetic MNIST, UniPROT outperforms $k$ -medoids on minority classes. ( Top Left ): Average minority-class accuracy vs. prototype count. ( Bottom Left ): CIFAR10 minority class-wise accuracy (avg. over 3 runs). ( Top Right ): MNIST with Induced Skew: (i) classes 0 and 5 both at 5% each (other classes at 90%), and (ii) classes 0 and 7 both at 3% each (other classes at 94%). ( Bottom Right ): UniPROT on average out-performs $k$ -medoids across different skew variations for minority classes.

Panel of Figure 6 — Figure 6: Perplexity dynamics on Pretraining LLaMA-3 500M and 60M for 20k steps.

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

Table 3: Comparison of performance of batch-wise selection across baselines for Phi-3 in SuperGlue ( Wang et al.,, 2019 ) .

Appendix figures & tables9 assets

Supplementary material from the paper’s appendix.

Appendix

Source Set Size	Max Iterations
64-200	100
500-1000	1000
1000-4000	2000
5000	4000

Appendix

Table 4: Source Size vs Maximum Iterations.

Panel of Figure 7 — Figure 7: Approximate vs. actual marginal gain as a function of the entropic regularization parameter $\lambda$ . Increasing $\lambda$ reduces the approximation error, with the marginal gain ratio converging to 1 and exhibiting low variance for larger $\lambda$ values.

Method	Avg	In-domain			Out-of-domain
		GSM8K	MATH	NumGLUE	SVAMP	Mathematics	SimulEq
COLM ( Nguyen et al.,, 2025 )	$59.86$	$74.13$	$36.70$	$63.14$	$86.50$	$36.40$	$62.30$
GREATS ( Wang et al.,, 2024 )	$60.79$	$78.62$	$37.90$	$63.90$	$85.50$	$36.90$	$61.90$
UniPROT -PB (Ours)	61.34	$78.20$	$37.60$	$66.03$	$84.90$	$37.70$	$63.60$

Appendix

Table 5: Comparison of UniPROT -PB performance for

|\mathcal{B}|=256

, following Yue et al., (2023) on MathInstruct dataset.

Panel of Figure 9 — Figure 9: Top: Per-batch performance of FT, GradNorm, COLM, GREATS, and UniPROT on MathInstruct with $|\mathcal{B}|=128$ , $k=64$ . Bottom: Validation perplexity when $|\mathcal{B}|=128$ , subset ratio 50%, and prototype selection is batch-wise.

Method	In-domain				Out-of-domain				Avg-All
	GSM8K	MATH	NumGLUE	Avg	SVAMP	Mathematics	SimulEq	Avg
FT	52.08	16.14	40.10	36.11	54.70	15.60	22.20	30.83	33.47
GradNorm	52.3 / 49.8	16.5 / 15.7	40.0 / 39.9	35.70	53.5 / 53.1	15.5 / 15.4	22.5 / 22.7	30.45	33.07
SBERT	51.3 / 21.7	14.6 / 14.01	41.6 / 27.44	28.44	56.3 / 34.5	16.2 / 13.6	21.98 / 13.8	26.06	27.25
COLM	50.6 / 49.88	21.42 / 17.4	40.15 / 39.9	36.56	55.8 / 54.1	16.7 / 15.4	22.17 / 21.01	30.86	33.71
GREATS	52.8 / 50.6	19.01 / 15.9	40.8 / 38.8	36.32	54.5 / 54.6	17.1 / 13.9	21.98 / 21.78	30.64	33.48

Appendix

Table 7: Performance across in-domain and out-of-domain datasets for Zephyr-3B on MathInstruct , batch size

|\mathcal{B}|=128

and total budget

k=32

. Results are reported under two configurations for all baselines: source-wise (left of “/”) and batch-wise (right of “/”).

Panel of Figure 10 — Figure 10: Zephyr-3b Validation perplexity when bs=128 and subset ratio 25% and prototype selection is batch wise.

Method	In-domain				Out-of-domain				Avg-All
	GSM8K	MATH	NumGLUE	Avg	SVAMP	Mathematics	SimulEq	Avg
FT	59.28	14.80	47.88	40.65	62.80	19.40	36.70	39.63	40.14
GradNorm	58.6 / 60.3	15.1 / 14.6	51.3 / 50.6	41.75	61 / 59	20.2 / 20.4	40.6 / 38.13	39.89	40.82
SBERT	58.4 / 45.5	12.43 / 9.83	52.01 / 48.4	37.76	61 / 64.2	18.1 / 16.6	36.5 / 33.8	38.37	38.06
COLM	57.01 / 58.8	14.86 / 14.6	51.05 / 50.8	41.19	60 / 63.9	18.9 / 19.8	32.4 / 31.9	37.82	39.50
GREATS	59.5 / 58.4	15.4 / 15.6	51.2 / 54.1	42.37	61.3 / 61.5	21.1 / 20.7	37.35 / 35.01	39.49	40.93

Appendix

Table 8: Performance across in-domain and out-of-domain datasets for Phi-2 on MathInstruct , batch size

|\mathcal{B}|=128

and total budget

k=32

. Results are reported under two configurations for all baselines: source-wise (left of “/”) and batch-wise (right of “/”).

Method	Avg	In-domain			Out-of-domain
		GSM8K	MATH	NumGLUE	SVAMP	Mathematics	SimulEq
COLM-PS ( Nguyen et al.,, 2025 )	$43.14$	$61.03$	$26.67$	$52.65$	$60.45$	$21.04$	$37.00$
GREATS-PS ( Wang et al.,, 2024 )	$43.73$	$61.92$	$27.21$	$52.40$	$62.00$	$20.80$	$38.03$
UniPROT -PS (Ours)	44.66	$62.40$	$27.63$	$53.97$	$64.72$	$20.30$	$38.91$

Appendix

Table 9: Comparison of Phi-2 per-source selection performance across in-domain and out-of-domain datasets, following Yue et al., (2023) on MathInstruct dataset, batch size

|\mathcal{B}|=256

and total budget

k=32

Class-Imbalanced Learning

Class-imbalanced learning develops classifiers that recognize rare classes without sacrificing performance on common ones.

See all 256 papers

Beyond Points: Spherical Distributional Part Prototypes for Interpretable Classification
Jun 25, 2026Duarte Leão, Diogo Pereira Araújo, Catarina Barata +1Prototype-Based Classification Image Classification
Instituto Superior Técnico, Universidade de Lisboa · Carnegie Mellon University
PRiSM: Prototype Regularization for Few-Shot VLMs
Jul 20, 2026Ghassen Baklouti, Omprakash Chakraborty, Jose Dolz +1VLM Adaptation Prototype Learning
ÉTS Montreal
DataProphet: Demystifying Supervision Data Generalization in Multimodal LLMs
Mar 20, 2026Xuan Qi, Luxi He, Dan Roth +1Multimodal Large Language Models Multimodal Dataset
University of Pennsylvania · Tsinghua University · Princeton University

Jun 25, 2026J/K move · Enter open · S save

Beyond Points: Spherical Distributional Part Prototypes for Interpretable Classification

Duarte Leão, Diogo Pereira Araújo, Catarina Barata, Carlos Santiago

Instituto Superior Técnico, Universidade de Lisboa · Carnegie Mellon University

Prototype-Based Classification Image Classification Prototype-Based Explanation

Open paper AskPDF arXiv

Prototype-based neural networks aim to provide intrinsic interpretability by grounding predictions in a small set of part prototypes. However, modern vision backbones typically operate in normalized, directional embedding spaces where each semantic part exhibits substantial intra-class variability. As a result, point prototypes often become redundant or unstable, hurting both explanation quality and robustness. We propose vMFProto, a distributional part-prototype framework that models each class as a mixture of von Mises-Fisher components on the hypersphere. Each prototype learns its own concentration, capturing part-specific variability, and we use entropic optimal transport (OT) to obtain structured patch-to-prototype assignments. A two-stage training schedule performs OT-driven prototype discovery followed by end-to-end refinement with patch-level distillation and distribution-aware diversity regularization. Experiments with frozen DINO backbones show that vMFProto achieves leading consistency and distinctiveness on CUB-200-2011 and competitive classification accuracy across CUB, Stanford Dogs, and Stanford Cars. Qualitative results confirm that vMFProto yields localized, non-redundant part evidence.

Abstract

Class-Imbalanced Learning

Explore similar work

Beyond Points: Spherical Distributional Part Prototypes for Interpretable Classification

PRiSM: Prototype Regularization for Few-Shot VLMs

DataProphet: Demystifying Supervision Data Generalization in Multimodal LLMs