From Classifier Heads to Encoder Prototypes: Personalized Federated Learning with Head-Derived Clustering

Amirardalan Dehghanpour; Arian Noori; Ehsan Tanghatari; Naimeh Omidvar; Hamid R. Rabiee

cs.LGOct 21, 2026

From Classifier Heads to Encoder Prototypes: Personalized Federated Learning with Head-Derived Clustering

Authors: Amirardalan Dehghanpour, Arian Noori, Ehsan Tanghatari, Naimeh Omidvar, Hamid R. Rabiee

Organizations: Sharif University of Technology · Tehran Institute for Advanced Studies (TeIAS)

Abstract

Personalized federated learning must decide which clients should share representations. We introduce HDC-FL (Head-Derived Clustering for Personalized Federated Learning), which reads that decision off a signal shared-encoder training already produces: the client-specific classifi- er heads. At periodic refresh rounds the server projects the cached heads into a normalized low-rank descriptor space, groups the descriptors, builds one encoder prototype per group from the cached client encoders, and returns to each client a soft mixture of those prototypes paired with its own unchanged head. Client-specific heads are never averaged or distributed to oth- er clients; the shared-encoder trajectory is left intact; and no second optimization trajectory, additional local objective, or post-training fine-tuning stage is required. Across 11 CIFAR-10 and CIFAR-100 configurations covering label-subset, Dirichlet and rotation heterogeneity with ResNet-18 and MobileNetV2 encoders, HDC-FL is competitive with the strongest personalized and clustered baselines in macro-averaged mean accuracy and attains the numerically highest macro-averaged worst-client and bottom-20% accuracy among the evaluated methods. The tail results distinguish methods that look similar when only their means are compared. A controlled donor-swap diagnostic further finds positive donor-ranking associations between head similarity and encoder compatibility in all three settings tested, most clearly under covariate shift.

Figures & tables

Figure 1: HDC-FL overview. At each refresh the server converts cached client-specific heads into low-rank client descriptors, groups them, constructs group-level encoder prototypes from cached client encoders, and returns a soft prototype mixture to each client. The mixture is paired with the client’s unchanged private head, and the shared training trajectory continues unmodified.

Method	C10 LS2	C10 LS3	C10 Rot	C100 $\alpha{=}0.07$	C100 $\alpha{=}0.3$	C100 LS2	C100 LS3	C100 Rot	Avg. Rank $\downarrow$
FedAvg-SE	97.17 $\pm$ 0.62	95.51 $\pm$ 0.71	86.42 $\pm$ 0.80	81.15 $\pm$ 1.06	72.09 $\pm$ 0.34	81.23 $\pm$ 0.63	76.88 $\pm$ 0.17	57.42 $\pm$ 0.68	7.25
FedBABU	96.99 $\pm$ 0.63	97.18 $\pm$ 1.35	90.05 $\pm$ 0.43	87.73 $\pm$ 1.51	82.19 $\pm$ 0.16	84.81 $\pm$ 4.02	82.33 $\pm$ 3.92	68.54 $\pm$ 1.91	2.12
FedALA	96.51 $\pm$ 0.65	95.50 $\pm$ 0.85	86.73 $\pm$ 0.03	81.45 $\pm$ 1.39	74.67 $\pm$ 0.03	80.96 $\pm$ 0.14	77.75 $\pm$ 0.40	63.86 $\pm$ 0.21	6.88
FedCE	96.89 $\pm$ 0.02	95.92 $\pm$ 0.12	86.52 $\pm$ 0.28	82.22 $\pm$ 0.41	77.19 $\pm$ 0.06	82.29 $\pm$ 0.76	78.98 $\pm$ 0.45	64.37 $\pm$ 0.26	4.50
FedSoft	96.40 $\pm$ 0.09	95.44 $\pm$ 1.39	84.41 $\pm$ 0.48	81.90 $\pm$ 0.94	75.41 $\pm$ 0.22	81.63 $\pm$ 0.25	79.23 $\pm$ 0.04	61.06 $\pm$ 0.59	6.88
IFCA	96.36 $\pm$ 1.05	94.38 $\pm$ 0.42	87.18 $\pm$ 0.13	80.12 $\pm$ 0.28	75.93 $\pm$ 0.55	80.61 $\pm$ 1.00	78.40 $\pm$ 0.23	63.69 $\pm$ 0.46	7.62

Table 1: Mean client test accuracy (%) on ResNet-18, mean

\pm

sample standard deviation over three seeds. Average rank is computed within each setting and then averaged over the eight settings. Bold and underline mark the best and second-best entry per setting.

Method	Mean	Worst	Bot.20%	Gap	Tail rank
FedAvg-SE	83.35	76.95	78.95	11.15	6.00
FedBABU	87.58	55.32	73.78	38.37	8.18
FedALA	84.19	77.68	79.75	11.07	5.73
FedCE	85.03	78.73	80.69	10.94	3.73
FedSoft	84.15	76.52	79.31	12.38	6.00
IFCA	83.64	71.25	77.58	17.90	7.27

Table 2: Client-level results macro-averaged over all 11 configurations.

Variant	Mean	Worst	Bot.20%
No prototype	77.57 $\pm$ 0.40	69.85 $\pm$ 3.97	72.20 $\pm$ 1.46
One prototype	80.30 $\pm$ 0.34	72.05 $\pm$ 2.90	74.93 $\pm$ 1.16
Head-derived, $G{=}4$	81.04 $\pm$ 0.73	72.57 $\pm$ 1.67	75.47 $\pm$ 1.10

Table 3: Controlled grouping ablation on CIFAR-10 rotation without pretrained initialization. Here

G=4

is fixed to match the four known rotation conditions and isolate the effect of head-derived grouping. Values are mean

\pm

sample standard deviation over three seeds.

	C10 LS3	C100 LS3	C100 Rot
Spearman $\rho$ (loss)	0.193	0.117	0.243
AUROC (loss)	0.665	0.565	0.783
Beneficial-donor rate	0.156	0.711	0.245
AP	0.513	0.782	0.642
AP lift over prevalence	0.357	0.071	0.397
Top-1 $-$ random (pp)	1.57	0.41	2.50

Table 4: Donor-compatibility diagnostic on MobileNetV2, averaged over targets and seeds. Spearman

\rho

measures rank association between head similarity and loss benefit, while AUROC and average precision (AP) measure discrimination between beneficial and non-beneficial donors. The beneficial-donor rate gives the positive-class prevalence against which AP should be interpreted. The final two rows report the accuracy gain from selecting the nearest-head donor rather than a random donor, and from the closest head-distance quintile relative to the farthest.

Appendix figures & tables14 assets

Supplementary material from the paper’s appendix.

Appendix

Method	Method-specific configuration
FedAvg-SE	The encoder is aggregated across participating clients using sample-count weighting, while each client retains its own classifier head. No additional personalization stage is applied.
FedBABU	The classifier is kept fixed during federated training and only the encoder/body is locally optimized and aggregated. Final personalization uses 5 epochs of full-model fine-tuning with SGD, learning rate $10^{-2}$ , momentum $0$ , and weight decay $0$ .
FedALA	ALA depth $p=2$ , local-data fraction $s=80\%$ , and weight-learning rate $\eta=1$ . Initial ALA weight optimization is capped at 50 epochs with patience 3 and tolerance $10^{-4}$ ; subsequent rounds use one ALA weight-update epoch.
IFCA	Uses $S=4$ cluster models. Client-to-cluster assignment is based on model loss evaluated on at most 512 local samples.
PACFL	Client subspace dimension $p=3$ , hierarchical-clustering threshold $1.5$ radians, average linkage, and one-shot client clustering.
FedSoft	Uses $S=3$ cluster centers, proximal coefficient $\lambda=0.1$ , importance-estimation interval $\tau=2$ , smoothing constant $\sigma=10^{-4}$ , and client-selection size $K_{\mathrm{sel}}=20$ .

Appendix

Table 5: Method-specific baseline configurations used in our experiments. FedAvg-SE denotes FedAvgSharedEncoder.

Method	CIFAR-100 LS3	CIFAR-10 $\alpha{=}0.07$	CIFAR-10 LS3	Avg. Rank $\downarrow$
FedAvg-SE	77.43 $\pm$ 0.37	95.85 $\pm$ 0.28	95.25 $\pm$ 0.74	6.67
FedBABU	81.27 $\pm$ 1.89	96.26 $\pm$ 1.19	96.05 $\pm$ 1.42	2.33
FedALA	77.77 $\pm$ 0.27	95.99 $\pm$ 0.10	94.87 $\pm$ 1.17	5.67
FedCE	78.87 $\pm$ 0.43	96.12 $\pm$ 0.38	95.96 $\pm$ 0.51	4.00
FedSoft	79.57 $\pm$ 0.48	95.07 $\pm$ 1.13	95.52 $\pm$ 0.39	5.33
IFCA	77.69 $\pm$ 0.49	90.90 $\pm$ 1.54	94.83 $\pm$ 0.65	8.67

Appendix

Table 6: MobileNetV2 mean client test accuracy (%), mean

\pm

sample standard deviation over seeds 42, 123 and 2026. Bold and underline mark the best and second-best entry per setting.

Method	C10 LS2	C10 LS3	C10 Rot	C100 $\alpha{=}.07$	C100 $\alpha{=}.3$	C100 LS2	C100 LS3	C100 Rot	M-C100 LS3	M-C10 $\alpha{=}.07$	M-C10 LS3
FedAvg-SE	88.73 $\pm$ 5.75	90.23 $\pm$ 1.03	79.53 $\pm$ 4.86	74.29 $\pm$ 0.82	67.10 $\pm$ 0.69	75.70 $\pm$ 0.69	68.06 $\pm$ 1.50	52.29 $\pm$ 1.28	72.45 $\pm$ 0.60	87.92 $\pm$ 2.67	90.17 $\pm$ 0.18
FedBABU	77.64 $\pm$ 10.88	82.66 $\pm$ 13.87	75.68 $\pm$ 3.07	55.35 $\pm$ 26.75	34.07 $\pm$ 6.89	51.94 $\pm$ 31.17	35.67 $\pm$ 5.29	20.04 $\pm$ 2.68	27.40 $\pm$ 6.79	76.28 $\pm$ 24.78	71.84 $\pm$ 18.52
FedALA	86.89 $\pm$ 7.23	90.20 $\pm$ 2.12	80.41 $\pm$ 4.74	72.97 $\pm$ 1.82	68.51 $\pm$ 1.17	74.21 $\pm$ 2.59	72.82 $\pm$ 0.84	57.50 $\pm$ 0.26	72.57 $\pm$ 0.79	88.87 $\pm$ 0.48	89.53 $\pm$ 2.33
FedCE	88.33 $\pm$ 5.50	91.50 $\pm$ 0.28	79.46 $\pm$ 4.10	75.33 $\pm$ 1.96	72.30 $\pm$ 1.19	75.92 $\pm$ 0.07	72.29 $\pm$ 2.25	57.76 $\pm$ 0.76	73.53 $\pm$ 1.19	88.28 $\pm$ 2.07	91.36 $\pm$ 0.79
FedSoft	87.55 $\pm$ 4.39	87.73 $\pm$ 3.01	75.76 $\pm$ 4.87	76.05 $\pm$ 0.55	66.92 $\pm$ 0.58	75.64 $\pm$ 0.14	72.14 $\pm$ 0.69	54.98 $\pm$ 1.36	73.90 $\pm$ 1.97	82.63 $\pm$ 6.61	88.40 $\pm$ 2.80
IFCA	86.15 $\pm$ 8.06	69.67 $\pm$ 26.74	81.83 $\pm$ 2.86	73.27 $\pm$ 0.01	69.68 $\pm$ 0.27	73.11 $\pm$ 0.49	68.78 $\pm$ 8.28	56.89 $\pm$ 3.31	70.65 $\pm$ 2.51	47.28 $\pm$ 13.44	86.41 $\pm$ 4.12

Appendix

Table 7: Final worst-client test accuracy (%), mean

\pm

sample standard deviation over three seeds.

Method	C10 LS2	C10 LS3	C10 Rot	C100 $\alpha{=}.07$	C100 $\alpha{=}.3$	C100 LS2	C100 LS3	C100 Rot	M-C100 LS3	M-C10 $\alpha{=}.07$	M-C10 LS3
FedAvg-SE	92.05 $\pm$ 1.36	91.84 $\pm$ 0.71	82.45 $\pm$ 2.19	75.96 $\pm$ 1.24	68.10 $\pm$ 0.16	77.20 $\pm$ 0.39	71.29 $\pm$ 0.29	53.68 $\pm$ 1.77	73.74 $\pm$ 0.20	90.26 $\pm$ 1.23	91.91 $\pm$ 0.54
FedBABU	89.45 $\pm$ 4.26	91.59 $\pm$ 5.49	82.78 $\pm$ 1.30	75.71 $\pm$ 10.75	68.89 $\pm$ 0.82	64.71 $\pm$ 20.65	59.18 $\pm$ 18.77	47.08 $\pm$ 10.85	56.46 $\pm$ 9.34	88.73 $\pm$ 5.46	87.03 $\pm$ 7.76
FedALA	91.64 $\pm$ 2.19	91.58 $\pm$ 2.75	82.78 $\pm$ 1.97	74.92 $\pm$ 1.38	69.90 $\pm$ 1.20	76.92 $\pm$ 1.30	73.85 $\pm$ 0.01	59.89 $\pm$ 0.10	74.27 $\pm$ 1.04	90.86 $\pm$ 0.99	90.67 $\pm$ 1.95
FedCE	92.11 $\pm$ 1.98	92.55 $\pm$ 0.37	81.61 $\pm$ 3.01	77.17 $\pm$ 0.79	73.39 $\pm$ 1.52	77.74 $\pm$ 0.62	74.15 $\pm$ 1.02	59.88 $\pm$ 0.18	74.99 $\pm$ 1.12	91.00 $\pm$ 1.45	92.94 $\pm$ 0.57
FedSoft	90.48 $\pm$ 1.67	91.20 $\pm$ 2.26	78.41 $\pm$ 2.87	77.01 $\pm$ 0.84	70.82 $\pm$ 1.10	77.36 $\pm$ 0.18	74.41 $\pm$ 0.36	56.73 $\pm$ 1.22	75.73 $\pm$ 0.92	88.32 $\pm$ 3.37	91.91 $\pm$ 1.09
IFCA	90.35 $\pm$ 4.31	86.74 $\pm$ 5.44	84.12 $\pm$ 1.78	75.84 $\pm$ 0.04	71.63 $\pm$ 0.09	74.58 $\pm$ 1.31	73.84 $\pm$ 2.27	59.78 $\pm$ 0.73	72.80 $\pm$ 0.63	73.45 $\pm$ 1.85	90.31 $\pm$ 0.43

Appendix

Table 8: Final bottom-20% client test accuracy (%), mean

\pm

sample standard deviation over three seeds.

Method	C10 LS2	C10 LS3	C10 Rot	C100 $\alpha{=}.07$	C100 $\alpha{=}.3$	C100 LS2	C100 LS3	C100 Rot	M-C100 LS3	M-C10 $\alpha{=}.07$	M-C10 LS3	Macro
FedAvg-SE	11.17	8.70	11.45	13.79	10.06	11.90	15.26	9.89	9.62	12.08	8.69	11.15
FedBABU	22.36	17.24	18.70	38.38	54.44	42.58	55.82	58.12	62.91	23.72	27.78	38.37
FedALA	12.84	8.05	11.55	16.47	11.59	12.13	9.16	10.51	9.35	11.13	8.95	11.07
FedCE	11.67	7.87	13.39	13.81	8.57	12.67	11.21	11.86	10.05	11.72	7.48	10.94
FedSoft	12.34	11.29	14.03	12.94	12.23	12.23	12.06	10.65	10.72	17.37	10.31	12.38
IFCA	13.54	28.94	10.52	13.26	10.80	14.32	14.13	12.46	14.31	52.72	11.91	17.90

Appendix

Table 9: Final fairness gap (best minus worst client accuracy, percentage points; lower is better). Each setting is first averaged over its seeds; the last column macro-averages the 11 setting-level means.

Figure 2: Mean client test accuracy over communication rounds for CIFAR-10 rotation, CIFAR-100 LS3 and CIFAR-100

\alpha=0.3

, averaged over three seeds. Before activation the HDC-FL curve reports the shared trajectory; afterwards it reports the personalized prototype mixtures constructed from it.

Figure 3: Learning curves for the configurations not shown in Figure 2 , averaged over three seeds. Shaded regions are one sample standard deviation.

Figure 4: Grouping diagnostics over communication rounds on CIFAR-10 rotation without pretrained initialization with four groups: cluster separation (silhouette), agreement between consecutive assignments (ARI), membership confidence (entropy), and personalization gain. Lines and shaded regions are the mean and one standard deviation over three seeds.

Descriptor	Mean	Sil.	ARI
Raw head	85.27 $\pm$ 0.09	0.027 $\pm$ 0.002	0.152 $\pm$ 0.065
Low rank ( $r{=}8$ )	85.24 $\pm$ 0.31	0.108 $\pm$ 0.012	0.401 $\pm$ 0.080

Appendix

Table 10: Descriptor. Sil. is the silhouette score of the resulting clustering and ARI the agreement between the assignments of consecutive refreshes; both are higher-is-better. Means

\pm

sample standard deviations over three seeds.

Rank $r$	Mean	Worst	Bot.20%	Silhouette $\uparrow$
8	40.85	35.57	36.71	0.195
20	40.96	35.38	36.36	0.126
40	40.24	34.02	35.91	0.090
60	40.50	34.05	35.71	0.068
80	40.23	34.39	35.84	0.059

Appendix

Table 11: Descriptor-rank sensitivity, single seed.

Membership	Mean	Worst	Bot.20%
Hard assignment	78.89 $\pm$ 0.46	73.55 $\pm$ 1.18	75.39 $\pm$ 0.52
Soft mixture	79.74 $\pm$ 0.40	74.63 $\pm$ 0.26	76.34 $\pm$ 0.46

Appendix

Table 12: Prototype membership. Means

\pm

sample standard deviations over three seeds.

$\lambda$	Mean	Worst	Bot.20%
0.0	76.95	71.93	72.58
0.5	83.72	79.92	80.34
1.0	85.22	80.12	81.77

Appendix

Table 13: Mixing coefficient

\lambda

. Single-seed diagnostic.

Setting	Q1	Q2	Q3	Q4	Q5
C10 LS3	$-2.96$	$-3.65$	$-4.11$	$-4.51$	$-4.63$
	(0.35)	(0.41)	(0.42)	(0.42)	(0.55)
C100 LS3	$+0.99$	$+0.97$	$+0.86$	$+0.55$	$+0.47$
	(0.27)	(0.26)	(0.25)	(0.24)	(0.26)
C100 Rot	$-1.16$	$-2.69$	$-2.99$	$-2.95$	$-3.13$
	(0.32)	(0.29)	(0.23)	(0.26)	(0.26)

Appendix

Table 14: Accuracy gain (percentage points) of a single donor-encoder swap, by low-rank head-distance quintile; Q1 is the closest fifth of donors. Values are means over the 60 target

\times

seed units of a setting, with standard errors.

	C10 LS3	C100 LS3	C100 Rot
Spearman $\rho$ , low rank	0.193	0.117	0.243
Spearman $\rho$ , raw	0.199	0.050	0.288
Difference	$-0.006$	$+0.067$	$-0.045$
AUROC, low rank	0.665	0.565	0.783
AUROC, raw	0.677	0.502	0.793
Difference	$-0.012$	$+0.064$	$-0.010$

Appendix

Table 15: Low-rank versus raw head similarity as a donor-ranking signal. Positive differences favour the low-rank descriptor.

Federated Learning

Federated learning trains shared models across data holders without requiring them to pool their raw data.

See all 961 papers

Personalized Federated Learning through Global Knowledge Distillation and Local Head Adaptation
Sep 15, 2026Polycarpo Souza Neto, José Mairton Barros da Silva Júnior, Charles Casimiro CavalcanteFederated Knowledge Distillation Non-IID Federated Learning
Universidade Federal do Cear´a, Fortaleza, CE, Brazil · Uppsala University, Uppsala, Sweden
Fed-BAC: Federated Bandit-Guided Additive Clustering in Hierarchical Federated Learning
May 12, 2026Satwat Bashir, Tasos Dagiuklas, Muddesar IqbalHierarchical Federated Learning Federated Learning Aggregation
Computer Science and Informatics · London South Bank University · London, UK +1
FedDES: Graph-Based Dynamic Ensemble Selection for Personalized Federated Learning
Mar 30, 2026Brianna Mueller, W. Nick StreetFederated Learning Non-IID Federated Learning
University of Iowa, Iowa City, Iowa

Sep 15, 2026J/K move · Enter open · S save

Personalized Federated Learning through Global Knowledge Distillation and Local Head Adaptation

Polycarpo Souza Neto, José Mairton Barros da Silva Júnior, Charles Casimiro Cavalcante

Universidade Federal do Cear´a, Fortaleza, CE, Brazil · Uppsala University, Uppsala, Sweden

Federated Knowledge Distillation Non-IID Federated Learning Knowledge Distillation Federated Learning

Open paper AskPDF arXiv

Statistical heterogeneity limits federated learning when a single global classifier cannot represent client-specific label distributions. In this work, we propose Personalized Federated Knowledge Distillation with Head Adaptation (pFedKDH), which aggregates only the shared backbone, keeps persistent client-specific heads, and uses a recalibrated global head as a teacher during local training. Across MNIST, Fashion-MNIST, CIFAR10, and CIFAR100 under class-wise Dirichlet partitions, pFedKDH obtains the best accuracy in most settings, with accuracy gaps up to 37.67% over the weakest baseline and consistently low standard deviation across repetitions. Component-wise diagnostics and convergence results support the role of persistent heads and distillation-guided local optimization under label-skewed data.

Abstract

Federated Learning

Explore similar work

Personalized Federated Learning through Global Knowledge Distillation and Local Head Adaptation

Fed-BAC: Federated Bandit-Guided Additive Clustering in Hierarchical Federated Learning

FedDES: Graph-Based Dynamic Ensemble Selection for Personalized Federated Learning