ScaleBiO：面向大语言模型数据重加权的可扩展双层优化
机器学习
2025-05-27 v2 最优化与控制
摘要
双层优化在各种机器学习场景中展现出重要作用，但大多数实际算法需要二阶信息，难以规模化。近期仅有理论文献中涌现出一类一阶算法范式，能够有效解决双层优化问题。然而，该范式在实际中的效率尚未得到验证，特别是在大型语言模型（LLM）的上下文中。本文引入首个该范式的可扩展实现，命名为 ScaleBiO，专注于大规模 LLM 数据重加权。通过结合近期提出的内存高效训练技术 LISA，我们的新算法使该范式能够在 8 张 H100 GPU 上规模化至约 300 亿参数的 LLM，标志着在实际场景下首次成功地将双层优化应用于大规模 LLM。经实验验证，ScaleBiO 在数据重加权任务中有效，包括 Llama-3-8B、Gemma-2-9B、Qwen-2-7B 和 Qwen-2.5-32B 等不同规模模型，双层优化在指令跟随和数学推理任务中取得优异性能，优于均匀抽样、影响感知数据过滤和参考模型抽样等多个流行基线方法。在理论上，ScaleBiO 确保所学习数据权重的最优性，并给出与传统一阶双层优化范式相符的收敛保证，适用于光滑和强凸目标。
引用
@article{arxiv.2406.19976,
title = {ScaleBiO: Scalable Bilevel Optimization for LLM Data Reweighting},
author = {Rui Pan and Dylan Zhang and Hanning Zhang and Xingyuan Pan and Minrui Xu and Jipeng Zhang and Renjie Pi and Xiaoyu Wang and Tong Zhang},
journal= {arXiv preprint arXiv:2406.19976},
year = {2025}
}
备注
ACL 2025