LLM 학습 동역학 분석을 위한 확장 가능한 손실 함수 경사도 곡률 측정 방법

A Scalable Measure of Loss Landscape Curvature for Analyzing the Training Dynamics of LLMs

신경망의 학습 동역학을 분석하는 데 있어 손실 함수 경사도 곡률의 변화를 이해하는 것은 매우 중요합니다. 가장 일반적으로 연구되는 측정 지표는 헤세 행렬의 최대 고유값인 '헤세 선명도'($λ_{ ext{max}}^H$)이며, 이는 로컬 학습 안정성을 결정하고 학습 과정 전반에 걸쳐 학습률과 상호 작용합니다. 헤세 선명도는 학습 동역학 분석에 중요한 역할을 하지만, 대규모 언어 모델(LLM)의 경우 높은 계산 비용으로 인해 직접 측정하기 어렵습니다. 본 연구에서는 업데이트 방향 $Δ extbf{θ}$이 주어졌을 때 10회 미만의 순전파 연산으로 계산 가능한, 계산 효율적인 측정 지표인 '임계 선명도'($λ_c$)를 분석합니다. 특히, 이 측정 지표는 점진적인 선명화 및 안정 영역과 같은 잘 알려진 헤세 선명도 현상을 잘 반영합니다. 본 연구에서는 이 측정 지표를 사용하여 최대 70억 개의 파라미터를 가진 OLMo-2 모델의 사전 학습 및 중간 학습 단계에서 이러한 선명도 현상을 대규모로 최초로 입증했습니다. 또한, 사전 학습에서 미세 조정으로의 전환을 분석하고 데이터 혼합 전략을 안내하기 위해, 하나의 손실 함수 경사도 곡률을 최적화하면서 다른 손실 함수 경사도 곡률을 측정하는 '상대적 임계 선명도'($λ_c^{1 o 2}$)를 도입했습니다. 임계 선명도는 실무자들에게 곡률 동역학을 진단하고 대규모 데이터 구성 선택에 대한 정보를 제공하는 실용적인 도구입니다. 더 넓은 의미에서, 본 연구는 확장 가능한 곡률 측정 방법이 대규모 학습에 대한 실행 가능한 통찰력을 제공할 수 있음을 보여줍니다.

Understanding the curvature evolution of the loss landscape is fundamental to analyzing the training dynamics of neural networks. The most commonly studied measure, Hessian sharpness ($λ_{\max}^H$) -- the largest eigenvalue of the loss Hessian -- determines local training stability and interacts with the learning rate throughout training. Despite its significance in analyzing training dynamics, direct measurement of Hessian sharpness remains prohibitive for Large Language Models (LLMs) due to high computational cost. We analyze $\textit{critical sharpness}$ ($λ_c$), a computationally efficient measure requiring fewer than $10$ forward passes given the update direction $Δ\mathbfθ$. Critically, this measure captures well-documented Hessian sharpness phenomena, including progressive sharpening and Edge of Stability. Using this measure, we provide the first demonstration of these sharpness phenomena at scale, up to $7$B parameters, spanning both pre-training and mid-training of OLMo-2 models. We further introduce $\textit{relative critical sharpness}$ ($λ_c^{1\to 2}$), which quantifies the curvature of one loss landscape while optimizing another, to analyze the transition from pre-training to fine-tuning and guide data mixing strategies. Critical sharpness provides practitioners with a practical tool for diagnosing curvature dynamics and informing data composition choices at scale. More broadly, our work shows that scalable curvature measures can provide actionable insights for large-scale training.