Knowledge Distillation in Federated Learning

In FKD, the student model (the final global model) and the client models (which can be teachers) are architecturally and functionally separate. Clients train their local models independently. The server never receives or aggregates these model parameters. Instead, it receives and aggregates knowledge, such as softened logits or embeddings, and uses this to train the student model via distillation loss. This breaks the direct parameter synchronization loop of standard Federated Averaging (FedAvg).

The core transmitted payload is not a gradient or parameter update, but a knowledge representation. Common forms include:

• Softened Logits: The output layer predictions (logits) softened by a high temperature parameter from the client's local model applied to a shared (or generated) calibration dataset.
• Embeddings: Intermediate feature representations from a specific layer of the client model.
• Label Distributions: For semi-supervised settings, clients may send aggregated soft labels for unlabeled data. This representation is often orders of magnitude smaller than a full model update, especially for large models.

FKD can naturally handle statistical heterogeneity (non-IID data) across clients. Because knowledge is shared in the output or feature space, clients can use:

• Different model architectures (heterogeneous FKD).
• Different private datasets. The server's distillation process focuses on aligning the behavior (output predictions) of the global student with the ensemble of client behaviors, rather than forcing parameter consensus. This can reduce client drift as local models are not penalized for diverging to fit their local data.

The server aggregates knowledge (e.g., averages softened logits from clients) to create a teacher ensemble. It then trains the global student model by minimizing a distillation loss (e.g., Kullback-Leibler divergence) between the student's predictions and the ensemble's predictions on an anchor dataset. This dataset can be:

• A small public proxy dataset.
• Synthetically generated data.
• Sampled from client-shared knowledge. This process transfers the collective knowledge of the clients into a single, compact student model without direct access to client parameters or raw data.

FKD offers inherent privacy benefits, though they are distinct from cryptographic guarantees:

• No Parameter Transmission: Client model parameters, which could be inverted to reveal training data, are never shared.
• Knowledge Obfuscation: Softened logits or aggregated embeddings contain less explicit information about individual data points than raw gradients.
• Composability with Stronger Techniques: The knowledge vector can be further protected with Differential Privacy (adding noise) or Secure Aggregation (masking values during transmission) to achieve formal privacy guarantees.

Implementing FKD involves key engineering decisions:

• Anchor Data: Requires a public or generated dataset for distillation, which may not be available in all domains.
• Knowledge Design: Choosing what knowledge to share (logits vs. embeddings) and at what temperature balances communication cost and information fidelity.
• Client-Server Synchronization: Can be performed in rounds (synchronous) or as knowledge arrives (asynchronous).
• Client Model Architecture: Homogeneous clients simplify aggregation; heterogeneous clients offer flexibility but complicate knowledge alignment. The primary trade-off is often between ultimate model performance (sometimes lower than FedAvg with full participation) and significant reductions in uplink communication cost and robustness to system heterogeneity.

The most common approach where clients send softened model outputs (logits) instead of parameter gradients. The server aggregates these logits (e.g., by averaging) and uses them as soft targets to train the global model via a distillation loss (e.g., Kullback-Leibler divergence). This is highly communication-efficient as logit vectors are typically much smaller than full model weights.

• Key Advantage: Drastically reduces uplink payload size.
• Challenge: Requires clients to have sufficient local data to produce meaningful, calibrated logits.
• Example: A client with 10 local images sends a 10x1000 logit matrix (for 1000 classes) instead of millions of weight parameters.

This variant distills knowledge from intermediate feature representations (activations) of a neural network, not just the final output layer. Clients send compressed representations from a designated hint layer. The server trains the global model to mimic these intermediate features using a mean squared error or similar loss.

• Key Advantage: Transfers richer, more generalizable knowledge about data structure.
• Use Case: Particularly effective for convolutional neural networks where early layers capture fundamental visual features.
• Communication Trade-off: Feature maps can be large; techniques like average pooling or bottleneck layers are used to compress them before transmission.

A technique where the server maintains and distills from an ensemble of client models. Instead of aggregating a single global model, the server treats the collection of client models as a committee. Knowledge is distilled from this ensemble's combined predictions into a single, compact global model.

• Key Advantage: Improves robustness and model performance by leveraging diverse client knowledge.
• Process: 1) Clients train local models. 2) Server collects models/logits. 3) Server uses unlabeled public data or a generator to produce inputs, queries the ensemble, and distills the outputs.
• Benefit: Can mitigate the effects of non-IID data across clients.

Also known as mutual distillation or co-distillation, this is a bidirectional process. Knowledge flows both from clients to server and from server to clients. Clients not only contribute knowledge but also receive and distill knowledge from the global model or other clients' models to improve their local performance.

• Key Advantage: Enables personalized federated learning; clients benefit from global knowledge while specializing for their local data.
• Mechanism: In each round, the server may broadcast a lightweight generalized model or meta-information that clients distill into their local models alongside their own training.
• Outcome: Produces a strong global model and a set of specialized, high-performance local models.

A critical technique for scenarios where the server cannot access a representative dataset. The server generates synthetic data (e.g., via a generative adversarial network) to act as input for the distillation process. It queries client models with this synthetic data to produce logits for aggregation and distillation.

• Key Advantage: Preserves data privacy absolutely; no real client data or samples are ever shared or reconstructed at the server.
• Challenge: Requires sophisticated generation to produce data that effectively covers the input distribution learned by the clients.
• Application: Essential for highly regulated domains like healthcare and finance.

This distinction defines the temporal relationship between client training and server distillation.

• Online Distillation: The distillation process is interleaved with federated training rounds. The global model is updated via distillation after each communication round, allowing for continuous, incremental learning. This is the standard approach in iterative FL.
• Offline Distillation: Clients perform complete local training independently. They then send their final model knowledge (e.g., logits on a shared public dataset) to the server in a single communication round. The server performs a one-time distillation to produce the final global model.
• Trade-off: Online enables finer convergence; offline is the ultimate in communication reduction, akin to one-shot federated learning.

A core compression technique where high-precision gradient values (e.g., 32-bit floating point) are mapped to a lower-bit representation (e.g., 8-bit integers) before transmission from client to server. This reduces the number of bits required per parameter update.

• Key Mechanism: Reduces the bit-width of communicated data.
• Trade-off: Introduces quantization noise, which must be managed to preserve convergence.
• Example: Converting gradients from FP32 to INT8 reduces communication payload by ~75%.

A technique where only a critical subset of gradient values is selected for transmission, while others are set to zero. This creates a sparse update tensor, which can be highly compressed.

• Common Method: Transmit only gradients with the largest magnitudes (top-k sparsification).
• Benefit: Can achieve compression ratios of 100x or more by sending <1% of values.
• Challenge: Often requires error feedback to accumulate and transmit discarded gradients in subsequent rounds, preventing information loss.

An aggressive compression method that combines sparsification and extreme quantization. Client updates are compressed into a format where each transmitted value is one of three levels: {-α, 0, +α}.

• Process: First, sparsify to select top-k values. Second, quantize each selected value to its sign (+1 or -1) multiplied by a scaling factor α.
• Efficiency: Achieves high compression by drastically reducing both the number of values sent and their bit representation.
• Use Case: Effective for scenarios with extremely constrained uplink bandwidth.

The foundational algorithm for federated learning. Clients perform multiple local Stochastic Gradient Descent steps on their data and send the updated model weights (or weight deltas) to the server for averaging.

• Core Idea: Reduces communication frequency by performing more local computation per round.
• Challenge: Can lead to client drift on non-IID data, where local models diverge from the global objective.
• Relation to Distillation: FedAvg communicates parameters; knowledge distillation communicates softened outputs or embeddings, offering a different representational format.

A fundamental challenge in federated learning where local client models diverge from the global optimum due to performing multiple steps of local optimization on statistically heterogeneous (non-IID) data.

• Cause: The local objective on a client's unique data distribution differs from the global objective.
• Impact: Degrades convergence speed and final model accuracy, as simple averaging of drifted models is suboptimal.
• Mitigation: Algorithms like FedProx (adds a proximal term) or SCAFFOLD (uses control variates) are designed to correct for drift. Knowledge distillation can also mitigate drift by aligning models on a shared output space.

A system-level constraint and design principle where only a subset of the total available clients is selected to participate in any given federated learning round.

• Necessity: Driven by practical limits on server bandwidth, client availability, and energy constraints.
• Implication: Inherently limits per-round communication volume but introduces sampling variance.
• Strategy: Adaptive client selection algorithms optimize which clients participate based on factors like data utility, compute capability, or network conditions to maximize learning efficiency per communication cost.