提出新方法提升大模型预训练收敛效率
Towards Efficient LLMs Annealing with Principled Sample Selection
- 基于损失曲面谱几何设计定向约束采样策略
- 在多规模模型上实现当前最佳训练效果
- 适合关注高效预训练与优化理论的研究者
大模型预训练中的退火阶段是决定最终模型质量的关键收敛期，但该阶段的有效数据选择仍面临挑战。现有方法依赖领域过滤或上下文扩展等经验性启发，缺乏优化理论支撑。本文从损失曲面谱几何角度分析退火阶段，指出最优收敛需在不同特征方向上满足异质梯度约束。基于此，提出DiReCT（Directionally-Restrained Constrained Training）框架，将退火阶段的数据选择建模为带方向约束的优化问题。通过利用海森矩阵的谱特性，对每样本梯度施加显式方向约束，识别出与最优曲率感知下降路径一致的样本。在多种模型规模上的实验表明，DiReCT持续取得领先性能。代码已开源：https://github.com/xuyj233/Direct。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。