[논문 리뷰] Taming Preconditioner Drift: Unlocking the Potential of Second-Order Optimizers for Federated Learning on Non-IID Data

Junkang Liu, Fanhua Shang외 4명|arXiv (Cornell University)|2026. 02. 22.

Stochastic Gradient Optimization Techniques인용 수 0

한 줄 요약

본 논문은 프리컨디셔너 드리프트를 연합 제2차 최적화의 비 IID 데이터에서 핵심 불안정성으로 식별하고, 로컬 프리컨디셔너를 정렬하고 보정하는 FedPAC를 도입하여 비전 및 언어 작업 전반에서 더 빠르고 더 안정적인 수렴을 얻습니다.

ABSTRACT

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 exttt{FedPAC}, a \emph{preconditioner alignment and correction} framework for reliable federated second-order optimization. exttt{FedPAC} explicitly decouples parameter aggregation from geometry synchronization by: (i) extbf{Alignment} (i.e.,aggregating local preconditioners into a global reference and warm-starting clients via global preconditioner); and (ii) extbf{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, exttt{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.

연구 동기 및 목표

• 비 IID 데이터에서 연합 제2차 최적화의 불안정성의 원인(프리컨디셔너 드리프트)을 식별한다.
• 전역 및 로컬 프리컨디셔너를 정렬하고 업데이트를 보정하는 통합 프레임워크 FedPAC를 제안한다.
• 비전과 언어 작업 전반에서의 안정성 및 정확도 향상에 대한 수렴 보장과 실증적 근거를 제공한다.

제안 방법

• 클라이언트 간 로컬과 글로벌 프리컨디셔너 간의 차이를 프리컨디셔너 드리프트로 정의하고 측정한다.
• 지오메트리 동기화(정렬)와 매개변수 집계(보정)를 분리하는 FedPAC를 제안한다.
• SOAP, Muon, Sophia 위에 FedPAC를 구축하여 FedPAC_Sophia, FedPAC_Muon, 및 FedPAC_SOAP를 얻는다.
• 로컬 프리컨디셔너를 글로벌 참조로 집계하여 정렬을 구현하고 이를 사용해 클라이언트를 워밍업 시작한다.
• 로컬 프리컨디셔드 업데이트와 글로벌 방향을 결합하는 트레이드오프 매개변수 beta를 사용하여 보정을 구현한다.
• 드리프트 결합 비선형 수렴 보장을 제공하여 드리프트 항을 감소시키고 더 빠른 수렴을 보인다.

Figure 1 : (a) In non-IID FL, first-order methods converge slowly, inducing little client drift. (b) Second-order methods converge faster locally and thus drift toward local optima, causing the aggregated global model to deviate from global optimum. (c) FedPAC corrects local second-order updates, yi

실험 결과

연구 질문

• RQ1 프리컨디셔너 드리프트가 비 IID 데이터에서 naive 연합 제2차 방법이 성능이 저조한 이유를 설명할 수 있는가?
• RQ2 FedPAC가 효과적으로 로컬 프리컨디셔너를 정렬하고 보정하여 이종 환경에서 1차 최적화 FL의 성능을 회복하거나 상회시킬 수 있는가?
• RQ3 표준 매끄러움성과 경계된 이질성 가정하에서 FedSOA와 FedPAC의 수렴 보장은 무엇인가?
• RQ4 FedPAC 변형들이 IID vs 비 IID 체제에서 CNN, Vision Transformer, 언어 모델에서 어떻게 성능을 발휘하는가?
• RQ5 FedPAC 성능에서 보정 강도 beta의 역할은 무엇인가?

주요 결과

• 2차 연합 제어 방법은 비 IID 데이터에서 글로벌 수렴을 저해하는 프리컨디셔너 드리프트를 겪는다.
• FedPAC는 프리컨디셔너 드리프트를 감소시키고 CNN, ViT, 언어 모델 전반에서 더 빠르고 안정적인 수렴을 제공한다.
• FedPAC 변형은 특히 강한 데이터 이질성(Dirichlet 분할) 하에서 로컬 2차 기본값에 비해 정확도를 일관되게 향상시킨다.
• FedPAC은 CIFAR-100 및 Tiny-ImageNet에서 기준선 대비 상당한 이점을 달성하고 LLaMA 모델의 C4 사전 학습에서 강력한 성능을 보인다.
• 이론적 결과는 드리프트 결합 수렴 보장을 보여주며, FedPAC은 명시적 이질성 항을 제거하고 드리프트와 관련된 노이즈를 감소시킨다.
• 적용 연구에서 정렬과 보정 둘 다 필요함이 확인되었고, beta가 대략 0.5일 때 강건한 성능을 제공한다.

(a) ResNet-18, IID