ScaleBiO: Scalable Bilevel Optimization for LLM Data Reweighting

📄 arXiv: 2406.19976v2 📥 PDF

作者: Rui Pan, Dylan Zhang, Hanning Zhang, Xingyuan Pan, Minrui Xu, Jipeng Zhang, Renjie Pi, Xiaoyu Wang, Tong Zhang

分类: cs.LG, math.OC

发布日期: 2024-06-28 (更新: 2025-05-25)

备注: ACL 2025

💡 一句话要点

提出ScaleBiO，通过可扩展双层优化实现LLM数据重加权，显著提升指令跟随和数学推理能力。

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

关键词: 双层优化 大型语言模型 数据重加权 一阶优化 内存高效训练 指令跟随 数学推理

📋 核心要点

  1. 现有双层优化算法依赖二阶信息，计算成本高昂，难以扩展到大型语言模型的数据重加权任务。
  2. ScaleBiO结合一阶双层优化范式和内存高效训练技术LISA，实现了在大型LLM上的可扩展数据重加权。
  3. 实验表明，ScaleBiO在指令跟随和数学推理任务中，显著优于均匀采样等多种数据选择基线方法。

📝 摘要（中文）

双层优化已在各种机器学习场景中展现出实用性，但大多数算法需要二阶信息，难以扩展。最近，理论文献中出现了一种一阶算法范式，能够有效解决双层优化问题。然而，这种范式的实际效率尚未得到验证，尤其是在大型语言模型（LLM）的背景下。本文介绍了该范式的第一个可扩展实例，名为ScaleBiO，专注于大规模LLM数据重加权的双层优化。通过结合最近提出的内存高效训练技术LISA，我们的算法使该范式能够在8个H100 GPU上扩展到约30B大小的LLM，标志着双层优化首次成功应用于大型LLM的实际场景。实验表明，ScaleBiO在不同规模的模型（包括Llama-3-8B、Gemma-2-9B、Qwen-2-7B和Qwen-2.5-32B）上的数据重加权验证了其有效性，在指令跟随和数学推理任务中优于几种流行的基线方法，包括均匀采样、影响感知数据过滤和基于参考模型的采样方法。理论上，ScaleBiO确保了学习到的数据权重的最优性，以及在平滑和强凸目标上的收敛保证，与传统的一阶双层优化范式相匹配。

🔬 方法详解

问题定义：论文旨在解决大型语言模型（LLM）数据重加权问题。现有双层优化方法由于需要计算二阶信息，计算复杂度高，难以扩展到参数量巨大的LLM上。这限制了双层优化在LLM数据选择和优化的应用。

核心思路：论文的核心思路是采用一阶双层优化算法范式，并结合内存高效的训练技术LISA，从而降低计算和内存需求，实现双层优化算法在大型LLM上的可扩展性。通过优化数据权重，提升LLM在特定任务上的性能。

技术框架：ScaleBiO的技术框架主要包含两个层面：外层优化和内层优化。外层优化负责更新数据权重，内层优化负责训练LLM。LISA技术用于在内层优化中减少内存占用。整个流程迭代进行，直至数据权重收敛。

关键创新：ScaleBiO的关键创新在于将一阶双层优化范式成功应用于大规模LLM的数据重加权任务。这是首次在如此大的模型上实现双层优化，并取得了显著的性能提升。结合LISA技术是实现可扩展性的关键。

关键设计：ScaleBiO的关键设计包括：1) 使用一阶梯度估计来近似二阶信息，降低计算复杂度；2) 采用LISA技术进行内存高效训练，允许在有限的GPU资源上训练大型模型；3) 精心设计的损失函数，用于指导数据权重的学习，以提升模型在目标任务上的性能。

📊 实验亮点

实验结果表明，ScaleBiO在Llama-3-8B、Gemma-2-9B、Qwen-2-7B和Qwen-2.5-32B等不同规模的LLM上均取得了显著的性能提升。在指令跟随和数学推理任务中，ScaleBiO优于均匀采样、影响感知数据过滤和基于参考模型的采样方法等基线方法，验证了其有效性。

🎯 应用场景

ScaleBiO可应用于各种需要数据选择和优化的LLM应用场景，例如指令微调、领域知识注入、对抗样本防御等。通过优化训练数据分布，提升LLM在特定任务上的性能和鲁棒性，具有广泛的应用前景。

📄 摘要（原文）

Bilevel optimization has shown its utility across various machine learning settings, yet most algorithms in practice require second-order information, making it challenging to scale them up. Only recently, a paradigm of first-order algorithms has emerged in the theoretical literature, capable of effectively addressing bilevel optimization problems. Nevertheless, the practical efficiency of this paradigm remains unverified, particularly in the context of large language models (LLMs). This paper introduces the first scalable instantiation of this paradigm called ScaleBiO, focusing on bilevel optimization for large-scale LLM data reweighting. By combining with a recently proposed memory-efficient training technique called LISA, our novel algorithm allows the paradigm to scale to $\sim$30B-sized LLMs on $8\times$H100 GPUs, marking the first successful application of bilevel optimization under practical scenarios for large-sized LLMs. Empirically, extensive experiments on data reweighting verify the effectiveness of ScaleBiO for different-scaled models, including Llama-3-8B, Gemma-2-9B, Qwen-2-7B, and Qwen-2.5-32B, where bilevel optimization succeeds in instruction-following and math reasoning tasks, outperforming several popular baselines, including uniform sampling, influence-aware data filtering, and reference-model-based sampling methods. Theoretically, ScaleBiO ensures the optimality of the learned data weights, along with a convergence guarantee matching the conventional first-order bilevel optimization paradigm on smooth and strongly convex objectives.