Class-Grouped Normalized Momentum and Faster Hyperparameter Exploration to Tackle Class Imbalance in Federated Learning

기술 요약: 연합 학습의 클래스 불균형을 위한 FedCGNM 및 FedHOO

문제 정의
클래스 불균형은 연합 학습(Federated Learning, FL)에서 중요한 과제로 남아 있으며, 특히 글로벌 라벨 분포가 롱테일(long-tailed) 형태를 띨 때 더욱 그러합니다. 이러한 시나리오에서 소수 클래스는 훈련 데이터에서의 과소 표현으로 인해 예측 성능이 저하됩니다. 가중치 재설정(loss reweighting), 고급 샘플링 또는 생성적 증강(generative augmentation)과 같은 표준 중앙 집중식 기법들은 연합 학습에 적용하기 어렵습니다. 개인정보 보호 제약으로 인해 데이터 교환이 불가능하며, 합성 생성기(synthetic generators)는 희소한 영역이나 도메인 특화 작업(예: 반도체 결함 탐지)에서 현실적인 샘플을 생성하는 데 실패하는 경우가 많습니다. 또한, 클래스별 정규화(Per-Class Normalization, PCN)와 같은 기존의 최적화 수준 접근 방식은 다중 클래스 설정에서 한계에 직면합니다. 즉, 소수 클래스로부터 발생하는 방향성 노이즈(directional noise)에 취약하며, 많은 단위 벡터를 합칠 때 스케일 불일치(scaling mismatch)를 유발하여 클래스 수가 증가함에 따라 수렴을 불안정하게 만듭니다.

방법론
본 논문은 이러한 문제를 해결하기 위해 클라이언트 측 옵티마이저인 FedCGNM(Federated Class-Grouped Normalized Momentum)과 보조 하이퍼파라미터 튜너인 FedHOO라는 두 갈래의 접근 방식을 제안합니다.

1. FedCGNM (옵티마이저):

   • 클래스 그룹화(Class Grouping): 모든 개별 클래스($C$)에 대해 그래디언트를 정규화하는 대신, FedCGNM은 클래스를 적은 수의 그룹($H \ll C$, 일반적으로 다수 클래스 대 소수 클래스)으로 분할합니다. 그룹화 규칙은 클래스 비율의 그룹 내 분산을 최소화하는 방식으로 도출되며, 이는 전수 조사(exhaustive search)를 대체할 수 있는 신뢰할 수 있는 프록시로 입증되었습니다.
   • 그룹별 정규화 모멘텀(Group-Wise Normalized Momentum): 각 그룹에 대해 알고리즘은 모멘텀 벡터를 유지합니다. 로컬 훈련 중에 그룹별로 확률적 그래디언트(stochastic gradients)가 계산되고 모멘텀이 업데이트됩니다. 결정적으로, 각 그룹의 모멘텀 벡터는 합산되기 전에 단위 길이(안정화 상수 $\delta$ 포함)로 정규화됩니다.
   • 업데이트 방향: 최종 업데이트 방향은 이러한 정규화된 그룹 모멘텀들의 합입니다. 이 설계는 다수 그룹과 소수 그룹 간의 그래디언트 크기를 균등하게 맞추는 동시에, 희귀 클래스 그래디언트에 내재된 방향성 노이즈를 완화합니다. 개별 클래스가 아닌 소수의 그룹 단위로 작동함으로써, PCN과 관련된 스케일 불일치 및 불안정성을 방지합니다.
   • 클라이언트 정렬(Client Alignment): 각 그룹의 업데이트가 단위 노름(unit norm)을 갖도록 강제함으로써, 이질적인 클라이언트들 간의 집계된 업데이트 방향이 더 잘 정렬되도록 하여 성능 저하의 주요 원인을 해결합니다.

2. FedHOO (하이퍼파라미터 튜너):

   • 맥락: 클라이언트 간의 조합 탐색 공간과 변동하는 스케줄로 인한 최적화 불안정성 위험 때문에, FL에서 적절한 리샘플링 비율을 선택하는 것은 쉽지 않습니다.
   • 메커니즘: FedHOO는 소규모 연합($K$가 작은 경우)을 위해 설계된 X-armed bandit(XAB) 기반 알고리즘입니다. 이는 FL의 선형적 집계 특성을 활용합니다. 각 통신 라운드에서 서버는 각 클라이언트에게 두 가지 후보 리샘플링 비율로 훈련하도록 지시합니다.
   • 효율성: 반환된 클라이언트 업데이트를 선형적으로 결합함으로써, 서버는 $2^K$번의 별도 훈련 라운드를 거치지 않고도 가능한 모든 $2^K$가지 로컬 비율의 조합에 대한 검증 성능을 추론할 수 있습니다. 이를 통해 개인정보를 유지하면서도(클라이언트는 스칼라 검증 지표만 반환) 훈련 초기 단계에서 효과적인 샘플링 비율을 빠르게 식별할 수 있습니다.

이론적 분석
본 논문은 시변(time-varying) 리샘플링 비율을 명시적으로 고려한 FedCGNM의 수렴 분석을 제공합니다. 매끄러움(smoothness), 유계 그래디언트(bounded gradients), 편향되지 않음(unbiasedness)에 대한 표준 가정을 바탕으로, 저자들은 FedCGNM이 표준적인 $O(T^{-1/2})$ 정지점 도달 속도를 달성함을 증명합니다. 분석 결과, 수렴이 성립하려면 라운드에 따른 샘플링 비율의 누적 변동량($V_T$)이 특정 조건을 만족해야 함을 보여주며, 이는 공격적이지만 불안정한 리샘플링이 수렴을 저해할 수 있음을 시사합니다.

주요 기여

1. FedCGNM 옵티마이저: 클래스를 그룹화하고 그룹당 단위 노름 모멘텀을 적용하는 최초의 클라이언트 측 옵티마이저입니다. 이는 다수/소수 클래스의 영향력을 균형 있게 조절하고, 노이즈를 줄이며, 스케일 불일치를 완화합니다.
2. 분산 인식 그룹화(Variance-Aware Grouping): 클래스 분할을 그룹 내 분산 최소화 문제로 프레이밍하여, 전수 조사와 유사한 최적의 분할을 생성하는 원칙적이고 경량화된 전략입니다.
3. 수렴 분석: FL, 그룹 정규화, 동적 샘플링 비율을 결합한 새로운 이론적 프레임워크를 제공하여 샘플링 스케줄이 수렴에 미치는 영향을 입증했습니다.
4. FedHOO: 연합 병렬성을 활용하여 효율적인 탐색을 수행하는 X-armed bandit 프레임워크를 통해, FL에서 조합적인 로컬 샘플링 비율을 결정하는 최초의 최적화 기반 알고리즘입니다.

실험 결과
제안된 방법들은 네 가지 공개 롱테일 벤치마크(CIFAR-10-LT, CIFAR-100-LT, Adult Income, UNSW-NB15)와 독자적인 반도체 칩 결함 데이터셋(CDD)에서 평가되었습니다.

• 성능: FedCGNM은 모든 설정에서 베이스라인(FedAvg, FedProx, CLIMB, FedGraB 및 PCN 변형 포함)을 일관되게 능가했습니다. 소규모 클라이언트 환경($K=5$)에서 FedCGNM과 FedHOO를 결합했을 때 가장 좋은 결과를 얻었으며, 리샘플링을 적용한 FedAvg 대비 최대 29%의 향상을 보였습니다.
• 산업적 응용: 결함률이 1.7%인 CDD 데이터셋에서 제안된 방법은 최상의 베이스라인 대비 16%의 개선을 달하여, 실제 개인정보 보호가 중요한 산업 현장에서의 효과를 입증했습니다.
• 강건성: 제안된 방식은 부분 참여가 발생하는 대규모 클라이언트 환경과 다양한 정도의 데이터 이질성(Non-IID) 하에서도 회복력을 보여주었습니다.
• 어블레이션(Ablation): 실험을 통해 두 개의 그룹($H=2$)이 최적의 편향-분산 트레이드오프를 제공하며, 중간 정도의 모멘텀 계수($\beta \approx 0.5$)가 최적의 성능을 낸다는 것을 확인했습니다.

의의 및 주장
본 논문은 FedCGNM과 FedHOO가 데이터 교환이나 합성 생성 없이도 개인정보가 보호되는 연합 환경에서 글로벌 클래스 불균형을 완화하기 위한 실용적인 프레임워크를 형성한다고 주장합니다. 본 연구의 의의는 클래스 불균형을 데이터 수준의 기법(개인정보 보호나 데이터셋 특성에 의해 제한될 수 있는)이 아닌 최적화 수준에서 해결한다는 점에 있습니다. 저자들은 제안된 방식이 특히 공동 하이퍼파라미터 튜닝이 가능하고 필수적인 소규모 연방(산업적 배포에서 흔함)에서 매우 효과적이며, 그룹화 전략이 기존의 정규화 기술이 실패하는 다중 클래스 시나리오에서 안정성과 정렬을 보장한다고 강조합니다. 이 연구는 롱테일 클래스 불균형이 존재하는 상황에서 이론적으로 근거가 있고 효율적인 샘pling 전략 튜닝 방법을 제공함으로써 해당 분야를 발전시켰습니다.