[Paper Review] Group Knowledge Transfer: Federated Learning of Large CNNs at the Edge

Chaoyang He, Murali Annavaram+1 more|arXiv (Cornell University)|Jul 28, 2020

Privacy-Preserving Technologies in DataComputer Science77 references190 citations

TL;DR

This paper reformulates federated learning for edge devices as FedGKT, a group knowledge transfer framework that trains small edge CNNs and synchronizes with a large server CNN via knowledge distillation, reducing edge compute and communication while achieving accuracy comparable to FedAvg.

ABSTRACT

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).

Motivation & Objective

• Motivate training large CNNs on resource-constrained edge devices without centralizing data.
• Propose FedGKT to shift computation to the server while keeping edge workloads light.
• Achieve comparable accuracy to FedAvg with significantly reduced edge FLOPs and parameters.
• Explore asynchronous training and knowledge distillation across many edge clients.

Proposed method

• Reformulate FL as an alternating minimization problem between an edge model and a server model.
• Use a small edge feature extractor plus a local classifier on each client and a large server CNN on the cloud.
• Introduce bidirectional knowledge distillation losses between edge and server models (KL/divergence terms) to stabilize training.
• Train server and edge components in alternating steps, exchanging hidden features and server logits.
• Employ an asynchronous training variant where the server begins training upon receiving inputs from any client.

Experimental results

Research questions

• RQ1 Can FedGKT match FedAvg accuracy while dramatically reducing edge computation and parameter counts?
• RQ2 Does bidirectional knowledge transfer improve convergence and performance on non-IID edge data?
• RQ3 Is asynchronous FedGKT training robust and time-efficient compared to synchronous approaches?
• RQ4 What are the trade-offs in communication, privacy, and scalability for FedGKT in practice?

Key findings

• FedGKT achieves comparable or slightly higher accuracy than FedAvg on CIFAR-10, CIFAR-100, and CINIC-10 across IID and non-IID data.
• Edge computation with FedGKT is 9 to 17 times lower (FLOPs) and edge CNN parameters are 54 to 105 times fewer than FedAvg baselines.
• Asynchronous FedGKT training does not degrade accuracy relative to synchronous training.
• Bidirectional knowledge transfer (server-to-edge and edge-to-server) improves performance, especially on more challenging non-IID datasets.
• FedGKT reduces edge memory and communication bandwidth by exchanging hidden features rather than full models.