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

📄 arXiv: 2604.10952v1 📥 PDF

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

分类: cs.LG

发布日期: 2026-04-13

备注: 25 pages, 31 figures. Accepted as a poster at AISTATS 2026

🔗 代码/项目: GITHUB

💡 一句话要点

UniPROT：基于部分最优传输和次模保证的均匀原型选择方法

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 原型选择 最优传输 次模函数 不平衡学习 子集选择

📋 核心要点

  1. 现有子集选择方法在不平衡数据集中易偏向多数类，导致少数类原型质量差，这是核心挑战。
  2. UniPROT的核心思想是最小化均匀加权原型分布与目标分布之间的最优传输距离，保证原型选择的公平性。
  3. 实验表明，UniPROT在不平衡分类和领域不平衡的LLM任务中，均能提升少数类表示和模型性能。

📝 摘要（中文）

本文提出了一种新的子集选择框架UniPROT，旨在解决从源分布中选择原型样本以代表目标数据分布的问题。现有方法通常依赖于隐式重要性得分，这可能偏向于多数类，导致少数类的原型质量较低。UniPROT通过最小化均匀加权原型分布与目标分布之间的最优传输(OT)距离来实现原型选择。该方法将OT边缘约束进行重新公式化，得到了一个基于部分最优传输的次模目标函数。理论上，证明了这种重新公式化使得贪婪算法能够达到相对于原始超可加最大化问题的(1-1/e)近似保证。实验表明，在不平衡分类基准测试中，UniPROT在不影响多数类准确性的前提下，持续改善少数类的表示。在领域不平衡的大型语言模型的微调和预训练中，UniPROT强制执行均匀的源贡献，从而产生稳健的性能提升。UniPROT是一种可扩展的、具有理论基础的均匀加权原型选择解决方案。

🔬 方法详解

问题定义：论文旨在解决原型选择问题，特别是在数据不平衡的情况下，现有方法倾向于选择代表多数类的原型，而忽略少数类。这导致模型在少数类上的性能较差。现有方法依赖于隐式的重要性得分，这些得分容易受到多数类的影响，无法保证原型选择的均匀性。

核心思路：UniPROT的核心思路是强制原型具有均匀的权重，并通过最小化均匀加权原型分布与目标分布之间的最优传输距离来进行原型选择。通过这种方式，确保每个原型对目标分布的贡献是相等的，从而避免了对多数类的偏向。这种设计旨在提高少数类的代表性，从而改善模型在不平衡数据集上的性能。

技术框架：UniPROT的整体框架包括以下几个主要步骤：1) 定义源分布和目标分布；2) 构建均匀加权的原型分布；3) 计算原型分布和目标分布之间的最优传输距离；4) 将最优传输问题转化为一个次模函数最大化问题；5) 使用贪婪算法选择原型子集。该框架的关键在于将原始的超可加目标函数转化为次模函数，从而可以使用贪婪算法进行高效的近似求解。

关键创新：UniPROT的关键创新在于将最优传输问题与次模函数最大化相结合，并提出了一种新的OT边缘约束的重新公式化方法。这种重新公式化使得可以使用贪婪算法以(1-1/e)的近似保证来解决原型选择问题。与现有方法相比，UniPROT能够保证原型选择的均匀性，从而更好地处理不平衡数据集。

关键设计：UniPROT的关键设计包括：1) 使用均匀权重来表示原型的重要性；2) 使用最优传输距离来衡量原型分布和目标分布之间的差异；3) 将最优传输问题转化为一个次模函数最大化问题，并使用贪婪算法进行求解。此外，论文还提出了一个基于部分最优传输的次模目标函数，以进一步提高原型选择的效率和准确性。具体的参数设置和损失函数取决于具体的应用场景，但核心思想是保持原型权重的均匀性。

📊 实验亮点

实验结果表明，UniPROT在不平衡分类任务中，能够显著提高少数类的分类精度，同时保持多数类的性能。在大型语言模型的预训练和微调中，UniPROT能够强制执行均匀的源贡献，从而产生稳健的性能提升。例如，在领域不平衡的LLM任务中，UniPROT能够带来显著的性能增益。

🎯 应用场景

UniPROT可应用于各种数据不平衡的机器学习任务，例如：医疗诊断、欺诈检测、罕见事件预测等。通过选择具有代表性的原型，可以提高模型在少数类上的性能，从而改善整体预测精度。此外，该方法还可用于大型语言模型的预训练和微调，以解决领域不平衡问题，提高模型的泛化能力。

📄 摘要（原文）

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}