기술 요약: Zeta – 행렬 최적화를 위한 좌표 적응형 프리컨디셔닝 기반의 이중 백색화(Dual Whitening)

문제 정의

대규모 신경망 학습은 요소별 적응(element-wise adaptation)을 넘어 가중치 파라미터의 구조적 특성을 활용하는 행렬 인식 옵티마이저(matrix-aware optimizers)로 점차 이동하고 있습니다. 대표적인 예로 Muon이 있는데, 이는 뉴턴-슐츠 반복법(Newton–Schulz iteration)을 통한 스펙트럼 백색화(spectral whitening)를 통해 싱귤러 값(singular values)을 균등하게 만들고 특정 방향이 업데이트를 독점하는 것을 방지함으로써 가중치 업데이트에 직교 제약(orthogonal constraints)을 부여합니다.

그러나 본 논문은 Muon 및 유사한 스펙트럼 방법론에서 간과되어 온 결정적인 취약점을 식별합니다: 뉴턴-슐츠 반복법은 입력 컨디셔닝(input conditioning)에 매우 민감하다는 점입니다. 뉴턴-슐츠 반복법은 싱귤러 값이 잘 분포되어 있을 때는 잘 작동하지만, 심층 네트워크 학습 과정에서의 가공되지 않은 모멘텀 행렬은 종종 심각한 **좌표별 스케일 불균형(coordinate-wise scale heterogeneity)**을 보입니다. 이러한 행렬에서는 일부 좌표가 불균형하게 큰 크기를 가짐으로써 싱귤러 스펙트럼을 왜곡합니다. 이렇게 컨디셔닝되지 않은 행렬이 뉴턴-슐츠 단계에 진입하면, 직교화 과정이 저하되어 그래디언트 기하학(gradient geometry)을 제대로 반영하지 못하는 왜곡된 업데이트 방향을 초래합니다. 저자들은 스펙트럼 백색화가 스케일이 정규화된 입력을 암묵적으로 요구하며, 가공되지 않은 모멘텀은 이 조건을 체계적으로 위반한다고 주장합니다.

방법론: Zeta 옵티마이저

이러한 구조적 격차를 해결하기 위해, 저자들은 좌표 백색화(coordinate whitening)와 스펙트럼 백색화(spectral whitening)를 엄격하게 순서가 정해진 비튜닝(non-tunable) 파이프라인으로 구성한 이중 백색화 옵티마이저인 Zeta를 제안합니다. 핵심 통찰은 스펙트럼 백색화가 통계적 안정화 위에 구축되어야 한다는 것입니다. 즉, 좌표 백색화가 스펙트럼 백색화가 안정적으로 기능하는 데 필요한 통계적 등방성(statistical isotropy)을 확립합니다.

Zeta 파이프라인은 다음과 같이 작동합니다:

1. 좌표 백색화 (Preconditioning):
   스펙트럼 단계에 진입하기 전, Zeta는 모멘텀 행렬에 좌표별 정규화를 적용합니다. 이는 Adam과 유사하게 1차 모멘트($M_t$)와 2차 모멘트($V_t$)의 지수 이동 평균을 유지합니다. 그 후 모멘텀 행렬은 실행 중인 2차 모멘트로 요소별 정규화됩니다:
   $$\tilde{G}_t = \frac{M_t}{\sqrt{V_t} + \epsilon}$$
   이 단계는 좌표 간의 스케일 격차를 균질화하여 행렬을 통계적으로 등방적인 상태로 변환합니다. 이론적으로 이는 근사적인 대각 프리컨디셔너(diagonal preconditioner) 역할을 수행하며, 싱귤러 스펙트럼을 왜곡할 수 있는 스케일 불균형 문제를 해결합니다.

2. 스펙트럼 백색화 (Orthogonalization):
   균질화된 행렬 $\tilde{G}_t$는 이후 뉴턴-슐츠 반복법에 투입됩니다. 입력값이 이제 극단적인 요소별 스케일 격차로부터 자유로워졌기 때문에, 싱귤러 스펙트럼의 컨디션이 개선되었습니다. 이 반복 과정은 싱귤러 값을 균등하게 만들어 균형 잡힌 업데이트 방향을 생성하는 충실한 직교 근사치로 수렴합니다.

3. 업데이트 구성:
   결과물인 직교 행렬 $U_t$는 스펙트럼 방향을 변경하지 않으면서 표준 옵티마이저 설정(예: AdamW 브랜치)과 업데이트 크기를 호환시키기 위해 RMS 정규화 계수를 사용하여 스케일링됩니다.

4. 이중 경로 실행 (Dual-Path Execution):
   행렬 형태의 파라미터(예: 어텐션 및 피드포워드 프로젝션)에 대해서는 전체 이중 파이프라인이 적용됩니다. 비행렬 파라미터(예: 바이어스, LayerNorm 스케일)의 경우, 스펙트럼 단계를 건너뛰고 AdamW 방식의 업데이트로 전환됩니다.

이론적 기여

본 논문은 설계의 근거가 되는 세 가지 핵심 이론적 보증을 제공합니다:

• 근사 백색화 (Approximate Whitening): 좌표 정규화 단계가 단순한 휴리스틱 스케일링을 넘어, 이상적인 백색화 변환에 대한 효율적인 대각 근사임을 증명합니다 (Theorem 1).
• 직교화 오차 감소 (Reduced Orthogonalization Error): 좌표 백색화가 뉴턴-슐츠 반복법에 들어가는 입력의 컨디션 넘버(condition number)를 엄격하게 개선함을 증명합니다. 결과적으로, 가공되지 않은 모멘텀을 사용하는 것보다 스펙트럼 단계의 직교화 오차가 엄격하게 감소함을 입증합니다 (Theorem 2).
• 스케일 보존 (Scale Preservation): 최종 RMS 스케일링은 업데이트의 스펙트럼 방향을 보존하면서 네트워크의 나머지 부분과 글로벌 크기를 동기화합니다 (Theorem 3).

실험 결과

Zeta는 언어 모델링(0.6B ~ 8B 파라미터), 혼합 전문가(MoE) 아키텍처, 비전 태스크에 걸쳐 평가되었으며, 강력한 베이스라인(AdamW, Muon, AdaMuon)과 일관되게 대등하거나 이를 능가하는 성능을 보였습니다.

• 수렴 속도: 사전 학습 태스크에서 Zeta는 Muon 및 AdamW보다 빠른 수렴을 보여주었습니다. 예를 들어, 1.7B 파라미터 모델에서 Zeta는 동일한 손실 수준에 도달하기 위해 AdamW 대비 1.67배 빠른 속도(상대적 속도)를 달켰으며, 이는 Muon(1.52배) 및 AdaMuon(1.61배)을 능가하는 수치입니다. 8B 모델에서도 Zeta는 Muon의 1.19배 대비 1.25배의 속도 향상을 달성했습니다.
• 효율성: Zeta의 반복당 시간 비용은 AdamW와 유사한 수준을 유지했습니다 (예: 8B 모델, 16개 GPU 환경에서 16개 GPU에서 25.47초 vs 24.46초). 이는 이중 백색화 과정이 유의미한 런타임 페널티를 발생시키지 않음을 나타냅니다.
• MoE 성능: 희소 전문가 활성화(sparse expert activation)가 상당한 그래디언트 불균형을 생성하는 MoE 아키텍처에서, Zeta는 경쟁 모델들보다 낮은 손실값으로 수렴하며 더 적은 변동성을 유지하는 명확한 우위를 보였습니다.
• 다운스트림 태스크: 12개의 언어 이해 벤치마크(MMLU, GSM8k 등)에서 Zeta로 학습된 모델은 1.7B 및 8B 규모 모두에서 가장 높은 평균 정확도를 기록했습니다. 마찬가지로 CIFAR-100 비전 태스크에서도 Zeta는 ViT-Tiny 및 ViT-Base 모두에서 가장 높은 Top-1 검증 정확도를 달성했습니다.
• 하이퍼파라미터 강건성: 모멘텀 계수($\beta_1, \beta_2$)에 대한 그리드 서치를 통해, Zeta가 하이퍼파라미터 선택에 매우 강건함을 확인했습니다. 테스트된 범위 내에서 최종 학습 손실의 변동은 1.2% 미만이었습니다.

의의 및 주장

본 논문은 기존의 행렬 인식 옵티마이저들이 가진 근본적인 불일치, 즉 제어되지 않은 좌표별 스케일을 가진 신호에 강력한 기하학적 연산(스펙트럼 백색화)을 적용하는 문제를 해결한다고 주장합니다. 저자들은 카이제곱 균등성 검사(chi-square uniformity tests)를 통해 행렬 내부의 스케일 불균형이 만연함을 검증하고, 좌표 백색화가 이를 효과적으로 교정함을 입증함으로써, 스케일 불균형을 직교화 이전에 해결하는 것이 안정적인 스펙트럼 방법론의 전제 조건임을 확립했습니다.

Zeta의 의의는 두 가지 이전에 경쟁 관계에 있던 최적화 패러다임(좌표별 적응형 스케일링과 스펙트럼 직교화)을 하나의 논리적으로 필연적인 파이프라인으로 통합했다는 점에 있습니다. 연산의 순서는 튜닝 가능한 하이퍼파라미터가 아니라 수학적 의존 관계입니다. 즉, 입력의 컨디션을 양호하게 유지하기 위해서는 좌표 백색화가 반드시 스펙트럼 백색화보다 선행되어야 합니다. 이러한 접근 방식은 행렬 수준 최적화의 기하학적 이점을 희생하지 않으면서도 더 빠른 수렴과 더 나은 일반화 성능을 이끌어냅니다.