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 · Microsoft India
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
Figure 2: Exact marginal gain versus the proposed approximate marginal gain ( 8 ) on MNIST.
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.
Figure 5
Figure 5: Validation log-pplx with changing prototype percentage (left) and ablation table (right) on Zephyr-3b .
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.
Figure 7: Approximate vs. actual marginal gain as a function of the entropic regularization parameter λ . Increasing λ reduces the approximation error, with the marginal gain ratio converging to 1 and exhibiting low variance for larger λ values.
Figure 11
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 ∣B∣=256 , following Yue et al., (2023) on MathInstruct dataset.
Figure 9: Top: Per-batch performance of FT, GradNorm, COLM, GREATS, and UniPROT on MathInstruct with ∣B∣=128 , k=64 . Bottom: Validation perplexity when ∣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 ∣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 “/”).
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 ∣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 ∣B∣=256 and total budget k=32 .
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.
Training-free few-shot adaptation methods have gained significant attention recently in the context of Vision-language Models (VLMs). Yet, current benchmarks rely on strong assumptions about the statistics of the adaptation data, e.g., class balance. We question these simplifying assumptions and introduce a more realistic benchmark that varies both the levels of class balance and the effective number of classes in few-shot tasks via Dirichlet sampling. Surprisingly, under our setting, we observe substantial drops in the performances of state-of-the-art methods, more so when the number of labeled samples increases. To mitigate this, we introduce PRiSM, a class-prototype regularization that can be deployed as a plug and play module on top of any existing baseline method, significantly improving performances. Our method optimizes a novel multi-term loss, which includes a regularizer maximizing inter-class pairwise distances, along with additional terms promoting support-feature alignment and fidelity to the baseline prototypes. Furthermore, we introduce an effective and computationally efficient block Majorize-Minimize optimizer for our objective. More specifically, we derive a valid blockwise Lipschitz constant (i.e., a bound on the Hessian's spectral norm), which can be computed efficiently via the Gershgorin circle theorem. Extensive experiments show that PRiSM improves several training-free baselines, with large gains when dealing with severe class imbalance and high numbers of classes.
Ghassen Baklouti, Omprakash Chakraborty, Jose Dolz +1
Conventional wisdom for selecting supervision data for multimodal large language models (MLLMs) is to prioritize datasets that appear similar to the target benchmark, such as text-intensive or vision-centric tasks. However, it remains unclear whether such intuitive similarity reliably predicts downstream performance gains. In this work, we take a first step toward answering a practical question: can we estimate the influence of a training dataset on a target benchmark before any training is performed? To investigate this question, we conduct an in-depth analysis of transfer across 14 vision-language datasets spanning 7 diverse tasks. Our results show that intuitive task similarity is an unreliable predictor of transferability, and that generalization depends more on the specific dataset than on its broad task category. Motivated by this finding, we propose DATAPROPHET, a simple and effective training-free metric that combines multimodal perplexity, similarity, and data diversity. Experiments show that DATAPROPHET produces supervision-data rankings that strongly correlate with rankings based on actual post-training performance gains, achieving a Kendall's tau of 86.0%. Moreover, DATAPROPHET enables better supervision-data selection, yielding up to 6.9% improvement over uniform selection, 1.4% over a state-of-the-art training-based baseline, and 0.2% above oracle selection based on experimental performance. Our code and data will be released.
Xuan Qi, Luxi He, Dan Roth +1
University of Pennsylvania · Tsinghua University · Princeton University