Perplexity를 넘어: 저차원 사전 학습(Low-Rank Pre-Training)에 대한 기하학적 및 스펙트럼 연구

대규모 언어 모델(Large Language Models)의 사전 학습(Pre-training)은 전체 랭크(Full-rank) 가중치, 그래디언트(Gradients), 그리고 옵티마이저 상태(Optimizer states)를 저장하는 메모리 비용에 의해 지배됩니다. 이를 해결하기 위해 저차원 사전 학습(Low-rank pre-training)이 등장했으며, 방법론의 범위가 급격히 확장되었습니다. 핵심적인 질문은 여전히 해결되지 않은 상태로 남아 있습니다: 저차원 방법론들이 전체 랭크 학습과 비교할 만한 일반화 성능을 가진 모델을 생성하는가, 아니면 랭크 제약(Rank constraint)이 도달하는 솔루션을 근본적으로 변화시키는가? 기존의 비교 연구들은 거의 전적으로 이전 문헌에서 이어져 온 단일 시드(Single-seed) 실행의 검증 퍼플렉시티(Validation perplexity)에 의존하고 있습니다. 그러나 퍼플렉시티(Perplexity)는 솔루션 품질을 나타내는 불충분한 대리 지표(Proxy)입니다. 두 방법론이 퍼플렉시티 측면에서는 일치하더라도, 서로 다른 손실 경관(Loss landscape) 영역과 내부 표현(Internal representations)으로 수렴할 수 있기 때문입니다.

본 연구에서는 세 가지 모델 규모(60M, 130M, 350M)에서 전체 랭크 학습과 대비하여 다섯 가지 저차원 사전 학습 방법론인 GaLore 및 Fira(메모리 효율적 옵티마이저), CoLA 및 SLTrain(아키텍처 재매개변수화), 그리고 ReLoRA(주기적 리셋을 포함한 어댑터 스타일 업데이트)가 찾아내는 솔루션의 특성을 규명함으로써 이 간극을 메웁니다. 우리는 네 가지 차원에 걸친 16가지 지표를 통해 각 방법론을 평가합니다: 무작위/Top-K PCA 방향을 따른 1차원 손실 경관(1-D loss landscape), 체크포인트 간의 1차원 보간(1-D interpolation), 가중치 및 학습된 업데이트의 스펙트럼 구조(Spectral structure), 그리고 전체 랭크 학습과의 활성화 유사성(Activation similarity)입니다.

연구 결과, 저차원 방법론들은 검증 퍼플렉시티(Validation perplexity)가 유사할 때조차 전체 랭크 학습과 동일하지 않으며, 서로 간에도 동일하지 않음을 보여줍니다. 전체 랭크 학습은 무작위 방향을 따라 저차원 방법론보다 더 날카로운 분지(Sharper basin)로 안착하는 반면, Top-1 PCA 방향에서는 그 반대의 현상이 나타납니다. 각 방법론은 기하학적으로 구별되는 분지로 수렴합니다. 저차원 활성화(Low-rank activations)는 학습이 진행됨에 따라 후기 레이어에서 전체 랭크 학습과 멀어지며, GaLore가 전체 랭크를 가장 밀접하게 추적합니다. 나아가, 검증 퍼플렉시티(Validation perplexity)가 모든 규모에서 다운스트림 성능(Downstream performance)으로 직결되지는 않습니다. 기하학적 및 스펙트럼 지표를 추가함으로써 예측 성능을 향상시킬 수 있습니다.