FedDr+: Stabilizing Dot-regression with Global Feature Distillation for Federated Learning
Authors: Seongyoon Kim
Minchan Jeong
Sungnyun Kim
Sungwoo Cho
Sumyeong Ahn
Se-Young Yun
Venue: Accepted by TMLR ()

Abstract: Federated Learning (FL) has emerged as a pivotal framework for the development of effective global models (global FL) or personalized models (personalized FL) across clients with heterogeneous, non-iid data distribution. A key challenge in FL is client drift, where data heterogeneity impedes the aggregation of scattered knowledge. Recent studies have tackled the client drift issue by identifying significant divergence in the last linear (classifier) layer. To mitigate this divergence, strategies such as freezing the classifier weights and aligning the feature extractor accordingly have proven effective. Although the local alignment between classifier and feature extractor has been studied as a crucial factor in FL, we observe that it may lead the model to overemphasize the observed classes and underestimate the unobserved classes within each client. Therefore, our goals are twofold: (1) improving local alignment and (2) maintaining the representation of unseen class samples, ensuring that the solution seamlessly incorporates knowledge from individual clients, thus enhancing performance in both global and personalized FL. To achieve this, we introduce a novel algorithm named FedDr+, which empowers local model alignment using dot-regression loss. FedDr+ freezes the classifier as a simplex ETF to align the features and improves aggregated global models by employing a feature distillation mechanism to retain information about unseen/missing classes. Our empirical results demonstrate that FedDr+ not only outperforms methods with a frozen classifier but also surpasses other state-of-the-art approaches, ensuring robust performance across diverse data distributions.

Published in Transactions on Machine Learning Research (03/2025)
FedDr+: Stabilizing Dot-regression with Global Feature Dis- tillation for Federated Learning
Seongyoon Kim Dept. ISysE, KAIST
curisam@kaist..c.k
Minchan Jcong
mejeongSkaastt.ou oc..r
KAIST AI
Sungnyun Kim
ksm..897@kaisttoo..k
KAIST AI
Sungwoo Cho
peter85260kaish..cc.
KAIST AI
Sumyeong Ahn"
sumyenggonn@hentech.ookk
KENTECH, Department of Energy Engineering / Energy A1
Sc- Yun'
yunseryoung@kaiss.ccckk
KAIST Af
Reviewed on
Abstract
Federated Learning (FL) has emerged as pivotal framework for the development of ef- fective global models (global FL) or personalized models (personalized FL) across clients with heterogeneous, non-iid data distribution. A key challenge in FL is client drift, where data hetcrogeneity impedes the aggregation of scattered knowledge. Recent studics have lackled the client drift issue by identilying significant divergence in the last linear (classi- fier) layer. To mitigate this divergence, strategies such as freezing the classifier weights and aligning the feature extractor accordingly have proven effective. Although the local align- ment between classifier and feature extractor has been studied observe that it may lead the model to overemphasize the observed classes and underestimate the unobserved classes within each client. Therefore, our goals are twofold: (1) improving local aligument and (2) maintaining the representation of unseen class samples, ensuring that the solution seamlessly incorporates knowledge from individual clients, thus enhancing performance in both global and personalized FI. To achieve this, we introduce novel al- gorithm named FedDr+, which empowers local model alignment using dot-regression loss. FedDr+ freezes the classifier
crucial factor in FL, we
simplex ETF to align the features and improves aggregated global models by employing feature distillation mechanism to retain information about unseen /missing classes. Our empirical results demonstrate that FedDr+ not only outper- forms methods with a frozen classifier but also surpasses other state-of-the-art approaches, ensuring robust performance across diverse data distributions. The code available at: http://a...oomeeessa//////a//e DaaaaammFFFDDDalaaeeee plus.
Corresponding authors. This work was done while Sumyeong Ahn was KAIST.

---

Published in Transactions on Machine Learning Research (03/2025)
1 Introduction
Federated Learning (FI.) (MeMahan et al., 2017; He el al., 2020b) is a distributed learning strategy that enables multiple clients to collaboratively train a model while preserving data privacy. The foundational method, FedAvg (MeMahan et al., 2017), involves distributing a global model, training local models on each client's private data, and aggregating these models without transmitting raw data. However, major challenge in FL is data heterogeneity, or non-iidness, where differing data distributions across clients lead to client drift, hindering the convergence and effectiveness of the global model. Addressing this challenge involves improving two key aspects: local aligummnt and global knowledge preservalion. Local alignment refers to the cosine similarity between the features extracted by the local model and the classifier's 's true class vectors, computed on the client's training data, aiming to maximize alignment for improved local training. Global knowledge preservation aims to retain the global model's knowledge of rare or unobserved classes in the client's 's training data, preventing forgetting during local updates.
While both local aligmment and global knowledge preservation are essential, they have generally been studied separately. Global knowledge preservation is crucial because it prevents the model from be- coming overly biased toward the data of individual clients, ensuring decisions are based on a broader, shared understanding. This approach enables better generalization across all clients, particularly for unseen classes (Lee et al., 2022; 2024). The challenge of balancing global and local knowledge in FI. resembles Catastrophic Forgetting in Continual Learning (CI.) (MeCloskey & Cohen, 1989), where learning new tasks can cause models to forget previously learned ones. To achieve this, strategies like FedProx (Li et al., 2020), MOON (Li et al., 2021a), and FedNTD (Lee et al., 2022) integrate global model regularization during local training. These methods align local models with the global objective through techniques like proximal terms, contrastive learning, and logit-based regularizers.
On the other direction, to improve the local alignment, recent studies have extensively fo- cused on freezing the last linear layer (called classifier) while updating only the feature ex- tractor. This approach is motivated by the fact that the classifier is most sensitive to data het- erogeneity (Luo et al., 2021; Li et al., 2023b; Fan et al., 2024); freezing it ensures that all local models aligu their features to a consis- tent classifier across clients. For example, Fed- BABU (Oh et al., 2022) fixes the classifier after initialization, allowing only the feature extrac- for to adapt. Other mothods (Dong el al., 2022; Li et al., 2023b; Fan et al., 2024; Huang et al., 2023; Xiao et al., 2024) further enhance local alignment by modifying the loss function or by using robust initialization techniques, such as the Equiangular Tight Frame (ETF) classifier.
Feature Extractor
f(x,00)
FD
Drt = BXLDR + (1 (1-11) -1- B) LFD
09
Ox,
Local Model
Global Model
ETF CLF
Feature Distillation
CE
C1
True label
L
C1
Pushing
C Pulling Not-brue
C4
C2
C4
Feature
C3 Discard "Pushing"
Frozen
Frozen'Clasiffen Classifier
Figure 1 Overview of the proposed method, FedDr+ trained with CD To enhance the local alignment, we employ dot- regression loss CDR, which discards the pushing term of cross-entropy loss, and propose a feature distillation LFD to preserve the knowledge imbued in the global model.
A frozen classifier is also extensively explored in other research areas, such as class imbal- ance (Yang et al., 2022) and class incremental learning (Yang et al., 2023), with consistent objective similar to aforementioned FL studies by introducing and utilizing novel type of loss, called dot-regression loss LDR, which aims to achieve align- ment rapidly. In summary, LDR originates from the decomposition analysis of cross-entropy (CE) loss, which includes pulling and pushing. As suggested in (Yang et al., 2022), the pulling component is a force that attracts fcatures to the target class, whereas the pushing component force that drives fcatures away from other non-target classes. LDR discards the pushing component, as it slows down convergence to the desired alignment (refer to Figure 1).
enhancing alignment. Recently, these fields have advanced
2

---

Published in Transactions on Machine Learning Research (03/2025)
However, our findings indicate that while dot-regression loss enhances local alignment intended, it does not lead to sufficient. performance improvement. of the aggregated server-side model. The main issue arises from the insufficient global knowledge preservation of unobserved classes during local training. Specifically, the focus on improving alignment for classes present in the local training dataset induces significant feature dynamics, which inadvertently disrupt the representation of features associated with unobserved classes. This disruption leads to forgetting and deteriorated alignment for these classes, highlighting the need for better preservation of global knowledge during local training.
We emphasize the importance of addressing both local alignment and global knowledge preservation together. FedDr+ provides a well-generalized global model by combining dot-regression loss with feature distillation, reducing the distance between feature vectors of local and global models. FedDr+ FT extends this by fine- tuning the FedDr+ global model using the same FedDr+ loss function, enhancing local alignment for client- specific data. Starting with a well-trained global model is essential for achieving effective personalization while maintaining global generalization (Nguyen et al., 2022; Chen et al., 2023).
Contributions. Our main contributions are summarized follows:
In high-hoterogeenctyy FL scttings, we observe a trado-off the classifier-frociing sctup: dot-regression loss improves local aligmment with observed classes but leads to lower global model performance compared to CE loss, due to a significant loss of information on unseen classes, which is critical for the global model.
To address this, we propose FedDr+, which preserves global knowledge through feature distillation while maintaining the advantages of dot-regression loss for local alignment. This contribution focuses on im- proving global federated learning (GFL).
We extend FedDr+ to personalized federated learning (PFL) via FedDr+ FT, which fine-tunes the FedDr+ global model using the FedDr+ loss function for client-specific data. This highlights the importance of starting with a well-generalized global model for personalization.
We demonstrate the superiority of our method across various datasets and non-iid settings.
Table 1: Main notations used throughout the paper.
Indices [C]
Index for class Index for FL round
re
E [N]
Index for client
Dataset
Derain
Training dataset for client Test dataset for client i Data on client sampled from distribution D (r: input data, class label)
Diess
(x.w) E Diminncest (2x Di
Oi
Dataset consists of observed classes in client
Dataset consists of unobserved classes in client
Parameters
0
Feature extractor weight parameters Classifier weight parameters (frozen during training) c-th row vector of V All model parameters Aggregated global model parameters round
V = =lui....,eel E ROxd
Ue, c E [C]
θ (0. V)
O (0g. V)
(0..
Trained model parameters on client at round
Model Forward p[=; 0) C
Softmax probability of imput ot element p(z;0) 0) Cross-entropy loss of input Featur vector of input æ
pz(x; 0), E [C] Lee(2;; 0) = (2:
[ax; 0) E IRd
3

---

Published in Transactions on Machine Learning Research (03/2025)
2 Preliminaries
In this section, we describe the basic settings of Federated Learning (FL), including the FedAvg pipeline (McMahan et al., 2017) and the dol-regression loss (Yang et al., 2022), both of which are utilized in our framework. For clarity, we summarize the main notations in Table 1.
2.1 Basic Setup of Conventional FedAvg Pipeline
Basic FL setup. Let [N] = (1,..., denote the indices of clients, each with a unique training dataset
Dimin (Imm ym) ) m=1 mosial where (mm: Vm) D' for the 7th client, 2m is the input data, and Vo E [C] is the corresponding label among C' classos. Importantly, FI. studies predominantly address the scenario where the data distributions are heterogencous, é.e., D' varies across clients. Knowledge distributed among clients is collected over R communication rounds. The general objective of FL is to train model fit to the aggregated knowledge, UiE[N D'. This objective can be seen as solving the optimization problem:
D
min -=
train
E(zg)>DD Clz,8;0,V
EjE[M]
DD Puain
NE[N
where L is the instance-wiss loss function, 0 is the weight parameter for the feature extractor, and V
[e,....,uc) uc] E RdxC is the classifier weight matrix. We use the notation θ to denote the entire set of model
parameters.
At the beginning of each round r E [R], the server has access only subset of clients S, C [M]] participating in the pth round. At each round r, the server transmits the global model parameters Oo to the participating clients. Each client then updates the parameters with their private data Dirain and uploads O to the global server. By incorporating the locally trained weights, the server then updates the global model parameters to Oz
Fed pipeline. Our study follows the conventional Fed (McMahan al., 2017) framework address the FL problem. FedAvg updates the global model parameters from locally trained parameters by aggregating these local models into O = Eies, W O where w = Di Ejes, |rrain is the importance weight of the ,th client.
2.2 Dot- Rot-Regression Loss for Feature Alignment
Dot-regression loss LDR- This loss (Yang et al., 2022) facilitates a faster alignment of feature vectors (penultimate layer outputs) f(x; 0) E Rdd to the true class direction of eys reducing the cosine angle as follows:
LoR(a, v, V. = [
(s(z;)), -1)
where cos(ve veoz) denotes the cosine of the angle between two vectors Z(veci, vec2).
The main motivation is that the gradient of the cross-entropy (CE) loss for the feature vector be decomposed into a pulling and pushing gradient, and recent work indicates that we can achieve better convergence by removing the pushing effect (Yang et al., 2022; Li & Zhan, 2021). The pulling gradient aligns f(x; 0) with Vy, while the pushing gradient ensures f(x; 0) does not align with ve for all o y (Appendix B details the exact form of pulling and pushing gradients). Since LDR directly attracts features to the true- class classifier, it drops the pushing gradient, thereby increasing the convergence speed for maximizing cos(f(;; 0), vy).
Frozen ETF classifier. Since LDR focuses on aligning feature vectors with the true-class classifier, the classifier is not required be trained. Instead, we construct the classifier to satisfy the simplex Equiangular Tight Frame (ETF) condition, constructive way achieve maximum angular separation between class vectors (Yang et al., 2022; 2023). Concretely, we initialize the classifier weight V follows and freeze it
4

---

Published in Transactions on Machine Learning Research (03/2025)
Alignment Gap
Local Alignment
0.6
Observed Unobserved
DR
0.8
Observed
DR
CE
Unobserved
CE
0.6
0.4
0.4
0.2
0.2
0
-0.2
-0.2
100
200
300
0
100
200
300
Round
Round
(a) Alignment
(b) Local Alignment
Figure 2: Comparison of (a) feature-classifer alignment gap and (b) feature-classffie alignment on the observed and unobserved classes test data for 0% trainod with LCe and LDn-
throughout training:
V V V- Lo (Ic-
where U c is randomly initialized orthogonal matrix. Note that each in the classifier weight V satisfics cos(1; U₁) c-1
C-1 for all + [C]
3 When Dot- Rot-Regression Loss Meets FL
Given our focus on applying LDR to FL, we first examine its impact on FI. models compared to the CE loss LcE. In sunmary, we find that while CDR improves aligument-related performance on observed class labels, it faces challenge with unobserved classes which are essential for the generalization objective. To address this issue, we propose FedDr+, which integrates CDR with a novel feature distillation loss. We then evaluate FedDr+ by analyzing the effect of feature distillation and compare it with various FL algorithms and regularizers.
Experimental configuration. this section, we conduct experiments on CIFAR-100 (Krizhevsky al., 2009) with a shard non-iid setting (s=10), where each client contains at most 10 classes. We additionally employ LDA setting (a=0.1) in Section 3.4. Refer to Section 4 for more details the dataset configuration. The model is trained for 320 communication rounds, randomly selecting 10% of clients in each round, and the learning rate is decayed at t6 and 24 rounds. The experimental configuration for this section is detailed in subsection 4.1.
3.1 Impact of LoR on Local and Global Models
We analyze the average performance of local and global models trained with LDR compared to LeE, focusing their ability generalize. In Figure 2a-2b, we evaluate statistics on two datasets: the observed class set Oi, containing classes in each client's training data Dirain and the unobserved class set Ui representing unseen classes. Separately, Figure 3 reports the evaluation all classes.
First, examine the amount of change from the given global model to each local model in every commu- nication round (Figure 2a). The alignment gap is denoted by cos(f(x; 0:),, Bt) - cos(f(x; 0%-1); v.). We then evaluate the featuro-classffier alignment cos(f(x: 0;),, Uy) of cach local model on the test data (Figure 2b). Finally, we observe the lest accuracy of the global model 0g (Figure 3).
"This relation cosines holds the vt's are symmetrically distributed such that ã 1 E vi = 0, and cos(u,, ey) are all the same for $ While use the " this context, also applies "rarely" existing classes.
5

---

Published in Transactions on Machine Learning Research (03/2025)
Alignment analysis of local models. As shown in Figure 2a-2b, LDR outperforms Lce on observed classes in terms of alignment gap and aligmment, while Lce achieves better results on unobserved classes for these metrics. The improvement on observed classes is attributed to CDR, which removes the pushing term present in LCE and concentrates its pulling effects on these classes. However, this design inherently overlooks unobserved classes, resulting in poorer performance on these classes compared to LCB. +
Accuracy result of global models. Figure 3 shows that in the shard setting (s = 10), Lce consistently outperforms LDR. This difference is due to the higher proportion of un- observed in this setting, where cach client has access to at most 10 out of 100 classes.
As observed in the alignment analysis, LDR performs par- ticularly poorly on the unobserved, with significantly negative alignment gap and lower aligument compared to models trained with LCE, which contributes to the lower overall accuracy of the global model when using CDR.
60 CE
Global Accuracy
DR
40
20
Importance of local alignment in observed classes. While LDR struggles with local alignment on unobserved classes, leading to lower global accuracy com- pared to LCE, local alignment in observed classes remains critical. At the last learming rate docay round (240th, both methods achieve higher or comparable performance unobserved classes than in the final round (320th)
0
100
200
300
Round
Figure 3: Comparison of the global test accuracy
of 0 on all classes trained using Lce and LDR-
However, the final round shows higher global accuracy due to improved local aligument on observed classes This underscores the need to preserve the advantages of LDR in observed classes while mitigating its degra- dation on unobserved classes.
3.2 FedDr+: CDR With Feature Distillation for FL
We propose FedDr+ to mitigate forgetting unobserved classes while retaining the strengths of dot-regression loss in aligning features of observed classes. Using LDR with the frozen classifier V, FedDr+ includes a regularizer that fully distills the global model's feature vectors [(x; 09) Rd to the client features f(x;; 0), to enhance generalization across all classes. The proposed loss function LD shown in (1), combines LDa with a regularizer LFD(I; 0, 09) = {=; 0)- [(;; 09) Unless specified, we use a scaling parameter 3 = 0.9
throughout the paper. The overall pseudocode of FedDr+ can be found in Appendix A.
0", V) = 3 (DR(I, 0, V) + (1 - 0,; L (2,1; 0,,
(1)
Why feature distillation? To address data heterogeneity in FL, various distillation methods have been explored, including model parameters (Oh et al., 2022; Li et al., 2020; He et al., 2020a; Li & Wang, 2019), logit-related measurement. (Ti & Wang, 2019; Tee et al., 2022; Ttahara et al., 2021; Ye et al., 2024; Tin et al., 2020; Chen et al., 2019; Qian et al., 2022), and co-distillation (Chen et al., 2024; Cho et al., 2023). In contrast, we utilize the feature distillation (Heo et al., 2019) technique because the feature directly concerns alignment. On the oer hand, logits lose information from features when projected onto a frozen ETF classifier (Heo et al., 2019; Li et al., 2017; 2023a; Ben- Baruch et al., 2022). By distilling features, we leverage the global, differentiated knowledge for each data input x. This approach aims to minimize drift towards observed classes, and hence, we expect it to enhance overall generalization.
3.3 Effect of Feature Distillation
Our findings from Section 3.1 indicate that CDR is unsuitable for the heterogeneous FL environment. This is primarily because there is notable gap in how features align with the fixed classifier between Oi and
AA theoretical justification for why LDR struggles with unobserved classes provided in Appendix C, where analyze feature gradients under the NTK framework.
6

---

Published in Transactions on Machine Learning Research (03/2025)
15
Observed Unobserved
DR
Observed
DR
Dr+
Feat. Angle Dist.
80
Unobserved
Dr+
Feat. Dist.
10
60
5
40
0
100
200
300
100
200
300
Round
Round
(a) Feature distance
(b) Feature angle distance
Figure 4:
We present (a) feature distance and (b) feature angle distance from 0 -- to 0¹. for observed and unobserved classes by training with LDR and LD+.
Local Alignment
Global Accuracy
0.6
Observed
CE
DR
Dr+
60
Unobserved
0.4
DR
40
Dr+
0.2
CE
20
0
0
0
100
200
100
200
300
Round
Round
Figure Comparison feature-classifier alignment Figure 6: Comparison of the global test accuracy of on the observed and unobserved classes test data for 0 on all classes trained using LCE, LDR and CDr+. of trained with LCE, LDn and LDt+-
Wt. To assess the effect of feature distillation (CrD), which imposes a constraint on the feature distance [>; 03) [x; 07-1)11 for E O³, measure this distance for both O' and 4 from the models trained with LDR and LD++. We additionally analyze the angle distance, (af(;; 0²,,, f(x; 0z_ it impacts feature- classifier aligunent. These values are averaged over the selected client set S,..
Feature distillation stabilizes the feature dynamics. By adding LFD, as revcaled in Figure 4a, the local model trained with LDD shows reduction in feature distance for observed classes, compared to the model trained with LDR- This reduction happens even for unobserved classes. As demonstrated in Figure 4b, the overall decrease in feature distance leads to reduction in feature angle distance for both class sets. In both local models trained with LDR and Lpt+. there is trend where the angle distance is significantly larger for 212 than for Of (Figure 4b). This large angle distance of ltd leads to the degradation of the feature-classifee alignment.
Feature distillation enhances local alignment and global accuracy. As in Figure 5 6, our proposed algorithm, i.e., LDt+, improves featuro-classifer alignment for both O' and 4t along with enhanced global accuracy compared to LDR. We altribute this improvement to the enhanced knowledge of the global model which is preserved by preventing the forgetting of previously trained knowledge. During the final convergence phase 2440th3300hh LD+ achieves better fcature-classffier alignment even on unobscrved classes compared LeE. As result, the final round (320th), LDt+ demonstrates the best local alignment across all classes, surpassing both LCE and LDR, contributing to its global model's superior performance. Even though the

---

Published in Transactions on Machine Learning Research (03/2025)
Table 2: Synergy of various FL algorithms and regularizers. Baseline indicates training FL models without regularizer. FD denotes feature distillation, which is the regularizer we use in FedDr+.
Sharding (a = 10)
LDA (o m
0.1)
Algorithun
Baseline
+ Prox
+KD
+NTD
+LD
+MOON
+FD
Buseline
+Prox
+KD
-NTD
+LD
+MONN
+DD
FedAvg
37.22
36.87
36.25
37.71
37.17
37.43
37.82
42.32
43,22
44.21
43.39
43.43
44.79
43.76
FedBABU
46,20
46.03
46.37
17:22
46.71
46.49
16.95
17.37
46.62
47.60
16.48
45.78
46.27
16.
SphereFed
43,90
44.94
43:47
43.95
43.13
45.21
46.98
43.77
47.76
47.25
47.01
46.81
49.74
FoIETF
32.42
31.87
32.76
32.65
32.25
34.30
32.77
46.27
45.71
46.67
46.16
45.91
45.98
48.47
FddGELA
29.17
28.09
29.11
28.84
29.36
28.80
30.33
27:11
29.03
28. 45
29.62
29.41
28.09
29.75 75.75
Pot-Regression
42.52
41.95
47.45
48.32
47.52
44.72
48.69
42.72
46.35
49.47
50.36
49.28
50.36
50.86
proposed regularizer demonstrates reasonable regularizing effect, question remains: "Is it superior to other previously used regularizers!"
3.4 Synergistic Effect with Different Types of FL Algorithms and Regularizers
We answer the above question by evaluating the synergy effect of various FL algorithms by maintaining their original training loss and incorporating specific regularizers, following the approach suggested in Equation 1. To address the issue of differing loss scales between the baseline FL algorithms and the regularizers, we thoroughly tune the coefficient within the range 0.3, 0.5, 0.7, 0.9, 0.99, 0.999, 0.9999), and report the resulting performance in Table 2. The selected 3 values are detailed in Appendix E.
FL algorithms and regularizers. We evaluate several baseline FL algorithms, including dot-regression. The baseline also include FodAvg (McMahan et al., 2017)-using an unfrozen classifier, FedBABU (Oh et al., 2022) exlending FedAvg by freezing the classifier during local training, and SphereFed et al., 2022)0022eaee enhancing feature-classifier alignment by using MSE loss between one-hot encoded labels and cosine similarity based logits. We also considered FodGELA and FedETF-bor applying client-specific adaptive loss functions tailored to data distribution.
Alongside the FD regularizer, we evaluate a range of regularizers, including Prox (Li et al., 2020) con- straining the distance between local and global model parameters; MOON (Li et al., 2021a) minimizing the angle distance between feature vectors of global and local models through contrastive learning; and several logit-based regularizers- KD (Hinton et al., 2015), NTD (Lee et al., 2022; Zhao et al., 2022), and LD (Kim et al., 2021) keeping logit-related measurements of local models closely aligned with the global model. Specifically, KD applies softened softmax probability from the logit vector, NTD does the same but excludes the true class dimension, and LD distills the entire logit vector.
FedDr+ shows best synergy. Table shows that FedDr+ (dot-regression + FD) achieves the strongest performance among the tested combinations. FD performs exceptionally well when combined with dot- regression, SphercFod, and FodBABU. These baselines freeze the classifier while not using client-speeific adaptive loss. Among them, the synergy is most effective in the order of dot-regression, SphereFed, and RedBABB-releectng how each optimizes feature-classffier alignment well. FD outperforms other regularizers by effectively sta stabilizing feature dynamics for observed classes, which helps mitigate the misalignment and performance degradation for unobserved classes. MOON, on the other hand, prioritizes the cosine similarity between feature vectors from local and global models but fails to adequately control feature norm dynamics. Prox applics uniform regularization that independent of specific data instances, leading to less refined control over feature dynamies. Logit-based regularizers lose effectiveness due to information loss when features are projected onto the classifier, they focus mitigating the dynamics of the less informative projected vectors rather than the richer original feature vectors.
For baselines like FedGELA and FedETF, which apply client-specific loss functions, none of the regularizers, including FD, lead to consistently lead to significant performance, particularly in the sharding setting. In non-frozen scttings, FodAvg, the FD regularizer does not offer significant performance boost. Under LDA, FedAvg combined with MOON outperforms FodAvg with FD, consistent with the claim of MOON (Li et al., 2021a).
8

---

Published in Transactions on Machine Learning Research (03/2025)
4 Experiments and Results
In this section, we first present the experimental results of FedDr+ in the context of global lederated learning (GFL). We then analyze the elapsed time and conduct sensitivity analysis for GFL, investigating the effects of varying local epochs, client sampling ratios, and different 8 values on the performance of FedDr+. Finally, we propose FedDr+ FT, which fine-tunes the FedDr+ GFL model with LDr+, and report the results comparing it with existing PFL methods.
4.1 Experimental Setup
Dataset and models. To simulate a realistic FL scenario involving 100 clients, we conduct extensive studies on three widely used datasets: CIFAR-10 (Krizhevsky et al., 2009), CIFAR-100 (Krizhevsky et al., 2009) and ImageNet-100 (Deng et al., 2009). We use VGG11 (Simonyan & Zisserman, 2014) for CIFAR-10, MobileNet (Howard al., 2017) for CIFAR-100, and ResNet-18 (He et al., 2016) for ImageNet-100, The training data is distributed among the 100 clients using sharding and the LDA (Latent Dirichlet Allocation) partition strategies.
Following the convention, sharding distributes the data into non-overlapping shards of equal size, each shard encompassing [Desal and [Daesl samples per class, where denotes the number of shards per 100xs client. On the other hand, LDA involves sampling probability vector from Dirichlet distribution, Pe (pe,1.Pe,,, ... Pe,,000 ~ Dir(a), and allocating a proportion Pe,k of instances of class E [C) to each client k E [100]. Smaller values of and a increase the level of data heterogeneity. For CIFAR-10 and CIFAR- 100, we explore a range of s and a values to assess the impact of different data heterogeneity levels. For ImageNet-100, we focus on experiments with 8 and = 0.1.
100xs
Implementation details. In each round of communication, a random 10% of clients are selected to participate in the training process. The total number of communication rounds is set to 320. The initial learning rate and the number of local epochs for CIFAR-10, CIFAR-100, and ImageNet-100 100 are determined through grid searches, with the detailed process and results provided in Appendix D. The learning rate rj is decayed by a factor of 0.1 the 160th and 240th communication rounds.
4.2 Global Federated Learning Results
We compare FedDr+ with range of GFL algorithms, considering both non-freezing and freezing classifier approaches. Among non-freezing classifiers, FedDr+ competes with FedAvg (McMahan et al., 2017), Fed- Prox (Li et al., 2020), SCAFFOLD (Karimireddy et al., 2020), MOON (Li et al., 2021a), FedNTD (Lee et al., 2022), FedExP (Jhunjhunwala et al., 2023), and FedSOL. (Lee et al., 2024). FedDr+ is also evaluated against freezing classifier algorithms such as FedBABU (Oh et al., 2022), SphereFed (Dong et al., 2022), FedETF (Li et al., 2023b), and FedGELA (Fan al., 2024). Among the bascline algorithms, SCAFFOLD incurs communication cost two times higher per round, denoted as (x2). Our experiments encompass heterogencous settings involving sharding and LDA non-IID environments.
Table 3 summarizes the accuracy comparison between various GFL methods proposed in recent literature and FedAvg under various conditions. While specific methods demonstrated effectiveness in particular scenarios, some of these underperformed relative to the robustness FedAvg. For cxample, SCAFFOLD shown strong performance in the less heterogeneous sharding setting on CIFAR-10; however, it failed model training under the highly heterogeneous LDA condition with 0.1. Notably, FedDr+ consistently outperformed all baselines in highly heterogeneous settings, achieving a 3.15% improvement in CIFAR-100 LDA with a 0.05 and 3.17% in ImageNet-100 LDA.
4.3 Elapsed Time Results
We compare FedDr+ with various GFL algorithms for the clapsed time per communication round on CIFAR- 100 (s=10). As shown in Table 4, incorporating a global model during the local update generally results in higher computation costs, leading to longer elapsed times compared to updates without. global model.
9

---

Published in Transactions on Machine Learning Research (03/2025)
Table 3: Accuracy comparison in the GFL setting. The entries are based on results obtained from three different seeds, indicating the mean and standard deviation of the accuracy of the global model, represented as X+v. The best performance in each case is highlighted in bold.
NIID Partition Strategy: Sharding
CIFAR-100
CIFAR-10
ImageNet-100
10
8=20
8=50
s= 100
a=2
8=5
s= 10
FedAvg
36. 63 0.22
42.25+ 1:42
45.57+ o.
48. 18.20+ 1:30
72. 08 0.57
81.53= u.ss
82.38= 0.40
67.78= 0.41
FedProx
37. 07+ 0.21
42.35+ 9.53
45.18+ 1.07
47.78+ 0.79
71.92= 0.51
81.29= 0.40
82.45+ 0.35
07.81+ 0.55
SCAFFOLD (x2)
08 0.57
18.15++ 1.21
49.31+ 0.02
50.73+ 0.42
75.49- 0.42
84.14+ 0.13
85.11+ 0.29
70.47+ 0.45
MOON
36.95 0.37
43. 13.00+ 0.27
43.95 0.12
46.92 c.08
67.35+ 1.16
80.90 0.26
82.62 0.31
08.19+ 0.32
FedNTD
34.05 1.19
41.78 0.31
46.42 0.5a
47.17 0.33
72.21 0.59
69.96
17.10
81.99 0.42
67.51 0.25
FedExP
36.85a 0.11
42.49 1:20
45.07 0:99
48.09 1:00
72.31 0.00
81.41 0.19
82.47 0.16
63.34 0.53
FedSOL
32.182 0.18
41.54g 1.00
47.421 0.76
47.70a 1.1a
54.93 a.5a
75.73g 0.10
77.00s 0.41
66.61 1.17
FedBABU
45.97+
45.53+ 0.79
46.52+
46,02 0:28
71.99+
81.07- 0.60
82.32
68,82
SphereFed
42.71
48.63
52.16: 0.3a
53.41:
76.33 0.3a
83.67=
84.36
69.71
FedETF
31.37
42.221
47.47
49.00
67.81;
80.78
82.60s
70.81 0.25
FedGELA
27. 97.95
38.631
44.67
47.95
63.774
79.05
81.56
67.08a
FedDr+ (Ours)
48.21
50.77
52.15a
52.41
76.571
83.22z 0.84
84.14g
71.471
NIID Partition Strategy: LDA
CIFAR-100
CIFAR-10
ImageNet-100
a=0.05
a=0.1
a=0.2
a=0.3
a=0.1
d=002
<=0.3
FedAvg
35. 58 1.55
42.10+ 0.00
44.78+ 0.72
45.73+ c:s8
68. 71+ 1.52
77.75- 0.25
80.76+ 0.51
05.11+ 0.25
FedProx
37.07+ 0.21
42.35+ p.sa
45.18+ 1:00
48.18+ 0.51
69.00 2.27
77.81
0.24
80.55+ 0.19
40.48 1.25
SCAFFOLD (x2)
40.54 0.45
46.14 a.mo
47.98 0:90
48. 33.00: 1:00
(Failed)
80.15 0.29
82.63 0.23
66.84 0.77
MOON
23.97 1.15
30.86, 0:21
33.60 0.ca
35.54 6:45
66.44 aiss
77.36 0.0s
80. 0.02
41.25 0.60
FedNTD
31.78 2.14
40.41 4.1...... 0:96
43.10. 2:00
43.04 o.sa
70.22 0.40
77.16 0.20
79.50 0.50
64.87 0.30
FedExP
34.39 1.77
40.85 1.30
44.47 0.98
45.44 0.14
70. 14g 0.53
78.09 0.21
80.40 0.54
59.40z 0.35
FedSOL.
34.49z 0.50
41.19a o
43.55 1.
44.85 0.54
59.51 1.97
67.55g 0.41
70.96 0.92
62.70z 0.59
FedBABU
41.97
45.77
44.28
44.80 oica
65. 15g
77.03
79.91
66.54
SphereFed
39.56g
46.54 bisa
49.41:
49.22
67.49,
80.05
82.621 0.60
67.031 0.30
FedETF
40.71
45.631 0.3a
46 (6.2881
46.691
70.75
77.86_
79.95
68.984
FedGELA
16.72s
27.121
33.682
36.171
50.692
66.04_
77.892
55.57
FedDr+ (Ours)
45.12a 1:00
49.48
50.67 0:88
51.15
72.07
80.90
82.422 0.10
70.20
Table 4: Elapsed time per round (in seconds) for various GFL algorithms.
Local update without global model
Local update with global model
FedAvg
FedRABU
FoIETF
SphereFed FeaGELA
FedExP
FedProx SCAFPOLD
FoINTD
FedSOr.
MOON Fedbre (Ours)
Elapsed time
20.4
20.5
20.7
20.4
20.5
20.1
22.6
21.4
21.2
27.2
56.7
24.9
FedDr+ exhibits a slightly longer clapsod time than most other algorithms, yet still requires less time than FedSOL and MOON.
4.4 Sensitivity Analysis
We explore the impact of varying client sampling ratio and local epochs on performance, as well as the effect of different & values in FedDr+, as detailed in Figure 7. All experiments are conducted on MobileNet using the CIFAR-100 dataset with LDA (a=0.1).
Effect of client sampling ratio and local epochs. We evaluate the sensitivity of hyperparameters in FedDr+ by comparing it to baselines under varying client sampling ratio and local epochs, starting from the default setting of client sampling ratio of 0.1 and local epoch of 3. Compared to FedAvg (without classifier freezing), FodBABU and SphereFed (all with classifier freczing) generally show performance improvements with increasing fraction ratios, but FedDr+ consistently outperforms the basclines. The number of local epochs is crucial in FL, too few epochs result in underfitting, while too many cause client drift, degrading
10

---

Published in Transactions on Machine Learning Research (03/2025)
+ FedAvg
FedBABU
SphereFed
FedDr+
Acc (LDA, a=0.1)
Acs (LDA, a=0.1)
50
Acc (LDA. a=0.1)
40
50
40
45
20
10
0.5
1.0
5
1.0
B
(a) Client sampling ratio
(b) Local epochs
(c) / sensitivity
Figure 7: Performance of baselines and FedDr+ on CIFAR-100 (a=0.1) with various analyses: (a) client sampling ratio, (b) the number of local epochs, and (c) sensitivity B.
Table 5: PFL accuracy comparison with MobileNet on CIFAR-100. Results are reported in the format Xty, representing the mean and standard deviation of the average personalized accuracies across all clients, computed over five seeds. The best performance in each case highlighted in bold.
Algorithm
8=10
8=20
8=100
a=0.05
a=0.11
a=0.3
Local only (Cce)
58.42; 0.22
42.37 0.29
19.02: 0.34
55.71 0.19
44.02: 0.39
27.98 0.0m
Local only (Ec=+ETF)
58.05: 0.25
41.721 0.25
19.06 0.18
55.41 0.23
43.561 0.31
27 27.691 0.17
Local only (CDR)
61. 11.05=
44.28
21.06=
58.56+
47.05= 0.15
31.16= 0.15
FodPer
70.62+
55 35.65+
25.57=
63.35+
51.90= 2.15
35.84+ 2.10
Por-FodAvg
31.71+ 1.08
38.64+ 0.40
45.71= 0.81
28.85 0.27
36.00+ 0.42
42.41 0.32
FodRep
62.59+ +9++0..0 0.30
51.18+ 1.00
26.51+ 0.27
57.734 73+ 0.43
49.59+ 0.40
36 16,22+ 0.89
Ditto
38.39+
42.16+
44.04=
34.86+
38.67+
42.05
FedAvg-FT
70.20+ 0.54
56.26+ 0.51
48.67 0.99
61.08+ 1.86
56.34+ 1.18
49.74+ 1.05
FedBABU-FT
80.73 0.65
71.02 0.54
51.70= 0.21
76. 12 0.55
69.94 0.34
57.40 1.50
SphereFed-FT
81.34
72.22:
56.58
74.49
69.39
59.51
FedETF-FT
53.32a 0.60
53.05
49.74= 0.85
32.31 0.40
53.70. 0.35
50.80 0.05
RaiGELA-FT
75.75 0.57
68.96 0.87
52.23 0.50
58.26 6.78
60.12z 0.71
53.09 0.82
FedDr+ FT (ours)
83.08 o.ar
74.80 0.08
56.56= 1.04
78.40 0.40
73.23 0.80
62.22: 0.98
global model performance. The default setting of local epochs 3 is optimal for all baselines, with FedDr+ achieving the best performance. Although performance generally declines when deviating from this peak point, FedDr+ remains the best highly competitive.
Weight ratio 3 analysis. We analyze the effect of scaling parameter in FedDr+ by varying while keeping other hyperparameters constant. The performance is evaluated for 8 E 0.,, 0.5, 0.7, 0.9, 1.0). When B 0, only feature distillation is applied, and when 3 1, only dot-regression is used. 3 E {0, 1) are generally less effective, whereas 8 0.7,0. show consistently good performance, indicating a balanced approach is beneficial.
4.5 Personalized Federated Learning Results
We introduce FedDr+ FT, inspired by prior work (Oh al., 2022; Dong al., 2022; Li et al., 2023b; Kim al., 2023; Fan et al., 2024), which enhances personalization by leveraging local data to fine-tune the global federated learning (GFL) model. We fine-tune the FedDr+ GFL model using LDt+ create FedDr+FT,, &e.,, 2-step approach. The overall pseudocode of FedDr+FT can be found in Appendix A. For comprchensive analysis, we compare FedDr+ FT with existing personalized foderated lcarning (PFL) methods, including 1-step approaches, f.e., creating PFL models from scratch, such as FedPer (Arivazhagan et al., 2019), Per-FedAvg (Fallah et al., 2020), FedRep (Collins et al., 2021), and Ditto (Li et al., 2021b), as well as 2-step methods such as FedAVG-FT, FodBABU-FT (Oh al., 2022), SphercFed-FT (Dong et al., 2022), FodETF-FT (Li et al., 2023b), and FedGELA-FT (Fan al., 2024). Additionally, we compare these methods with various simple local models that have not undergone federated learning: (1) Local only (CCE),
11

---

Published in Transactions on Machine Learning Research (03/2025)
trained with LeE, (2) Local only (CCE + ETF), trained with Lee and initializing the classifier with an ETF classifier, and (3) Local only (LDR), trained using LDR-
In Table 5, we first compare the performance of simple local models in PFL by examining CDR and LCE. While methods using LC show no significant differences, utilizing CDR leads to substantial performance improvements in PFL across all settings. The *Local only (ECr)" and "Local only (CCE + +ETF)* methods exhibit similar performance due to the nearly classwise orthogonal nature of randomly initialized classi- fiers (Oh et al., 2022; Saxe et al., 2013; Glorot & Bengio, 2010a; He et al., 2015; Lezama al., 2018). With large number of classes (C=100), the ETF classifier, which is also nearly classwise orthogonal, performs similarly to random initialization. When comparing FedDr+ FT with other 2-step methods, FedDr+ FT con- sistently demonstrates superior performance. This aligns with previous research (Nguyen et al., 2022; Chen et al., 2023) suggesting that fine-tuning from a well-initialized model yields better PFL performance. Ad- ditionally, compared with 1-step algorithms, FedDr+ FT continues to show superiority, outperforming all bascline methods across all settings.
5 Related Work
Federated learning. Federated Learning (FL) is a decentralized approach to deep learning where multiple clients collaboratively train a global model using their own datasets (MeMahan et al., 2017; Li et al., 2020). This approach faces challenges due to data heterogeneity across clients, causing instability in the learning process (Karimireddy et al., 2020; Luo et al., 2021). To address this problem, strategies like classifier variance reduction in FedPVR (Li al., 2022) and virtual features in CCVR (Luo et al., 2021) have been proposed. Additionally, it is essential to distinguish between Global Federated Learning (GFL) and Personalized Fed- erated Learning (PFL), as these are crucial concepts in FL. GFL aims to improve single global model's performance across clients by addressing data heterogeneity through methods like client drift mitigation (Li et al., 2020; Karimireddy et al., 2020; Jhunjhunwala et al., 2023), enhanced aggregation schemes (Wang et al., 2020a;b), and data sharing techniques using public or syuthesized datasets (Lin et al., 2020; Luo et al., 2021). Otherwise, PFT. focuses on creating personalized models for individual clients by decoupling feature extractors and classifiers for unique updates (Oh et al., 2022; Arivazhagan et al., 2019; Collins et al., 2021), modifying local loss functions (Fallah et al., 2020; Li et al., 2021b), and using prototype communication techniques (Tan et al., 2022; Xu et al., 2023).
Frozen classifter in FL. By focusing on aligument, previous studies have attempled to mitigate data heterogeneity by froezing the classifier (Oh et al., 2022; Dong et al., 2022; Li et al., 2023b). Nevertheless, these methods have yet to effectively improve the alignment between features and their corresponding classifier weights. Motivated by this, we integrated the dot-regression method into FL to achieve better-aligned local model by freezing the classifier. Dot-regression, proposed to address class imbalance, focuses on aligning feature vectors to a fixed classifier, demonstrating superior aligmment performance compared to previous approaches. However, optimizing the dot-regression loss to align feature vectors with fixed classifier caused the local model to lose information on unobserved classes, thereby degrading global model performance. To address these issues, FodLoGe (Xiao et al., 2024) employing realignment tochniques to ensure the well-aligned local model's performance translated to the global model. Additionally, in FedGELA (Fan el al., 2024), the classifier is globally fixed as simplex ETF while being locally adapted to personal distributions. Also, FedPAC Xu et al. (2023) addressed these challenges by leveraging global semantic knowledge for explicit local-global feature alignment. Besides alignment-focusee methods, there have been various attempts to maintain good local model performance in the global model Jiang et al. (2023); An et al. (2024); Chen & Chao (2022).
Knowledge distillation in FL. Knowledge distillation (KD) has been widely studied in FL settings, such as in FedMD (Li & Wang, 2019) and FedDF (Lin et al., 2020), where pretrained teacher model transfers knowledge to a student model. Additional distillation-based methods, such as FedFed (Yang el al., 2024) and co-distillation framework for PFL (Chen et al., 2024; Cho et al., 2023), have also been explored. In contrast to existing mothods, we propose a loss function incorporating fcature distillation to maintain the performance of both local and global models. To our knowledge, this is the first application of feature distillation in FL. This approach highlights the importance of distinguishing between GFL and PFI..
12

---

Published in Transactions on Machine Learning Research (03/2025)
6 Limitations
While our proposed method, FedDrt, effectively enhances both local alignment and global knowledge preser- vation in Federated Learning (FL), it has certain limitations. First, our approach builds upon dol-regression to improve local aligmment, but this is just one possible strategy. Alternative methods, such as directly max- imizing local alignment without relying on dot-regression, could explored to further enhance performance in FL scttings. Second, although FedDr+ effectively mitigates forgetting of unobserved classes by incorpo- rating feature distillation, dot-regression alone remains less effective in preserving alignment for unobserved classes. While empirical results demonstrate that FedDr+ alleviates this issuc, further theoretical in- vestigation is needed to develop more principled approach to ensuring aligument across both observed and unobserved classes. These limitations highlight opportunities for future work to extend and refine our method, improving its robustness and generalizability in diverse FL environments.
7 Conclusion
Motivated by the recent FL methods enhancing feature alignment with fixed classifier, we first investigate the effects of applying dot-regression loss for FL. Since the dot-regression is the most direct method for feature-classffier alignment, we find improves alignment and accuracy local models but degrades the performance of the global model. This happens because local clients trained with dot-regression tend forget classes that have not been observed. To address this, we propose FedDr+, combining dol-regression with feature distillation method. By regularizing the deviation of local features from global features, FedDr+ allows local models to maintain knowledge about all classes during training, thereby ultimately preserving general knowledge of the global model. Our method achieves top performance in global and personalized FL experiments, even when data is distributed unevenly across devices (non-IID settings).
Acknowledgement
This work was supported by Institute of Information & communications Technology Planning & Eval- uation (ITP) grant funded by Korca government (MSIT) [No. 2021-0-00907, Development of Adaptive and Lightweight Eola-Collaborative Analysis Technology for Enabling Proactively Immediate Response and Rapid Learning, 90%] and [No. 2019-0-00075, Artificial Intelligence Graduate School Program (KAIST), 10%].
13

---

Published in Transactions on Machine Learning Research (03/2025)
References
Xuming An, Li Shen, Han Hu, and Yong Luo. Federated learning with manifold regularization and normalized update reaggregation. Advonces in Neural Information Processing Systems, 36, 2024.
Manoj Ghuhan Arivazhagan, Vinay Aggarwal, Aaditya Kumar Singh, and Sunav Choudhary. Federated learning with personalization layers. arXiv preprint arXim:1912.00811 2019.
Yuval Belfer, Amnon Geifman, Meirav Galun, and Ronen Basri. Spectral analysis of the neural tangent 00300000
kernel for deep residual networks, 2021. URL
Emanuel Con-Baruch, Matan Karklinsky, Yossi Biton, Avi Ben-Cohen, Hussam Lawen, and Nadav Za- mir. It's all in the head: Representation knowledge distillation through classifier sharing, arXiv preprint orXiv:2201.0694: 2022.
Hong-You You Chen and Wui-Lun Chao. On bridging generic and personalized federated learning for image clas- sification. In International Conference on Learning Representations, 2022. URL https://openrevieww net/forrumiccII1h0xx000rr
Hong- You Chen, Cheng- Hao Tu, Ziwei Li, Han Wei Shen, and Wei-Lun Chao. On the importance and applicability of pre-training for federated learning. In The Fleventh International Conference on Learning Representations, 2023. URL https: / /twp///eie/eeeeeee//e///ee net net//forum?ddffNNVV-. P0XP.
Wei-Chun Chen, Chia-Che Chang, and Che-Rung Lee. Knowledge distillation with feature maps for im- age classification, In Computer Vision-ACCV 2018: 14th Asian Conference on Computer Vision, Perth, Australio, December 6, 2018, Revised Selected Papers, Port III 14, pp. 200 215. Springer, 2019.
Zihan Chen, Howard Yang, Tony Quek, and Kai Fong Ernest Chong. Spectral co-distillation for personalized federated learning. Advances in Neural Information Processing Systems, 36, 2024.
Yac Joe Cho, Jianyu Wang, Tarun Chirvolu, and Gauri Joshi. Commuicattion-dffrienttt and model- heterogeneous personalized federated learning via clustered knowledge transfer. IEEE Journal of Selected Topics in Signal Processing, 17(1):234 227, 2023.
Liam Collins, Hamed Hassani, Aryan Mokhtari, and Sanjay Shakkottai. Exploiting shared representations for personalized federated learning. In International conference on machine learning, pp. 1089-2099) PMLR, 2021.
Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large-scale hierarchical image database. In 2009 IEEE conference on computer vision and pattern recognition, pp. 248-255. Teee, 2009.
Xin Dong, Sai Qian Zhang, Ang Li, and HT Kung. Spherefed: Hyperspherical federated learning. In European Conference on Computer Vision, pp. 165-184. Springer, 2022.
Alireza Fallah, Aryan Mokhtari, and Asuman Ozdaglar. Personalized federated learning with theoretical guarantees: A model-agnostic meta-learning approach. Advances in Neural Information Processing Sys- lems, 33:3557-3566 2020.
Ziqing Fan, Jiangchao Yao, Bo Han, Ya Zhang, Yanfeng Wang, et al. Federated lcarning with bilateral curation for partially class-disjoint data. Advances in Newral Information Processing Systems, 36, 2024.
Xavier Glorot and Yoshua Bengio. Understanding the difficulty of training deep feedforward ncural networks. In Proceedings of the thirleenth international conference on artificial intelligence and slatistics, pp. 249
256. JMLR Workshop and Conference Proceedings, 2010a.
Xavier Glorot and Yoshua Bengio. Understanding the difficulty of training deep feedforward neural networks. In Yee Whye Teh and Mike Titterington (eds.), Proceedings of the Thirteenth International Conference on Artificial Intelligence and Statistics, volume of Proceedings of Machine Learning Research, pp. 249- 256, Chia Laguna Resort, Sardinia, Italy, 13-15 May 2010b. PMLR. URL https:// mlr. press/v9/glorotloa. html.
14

---

Published in Transactions on Machine Learning Research (03/2025)
Chaoyang He, Murali Annavaram, and Salman Avestimehr. Group knowledge transfer: Federated learning of large cons at the edge. Advances in Neural Information Processing Systems, 33:14068-1444 14480, 2020a.
Chaoyang He, Songze Li, Jinhyun So, Xiao Zeng, Mi Zhang, Hongyi Wang, Xiaoyang Wang, Praneeth Vepakomma, Abhishek Singh, Hang Qiu, et al. Fedml: A research library and benchmark for federated machine learning. arXiv preprint arXi::207. 13518, 2020b.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Delving deep into rectifiers: Surpassing human- level performance on imagenet classification. In Proceedings of the IEEE international conference on computer vision, pp. 1026-10344 2015.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Identity mappings in deep residual networks. In Computer Vision ECCV 2016: 14th European Conference, Amsterdam, The Netherlands, October 11 14, 2016, Proceedings, Part IV 14, pp. 630 645. Springer, 2016.
Byeongho Heo, Jeesoo Kim, Sangdoo Yun, Hyojin Park, Nojun Kwak, and Jin Young Choi. A comprehensive overhaul of feature distillation. In Proceedings of the IEEE/CVE International Conference Computer Vision, pp. 1921-1930, 2019.
Geoffrey Hinton, Oriol Vinyals, and Jeff Dean. Distilling the knowledge a neural network. arXio preprinl arXiv:1503. 02531, 2015.
Andrew G Howard, Menglong Zhu, Bo Chen, Dmitry Kalenichenko, Weijun Wang, Tobias Weyand, Marco Andreetto, and Hartwig Adam. Mobilenets: Efficient convolutional neural networks for mobile vision applications. arXiv preprint #xXiv:104.04886. 2017.
Chenxi Huang, Liang Xie, Yibo Yang, Wenxiao Wang, Binbin Lin, and Deng Cai. Neural collapse inspired federated learning with non-iid data, 2023.
Sohei Itahara, Takayuki Nishio, Yusuke Koda, Masahiro Morikura, and Koji Yamamoto. Distillation- based semi-supervised federated lcarning for communication-effiinn collaborative training with non-iid private data. IEEE Transactions on Mobile Computing, 22(1):191 205, 2021.
Arthur Jacot, Franck Gabriel, and Clement Hongler. Neural tangent kernel: Convergence and gen- cralization in neural networks. In S. Bengio, H. Wallach, H. Larochelle, K. Grauman, N. Cesa- Bianchi, and R. Garnett (eds.), Advances in Newral Information Processing Systems, volume 31. Cur- ran Associates, Inc., 2018 URL https:// 5abbe21f34662bb336c66102447767 Papee.apprpppr pdf..
Divyansh Jhunjhunwala, Shiqiang Wang, and Gauri Joshi. Fedexp: Speeding up federated averaging via extrapolation. In The Eleventh International Conference on Learning Representations, 2023. URL https:
//openreview. http:///wwwee/foooonndddo//ooo/
Meirui Jiang, Anjie Le, Xiaoxiao Li, and Qi Dou. Heterogeneous personalized federated learning by local- global updates mixing via convergence rate. In The Twelfth International Conference on Learning Repre- sentations, 2023.
Sai Praneeth Karimireddy, Satyen Kale, Mehryar Mohri, Sashank Reddi, Sebastian Stich, and Ananda Theertha Suresh. Scaffold: Stochastic controlled averaging for federated learning. In Interna- tional conference machine learning, pp. 5132 5143. PMLR, 2020.
Seongyoon Kim, Gihun Lee, Jaehoon Oh, and Se- Young Yun. Fedfn: Feature normalization for alleviating data heterogeneity problem in foderated learning. arXiv preprint arXiv:2311.1322 2023.
Tachyeon Kim, Jachoon Oh, NakYil Kim, Sangwook Cho, and Sc- So-Young Yun. Comparing kullback-lcibler divergence and mean squared error loss in knowledge distillation. arXiv preprint arXiv:2105.080 2021.
Alex Krizhevsky, Vinod Nair, and Geoffrey Hinton. Cifar-10 and cifar-100 100 datasets. toronto. talukriz/eifar himl, 6(1):1, 2009.
URl: https:/ https://wwww. cs.
15

---

Published in Transactions on Machine Learning Research (03/2025)
Gihun Lee, Minchan Jeong, Yongjin Shin, Sangmin Bac, and Se- Young Yun. Preservation of the global knowl- edge by not-true distillation in federated learning. Advances in Newral Information Processing Systems, 35:38461 14484461-3847,, 2022.
Gihun Lee, Minchan Jeong, Sangmook Kim, Jaehoon Oh, and Se- Se-Young Yun. Fedsol: Stabilized orthogonal learning with proximal restrictions in federated learning, 2024.
José Lezama, Qiang Qiu, Pablo Musé, and Guillermo Sapiro. Ole: Orthogonal low-rank embedding-a plug and play geometric loss for deep learning. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, pp. $109-8118, 2018.
Bo Li, Mikkel N Schmidt, Tommy S Alstrom, and Sebastian U Stich. On the effectiveness of partial variance reduction in federated learning with heterogeneous data. arXiv preprint arXiv:2212.00194 2022.
Daliang Li and Junpu Wang. Fedmd: Hetcrogenous federated learning via model distillation. arXiv preprint arXiv:1910.0384 033581, 2019.
Jingzhi Li, Zidong Guo, Hui Li, Seungju Han, Ji-won Back, Min Yang, Ran Yang, and Sungjoo Suh. Rethink- ing feature-based knowledge distillation for face recognition. In Proceedinggs of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 20156 75-60015, 2023a.
Qinbin Li, Bingsheng He, and Dawn Song. Model-contrastive federated learning. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp. 10713-10722, 2021a.
Quanquan Li, Shengying Jin, and Junjie Yan. Mimicking very efficient network for object detection. In Proceedings of the ieee conference on computer vision and pottern recognition, pp. 1556-6364, 2017.
Tian Li, Anit Kumar Sahu, Manzil Zaheer, Maziar Sanjabi, Ameet Talwalkar, and Virginia Smith. Federated optimization in heterogeneous networks. Proceedings of Machine learning and systems, 2:429-450, 2020.
Tian Li, Shengyuan Hu, Ahmad Beirami, and Virginia Smith. Ditto: Fair and robust federated learning through personalization. In International Conference on Machine Learning, pp. (357-6368. PMLR, 2021b.
Xin-Chun Li and De-Chuan Zhan. Fedrs: Federated learning with restricted softmax for label distribution non-iid data. In Proceedings of the 27th ACM SIGKDD Conference on Knowledge Discovery & Data Mining, pp. 995-1005, 2021.
Zexi Li, Xinyi Shang, Rui He, Tao Lin, and Chao Wu. No fear of classifier biases: Neural collapse inspired federated learning with synthetic and fixed classifier. In Proceedings of the IEEE/CVP International Conference Computer Vision, pp. 5319-5322 2023b.
Tao Lin, Lingjing Kong, Scbastian U Stich and Martin Jaggi. Ensemble distillation for robust model fusion in federated learning. Advances in Neural Information Processing Systems, 33:2351-2233 2263, 2020.
Mi Luo, Fei Chen, Dapeng Hu, Yifan Zhang, Jian Liang, and Jiashi Feng. No fear of heterogencity: Classifier calibration for federated learning with non-iid data. Advances in Neural Information Processing Systems, 14:5972-5984, 2021.
Michael McCloskey and Neal J Cohen. Catastrophic interference in connectionist networks: The sequential learning problem. In Psychology of learning and motivation, volume 24, pp. 109-165. Elsevier, 1989.
Brendan McMahan, Fider Moore, Daniel Ramage, Seth Hampson, and Blaise Aguera y Arcas. Conmmiication-fffciennt learning of deep networks from decentralized data. In Artificial intelligence and statistics, pp. 1273-1282 PMLR, 2017.
John Nguyen, Jianyu Wang, Kshitiz Malik, Maziar Sanjabi, and Michael Rabbat. Where to begin? on the impact of pre-training and initialization in federated learning. arXio preprint arXiv:2206. 153877776..1387, 2022.
Jachoon Oh, SangMook Kim, and Sc Young Yun. FodBABU: Toward enhancod representation for federated image classification. In International Conference on Learning Representations, 2022. URL https:// ppenreoviw...t//oooonnrtttttttto
16

---

Published in Transactions on Machine Learning Research (03/2025)
Biao Qian, Yang Wang, Hongzhi Yin, Richang Hong, and Meng Wang. Switchable online knowledge distil- lation. In European Conference on Compuler Vision, pp. 449-466. Springer, 2022.
Andrew M Saxe, James L McClelland, and Surya Ganguli. Exact solutions to the nonlinear dynamics of learning in deep linear neural networks. arviv preprint 1xXio:132.61122 2013.
Karen Simonyan and Andrew Zisserman. Very deep convolutional networks for large-scale image recognition.
arXiv preprint trXii:1409.. 1556, 2014.
Yue Tan, Guodong Long, Lu Liu, Tianyi Zhou, Qinghua Lu, Jing Jiang, and Chengqi Zhang. Fedproto: Federated prototype learning across heterogeneous clients. In Proceedings of the AAAI Conference on Arlificial Intelligence, pp. 8432 8440, 2022.
Hongyi Wang, Mikhail Yurochkin, Yuckai Sun, Dimitris Papailiopoulos, and Yasaman Khazaeni. Federated learning with matched averaging. In International Conferencee Learving Representations, 2020a. URL
Jianyu Wang, Qinghua Liu, Hao Liang, Gauri Joshi, and Vincent Poor. Tackling the objective inconsistency problem in heterogeneous federated optimization. Advances in neural information processing systems, 33: 7611 7623, 2020b.
Zikai Xiao, Zihan Chen, Liyinglan Liu, YANG FENG, Joey Tianyi Zhou, Jian Wu, Wanlu Liu, Howard Hao Yang, and Zuozhu Liu. Fedloge: Joint local and generic federated learning under long tailed data. In The Twelfth International Conference on Learning Representations, 2024. URL https://openreview.nr forum?id-v3j 5d0GQgHH.
Jian Xu, Xinyi Tong, and Shao-Lun Huang. Personalized federated learning with feature alignment and classifier collaboration. In The Eleventh International Conference on Learning Representations, 2023. URL
net/forummiddsxXrrBBkK
Greg Yang. Scaling limits of wide neural networks with weight sharing: Gaussian process behavior, gradient independence, and neural tangent kernel derivation, 2020. URL https://arxiee www.:oggg//////.ooooooooooos http/1992..44
https://araw..orr/ab/////// httgg/abs/192..4/6/// 04760.
Yibo Yang, Shixiang Chen, Xiangtai Li, Liang Xie, Zhouchen Lin, and Dacheng Tao. Inducing neural collapse in imbalanced learning: Do we really need a learnable classifier at the end of deep neural network? Advances in Neural Information Processing Systems, 15:37991-38002, 2022.
Yibo Yang, Haobo Yuan, Xiangtai Li, Zhouchen Lin, Philip Torr, and Dacheng Tao. Neural collapse inspired feature-classifeer alignment for few-shot class incremental learning. Xiv preprint atl::302203004, 2023.
Zhiqin Yang, Yonggang Zhang, Yu Zheng, Xinmei Tian, Hao Peng, Tongliang Liu, and Bo Han. Fedfed: Feature distillation against data heterogeneity in federated learning. Advances in Neural Information Processing Systems, 36, 2024.
Rui Ye, Yaxin Du, Zhenyang Ni, Yanfeng Wang, and Siheng Chen. Fake it till make it: Federated learning with consensus-oriented generation. In The Twelfth International Conference on Learning Representations, 2024. URL
Borui Zhao, Quan Cui, Renjie Song, Yiyu Qiu, and Jiajun Liang. Decoupled knowledge distillation. In Proceedings of the IEEE/CVF Conference on computer vision and pattern recognition, pp. 11953 -1962, 2022.
17

---

Published in Transactions on Machine Learning Research (03/2025)
Appendix
FedDr+: Stabilizing Dot-regression with Global Feature Distillation for Federated Learning
The notations and pseudo code of FedDr+ and FedDr+ FT are organized in Appendix A. In Appendix B, we provide a detailed explanation of the pulling and pushing gradients of the CE loss. In Appendix C, we provide theoretical analysis of dol-regression, focusing the feature vector gradient of the loss and its implications under the NTK framework, particularly for unobserved classes. The experimental setup is described in Appendix D, which includes code implementation, dataset descriptions, model specifications, optimizer settings, NIID partition, and the hyperparameter search process. Additional experimental results, including further analysis on the synergy effect and PFT,, as well as results on IID dataset performance, scalability experiments, and stochastic client data settings, are presented in Appendix
A Notations, Pseudo Code of FedDr+ and FedDr+
In this section, we first introduce the key notations used in our method and then present the pscudocode for FedDr+ and FedDr+ FT. The pseudocode provides a clear and concise implementation guide for both global federated learning (GFL) with FedDr+ and personalized federated learning (PFL) with FedDr+ FT.
A.1 Main Notations
To maintain clarity, Table 6 defines key indices, datasets, model parameters, and computations in alg and FedDr+ FT, forming the basis for our method and analysis.
Table 6: Notations used throughout the paper.
Indices
[C]
Index for class Index for FL round
r R
ie [N]
Index for client
Dataset
Drain
Training dataset for client Test dataset for client i Data on client sampled from distribution D (x: input data, y: class label)
D Dist
(x,9) E D Diminntes (x. ~ D
Oi
Dataset consists of observed classes in client
Dataset consists of unobserved classes in client
Parameters
Feature extractor weight parameters Classifier weight parameters (frozen during training) o-th row vector of All model parameters Aggregated global model parameters round r Trained model parameters client at round
V = -lu....,eell e Bxxd
Ve, e E [C]
θ (0.V)
O (0g, V)
θ (0.. V)
Model Forward p(=;0) R
Softmax probability of input c-th element of p(r; 0)
0)( E [C]
Loe( 0)
(x:
Cross-entropy loss of input
f(x;; 0) E Rd
Feature vector of input x
=(>; f(:; V RC
Logit vector of imput x
=a(1; 0, E [C]
o-th element of 2(2; 0)
18

---

Published in Transactions on Machine Learning Research (03/2025)
A.2 Pseudo Code of FedDr+ and FedDr+ FT
We now present the pseudocode for FedDr+ and FedDr+ FT, outlining their key operations for global and personalized federated learning. The algorithm consists of two main stages:
Algorithm FedDr+, FedDr+ FT
Input: Total rounds H, local epochs E, training dataset Dinin for client sampled client set N [N] at round learning rate x/*) at round
Initial Parameters: ETF Classifier V, Initial global model parameters 6 (0% V)
for == I,...,N N do Server broadcasts V clicut
/** STEP 1: Get GFL Model 0º of FedDr+ +*/ for L...., do
5 Server samples clients NOr) broadcasts 0 0.-1 for each elient C N parallel do for Local Steps 1......... do for Batches + 0 of VDr+ (ID (i 15 0 0 V)
Using [Equation (1)
9
Uplond 0f to server
Noopl Eie 0: GFL output: 6% (0
Server Aggregation: 0g
STEP 2: Get PFL Models (0R1 FedDr+FT +*/
for 1.....N .
Server broadeasts 0+++= < 0f to clicut for Local Steps e I.... E'
for Batches j I.....
15
0 0 0+1 V (ID (iiiii 0% V)
Using [Equation (1)
16 PFL outputs: (0 (0 VI
B Preliminaries: Pulling and Pushing Feature Gradients in CE
In this section, we first compute the classifier's gradient with respect to the features. Next, explain how the cross-entropy loss draws the pulling and pushing effects.
B.1 Feature Gradient of Lce
We begin by presenting two lemmas that support. Proposition 1 and clarify pulling and pushing feature gradients in the cross-entropy (CE) loss.
Lemma 1. For all e, d [C], (p; 0) 02, 0) -p((2;;) 0) (x; otherwise
[p 0) (1 p=(2; 0)) ic=d
(1)
Proof. Note that p(z; 0) = (x; 0))
E C. Then,
D (x;
(i) c' case:
exp(zc(1;0)) (EC, exp(2((2;)) 01) exp(2c 0))²
0p. 0)
exp(2o (x: 0))
O2. 0)
(z: 0)
D exp(2; (2; 0))
E exp(2, 0))
=pe(2; 0) pa(1;; 0) pe(x; 0) -=(
19

---

Published in Transactions on Machine Learning Research (03/2025)
(ii c d' case:
Ope (z; 0)
exp(zer (x;0))
exp(zo(1;; 0)) exp (x;0)
(z; 0)
(z; 0)
Lisl exp(2;
E exp(2 01)
=>p(<z; ()pe-
Lemma 2. V=(E:0) (;;0) p(x; 0) By, where By E R is the unit vector with its y-th element as 1.
Proof.
@CCE(1, 0)
0
1
Opy (2; 0)
log log (x; 0) = Py Py(x; 0) (1;
(zo
Ōzo(2;0 0)
(c(x;; 0) -1 if cc po(=; 0) 1 y)
Do (e 0)
else
The last equality holds by Lemma 1. Therefore, the desired result is satisfied.
Proposition 1. Given (I, the gradient of the Lce with respect to f(x;; 0) is given by:
v CoEl), v; = -(1 Py (2: 1) +
Σ
Pe(x;
(2)
cE(ONVE
Proof.
(1; (=)o) V /
V (1:00) CCE(E, u:
- v)
0Len(1, (1;0)=((;;00)
(zc(1; 0)
OLer(I, y/ (x;;))
OCe (1:;0)=
Ozy( (2; 0)
(2,(2;6)
<E(C)((V
θLCr y/: 0)
ÓCEE(F, y;
"y +
0zy (x; 0)
0z.
<E(C]M()
-1 (x; 1))
Σ
Pe(a;
Oyce-
<E(C)Mb)
Applying the chain rule for the second stop and invoking Lemma for the final equality confirms the result.
B.2 Physical Meaning of V ((mo) Lee( y: 8)
The gradient Lce:) yy 0) consists of two components:
Fruil =1 Py (z: 0))
F Push = =- Σ pe(2; 0)ve.
<E(O)(())
FPull moves the feature vector towards the classifier vector by of the true class, promoting alignment. contrast, F Frush moves it away from the classifier vectors for [C] \ {v), inducing misalignment.
20

---

Published in Transactions on Machine Learning Research (03/2025)
C Theoretical Perspective of Dot- Dot-Regression (DR)
In this section, we provide thooretical analysis of dot-regression (DR) loss in the context of feature-classifier alignment. We first derive the feature gradient of LDR and analyze its effect on feature updates. We then present an NTK-based perspective explaining why dot-regression struggles with unobserved classes in FL. Finally, we compare DR with cross-entropy (CE) loss to highlight its limitations and the necessity of feature distillation.
C.1 Feature Gradient of LDR
In this subsection, we derive the gradient of dot-regression loss with respect the feature vector on the observed classes.
Theorem C.1. Given (x, v), the gradient of the LDR with respect to s(x; 0y) is given by:
1 - cosa {
[(z; fy)
((x0,) CDn(,, 9; [f((:;
Vy cos a
(
f(z; 0,)
where cos
[1; 01)112 Vyy
Proof.
f(x;0) 0,)
1/, (1-1) 2)
}
((R:0)) v; = s(m00)
[I;;;; 0r)1l2
f(x; 0,)
f(z; 01)
V/, 1) V {(x00))) (1; 01)
(1; 05)112
0() 0,)
V/ -1)
1
f(x; (z; 0,)
V
(x; 05)
[x; 01)
(1; 0
1 cos [1((;;
f(x; 0,)
Vy,
[|[(x;; 0,)
C.1.1 Physical Meaning of V ((:00) (I,
According to Theorem C.1, the change in the feature vector Af(x; 0,) is given by:
Af(x; - COs
11(x;0,)1a (Vr (/-- (u-coos cos [x)
[f(=; (f)xx00)]1
where 7j is the learning rate and a is the angle between the feature vector [(x; 0y) and the larget vector Vy f(x:0,)
The term inside the parentheses, Vy, cos a [T<=0; represents component orthogonal to f(z;; 0y) that points towards Vy, This component adjusts f(x; 0,) to increase its cosine similarity with V,, while also expanding its norm.
The scaling factor 1-cosa determines the update magnitude. As training progressos, f(x; 0,) aligns more I
closely with Vy, reducing 1 - cos a and increasing [>; 02) Consequently, Af 0,) diminishes over time, reflecting convergence as the cosine similarity with V,, approaches its maximum.
Figure illustrates this process, showing how the orthogonal component drives both the rotation and scaling of [a; 0y) toward alignment with Vy,
21

---

Published in Transactions on Machine Learning Research (03/2025)
f
1-cos (V-- COS a
Af n
Ifll
II
f
f
Af
COS a
Iffl
COS a 1
a
Iifl
V
Figure 8: Feature gradient of CDR. The gradient update rotates fa; 0,) toward Vy while increasing its norm. As training progresses, the update magnitude decreases, leading convergence.
C.2 NTK Perspective: Why Dot- not-Regression in FL Fails on Unobserved Classes
In this subsection, we analyze why dot-regression struggles with unobserved classes under the Neural Tangent Kernel (NTK) regime (Jacot et al., 2018). In the NTK regime, the feature gradient of any input is a weighted sum of the feature gradients from training samples. Assuming the network width is sufficiently wide, these weights depend only on the pair of inputs, the initialization distribution, such as He initialization (Glorot & Bengio, 2010b), and the activation functions, The NTK regime holds when the setting where every layer in the neural network has infinite width, with parameters initialized i.dd. This section explains how NTK- based gradient updates fail to align feature vectors with unobserved class directions, which leads to poor gencralization FI.
C.2.1 Gradient Flow in the NTK Regime
We treat gradient descent continuous process. P is the number of trainable parameters in the feature extractor, and Op (p [P]) denote each parameter. We focus specific client, denoted by
During training, gradient descent updates the model parameters to minimize the loss function. As below, we can see the evolution of the function f(x; O(t)) can be analyzed using the kernel (t) mi), which evolves along the training process:
df(x;; O(t)) dt
(f:x(((t) T ddp
(Chain Rule)
D0p
dt
Of (x; 0()))
T
0f(;; 0(1)
(( ((,,0:0)
(Gradient Descent)
30p
Dc train (E.ū) DJ
00p
pml
Grain
O5 (()))
Ta) 0(1))
Σ
((2:0) ((E,ū:0)
Di frain
08p
09p
(2.0) irain
1=1
a @(())aa))
Σ
@ll (1)
Dirain
(
9nn practice, width effects deviations from the ideal NTK behavior.
22

---

Published in Transactions on Machine Learning Research (03/2025)
In the NTK regime with infinitely large widths, the matrix (t) a) converges to scalar multiple of the identity matrix, O (a, E)I. Furthermore, with the same condition, this scalar kernel remains constant throughout training (Jacot et al., 2018; Yang, 2020; Belfer al., 2021), though fmite-width effects may introduce small variations. In the NTK regime, the gradient descent dynamics are given by:
(f(2;0(t) dt
Σ
0 (z, 2) C(2, 0)
(in NTK Regime)
D Praiin
(2.0) train
Ex
Thus, O (ar,a) .) determines how each training sample Æ influences an arbitrary input æ, and NTK regime, this weight depends only on the initialization distribution.
In Federated Learning (FL), local models are independently updated on different clients before aggregation. Under the NTK regime, each client follows the gradient flow during local training. FL aggregation then combines feature representations Iearned from different data distributions, leading to shifts in the global feature representation. By aggregating updates from multiple clients, FL integrates feature information from clients that have observed missing classes, thereby improving feature alignment.
C.2.2 Limitations of Dot- not-Regression loss in FL under the NTK Regime
Dot-regression loss (Yang et al., 2022) speeds up the alignment of feature vectors fax 0) E IR (pre-classifier layer outpuls) with the true class direction uy by minimizing the cosine angle:
2 CDR(E, V) = (( (coss (f(x=)), - -1)²
This loss function is motivated by the decompasition of cross-entropy (CF) loss gradients into pulling and pushing components. Prior work suggests that removing the pushing effect in CE can improve conver- gence (Yang et al., 2022; Li & Zhan, 2021).
Let e be an unobserved class for specific client with the classifier vector ve. From Theorem C.1, follows that under the NTK regime, the gradient descent process on the client i is independent of ve for arbitrary input z.
To analyze this, we first express the feature gradient under the dot-regression loss LDn in the local learning stage. For simplicity, we omit the dependence on O(t) in the feature notation and write ff; 0(t), uy) as cos(f($), vy). The feature gradient given by:
dx( dt Diruial yEO* (E.D)E train
(x, 1 cos ())()()))))) [
f(i) ) (2y cos(((()))))
Σ Σ
(in NTK Regime)
[[
Since c & Oi, the feature gradient evaluated on training data does not depend on Ve- Given that feature gradients are weighted sum over training data in the NTK regime, this implies that the learned feature representation for arbitrary input remains unaffected by ve during local training.
Therefore, dot-regression cannot align features with unobserved classes in local training. To examine this effect more closely, consider two cases fi(x) and f2(x) with the same input & with label e, whose settings and initialization at time = are identical except for the classifier vector Ue of class c, fixed with w and ([|w) 1, Vy E O²: : w 1 1y).. In the NTK regime under the dot-regression loss, we have:
d L (1()),
Σ YEO! (f. irain
O cos( (f)(f (())) (())),o))
D
[
(in NTK Regime)
Since every term in the update equation identical for f1( and fa(x), except for (f(#). ve), which takes opposite values in each case, it follows that U1(x), (ff22)),,) v.) for all time & 2 0. This demonstrates
23

---

Published in Transactions on Machine Learning Research (03/2025)
that classifier initialization strongly determines alignment in the local learning stage. Consequently, the global aggregation stage is the only way to generalize to classes that haven't 't been observed yet. This slows down the overall accuracy of the FL server.
C.2.3 Cross-Entropy Loss and Classse- Alignment
In contrast, cross-entropy (CE) loss explicitly guides feature gradients toward Ver weighted by the softmax probability Pe and the NTK weight. This ensures that even when class is absent, local training still produces meaningful updates. After cach global aggregation, the refined Pe further strengthens alignment, allowing CE to maintain consistent feature-classifier aligmment across all classes.
This observation aligns with our empirical findings: without feature distillation, dot-regression struggles to generalize unobserved classes, whereas CE enables continuous feature updates, leading improved generalization.
24

---

Published in Transactions on Machine Learning Research (03/2025)
D Experimental Setup
This section details the code implementation, dataset descriptions, model specifications, optimizer settings, non-IID (NIID) partitioning, and hyperparameter search process used in our experiments.
D.1 Code Implementation
Our implementations are conducted using the PyTorch framework. Specifically, the experiments presented in Table and Table 4 are executed a single NVIDIA RTX 3090 GPU, based on the code structure from the following repository: ottps://gitiin out on single NVIDIA A5000 GPU, utilizing the code framework from http:://githhne cot/jooooooooonoooo oh/ FedBABU.
com/Lee- Gihun/FedNTD. The other parts of our study are carried
D.2 Datasets, Model, and Optimizer
To simulate realistic FL scenario, we conduct extensive studics three widely used datasets: CIFAR- 10 (Krizhevsky et al., 2009), CIFAR-100 (Krizhevsky et al., 2009) and ImageNet-100 Deng el al. (2009). For each dataset, appropriate models are employed: VGG11 (Simonyan & Zisserman, 2014) for CIFAR- 10, MobileNet (Howard et al., 2017) for CIFAR-100, and ResNet-18 (He et al., 2016) for ImageNet-1000 A momentum optimizer is utilized for all experiments. The data preprocessing pipeline for the training phase includes RandomResizedCron
RandomHorrzzottllFiip and Normalize transformations for all datasets. During testing, only the Normalize transformation applied for CIFAR-10 and CIFAR-100, while for ImageNet-100, Resize, CenterCrop, and Normalize are applied. Unless otherwise noted, the basic setting of our experiments follows the dataset statisties, FL scenario specifications, and optimizer hyperparameterss summarized in Table 7.
Table Summary of Dataset, Model, FL System, and Optimizer Specifications
Datasets
|
[Prest]
N
A
CIFAR-10
50000
10000
100
320
0.1
50
0.9
le-5
CIFAR-100 100
100
50000
10000
100
320
0.1
50
0.9
1e-5
ImageNet-110
100
130000
5000
100
320
0.1
50
0.9
1e-5
Note: In terms of dataset information, C represents the number of classes in the dataset, with |Dtrain| and Drest indicating the total numbers of training and test data used, respectively. For the federated learning (FL) system specifics, R indicates the total number of FL rounds, is the ratio of clients selected for each round, and E denotes the number of local epochs. Local model training utilizes a momentum optimizer where is the batch size, and m and À represent the momentum and weight decay parameters, respectively. The initial learning rale v is decayed by a factor 0.1 at the 160th and 240th communication rounds. The initial learning rate 7 and the number of local epochs E were determined via extensive grid search for each algorithm, with details outlined in Appendix D.4.
D.3 Non-IID Partition Strategies
To induce heterogencity in each client's training and test data (Dimain Diest), we distribute the entire class- balanced datasets, Darain and Drest, among 100 clients using both sharding and Latent. Dirichlet Allocation (LDA) partitioning strategies:
Sharding (MeMahan al., 2017; Oh et al., 2022): We organize the Durain and Dreat by label and divide [eraal and [Deatl samples of the 100xs
them into non-overlapping shards equal size. Each shard encompasses 100xs same class, where denotes the number of shards per client. This sharding technique used to create Dimin and D D which are then distributed to each client i, ensuring that each client has the same number of training and test samples. The data for each client. is disjoint. As result, each client has access to
25

---

Published in Transactions on Machine Learning Research (03/2025)
maximum of s different classes. Decreasing the number of shards per user increases the level of data heterogeneity among clients.
Latent Dirichlet Allocation (LDA) (Luo et al., 2021; Wang et al., 2020a): We utilize the LDA technique to create Dirain from Dreain- This involves sampling a probability vector P:: = (Pc.1 Pe.2, ... PP,,10) ~ Dir(a) and allocating a proportion Po,k of instances of class c E [C] to each client. k E [100]. Here, Dir(a) represents the Dirichlet distribution with the concentration parameter a. The parameter a controls the strength of data heterogeneity, with smaller values leading to stronger heterogencity among clients. For Diest we randomly sample from Diest to match the class frequency of Dirain and distribute it to cach client i.
D.4 Hyperparameter Search for and E
To optimize the initial learning rate (n) and the number of local epochs (E) for our algorithm, we conduct a grid search on the CIFAR-10, CIFAR-100, and ImageNet-100 datasets. The process and reasoning are outlined below.
D.4.1 Rationale for Varying Initial Learning Rate (1)
The algorithms used in our experiments differ in handling feature normalization within the loss function. Some algorithms apply feature normalization, while others do not. When features f(; 0) are normalized, the resulting gradient is scaled by T( This scaling effect necessitates a grid search across various learning rates to account for the differences in learning behavior.
D.4.2 Rationale for Varying Local Epochs E
In FL, choosing the appropriate number of local epochs is crucial. Too few epochs can lead to underfitting while too many can cause client drift. Therefore, finding the optimal number of local epochs essential by exploring a range of values.
D.4.3 Grid Search Process and Results
Considering the above reasons, perform grid search for n and E on CIFAR-10, CTFAR-100, and ImageNet- 100 datasets. The grid search for CIFAR-10 uses shard size of 2, while for CIFAR-100, shard size of 10 is used. Additionally, for ImageNet-100,, shard size of 20 is used. The detailed procedures for each dalaset are provided below. These optimal settings have also been confirmed to yield good performance in less heterogencous scttings.
CIFAR-10. We examine 7 values from (0.01, 0.05, 0.1, 0.15, 0.2, 0.25, 0.3, 0.35, 0.4, 0.45, 0.5, 0.55, 0.6). For E, we consider {1, 3, 5, 10, 15). A default initial learning rate of 0.01 is used unless specified otherwise. The optimal learning rates vary by algorithm, and the results are summarized in Table 8. Table 8 also includes the additional hyperparameters used for each algorithm. The notation for these additional hyperparameters follows the conventions used throughout this paper (Li et al., 2020; Lee et al., 2022; Jhunjhunwala et al., 2023; Li et al., 2023b; Lee al., 2024). The optimal number of local epochs found to be for every algorithm.
CIFAR-100. We examine values from {0.1, 0.5, 1.0, 1.5, 2.0, 2.5, 3.0, 3.5, 4.0, 4.5, 5.0, 5.5, 6.0, 6.5, 7.0). For E, we consider (1, 3, 5, 10). default initial learning rate of 0.1 is used unless specified otherwise. The optimal learning rates differ by algorithm, and the results are listed in Table 9. Table 9 also includes the additional hyperparameters used for each algorithm. The notation for these additional hyperparameters follows the conventions used throughout this paper (Li et al., 2020; Lee et al., 2022; Jhunjhunwala et al., 2023; Li et al., 2023b; Lee et al., 2024). The optimal number of local epochs is found to be 3 for every algorithm.
ImageNet-100. We examine 17 values from (0.01, 0.1, 1.0, 10.0), which are chosen to maintain a consistent logarithmic scale difference. A default initial learning rate of 0.1 is used unless specified otherwise. The
26

---

Published in Transactions on Machine Learning Research (03/2025)
optimal learning rates differ by algorithm, and the results are listed in Table 10. Table 10 also includes the additional hyperparameters used for each algorithm. The notation for these additional hyperparameters follows the conventions used throughout this paper (Li et al., 2020; Lee et al., 2022; Jhunjhunwala et al., 2023; et al., 2023b; Lee et al., 2024). The optimal number of local epochs is fixed at 5, following the sctting of (Lee et al., 2024).
Table 8: Hyperparameters for VGG11 training on CIFAR-10.
Feature mnstrrmalianl alyecithmm
normaliaed alguritlin.
Hyperperamnteo
BulAdg
RalBABU
PwiPPre
SCAFFOLD
MODN
BdNND
FwiRRsP
FwIGRLA
PwiRTF
SplaniFull
FedDr (Oun)
0.01
001
0.01
0ut
0.1
0.01
0.01
0.01
0.01
000
6:35
Additiionl
Noew
Neew
(1.05)
(8.3)=(1.3)
08, ()
Table Hyperparameters for MobileNet training on CIFAR-100.
Fosture
algoeitluue
Focture nocualiard algorithms
Reprepanamtte
Ballig
HaBBABU
FodPPrx
SCAFFULD
HadNTD
FoRGELA 0.1
Splacdtiad
(Uus)
0.1
0.1
0.1
0.1
0.1
0.1
0.1
0.1
0.5
6.5
5.0
Additiina
Nome
Norw
j=000001
Now
(1.05)
(9.3)-(133
-0.001
2.2.0
Nome
(8-))((11
Nov
3-009
Table 10: Hyperparameters for ResNet-18 training on ImageNet-100.
Feature un normalized algoctciams
Feazuze normallaed algorithms
Hyperporametos
Fodåve
FdBBABU
FedProx
SCAFFOLD
MOON
ReIND
FodRsP
PoaSOL
PodGELA
PocEFT
Spluzeeed
FedDr+ (Ouss)
7
0.1
1.0
Additional
Nome
None
--100001
Noe
(a. (1.08)))
(8.1)-(133
.000001
2.0.
Now
(#.+)=(111
Nove
g-0.
27

---

Published in Transactions on Machine Learning Research (03/2025)
E
Additional Experiment Results
This section provides additional experimental results, including analysis on the synergy effect, personalized FL (PFL), IID dataset performance, stochastic client data settings, and scalability experiments.
E.1
Synergy Effect Details
Table 11: Synergy of various FL algorithms and regularizers. Bascline indicates training FL models without a regularizer. FD denoles feature distillation, which is the regularizer we use in FedDr+.
Sharding (s 10)
LDA (a 0.1)
Algorithm
Baseline
+ Prox
+MOON
+KD
+NTD
+LD
-FD
Baseline
Prox
+MOON
+KD
+NTD
+LD
+FD
Avd
37.22
36.87
37.43
36.25
37.71
37.17
37.82
42.52
43.22
44.79
44.21
43.39
43.43
43.76
PedBABU
46.20
46.03
46.49
47.22
46.95
47.60
46.48
45.78
SphereFed
13.90
11.96
13.13
13.47
43.95
45.21
13.77
16.81
17.76
47.25
17.01
19.74
FedETF
32.42
31.87
31.30
32.76
32.65
32.25 25
32.77
16,27
45.71
45.98
46.67
46.16
45.91
46.47
FedGELA
29.17
28.69
28.80
29.11
28.84
29.36
30.33
27.11
29.03
28.09
28.45
29.62
29. 41
29.75
Dot- Dot- Regression
42.52
41.95
44.72
47.45
48.32
47.52
48.69
42.72
46.35
50.36
49.47
50.36
49.28
50.86
Table 12: Optimal 33 value selected through grid scarch to achieve the best synergy of various FL algorithms
and regularizers.
Sharding
10)
LDA (a 0.1)
Algorithu
Baseline
-Prox
+MOON
+KD
+NTD
+LD
+PD
Baseline
+Prux
MOON
+KD
+NTD
+LD
-FD
FodAvg
None
0.999
0.5
0.9999
0.9999
0.909
0.9
Nome 0.999
0.99
0.999
0.99
0.999
0.9999
FedBABB
None
0.9999
0.999
0.99
None 0.999
0.999
0.999
0.99
0.9999
Sphereleed
None
0.9999
0.9099
0.9999
0.9
0.99
None
0.9999
0.999
0.999
0.9999
0.939
0.99
FodETF
None
0.999
0.3
0.5
0.999
0.9
0.9
None
0.9909
0,9999
0.5
0.999
0.99
0.99
FedGELA
None
0.9999
0.7
0.5
0.7
0.5
0.7
None 0.99
0.9
0.5
0.5
0.5
0.3
Dot-Regressin
None
0.9999
0.9
0.5
0.5
0.5
0.9
None
0.9999
0.5
0.5
0.5
0.5
0.9
Table 13: Synergy of various FL algorithms and regularizers at / = 0.9.
Sharding (*
10)
LDA (a 0.1)
Algorithu
Baseline
+Prux
+MOON
+KD
+NTD
+LD
-FD
Baseline
+Prux
+MOON
+KD
+NTD
+LD
+FD
FedAvg
37.22
30.27
36.67
35.14
35.56
31.83
37.82
12.52
36.09
12.09
41.48
41.34
43.36
43.10
FedBABU
46.20
36.71
46.49
45.50
45.09
45.81
45.31
47.37
39.04
45.92
45.58
45.56
46.46
44.77
SphereFed
43.90
1.36
1.89
41.01
43.47
41.73
45.21
16.98
1.16
2.21
45.22
46.25
43.84
48.61
ReIETF
32.42
25.18
32.58
32.76
31.98
32.25
32.77
46.27
34,92
45.38
44.94
45.77
44.36
45.92
RedGELA
29.17
25.52
28.57
28.84
28.67
28.37
29.07
27.11
26.84
28.09
27.78
28.27
27.97
27.60
Dot- Regression
42.52
5.42
44.72
46.60
45.78
47.52
48.69
42.72
7.47
30.69
48.19
33.08
49.09
50.79
We evaluate the synergy ellect of various FL algorithms by maintaining their original training loss while incorporating specific regularizers, detailed in Equation 1 of the main text. To manage the differing loss scales between the baseline FL algorithms and the regularizers, range of values (0.1, 0.3. 0.5, 0.7, 0.9, 0.99, 0.999, 0.9999). The resulting performance and optimal / values are shown in Table 11 and Table 12. However, when we set B = 0.9 without addressing the issue of differing loss scales, the performance results, presented in Table 13, reveal that several synergies are significantly inferior due this oversight.
systematically tune the coefficient 3 across
28

---

Published in Transactions on Machine Learning Research (03/2025)
Table 14: PFL accuracy comparison with MobileNet on CIFAR-100. For PFL, we denote the entries in the form of X representing the mean and standard deviation of personalized accuracies across all clients derived from single seed.
Algorithm
8=10 10
s=20
s=100
a=0.05
a=0.1
a=0.33
Dot- Regresggesson
42.52
49.02
52.86
30.31
37 77.52=5.m
47.08=3.60
Dot-Regression FT (LDR)
(0.84
74..1=<(o.r))
56.84
72.02=
66.96
(s.o4
Not-Regression FT (Cr++)
+0.82+(e.12)
(3.73=(0.78))
56.69+
71.85
(6.59+
59.87 (3.00))
FedDr+ (ours)
48.69
51.00
53.23
39.63+t.12
45.83=a.18
48.04=3.44
FedDr+ FT (CDR) (ours)
84.23
73.732(170)
54 95)
78.65
74.86=(4:77)
62.47
FedDr+ FT (Cn++) (ours)
84.10=(s.48)
75.42=((.s0)
56.76+(4.91)
(8.55
(4.75+((4755)
62.16++a.77))
E.2 Personalized Federated Learning Results
We introduce FedDr+ FT and dot-regression FT, inspired by prior work (Oh et al., 2022; Dong et al., 2022; Li et al., 2023b; Kim et al., 2023). These methods enhance personalization by leveraging local data to fine-tune the GFL model. We investigate the impact of fine-tuning using LDD+ and LDR loss for each GFL model lo assess their effectiveness on personalized accuracy. Performance metrics without standard deviations indicate results on Diest obtained from the GFL model after the initial step the 2-step method. Our experiments involve heterogeneous settings with sharding and LDA non-IID environments, using MobileNet on CIFAR- 100 datasets. We set as 10, 20, and 100, and the LDA concentration parameter (a) as 0.05, 0.1, and 0.3. Table 14 provides detailed personalized accuracy results.
Our 2-step process involves first developing the GFL model either using dot-regression or FedDr+. In the second step, we fine-tune this model to create the PFT. model, again using LDR or LDT+- This results in four combinations: Dot-Regression FT (CDR), Dot- Not-Regression FT (CD++), FedDr+ FT (CDR), and FedDr+ FT (CDr+). When the GFL model is fixed, using LDR for fine-tuning consistently outperforms D across all settings, because dot-regression focuses on local alignment which advantages personalized fine-tuning. Conversely, when the finc-tuning method is fixed, employing LDr+ for the GFL model consistently outper- forms LDR across all settings. This aligus with previous research (Nguyen et al., 2022; Chen et al., 2023) suggesting that fine-tuning from a well-initialized model yields better PFL performance.
E.3 IID Data Performance
To address the question regarding the performance of FedDr+ or dot-regression loss in Foderated Learning (FL) settings with IID data, we conducted experiments on CIFAR-100 with 100 clients, distributing data IID and ensuring fair number of samples per client. We evaluated FedAvg, FedBABU, Dot-regression, and FedDr+ across 5 seeds, calculating the mean and standard deviation of the global model accuracy for each algorithm.
Table 15: Global model accuracy (%) in IID data settings.
Algorithm
Accuracy (mean = std)
FedAvg
47.19 1.06
FedBABU
45.18 + 0.61
Dot-regression
51.48 + 0.99
FedDr+
51.10 + 0.61
From the Table 15, evident that Dot-regrossion and FedDr+ achieve the highost performance, significantly outperforming both FedAvg and FedBABU, The performance of Dot-regression and FedDr+ is nearly identical under ITD settings.
29

---

Published in Transactions on Machine Learning Research (03/2025)
This similarity arises because, in the IID scenario, there are no unobserved classes across clients. As result, the feature distillation mechanism in FedDr+, which is specifically designed to mitigate lorgetting on unobserved classes, does not provide additional benefits. Instead, both Dot-regression and FedDr+ excel in improving local alignment across all classes, fully achieving the global modcl's objective of enhancing local alignment for all clients.
E.4
Performance in Stochastic Client Data Settings
While original experiments on CIFAR-100 (s=10) with 100 clients assumed static client dataset, conducted additional experiments where each client randomly removed one class from its dataset every 10 FL rounds. As expected, global model accuracy decreased for all methods, as shown in Table 16. However, FedDr+ consistently outperformed CE and Dot-regression, demonstrating its robustness in handling dynamic class distributions. The round-wisc global test accuracy trends for CE, Dot-regrossion, and FedDr+ in the stochastic selting are presented Figure 9c, further confirming FedDr+ stability and superior performance across training rounds.
Table 16: Global model accuracy (%) in static and stochastic client data settings.
Algorithm
Static Setting
Stochastic Setting
CE
46.20
43.59
Dot-regression
42.52
38. 13
FedDr+
48.69
44.96
Local Alignment
Local Alignment Erased Class
Global Accuracy
0.6
0.4
Observed Unobserved
CE
DR
Dr+
CE
DR
Dr+
0.3
40
Dr+
0.2
0.1
0-
0-
100
200
100
200
300
100
200
300
Round
Round
Round
(a) Local alignment
(b) Local alignment erased class
(c) Global test accuracy
Figure 9: Comparison of (a) feasure-classifier alignment the observed and unobserved classes test data, (b) ceasurer alignment on erased-class test data for 0 and (c) global test accuracy of 0ø on all classes. Models are trained using LCE, LDR, and LDTt-
To further investigate why FedDr+ maintains superior global accuracy in the stochastic setting, we analyzed the feature-classffier alignment for both observed unnood//vvdbd classes and crased classes.
Local alignment for Abserved//unobsevv classes (Fig 9a):
FedDr+ maintains superior feature-classifeer alignment for both observed and unobserved classes compared to Dot-regression, consistently outperforming it across all rounds.
During the final convergence phase, FedDr+ surpasses even CE in unobserved class aligmment, confirming its effectiveness in preserving global knowledge.
Local alignment for erased class (Fig 9b):
Even for erased class (those removed during training), FedDr+ retains stronger feature-classifer
alignment than Dot-regression.
During the final convergence phase, FedDr+ also surpasses CE in erased class alignment, further demonstrating its ability to mitigate forgetting of removed class knowledge.
30

---

Published in Transactions on Machine Learning Research (03/2025)
These results suggest that the feature distillation mechanism in FedDr+ effectively enhances global knowledge preservation while also enabling effective learning of observed classes, even when class distributions change dynamically.
E.5 Scaling to Larger Numbers of Clients and Training Rounds
We conducted experiments CIFAR-100 (s=10) with 1,000 communication rounds, increasing the number of clients to 100, 200, 500, and 1,000. All algorithms used previously grid-searched optimal hyperpa- rameters, and results are averaged over three independent seeds. All algorithms used previously grid-searched optimal hyperparameters, and results are averaged over three independent seeds.
Table 17: Global model accuracy (%) for different numbers of clients with 1,000 communication rounds.
Algorithm
N=100
N-200
N-500
N=1,000
AvdAvg
50.50 0.57
42.51 + 1.47
33.02 + 0.74
26.63 1.31
FedBABU
58.19 + 1.07
48.75 + 1.99
37.40
+ 0.41
25.10 + 1.08
FedDr+
64.21 + 1.24
59.78 + 0.71
43.27 + 0.31
28.99 + 0.98
Table 17 confirms that FedDr+ consistently outperforms FedAvg and FedBABU across all settings, demon- strating robust scalability in large-scale FT..
31


Reviews:
[{"review_id": "6re04St8Ru", "reviewer": "Reviewer_f4TB", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths\n\n* The paper is well-written and easy to follow.\n* Improving performance in both global and personalized FL with heterogeneous data is an important open problem, and the authors provide a sound method.\n* Each component of FedDr+ is introduced with clear motivation and supported by sufficient experimental evidence.\n* The evaluations are comprehensive and well-executed.\n\nWeaknesses and Questions\n\n* Beyond the intuitive explanation, is there any theoretical justification for why CE performs better than DR on unobserved data?\n* The number of clients and training rounds is small. What happens if training rounds exceed 1k and the number of clients exceeds 1k?\n* Does the observed/unobserved data class change in each training round?\n* What happens if some data classes appear only at the beginning and then never reappear? Does FedDr+ still retain those features?", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The paper introduces FedDr+, an approach designed to improve performance in both global and personalized federated learning (FL) scenarios with heterogeneous client data. The method focuses on enhancing local alignment while preserving the representation of unseen class samples. The authors adopt a dot-regression loss to strengthen local alignment, while they also identify a drawback of this approach: a significant loss of information related to unseen classes. To maintain global knowledge, FedDr+ incorporates a feature distillation mechanism during local model training. This ensures the model does not overly prioritize local alignment at the expense of broader generalization. The paper includes several experiments to validate the effectiveness of FedDr+.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "Please see above", "broader_impact_concerns": "NA", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "Kd827DrpaA", "reviewer": "Reviewer_TUZs", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n1. The paper investigates the effects of applying dot-regression loss for FL\n2. The paper proposes FedDr+ algorithm that combines two of the existing techniques to deal with data heterogeneity in FL setups.\n3. The paper presents exhaustive experimental results to demonstrate the improvements attained by FedDr+\n4. The paper also presents results on personalized FL.\n5. Ablation study on the newly introduced hyper-parameter i.e., $\\beta$ is provided.\n6. The paper is well written and easy to follow. The paper address an important concern in Federated Learning i.e., data heterogeneity.\n\nWeaknesses/questions:\n1. Contribution 1 indicates that \"We find that dot-regression loss is not easily compatible with FL\". However, column 2 and 9 of Table 2 shows that dot regression performs better than FedAvg which uses cross-entropy loss. Can you please explain why dot-regression loss is not easily compatible with FL? \n2. Will FedDr+ or dot-regression loss outperform  FedAvg in the FL settings with IID data? \n3. ETF initialization forces a prior distribution on the class similarities/dissimilarities i.e., all the classes are equally distinct from each other. This prior distribution can be very distinct from the original distribution and can result in poor model performance. How does the algorithm deal with this issue?", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The paper proposed FedDr+ which combines the dot-regression loss [1] with Feature Distillation [2] for federated learning over non-iid data. The main idea of the algorithm is as follows:\n1. Initialize the classifier to satisfy the simplex Equiangular Tight Frame (ETF) condition and freeze it.\n2. Train the remaining network using FedAvg algorithm where the local loss on each client is given as $\\beta L_{DR} + (1-\\beta) L_{FD}$.\n \n    a. $L_{DR}$ is the dot regression loss which tries to minimize the cosine difference between the feature extractor output $f(x,\\theta)$ and the classifier's weight vector corresponding to the ground truth class $v_y$. Here, $\\theta$ is the trainable parameters set and $V$ is frozen i.e., only the feature extractor is trained during back-propagation.\n    \n    b. $L_{FD}$ is the feature distillation loss which reduces the $L_2$ distance between the  global model’s feature vectors $f(x, \\theta^g_{r-1})$ and local model's feature vector  $f(x, \\theta^i_{r})$ at each client $i$. This regularizes the local feature vectors to be closer to the received global feature vectors at each round. \n\nThe experimental results show that combining dot regression loss with feature distillation loss results in performance improvements for FL with non-IID data in both global and personalized settings. \n\nReferences:\n1. Yibo Yang, Shixiang Chen, Xiangtai Li, Liang Xie, Zhouchen Lin, and Dacheng Tao. Inducing neural collapse in imbalanced learning: Do we really need a learnable classifier at the end of deep neural network? Advances in Neural Information Processing Systems, 35:37991–38002, 2022\n2. Byeongho Heo, Jeesoo Kim, Sangdoo Yun, Hyojin Park, Nojun Kwak, and Jin Young Choi. A comprehensive overhaul of feature distillation. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pp. 1921–1930, 2019.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "1. Including the results on IID data setting can provide valuable insights into the proposed FedDr+\n2. Please add a limitation section addressing the trade-offs of the proposed algorithm.\n3. Conducting experiments on multiple seeds and reporting the mean and standard deviation of the accuracy is important as FL with non-IID data usually has high variance in the performance.\n\nMinor changes:\n1. Add a pseudo code for FedDr+ in the appendix to provide a concise implementation guide for the readers.", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "QVO4OvDT6g", "reviewer": "Reviewer_BYv6", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n1. The article provides interesting (and new for the FL system) loss functions that enhance feature alignment in FL (for the local models) while further ensuring global knowledge is not lost.\n2. The two major proposed loss functions are simple to add to the existing objectives ensuring practical use cases.\n3. Furthermore, both proposed objectives are intuitive to understand in the context of the current problem.\n4. Experimental results verify the utility of the proposed solutions.\n\nWeaknesses:\n1. The paper is hard to read. Although there is sufficient focus on FL systems and related nomenclature, some key topics such as local alignment are not properly elaborated upon. Ideally, adding a summary about local alignment and global knowledge preservation is necessary to improve the paper’s readability. \n2. Furthermore, it is unclear from the article whether the final aim is to solve local and global problems simultaneously or via two different models. For example, if we learn two models from the same FL setup, why are such two models not superior to a single solution? Intuition about the ‘why’ behind the proposed primary problem would be helpful to further appreciate the paper’s contributions. \n3. Can the authors elaborate upon the exact changes required for demonstrating the convergence of the newly proposed federated algorithm? I believe the primary updates are to the loss function, but are there certain important changes to how the algorithm converges in comparison to FedAvg (or similar algorithms)?\n4. One of the use cases of FL is to limit the data sent to the server and provide some form of data-minimization privacy. Given the form of the new loss function incorporated in the article, can the authors elaborate on the privacy impact of the updated method? Furthermore, it is the new method trivially extensible to differential privacy or other aggregation methods such as secure party computation?", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The paper proposes a mechanism for solving the two important problems of locally optimized FL (personalized FL) and maintaining global knowledge (global FL). By leveraging the dot-regression loss function alongside other enhancements, the article generates a novel FL-compatible loss function geared towards the objective above. The dot-regression loss ensures that the feature vector per input and the corresponding classifier parameters vector have a similar alignment (in terms of their angle measured by cosine). Such a method prompts the development of classifiers in tune with the local client’s features. Furthermore, to ensure that the local models maintain higher predictivity for unobserved classes (not available in client training data) the global model features are distilled to the client features by means of a mean squared error loss function between the global and local model parameters. Overall, combining these two loss functions in a weighted fashion provides the base loss function for both the global and local FL systems. Finally, experiments on CIFAR-10/100 and ImageNet-100 demonstrate the empirical efficacy of both approaches FedDr+ and FedDr+FT for global and local FL respectively.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "The changes requested are highlighted in the weaknesses section above. Improving the overall writing, and identifying the convergence analysis for the newly proposed FL method alongside quantifying the impact of FL on privacy are the main suggested changes.", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: The authors experimentally validate their federated learning algorithm. In particular the reviewers are all convinced, especially after response period, that there are good performance improvements in the case of heterogeneous data setting of federated learning. The authors have shown that the improvement persists both in the global sense, and also in a personalized (locally measured) sense. While the theoretical justification is lacking, the method seems to be promising to federated learning practitioners.
Decision: Accept (poster)