[논문 리뷰] Towards Model Agnostic Federated Learning Using Knowledge Distillation

A. N. Afonin, Sai Praneeth Karimireddy|arXiv (Cornell University)|2021. 10. 28.

Privacy-Preserving Technologies in Data참고 문헌 35인용 수 4

한 줄 요약

이 논문은 원시 데이터를 공유하지 않고도 이질적인 모델과 데이터를 가진 에이전트 간의 협업을 가능하게 하는 지식 증류(KD)를 활용한 모델에 종속되지 않는 분산 학습 프레임워크를 제안한다. 이는 분산 커널 리지 회귀 기반의 이론적 프레임워크를 도입하여, 데이터 이질성 하에서 교차하는 KD는 성능을 떨어뜨리지만, 평균화된 KD(AvgKD)는 성능 저하를 피하고 실제 딥러닝 실험에서 이론적 예측에 가까운 성능을 보임을 밝혀낸다.

ABSTRACT

Is it possible to design an universal API for federated learning using which an ad-hoc group of data-holders (agents) collaborate with each other and perform federated learning? Such an API would necessarily need to be model-agnostic i.e. make no assumption about the model architecture being used by the agents, and also cannot rely on having representative public data at hand. Knowledge distillation (KD) is the obvious tool of choice to design such protocols. However, surprisingly, we show that most natural KD-based federated learning protocols have poor performance. To investigate this, we propose a new theoretical framework, Federated Kernel ridge regression, which can capture both model heterogeneity as well as data heterogeneity. Our analysis shows that the degradation is largely due to a fundamental limitation of knowledge distillation under data heterogeneity. We further validate our framework by analyzing and designing new protocols based on KD. Their performance on real world experiments using neural networks, though still unsatisfactory, closely matches our theoretical predictions.

연구 동기 및 목표

• 공유되는 모델 아키텍처나 공개 데이터가 필요 없이도 협업이 가능한 모델에 종속되지 않는 분산 학습 프rotocol를 설계하는 것.
• 특히 데이터 및 모델 이질성 하에서 지식 증류의 한계를 조사하는 것.
• KD 기반의 분산 학습 프로토콜의 행동을 분석하고 예측하기 위한 이론적 프레임워크를 개발하는 것.
• 이질적 환경에서 성능 저하를 완화하는 개선된 KD 기반 프로토콜를 설계하는 것.
• 실세계 데이터셋과 모델(MLP 및 랜덤 포레스트 포함)을 활용한 실증적 평가를 통해 이론적 통찰을 검증하는 것.

제안 방법

• 모델에 종속되지 않는 분산 학습을 병합된 데이터셋 위의 커널 리지 회귀로 공식화하여 기능 유사성과 이질성을 분석할 수 있도록 한다.
• 데이터 및 모델 이질성을 모두 모델링할 수 있는 이론적 프레임워크인 분산 커널 리지 회귀를 도입한다.
• 교차 지식 증류(AKD)를 볼록 집합 위의 교차 투영으로 해석하여 순차적 정보 손실을 드러낸다.
• 안정적인 대안으로 평균화된 지식 증류(AvgKD)를 제안하며, 성능 저하를 피함을 입증한다.
• MNIST 및 CIFAR10에서 실세계 실험을 통해 이론적 예측을 검증하며, 교차 엔트로피 및 MSE 손실 함수를 사용한다.
• M에이전트 환경으로 AvgKD를 확장하여, 다양한 데이터 이질성 수준에서 안정성과 조기 커뮤니케이션의 이점을 입증한다.

실험 결과

연구 질문

• RQ1 표준 지식 증류 기반의 분산 학습 프로토콜가 왜 데이터 이질성 하에서 성능이 떨어지는가?
• RQ2 공유되는 모델 아키텍처가 필요 없이 '피팅'과 '예측'만을 원천으로 사용하는 모델에 종속되지 않는 분산 학습 프로토콜를 설계할 수 있는가?
• RQ3 데이터 이질성이 분산 환경에서 지식 증류의 수렴성과 성능에 어떤 영향을 미치는가?
• RQ4 KD 기반의 분산 학습 프로토콜의 행동을 예측하고 분석할 수 있는 이론적 프레임워크를 개발할 수 있는가?
• RQ5 평균화된 지식 증류(AvgKD)가 이질적 데이터 및 모델 환경에서 안정성과 최종 정확도 측면에서 교차 지식 증류(AKD)를 능가하는가?

주요 결과

• 교차 지식 증류(AKD)는 교차 투영 과정에서 순차적 정보 손실로 인해 라운드가 진행될수록 성능이 떨어지며, 특히 높은 데이터 이질성 하에서 심화된다.
• 평균화된 지식 증류(AvgKD)는 테스트된 모든 설정에서 성능 저하가 없음을 보이며, MLP 및 랜덤 포레스트와 같은 다양한 모델에서도 동일하게 적용된다.
• MNIST 및 CIFAR10에서 AvgKD는 모든 라운드 동안 안정적인 성능 유지를 보이며, AKD는 일관되게 성능이 떨어지며, 특히 정규화 강도가 높거나 모델 이질성 또는 데이터 불균형이 심할수록 떨어짐이 가속화된다.
• 실증 결과는 분산 커널 리지 회귀 프레임워크로부터 도출된 이론적 예측과 매우 유사하게 일치하여, 이 프레임워크가 프로토콜 설계에 실용적 가치가 있음을 입증한다.
• 전체 데이터 이질성(Alpha=0)이 있는 M에이전트 환경에서 조기 정지가 성능 향상에 기여하며, AvgKD에서는 모든 에이전트가 조기 커뮤니케이션 라운드에서 이점을 얻는다.
• AvgKD에서 두 에이전트의 성능은 데이터 이질성에 크게 의존하지만, 라운드가 진행되면서 성능 저하가 발생하지 않으며, MLP와 랜덤 포레스트와 같이 근본적으로 다른 모델 간에도 동일하게 적용된다.