BLADE: Scalable Bi-level Adaptive Data Selection for LLM Training

2026-06-16

AI Summary

This work addresses the inefficiency and computational infeasibility of static data selection in training trillion-token-scale large language models by introducing BLADE, a novel framework that reformulates bilevel optimization into a single-level objective with a penalty term via Lagrange multipliers. This transformation circumvents the explicit computation of the Hessian inverse, enabling efficient adaptive data selection while preserving theoretical rigor. BLADE dynamically synchronizes a reference model with the training model and establishes a formal connection to excess-loss-based methods. Coupled with a memory-efficient stochastic block-coordinate Frank-Wolfe algorithm, the approach consistently outperforms existing data selection strategies across multiple experiments, offering a scalable and practical solution for data curation in large-scale language model training.

Abstract

As Large Language Model (LLM) datasets scale to trillions of tokens, data selection has emerged as a critical frontier to filter out uninformative noise and construct adaptive learning trajectories. Beyond static heuristic filtering, advanced data selection methods for LLM training largely follow two paradigms, each with fundamental limitations. Influence-based methods provide principled bi-level objectives but require intractable inverse-Hessian computations, while excess-loss methods are computationally efficient but rely on a static reference model that becomes misaligned with the evolving proxy model during training. We propose BLADE (Bi-Level Adaptive Data sElection), a Hessian-free framework for data selection. BLADE reformulates the bi-level optimization problem underlying influence-based methods as a penalized single-level objective via Lagrange multipliers, avoiding inverse-Hessian computation while revealing a principled connection to excess-loss based data selection. The resulting objective recovers an excess-loss form but replaces the static reference model with a dynamic one that stays synchronized with training. Theoretically, we prove that this penalized formulation guarantees first-order convergence. For efficient online batch selection, we instantiate BLADE as a memoryless randomized block-coordinate Frank-Wolfe algorithm. Extensive experiments show that BLADE consistently outperforms state-of-the-art data selection baselines, providing a practical recipe for LLM training.

Problem

• data selection
• large language models
• bi-level optimization
• adaptive learning
• scalable training

Innovation

• bi-level optimization
• data selection
• Hessian-free
• adaptive training
• excess loss