거칠기 기반 연합 학습 알고리즘 RI FedAvg

초록

본 논문은 비동질적 데이터 환경에서 발생하는 클라이언트 드리프트 문제를 완화하기 위해 손실 함수의 거칠기를 정량화하는 Roughness Index(RI)를 활용한 정규화 항을 로컬 목표에 추가한 새로운 연합 학습 알고리즘 RI‑FedAvg을 제안한다. 비볼록 목표에 대한 수렴성을 이론적으로 증명하고, MNIST, CIFAR‑10, CIFAR‑100 데이터셋에서 기존 FedAvg, FedProx, FedDyn, SCAFFOLD, DP‑FedAvg 등을 능가하는 정확도와 수렴 속도를 보인다.

상세 분석

RI‑FedAvg은 기존 연합 학습에서 가장 큰 난제 중 하나인 비IID(Non‑IID) 데이터로 인한 클라이언트 드리프트를 손실 함수의 “거칠기”라는 새로운 메트릭을 통해 정량화하고, 이를 정규화 강도로 매핑함으로써 해결한다. Roughness Index는 무작위 방향으로 손실 함수를 1차원으로 투사한 뒤, 총 변동량(Total Variation)을 정규화하여 표준편차와 평균의 비율로 정의한다. 높은 RI 값은 손실 지형이 급격히 변동하고 최적화가 어려움을 의미한다. 논문은 이 RI를 λ·I_k 형태의 가중치와 함께 ‖w−w_t‖² 정규화 항에 곱해 로컬 손실에 추가함으로써, 거칠기가 큰 클라이언트는 더 강하게 전역 모델에 묶이고, 거칠기가 작은 클라이언트는 자유롭게 학습하도록 설계하였다.

이 설계는 두 가지 중요한 장점을 제공한다. 첫째, FedProx와 달리 정규화 강도가 클라이언트마다 자동으로 조정되므로, 데이터 이질성이 심한 경우에도 과도한 제약이나 과소 제약을 방지한다. 둘째, SCAFFOLD가 필요로 하는 추가 메모리와 제어변수 전송을 요구하지 않으며, 통신 오버헤드가 동일하게 유지된다.

이론적 분석에서는 L‑smoothness, bounded variance, bounded non‑IID dissimilarity 등 기존 연합 학습 문헌에서 채택한 가정을 그대로 적용하면서, RI 기반 정규화가 추가된 로컬 손실에 대해 기대값과 변동을 엄격히 바인딩한다. 주요 정리는 모든 클라이언트가 동일한 λ·I_k 정규화를 사용하더라도, 전체 기대 손실 함수 f(w)의 그라디언트 노름이 점차 감소하여 결국 stationary point에 수렴한다는 것이다. 증명 과정에서 RI 자체가 확률적 추정치임을 고려해, M개의 무작위 방향과 m개의 샘플링 포인트에 대한 편차가 충분히 작을 경우 RI가 충분히 정확하게 추정된다고 가정한다.

실험 부분에서는 비IID 파티션을 적용한 MNIST(10클래스, 2‑5개 클라이언트), CIFAR‑10(10클래스, 20개 클라이언트), CIFAR‑100(100클래스, 30개 클라이언트)에서 학습을 수행하였다. 결과는 RI‑FedAvg이 평균 정확도에서 FedAvg 대비 35%p 향상, 수렴 에폭에서는 3040% 감소를 보였으며, 특히 데이터 불균형이 심한 경우(예: 라벨 비율 1:10)에도 안정적인 성능을 유지한다. 또한, 정규화 파라미터 λ와 RI 추정 파라미터(M, m, l)의 민감도 분석을 통해 λ=0.1, M=10, m=100, l=0.01이 대부분의 시나리오에서 최적임을 확인하였다.

한계점으로는 RI 계산 비용이 O(M·m·d)이며, 모델 차원이 매우 큰 경우(예: 대규모 Transformer)에는 추가 연산 부담이 발생한다는 점이다. 또한, RI가 손실 지형의 “거칠기”만을 포착하므로, 평탄하지만 최적점이 멀리 있는 경우에는 정규화가 충분히 효과적이지 않을 수 있다. 향후 연구에서는 RI를 다중 스케일로 확장하거나, 메타러닝을 통해 λ와 RI 매핑 함수를 자동 튜닝하는 방안을 제시한다.

전반적으로 RI‑FedAvg은 손실 지형 정보를 활용한 적응형 정규화라는 새로운 패러다임을 제시하며, 비IID 연합 학습에서의 클라이언트 드리프트 완화와 수렴 가속에 실질적인 기여를 한다.