프리컨디셔너 드리프트 제어: Non-IID 데이터 연합 학습을 위한 2차 최적화기의 잠재력 실현

Taming Preconditioner Drift: Unlocking the Potential of Second-Order Optimizers for Federated Learning on Non-IID Data

2차 최적화기는 대규모 학습을 크게 가속화할 수 있지만, 이를 단순히 연합 학습에 적용할 경우 Non-IID 데이터에서 불안정하거나 발산하는 문제가 자주 발생한다. 본 연구에서는 그 주요 원인이 '프리컨디셔너 드리프트(preconditioner drift)'임을 밝힌다. 즉, 클라이언트 측의 2차 학습은 서로 다른 '곡률로 정의된 기하학적 구조'(즉, 프리컨디셔너 좌표계)를 유도하며, 서버 측의 모델 평균화 과정에서 호환되지 않는 척도로 계산된 업데이트가 결합되어 글로벌 하강 방향을 훼손하게 된다. 이러한 기하학적 불일치를 해결하기 위해, 우리는 신뢰할 수 있는 연합 2차 최적화를 위한 '프리컨디셔너 정렬 및 보정' 프레임워크인 FedPAC을 제안한다. FedPAC은 다음 두 가지 과정을 통해 파라미터 집계와 기하학적 동기화를 명시적으로 분리한다. (i) 정렬(Alignment): 로컬 프리컨디셔너를 글로벌 참조로 집계하고 글로벌 프리컨디셔너를 통해 클라이언트를 웜 스타트(warm-start)한다. (ii) 보정(Correction): 장기적인 드리프트를 억제하기 위해 글로벌 전처리 방향을 사용하여 로컬 전처리 업데이트를 조정한다. 또한, 우리는 부분 참여 환경에서 선형 속도 향상을 포함한 드리프트 결합 비볼록 수렴 보장을 제공한다. 실험 결과, FedPAC은 비전 및 언어 태스크 전반에서 안정성과 정확도를 일관되게 개선하였으며, 특히 ViT를 적용한 CIFAR-100에서 최대 5.8%의 절대 정확도 향상을 달성하였다. 코드는 https://anonymous.4open.science/r/FedPAC-8B24 에서 이용 가능하다.

Second-order optimizers can significantly accelerate large-scale training, yet their naive federated variants are often unstable or even diverge on non-IID data. We show that a key culprit is \emph{preconditioner drift}: client-side second-order training induces heterogeneous \emph{curvature-defined geometries} (i.e., preconditioner coordinate systems), and server-side model averaging updates computed under incompatible metrics, corrupting the global descent direction. To address this geometric mismatch, we propose \texttt{FedPAC}, a \emph{preconditioner alignment and correction} framework for reliable federated second-order optimization. \texttt{FedPAC} explicitly decouples parameter aggregation from geometry synchronization by: (i) \textbf{Alignment} (i.e.,aggregating local preconditioners into a global reference and warm-starting clients via global preconditioner); and (ii) \textbf{Correction} (i.e., steering local preconditioned updates using a global preconditioned direction to suppress long-term drift). We provide drift-coupled non-convex convergence guarantees with linear speedup under partial participation. Empirically, \texttt{FedPAC} consistently improves stability and accuracy across vision and language tasks, achieving up to $5.8\%$ absolute accuracy gain on CIFAR-100 with ViTs. Code is available at https://anonymous.4open.science/r/FedPAC-8B24.