불안정성을 제어하는 강인한 2차 최적화기 연합학습

본 논문은 비IID 데이터 분포가 지배적인 연합학습 환경에서 2차 최적화 기법이 직면하는 핵심 문제를 체계적으로 분석하고, 이를 해결하기 위한 새로운 프레임워크인 Federated Robust Curvature Optimization(FedRCO)를 제안한다. 서론에서는 연합학습이 데이터 프라이버시를 보장하면서도 대규모 모델 학습을 가능하게 하지만, 기존의 FedAvg와 같은 1차 SGD 기반 방법은 통신 비용과 수렴 속도에서 한계가 있음을 지적한다. 특히 비IID 상황에서 클라이언트 간 손실 함수의 차이로 인한 클라이언트 드리프트가 발생하고, 이는 수렴 안정성을 크게 저하시킨다. 이러한 문제를 완화하기 위해 최근 2차 최적화, 특히 Kronecker‑Factored Approximate Curvature(K‑FAC)를 활용한 연구가 진행되고 있으나, 연합학습 특성상 (1) 작은 배치로 인한 FIM의 랭크 결핍, (2) 로컬 곡률과 글로벌 손실 곡률 간의 불일치라는 두 가지 주요 불안정 요인이 존재한다. 이론적 섹션에서는 두 가지 실패 모드에 대한 정량적 분석을 제공한다. Proposition 3.1은 배치 크기 B가 파라미터 차원 d보다 현저히 작을 경우 경험적 FIM이 랭크 결핍 상태가 되며, 이는 널 공간에 잡음이 축적돼 무한대 업데이트를 야기한다는 것을 증명한다. Proposition 3.2는 비IID 데이터에서 로컬 곡률이 글로벌 곡률을 과소/과대 추정하면, 사전조건된 업데이트가 비정상적으로 확대돼 발산한다는 점을 보인다. 이어서 Proposition 3.3·3.4는 2차 최적화가 조건수 κ를 1에 가깝게 만들면 SGD 대비 수렴 속도가 급격히 개선된다는 기존 이론을 재정리한다. 이러한 분석을 바탕으로 FedRCO는 불안정성을 억제하면서 2차 최적화의 장점을 살리는 설계를 제시한다. FedRCO의 핵심 구성은 다음과 같다. (1) Gradient Anomaly Monitor: 클라이언트는 사전조건된 그래디언트 노름의 슬라이딩 윈도우 평균을 기준으로 이상 점수 Sₖ를 계산한다. τ_low와 τ_high 두 임계값을 두어 점진적 드리프트와 급격한 폭발을 구분한다. (2) Fail‑Safe Resilience Protocol: 이상 점수가 τ_low 이하이면 정상 업데이트, τ_low와 τ_high 사이이면 사전 정의된 안정적인 그래디언트 상한(˜∇_stable)으로 소프트 롤백, τ_high를 초과하거나 연속적으로 τ_low를 초과하면 모든 곡률 통계와 옵티마이저 상태를 초기화하고 글로벌 모델을 재동기화하는 하드 리셋을 수행한다. (3) Curvature‑Preserving Adaptive Aggregation: 전통적인 평균 집계 대신 로컬 정확도와 글로벌 정확도를 이용해 동적 가중치 γ = Acc′/(Acc′+Acc)를 계산하고, θ_new_local = γ·θ_global + (1−γ)·θ_local 형태로 보간한다. 이는 로컬 곡률을 보존하면서도 글로벌 모델과의 일관성을 유지한다. 클라이언트 측 최적화는 블록‑다이아고날 K‑FAC 근사를 사용한다. 각 레이어에 대해 FIM을 Ω⊗Γ 형태로 근사하고, Tikhonov 정규화(ε)와 π‑보정(π = tr(A)/tr(G))을 적용해 역행렬 계산의 수치적 안정성을 확보한다. 또한, 작은 배치에서도 통계적 노이즈를 완화하기 위해 EMA(α)로 Ω와 Γ를 업데이트한다. 최종 사전조건된 자연그라디언트 업데이트는 vec(˜∇θ) = (Ω+π√ε I)⁻¹ ∇L (Γ+√ε/π I)⁻¹ 형태로 구현된다. 실험에서는 CIFAR‑10, FEMNIST, Shakespeare 등 다양한 비IID 시나리오에서 FedRCO를 기존 FedAvg, FedProx, FedNL, FedPM 등과 비교한다. 결과는 다음과 같다. (1) 수렴 속도: FedRCO는 동일 정확도에 도달하기 위해 평균 30%~45% 적은 통신 라운드를 필요로 한다. (2) 최종 정확도: 비IID 환경에서도 기존 2차 방법보다 1%~2% 높은 정확도를 달성한다. (3) 안정성: 그래디언트 폭발 발생 빈도가 기존 방법 대비 70% 이상 감소하고, 하드 리셋 횟수도 크게 줄어든다. (4) 연산 비용: 블록‑다이아고날 K‑FAC와 EMA를 사용함으로써 클라이언트 측 연산량은 1차 SGD 대비 1.5배 수준에 머물며, 모바일 디바이스에서도 실용적으로 실행 가능함을 보인다. 결론적으로, FedRCO는 비IID 연합학습에서 2차 최적화의 핵심 장애인 곡률 불일치와 랭크 결핍을 실시간 감시와 복구 메커니즘으로 완화하고, 곡률을 보존하는 적응형 집계 방식을 통해 글로벌 모델과 로컬 모델 간의 기하학적 일관성을 유지한다. 이론적 증명과 실험적 검증을 통해 제안된 프레임워크가 기존 1차·2차 방법보다 뛰어난 수렴 효율과 안정성을 제공함을 입증한다. 향후 연구에서는 FedRCO를 더욱 경량화된 근사 기법과 결합하거나, 프라이버시 보호를 위한 암호화된 곡률 집계와 연계하는 방향을 제시한다.