Beyond Perplexity: A Geometric and Spectral Study of Low-Rank Pre-Training

기술 요약: 퍼플렉시티를 넘어: 저랭크 사전 학습의 기하학적 및 스펙트럼 연구

문제 제기
대규모 언어 모델 (LLM) 의 사전 학습은 풀랭크 가중치, 기울기, 옵티마이저 상태를 저장하는 데 따른 메모리 비용에 의해 제약받습니다. 이러한 비용을 줄이기 위한 해결책으로 저랭크 사전 학습이 등장했으며, 그 방법들은 옵티마이저 기반 (예: GaLore, Fira), 아키텍처 기반 (예: CoLA, SLTrain), 어댑터 기반 (예: ReLoRA) 의 세 가지 범주로 진화했습니다. 그러나 이러한 방법들에 대한 평가는 여전히 제한적입니다. 현재의 관행은 거의 전적으로 단일 시드 실행에서 얻은 검증 퍼플렉시티에 의존합니다. 저자들은 퍼플렉시티가 솔루션의 질을 나타내는 나쁜 대리 지표라고 주장합니다. 동일한 퍼플렉시티를 가진 모델들도 손실 지형의 근본적으로 다른 영역으로 수렴하여 고유한 내부 표현을 가질 수 있기 때문입니다. 또한, 기존 비교들은 해당 방법으로 사전 학습된 모델이 얼마나 잘 미세 조정되는지와 해당 방법으로 사전 학습된 모델이 하류 작업에서 얼마나 잘 일반화되는지 사이의 차이를 구분하지 못하는 경우가 많습니다.

방법론
저자들은 세 가지 모델 규모 (60M, 130M, 350M) 에 걸쳐 풀랭크 학습 대비 다섯 가지 저랭크 사전 학습 방법 (GaLore, Fira, CoLA, SLTrain, ReLoRA) 이 발견한 솔루션을 특성화하기 위한 진단 프레임워크를 도입합니다. 본 연구는 네 가지 차원에 걸쳐 16 가지 지표를 사용하여 이러한 방법들을 평가합니다:

1. 1 차원 손실 지형 (무작위 방향): 저자들은 모델 파라미터를 100 개의 무작위 방향으로 교란시켜 검증 손실의 평균과 분산을 측정합니다. 이를 통해 날카로움 (Sharpness) (체크포인트에서 멀어질 때 손실 증가율) 과 방향 분산 (Direction Variance) (손실 표면의 이방성) 을 나타내는 스칼라 지표를 유도합니다.
2. 1 차원 손실 지형 (Top-K PCA 방향): 가중치 행렬의 상위 특이 방향을 따라 교란을 적용하여 가장 중요한 축을 따른 기하학을 분석합니다.
3. 1 차원 보간: 저자들은 체크포인트 간의 선형 경로를 따라 손실을 측정합니다. 두 가지 구체적인 지표를 도입합니다:
   • 연속 체크포인트 장벽 높이 (CCBH): 동일한 방법의 순차적 체크포인트 간의 장벽을 측정합니다.
   • 방법 간 장벽 높이 (IMBH): 동일한 학습 단계에서 서로 다른 학습 방법 간의 기하학적 분리를 정량화하는 새로운 지표입니다. 이는 서로 다른 방법들이 동일한 분지로 수렴하는지 아니면 기하학적으로 구별된 솔루션으로 수렴하는지를 보여줍니다.
4. 스펙트럼 및 랭크 지표: 연구는 가중치와 업데이트의 스펙트럼 구조를 추적하며, 유효 랭크 (Effective Rank), 안정 랭크 (Stable Rank), 스펙트럼 갭 (Spectral Gap), 임계값 랭크 (Threshold Rank) 를 포함합니다.
5. 활성화 유사성: 저자들은 활성화 L2 거리, 코사인 유사도, 선형 CKA (Centered Kernel Alignment) 를 사용하여 저랭크 모델의 은닉 상태를 풀랭크 기준과 비교합니다.

마지막으로, 저자들은 이러한 기하학적 및 스펙트럼 특징들이 퍼플렉시티 단독보다 하류 제로샷 성능을 더 잘 예측할 수 있는지, LOSO(Leave-One-Size-Out) 및 LOMO(Leave-One-Method-Out) 교차 검증을 사용한 선형 회귀를 통해 테스트합니다.

주요 기여

1. 진단 프레임워크: 본 논문은 시드나 에포크 간뿐만 아니라 학습 알고리즘 간의 기하학적 분리를 정량화하는 새로운 IMBH 지표를 포함한 포괄적인 평가 세트를 제안합니다.
2. 크로스 패밀리 실증적 발견: 본 연구는 저랭크 방법들이 풀랭크 학습이나 서로 다른 방법들과 동등하지 않음을 확립합니다. 주요 발견 사항은 다음과 같습니다:
   • 기하학적 분리: 방법들은 기하학적으로 구별된 분지로 수렴하며, IMBH 는 학습이 진행됨에 따라 단조 증가합니다.
   • 스펙트럼 서명: GaLore 는 풀랭크 업데이트의 스펙트럼 구조와 밀접하게 일치하는 반면, CoLA 는 그렇지 않습니다.
   • 표현 발산: 저랭크 활성화는 후기 레이어에서 풀랭크와 점점 더 발산합니다. CoLA 는 직교하는 활성화를 보이는 반면, Fira 는 최종 레이어에서만 정렬을 회복합니다.
3. 개선된 하류 예측: 검증 퍼플렉시티만으로는 하류 성능의 불완전한 예측자입니다. 퍼플렉시티에 기하학적 및 스펙트럼 특징을 추가하면 예측 정확도가 크게 향상됩니다 (LOSO 피어슨 상관계수가 0.873 에서 0.913 으로 증가).
4. 통합 코드베이스: 저자들은 추가 연구를 촉진하기 위해 사전 학습, 지표 계산, 하류 평가를 통합한 오픈 소스 파이프라인을 공개합니다.

결과

• 손실 지형 기하학: 풀랭크 학습은 무작위 방향에서는 더 날카로운 분지에 정착하지만, Top-1 PCA 방향에서는 더 평평한 분지에 정착합니다. 반면, Fira 와 SLTrain 과 같은 저랭크 방법들은 무작위 방향에서 더 평평한 분지로 수렴하는 반면, CoLA 와 ReLoRA 는 가장 날카로운 분지에 도달합니다. GaLore 는 혼합된 프로필을 보입니다: 무작위 방향에서는 중간 정도 평평하지만 Top-1 PCA 방향에서는 날카롭습니다.
• 방법 간 장벽: 모든 방법은 기하학적으로 분리된 분지를 차지합니다. 60M 규모에서 풀랭크와 임의의 저랭크 방법 간의 장벽은 저랭크 방법들 간의 장벽보다 높습니다. Fira 는 풀랭크에 기하학적으로 가장 가까운 분지로 수렴하는 반면, SLTrain 은 가장 고립되어 있습니다.
• 스펙트럼 분석: 풀랭크는 모든 랭크 지표에서 우세합니다. GaLore 는 유효 랭크와 안정 랭크에서 풀랭크를 밀접하게 추적하지만 더 작은 스펙트럼 갭을 보입니다. Fira 와 GaLore 가 풀랭크 가중치 업데이트의 스펙트럼 분포를 가장 충실히 재현합니다.
• 활성화 정렬: ReLoRA 는 모든 규모에서 풀랭크 표현을 가장 밀접하게 추적합니다. CoLA 는 가장 큰 편차를 보이며, 후기 레이어에서 활성화가 기준과 직교하게 됩니다. Fira 는 높은 레이어 평균 드리프트를 보이지만 놀랍게도 최종 레이어에서는 편차가 낮습니다.
• 하류 성능: 검증 퍼플렉시티는 하류 순위 일관적으로 예측하지 못합니다. 예를 들어, Fira 는 모든 규모에서 가장 낮은 퍼플렉시티를 달성하지만, CoLA 는 퍼플렉시티 상위 두 위를 차지하지 않았음에도 130M 및 350M 에서 여러 하류 카테고리에서 선두를 차지합니다.

의의
본 논문은 저랭크 사전 학습 방법들이 풀랭크 학습의 교체 가능한 근사가 아니라, 도달하는 솔루션을 근본적으로 변화시킨다고 결론지었습니다. 연구는 검증 퍼플렉시티가 모델 품질에 대한 불완전한 그림을 제공함을 보여줍니다. 기하학적 및 스펙트럼 지표를 통합함으로써 연구자들은 하류 성능을 더 잘 예측하고 서로 다른 저랭크 전략이 도입하는 고유한 귀납적 편향을 이해할 수 있습니다. 이 작업은 일부 방법 (예: Fira) 은 풀랭크에 기하학적으로 가까운 분지로 수렴하는 반면, 다른 방법들 (예: CoLA) 은 구별된 골짜기를 차지한다는 점을 강조하며, 저랭크 방법의 선택이 모델이 거주하는 손실 지형의 특정 영역을 결정함을 시사합니다.