Towards Efficient LLMs Annealing with Principled Sample Selection

2026-05-29

AI Summary

This work addresses the lack of theoretically grounded, efficient data selection methods during the annealing phase of large language model pretraining, where existing heuristic strategies often fail to ensure optimal convergence. Drawing from the spectral geometry of the loss landscape, the study introduces DiReCT (Directionally Constrained Optimization Framework), which uniquely integrates data selection with the spectral properties of the Hessian matrix. By leveraging spectral decomposition, DiReCT enables curvature-aware gradient alignment and constrained sample selection. Empirical evaluations across multiple model scales demonstrate that this approach significantly enhances training efficiency and final model performance during the annealing stage, achieving state-of-the-art results.

Abstract

The annealing phase is a pivotal convergence stage in LLM pre-training that ultimately determines final model quality. However, effectively selecting training data during this phase remains a key challenge. Current strategies rely on empirical heuristics, such as domain filtering or context extension, which lack a principled grounding in optimization theory. In this work, we characterize the annealing phase through the lens of the loss landscape's spectral geometry. We argue that optimal convergence requires gradient updates to satisfy heterogeneous constraints across different eigen-directions. Building on this insight, we formulate data selection as a problem of satisfying these directional constraints. To this end, we propose DiReCT (Directionally-Restrained Constrained Training), a novel framework that reformulates sample selection in the annealing stage as a constrained optimization problem. By imposing explicit directional constraints on per-sample gradients based on the spectral properties of the Hessian, DiReCT identifies samples that align with the optimal curvature-aware descent path. Extensive experiments across various model scales demonstrate that DiReCT consistently achieves state-of-the-art performance. For future research, code is available at https://github.com/xuyj233/Direct.

Problem

• LLM pre-training
• annealing phase
• data selection
• loss landscape
• gradient constraints

Innovation

• spectral geometry
• constrained optimization
• gradient directionality
• annealing phase
• curvature-aware training