기술 요약: 네이티브 저차원(Low-Rank) LLM 사전 학습의 안정화

1. 문제 정의

파운데이션 모델은 상당한 성공을 거두었으나, 파라미터 수의 증가로 인해 상당한 계산 및 메모리 병목 현상이 발생하고 있습니다. 저차원 인수분해(가중치 행렬을 $W = AB^\top$로 표현)는 미세 조정(예: LoRA)에는 효과적임이 입증되었지만, 이를 처음부터(from scratch) 사전 학습에 적용하는 것은 여전히 난제로 남아 있습니다.

보조적인 풀-랭크(full-rank) 가중치 없이 네이티브 저차원 체제에서 모델을 학습시키려는 기존의 시도들은 심각한 불안정성을 겪습니다. 본 연구에서 식별한 핵심 문제는 가중치 업데이트의 스펙트럼 노름(spectral norm, 최대 특잇값)의 통제되지 않는 성장입니다. 밀집(dense) 학습에서는 AdamW와 같은 옵티마이저가 스펙트럼 노름을 유계(bounded)로 유지하는 것과 달리, 저차원 인수분해는 무한한 등가 표현($W = (\lambda A)(\frac{1}{\lambda} B)$)을 허용합니다. 이러한 스케일링 불변성(scaling invariance)은 스케일링 인자 $\lambda$가 무한히 커질 수 있게 하여, 활성화 값의 폭주와 학습 발산을 초래합니다. 기존 방법들(예: Wei et al., 2024a; Huh et al., 2024)은 보조 풀-랭크 가중치를 유지하거나 하이브리드 아키텍처를 사용하여 이를 우회하지만, 이는 상당한 메모리 오버헤드를 발생시키거나 진정한 엔드 투 엔드(end-to-end) 인수분해 학습을 달야하지 못합니다.

2. 방법론: Spectron

저자들은 보조적인 풀-랭크 구성 요소 없이도 안정적이고 엔드 투 엔드 방식의 네이티브 저차원 사전 학습을 가능하게 하는 방법인 Spectron을 제안합니다. Spectron은 **스펙트럼 재정규화(Spectral Renormalization)**와 **그래디언트 직교화(Gradient Orthogonalization)**라는 두 가지 결합된 메커니즘을 통해 스펙트럼 불안정성을 해결합니다.

핵심 메커니즘

이 방법은 합성 가중치 업데이트 $\Delta W$의 스펙트럼 노름을 규정된 반경 $\eta$(학습률에서 유도됨) 내로 제한하는 것을 목표로 합니다.

1. 그래디언트 직교화: Newton-Schulz 반복법을 사용하여 인자 $A$와 $B$에 대한 모멘텀 업데이트를 직교화합니다. 이를 통해 업데이트가 국소적 제약 반경 내에 머물도록 보장하여, 임의적인 방향성 성장을 방지합니다.
2. 적응형 스펙트럼 재정규화: 전역 제약 조건 $\|\Delta W\|_2 \leq \eta$를 만족하기 위해, 현재 $A$와 $B$의 스펙트럼 노름을 기반으로 인자의 학습률을 동적으로 조절합니다.
   인자의 업데이트 규칙은 다음과 같습니다:
   $$\Delta A \leftarrow \frac{\eta}{\|A\|_2 + \|B\|_2 + 1} \cdot \text{Ortho}(M_A)$$
   $$\Delta B \leftarrow \frac{\eta}{\|A\|_2 + \|B\|_2 + 1} \cdot \text{Ortho}(M_B)$$
   여기서 $\|A\|_2$와 $\|B\|_2$는 단일 파워 이터레이션(power iteration) 단계를 통해 효율적으로 추정됩니다.

구현 효율성

• 직교화: 5회의 Newton-Schulz 반복을 통해 구현되었습니다.
• 스펙트럼 추정: 1회의 파워 이터레이션 단계를 통해 구현되었습니다.
• 오버헤드: 총 계산 오버헤드는 FLOPs의 1% 미만이며, 이는 약 25%의 추가 FLOPs를 요구하는 "셀프 가이드(self-guided)" 학습 방법보다 크게 개선된 수치입니다.

3. 주요 기여

1. Spectron 알고리즘: 직교화와 스펙트럼 재정규화를 결합한 적응형 저차원 인자 업데이트 방법입니다. 이는 결과적인 가중치 업데이트의 스펙트럼 노름을 제한하여, 보조 밀집 구성 요소 없이도 무작위 초기화로부터 안정적인 학습을 가능하게 합니다.
2. 실증적 검증: Spectron으로 학습된 인수분해 모델이 FineWeb 데이터셋 상의 세 가지 모델 규모(94M, 297M, 454M 파라미터)에서 밀집 베이스라인과 대등하거나 이를 능가하는 성능을 달성함을 입증했습니다.
3. 연산 최적화 스케일링 법칙(Compute-Optimal Scaling Laws): IsoFLOP 분석을 통해 저자들은 네이티브 저차원 트랜스포머에 대한 특정 스케일링 법칙을 도출했습니다:
   • 최적 모델 크기: $N_{opt} \propto C^{0.479}$
   • 최적 학습 토큰 수: $D_{opt} \propto C^{0.521}$
     이 지수들은 Chinchilla 법칙($N \propto C^{0.49}, D \propto C^{0.51}$)과 약간 다르며, 이는 연산 최적화된 저차원 모델이 비례적으로 더 큰 데이터셋으로 학습된 더 작은 아키텍처를 선호함을 나타냅니다.

4. 결과

• 안정성: Spectron은 나이브한 저차원 학습(AdamW) 및 셀프 가이드 학습에서 관찰되는 손실 스파이크(loss spikes)와 발산을 제거합니다. 학습 전반에 걸쳐 스펙트럼 노름을 유계로 유지합니다.
• 성능:
  • 454M 파라미터 모델에서 Spectron은 12.11의 검증 퍼플렉시티(perplexity)를 달성하여, 셀프 가이드 베이스라인(13.70)과 나이브 AdamW(14.57)를 능가했습니다.
  • 동일한 FLOPs로 학습했을 때, 454M 인수분해 모델은 780M 밀집 모델의 검증 손실과 일치하였으며, 이는 약 42% 더 많은 파라미터를 가진 모델의 성능을 효과적으로 회복했음을 의미합니다.
  • 다운스트림 태스크(HellaSwag, PIQA, ARC-easy)에서도 동일한 파라미터 수를 가진 베이스라인 및 밀집 모델 대비 일관된 개선을 보였습니다.
• 추론 효율성: 동일한 퍼플렉시티를 달로하기 위해 필요한 파라미터 수가 줄어듦에 따라, Spectron 모델은 추론 비용을 대폭 절감할 수 있습니다(현대적 연산 예산 기준, Chinchilla 최적 밀집 트랜스포머 대비 최대 50% 절감 예상).

5. 의의 및 주장

본 논문은 적절한 스펙트럼 제약이 적용될 때 네이티브 저차원 학습이 밀집 사전 학습의 실행 가능하고 안정적인 대안이 될 수 있다고 주장합니다.

• 과학적 통찰: 본 연구는 최적화 과정에서 풀-랭크 표현이 반드시 필요한 것은 아니라는 점을 밝혀냈습니다. 이전에 관찰된 불안정성은 저차원 구조 자체의 내재적 한계가 아니라 통제되지 않은 스펙트럼 성장 때문이었습니다.
• 실질적 영향: 보조 풀-랭크 가중치의 필요성을 제거함으로써, Spectron은 파운데이션 모델 개발을 민주화하고 잠재적으로 메모리 소비와 하드웨어 요구 사항을 줄일 수 있습니다.
• 스케일링 동작: 도출된 스케일링 법칙은 고정된 연산 예산 내에서 저차원 모델을 위한 최적의 전략이 더 작은 모델을 더 많은 데이터로 학습시키는 것이며, 이를 통해 본질적으로 추론 효율적인 모델을 얻을 수 있음을 시사합니다.

인정된 한계점:

• 실험은 1.5B 파라미터까지의 모델을 대상으로 수행되었으며, 이는 조 단위 토큰을 다루는 프런티어 급 규모보다는 작습니다.
• 현재 구현은 커스로 퓨즈드 커널(custom fused kernel)을 사용하지 않아 밀집 학습에 비해 실제 실행 시간(wall-clock time) 격차가 존재하지만, 저자들은 이를 해결할 수 있다고 언급했습니다.
• 본 방법론은 LLM에 초점을 맞추고 있으나, 그 원리는 멀티모달 아키텍처로 확장될 수 있습니다.