기술 요약: Fed-BAC

문제 정의

계층적 연방 학습 (HFL) 은 엣지 서버를 활용하여 부분적 집계를 수행함으로써 다중 액세스 엣지 컴퓨팅 (MEC) 환경에서의 통신 병목 현상을 완화합니다. 그러나 기존 HFL 방법들은 실제 세계의 데이터 이질성을 다룰 때 세 가지 중요한 한계에 직면합니다:

1. 정적 가정: 대부분의 방법은 엣지 서버 간에 동질적인 데이터 분포를 가정하여, 레이블 분포가 다른 클라이언트들로부터의 업데이트를 집계할 때 "클라이언트 드리프트"와 정확도 저하 (최대 20%) 를 초래합니다.
2. 고립된 클러스터링: 클러스터링된 FL 접근법 (예: IFCA) 은 유사한 클라이언트들을 그룹화하지만, 일반적으로 완전히 고립된 모델을 훈련시켜 클러스터 간 지식 이전을 방해합니다.
3. 비효율적인 선택: 클라이언트 선택은 종종 무작위적이거나 정적입니다. 기존 밴딧 기반 접근법들은 평평한 아키텍처에서 작동하며, 3 계층 토폴로지 (클라우드 - 엣지 - 클라이언트) 전반에 걸쳐 클러스터 할당과 클라이언트 선택을 공동으로 최적화하지 못합니다.

또한, 클라이언트별 잔차를 유지하는 개인화 방법들 (예: Ditto, APFL) 은 $O(N)$의 파라미터 오버헤드를 발생시켜 대규모 환경에서는 비실용적이며 3 계층 시스템에서 통신 비용을 악화시킵니다.

방법론: Fed-BAC 프레임워크

본 논문은 가법적 클러스터 개인화와 2 단계 밴딧 의사결정 아키텍처를 통합한 Fed-BAC(Federated Bandit-Guided Additive Clustering) 프레임워크를 제안합니다.

1. 가법적 클러스터 개인화

고립된 클러스터 모델을 훈련하는 대신, Fed-BAC 는 로짓 공간에서 가법적 분해를 사용합니다. 클러스터 $k$의 클라이언트에 대한 예측 로짓 $\hat{y}$는 다음과 같이 정의됩니다:
$$\hat{y} = h(x; \Theta_{global}) + f(x; \Theta_k)$$

• $\Theta_{global}$: 모든 서버 간의 공통 지식을 포착하는 공유 글로벌 네트워크.
• $\Theta_k$: 분포 변이를 포착하는 클러스터별 네트워크.
  이 아키텍처는 모든 $M$개의 엣지 서버가 매 라운드 글로벌 모델에 기여할 수 있게 하여, 클러스터별 잔차를 통한 전문성을 유지하면서도 클러스터 간 지식 이전을 가능하게 합니다.

2. 2 단계 밴딧 프레임워크

Fed-BAC 는 3 계층 토폴로지를 해결하기 위해 의사결정을 두 단계로 분리합니다:

• 클라우드 레벨 (클러스터 할당): Contextual LinUCB를 사용하여 엣지 서버를 클러스터에 동적으로 할당합니다.

  • 맥락적 특징: 밴딧은 네 가지 FL 특화 특징을 활용합니다: 손실 비율 (현재 대안 클러스터 대비), 클러스터 균형, 할당 안정성, 그리고 훈련 진행 상황.
  • 메커니즘: 매 $\tau_{re}$ 라운드마다 클라우드는 각 서버가 현재 클러스터에 적합한지 대안 클러스터에 적합한지 평가합니다. 서버의 손실 비율이 다른 곳에서 더 적합함을 나타내면 할당이 업데이트됩니다. 이는 "무리 행동"을 방지하고 시스템이 최적의 활성 클러스터 수 ($K \le M$) 를 발견할 수 있게 합니다.

• 엣지 레벨 (클라이언트 선택): 각 엣지 서버에서 **Thompson Sampling (TS)**을 사용하여 참여할 클라이언트 부분집합 ($N_m$개 중 $B$개) 을 선택합니다.

  • 메커니즘: 각 서버는 모든 클라이언트에 대해 베타 분포 사전 확률을 유지합니다. 서버의 클러스터 모델의 단일 라운드 정확도 변화 ($r_{TS}$) 를 기반으로 사후 확률이 업데이트됩니다.
  • 목표: 고기여 클라이언트를 동적으로 식별하여 통신 오버헤드를 줄이면서도 모델 성능을 유지합니다.

3. 훈련 프로토콜

시스템은 다음을 포함하는 동기식 라운드로 운영됩니다:

1. 배포: 클라우드는 $\Theta_{global}$와 특정 $\Theta_k$를 각 엣지 서버로 전송합니다.
2. 선택 및 훈련: 엣지 서버는 TS 를 사용하여 클라이언트를 선택하고, 클라이언트들은 두 네트워크 모두에서 로컬 SGD 를 수행합니다.
3. 집계:
   • 엣지 집계: 클라이언트 업데이트의 가중 평균.
   • 클라우드 집계: 2 단계 업데이트. 첫째, 모든 서버로부터의 $\Theta_{global}$에 대한 글로벌 집계 (선택된 부분집합이 아닌 총 서버 데이터 크기를 사용). 둘째, 각 활성 클러스터에 대한 $\Theta_k$의 클러스터별 집계.

주요 기여

1. 2 단계 밴딧 아키텍처: 클러스터 할당 (클라우드의 LinUCB) 과 클라이언트 선택 (엣지의 Thompson Sampling) 을 분리하는 새로운 계층적 의사결정 시스템. 이는 각 의사결정의 시간 척도에 적합한 서로 다른 업데이트 주기를 허용합니다.
2. 통합된 가법적 개인화: 3 계층 HFL 아키텍처 내에서 가법적 클러스터 개인화를 통합한 최초의 사례. 클러스터를 고립시키는 IFCA 와 달리, Fed-BAC 는 글로벌 백본을 공유하여 그룹 간 지식 이전을 촉진합니다.
3. 공동 최적화: 기존 방법들이 이러한 작업을 별도로 처리하거나 정적 구성을 가정하는 격차를 해결하기 위해 클러스터 할당과 클라이언트 선택을 공동으로 최적화합니다.

실험 결과

실험은 중등도 ($\alpha=0.5$) 와 심한 ($\alpha=0.1$) 디리클레 비-IID 분할 하에서 세 가지 벤치마크 (CIFAR-10, SVHN, Fashion-MNIST) 에서 수행되었습니다.

• 정확도 향상: Fed-BAC 는 HierFAVG 대비 최대 +35.5 퍼센트 포인트, IFCA 대비 +8.4 퍼센트 포인트의 분산 정확도 향상을 달성했습니다.
• 이질성 민감도: Fed-BAC 의 성능 이점은 데이터 이질성의 심화 정도에 따라 증가했습니다. 심한 이질성 ($\alpha=0.1$) 하에서 HierFAVG 는 크게 저하된 반면, Fed-BAC 는 높은 정확도를 유지했습니다.
• 효율성: 80% 클라이언트 참여($p=0.8$) 만 필요함에도 불구하고, Fed-BAC 는 데이터셋과 목표 정확도에 따라 기준선보다 1.5~4.8 배 빠르게 수렴했습니다. 이 속도 향상은 대표성 있는 정확도 목표 달성을 위한 총 통신 바이트를 3.0 배에서 5.9 배까지 감소시켰습니다.
• 확장성: 500 개 클라이언트 (5 배 확장) 를 사용한 CIFAR-10 에 대한 확장성 연구는 Fed-BAC 가 HierFAVG 대비 (+12.6pp) 와 IFCA 대비 (+3.3pp) 우위를 유지함을 확인하여, 더 큰 배포 규모에서의 견고성을 검증했습니다.
• 공정성: Fed-BAC 는 특히 심한 이질성 하에서 IFCA 와 HierFAVG 에 비해 서버별 정확도 분산이 낮아 개선된 서버 간 공평성을 보여주었습니다.

중요성과 주장

본 논문은 Fed-BAC 가 계층적 연방 학습에서 통계적 이질성과 통신 효율성의 중요한 과제를 해결한다고 주장합니다. 가법적 분해(지식 공유 가능) 와 밴딧 기반 적응적 의사결정(클러스터 구조와 클라이언트 참여 모두 최적화) 을 결합함으로써, 이 프레임워크는 기존 최첨단 방법들보다 우수한 성능을 발휘합니다.

저자들은 데이터 분포가 갈라질수록 가법적 개인화 접근법의 가치가 커진다고 강조하며, 클러스터별 잔차를 유지하면서 글로벌 백본을 공유하는 것이 완전히 고립된 클러스터링보다 우월함을 확인했습니다. 결과는 Fed-BAC 가 데이터 이질성이 심하고 통신 자원이 제한된 MEC 배포를 위한 확장 가능한 솔루션을 제공함을 시사합니다.

논문은 현재 평가가 단일 참여율을 사용한다는 점을 겸손하게 언급하며, 향후 작업은 참여율의 체계적 스윕, 실제 세계 엣지 하드웨어 테스트베드, 그리고 결합된 비정상 밴딧 시스템에 대한 공식 수렴 보장을 탐구해야 한다고 결론지었습니다.