Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

ScaleBiO: Scalable Bilevel Optimization for LLM Data Reweighting

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

TL;DR

ScaleBiO addresses the challenge of learning optimal data-source weights for large-scale LLM training by formulating the problem as bilevel optimization and introducing a scalable, fully first-order minimax approach. The method leverages a memory-efficient, randomized block-coordinate scheme (inspired by LISA) to enable training on models up to ~30B parameters, with a minimax reformulation that yields first-order hypergradients and theoretical convergence guarantees. Empirically, ScaleBiO outperforms strong baselines in instruction following and mathematical reasoning across multiple models and data sources, including a real-world MT-Bench evaluation on about 4.2 million samples from 18 sources. The work demonstrates both practical scalability and theoretical optimality of learned data weights, offering a principled pipeline for data filtering and selection in large-scale LLM fine-tuning, while acknowledging limitations related to pre-training-scale validation and potential multi-objective extensions.

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.

ScaleBiO: Scalable Bilevel Optimization for LLM Data Reweighting

TL;DR

ScaleBiO addresses the challenge of learning optimal data-source weights for large-scale LLM training by formulating the problem as bilevel optimization and introducing a scalable, fully first-order minimax approach. The method leverages a memory-efficient, randomized block-coordinate scheme (inspired by LISA) to enable training on models up to ~30B parameters, with a minimax reformulation that yields first-order hypergradients and theoretical convergence guarantees. Empirically, ScaleBiO outperforms strong baselines in instruction following and mathematical reasoning across multiple models and data sources, including a real-world MT-Bench evaluation on about 4.2 million samples from 18 sources. The work demonstrates both practical scalability and theoretical optimality of learned data weights, offering a principled pipeline for data filtering and selection in large-scale LLM fine-tuning, while acknowledging limitations related to pre-training-scale validation and potential multi-objective extensions.

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 30B-sized LLMs on 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.
Paper Structure (36 sections, 7 theorems, 50 equations, 3 figures, 11 tables)

This paper contains 36 sections, 7 theorems, 50 equations, 3 figures, 11 tables.

Table of Contents

  1. Introduction
  2. Related Work
  3. Bilevel Optimization
  4. Data Reweighting
  5. Methods
  6. Problem Formulation
  7. Fully First-order Hypergradient Method
  8. Proposed Algorithm
  9. Theoretical Results
  10. Experiments
  11. Small Scale Experiments
  12. Data Denoising
  13. Multilingual Training
  14. Instruction Following
  15. Real-World Application Experiments
  16. ...and 21 more sections

Key Result

Lemma 1

Under Assumption assumpt:v2, if $\alpha > 2\ell_{11}/\mu_2$, we have where the condition number $\kappa$ is defined by $\max\left\lbrace \ell_{10}, \ell_{11}, \ell_{21}, \ell_{22} \right\rbrace / \mu_2$.

Figures (3)

  • Figure 1: Data denoising with GPT-2: weights for noisy data and clean data.
  • Figure 2: Multilingual reweighting with GPT-2: weights of Chinese and English. Training set: 1:1; Validation set: 6:4.
  • Figure 3: Instruction Following with GPT-2: weights for Alpaca-GPT4 and Alpaca.

Theorems & Definitions (8)

  • Lemma 1
  • Theorem 1
  • Lemma 2
  • Lemma 3
  • Lemma 4
  • Lemma 5
  • Lemma 6
  • proof