WHERE to Generate Matters: Budget-Aware Synthetic Augmentation for Label Skewed Federated Learning

기술 요약: 생성 위치의 중요성: 레이블 왜곡된 연합 학습을 위한 예산 인식형 합성 증강

문제 정의

클라이언트가 로컬하게 불균형한 클래스 분포를 보유하는 연합 학습(FL) 환경에서 레이블 왜곡(Label skew)은 클라이언트 드리프트(Client drift)를 유발하고 글로벌 모델의 정확도를 저하시키는 근본적인 과제입니다. 기존의 합성 데이터 증강 방법들은 누락된 샘플을 생성하여 이를 교정하려고 시도하지만, 지배적인 정책인 **전체 클래스 균형화(Full Class Balancing)**는 모든 클래스가 클라이언트의 가장 큰 클래스 크기에 도달할 때까지 모든 클라이언트가 샘플을 생성하도록 요구합니다.

이러한 접근 방식은 계산적으로 매우 부담스럽습니다. 본 논문에서 입증되었듯이, 표준 CIFAR-10 파티션(클라이언트 20개, 높은 왜곡 $\alpha=0.1$)에서 완전한 균형을 달吻하기 위해서는 226,000개 이상의 합성 샘플을 생성해야 합니다. 또한, 기존의 예산 제한 방식들은 생성 자원을 균등하게 할당하거나 클라이언트 데이터 크기에 따라 할당하며, 구체적인 레이블 분포의 왜곡은 무시합니다. 저자들은 단순히 얼마나 많은 데이터를 생성할 것인가가 아니라, 고정된 예산 내에서 정확도를 극대화하기 위해 샘플이 어디로(WHERE) 가야 하는가가 핵심적인 질문이라고 주장합니다.

방법론: FedEAS

본 논문은 FedEas(Federated Entropy-Adaptive Synthesis)를 제안합니다. 이는 각 클라이언트의 로컬 레이블 분포에 따라 클래스별 생성 예산을 동적으로 할당하는 정책입니다. 전체 생성 예산을 사전에 고정하는 기존 방식과 달리, FedEAS는 개별 클라이언트 예산의 합으로부터 총 생성 예산을 도출합니다.

핵심 메커니즘

FedEAS는 클라이언트 드리프트 및 불균형 감소 분석에서 도출된 세 가지 주요 관찰 결과에 기반하여 작동합니다:

1. 불균형은 집중되어 있다: 매우 왜곡된 일부 서브셋의 클라이언트가 글로벌 불균형에 불균형하게 기여합니다.
2. 타겟팅된 생성이 효율적이다: 불균형은 결핍되거나 부재하는 클래스에 샘플이 추가될 때만 감소하며, 이미 풍부한 클래스에 샘플을 추가하는 것은 수익 체감 또는 부정적인 결과를 낳습니다.
3. 수익 체감: 분포가 더 균형 잡힐수록 불균형을 제거하는 비용은 증가합니다. 즉, 완전한 균형화는 상당한 정확도 이득을 얻기 위해 반드시 필요한 것은 아닙니다.

엔트로피 적응형 예산 규칙

각 클라이언트 $k$에 대해, FedEAS는 다음 공식을 사용하여 클래스별 생성 예산 $b_k$를 계산합니다:

$$b_k = \lfloor \sqrt{\bar{n}_k} \cdot (1 - \tilde{H}_k) \cdot \beta \rfloor$$

여기서:

• $\bar{n}_k = N_k / C$ 는 클라이언트 $k$의 클래스당 평균 개수입니다.
• $\tilde{H}_k$ 는 클라이언트의 로컬 레이블 분포에 대한 정규화된 샤논 엔트로피(0은 매우 왜곡됨, 1은 완벽하게 균형 잡힘을 의미)입니다.
• $\beta$ 는 모든 클라이언트가 공유하는 스칼라 예산 파라미터입니다.
• $(1 - \tilde{H}_k)$ 항은 왜곡이 심한 클라이언트에게 더 큰 예산을 할당하는 반면, 거의 균형 잡혔거나 IID인 클라이언트는 거의 할당하지 않도록 보장합니다(극한 상황에서 표준 FedAvg로 수렴).
• $\sqrt{\bar{n}_k}$ 항은 대규모 클라이언트에 과도하게 할당되는 것을 방지하기 위해 클라이언트 크기에 따라 예산을 하위 선형적으로 스케일링합니다.

생성 및 학습

$b_k$가 계산되면, 클라이언트는 로컬 카운트 $n_k^c < b_k$인 클래스 $c$에 대해서만 합성 샘플을 생성합니다. 클래스 $c$에 대해 생성되는 샘플 수는 $n_{k, \text{gen}}^c = \max\{0, b_k - n_k^c\}$입니다. 이 "임계값 채우기(fill-to-threshold)" 전략은 샘플이 오직 부재하거나 희소한 클래스로만 향하도록 보장합니다. 생성은 사전 학습된 클래스 조건부 확산 확률 모델(DDPM)을 사용하여 클라이언트당 한 번 수행되며, 생성된 샘플은 이후 모든 학습 라운드에서 사용하기 위해 캐싱됩니다.

주요 기여

1. 엔트로피 적응형 예산 할당: FedEAS는 로컬 레이블 왜곡에 기반하여 합성 데이터의 양과 분포("어디로")를 결정하는 정책을 도입합니다. 왜곡이 심한 클라이언트에 더 큰 예산을 할당하고 희소한 클래스만을 타겟팅합니다.
2. 생성 효율적 정확도: FedEAS는 전체 클래스 균형화의 정확도 이득 대부분을 회복하면서도 총 생성 예산을 94.1% 감소시킵니다 (예: CIFAR-10에서 226,783개 대신 13,437개 샘플 생성). 이는 단일 GPU 기준 엔드 투 엔드 실행 시간을 3.7배 단축하는 효과를 가져옵니다.
3. 할당의 우위: 동일한 총 생성 샘플을 사용하는 매치드 예산(matched-budget) 프로토콜 하에서, FedEAS는 CIFAR-10 및 CIFAR-100에서 균등(Uniform) 할당 방식보다 최대 18.82% 더 높은 성능을 보입니다. 또한 "결핍 전용(Missing-only)" 생성 정책보다 최대 2배 빠르게 수렴합니다.

실험 결과

저자들은 Dirichlet 레이블 왜곡($\alpha \in \{0.05, 0.1, 0.3\}$) 조건 하의 CIFAR-10 및 CIFAR-100에 대해 다양한 클라이언트 수($K=20, 100$)로 FedEAS를 평가했습니다.

• 정확도 vs 예산: CIFAR-10 ($\alpha=0.1, K=20$)에서 FedEAS는 **71.65%**의 정확도를 달성하여, 전체 균형화(79.91%)가 제공하는 이득의 2/3를 회복하면서 합성 샘플은 단 5.9%만 사용했습니다.
• 베이스라인 비교: FedEAS는 표준 FL 베이스라인(FedAvg, FedProx, SCAFFOLD, MOON) 및 증강 베이스라인(Uniform, Missing-only)보다 일관되게 우수한 성능을 보였습니다.
• 예산 민감도: 정확도는 예산 파라미터 $\beta$에 따라 단조 증가함을 보였습니다. 저자들은 $\beta=12$를 기본값으로 선택했으며, 이는 초기 불균형 점수의 최소 50%를 제거하면서 합성 샘플을 전체 학습 데이터의 25% 미만으로 유지한다는 두 가지 조건을 충족했습니다.
• 절제 연구(Ablation Studies): 엔트로피 요소를 제거하거나 선형 규칙으로 하위 선형 크기 스케일링을 대체할 경우 정확도가 낮아지고 생성 비용이 높아졌으며, 이는 예산 공식에서 두 구성 요소의 필요성을 확인시켜 줍니다.

의의 및 주장

본 논문은 레이블 왜곡된 FL에서 합성 증강의 **주요 설계 문제는 할당(allocation)**이라고 주장합니다. 저자들은 총 생성 예산이 사전에 결정된 고정된 제약 조건이 아니라, 로컬 분포에서 도출된 클라이언트별 필요에 따른 결과여야 한다고 주장합니다.

"전체 균형화"에서 "엔트로피 적응형 예산 책정"으로 초점을 전환함으로써, FedEAS는 최소한의 합성 데이터 생성만으로 상당한 정확도 향상을 달성할 수 있음을 입증했습니다. 이 방법은 생성기 독립적(generator-agnostic)이며(어떠한 클래스 조건부 생성기와도 호환 가능), 표준 FL 통신 프로토콜이나 집계 가중치에 변경을 요구하지 않아 이질적인 환경에서 클라이언트 드리프트를 완화하기 위한 실용적이고 효율적인 솔루션이 됩니다.