Two-Stage Optimizer-Aware Online Data Selection for Large Language Models

1. Problem Statement

The paper addresses the challenge of online data selection for fine-tuning Large Language Models (LLMs). Unlike offline settings where the full dataset is available for pre-computation, online fine-tuning involves data arriving sequentially (e.g., continual instruction tuning) with strict memory and latency constraints.

Key Challenges Identified:

• Optimizer Ignorance: Existing gradient-based selection methods (e.g., TracIn, LESS) often assume simple Stochastic Gradient Descent (SGD). They fail to account for the non-linear update rules of adaptive optimizers like Adam/AdamW, which are standard in LLM training.
• Static vs. Dynamic Utility: Offline methods treat sample utility as static. In online settings, a sample's utility depends on the current model state and the specific optimizer geometry at that step.
• Redundancy and Interactions: Simple gradient alignment (inner products) ignores interactions between selected samples. In adaptive optimizers, selecting highly correlated samples can lead to diminishing returns or unstable updates.
• Computational Cost: Computing full per-sample gradients for LLMs is infeasible due to memory and time constraints, especially with long-context data.

2. Methodology

The authors propose an Optimizer-Aware Framework that reframes data selection not as static ranking, but as shaping the next training update to match a target task under the specific geometry induced by the optimizer.

A. Theoretical Formulation

• Optimizer-Aware Update Matching: The goal is to maximize the alignment between the weighted training gradient and the target validation gradient, considering the optimizer's transformation function $P_t$.
  • Objective: $\max_w \langle \nabla L_{val}, P_t(\sum w_i \nabla l_i) \rangle$.
• Distance-Based Subset Utility: Instead of simple alignment, the authors formulate the problem as minimizing the distance between the target gradient and the optimizer-induced update:
  • $\min_{w \ge 0} \| \nabla l_{val} - P_t(\nabla l_{tr}) \|_2^2 + \lambda \|w\|_2^2$.
  • Second-Order Interpretation: This distance-based objective approximates a second-order Taylor expansion of the target loss, naturally penalizing redundant gradients (pairwise interactions).
  • Non-Negative Constraint: Weights $w$ are constrained to be non-negative ($w \ge 0$) to prevent "destructive error cancellation," where the optimizer subtracts large opposing vectors to fit the target, which destabilizes training.

B. Efficient Gradient Representation

To make this tractable for LLMs, the paper introduces:

1. LoRA Integration: Using Low-Rank Adaptation to reduce the dimensionality of trainable parameters.
2. Random Projection: Projecting high-dimensional gradients into a lower-dimensional subspace while preserving inner products (Johnson-Lindenstrauss Lemma).
3. Factorized Outer-Product Representation: Exploiting the structure of linear layer gradients ($\nabla l = g a^T$). Instead of storing full gradients, the method stores activations ($a$) and output gradients ($g$).
   • Inner products are computed as $\langle g_j, g_i \rangle \cdot \langle a_j, a_i \rangle$.
   • This reduces memory complexity from $O(d_1 d_2)$ to $O(d_1 + d_2)$ and handles long sequences efficiently.

C. Optimizer-Aware Preconditioning

For adaptive optimizers like Adam, the update rule involves a preconditioning matrix derived from the second moment estimate ($\hat{v}$).

• Linearized Approximation: The authors approximate the Adam update operator as a fixed linear transformation $D_{t-1} \odot g$ using the state from the previous step.
• Asymmetric Preconditioning: To avoid re-evaluating the second moment for every subset, they transfer the preconditioner to the validation side: $\langle \nabla l_{val}, D_{t-1} \odot \nabla l_{tr} \rangle = \langle D_{t-1} \odot \nabla l_{val}, \nabla l_{tr} \rangle$. This aligns candidates with the actual optimization direction.

D. Two-Stage Filter-then-Weight Algorithm

The selection process is decoupled into two stages to ensure stability and efficiency:

1. Stage 1: Greedy Filtering (Candidate Selection):
   • Uses a greedy residual search (similar to Orthogonal Matching Pursuit) to select a diverse subset of candidates.
   • Assumes unit weights during this phase to prioritize geometric diversity and avoid overfitting to noisy gradient directions.
2. Stage 2: Precise Weighting (Coefficient Optimization):
   • Solves a Non-Negative Least Squares (NNLS) problem on the filtered subset to find optimal continuous weights.
   • This stage refines the composite gradient to strictly align with the optimizer-aware target.

3. Key Contributions

1. Optimizer-Aware Framework: The first framework to explicitly model the interaction between data selection and adaptive optimizers (like Adam) in an online setting, moving beyond SGD approximations.
2. Theoretical Insight: Establishes a connection between gradient matching and second-order target utility, providing a principled justification for using distance-based objectives and non-negative constraints.
3. Efficient Implementation: Introduces a factorized gradient representation and random projection strategy that enables online selection for LLMs with long contexts without prohibitive memory costs.
4. Two-Stage Decoupling: Proposes a robust "Filter-then-Weight" algorithm that separates candidate identification from weight optimization, solving the instability issues found in coupled selection-weighting methods (like standard OMP).

4. Experimental Results

The method was evaluated on Llama-3.2-1B and Qwen3-0.6B using MMLU (knowledge) and TyDiQA (multilingual QA) benchmarks.

• Performance: The proposed method consistently outperformed existing baselines (TracIn, LESS, GREATS, GRAD-MATCH) and often surpassed the "Full Data" baseline, despite using only 5% of the training data.
  • Example: On TyDiQA with Llama-3.2-1B, the method achieved 48.86 F1, compared to 47.50 for GREATS and 42.47 for Full Data.
• Data Efficiency: The method showed superior performance, particularly at low data ratios, indicating it effectively identifies informative samples when training signals are scarce.
• Ablation Studies:
  • Optimizer Awareness is Critical: Using reweighting with raw gradients (ignoring optimizer state) performed worse than simple selection. The preconditioning step is essential for the weighting stage to be beneficial.
  • Non-Negative Constraints: Removing the non-negative constraint led to catastrophic failure due to gradient cancellation.
  • Two-Stage vs. Hard Filter: The two-stage approach (Greedy Filter + NNLS Weighting) provided better long-term stability and final performance compared to simple Top-k filtering.

5. Significance

This paper provides a critical bridge between data-centric AI and optimization theory for LLMs.

• Practical Impact: It offers a scalable solution for fine-tuning LLMs in resource-constrained, streaming environments (e.g., continual learning, edge devices) where storing full datasets or checkpoints is impossible.
• Theoretical Advancement: It challenges the assumption that data selection can be treated as a static ranking problem, demonstrating that the "geometry" of the optimizer fundamentally changes which data is useful.
• Future Direction: The work suggests that future data selection algorithms must be co-designed with the specific optimizer used, rather than being optimizer-agnostic.