빠르게 수렴하고 통신은 최소화: Hessian 기반 연합 제로오더 최적화

초록

HiSo는 전역 대각 Hessian 근사를 이용해 제로오더 연합 학습의 수렴 속도를 크게 높이면서, 스칼라만 전송하는 차원‑자유 통신을 그대로 유지한다. 비볼록 함수에 대해 Lipschitz 상수와 차원에 독립적인 가속 수렴률을 이론적으로 증명하고, LLM 파인튜닝 벤치마크에서 기존 ZO‑FL 방법보다 1∼5배 적은 라운드로 동일 혹은 더 높은 정확도를 달성한다.

상세 분석

본 논문은 연합 학습(Federated Learning, FL) 환경에서 대규모 언어 모델(LLM) 파인튜닝 시 발생하는 통신 병목을 해결하기 위해 제로오더(Zeroth‑Order, ZO) 최적화를 차원‑자유 방식으로 활용한다는 점에서 출발한다. 기존 ZO‑FL 방법인 DeComFL은 스칼라(그라디언트 스칼라와 난수 시드)만을 전송함으로써 통신량을 모델 차원 d와 무관하게 MB 수준으로 축소했지만, 무작위 탐색 방향에만 의존하는 ZO‑SGD는 수렴 속도가 매우 느리다는 한계가 있었다.

HiSo는 이러한 한계를 극복하기 위해 두 가지 핵심 아이디어를 도입한다. 첫째, 전역 대각 Hessian 근사를 사전조건(precondition)으로 사용한다. 정확한 Hessian을 전송하지 않고도, 각 라운드마다 서버가 전역 대각 Hessian 추정치를 유지하고, 클라이언트는 이를 이용해 “H⁻¹/₂ u” 형태의 변형된 탐색 방향을 생성한다. 여기서 u는 기존과 동일하게 표준 정규분포에서 샘플링된 난수 벡터이며, H⁻¹/₂는 대각 근사이므로 스칼라 연산만으로 구현 가능하다. 결과적으로 업데이트는

Δxᵢ = (fᵢ(xᵢ+μH⁻¹/₂u)−fᵢ(xᵢ))/μ · H⁻¹/₂u

와 같이 표현되며, 이는 기대값이 H⁻¹∇fᵢ와 동일한 Newton‑style 업데이트가 된다.

둘째, 스칼라‑전송 구조를 유지하기 위해 Hessian 근사는 전역적으로 공유되는 “스칼라만” 전송되는 형태로 설계된다. 구체적으로, 각 클라이언트는 로컬에서 두 번의 함수 평가(μ·u와 −μ·u)만 수행해 대각 Hessian 원소를 추정하고, 이를 평균하여 서버에 전달한다. 이 과정에서 전송되는 정보는 여전히 스칼라(각 파라미터별 추정값)와 난수 시드뿐이며, 차원에 비례하는 부하가 발생하지 않는다.

이론적 분석에서는 저효율 랭크와 화이트닝(whitening) 가정 하에 Hessian‑inform ZO 그라디언트의 분산을 엄격히 제한한다. 이를 통해 비볼록 함수에 대해 L(리프시츠 상수)와 d에 독립적인 수렴률 O(1/√T)·(1/√K) (T: 라운드 수, K: 로컬 업데이트 수)를 도출한다. 특히, 기존 ZO‑FL의 최악‑사례 O(d)와 대비해 차원에 무관한 상수‑레벨 수렴을 보장한다는 점이 혁신적이다.

실험에서는 OPT‑1.3B, LLaMA‑7B 등 다양한 LLM을 여러 데이터셋(예: 텍스트 분류, 질문‑응답, 요약)에서 파인튜닝하였다. HiSo는 DeComFL 대비 평균 2.3배, 최악‑사례 5배까지 라운드 수를 절감했으며, 최종 정확도·BLEU·ROUGE 등에서도 동등하거나 약간 우수한 결과를 기록했다. 또한, 첫 번째‑순서 기반 연합 방법(FedAvg, FedAdam 등)과 비교했을 때 통신량은 10⁸∼10⁹ 배 절감되는 동시에 수렴 속도는 경쟁 수준을 유지한다.

결과적으로 HiSo는 “Hessian‑informed”라는 두 번째 차원의 정보를 활용하면서도 “scalar‑only communication”이라는 제로오더 연합 학습의 핵심 장점을 손상시키지 않는다. 이는 차원‑자유 통신이 필수적인 초대규모 모델 파인튜닝에 있어, 수렴 속도와 통신 효율을 동시에 최적화할 수 있는 새로운 패러다임을 제시한다.