[논문 리뷰] FedZKT: Zero-Shot Knowledge Transfer towards Resource-Constrained Federated Learning with Heterogeneous On-Device Models

Lan Zhang, Dapeng Wu외 1명|arXiv (Cornell University)|2021. 09. 08.

Privacy-Preserving Technologies in Data인용 수 8

한 줄 요약

FedZKT는 자원이 제한된 이질적인 디바이스 모델을 가진 디바이스들이 원시 데이터를 공유하지 않고도 공동으로 글로벌 모델을 훈련시킬 수 있는 데이터 프리, 제로샷 지식 전이 프레임워크를 제안한다. 계산이 집약적인 디스틸레이션을 서버로 이관하고, 앙상블된 디바이스 모델들로부터 중심 지식을 학습하는 생성자(generator)를 사용함으로써, FedZKT는 모델 이질성, 비IIDs 데이터, 느린 작업자(straggler) 효과 등에 대해 뛰어난 성능을 보이며, 디바이스 측의 계산 부담을 최소화하고 사전에 개인 데이터가 필요로 하지 않는다.

ABSTRACT

Federated learning enables multiple distributed devices to collaboratively learn a shared prediction model without centralizing their on-device data. Most of the current algorithms require comparable individual efforts for local training with the same structure and size of on-device models, which, however, impedes participation from resource-constrained devices. Given the widespread yet heterogeneous devices nowadays, in this paper, we propose an innovative federated learning framework with heterogeneous on-device models through Zero-shot Knowledge Transfer, named by FedZKT. Specifically, FedZKT allows devices to independently determine the on-device models upon their local resources. To achieve knowledge transfer across these heterogeneous on-device models, a zero-shot distillation approach is designed without any prerequisites for private on-device data, which is contrary to certain prior research based on a public dataset or a pre-trained data generator. Moreover, this compute-intensive distillation task is assigned to the server to allow the participation of resource-constrained devices, where a generator is adversarially learned with the ensemble of collected on-device models. The distilled central knowledge is then sent back in the form of the corresponding on-device model parameters, which can be easily absorbed on the device side. Extensive experimental studies demonstrate the effectiveness and robustness of FedZKT towards on-device knowledge agnostic, on-device model heterogeneity, and other challenging federated learning scenarios, such as heterogeneous on-device data and straggler effects.

연구 동기 및 목표

• 기존 방법이 동일한 모델 아키텍처를 요구하는 바에 비해, 자원이 제한되고 이질적인 디바이스 모델을 가진 디바이스들을 페더레이티드 러닝에 통합하는 데 도전하는 것.
• 지식 전이를 위해 개인 데이터나 공개 데이터셋에 의존하지 않고, 데이터 프리 협업을 가능하게 하기 위해 이를 제거하는 것.
• 강력한 디바이스 측 계산 부담을 줄이기 위해 계산이 집약적인 디스틸레이션 과정을 서버로 이관하는 것.
• 실제 페더레이티드 환경에서 모델 이질성, 비IIDs 데이터 분포, 느린 작업자 효과에 대한 강건성을 확보하는 것.
• 글로벌 모델 수렴성과 성능을 유지하면서도 디바이스 측에서 독립적인 모델 설계를 가능하게 하는 것.

제안 방법

• 디바이스들은 로컬 하드웨어 제약 조건에 따라 독립적으로 작고 자원 효율적인 모델을 훈련시킨다.
• 서버에서 학습된 생성자는 수집된 디바이스 모델의 앙상블과 함께 적대적으로 훈련되어, 개인 데이터가 필요 없이 중심 지식을 디스틸레이션한다.
• 구조적으로 이질적인 모델들 간의 제로샷 지식 디스틸레이션을 가능하게 하기 위해 새로운 SL(소프트 레이블) 손실 함수를 도입한다.
• 디스틸레이션된 지식은 디바이스 모델 파라미터로 다시 전송되어, 디바이스들이 쉽게 글로벌 지식을 수용하고 미세조정할 수 있도록 한다.
• 프레임워크는 계산이 집약적인 디스틸레이션을 서버에 할당하여 디바이스의 자원 효율성을 유지한다.
• L2 정규화는 비IIDs 데이터 분포 하에서의 강건성을 향상시키기 위해 디바이스 훈련 중 적용된다.

Figure 1 : Overview of FedZKT: zero-shot knowledge transfer for federated learning with heterogeneous on-device models. Most compute-intensive workloads of FedZKT are done at the server. Participating devices only need to update the on-device models based on their local data.

실험 결과

연구 질문

• RQ1 공유 데이터나 모델 아키텍처 정렬이 필요 없이, 매우 이질적인 디바이스 모델을 가진 디바이스들 간에 페더레이티드 러닝을 효과적으로 구현할 수 있는가?
• RQ2 공개 데이터셋이나 사전 학습된 데이터 생성기 없이, 개인 데이터에 의존하지 않고 이질적인 모델들 간에 지식을 전이할 수 있는가?
• RQ3 디스틸레이션 작업이 서버로 이관될 경우, 자원이 제한된 디바이스들이 페더레이티드 러닝에 얼마나 효과적으로 참여할 수 있는가?
• RQ4 실제 구현 환경에서 비IIDs 데이터 분포와 느린 작업자 효과에 대해 이 프레임워크의 강건성은 어느 정도인가?
• RQ5 디바이스 모델 아키텍처와 디바이스 수가 글로벌 모델의 수렴성과 정확도에 어떤 영향을 미치는가?

주요 결과

• FedZKT는 다양한 디바이스 모델에서 뛰어난 성능을 달성한다: CIFAR-10 (IID)에서 Model C를 사용하는 Device 3는 88.63% 정확도를 기록했고, Device 8 역시 88.34%를 기록하여 모델 이질성에 대한 높은 적응성을 입증한다.
• 각 라운드에 20%의 디바이스만 활성화된 상태(p=0.2)에서도 안정적인 훈련을 유지하여 느린 작업자 효과에 대한 강건성을 입증한다.
• L2 정규화는 비IIDs 데이터에서 성능 향상을 이끈다: C=5 불균형 상황에서 정확도는 56.58%에서 63.89%로 상승했고, β=0.5 분포 불균형 상황에서는 66.17%에서 69.39%로 향상되었다.
• 디바이스 수(K ∈ {5,10,15,20})는 성능에 미미한 영향을 미치며, 다양한 K 값에서 평균 정확도의 변동은 ±2% 이내이다.
• 디스틸레이션 과정이 서버 최적화되어 있고 지식 전달이 경량 모델 파라미터 업데이트 방식으로 이루어지므로, 자원이 극히 제한된 디바이스들도 효과적으로 기여할 수 있다.
• 프레임워크는 개인 데이터나 공개 데이터에 의존하지 않으며 효과적인 지식 전달을 달성함으로써, 이질적인 페더레이티드 환경에서 데이터 프리, 제로샷 디스틸레이션의 가능성과 타당성을 입증한다.

Figure 2 : Norm of gradients w.r.t input data (MNIST, IID). The gradients for the KL-divergence loss tend to vanish, while the gradients for the $\ell_{1}$ norm loss are much larger and unstable during the learning process. The proposed SL loss overcomes both problems in the federated learning.