Group Knowledge Transfer: Federated Learning of Large CNNs at the Edge

Scaling up the convolutional neural network (CNN) size (e.g., width, depth, etc.) is known to effectively improve model accuracy. However, the large model size impedes training on resource-constrained edge devices. For instance, federated learning (FL) may place undue burden on the compute capability of edge nodes, even though there is a strong practical need for FL due to its privacy and confidentiality properties. To address the resource-constrained reality of edge devices, we reformulate FL as a group knowledge transfer training algorithm, called FedGKT. FedGKT designs a variant of the alternating minimization approach to train small CNNs on edge nodes and periodically transfer their knowledge by knowledge distillation to a large server-side CNN. FedGKT consolidates several advantages into a single framework: reduced demand for edge computation, lower communication bandwidth for large CNNs, and asynchronous training, all while maintaining model accuracy comparable to FedAvg. We train CNNs designed based on ResNet-56 and ResNet-110 using three distinct datasets (CIFAR-10, CIFAR-100, and CINIC-10) and their non-I.I.D. variants. Our results show that FedGKT can obtain comparable or even slightly higher accuracy than FedAvg. More importantly, FedGKT makes edge training affordable. Compared to the edge training using FedAvg, FedGKT demands 9 to 17 times less computational power (FLOPs) on edge devices and requires 54 to 105 times fewer parameters in the edge CNN. Our source code is released at FedML ( https://fedml.ai ). Related Works Federated Learning. Existing FL methods such as FedAvg [6], FedOpt [26], and FedMA [8] face significant hurdles in training large CNNs on resource-constrained devices. Recent works FedNAS [27, 3] and [28] work on large CNNs, but they rely on GPU training to complete the evaluations. Others [29, 30, 31, 32, 33, 34, 35, 36, 37] optimize the communication cost without considering edge computational limitations. Model parallelism-based split learning [11, 12] attempts to break the computational constraint, but it requires frequent communication with the server. Knowledge Distillation (KD). We use KD [19] in a different manner from existing and concurrent works [38, 39, 40, 41, 42, 43, 44, 45] . Previous works only consider transferring knowledge from a large network to a smaller one [19, 20, 21, 22] , or they transfer knowledge from a group, but each member in the group shares the same large model architecture or a large portion of the neural architecture with specific tail or head layers [46, 47, 48, 49, 50, 51] . Moreover, all teachers and students in distillation share the same dataset [50, 52, 53, 54] , while in our setting each member (client) can only access its own independent dataset. Previous methods use centralized training, but we utilize an alternating training method. Efficient On-device Deep Learning. Our work also relates to efficient deep learning on edge devices, such as model compression [55, 56, 57] , manually designed architectures (MobileNets [58], ShuffeNets [59], SqueezeNets [60]), or even efficient neural architecture search (EfficientNets [4], FBNet [61]). However, all of these techniques are tailored for the inference phase rather than the training phase.