ScaleBiO: Scalable Bilevel Optimization for LLM Data Reweighting

• URL: http://arxiv.org/abs/2406.19976v2
• Date: Sun, 25 May 2025 05:03:51 GMT
• Title: ScaleBiO: Scalable Bilevel Optimization for LLM Data Reweighting
• Authors: Rui Pan, Dylan Zhang, Hanning Zhang, Xingyuan Pan, Minrui Xu, Jipeng Zhang, Renjie Pi, Xiaoyu Wang, Tong Zhang,
• Abstract summary: Bilevel optimization has shown its utility across various machine learning settings, yet most algorithms in practice require second-order information.<n>This paper introduces the first scalable instantiation of this paradigm called ScaleBiO, focusing on bilevel optimization for large-scale LLM data reweighting.<n>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 $8times$H100 GPU.
• Score: 26.931958430968024
• License: http://creativecommons.org/licenses/by/4.0/
• 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.

Related papers

• Highly Efficient and Effective LLMs with Multi-Boolean Architectures [1.4195677954898822]
  Weight binarization has emerged as a promising strategy to drastically reduce the complexity of large language models (LLMs)<n>We introduce a novel framework that effectively transforms LLMs into multi- kernel Boolean parameters, for the first time, finetunes them directly in the Boolean domain, eliminating the need for expensive latent weights.<n>Our method outperforms recent ultra low-bit quantization and binarization methods.
  arXiv  Detail & Related papers  (2025-05-28T19:40:34Z)
• GOLLuM: Gaussian Process Optimized LLMs -- Reframing LLM Finetuning through Bayesian Optimization [0.4037357056611557]
  Large Language Models (LLMs) can encode complex relationships in their latent spaces. We introduce LLM-based deep kernels, jointly optimized with GPs to preserve the benefits of both. Our method nearly doubles the discovery rate of high-performing reactions compared to static LLM embeddings.
  arXiv  Detail & Related papers  (2025-04-08T17:59:57Z)
• Make Optimization Once and for All with Fine-grained Guidance [78.14885351827232]
  Learning to Optimize (L2O) enhances optimization efficiency with integrated neural networks. L2O paradigms achieve great outcomes, e.g., refitting, generating unseen solutions iteratively or directly. Our analyses explore general framework for learning optimization, called Diff-L2O, focusing on augmenting solutions from a wider view.
  arXiv  Detail & Related papers  (2025-03-14T14:48:12Z)
• Provably Faster Algorithms for Bilevel Optimization via Without-Replacement Sampling [96.47086913559289]
  gradient-based algorithms are widely used in bilevel optimization. We introduce a without-replacement sampling based algorithm which achieves a faster convergence rate. We validate our algorithms over both synthetic and real-world applications.
  arXiv  Detail & Related papers  (2024-11-07T17:05:31Z)
• Memory-Efficient Gradient Unrolling for Large-Scale Bi-level Optimization [71.35604981129838]
  Traditional gradient-based bi-level optimization algorithms are ill-suited to meet the demands of large-scale applications. We introduce $(textFG)2textU$, which achieves an unbiased approximation of the meta gradient for bi-level optimization. $(textFG)2textU$ is inherently designed to support parallel computing, enabling it to effectively leverage large-scale distributed computing systems.
  arXiv  Detail & Related papers  (2024-06-20T08:21:52Z)
• Bypass Back-propagation: Optimization-based Structural Pruning for Large Language Models via Policy Gradient [57.9629676017527]
  We propose an optimization-based structural pruning on Large-Language Models. We learn the pruning masks in a probabilistic space directly by optimizing the loss of the pruned model. Our method operates for 2.7 hours with around 35GB memory for the 13B models on a single A100 GPU.
  arXiv  Detail & Related papers  (2024-06-15T09:31:03Z)
• Revisiting Zeroth-Order Optimization for Memory-Efficient LLM Fine-Tuning: A Benchmark [166.40879020706151]
  This paper proposes a shift towards BP-free, zeroth-order (ZO) optimization as a solution for reducing memory costs during fine-tuning. Unlike traditional ZO-SGD methods, our work expands the exploration to a wider array of ZO optimization techniques. Our study unveils previously overlooked optimization principles, highlighting the importance of task alignment, the role of the forward gradient method, and the balance between algorithm complexity and fine-tuning performance.
  arXiv  Detail & Related papers  (2024-02-18T14:08:48Z)
• Contextual Stochastic Bilevel Optimization [50.36775806399861]
  We introduce contextual bilevel optimization (CSBO) -- a bilevel optimization framework with the lower-level problem minimizing an expectation on some contextual information and the upper-level variable. It is important for applications such as meta-learning, personalized learning, end-to-end learning, and Wasserstein distributionally robustly optimization with side information (WDRO-SI)
  arXiv  Detail & Related papers  (2023-10-27T23:24:37Z)
• BiSLS/SPS: Auto-tune Step Sizes for Stable Bi-level Optimization [33.082961718280245]
  Existing algorithms involve two coupled learning rates that can be affected by approximation errors when computing hypergradients. We investigate the use of adaptive step-size methods, namely line search (SLS) and Polyak step size (SPS), for computing both the upper and lower-level learning rates. New algorithms, which are available in both SGD and Adam versions, can find large learning rates with minimal tuning and converge faster than corresponding vanilla BO algorithms.
  arXiv  Detail & Related papers  (2023-05-30T00:37:50Z)
• Fast Adaptive Federated Bilevel Optimization [14.579475552088692]
  We propose a novel adaptive federated bilevel optimization algorithm (i.e.,AdaFBiO) to solve the distributed bilevel optimization problems. AdaFBiO uses the unified adaptive matrices to flexibly incorporate various adaptive learning rates to update variables in both UL and LL problems. We provide a convergence analysis framework for our AdaFBiO algorithm, and prove it needs the sample of complexity of $tildeO(epsilon-3)$ with communication complexity of $tildeO(epsilon-2)$ to obtain an $
  arXiv  Detail & Related papers  (2022-11-02T13:55:47Z)
• A Generic Descent Aggregation Framework for Gradient-based Bi-level Optimization [41.894281911990554]
  We develop a novel Bi-level Descent Aggregation (BDA) framework for bi-level learning tasks. BDA aggregates hierarchical objectives of both upper level and lower level. We propose a new proof recipe to improve the convergence results of conventional gradient-based bi-level methods.
  arXiv  Detail & Related papers  (2021-02-16T06:58:12Z)
• A Generic First-Order Algorithmic Framework for Bi-Level Programming Beyond Lower-Level Singleton [49.23948907229656]
  Bi-level Descent Aggregation is a flexible and modularized algorithmic framework for generic bi-level optimization. We derive a new methodology to prove the convergence of BDA without the LLS condition. Our investigations also demonstrate that BDA is indeed compatible to a verify of particular first-order computation modules.
  arXiv  Detail & Related papers  (2020-06-07T05:18:50Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.