연합 학습에서 모멘텀과 SAM의 갈등을 해소하는 FedWMSAM
📝 원문 정보
- Title: A Fast and Flat Federated Learning Method via Weighted Momentum and Sharpness-Aware Minimization
- ArXiv ID: 2511.22080
- 발행일: 2025-11-27
- 저자: Tianle Li, Yongzhi Huang, Linshan Jiang, Chang Liu, Qipeng Xie, Wenfeng Du, Lu Wang, Kaishun Wu
📝 초록 (Abstract)
연합 학습에서는 제한된 통신 예산 하에서 모델이 빠르게 수렴하면서 비동질적인 클라이언트 분포 전반에 일반화되어야 한다. 이러한 이중 요구는 자연스럽게 두 가지 널리 사용되는 기법, 즉 진행을 가속화하는 클라이언트·서버 모멘텀과 평탄한 해를 선호하는 Sharpness‑Aware Minimization, 약칭 SAM을 도입하게 만든다. 그러나 모멘텀과 SAM을 단순히 결합하면 비동질적인 연합 학습 환경에서 두 가지 구조적 문제점이 남는다. 우리는 이를 ‘지역‑전역 곡률 불일치’와 ‘모멘텀‑에코 진동’이라는 두 실패 모드로 규정하고 공식화한다. 전자는 지역 SAM 방향이 전역 손실 지오메트리를 반영하지 못함을, 후자는 누적된 모멘텀이 후기 학습 단계에서 불안정을 초래함을 의미한다. 이 두 실패 모드를 동시에 다루는 방법은 기존 연구에서 제시된 바 없다. 우리는 FedWMSAM을 제안한다. 첫째, 서버가 집계한 모멘텀으로부터 전역 교란을 생성해 클라이언트의 SAM 방향을 전역 하강 지오메트리와 정렬시킨다. 이를 통해 효율성을 유지하면서 단일 역전파 SAM 근사를 가능하게 한다. 둘째, 모멘텀과 SAM을 코사인 유사도 기반 적응 규칙으로 결합해 초기에는 모멘텀을, 후기에는 SAM을 적용하는 두 단계 학습 스케줄을 만든다. 이론적으로는 비동질성 하에서 교란이 초래하는 분산 σ²_ρ 를 명시적으로 모델링한 수렴 경계를 제시하고, 그 경계가 클라이언트 수, 로컬 업데이트 횟수, 라운드 수, 전체 데이터 양에 어떻게 의존하는지를 분석한다. 실험적으로는 다양한 데이터셋과 모델 구조에 걸쳐 광범위한 평가를 수행했으며, 제안 방법이 최적화 난제들을 효과적으로 해결하고 적응성과 견고성을 동시에 갖추어 기존 방법보다 우수함을 입증하였다.💡 논문 핵심 해설 (Deep Analysis)
첫 번째 결함인 ‘지역‑전역 곡률 불일치’는 SAM이 각 클라이언트에서 로컬 손실에 대한 교란을 생성한다는 점에서 비롯된다. 비동질적인 데이터 환경에서는 로컬 손실의 곡률이 전역 손실의 곡률과 크게 다를 수 있다. 따라서 로컬 SAM이 탐색하는 평탄한 영역이 실제 전역 최적화 경로와 일치하지 않아, 전체 모델이 비효율적인 방향으로 움직이게 된다. 이 문제를 해결하려면 로컬 SAM 방향을 전역 손실의 기하학에 맞추어야 하는데, 기존 방법은 서버와 클라이언트 간에 추가적인 고비용 통신이나 복잡한 2차 미분 계산을 요구한다.
두 번째 결함인 ‘모멘텀‑에코 진동’은 학습 후기, 특히 SAM이 적용된 후에 모멘텀 값이 누적되면서 발생한다. 모멘텀은 과거 그라디언트 정보를 보존해 학습을 가속화하지만, SAM에 의해 손실 표면이 인위적으로 변형되면 이 보존된 정보가 현재 손실 지형과 충돌한다. 결과적으로 학습이 진동하거나 발산하는 현상이 나타나며, 이는 특히 통신 라운드가 제한된 상황에서 심각한 성능 저하를 초래한다.
FedWMSAM은 이러한 두 문제를 동시에 해결하도록 설계되었다. 첫 번째 단계에서는 서버가 각 라운드에서 집계한 모멘텀을 이용해 ‘전역 교란’(global perturbation)을 생성한다. 이 교란은 서버‑집계 모멘텀의 방향과 크기를 반영하므로, 클라이언트가 로컬 SAM을 수행할 때 전역 손실의 곡률을 내재화한다. 이렇게 하면 별도의 2차 미분 연산 없이도 단일 역전파 단계만으로 전역‑정렬된 SAM을 구현할 수 있어 통신·연산 비용이 크게 늘어나지 않는다.
두 번째 단계에서는 모멘텀과 SAM을 코사인 유사도 기반 적응 규칙으로 결합한다. 초기 학습 단계에서는 모멘텀과 SAM 사이의 유사도가 낮아 모멘텀 중심의 업데이트가 진행되고, 학습이 진행될수록 두 벡터의 방향이 점차 일치하게 된다. 이때 적응 규칙은 SAM의 비중을 점진적으로 높여 ‘초기‑모멘텀, 후기‑SAM’ 스케줄을 구현한다. 결과적으로 초기에는 빠른 수렴을, 후기에는 평탄한 최적점을 향한 미세 조정을 동시에 달성한다.
이론적 분석에서는 비동질성 하에서 교란이 초래하는 추가 분산 σ²_ρ 를 명시적으로 모델링하고, 클라이언트 수(S), 로컬 업데이트 횟수(K), 라운드 수(R), 전체 데이터 양(N)과의 관계를 정량화한 수렴 경계를 제시한다. 이 경계는 기존 모멘텀·SAM 결합 방식보다 더 타이트하며, 특히 σ²_ρ 가 작을수록 빠른 수렴을 보장한다는 점에서 실용적 의미가 크다.
실험 부분에서는 이미지 분류(CIFAR‑10, CIFAR‑100), 자연어 처리(SST‑2) 등 다양한 데이터셋와 ResNet, Transformer 등 여러 모델 구조에 대해 광범위한 비교를 수행했다. FedWMSAM은 통신 라운드가 제한된 상황에서도 정확도와 수렴 속도 모두에서 기존 최첨단 방법들을 앞섰으며, 특히 데이터 비동질성이 심한 경우에 그 우수성이 두드러졌다. 또한, 하이퍼파라미터 민감도 실험을 통해 제안된 적응 규칙이 다양한 설정에 강건함을 확인했다.
요약하면, FedWMSAM은 전역‑정렬된 SAM 교란과 모멘텀‑SAM 적응 스케줄을 통해 연합 학습의 두 핵심 난제인 빠른 수렴과 일반화 성능을 동시에 달성한다는 점에서 중요한 진전을 제공한다. 이는 앞으로 비동질성 환경에서 효율적인 연합 학습을 설계하는 데 새로운 패러다임을 제시한다.
📄 논문 본문 발췌 (Translation)
📸 추가 이미지 갤러리
