기술적 요약: 옵티마이저 설계를 위한 대칭성 호환 원칙

1. 문제 제기

현대 딥러닝 관행에는 근본적인 기하학적 불균형이 존재합니다. 신경망 아키텍처 (임베딩, LM 헤드, SwiGLU MLP, MoE 라우터) 는 풍부한 고유 대칭성과 공변성 (equivariance) 속성을 지니고 있는 반면, Adam, AdamW, RMSprop 과 같은 지배적인 최적화 방법들은 **좌표별 (coordinate-wise)**로 작동합니다. 이러한 방법들은 행렬 매개변수를 독립적인 좌표의 평탄화된 벡터로 취급하여 행렬 구조를 무시하고 매개변수 공간의 자연스러운 공변성을 깨뜨립니다.

이러한 "기하학 무지 (geometry-blind)" 접근법은 서로 다른 레이어 유형의 고유한 기하학을 구분하지 못합니다. 그 결과, 인기 있는 옵티마이저들은 스펙트럼 구조를 폐기하고, 불일치하는 기하학을 유발하며, 특정 매개변수 블록에 내재된 불변성 및 공변성 구조 (예: 어휘의 치환 불변성 또는 MoE 의 전문가 교환 가능성) 를 존중하지 못합니다. 좌표별 적응적 방법과 스펙트럼 옵티마이저 간의 현재 이론적 이해와 실증적 비교는 제한적이며, 종종 기하학적 행동에 대한 심층 분석 없이 대규모 벤치마크에 의존합니다.

2. 방법론: 대칭성 호환 원칙

이 논문은 옵티마이저 설계를 위한 레이어별 대칭성 호환 원칙을 제안합니다. 핵심 논지는 특정 블록에 작용하는 대칭군 $G$ 하에서 매개변수 블록 $W$에 대한 옵티마이저 업데이트 규칙 $\mathcal{U}$가 **공변적 (equivariant)**이어야 한다는 것입니다.

2.1 이론적 프레임워크

$W \in \mathbb{R}^{m \times n}$를 매개변수 행렬이라고 합시다. 입력 및 출력 공간의 좌표가 가역 행렬 $P$와 $Q$로 변환될 때, 선형 맵은 $\tilde{W} = P W Q^{-1}$로 변환됩니다. 이에 따라 기울기는 반변량/공변량 방식으로 변환됩니다. 옵티마이저 업데이트 맵 $\mathcal{U}$는 다음을 만족해야 합니다:
$$\mathcal{U}(g \cdot D) = g \cdot \mathcal{U}(D)$$
여기서 $D$는 업데이트 방향 (예: 기울기 또는 모멘텀) 이고 $g$는 대칭군의 원소입니다.

이 논문은 서로 다른 레이어 유형의 대칭군에 기반하여 특정 옵티마이저 클래스를 유도합니다:

1. 일반 행렬 레이어 (선형/어텐션):

   • 대칭성: 이직교 (Bi-orthogonal, $O_m \times O_n$).
   • 옵티마이저 클래스: 스펙트럼 옵티마이저. 업데이트는 특이 벡터를 보존하면서 특이값에 작용하는 스펙트럼 연산자여야 합니다. 이는 확률적 스펙트럼 강하 (SSD), Muon, Scion, 극성 기울기 (Polar Gradient) 방법을 통합합니다. 업데이트는 일반적으로 모멘텀 방향의 직교 극성 인자 (orthogonal polar factor) 입니다.

2. 임베딩 및 LM 헤드:

   • 대칭성: 왼쪽 치환 ($P_v$) 및 오른쪽 직교 ($O_d$). 행은 어휘 항목 (이산 치환) 을 인덱싱하고, 열은 숨겨진 특징 (연속 회전) 을 인덱싱합니다.
   • 옵티마이저 클래스: 왼쪽 치환 - 오른쪽 직교 (LPRO) 공변적 업데이트.
     • 오른쪽 스펙트럼: 오른쪽 그람 행렬 ($D^\top D$) 을 통해 전역적으로 작용합니다.
     • 행 노름 (Row-Norm): 행 노름 (어휘 항목) 에 로컬하게 작용합니다.
     • 하이브리드: 행 노름 및 오른쪽 스펙트럼 업데이트의 조합.

3. SwiGLU MLP 프로젝션:

   • 대칭성: 중간 뉴런 치환 ($P_{d_{ff}}$). 게이트 및 업 프로젝션은 중간 뉴런에 해당하는 행을 가지며, 다운 프로젝션은 이에 해당하는 열을 가집니다.
   • 옵티마이저 클래스: 행 인식 (Row-aware) 및 열 인식 (Column-aware) 변형. 게이트/업 프로젝션은 LPRO 호환 업데이트 (행 노름/오른쪽 스펙트럼) 를 사용하고, 다운 프로젝션은 전치된 기하학을 사용합니다.

4. MoE 라우터:

   • 대칭성: 전문가 치환 ($P_e$) 및 공유 로그 - 시프트 불변성 (모든 로그에 상수를 더해도 소프트맥스 분포는 변하지 않음).
   • 옵티마이저 클래스: 라우터 호환 업데이트. 이는 중심화된 방향 $D_c = \Pi^\perp D$ (공유 시프트 제거) 에 작용합니다.
     • 중심화된 왼쪽 스펙트럼: 중심화된 전문가 그람 행렬에 작용합니다.
     • 중심화된 행 노름: 중심화된 전문가 행 노름에 작용합니다.

2.2 실용적 구현

저자들은 이러한 원칙을 실용적인 모멘텀 기반 옵티마이저로 구체화했습니다:

• RightPolarGradM / LeftPolarGradM: 행렬 역제곱근을 위한 뉴턴 - 슈르츠 (Newton-Schulz) 반복을 사용하는 한쪽 스펙트럼 업데이트.
• RowNormM: 행별 정규화 업데이트.
• HybridPolarGradM: 행 노름 및 한쪽 스펙트럼 단계의 조합.
• 중심화 (Centering): MoE 라우터의 경우, 업데이트는 중심화된 모멘텀 $M_{k,c} = \Pi^\perp M_k$에 적용됩니다.

3. 주요 기여

1. 대칭성 호환 원칙: 좌표별 옵티마이저가 행렬 값 매개변수에 대해 기하학적으로 불일치한다는 형식적 논증. 논문은 업데이트 규칙이 매개변수 블록의 대칭군과 일치해야 함을 확립합니다.
2. 스펙트럼 옵티마이저의 통합적 관점: 스펙트럼 옵티마이저 (Muon, SSD 등) 를 이직교 공변적 업데이트 맵의 고유 클래스로 해석하는 대칭성 기반 해석 제공.
3. 새로운 공변적 옵티마이저 클래스: 비표준 대칭성을 가진 레이어에 대한 구체적인 업데이트 규칙 유도:
   • 임베딩 및 LM 헤드를 위한 LPRO-공변적 업데이트.
   • SwiGLU MLP 를 위한 행/열 인식 업데이트.
   • MoE 라우터를 위한 중심화된 전문가 치환 업데이트.
4. 종단 간 사전 학습 증거: 모든 주요 행렬 클래스가 대칭성 일치 업데이트를 받는 전체 레이어별 옵티마이저 스택 구축. 이를 밀집 및 희소 MoE 언어 모델 (Qwen3-0.6B, Gemma 3 1B, OLMoE-1B-7B, 축소된 gpt-oss) 에서 평가함.

4. 실험 결과

저자들은 FineWeb-Edu 의 10B 토큰 하위 집합을 대상으로 네 가지 아키텍처에서 사전 학습 실험을 수행했습니다.

• 밀집 모델 (Qwen3-0.6B, Gemma 3 1B):

  • 어휘 인덱스 행렬 (임베딩/LM 헤드) 에 대한 AdamW 를 RowNormM 또는 HybridPolarGradM으로 대체하는 것이 최종 검증 손실을 일관되게 개선했습니다.
  • 더 작은 Qwen3 모델에서는 이득이 미미했으나, 더 큰 Gemma 3 모델에서는 더 두드러졌으며, 어휘 행렬이 커지고 기울기가 더 이방성 (anisotropic) 이 됨에 따라 기하학 인식 업데이트가 더 중요해짐을 시사합니다.
  • SwiGLU MLP 프로젝션에 HybridPolarGradM을 사용하면 성능이 추가로 개선되었는데, 이는 키가 크고 얇은 행렬 (tall-skinny matrices) 의 행 스케일 불균형을 교정했기 때문일 것입니다.

• 희소 MoE 모델 (OLMoE-1B-7B, 축소된 gpt-oss):

  • 임베딩, LM 헤드, 라우터에 대칭성 호환 업데이트를 적용한 구성 (예: RowNormM + LeftPolarGradM) 은 이러한 구성 요소에 AdamW 를 사용한 구성보다 우수한 성능을 보였습니다.
  • 완전히 좌표별인 베이스라인 (전체 AdamW) 은 가장 낮은 성능을 보였습니다.
  • 대칭성 호환 라우터 업데이트는 좌표별 라우터 업데이트보다 개선되었으며, 일부 경우 훈련 손실 급증을 감소시켜 안정성 향상을 시사했습니다.

• 일반적 경향: 모든 모델에서 대규모 어휘 인덱스 행렬에 대한 AdamW 를 대칭성 호환 옵티마이저로 대체하는 것이 최종 검증 손실을 일관되게 낮췄습니다. 이러한 이점은 다양한 아키텍처와 희소성 수준에서 유지되었습니다.

5. 중요성 및 주장

이 논문은 옵티마이저 설계에 있어 "기하학 무지"인 좌표별 관점에서 "대칭성 인식" 행렬 관점으로의 원칙적 전환을 제시한다고 주장합니다.

• 아키텍처 - 옵티마이저 공동 설계: 이 작업은 옵티마이저 설계가 아키텍처 인식 문제여야 한다고 주장합니다. 서로 다른 매개변수 클래스 (임베딩, 라우터, 선형 레이어) 는 고유한 자연 대칭성을 가지며 서로 다른 옵티마이저 기하학을 요구합니다.
• 이론적 기반: 스펙트럼 옵티마이저가 선형 레이어에서 작동하는 이유에 대한 엄밀한 수학적 특성화 (공변성을 통해) 를 제공하고, 이 논리를 확장하여 다른 레이어에 대한 새로운 옵티마이저를 유도합니다.
• 실용적 유용성: 결과는 대칭성 호환 할당이 이론적으로 타당할 뿐만 아니라 실용적으로 유익하며, 대규모 사전 학습에서 최종 손실과 훈련 안정성을 개선함을 보여줍니다.
• 겸손함: 저자들은 명시적으로 공변적 옵티마이저가 모든 영역에서 좌표별 방법을 우세하다고 주장하지 않습니다. 대신, 레이어별 대칭성 호환 원칙이 대표적인 사전 학습 설정에서 경쟁력 있고 종종 유익한 할당으로 이어진다는 증거를 제시합니다. 그들은 결과를 원칙의 유용성에 대한 증거로 보며 포괄적인 벤치마크로 보지 않습니다.

이 논문은 대칭성과 공변성을 옵티마이저 설계에 도입함으로써 더 안정적이고, 매개변수화 일관성이 있으며, 이론적으로 기반을 둔 대규모 모델 사전 학습을 향한 길을 연다고 결론지었습니다.