얼어붙은 부분 공간을 깨다: 중요도 샘플링을 통한 LLM 사전 학습의 저랭크 최적화

초록

대규모 언어 모델(LLM) 사전 학습의 메모리 효율성을 높이기 위한 저랭크 최적화 기법에서, 기존의 ‘지배적 부분공간’ 선택 방식은 학습 중에 부분공간이 고정되어 다양한 가중치 업데이트를 방해하는 문제가 있습니다. 본 연구는 중요도 샘플링을 활용해 부분공간을 선택하는 새로운 방법(SARA)을 제안합니다. 이 방법은 다양한 부분공간을 탐색하도록 유도하여 성능을 향상시키고, 이론적으로 수렴을 보장합니다. 실험을 통해 SARA가 기존 방법들을 크게 능가함을 입증했습니다.

상세 분석

본 논문은 LLM 사전 학습에서 Adam과 같은 옵티마이저의 메모리 사용량을 획기적으로 줄이기 위한 저랭크 최적화(Low-Rank Optimization) 방법론의 핵심 한계를 지적하고, 이를 해결하는 혁신적인 기법을 제안합니다.

기존 방법론(GaLore, Fira 등)의 핵심은 그래디언트를 저랭크 부분공간에 투영(Projection)하여, 풀랭크 옵티마이저 상태를 저장하는 대신 축소된 랭크의 상태만 저장하는 것입니다. 이때 부분공간을 어떻게 선택하느냐가 성능의 관건입니다. 일반적으로 그래디언트 정보를 최대한 보존하기 위해 가장 큰 특이값(Singular Value)에 해당하는 특이벡터들로 구성된 ‘지배적 부분공간(Dominant Subspace)‘을 선택합니다. 직관적으로 최선의 근사치를 제공할 것 같지만, 논문은 이 방식이 가진 치명적인 문제를 실증적으로 발견했습니다.

그 문제는 ‘얼어붙은 부분공간(Frozen Subspace)’ 현상입니다. 사전 학습이 진행됨에 따라 많은 레이어에서 그래디언트의 지배적 부분공간이 더 이상 진화하지 않고 고정됩니다. 결과적으로 가중치 업데이트가 거의 동일한 저랭크 방향으로 반복되어, 누적된 업데이트 자체도 저랭크에 머물게 되는 ‘저랭크 병목(Low-Rank Bottleneck)‘이 발생합니다. 이는 모델의 표현력을 심각하게 제한합니다.

이를 해결하기 위해 제안된 SARA(Importance SAmpling for Low-RAnk Optimization) 방법은 부분공간 선택에 유의미한 무작위성을 도입합니다. 구체적으로, 그래디언트의 특이값 분해(SVD) 후, 각 왼쪽 특이벡터를 해당 특이값에 비례하는 확률로 가중치를 부여한 중요도 샘플링(Weighted Sampling Without Replacement)을 수행하여 r개의 기저 벡터를 선택합니다. 이는 가장 큰 특이값만을 고집하는 것이 아니라, 상대적으로 작은 특이값을 가진 방향도 적절한 확률로 탐색할 기회를 부여합니다.

이 접근법의 핵심 통찰은 두 가지입니다. 첫째, 인접한 학습 단계 간의 부분공간 중첩(Overlap)을 줄여 업데이트 방향의 다양성(Diversity)을 확보합니다. 둘째, 이론적으로 수렴을 보장합니다(Theorem 3.4). 기존의 지배적 부분공간 선택 방식은 수렴 보장이 없으나, SARA는 투영 오차를 제어할 수 있는 샘플링 확률 하한(δ)을 통해 안정적인 수렴 속도를 증명했습니다.

SARA는 GaLore나 Fira 같은 기존 저랭크 옵티마이저 프레임워크에 ‘플러그 앤 플레이’ 방식으로 쉽게 통합될 수 있으며, SVD 계산 후 추가되는 샘플링 오버헤드는 무시할 수 있을 정도로 작습니다. 실험에서는 LLaMA 모델 사전 학습에서 기존 방법 대비 최대 46.05% 성능 격차를 줄이며, 풀랭크 Adam에 근접하는 성능을 달성했습니다. 이는 단순한 부분공간 선택 전략의 변화가 저랭크 최적화의 실용적 성능에 미치는 영향이 매우 큼을 보여주는 결과입니다.