Towards Efficient LLMs Annealing with Principled Sample Selection

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.

The final "annealing" phase of large language model (LLM) training critically determines model quality, yet selecting optimal data for this phase remains an open challenge — current methods rely on heuristic rules without theoretical grounding. We characterize this phase through the geometry of the loss landscape, showing that optimal convergence requires gradient updates to satisfy heterogeneous constraints across different parameter directions. We propose DiReCT (Directionally-Restrained Constrained Training), which reformulates sample selection as a constrained optimization problem. By imposing directional constraints on per-sample gradients based on the Hessian's spectral properties, DiReCT selects samples aligned with the optimal curvature-aware descent path. Experiments across multiple model scales show DiReCT achieves state-of-the-art performance. Code: https://github.com/xuyj233/Direct.