PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Chaoyang He, Murali Annavaram, Salman Avestimehr

arXiv:2007.14513v4cs.LGcs.CV

TL;DR

Large CNN training can exceed edge-device compute capacity, even where federated learning is needed to keep data decentralized. FedGKT alternates training of compact edge CNNs and a large server CNN through knowledge distillation; across the reported settings, accuracy is comparable to FedAvg while edge computation is substantially lower.

  • Problem

    Large CNNs can improve accuracy, but their training burdens resource-constrained edge devices, where federated learning is useful without centralizing device data.

    Source: §1 · p. 1

  • Method

    FedGKT alternates training compact edge CNNs and a large server CNN, periodically transferring knowledge between them through distillation.

    Source: §1 · p. 2, §3.3 · p. 5

  • Results

    Across I.I.D. and non-I.I.D. datasets, FedGKT achieves accuracy comparable to or better than FedAvg while requiring 9 to 17 times less computational power (FLOPs) on edge devices.

    Source: §1 · p. 2, §4.2 · p. 7

  • Takeaways & Limitations

    FedGKT makes edge training more affordable, with lower communication cost and support for asynchronous training.

    Source: §1 · p. 2, §6 · p. 9

    The potential for personalized client models is not demonstrated experimentally; the authors leave fine-tuning such models for future work.

    Source: §5 · p. 9

Abstract

from arXiv · show

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

1 Introduction

Large CNNs can improve accuracy, but their training demands exceed the capabilities of many edge devices. FedGKT trains compact edge CNNs and transfers their knowledge to a server model, reducing edge computation while maintaining accuracy comparable to FedAvg.

1 Introduction

  • Scaling CNN size can improve accuracy, but resource-constrained devices struggle to train large CNNs; FedAvg evaluations have focused on small CNNs or assumed ample client compute.Split learning offloads part of a model but requires multiple server-edge communication rounds for each mini-batch iteration.Source: §1 · pp. 1–2
  • FedGKT alternates local training of compact edge CNNs with periodic knowledge distillation to a large server CNN, then sends server predictions back to improve edge models.The final model combines the local feature extractor with the shared server model, shifting computing burden from edges to the server.Source: §1 · p. 2
  • 9 to 17 times less computational power (FLOPs) on edge devices and 54 to 105 times fewer parameters in the edge CNN are required by FedGKT than by FedAvg.The introduction reports these reductions alongside accuracy comparable to FedAvg.Source: §1 · p. 2
  • FedGKT combines local-SGD training, hidden-feature exchange instead of whole-model exchange, and asynchronous server training that can begin with any client's input.These features reduce communication frequency and bandwidth and avoid split learning's synchronization issue.Source: §1 · p. 2

2 Related Works

Prior work addresses federated optimization, knowledge distillation, and efficient on-device learning, but the paper distinguishes FedGKT from these approaches by its training setting and knowledge-transfer design.

2 Related Works

  • Prior federated methods including FedAvg [6], FedOpt [26], and FedMA [8] face challenges training large CNNs on constrained devices; other approaches rely on GPU training, optimize communication alone, or require frequent server contact.The cited approaches include FedNAS [27] [3] and [28], communication-focused works [29] [33] [34] [36] [37], and split learning [12].Source: §2 · p. 3
  • Unlike prior KD methods that transfer from a large network to a smaller one or use shared large architectures and datasets, FedGKT transfers across independent client datasets with alternating training.The passage cites prior large-to-small KD [19] [22], group methods [47] [49] [50] [51], and shared-dataset approaches [50] [52].Source: §2 · p. 3
  • On-device methods such as compression, MobileNets [58], ShuffeNets, and SqueezeNets [60] are described as targeting inference rather than training.The passage also lists EfficientNets [4] and FBNet as efficient architecture-search approaches tailored to inference.Source: §2 · p. 3

3 Group Knowledge Transfer

The section reframes federated CNN training around small edge models and a large server model, then trains them through alternating, bidirectional knowledge transfer.

3.1 Preliminary

  • The formulation targets supervised CNN training across K devices without centralizing data, while conventional FL methods struggle to train large CNNs on edge devices lacking GPUs and memory.Source: §3.1 · p. 3

3.2 Reformulation

  • The reformulation splits the global CNN into an edge feature extractor and classifier and a larger server model trained using edge feature maps.Source: §3.2 · p. 4
  • When edge models are much smaller than the server model, edge training is affordable, and transferring feature maps requires less bandwidth than exchanging all model parameters.Source: §3.2 · p. 4
  • The approach depends on clients adequately training their feature extractors, but small local datasets may be insufficient and the server optimization depends on feature quality.Source: §3.2 · p. 4

3.3 Group Knowledge Transfer (FedGKT)

  • FedGKT alternates server and edge optimization, using cross-entropy and KL-divergence losses to transfer knowledge from many edges to the server and back to each edge.Source: §3.3 · p. 5
  • In each round, clients train locally and send feature maps and logits to the server, which trains its larger CNN and returns global logits to the clients.Source: §3.3 · p. 6

4 Experiments

Across CIFAR-10, CIFAR-100, and CINIC-10, FedGKT achieves accuracy comparable to or better than FedAvg while reducing edge computation and communication costs. The experiments also examine knowledge transfer, asynchronous training, client counts, and smaller edge models.

4.1 Experimental Setup

  • Experiments use ResNet-56 and ResNet-110 baselines, with FedGKT pairing compact ResNet-8 edge models with ResNet-55 or ResNet-109 server models.The evaluation covers CIFAR-10, CIFAR-100, and CINIC-10, including non-I.I.D. variants; FedAvg and centralized training are baselines.Source: §4.1 · pp. 6–7

4.2 Result of Model Accuracy

  • FedGKT obtains comparable or better accuracy than FedAvg in both I.I.D. and non-I.I.D. settings across the evaluated datasets and ResNet models.Table 1 summarizes ResNet-56 and ResNet-110 test accuracy; Figure 3 plots ResNet-56 accuracy during training with 16 clients.Source: Table 1 · p. 7, §4.2 · p. 7, Figure 3 · p. 7

4.3 Efficiency Evaluation

  • On CIFAR-100, FedGKT ResNet-8 uses 9 times less edge computation than FedAvg ResNet-56 and 17 times less than FedAvg ResNet-110.ResNet-8 has 11K parameters, 54 times fewer than ResNet-56 and 105 times fewer than ResNet-110; on an Intel i7 it takes 30 ms per mini-batch versus 950 ms for ResNet-110.Source: §4.3 · p. 8, Figure 4 · p. 8
  • For ResNet-56, FedGKT uses fewer feature-map exchanges with the server than split learning.The communication-cost comparison does not use data compression techniques.Source: §4.3 · p. 8, Figure 5 · p. 8

4.4 Ablation Study: Understanding FedGKT under Different Settings

  • Without knowledge-distillation loss, client and server training with only cross-entropy can yield low accuracy or diverge; server-to-edge transfer helps, while bidirectional transfer is more effective on CIFAR-100.The bidirectional-transfer advantage is more pronounced as the dataset becomes more difficult.Source: §4.4 · p. 8, Table 2 · p. 8
  • Asynchronous training does not negatively affect accuracy, and the server can train without waiting for updates from all clients.This avoids the repeated synchronization required by split learning for each mini-batch iteration.Source: §4.4 · p. 8, Table 3 · p. 8
  • Adding more edge nodes generally does not negatively affect FedGKT test accuracy.The study evaluates accuracy with varying numbers of edge nodes.Source: §4.4 · p. 8, Table 4 · p. 8
  • Reducing edge models below ResNet-8 to ResNet-4 or ResNet-6 on CIFAR-10 reduces overall accuracy.ResNet-4 and ResNet-6 use one and two BasicBlock components, respectively.Source: §4.4 · p. 8, Table 5 · p. 8

5 Discussion

FedGKT’s discussion identifies limits in privacy analysis, communication cost, supervised-learning scope, scalability, and demonstrated personalization. These boundaries constrain what the method establishes across settings.

5 Discussion

  • FedGKT does not analyze or compare privacy leakage from exchanging gradients, models, and hidden maps.The authors note that hidden-map exchange may be safer, but that claim is not a comparative analysis.Source: §5 · p. 9
  • When sample counts and image resolution are both extremely large, FedGKT and split learning would have high total communication costs.FedGKT’s communication cost depends on the number of data points, despite lower bandwidth requirements than gradient or model exchange.Source: §5 · p. 9
  • FedGKT supports supervised learning only, while many applications lack sufficient labels.The authors connect label shortages to the difficulty of incentivizing users to label private local data.Source: §5 · p. 9
  • FedGKT does not claim an advantage in scalability, and uniform client sampling required many more training rounds to converge in the reported experiments.The discussion notes that scalability remains a common limitation across existing works.Source: §5 · p. 9
  • Personalized client-and-server model combinations are proposed as a possibility, but their effectiveness was not demonstrated in the experiments.The authors suggest fine-tuning client models for several epochs as a future exploration.Source: §5 · p. 9

6 Conclusion

FedGKT trains small edge CNNs and transfers their knowledge to a larger server-side CNN. Across the evaluated architectures and datasets, it achieves accuracy comparable to or slightly higher than FedAvg while substantially reducing edge resource demands.

6 Conclusion

  • FedGKT achieved comparable or slightly higher accuracy than FedAvg while using 9–17 times less edge computation and 54–105 times fewer edge-CNN parameters.The evaluation used ResNet-56 and ResNet-110 on CIFAR-10, CIFAR-100, CINIC-10, and their non-I.I.D. variants.Source: §6 · p. 9
  • FedGKT trains small CNNs on edges and periodically transfers their knowledge by distillation to a high-capacity server-side CNN.The framework is designed to reduce edge computation and communication cost while supporting asynchronous training.Source: §6 · p. 9

Broader Impact

FedGKT is presented as a way for resource-limited organizations and device users to collaboratively train large CNNs without centralizing private datasets. The authors also identify risks from malicious client updates and unbalanced benefits across clients.

Broader Impact

  • Organizations with limited training resources could collaboratively train large CNNs without revealing their datasets, while meeting confidentiality and related constraints.The authors also describe potential benefits for smartphone users seeking private data protection and model services.Source: Broader Impact · p. 10
  • Malicious clients may send incorrect hidden features and soft labels that could affect model accuracy, while the framework does not balance differing client interests.The discussion gives dataset size and fairness as an example of benefits varying across clients.Source: Broader Impact · p. 10

A.1 A Summary of Dataset Used in Experiments

The experiments use CIFAR-10, CIFAR-100, and CINIC-10, including non-I.I.D. settings. For heterogeneous comparisons, the authors fix a client distribution and report that FedGKT also outperforms baselines under other tested distributions.

A.1 A Summary of Dataset Used in Experiments

  • The experiments use CIFAR-10, CIFAR-100, and CINIC-10, with CINIC-10 combining ImageNet and CIFAR-10 sources whose distributions are not guaranteed to match.CINIC-10 contains 270,000 images; experiments train on its training split and test on its testing split, leaving validation unused.Source: §A.1 · p. 14
  • The non-I.I.D. partition is unbalanced, allocating each class’s training samples across clients according to sampled proportions.The passage specifies sampling those proportions from a Dirichlet distribution with parameter 0.5.Source: §A.1 · p. 14

A.2 Heterogeneous Distribution (non-I.I.D.) in Each Client

  • The authors fix a non-I.I.D. client distribution for fair method comparisons and report that FedGKT also outperforms baselines under other tested distributions.The cited table identifies the displayed example as a heterogeneous CIFAR-10 distribution.Source: §A.2 · p. 14, Table 6 · p. 14

B Extra Experimental Results and Details

This section presents figures on edge computational efficiency and specifies how communication cost is calculated for split learning and FedGKT.

B.1 Computational Efficiency on CIFAR-10 and CINIC-10

  • Figures 6 and 7 cover edge computational efficiency on CIFAR-100 and CINIC-10, respectively.The captions identify the datasets but do not state comparative outcomes.Source: Figure 6 · p. 14, Figure 7 · p. 14

B.2 The Method of Communication Cost Calculation

  • For split learning, communication cost is the hidden feature map size plus the gradient size in the split layer.Source: §B.2 · p. 14
  • For FedGKT, communication cost is the hidden feature map and received soft-label sizes multiplied by the dataset’s number of samples.Source: §B.2 · p. 15

B.3 Details of Convolutional Neural Architecture on Edge and Server

The appendix describes the compact ResNet-8 edge CNN and points to architecture and hyperparameter tables for the experiments. It recommends tuning hyperparameters when applying FedGKT to new CNN architectures and datasets.

B.3 Details of Convolutional Neural Architecture on Edge and Server

  • ResNet-8 uses its head convolutional layer, including batch normalization and ReLU, as the feature extractor; two Bottlenecks and the final fully connected layer form the classifier.Each Bottleneck contains three convolutional layers.Source: §B.3 · p. 15
  • Architecture details are tabulated for ResNet-8, ResNet-55, and ResNet-109.Source: Table 7 · p. 15, Table 8 · p. 15, Table 9 · p. 16

B.4 Hyperparameters

Related papers

Loading 2007.14513v4…