노이즈 적응형 레이어별 학습률을 통한 기하학 인식 최적화 가속화

초록

본 논문은 Muon·Scion 등 기존 기하학 인식 옵티마이저가 동일 그룹 내 레이어에 고정 학습률을 적용하는 한계를 지적하고, 각 레이어의 그래디언트 노이즈를 실시간으로 추정해 층별 학습률을 동적으로 조정하는 LANTON 알고리즘을 제안한다. 이 방법은 이중 노름(dual norm) 기반 노이즈 추정, 모멘텀 기반 변동 추적, 그리고 그룹별 최적 LMO 선택을 결합해 수렴 속도를 이론적으로 향상시키며, LLaMA·GPT 등 대형 트랜스포머와 CNN 실험에서 기존 최첨단 옵티마이저 대비 학습 효율과 샘플 효율을 크게 높였다.

상세 분석

LANTON은 기존 기하학 인식 최적화 프레임워크(예: Muon, Scion)의 핵심 구조를 그대로 유지하면서, 레이어별 노이즈 적응형 학습률 스케일링을 추가한다. 먼저 각 레이어 ℓ에 대해 스토캐스틱 그래디언트 Gℓ,t와 이전 단계 그래디언트 Gℓ,t‑1(또는 독립적인 ˜Gℓ,t)을 이용해 이중 노름 ‖·‖*에 대한 노이즈 추정값 Hℓ,t를 β₂ 모멘텀 방식으로 누적한다. 여기서 사용되는 이중 노름은 해당 레이어가 속한 파라미터 그룹에 따라 달라지며, 예를 들어 hidden‑matrix 레이어는 RMS→RMS 연산자 노름의 듀얼인 nuclear norm, 임베딩·LM‑head 레이어는 ℓ₁→ℓ_∞ 노름의 듀얼인 ℓ₁→ℓ₁ 노름을 사용한다.

노이즈 추정값 Hℓ,t가 커질수록 해당 레이어는 더 큰 변동성을 보인다고 판단하고, αℓ,t = α / √(α² + Hℓ,t) 로 스케일링 팩터를 계산한다. 이후 그룹 내 최대 스케일링 α_m,t 로 정규화한 비율 αℓ,t/α_m,t 를 기본 학습률 η_t(코사인 디케이 일정)와 곱해 최종 레이어별 학습률 ηℓ,t를 얻는다. 이 설계는 “노이즈가 큰 레이어는 작은 스텝을, 노이즈가 작은 레이어는 큰 스텝을” 취함으로써, 전통적인 SGD/Adam 계열이 가정하는 균일 노이즈 모델을 넘어선 보다 정교한 학습 동역학을 구현한다.

이론적 분석에서는 ‖∇f(X_t)‖_dual 기준으로 ˜O(1/√T + p·P_ℓ·σ̄_ℓ / T^{1/4}) 의 수렴 속도를 증명한다. 여기서 σ̄_ℓ는 레이어 ℓ의 노이즈 상한, p는 파라미터 그룹 수, P_ℓ는 해당 그룹의 차원 비율이다. 기존 기하학 옵티마이저가 보여준 ˜O(1/√T)와 비교해, 레이어별 노이즈를 명시적으로 고려함으로써 추가적인 T^{-1/4} 향상을 얻는다.

실험에서는 LLaMA‑7B, GPT‑2‑XL 등 대규모 트랜스포머와 ResNet‑50 기반 이미지 분류 모델에 LANTON을 적용했다. 동일한 배치 크기·학습 스케줄 하에서, LANTON은 동일 에폭당 손실 감소율이 10‑15% 정도 빠르고, 최종 퍼플렉시티·정확도에서도 0.3‑0.5%의 개선을 기록했다. 특히 초기 학습 단계에서 노이즈가 급격히 변하는 레이어(예: QK 매트릭스)의 학습률이 자동으로 감소하면서 불안정성을 크게 완화시켰다.

알고리즘 구현 측면에서는 기존 LMO 호출을 그대로 재사용하고, 노이즈 추정에 필요한 추가 연산은 벡터 ‖·‖₂ 차원에서의 간단한 차이 계산과 EMA 업데이트 정도에 불과해, 전체 연산량 및 메모리 오버헤드가 5% 이하로 제한된다. 또한 옵션 II(독립적인 ˜Gℓ,t 사용)는 이론적 보장을 제공하지만, 실제 실험에서는 옵션 I(이전 단계 그래디언트 사용)만으로도 충분히 좋은 성능을 얻었다.

요약하면, LANTON은 “기하학 인식 + 노이즈 적응”이라는 두 축을 결합해, 레이어별 곡률·노이즈 이질성을 정량화하고 이를 학습률에 직접 반영함으로써, 대규모 모델 훈련 시 기존 최첨단 옵티마이저 대비 효율성을 크게 끌어올린 혁신적인 접근법이다.