원칙적인 샘플 선택을 통한 효율적인 LLM 어닐링(Annealing)을 향하여

어닐링 (Annealing) 단계는 LLM 사전 학습 (Pre-training)에서 최종 모델의 품질을 궁극적으로 결정하는 핵심적인 수렴 단계입니다. 그러나 이 단계 동안 학습 데이터를 효과적으로 선택하는 것은 여전히 주요한 과제로 남아 있습니다. 현재의 전략들은 도메인 필터링 (Domain filtering)이나 컨텍스트 확장 (Context extension)과 같은 경험적인 휴리스틱 (Heuristics)에 의존하고 있으며, 이는 최적화 이론 (Optimization theory)에 기반한 원칙적인 근거가 부족합니다. 본 연구에서는 손실 경관 (Loss landscape)의 스펙트럼 기하학 (Spectral geometry) 관점을 통해 어닐링 단계를 규명합니다. 우리는 최적의 수렴을 위해 그래디언트 업데이트 (Gradient updates)가 서로 다른 고유 방향 (Eigen-directions)에 걸쳐 이질적인 제약 조건 (Heterogeneous constraints)을 충족해야 한다고 주장합니다. 이러한 통찰을 바탕으로, 우리는 데이터 선택을 이러한 방향성 제약 조건을 충족하는 문제로 공식화합니다. 이를 위해, 우리는 어닐링 단계의 샘플 선택을 제약 최적화 문제 (Constrained optimization problem)로 재정의하는 새로운 프레임워크인 DiReCT (Directionally-Restrained Constrained Training)를 제안합니다. 헤시안 (Hessian)의 스펙트럼 특성에 기반하여 샘플별 그래디언트에 명시적인 방향성 제약을 부과함으로써, DiReCT는 최적의 곡률 인식 하강 경로 (Curvature-aware descent path)와 일치하는 샘플을 식별합니다. 다양한 모델 규모에 걸친 광범위한 실험을 통해 DiReCT가 일관되게 최첨단 (State-of-the-art) 성능을 달성함을 입증했습니다. 향후 연구를 위해 코드는 https://github.com/xuyj233/Direct 에서 확인할 수 있습니다.