PaperlayerView exact revision on arXiv

Source-linked AI summary

Federated Knowledge Distillation

Hyowoon Seo, Jihong Park, Seungeun Oh, Mehdi Bennis, Seong-Lyun Kim

arXiv:2011.02367v1cs.LGcs.DCcs.ITcs.NImath.IT

TL;DR

Co-distillation requires workers to observe common samples for each loss calculation, creating communication and privacy concerns. This chapter analyzes KD and CD with NTK and studies FD applications, reporting CD convergence and higher accuracy than FL in selected applications.

  • Problem

    Online KD/CD requires workers to observe common samples for each loss calculation, prompting sample exchanges that raise communication and privacy concerns.

    Source: §1.2 · p. 9

  • Method

    The chapter uses NTK to analyze KD and CD, then studies FD applications using wireless-channel characteristics and proxy datasets, including reinforcement learning.

    Source: §1.5 · p. 27

  • Results

    CD asymptotically reaches ground-truth predictions even with two workers and converges faster with more workers; selected FD applications report higher accuracy than FL.

    Source: §1.5 · p. 27

  • Takeaways & Limitations

    The chapter concludes that FD remains communication-efficient and can achieve higher accuracy than FL in its advanced applications.

    Source: §1.5 · p. 27

    Under policy distillation’s described operations, agents and an honest-but-curious server may inspect visited states and actions, while global experience memory grows with the number of agents.

    Source: §1.4.1 · p. 22

Abstract

from arXiv · show

Distributed learning frameworks often rely on exchanging model parameters across workers, instead of revealing their raw data. A prime example is federated learning that exchanges the gradients or weights of each neural network model. Under limited communication resources, however, such a method becomes extremely costly particularly for modern deep neural networks having a huge number of model parameters. In this regard, federated distillation (FD) is a compelling distributed learning solution that only exchanges the model outputs whose dimensions are commonly much smaller than the model sizes (e.g., 10 labels in the MNIST dataset). The goal of this chapter is to provide a deep understanding of FD while demonstrating its communication efficiency and applicability to a variety of tasks. To this end, towards demystifying the operational principle of FD, the first part of this chapter provides a novel asymptotic analysis for two foundational algorithms of FD, namely knowledge distillation (KD) and co-distillation (CD), by exploiting the theory of neural tangent kernel (NTK). Next, the second part elaborates on a baseline implementation of FD for a classification task, and illustrates its performance in terms of accuracy and communication efficiency compared to FL. Lastly, to demonstrate the applicability of FD to various distributed learning tasks and environments, the third part presents two selected applications, namely FD over asymmetric uplink-and-downlink wireless channels and FD for reinforcement learning.

1.1 Preliminaries: Knowledge Distillation and Co-Distillation

The chapter frames FD's foundations through KD, which transfers a pre-trained teacher's knowledge, and CD, which uses other students' ensemble predictions. Its NTK analysis finds CD converges asymptotically to ground truth, even with two workers, and that more workers accelerate convergence.

1.1.1 Knowledge Distillation

  • KD trains a student to match a pre-trained teacher’s logits alongside its task loss, rather than only fitting one-hot labels.The NTK analysis uses mean squared error for mathematical tractability, while the rest of the chapter uses cross entropy.Source: §1.1.1 · pp. 3–4
  • In the kernel regime, the student’s prediction error decreases as the pre-trained teacher’s prediction approaches ground truth.Source: §1.1.1 · p. 6

1.1.2 Co-Distillation

  • CD replaces KD’s pre-trained teacher with an ensemble of the other workers’ logits, making each worker’s teacher knowledge dependent on the other workers.Source: §1.1.2 · p. 6
  • With C=100, numerical simulations show convergence in one communication round; increasing C also speeds convergence and reduces variance.Source: Figure 1.4 · p. 8, §1.1.2 · p. 8
  • CD converges asymptotically to ground-truth predictions even with two workers; more workers accelerate convergence and can enable one-shot CD as C→∞.The chapter contrasts this with KD, where zero prediction error requires an ideally pre-trained teacher.Source: §1.1.2 · p. 8

1.2 Federated Distillation

FD addresses co-distillation’s dependence on shared sample observations by grouping samples by label and exchanging averaged logits. In classification, workers aggregate logits locally, receive global label-specific averages, and use them as teacher targets during local training.

1.2 Federated Distillation

  • FD avoids co-distillation’s shared-sample requirement by grouping samples according to labels, enabling knowledge exchange without common sample observations.The shared-observation requirement was identified as a source of extensive sample exchanges that may also violate local data privacy.Source: §1.2.1 · p. 9, §1.2 · p. 9

1.2.1 Federated Distillation for Classification

  • Workers upload local-average logits by label; the server averages them per label and sends the resulting global-average logits back to each worker.During local KD, each worker uses the downloaded global-average logit associated with its current sample’s ground-truth label as the teacher target.Source: §1.2.1 · p. 9
  • The classification implementation uses a quadratic loss for both task training and distillation, with a constant learning rate and a weighting constant for the distillation regularizer.The quadratic loss can be replaced by another well-defined loss, such as cross-entropy.Source: §1.2.1 · p. 10
  • 4.3x faster convergence than FL with less than 10% accuracy compromise on MNIST, using a 5-layer convolutional NN and 2 workers.Across different worker counts, FD reduces communication payload by around 10,000x per round and total communication cost until convergence by over 40,000x compared with FL; accuracy compromise is particularly pronounced under non-IID data.Source: §1.2.1 · p. 10

1.2 Federated Distillation

Recent work considers FD in wireless environments and other tasks, while identifying accuracy trade-offs and approaches for asymmetric channels and proxy data. Figure 1.6 compares test accuracy and communication cost under IID and non-IID MNIST data.

1.2.2 Recent Progress and Future Direction

  • In wireless fading channels, FD’s reduced payload can mean more successful receptions or lower latency, and may yield higher accuracy than FL under low signal-to-noise ratio or bandwidth.The chapter identifies more realistic wireless settings and advanced physical-layer and multiple-access techniques as directions for further study.Source: §1.2.2 · p. 11
  • FD is more vulnerable than FL to non-IID data, and its higher communication efficiency can come at the cost of accuracy.When a worker lacks samples of a target class, the global average logits may rarely transfer that class’s knowledge into its local model.Source: §1.2.2 · p. 11
  • For asymmetric channels, one approach uses FL in the downlink and FD in the lower-capacity uplink; proxy-data approaches use a public dataset or mean samples per label.Figure 1.6 compares test accuracy and the sum communication cost of all workers per epoch under IID and non-IID MNIST data.Source: §1.2.2 · p. 12, Figure 1.6 · p. 12

1.3 Application: FD Under Uplink-Downlink Asymmetric Channels

Mix2FLD combines federated learning after distillation with Mixup and inverse-Mixup to address uplink–downlink asymmetry while preserving privacy and accuracy. Its evaluation reports accuracy, convergence, scalability, and privacy trade-offs across channel and data settings.

1.3.1 Baseline: MixFLD

  • MixFLD combines FD's local-average-logit uploads with FL's global-weight downloads, using server-side KD to convert the logits into a global model.Workers upload locally mixed seed samples before collection; the stated differential-privacy analysis guarantees privacy for sufficiently large local datasets under γ = 0.5 with additive noise.Source: §1.3.1 · pp. 14–16

1.3.2 Proposed: Mix2FLD

  • Mix2FLD adds inverse-Mixup across devices to convert mixed samples back to hard-labeled synthetic samples without restoring raw samples.Pairing samples with symmetric labels permits generation of more inversely mixed samples than uploaded mixed samples.Source: §1.3.2 · pp. 16–17

1.3.3 Numerical Evaluation and Discussions

  • Under asymmetric channels, Mix2FLD achieves up to 16.7% higher accuracy and 1.2 times faster convergence than FL; output uploads reduce uplink payload size by up to 42.4 times.The comparison is shown for asymmetric channels with IID and non-IID datasets.Source: §1.3.3 · p. 18
  • Increasing devices from 10 to 50 raises average test accuracy by 5.7% and cuts variance by 50% on IID data; non-IID data show a smaller accuracy gain with the same tendency.The authors conclude that Mix2FLD is scalable under both IID and non-IID distributions.Source: Figure 1.9 · p. 19, §1.3.3 · p. 19
  • On non-IID data, Mix2FLD reduces the unnecessary noise that MixFLD's mixed-up samples introduce, and achieves higher accuracy than FL under symmetric channels.The text attributes this to Mix2FLD incorporating the data distribution in global training.Source: §1.3.3 · p. 19
  • Increasing the number of generated inverse-Mixup samples improves accuracy by up to 1.7% at the same uploaded seed amount, without additional communication latency.The chapter also reports a privacy–accuracy trade-off: mixing closer to γ = 0.5 improves sample privacy but may compromise accuracy.Source: §1.3.3 · p. 20, Table 1.1 · p. 20, Table 1.2 · p. 20

1.4 Application: FD for Reinforcement Learning

The chapter extends federated distillation to reinforcement learning through proxy experience memories, addressing privacy and scaling concerns of policy distillation. In CartPole, FRD matches PD’s mission completion time with lower communication cost, while trading some speed against FRL.

1.4 Application: FD for Reinforcement Learning

  • FRD adapts federated distillation beyond classification by using proxy experience memories for distributed reinforcement learning with interactive agents.The chapter presents it as a communication-efficient, privacy-preserving alternative to exchanging raw experience memories.Source: §1.4 · p. 21

1.4.1 Policy Distillation With Experience Memory

  • Policy distillation uploads agents’ experience memories for server aggregation and redistribution, but exposes visited states and actions and scales its global memory with agent count.Agents train local models by minimizing cross-entropy between their policies and the global experience-memory policy.Source: §1.4 · p. 21, §1.4.1 · p. 22

1.4.2 Federated Reinforcement Distillation With Proxy Experience Memory

  • FRD exchanges proxy experience memories containing representative states and averaged policies, reducing payloads through quantization while avoiding disclosure of raw experiences.Agents map observed states to proxy states, average policies within clusters, and fit local models to the downloaded global proxy memory.Source: §1.4.2 · pp. 22–24

1.4.3 Experiments and Discussions

  • The CartPole experiments cluster each agent’s four-dimensional state using 30 subdivisions per observation dimension and map raw states to the nearest proxy state.The evaluation compares mission completion time and communication cost for FRD, PD, and FRL.Source: §1.4.3 · p. 25
  • FRD completes the mission slightly slower than FRL, especially with few agents, but avoids FRL’s larger payloads above 100 neurons per layer.FRL also requires identical critic NN architectures, whereas FRD’s communication cost is bounded by the number of state clusters and does not constrain NN architecture.Source: §1.4.3 · p. 25
  • For 2 agents, FRD saves around 50% of communication cost compared with PD while achieving mission completion time as fast as PD, as shown in Figure 1.12.Source: §1.4.3 · p. 25

1.5 Conclusion

The conclusion presents FD as an output-exchange alternative to parameter-exchange FL and summarizes theoretical analysis and applications. It identifies unsupervised and self-supervised learning under realistic wireless conditions and changing topologies as future research directions.

1.5 Conclusion

  • The chapter’s NTK analysis shows CD converges asymptotically to ground-truth prediction with two workers, with more workers accelerating convergence.The chapter also reports that its vanilla FD implementation falls short of CD’s upper-bound accuracy, while advanced applications can achieve higher accuracy than FL.Source: §1.5 · p. 27
  • Future work could study FD for unsupervised and self-supervised learning under more realistic wireless channels and time-varying network topologies.Source: §1.5 · p. 27

Related papers

Loading 2011.02367v1…