Get more for less: Principled Data Selection for Warming Up Fine-Tuning in LLMs
Proposes GOT-D, an OT-gradient-based data selection method, improving fine-tuning efficiency and performance by up to 13.9%.
Key Findings
Methodology
The approach leverages the gradient of the optimal transport (OT) distance between candidate and target datasets to guide sample selection. It involves computing OT distances using entropy-regularized algorithms (e.g., Sinkhorn), deriving gradients via dual solutions, and iteratively selecting samples with the largest negative gradients. This process efficiently identifies samples that most effectively align the pretraining distribution with the target task distribution. The method integrates GPU-accelerated OT solvers, enabling the processing of millions of samples within minutes. The selected subset is then used for pre-fine-tuning, enhancing downstream task performance.
Key Results
- Across diverse tasks, the method outperforms baseline sampling strategies, achieving an average performance increase of 1.13% on 8 domain-specific NLP tasks and 13.9% on zero-shot tasks with models up to 2.7B parameters. Toxicity metrics for GPT-2 are also improved, with maximum toxicity scores reduced by 0.21 and toxicity probability dropping from 0.67 to 0.21. The approach scales efficiently, processing millions of samples within an hour on a single GPU.
- In low-budget scenarios (e.g., 50K samples), the data efficiency gains are significant, with performance improvements surpassing traditional methods. The method demonstrates robustness across tasks, model sizes, and data domains, validating its broad applicability.
- The experimental results confirm that OT-gradient-guided sample selection effectively reduces distribution mismatch, enhances model safety, and accelerates fine-tuning, making large-scale model adaptation more practical and cost-effective.
Significance
This work addresses a core challenge in large-scale NLP: how to efficiently adapt pre-trained models to specific tasks with limited data and computational resources. By theoretically grounding the sample selection process in OT gradient optimization, it offers a principled, scalable solution that improves data efficiency and model performance. The method's ability to rapidly process vast datasets and select the most impactful samples paves the way for more accessible, cost-effective deployment of large models in industry and research, especially where data annotation is expensive or limited. It also opens new avenues for integrating distributional metrics into active learning and domain adaptation workflows, fostering more intelligent and resource-efficient AI systems.
Technical Contribution
The paper introduces a novel sample selection algorithm based on the gradient of the OT distance, combining theoretical guarantees with practical efficiency. It formalizes the problem of minimizing the OT distance between the effective data distribution (a weighted combination of pretraining and selected samples) and the target distribution. The core innovation lies in deriving the OT gradient via dual solutions, enabling fast, large-scale computation with entropy regularization and GPU acceleration. The method extends existing distribution matching techniques by explicitly considering the weighted influence of pretraining data, providing a more accurate proxy for downstream performance. Empirical results across multiple NLP tasks demonstrate the superiority of this approach over traditional sampling and distribution-matching methods.
Novelty
This work is the first to incorporate OT gradient information directly into large-scale data selection for fine-tuning pre-trained language models. Unlike prior methods that focus solely on target distribution similarity, it emphasizes pulling the pretraining distribution closer to the target via gradient-guided sample selection. This shift addresses the mismatch between training and downstream performance, especially in low-data regimes. The integration of entropy-regularized OT solvers with gradient-based selection in a scalable manner constitutes a significant innovation, bridging theoretical insights with practical efficiency.
Limitations
- The method assumes candidate data roughly matches the pretraining distribution; if this assumption fails, effectiveness diminishes.
- In tasks requiring domain-specific knowledge far from pretraining data, the approach may underperform due to limited transferability.
- While GPU acceleration reduces computation time, OT problem solving at very large scales still incurs non-negligible costs, especially for highly specialized datasets.
Future Work
Future directions include extending the framework to multi-modal and multi-task settings, integrating active learning strategies, and exploring adaptive sampling during training. Improving OT solver efficiency further, especially for ultra-large datasets, is also a key goal. Additionally, incorporating domain knowledge and semi-supervised signals could enhance sample relevance, broadening applicability to specialized fields like biomedical or legal NLP. Developing dynamic, iterative sampling schemes that adapt during training could further optimize data efficiency and model robustness.
AI Executive Summary
Large-scale pre-trained language models (LLMs) have revolutionized natural language processing, but fine-tuning these models efficiently remains a significant challenge due to high data and computational costs. Traditional data selection strategies, such as random sampling or target distribution matching, often fall short in low-resource scenarios, failing to leverage the structural information embedded in pretraining data. This inefficiency hampers rapid adaptation, especially when dealing with limited labeled data or strict resource constraints.
This study introduces GOT-D, a novel data selection algorithm grounded in the gradient of the optimal transport (OT) distance. By measuring how each candidate sample influences the OT distance between the candidate dataset and the target distribution, the method identifies samples that most effectively align the pretraining distribution with the target task. The approach employs entropy-regularized OT solvers (e.g., Sinkhorn algorithm) combined with GPU acceleration, enabling rapid processing of millions of samples within minutes. The core idea is to select samples with the largest negative gradients, which most effectively reduce the distributional gap.
Empirical results across multiple NLP tasks—including domain-specific natural language understanding, generation, and zero-shot learning—demonstrate the method’s superiority over baseline techniques. For instance, in eight domain-specific tasks, performance improved by an average of 1.13%, and in zero-shot tasks with models up to 2.7B parameters, accuracy increased by 13.9%. Toxicity reduction experiments on GPT-2 further showcase the method’s ability to enhance safety, lowering toxicity scores significantly. The algorithm’s scalability, speed, and effectiveness make it a promising tool for cost-effective, high-performance model adaptation.
Overall, this work bridges theoretical insights from optimal transport with practical engineering, offering a scalable, principled solution to the data efficiency bottleneck in large-scale NLP. It opens pathways for more accessible deployment of powerful models, especially in resource-constrained environments, and sets a foundation for future research integrating distributional metrics into active learning and domain adaptation strategies.
Deep Analysis
Background
近年来，预训练语言模型（如GPT、BERT）在自然语言处理领域取得了巨大突破，但其微调成本不断攀升，尤其在数据有限或标注昂贵的场景中。传统方法多采用随机采样或简单的目标分布匹配，存在样本利用率低、效率不足的问题。随着模型规模的扩大，如何高效、有效地选择训练样本成为研究热点。最优传输（OT）作为一种衡量概率分布差异的数学工具，已在图像处理和统计学中广泛应用，逐渐被引入到深度学习中的数据选择策略中。该背景推动了结合OT理论设计高效样本筛选算法的研究需求。
Core Problem
当前大规模预训练模型微调面临两个核心难题：一是如何在有限预算下选择最具代表性和有效性的数据子集，二是如何确保所选样本能最大程度缩小预训练分布与目标任务分布的差异。传统方法多依赖于目标分布匹配，忽视了预训练数据的结构信息，导致样本效率低，效果有限。此外，处理超大规模数据时计算成本高，难以快速响应实际需求。如何在保证性能提升的同时，降低计算成本和数据需求，成为亟待解决的问题。
Innovation
本研究的创新点在于：1）引入最优传输（OT）距离的梯度信息，用于指导样本选择，2）结合熵正则化和GPU加速，实现百万样本的快速筛选，3）提出在低数据量条件下，最大负梯度样本的选择策略，有效缩小预训练与目标分布差异。与传统的目标分布匹配不同，该方法强调利用预训练数据的结构信息，优先拉近预训练分布与目标任务的距离，从而提升微调效果。这一策略在理论上得到证明，实验证明在多项任务中优于现有方法，显著降低成本，提高效率。
Methodology
- �� 构建候选数据集（如OpenWebTextCorpus）作为预训练分布的代理。• 计算候选集与目标任务数据的OT距离及其梯度，利用Kantorovich-Rubinstein对偶定理获得梯度信息。• 通过梯度下降，筛选出最大负梯度样本子集，优先加入训练。• 结合熵正则化，确保样本多样性和稳定性。• 利用GPU并行计算，快速完成大规模OT问题的求解。• 迭代优化样本子集，直到满足预算限制。• 最终用选中的样本进行预调优，提升模型对目标任务的适应性。
Experiments
采用多任务、多领域数据，包括NLU、NLG和零样本任务，数据集涵盖GLUE、毒性控制等。对比基线包括随机采样和传统分布匹配。指标包括性能提升（准确率、F1、毒性指标）和运行时间。超参数设置：样本规模（10K、20K）、训练轮次（3轮）、学习率（2e-5）。通过消融实验验证OT梯度引导的有效性，分析不同样本规模和模型参数对结果的影响。
Results
在多个任务中，筛选后模型性能优于对比方法。例如，8个领域任务中平均提升1.13%，零样本任务中性能提升达13.9%。毒性控制方面，筛选样本显著降低最大毒性得分0.21，毒性概率从0.67降至0.21。算法运行速度快，能在单GPU小时内处理百万样本，显示出极强的实用性和扩展性。这些结果验证了OT梯度引导样本选择在提升模型性能和安全性方面的有效性。
Applications
该方法适用于需要快速适应新任务或控制模型偏见的场景，如内容过滤、模型安全、定制化应用。只需少量未标注数据，即可大幅提升模型性能和安全性，降低标注成本。未来可结合主动学习和领域知识，进一步优化样本选择策略，推动大规模模型的高效微调。
Limitations & Outlook
该方法假设候选数据与预训练数据分布相近，若偏离较大，效果可能减弱。对极端专业领域或跨模态任务的适用性有限。此外，OT求解虽高效，但在超大规模数据集上仍存在计算成本，未来需优化算法以降低硬件依赖。
Plain Language Accessible to non-experts
想象你在准备一场大餐，但食材都来自不同的超市。为了做出最美味的菜，你希望挑选那些最能补充你厨房中缺少的食材。传统方法可能只是随机挑一些食材，或者只看它们是否来自你熟悉的超市。而这项研究就像用一种聪明的测量工具，衡量每样食材能多大程度上弥补你厨房的不足，然后优先挑选那些最能帮你做出美味佳肴的食材。这样，你花少的钱、用少的食材，就能做出更棒的菜。这种方法让你在有限的预算内，最大化你的厨艺潜力，既省钱又高效。
ELI14 Explained like you're 14
想象你在玩一个超级复杂的游戏，你的目标是让你的角色变得更厉害。可是，你只有有限的金币可以用来买装备。你可以随机买，也可以用一种聪明的方法，挑选那些最能帮你弥补技能短板的装备。这个方法就像用一个特殊的指南，告诉你哪些装备能让你的角色变得更强。它会分析所有可能的装备，找到那些最能弥补你缺少技能的装备，然后优先买它们。这样，你用少少的钱，就能变得更厉害，赢得比赛！这就像在大规模训练AI模型时，用一种聪明的策略，挑选最重要的数据，让模型变得更强、更安全、更聪明。
Abstract
This work focuses on leveraging and selecting from vast, unlabeled, open data to pre-fine-tune a pre-trained language model. The goal is to minimize the need for costly domain-specific data for subsequent fine-tuning while achieving desired performance levels. While many data selection algorithms have been designed for small-scale applications, rendering them unsuitable for our context, some emerging methods do cater to language data scales. However, they often prioritize data that aligns with the target distribution. While this strategy may be effective when training a model from scratch, it can yield limited results when the model has already been pre-trained on a different distribution. Differing from prior work, our key idea is to select data that nudges the pre-training distribution closer to the target distribution. We show the optimality of this approach for fine-tuning tasks under certain conditions. We demonstrate the efficacy of our methodology across a diverse array of tasks (NLU, NLG, zero-shot) with models up to 2.7B, showing that it consistently surpasses other selection methods. Moreover, our proposed method is significantly faster than existing techniques, scaling to millions of samples within a single GPU hour. Our code is open-sourced (Code repository: https://anonymous.4open.science/r/DV4LLM-D761/ ). While fine-tuning offers significant potential for enhancing performance across diverse tasks, its associated costs often limit its widespread adoption; with this work, we hope to lay the groundwork for cost-effective fine-tuning, making its benefits more accessible.