대규모 배치 학습을 위한 피셔 직교 투영 기반 자연 기울기 최적화

초록

본 논문은 초대형 배치 학습에서 기존 자연 기울기(K-FAC)의 고감쇠 문제를 해결하기 위해 두 개의 서브 배치 그래디언트를 이용해 피셔 메트릭 상에서 직교화된 보정 성분을 추가하는 Fisher‑Orthogonal Projection(FOP) 기법을 제안한다. 평균 그래디언트에 변동성을 반영한 orthogonal component를 결합함으로써 대규모 배치에서도 안정적인 2차 정보 활용이 가능해지고, 실험에서 SGD/AdamW 대비 1.5‑7.5배 빠른 수렴 및 장기 불균형 데이터셋에서 2‑3% 향상된 Top‑1 정확도를 보였다.

상세 분석

FOP는 대규모 배치에서 발생하는 두 가지 근본적인 문제를 동시에 해결한다. 첫째, 배치 크기가 커짐에 따라 그래디언트 노이즈가 급격히 감소하면서 SGD와 같은 1차 최적화기가 평탄한 최소점으로 탈출하기 어려워진다. 둘째, K‑FAC과 같은 자연 기울기 방법은 피셔 행렬이 고조건수화되면서 수치적 불안정성이 커지고, 이를 억제하기 위해 과도한 damping을 적용하면 곧바로 2차 정보가 사라져 단순 GD와 동일한 성능을 보인다.

FOP는 동일 파라미터 θ에서 두 개의 독립 서브 배치 L₁, L₂를 사용해 각각의 그래디언트 g₁, g₂를 계산한다. 평균 그래디언트 g_avg = (g₁+g₂)/2와 차이 그래디언트 g_diff = g₁−g₂를 구한 뒤, 피셔 메트릭 F에 대한 내적 ⟨·,·⟩F를 이용해 g_diff를 g_avg에 직교화한다. 구체적으로 스칼라 s_proj = (g_diffᵀF g_avg)/(g_avgᵀF g_avg+ε) 로 정규화하고, g⊥diff = g_diff − s_proj·g_avg 를 얻는다. 이렇게 하면 ⟨g_avg, g_⊥diff⟩_F = 0 가 보장되어, 평균 그래디언트가 이미 포착한 정보와 중복되지 않는 순수한 변동성 성분만을 추가한다.

결합된 업데이트 방향은 g_combined = g_avg + β·g_⊥diff 로 정의되며, β는 각 레이어별로 최적화된 값 β* = D/E 로 자동 조정된다. 여기서 D = g_avgᵀF⁻¹ g_⊥diff, E = g_⊥diffᵀF⁻¹ g_⊥diff 로 정의되며, 이는 2차 테일러 전개를 통한 손실 감소량을 최소화하는 closed‑form 해이다. β가 0에 가까워지면 FOP는 기존 K‑FAC과 동일하게 동작하고, β가 양수이면 변동성 보정이 강화되어 대규모 배치에서도 충분한 탐색성을 유지한다.

또한 레이어별 학습률 η*_ℓ 를 g_totᵀF⁻¹ g_combined / (g_combinedᵀF⁻¹ g_combined) 로 정의해, 각 레이어의 곡률과 정렬 정도에 따라 자동으로 스케일링한다. 이는 고감쇠 상황에서도 과도한 스텝을 방지하고, 필요한 경우 전체 학습률 η₀ 와 곱해 최종 업데이트 d_ℓ = η₀·η*_ℓ·F⁻¹_ℓ·g_combined_ℓ 를 산출한다.

KL‑norm 분석에서는 FOP 업데이트의 KL 거리 ‖F^{1/2}Δ_FOP‖² 가 기본 K‑FAC 항(∝ λ⁻²)과 교차·보정 항(∝ λ⁻¹)으로 분해됨을 보인다. 즉, damping λ 가 커질수록 기본 항은 급격히 감소하지만 보정 항은 완만히 감소하므로, 큰 λ 상황에서도 변동성 보정이 상대적으로 더 큰 영향을 미친다. 이는 초기 학습 단계에서 큰 g_avg 와 의미 있는 g_⊥diff 가 존재할 때, β가 음수가 되어 교차 항이 KL 거리를 감소시켜 damping을 완화시키는 메커니즘을 제공한다.

시스템 구현 측면에서는 두 개의 독립 GPU 그룹이 각각 g₁, g₂ 를 AllReduce 로 계산하고, 레이어별 전문 GPU가 자신의 피셔 블록을 업데이트·역산한다. 이후 각 전문 GPU가 로컬 F⁻¹ 와 두 그래디언트를 이용해 FOP을 적용하고, 결과를 브로드캐스트한다. 이중 그래디언트 전략은 통신 비용을 크게 증가시키지 않으면서도 서브 배치 간 변동성을 효율적으로 활용한다.

실험에서는 ResNet‑18, Vision Transformer 등 다양한 모델에 대해 배치 크기 2¹¹2¹⁶(204865536)까지 확장했으며, FOP는 K‑FAC 대비 1.2‑1.3배, SGD/AdamW 대비 1.5‑1.7배 빠른 수렴을 보였다. 특히 32768~50000 배치에서는 최대 7.5배 속도 향상을 기록했고, CIFAR‑10/100 장기 불균형 벤치마크에서는 Top‑1 오류를 2.3‑3.3% 감소시켰다. 작은 배치(≤1024)에서도 정확도 저하 없이 기존 최적화기와 동등한 성능을 유지한다. 코드가 pip 패키지로 제공되어 기존 파이프라인에 한 줄만 추가하면 적용 가능하다.

요약하면, FOP는 피셔 메트릭 기반 직교 보정을 통해 대규모 배치에서도 유용한 2차 정보를 보존하고, 자동 β·η 조정 메커니즘으로 안정성을 확보한다. 이는 고감쇠 문제를 근본적으로 완화하고, 대규모 데이터센터 환경에서 자연 기울기 최적화를 실용화하는 중요한 진전이다.