PaperlayerView exact revision on arXiv

Source-linked AI summary

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

arXiv:2406.19976v2cs.LGmath.OC

TL;DR

Bilevel optimization has been difficult to scale because existing algorithms often require costly second-order information, and the practical efficiency of first-order methods on large LLMs remains unverified. ScaleBiO applies bilevel optimization to large-scale LLM data reweighting. The paper reports scalable data reweighting and performance improvements on downstream tasks, with theoretical guarantees under stated assumptions.

  • Problem

    Bilevel optimization algorithms often require second-order information, and the practical performance of first-order methods on large-scale LLM data reweighting has not been thoroughly investigated.

    Source: §1 · p. 2

  • Method

    ScaleBiO is a bilevel optimization approach for learning LLM data weights that uses a single-loop framework and randomized block-coordinate updates.

    Source: §3.3 · p. 5

  • Results

    ScaleBiO enables data reweighting on models over 30B parameters, and its learned weights are reported to improve downstream performance.

    Source: §6 · p. 8

  • Takeaways & Limitations

    The authors present ScaleBiO as an efficient data-filtering and selection pipeline for improving model performance on downstream tasks.

    Source: §6 · p. 8

    ScaleBiO has not yet been verified in large-scale pre-training settings because those experiments require substantial computational resources.

    Source: Limitations · p. 9

Abstract

from arXiv · show

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.

1 Introduction

ScaleBiO addresses the difficulty of learning data-source weights for large LLMs, where bilevel optimization has been hard to scale. It combines a scalable bilevel approach with theoretical guarantees and reports gains over popular data-reweighting baselines.

1 Introduction

  • Data reweighting is used across goals including fairness, faster pre-training, training robustness, and transfer learning.The introduction notes that reweighting and filtering can improve performance across diverse tasks.Source: §1 · p. 1
  • Prior bilevel algorithms often require Hessian, Jacobian, or vector-product computations, which can be expensive and memory-intensive at large scale.A recently proposed fully first-order stochastic method uses only first-order gradient information.Source: §1 · p. 2
  • The practical gap is learning useful sampling weights across LLM training data sources, since some datasets can degrade task-specific performance.The paper frames this as data-task misalignment and asks how to balance sources for optimal performance.Source: §1 · p. 2, How to balance each data source in the… · p. 2

How to balance each data source in the training

  • Existing data-weighting approaches rely on preset choices or lack theoretical guarantees, motivating tests of theoretically optimal bilevel optimization for LLM reweighting.The paper identifies these approaches as leading to suboptimal sampling weights.Source: How to balance each data source in the… · p. 2
  • ScaleBiO is presented as the first scalable, theoretically optimal bilevel optimization approach for large LLM training, capable of scaling to models with about 30 billion parameters.The authors say it makes optimal data weights learnable for large-scale LLMs.Source: How to balance each data source in the… · p. 2
  • ScaleBiO surpasses uniform sampling, LESS, and RHO-LOSS by 1%–9% on GSM8K and MATH, with superiority also reported on instruction-following tasks.Its convergence guarantee matches Kwon et al. (2023) on smooth and strongly convex objectives.Source: How to balance each data source in the… · p. 2

2 Related Work

Prior bilevel methods face computational or scalability limits, while recent first-order theory has not been thoroughly tested on large LLMs. Existing data-reweighting methods also do not ensure optimal learned weights.

2.1 Bilevel Optimization

  • Traditional and stochastic bilevel methods can remain impractical at scale because they require second-order information, including Jacobian- and Hessian-vector products.Finite-difference Jacobian approximations may be sensitive to the smoothing constant and numerically problematic.Source: §2.1 · pp. 2–3
  • Fully first-order penalty-based methods avoid Hessian or Jacobian estimation, but their practical use in large-scale LLM settings remains unexplored.The cited work includes stochastic bilevel methods and a deterministic method with near-optimal sample complexity.Source: §2.1 · p. 3
  • Earlier practical NLP applications used finite differences or Hessian-vector products, and the authors report that no approach listed in Table 1 had been applied to over 3B-sized LLMs.Table 1 compares implemented model sizes, methods, and tested tasks.Source: §2.1 · p. 3, Table 1 · p. 3

2.2 Data Reweighting

  • Prior data-reweighting methods use signals such as loss gaps, self-influence, reference losses, or distributionally robust optimization, but do not ensure optimal learned weights.The cited approaches include methods for online weighting, validation-based adjustment, and domain weighting without downstream-task knowledge.Source: §2.2 · p. 3

3 Methods

ScaleBiO formulates source-mixture selection as bilevel optimization of validation loss after training on weighted data, then solves a minimax reformulation with a single-loop method. Randomized block updates target scale, and the theory provides a stationarity complexity bound under a constant batch size.

3.1 Problem Formulation

  • Each data source receives a sampling probability, and the source mixture is optimized to minimize validation loss after training on the weighted data.The source probabilities form a distribution over the available sources.Source: §3.1 · p. 4

3.2 Fully First-order Hypergradient Method

  • An auxiliary variable u detaches the inner–outer dependency by expressing the inner optimum as a training-loss-gap constraint in a minimax formulation.As α tends to infinity, the minimax and bilevel problems are equivalent under certain smoothness assumptions.Source: §3.2 · p. 4

3.3 Proposed Algorithm

  • ScaleBiO is a single-loop solver for the reformulated min-max problem, using randomized block updates for u and w to reduce memory consumption.The penalty coefficient α is set large to make the min-max solution approximate the original bilevel problem.Source: §3.3 · p. 5

3.4 Theoretical Results

  • For B=O(1), Algorithm 1 has O(ε^-7) complexity for finding an ε-stationary point; Γα is O(κ^3)-smooth independently of α.The complexity matches that reported by Kwon et al. (2023).Source: §3.4 · p. 5

4 Experiments

Experiments span GPT-2 toy tasks and real-world LLM instruction-following and mathematical reasoning. ScaleBiO adapts source weights in toy settings, outperforms baselines on real-world tasks, and scales to Qwen-2.5-32B.

4.1.1 Data Denoising

  • ScaleBiO assigns minimal weight to noisy samples, demonstrating automatic denoising in the GPT-2 experiment.The training data combined clean Alpaca examples with examples whose outputs were replaced by periods.Source: §4.1.1 · p. 6, Figure 1 · p. 6

4.1.2 Multilingual Training

  • 6:4: ScaleBiO nearly recovers the validation-set Chinese-to-English mix from training data sampled at 1:1.The validation set contains 600 Chinese and 400 English examples; the training set contains 40,000 examples from each source.Source: §4.1.2 · p. 6, Figure 2 · p. 6

4.1.3 Instruction Following

  • 10% of training data came from Alpaca-GPT4, which ScaleBiO up-weighted as the higher-quality source.Alpaca-GPT4 shares Alpaca's instructions and inputs but uses outputs generated by GPT-4.Source: §4.1.3 · pp. 6–7, Figure 3 · p. 6

4.2.1 Instruction Following

  • ScaleBiO is the only bilevel algorithm yielding meaningful source weights and outperforms LESS and RHO-LOSS on MT-Bench instruction following.MT-Bench scores are evaluated with a GPT-4o judge on a 0-to-10 scale.Source: §4.2.1 · p. 7, Table 2 · p. 7

4.2.2 Mathematical Reasoning

  • By 1%–9%, ScaleBiO consistently outperforms all baselines across models and benchmarks on mathematical reasoning.The comparison covers GSM8K and MATH accuracies.Source: §4.2.2 · p. 8, Table 4 · p. 8
  • On Qwen-2.5-32B, ScaleBiO is the only data reweighting implementation to scale; LESS and RHO-LOSS run out of GPU memory.ScaleBiO has the same space complexity as full-parameter fine-tuning.Source: §4.2.2 · p. 8, Table 5 · p. 8

5 Discussion

ScaleBiO’s learned weights can transfer from Llama-3-8B to Llama-3-70B with some improvement, while weights vary substantially across model families and less within a family.

5 Discussion

  • Learned data weights vary significantly across model families, but the difference is much smaller within the same family.The authors conjecture that differences in pretraining-data distributions and model strengths may drive distinct dataset needs for the same downstream task.Source: §5 · p. 8
  • Llama-3-8B-learned weights yield certain improvement when transferred to Llama-3-70B.Table 6 reports MT-Bench results for Llama-3-70B using transferred weights.Source: Table 6 · p. 8, §5 · p. 8
  • With small datasets, optimal weights shared across models may be impossible, leaving model-dependent reweighting as the only choice.The authors present this as a consequence of the hypothesized differences in pretraining distributions and model strengths.Source: §5 · p. 8

6 Conclusion

ScaleBiO scales bilevel optimization for LLM data reweighting to models over 30B parameters, with theoretical guarantees and empirical downstream-task benefits.

6 Conclusion

  • ScaleBiO is the first bilevel optimization instantiation reported to scale beyond 30B-sized LLMs for data reweighting.The paper reports improved performance on various downstream tasks through an efficient data filtering and selection pipeline.Source: §6 · p. 8
  • ScaleBiO ensures optimality of learned data weights and matches conventional first-order penalty-based bilevel optimization convergence guarantees on smooth, strongly convex objectives.Source: §6 · p. 8
  • Empirically, ScaleBiO enables data reweighting on models larger than 30B parameters and improves performance across downstream tasks.Source: §6 · p. 8

Limitations

The paper identifies unverified large-scale pretraining and single-loss risks as limitations, while noting potential training-cost reductions from more effective reweighting.

Limitations

  • ScaleBiO has not been verified in large-scale pretraining, which would require substantial computational resources.The authors hope large-scale fine-tuning results can be a first step toward that setting.Source: Limitations · p. 9
  • Optimizing sampling weights on a single loss metric may lead models to neglect other aspects, including safety or ethics.The authors recommend multi-objective losses and post-training alignment to compensate.Source: Limitations · p. 9
  • More effective data reweighting may further reduce the training cost of large language models.Source: Limitations · p. 9

Ethical Considerations

The paper reports licensing compliance and public availability for its experimental datasets, and lists data-weight tables for instruction-following and mathematical-reasoning settings.

Ethical Considerations

  • The instruction-following experiments use publicly available datasets, list their licenses in the appendix, and report no privacy issues.The named datasets include ShareGPT, OpenOrca, WildChat, AlpacaChat, LMSYS-Chat, and Airoboros.Source: Ethical Considerations · p. 9

A.1 Data Weights across Model Families

  • Table 7 lists the top-three weighted data sources for six backbone models in instruction-following tasks.Source: Table 7 · p. 14, §A.1 · p. 14
  • Table 8 lists the top-three weighted data sources for LLaMA-3-8B, Qwen-2-7B, and Gemma-2-9B in mathematical-reasoning tasks.Source: Table 8 · p. 14

A.2 Mathematical Reasoning: Stronger Benchmarks

On stronger mathematical-reasoning benchmarks, ScaleBiO consistently outperforms baselines on pass@1, while its Cons@64 performance is comparable. The experiments use a smaller, higher-quality dataset and evaluate a fine-tuned model on AIME24, AIME25, and AIMO25.

A.2 Mathematical Reasoning: Stronger Benchmarks

  • The setup uses 8K prompts from four math datasets, generates question-answer pairs with Deepseek-R1, selects 4K training samples, and evaluates on three reference sets.AIME24 and AIME25 each contain 30 questions, while AIMO25 contains 10.Source: §A.2 · p. 14
  • ScaleBiO consistently outperforms baselines on pass@1 across AIME2024, AIME2025, and AIMO2025.Table 9 compares methods on these benchmarks for DeepSeek-R1-Distill-Qwen-1.5B.Source: Table 9 · p. 14, §A.2 · p. 14
  • On Cons@64, ScaleBiO is comparable to baselines; the authors conjecture the narrowing gap reflects limited diversity in the small dataset.They expect a larger amount of data to improve the gap.Source: §A.2 · pp. 14–15

B.1 Small Scale Experiments

  • Small-scale experiments use GPT-2 with 124 million parameters and train for 3 epochs with batch size 64 and alpha 10.The sampling-weight learning rate is 10^-2; model learning rates are 10^-5.Source: §B.1 · p. 15

B.2 Large Scale Experiments

The instruction-following experiments use 18 data sources, an MT-Bench-style reference set, and data reweighting to select samples for backbone training. The appendix also specifies a separate validation split for mathematical reasoning.

Instruction Following

  • The instruction-following portfolio has 18 sources: nine high-quality datasets and nine multilingual Alpaca datasets treated as irrelevant sources.Conversations or instructions longer than 1,024 tokens are filtered out.Source: Table 10 · p. 15, §B.2 · p. 16
  • The reference set contains 1,200 MT-Bench-distributed samples, with 600 similarly generated samples reserved for hyperparameter tuning.GPT-4 prompts generate three two-turn instruction sets for each of MT-Bench’s eight categories.Source: §B.2 · p. 16
  • ScaleBiO reweights 3,000 samples per source, then samples 10,000 training examples using the learned weights; baselines may use the additional 1,800 MT-Bench-styled samples.Fine-tuning uses one epoch, an initial learning rate of 8 × 10^-6, and global batch size 64.Source: §B.2 · p. 16

Mathematical Reasoning

  • For mathematical reasoning, 280 samples from each of six validation sources form a 1,680-sample set, split by ScaleBiO into 1,400 reweighting and 280 tuning samples.Other baselines use the full validation set.Source: §B.2 · pp. 16–17, Table 11 · p. 17

C Important Lemmas

Under Assumption 1, the lemmas establish regularity properties for the penalized objective and its solution mappings. The proof of Theorem 1 estimates inner-variable errors using randomized block-coordinate updates and curvature conditions.

C Important Lemmas

D Proofs of Theorem 1

  • The proof bounds inner-variable errors with randomized block-coordinate updates, using strong concavity in u and strong convexity in w under its stated alpha condition.The analysis also invokes independent training and validation samples and a bound on minibatch-gradient variance.Source: D Proofs of Theorem 1 · pp. 18–21
Loading 2406.19976v2…