# ScaleBiO: Scalable Bilevel Optimization for LLM Data Reweighting

arXiv: https://arxiv.org/abs/2406.19976 · HTML: https://mlsanity.com/paper/2406.19976/

## Key facts

- Cited by 37 papers in the map (data through 14 Oct 2026).
- Builds on 57 of its 87 references that are in the map.
- First posted on arXiv 28 Jun 2024.

## Builds on (most cited first, 30 of 57)

- [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018)
- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [GPT-4 Technical Report](https://mlsanity.com/paper/2303.08774/) (2023)
- [Training Verifiers to Solve Math Word Problems](https://mlsanity.com/paper/2110.14168/) (2021)
- [Judging LLM-as-a-Judge with MT-Bench and Chatbot Arena](https://mlsanity.com/paper/2306.05685/) (2023)
- [Measuring Massive Multitask Language Understanding](https://mlsanity.com/paper/2009.03300/) (2020)
- [Measuring Mathematical Problem Solving With the MATH Dataset](https://mlsanity.com/paper/2103.03874/) (2021)
- [Think you have Solved Question Answering? Try ARC, the AI2 Reasoning Challenge](https://mlsanity.com/paper/1803.05457/) (2018)
- [Finetuned Language Models Are Zero-Shot Learners](https://mlsanity.com/paper/2109.01652/) (2021)
- [LIMA: Less Is More for Alignment](https://mlsanity.com/paper/2305.11206/) (2023)
- [GLM: General Language Model Pretraining with Autoregressive Blank Infilling](https://mlsanity.com/paper/2103.10360/) (2021)
- [GLaM: Efficient Scaling of Language Models with Mixture-of-Experts](https://mlsanity.com/paper/2112.06905/) (2021)
- [Aligning AI With Shared Human Values](https://mlsanity.com/paper/2008.02275/) (2020)
- [Yi: Open Foundation Models by 01.AI](https://mlsanity.com/paper/2403.04652/) (2024)
- [GPT-NeoX-20B: An Open-Source Autoregressive Language Model](https://mlsanity.com/paper/2204.06745/) (2022)
- [Instruction Tuning with GPT-4](https://mlsanity.com/paper/2304.03277/) (2023)
- [WildChat: 1M ChatGPT Interaction Logs in the Wild](https://mlsanity.com/paper/2405.01470/) (2024)
- [Learning to learn by gradient descent by gradient descent](https://mlsanity.com/paper/1606.04474/) (2016)
- [The Falcon Series of Open Language Models](https://mlsanity.com/paper/2311.16867/) (2023)
- [LESS: Selecting Influential Data for Targeted Instruction Tuning](https://mlsanity.com/paper/2402.04333/) (2024)
- [Meta-Learning with Implicit Gradients](https://mlsanity.com/paper/1909.04630/) (2019)
- [MAmmoTH: Building Math Generalist Models through Hybrid Instruction Tuning](https://mlsanity.com/paper/2309.05653/) (2023)
- [LMSYS-Chat-1M: A Large-Scale Real-World LLM Conversation Dataset](https://mlsanity.com/paper/2309.11998/) (2023)
- [Sheared LLaMA: Accelerating Language Model Pre-training via Structured Pruning](https://mlsanity.com/paper/2310.06694/) (2023)
- [How Far Can Camels Go? Exploring the State of Instruction Tuning on Open Resources](https://mlsanity.com/paper/2306.04751/) (2023)
- [Bilevel Programming for Hyperparameter Optimization and Meta-Learning](https://mlsanity.com/paper/1806.04910/) (2018)
- [DoReMi: Optimizing Data Mixtures Speeds Up Language Model Pretraining](https://mlsanity.com/paper/2305.10429/) (2023)
- [Gradient-based Hyperparameter Optimization through Reversible Learning](https://mlsanity.com/paper/1502.03492/) (2015)
- [Magpie: Alignment Data Synthesis from Scratch by Prompting Aligned LLMs with Nothing](https://mlsanity.com/paper/2406.08464/) (2024)
- [Data Selection for Language Models via Importance Resampling](https://mlsanity.com/paper/2302.03169/) (2023)

## Built on by (most cited first, 30 of 37)

- [SEAL: Safety-enhanced Aligned LLM Fine-tuning via Bilevel Data Selection](https://mlsanity.com/paper/2410.07471/) (2024)
- [The Best Instruction-Tuning Data are Those That Fit](https://mlsanity.com/paper/2502.04194/) (2025)
- [Unleashing the Power of Data Tsunami: A Comprehensive Survey on Data Assessment and Selection for Instruction Tuning of Language Models](https://mlsanity.com/paper/2408.02085/) (2024)
- [DataRater: Meta-Learned Dataset Curation](https://mlsanity.com/paper/2505.17895/) (2025)
- [More is not always better? Enhancing Many-Shot In-Context Learning with Differentiated and Reweighting Objectives](https://mlsanity.com/paper/2501.04070/) (2025)
- [Distributionally Robust Reinforcement Learning with Human Feedback](https://mlsanity.com/paper/2503.00539/) (2025)
- [ExeSQL: Self-Taught Text-to-SQL Models with Execution-Driven Bootstrapping for SQL Dialects](https://mlsanity.com/paper/2505.17231/) (2025)
- [AutoMixAlign: Adaptive Data Mixing for Multi-Task Preference Optimization in LLMs](https://mlsanity.com/paper/2506.00569/) (2025)
- [Safety at One Shot: Patching Fine-Tuned LLMs with A Single Instance](https://mlsanity.com/paper/2601.01887/) (2026)
- [Unbiased Gradient Low-Rank Projection](https://mlsanity.com/paper/2510.17802/) (2025)
- [StoSignSGD: Unbiased Structural Stochasticity Fixes SignSGD for Training Large Language Models](https://mlsanity.com/paper/2604.15416/) (2026)
- [A Survey of LLM × DATA](https://mlsanity.com/paper/2505.18458/) (2025)
- [Stochastic Bilevel Optimization with Heavy-Tailed Noise](https://mlsanity.com/paper/2509.14952/) (2025)
- [Unlocking TriLevel Learning with Level-Wise Zeroth Order Constraints: Distributed Algorithms and Provable Non-Asymptotic Convergence](https://mlsanity.com/paper/2412.07138/) (2024)
- [BLISS: A Lightweight Bilevel Influence Scoring Method for Data Selection in Language Model Pretraining](https://mlsanity.com/paper/2510.06048/) (2025)
- [On the Condition Number Dependency in Bilevel Optimization](https://mlsanity.com/paper/2511.22331/) (2025)
- [ADAPT: Learning Task Mixtures for Budget-Constrained Instruction Tuning](https://mlsanity.com/paper/2512.04555/) (2025)
- [TANDEM: Bi-Level Data Mixture Optimization with Twin Networks](https://mlsanity.com/paper/2606.04401/) (2026)
- [Differentially Private Bilevel Optimization](https://mlsanity.com/paper/2409.19800/) (2024)
- [Rethinking Data Mixture for Large Language Models: A Comprehensive Survey and New Perspectives](https://mlsanity.com/paper/2505.21598/) (2025)
- [Stable Preference Optimization: A Bilevel Approach to Catastrophic Preference Shift](https://mlsanity.com/paper/2507.07723/) (2025)
- [Computational Budget Should Be Considered in Data Selection](https://mlsanity.com/paper/2510.16806/) (2025)
- [Faster Gradient Methods for Highly-Smooth Stochastic Bilevel Optimization](https://mlsanity.com/paper/2509.02937/) (2025)
- [Holdout-Loss-Based Data Selection for LLM Finetuning via In-Context Learning](https://mlsanity.com/paper/2510.14459/) (2025)
- [Bridging Constraints and Stochasticity: A Fully First-Order Method for Stochastic Bilevel Optimization with Linear Constraints](https://mlsanity.com/paper/2511.09845/) (2025)
- [A Tale of Two Problems: Multi-Task Bilevel Learning Meets Equality Constrained Multi-Objective Optimization](https://mlsanity.com/paper/2605.09094/) (2026)
- [DomainPilot: Domain-Level Loss-Guided Two-Stage Data Mixture Optimization for Efficient Language Model Fine-Tuning](https://mlsanity.com/paper/2607.22769/) (2026)
- [FocuSFT: Bilevel Optimization for Dilution-Aware Long-Context Fine-Tuning](https://mlsanity.com/paper/2605.09932/) (2026)
- [BROS: Bias-Corrected Randomized Subspaces for Memory-Efficient Single-Loop Bilevel Optimization](https://mlsanity.com/paper/2605.10288/) (2026)
- [Bilevel Optimization of Synthetic Trajectories for Multi-Turn LLM Fine-Tuning](https://mlsanity.com/paper/2605.24743/) (2026)

## Abstract (arXiv)

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 8×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.