ScaleBiO: Scalable Bilevel Optimization for LLM Data Reweighting

TL;DR

This paper presents ScaleBiO, a scalable first-order bilevel optimization algorithm for large language model data reweighting, demonstrating its effectiveness and efficiency on models up to 30 billion parameters.

Contribution

It introduces ScaleBiO, the first practical, scalable bilevel optimization method for large-scale LLM data reweighting, combining with LISA for efficiency and demonstrating success on models up to 30B parameters.

Findings

ScaleBiO effectively improves instruction-following and math reasoning tasks.

It outperforms baseline data reweighting methods.

Theoretical guarantees include optimality and convergence for smooth, strongly convex objectives.

Abstract

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.