기술적 요약: 옵티마이저 유도 스펙트럴 스케일링 법칙

문제 정의

대형 언어 모델 (LLM) 의 고전적 스케일링 법칙은 검증 손실과 모델 크기, 데이터 양, 연산량과 같은 자원 간의 예측 가능한 멱함수 관계를 확립해 왔습니다. 그러나 이러한 프레임워크들은 일반적으로 옵티마이저를 고정된 훈련 세부 사항으로 간주하며, 이는 수렴 속도에만 영향을 미치거나 곱셈적 효율성 인자로만 작용한다고 가정합니다. 본 논문은 이러한 가정에 도전하여, 옵티마이저가 표현 스케일링의 근본적인 축임을 주장합니다. 구체적으로 저자들은 서로 다른 옵티마이저가 증가된 피드포워드 네트워크 (FFN) 폭을 얼마나 효과적으로 활용된 스펙트럴 용량(잠재 공간의 유효 차원성) 으로 전환하는지 조사합니다. 동일한 검증 손실이 동일한 내부 표현 기하학을 의미하는지, 아니면 서로 다른 옵티마이저가 근본적으로 다른 스펙트럴 구조를 통해 동일한 아키텍처 용량을 실현하는지 질문합니다.

방법론

이 연구는 트랜스포머 아키텍처 내 FFN 표현의 고유 스펙트럼을 중심으로 한 스펙트럴 분석 프레임워크를 사용합니다.

• 스펙트럴 지표: 저자들은 스펙트럴 확산을 정량화하기 위해 Rényi 유효 랭크 계열을 활용합니다. 분석은 두 가지 주요 지표를 기반으로 합니다:
  • 소프트 랭크 ($R_1$): $\alpha=1$(Shannon 엔트로피) 에 해당하며, 엔트로피 가중 스펙트럴 확산을 측정하고 확산된 용량을 포착합니다.
  • 하드 랭크 ($R_2$): $\alpha=2$(참여 비율) 에 해당하며, 집중 민감성 유효 차원성을 측정하고 지배적 모드 용량을 포착합니다.
  • 하드 - 소프트 비대칭 ($\Delta_{1,2}$): $\log R_1 - \log R_2$로 정의되며, 이 지표는 확산된 용량과 지배적 용량 간의 격차를 정량화합니다.
• 실험 설정: 실험은 FineWeb-Edu 데이터셋으로 훈련된 160M 및 350M 파라미터의 GPT 스타일 디코더 전용 트랜스포머에서 수행되었습니다. 스케일링 법칙 ($R(D) \propto D^\beta$) 을 확립하기 위해 FFN 은닉 차원 ($D$) 을 변경했습니다.
• 옵티마이저 비교: 표준 AdamW를 행렬 인식 옵티마이저 (Muon, NorMuon) 및 랭크 제약 변형 (Dion, 다양한 랭크 분수 $r$) 과 비교했습니다.
• 층화: 언어 데이터의 Zipfian 분포를 해결하기 위해 FFN 표현을 토큰 빈도에 따라 HEAD(빈번한), MID, TAIL(희귀한) 영역으로 층화했습니다.
• 아키텍처 제어: 저자들은 옵티마이저 유도 변화를 헤드 수 감소를 통한 헤드당 어텐션 랭크 증가 및 회전 위치 임베딩 (RoPE) 제거와 같은 아키텍처 개입과 비교했습니다.

주요 결과

1. 옵티마이저 의존적 스펙트럴 스케일링 법칙

동일한 트랜스포머 아키텍처는 옵티마이저에 따라 현저히 다른 스펙트럴 스케일링 법칙을 보입니다.

• 하드 랭크 스케일링: AdamW 는 특히 희귀 (TAIL) 토큰의 경우 약한 하드 랭크 스케일링 ($\beta_{hard} \approx 0.29$~$0.44$) 을 보입니다. 반면 Muon 은 동일한 영역에서 거의 선형 스케일링 ($\beta_{hard} \approx 1.02$) 을 달성합니다. 이는 스케일링 지수에서 2.3$\times$ 증가를 의미합니다.
• 하드 - 소프트 비대칭: AdamW 는 큰 양의 비대칭 ($\Delta_{1,2} \approx 0.37$) 을 보여주며, 이는 추가된 FFN 폭이 지배적 표현 방향보다는 확산된 스펙트럴 질량을 확장함을 나타냅니다. Muon 과 NorMuon 은 이 비대칭을 거의 0 으로 줄여 ($\Delta_{1,2} \approx 0.14$), 추가된 폭이 효율적으로 지배적 모드 용량으로 전환됨을 의미합니다.

2. 일치된 손실 $\neq$ 일치된 기하학

본 논문은 검증 손실이 표현 구조를 위한 불충분한 대리 지표임을 입증합니다.

• 연장 훈련: AdamW 를 연장된 단계 (12K 대 6K) 로 훈련하여 Dion(1/16) 변형의 퍼플렉시티와 일치시킬 수 있었습니다. 그러나 스펙트럴 기하학은 여전히 구별되었습니다: AdamW 의 하드 랭크 스케일링은 거의 소멸 ($\beta_{hard} \to 0.03$) 된 반면, Dion 은 강력한 멱함수 스케일링을 유지했습니다.
• 학습률: 학습률을 스윕한 결과, AdamW-Muon 스펙트럴 격차는 하이퍼파라미터 인공물이 아님을 확인했습니다. 스윕 전반에 걸쳐 Muon 의 최저 하드 랭크 지수가 AdamW 의 최고 지수를 일관되게 초과했습니다.

3. 빈도 및 랭크 의존성

• 토큰 빈도: 옵티마이저 효과는 MID 및 TAIL 토큰 영역에서 가장 두드러집니다. AdamW 는 희귀 토큰에 대한 하드 용량 스케일링에 어려움을 겪는 반면, Muon 은 선형적으로 스케일링합니다.
• 업데이트 랭크: Dion 을 사용하여 업데이트 랭크의 효과를 분리했습니다. 랭크 분수 ($r$) 를 $1/2$에서 $1/16$으로 감소시키면 하드 랭크 스케일링이 $\beta \approx 0.88$에서 $\beta \approx 0.40$(AdamW 수준에 근접) 으로 하락하는 반면, 소프트 랭크 스케일링은 더 점진적으로 저하되었습니다. 이는 업데이트 랭크가 폭을 지배적 모드 용량으로 전환하는 조절 노브 역할을 함을 나타냅니다.

4. 아키텍처 개입에 대한 우세

옵티마이저 유도 스펙트럴 변화는 종종 아키텍처 변경의 효과를 초과하거나 재형성합니다:

• 어텐션 랭크: 헤드 수 감소를 통한 헤드당 어텐션 랭크 증가는 30 개 빈도/지표 영역 중 28 개 영역에서 옵티마이저 전환보다 작은 스펙트럴 스케일링 변화를 발생시켰습니다.
• 위치 인코딩: RoPE 제거는 옵티마이저 의존적 재분포를 유도했습니다. Muon 의 경우 NoPE 는 HEAD 하드 랭크 스케일링을 증가시켰으나 MID/TAIL 스케일링을 감소시켰고, NorMuon 의 경우 모든 영역에서 스케일링을 감소시켰습니다. 이는 아키텍처 신호가 옵티마이저 기하학을 통해 표현됨을 보여줍니다.

중요성 및 주장

본 논문은 옵티마이제이션을 모델 크기, 데이터, 연산량과 함께 표현 스케일링의 1 급 축으로 취급해야 한다고 주장합니다.

• 아키텍처 - 옵티마이저 공동 설계: 실현된 스펙트럴 용량은 아키텍처 단독이 아닌 아키텍처 - 옵티마이저 쌍의 속성입니다. 아키텍처 개입 (어텐션 랭크 또는 위치 인코딩 변경과 같은) 은 고정된 옵티마이저 독립적 변화를 유도하지 않으며, 그 효과는 옵티마이저 기하학에 의해 매개됩니다.
• 표현 구조: 옵티마이저 선택은 모델 용량이 고유 모드에 걸쳐 어떻게 구조화되고 토큰 빈도 영역에 걸쳐 어떻게 할당되는지를 결정합니다. 이는 단순히 수렴 속도를 위한 효율성 노브가 아닙니다.
• 한계점: 저자들은 이러한 발견이 통제된 경험적 영역 (최대 350M 파라미터의 GPT 스타일 모델) 내에서 확립되었음을 지적합니다. 350M 결과가 160M 패턴을 재현하지만, 프론티어 규모 (1B 이상) 에서의 추가 검증이 필요합니다. 또한 스펙트럴 랭크는 용량 활용을 정량화하지만 하류 작업 행동을 완전히 특징짓지는 못합니다.

결론적으로, 본 연구는 옵티마이저 - 아키텍처 공동 설계로의 전환을 고무하며, 미래의 스케일링 법칙은 모델 용량이 어떻게 실현되는지의 주요 결정 요인으로서 옵티마이저의 암시적 귀납적 편향을 고려해야 함을 시사합니다.