연합 균형 학습

초록

연합 학습에서 비IID 데이터로 인한 클라이언트 드리프트를 근본적으로 막기 위해, 각 클라이언트가 고정된 샘플 수 안에서 클래스별 균형을 맞추도록 엣지‑사이드 생성 모델을 활용한다. 부족한 클래스는 합성 이미지로 채우고(knowledge filling), 과잉 클래스는 손실 기반 샘플링으로 줄이며(knowledge sampling), 합성 데이터와 실제 데이터 간 도메인 차이를 완화하기 위해 클래스별 학습 가능한 임베딩(knowledge alignment)과 임베딩 드롭(knowledge drop) 전략을 도입한다. 다양한 멀티미디어 데이터셋에서 기존 연합 학습 방법들을 크게 능가한다.

상세 분석

본 논문은 연합 학습(Federated Learning, FL)에서 가장 심각한 문제 중 하나인 비IID 데이터 분포가 초래하는 클라이언트 드리프트를, 기존 연구가 사후에 손실 함수나 그래디언트를 조정해 보완하는 방식과 달리, 클라이언트 측 데이터 자체를 사전에 균형 잡는 접근으로 전환한다. 핵심 아이디어는 “고정된 샘플 수”라는 현실적 제약 하에, 각 클라이언트가 보유한 클래스별 샘플 수를 평균 균형점 B_k (클래스 수 평균)와 비교해 데이터‑과잉, 데이터‑부족, 데이터‑결손 세 그룹으로 분류한다. 데이터‑부족·결손 클래스는 엣지 디바이스에 탑재 가능한 경량화 Stable Diffusion 등 생성 모델을 이용해 합성 이미지를 생성하고, 이를 “knowledge filling”이라 부른다. 반면 데이터‑과잉 클래스는 무작위 샘플링이 아닌, 각 샘플의 손실값을 기준으로 높은 손실을 가진 샘플을 우선 보존하는 “loss‑based knowledge sampling”을 적용해 균형점까지 다운샘플링한다.

합성 데이터와 실제 데이터 사이의 도메인 격차를 해소하기 위해, 논문은 클래스별 학습 가능한 임베딩 P_i 를 도입한다. 합성 이미지 I_ij 를 특징 추출기 F 에 통과시킨 후 얻은 프로토타입에 P_i 를 요소별 더해 분류 헤드 H 에 입력함으로써, 합성 이미지가 실제 데이터와 유사한 표현을 갖도록 유도한다(knowledge alignment). 과적합 방지를 위해 일정 비율 ζ 의 샘플에 대해 P_i 를 무작위로 제외하는 knowledge drop 전략을 추가한다.

또한, 클라이언트마다 연산 자원이 다를 수 있음을 고려해, 연산량이 풍부한 클라이언트는 더 많은 합성 데이터를 생성하고, 제한된 클라이언트는 기존 샘플을 중심으로 샘플링만 수행하도록 유연한 프레임워크를 설계하였다. 실험에서는 이미지·비디오·위성 데이터 등 멀티모달 데이터셋을 사용해 FedProx, SCAFFOLD, MOON, FedGen 등 최신 SOTA와 비교했으며, 전반적으로 2~7%p 이상의 정확도 향상을 기록하였다. 특히, 데이터‑부족 클래스가 심하게 편중된 시나리오에서 합성 데이터가 모델의 일반화 능력을 크게 끌어올리는 것이 확인되었다.

이러한 설계는 프라이버시 위험을 최소화하면서도, 엣지 디바이스의 점진적인 연산 능력 향상을 활용해 연합 학습의 근본적인 비IID 문제를 데이터 레벨에서 사전에 완화한다는 점에서 의미가 크다.