RL-Guided Data Selection for Language Model Finetuning
作者: Animesh Jha, Harshit Gupta, Ananjan Nandi
分类: cs.LG
发布日期: 2025-09-30
备注: To appear in NeurIPS 2025 Constrained Optimization for ML Workshop
💡 一句话要点
提出基于强化学习的数据选择方法，提升大语言模型微调效率与性能。
🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture) 支柱九：具身大模型 (Embodied Foundation Models)
关键词: 强化学习 数据选择 大语言模型 微调 马尔可夫决策过程
📋 核心要点
- 现有数据选择方法难以有效应用于大语言模型微调，无法在有限预算下最大化模型性能。
- 将数据选择问题建模为马尔可夫决策过程，利用强化学习训练智能体学习最优选择策略。
- 实验表明，仅使用5%的数据子集，即可达到甚至超过全量数据微调的性能，并显著减少训练时间。
📝 摘要（中文）
本文将大语言模型（LLM）微调的数据选择问题建模为带预算约束的优化问题，目标是在严格的训练数据预算下最大化模型的下游性能。由于直接求解该问题通常是难以处理的，并且现有的近似方法主要面向预训练，迁移到微调设置时效果不佳。因此，本文将该问题重新表述为一个可处理的马尔可夫决策过程（MDP），并使用各种强化学习（RL）方法训练智能体，以学习最优的数据选择策略，该策略由高效的、基于代理模型的奖励信号引导。在四个数据集上的实验表明，使用本文方法选择的5%子集进行训练，其性能与在完整数据集上进行微调的性能相匹配甚至优于后者，准确率最高提升达10.8个百分点，同时减少了高达2倍的实际训练时间，突显了强化学习引导的数据选择的潜力。
🔬 方法详解
问题定义：论文旨在解决大语言模型微调过程中，如何在有限的数据预算下，选择最具信息量的子集，以最大化模型在下游任务上的性能。现有方法，尤其是那些面向预训练的数据选择方法，在微调场景下表现不佳，无法有效利用数据，导致模型性能提升受限，且计算成本高昂。
核心思路：论文的核心思路是将数据选择过程建模为一个马尔可夫决策过程（MDP）。智能体通过与环境交互，学习选择哪些数据样本进行微调。这种方法允许智能体根据模型的反馈动态调整选择策略，从而找到最优的数据子集。核心在于利用强化学习的探索能力，在庞大的数据空间中高效地搜索有价值的样本。
技术框架：整体框架包含以下几个主要模块：1) 环境：由数据集和代理模型组成，代理模型用于快速评估数据样本对模型性能的潜在影响。2) 智能体：使用强化学习算法（如PPO）训练，负责根据当前状态（例如，已选择的数据样本）选择下一个要添加的数据样本。3) 奖励函数：基于代理模型的性能变化，为智能体的选择行为提供反馈。奖励函数的设计至关重要，它引导智能体学习选择能够最大化模型性能的数据样本。4) 训练循环：智能体与环境交互，不断优化数据选择策略。
关键创新：论文的关键创新在于将数据选择问题转化为强化学习问题，并利用代理模型加速奖励信号的评估。与传统的基于启发式规则或预训练目标的数据选择方法不同，该方法能够根据微调任务的特定需求，动态地学习最优的数据选择策略。此外，使用代理模型可以显著降低评估数据样本价值的计算成本。
关键设计：论文中，代理模型通常是一个规模较小的语言模型，它被训练来近似完整模型的行为。奖励函数的设计考虑了模型在验证集上的性能提升，以及选择数据样本的成本。智能体使用近端策略优化（PPO）算法进行训练，该算法能够稳定地学习复杂的策略。具体参数设置（如学习率、折扣因子等）需要根据具体任务进行调整。
🖼️ 关键图片
📊 实验亮点
实验结果表明，使用强化学习选择的5%数据子集进行微调，在四个数据集上均能达到或超过全量数据微调的性能，准确率最高提升10.8个百分点，同时训练时间缩短高达2倍。这证明了该方法在提升模型性能和效率方面的显著优势。
🎯 应用场景
该研究成果可广泛应用于自然语言处理领域，尤其是在资源受限的场景下，例如移动设备上的模型部署、低成本的AI服务等。通过智能地选择训练数据，可以显著降低模型训练成本，提高模型性能，加速AI技术的普及和应用。
📄 摘要（原文）
Data selection for finetuning Large Language Models (LLMs) can be framed as a budget-constrained optimization problem: maximizing a model's downstream performance under a strict training data budget. Solving this problem is generally intractable, and existing approximate approaches are pretraining-oriented and transfer poorly to the fine-tuning setting. We reformulate this problem as a tractable Markov Decision Process (MDP) and train agents using various Reinforcement Learning (RL) methods to learn optimal data selection policies, guided by an efficient, proxy-model-based reward signal. Across four datasets, training on a $5\%$ subset selected by our approach matches or outperforms fine-tuning on the full dataset by up to $10.8$ accuracy points, while cutting wall-clock training time by up to $2 \times$, highlighting the promise of RL-guided data selection.