FedProtoKD: Dual KD for Heterogeneous Federated Learning

• The paper introduces FedProtoKD, which combines dual knowledge distillation, projection-based dimensional alignment, and trainable server prototypes to address prototype margin shrinking in heterogeneous federated learning.
• The method employs adaptive class-wise margins and importance-weighted public sample logits to maintain inter-class separability and facilitate effective knowledge transfer among diverse client models.
• Experimental results on CIFAR-10 and CIFAR-100 show significant improvements in server and client accuracy over state-of-the-art methods under extreme non-IID and model heterogeneity conditions.

Searching arXiv for the specified paper and closely related federated prototype/distillation methods. FedProtoKD is a heterogeneous federated learning (HFL) framework designed to address a specific failure mode of prior prototype-based HFL methods: prototype margin shrinking. In the formulation reported in "FedProtoKD: Dual Knowledge Distillation with Adaptive Class-wise Prototype Margin for Heterogeneous Federated Learning" (Hossen et al., 26 Aug 2025), clients may use different model architectures and hold highly non-IID private data, while communication is based on class-representative prototypes and logits rather than full model weights. The method combines dual knowledge distillation, projection-based dimensional alignment, a trainable server prototype mechanism called Adaptive Class Wise Margin-Based Trainable Server Prototype (ACTSP), and class-wise adaptive margins. Its stated objective is to preserve inter-class separability and improve both server and client performance under model heterogeneity and extreme statistical heterogeneity (Hossen et al., 26 Aug 2025).

1. Problem setting and motivation

FedProtoKD is situated within prototype-based HFL, where clients exchange class representatives rather than full parameter tensors. The paper identifies three jointly interacting sources of degradation in prior methods: non-IID data, heterogeneous client models, and prototype shrinking during aggregation (Hossen et al., 26 Aug 2025). In this setting, each client may observe only a skewed subset of classes or class frequencies, so local prototypes can differ substantially in geometry and confidence. At the same time, clients may instantiate different architectures and different prototype dimensions, making direct aggregation nontrivial.

The central criticism is directed at simple server-side weighted averaging of prototypes. According to the paper, such averaging can compress inter-class geometry, yielding global prototypes that are less separable than the client-side prototypes from which they were formed. This loss of inter-class distance is termed the prototype margin shrinking problem, and it is presented as particularly harmful when models are architecturally heterogeneous and the data distributions are extremely non-IID (Hossen et al., 26 Aug 2025).

FedProtoKD is positioned as a remedy through the combination of four components: dual knowledge distillation from clients to server and server to clients, projection-based dimensional alignment for heterogeneous model outputs, a trainable server prototype generated by a contrastive-learning-style module, and an adaptive class-wise margin that aims to prevent over-regularization while preserving inter-class separability. The paper further argues that logits from public samples provide model-agnostic output-level knowledge, whereas prototypes encode representation geometry from private data; the method therefore treats the two as complementary forms of knowledge (Hossen et al., 26 Aug 2025).

2. End-to-end training pipeline

FedProtoKD operates in communication rounds. On the client side, each client first performs local supervised training on private data using cross-entropy:

$$\min_{\theta_c} \sum_{(x_i, y_i) \in D_c} \mathcal{L}_{\text{CE}\left(\mathcal{F}_{\theta_c}(x_i), y_i \right)}.$$

After this local phase, each client runs its model on the unlabeled public dataset $\mathcal{D}_p$, extracts logits for public samples, computes class prototypes from private data, and sends both logits and prototypes to the server (Hossen et al., 26 Aug 2025).

The server-side phase comprises six reported steps: aggregation of client logits using a confidence/variance-based weighting strategy, pseudo-label inference for public samples, generation of trainable server prototypes using ACTSP, computation of importance weights for public samples based on prototype closeness, training of the global server model on public data using both KD and prototype losses, and transmission of server logits and server prototypes back to the clients (Hossen et al., 26 Aug 2025).

Clients then incorporate the returned server knowledge through KL distillation on server logits, CE loss on pseudo-labels, and, after the first round, prototype regularization against server prototypes. The paper therefore characterizes the method as “dual KD” in two senses: logits distillation for output-level knowledge and prototype distillation for feature-level knowledge (Hossen et al., 26 Aug 2025).

A plausible implication is that FedProtoKD uses public data as a common cross-client substrate for knowledge transfer, while reserving client-private data for prototype extraction and local supervised adaptation. This interpretation is consistent with the paper’s distinction between public-sample logits and private-data representation geometry.

3. Prototype representation and heterogeneous architecture alignment

For a client $c_i$ and class $c$, the class prototype is defined as the mean of the feature vectors for samples in that class:

$P_{c_i}^c = \frac{1}{|D_{c_i}^{c}|} \sum_{(x_i, y_i) \in D_{c_i}^{c} \mathcal{F}_{\theta_{c_i}(x_i).$

This is the standard prototype construction used as the foundation for the framework (Hossen et al., 26 Aug 2025).

The paper contrasts this with the vanilla aggregation used in prior methods:

$$\mathcal{\tilde{P}^c = \frac{1}{|\mathcal{N}_c|} \sum_{i \in \mathcal{N}_c} \frac{|D_{c_i}^{c}|}{N_c} P_{c_i}^c }.$$

Here, $N_c$ denotes the number of class-$c$ samples on one client and $\mathcal{N}_c$ the total number across clients. FedProtoKD argues that such aggregation can overweight weak clients or compress the geometry of stronger clients, thereby shrinking the global prototype margin (Hossen et al., 26 Aug 2025).

A key architectural contribution is a learnable projection layer that aligns prototype feature dimensions across heterogeneous models. The reported backbone configurations have different native prototype widths: ResNet18 and ResNet34 use 512-dimensional prototype layers, whereas ResNet50, ResNet101, and ResNet152 use 2048-dimensional layers. FedProtoKD maps all representations into a common $K$-dimensional space via a learnable projection layer placed after the feature extractor, making prototype exchange possible even when client architectures differ substantially (Hossen et al., 26 Aug 2025).

The paper specifically distinguishes this design from Fed2PKD, which assumes a uniform 2048-dimensional prototype space and therefore does not fully solve architectural heterogeneity. This suggests that dimensional alignment is treated not as a peripheral engineering convenience but as a prerequisite for meaningful prototype interaction across heterogeneous clients (Hossen et al., 26 Aug 2025).

4. ACTSP and the adaptive class-wise margin

The server-side prototype mechanism is named Adaptive Class Wise Margin-Based Trainable Server Prototype (ACTSP). For each class $\mathcal{D}_p$0, the server initializes a trainable vector

$\mathcal{D}_p$1

which is passed through a shared server-side neural network $\mathcal{D}_p$2 to obtain a refined global prototype:

$\mathcal{D}_p$3

The generator is described as a small MLP, typically two fully connected layers with ReLU in between, and the same network parameters are shared across all classes (Hossen et al., 26 Aug 2025).

The paper’s stated intuition is that this trainable mechanism can maintain or enlarge inter-class gaps better than direct averaging. Instead of imposing a single fixed maximum margin for all classes, FedProtoKD defines a class-specific adaptive margin:

$\mathcal{D}_p$4

where

$\mathcal{D}_p$5

Here $\mathcal{D}_p$6 is the cluster center for class $\mathcal{D}_p$7, $\mathcal{D}_p$8 is a distance between class centers, and $\mathcal{D}_p$9 is a threshold that caps margin growth (Hossen et al., 26 Aug 2025).

The paper emphasizes three consequences of this formulation: classes with tight clusters should not be pushed too far apart, classes with more dispersion can be regularized more strongly, and the margin changes over rounds $c_i$0 as separability evolves. The associated prototype contrastive objective is defined as

$c_i$1

The distance $c_i$2 is Euclidean in the proposed method (Hossen et al., 26 Aug 2025).

A fixed-margin comparison variant is also reported:

$c_i$3

FedProtoKD replaces the fixed $c_i$4 with the adaptive $c_i$5 (Hossen et al., 26 Aug 2025). In the paper’s framing, this modification is crucial under extreme heterogeneity because fixed margins may over-regularize some classes and under-regularize others.

5. Dual knowledge distillation and public-sample weighting

FedProtoKD distills knowledge in two directions. On the server side, clients send logits on public samples, and the server aggregates them using a variance-based importance weighting:

$c_i$6

The paper interprets higher variance of logits as higher confidence. Pseudo-labels are then inferred as

$c_i$7

The server KD loss is reported as

$c_i$8

In parallel, prototype feature representation distillation matches server embeddings to the server prototype of the pseudo-label:

$c_i$9

The total server objective is

$c$0

with $c$1 (Hossen et al., 26 Aug 2025).

On the client side, the server transmits server logits $c$2 and server prototypes $c$3. Clients optimize

$c$4

$c$5

$c$6

with $c$7. Starting from round $c$8, clients also regularize private training with prototype matching:

$c$9

$P_{c_i}^c = \frac{1}{|D_{c_i}^{c}|} \sum_{(x_i, y_i) \in D_{c_i}^{c} \mathcal{F}_{\theta_{c_i}(x_i).$0

$P_{c_i}^c = \frac{1}{|D_{c_i}^{c}|} \sum_{(x_i, y_i) \in D_{c_i}^{c} \mathcal{F}_{\theta_{c_i}(x_i).$1

where $P_{c_i}^c = \frac{1}{|D_{c_i}^{c}|} \sum_{(x_i, y_i) \in D_{c_i}^{c} \mathcal{F}_{\theta_{c_i}(x_i).$2 (Hossen et al., 26 Aug 2025).

FedProtoKD also modifies the treatment of public samples. Instead of FedPKD’s filtering strategy, it keeps all public samples but weights them according to closeness between a public sample’s feature and its pseudo-label prototype. The importance factor is

$P_{c_i}^c = \frac{1}{|D_{c_i}^{c}|} \sum_{(x_i, y_i) \in D_{c_i}^{c} \mathcal{F}_{\theta_{c_i}(x_i).$3

where

$P_{c_i}^c = \frac{1}{|D_{c_i}^{c}|} \sum_{(x_i, y_i) \in D_{c_i}^{c} \mathcal{F}_{\theta_{c_i}(x_i).$4

Here $P_{c_i}^c = \frac{1}{|D_{c_i}^{c}|} \sum_{(x_i, y_i) \in D_{c_i}^{c} \mathcal{F}_{\theta_{c_i}(x_i).$5 is the normalized inverse $P_{c_i}^c = \frac{1}{|D_{c_i}^{c}|} \sum_{(x_i, y_i) \in D_{c_i}^{c} \mathcal{F}_{\theta_{c_i}(x_i).$6 distance, $P_{c_i}^c = \frac{1}{|D_{c_i}^{c}|} \sum_{(x_i, y_i) \in D_{c_i}^{c} \mathcal{F}_{\theta_{c_i}(x_i).$7 is the median-like center, $P_{c_i}^c = \frac{1}{|D_{c_i}^{c}|} \sum_{(x_i, y_i) \in D_{c_i}^{c} \mathcal{F}_{\theta_{c_i}(x_i).$8 controls the sigmoid steepness, and $P_{c_i}^c = \frac{1}{|D_{c_i}^{c}|} \sum_{(x_i, y_i) \in D_{c_i}^{c} \mathcal{F}_{\theta_{c_i}(x_i).$9 controls the fraction or emphasis of samples. The inverse distance is defined conceptually as

$$\mathcal{\tilde{P}^c = \frac{1}{|\mathcal{N}_c|} \sum_{i \in \mathcal{N}_c} \frac{|D_{c_i}^{c}|}{N_c} P_{c_i}^c }.$$0

According to the paper, this strategy down-weights noisy or far samples without discarding them, thereby preserving data diversity while emphasizing informative public examples (Hossen et al., 26 Aug 2025).

6. Experimental protocol and empirical findings

The reported evaluation uses CIFAR-10 and CIFAR-100, each with 50,000 training images and 10,000 test images. A public unlabeled dataset of 2,500 samples is constructed, and the remainder is distributed among clients (Hossen et al., 26 Aug 2025).

Two heterogeneity regimes are considered. The first is practical or Dirichlet heterogeneity, with $$\mathcal{\tilde{P}^c = \frac{1}{|\mathcal{N}_c|} \sum_{i \in \mathcal{N}_c} \frac{|D_{c_i}^{c}|}{N_c} P_{c_i}^c }.$$1 for extreme non-IID and $$\mathcal{\tilde{P}^c = \frac{1}{|\mathcal{N}_c|} \sum_{i \in \mathcal{N}_c} \frac{|D_{c_i}^{c}|}{N_c} P_{c_i}^c }.$$2 for moderate non-IID. The second is pathological split heterogeneity, with CIFAR-10 using $$\mathcal{\tilde{P}^c = \frac{1}{|\mathcal{N}_c|} \sum_{i \in \mathcal{N}_c} \frac{|D_{c_i}^{c}|}{N_c} P_{c_i}^c }.$$3 and $$\mathcal{\tilde{P}^c = \frac{1}{|\mathcal{N}_c|} \sum_{i \in \mathcal{N}_c} \frac{|D_{c_i}^{c}|}{N_c} P_{c_i}^c }.$$4, and CIFAR-100 using $$\mathcal{\tilde{P}^c = \frac{1}{|\mathcal{N}_c|} \sum_{i \in \mathcal{N}_c} \frac{|D_{c_i}^{c}|}{N_c} P_{c_i}^c }.$$5 and $$\mathcal{\tilde{P}^c = \frac{1}{|\mathcal{N}_c|} \sum_{i \in \mathcal{N}_c} \frac{|D_{c_i}^{c}|}{N_c} P_{c_i}^c }.$$6 (Hossen et al., 26 Aug 2025). Model heterogeneity includes ResNet18, ResNet34, ResNet50, ResNet101, and ResNet152, with ResNet34 used as the server model in all experiments.

The baselines are FedAvg, FedProto, FedProx, FedPKD, Fed2PKD, FedTGP, and FedProtoKD-$$\mathcal{\tilde{P}^c = \frac{1}{|\mathcal{N}_c|} \sum_{i \in \mathcal{N}_c} \frac{|D_{c_i}^{c}|}{N_c} P_{c_i}^c }.$$7, the latter being a FedProtoKD variant with fixed max margin and all public samples. The evaluation metrics are server test accuracy $$\mathcal{\tilde{P}^c = \frac{1}{|\mathcal{N}_c|} \sum_{i \in \mathcal{N}_c} \frac{|D_{c_i}^{c}|}{N_c} P_{c_i}^c }.$$8, client test accuracy $$\mathcal{\tilde{P}^c = \frac{1}{|\mathcal{N}_c|} \sum_{i \in \mathcal{N}_c} \frac{|D_{c_i}^{c}|}{N_c} P_{c_i}^c }.$$9, and prototype margin or separability analysis. Reported hyperparameters include clients $N_c$0, participation ratio $N_c$1, client local epochs $N_c$2, server epochs $N_c$3, batch size $N_c$4, rounds $N_c$5, ACTSP training epochs $N_c$6, margin growth threshold $N_c$7, prototype dimension $N_c$8 default, and $N_c$9; the FedPKD filter ratio is $c$0 (Hossen et al., 26 Aug 2025).

The paper reports that FedProtoKD achieves average improvements of 1.13% to 34.13% in accuracy across settings and significantly outperforms existing state-of-the-art HFL methods (Hossen et al., 26 Aug 2025). In heterogeneous model settings, compared with Fed2PKD, server accuracy improves by 6.82% to 34.13% under extreme heterogeneity and by 2.59% to 29.03% under moderate heterogeneity. Client accuracy improves by about 9.32% to 15.1% over baselines in extreme heterogeneous cases and by up to about 3% over FedPKD and 1.39% over FedTGP in many settings (Hossen et al., 26 Aug 2025).

In homogeneous model settings, the method still improves over FedPKD, FedTGP, Fed2PKD, FedAvg, and FedProx, although the gains are smaller than in heterogeneous-model settings. The paper reports server gains up to 7.11% to 34.71% in extreme heterogeneity and 10.38% to 35.18% in moderate heterogeneity (Hossen et al., 26 Aug 2025). The strongest empirical pattern reported is that the method is especially effective when both model heterogeneity and data skew are severe.

7. Ablations, comparisons with prior methods, and limitations

The ablation study isolates two major components. FedProtoKD-$c$1 uses a fixed maximum margin for all classes together with public-sample importance weighting, thereby isolating the effect of the adaptive class-wise margin. A second variant, W/O $c$2, uses the adaptive class-wise margin but all public samples without importance weighting, thereby isolating the effect of sample weighting. The full FedProtoKD combines both. The reported conclusion is that the adaptive class-wise margin improves prototype separability and accuracy, public-sample importance weighting further boosts performance, and the full model consistently performs best (Hossen et al., 26 Aug 2025).

The paper also compares prototype margin trends and t-SNE visualizations. These comparisons show that FedProto and FedTGP suffer from margin shrinking, while FedProtoKD maintains larger inter-class gaps and learned prototypes that are more clearly separated (Hossen et al., 26 Aug 2025). This suggests that the method’s principal empirical contribution is not only accuracy improvement but also stabilization of prototype geometry under aggressive heterogeneity.

Relative to prior prototype-based HFL methods, the paper draws several distinctions. Versus FedProto, FedProtoKD adds dual KD, prototype projection, trainable server prototypes, and adaptive margins. Versus FedPKD, it keeps all public samples and weights them by importance instead of filtering them, while also replacing simple aggregation with ACTSP. Versus Fed2PKD, it explicitly addresses varying feature dimensions with a projection layer and supports broader architectural diversity. Versus FedTGP, it replaces a global fixed maximum margin with a class-wise adaptive margin intended to avoid over-regularization (Hossen et al., 26 Aug 2025).

The paper does not provide a long explicit limitations section, but several assumptions and constraints are stated or directly implied. FedProtoKD assumes access to a public unlabeled dataset shared across clients and server. It introduces additional model complexity through the prototype dimension projection layer and extra server-side overhead through ACTSP. The experiments are limited to CIFAR-10 and CIFAR-100, so scalability to larger datasets and foundation models is left for future work. The method also assumes that prototype closeness is a good proxy for public-sample quality; the paper notes that this may be less reliable if public samples are very far from client distributions (Hossen et al., 26 Aug 2025).

Taken together, FedProtoKD is presented as a prototype-based HFL method that treats knowledge as dual entities—logits and prototypes—aligns heterogeneous feature spaces through learnable projection, and counters prototype margin shrinkage through ACTSP and adaptive class-wise margins. A plausible implication is that its main significance lies in reframing prototype aggregation as a trainable geometric alignment problem rather than a purely averaging-based communication primitive.