#  FedDr+: Stabilizing Dot-regression with Global Feature Distillation for
Federated Learning

Seongyoon Kim
Dept. ISysE, KAIST
<EMAIL_ADDRESS>
Minchan Jeong
Graduate School of AI, KAIST
<EMAIL_ADDRESS>
Sungnyun Kim
Graduate School of AI, KAIST
<EMAIL_ADDRESS>
Sungwoo Cho
Graduate School of AI, KAIST
<EMAIL_ADDRESS>
Sumyeong Ahn
CSE, Michigan State University
<EMAIL_ADDRESS>
Se-Young Yun11footnotemark: 1
Graduate School of AI, KAIST
<EMAIL_ADDRESS>
corresponding authors

###### 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 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 within each client. Thus, our objectives
are twofold: (1) enhancing local alignment while (2) preserving the
representation of unseen class samples. This approach aims to effectively
integrate knowledge from individual clients, thereby improving performance for
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.
Consequently, we provide empirical evidence demonstrating that our algorithm
surpasses existing methods that use a frozen classifier to boost alignment
across the diverse distribution.

##  1 Introduction

Federated Learning (FL) [1, 2, 3] is a privacy-aware distributed learning
strategy that employs data from multiple clients while ensuring their data
privacy. A foundational method in FL, known as FedAvg [1], involves four
iterative phases: (1) distributing a global model to clients, (2) training
local models using each client’s private dataset, (3) transmitting the locally
trained models back to the server, and (4) aggregating these models. This
method effectively protects privacy without requiring the transmission of raw
data to the server. However, a significant challenge in FL is data
heterogeneity, called _non-iidness_ , which refers to the different underlying
data distribution across clients. Such variance can cause _client drift_
during training, obstructing the convergence of the aggregated model and
significantly reducing its effectiveness.

To address client drift in non-iid scenarios, recent works [2, 4, 5, 6] have
identified that the last classifier layer in neural networks is particularly
vulnerable to this issue. Hence, they suggest strategies that freeze the
classifier while updating only the feature extractor. These approaches aim to
enhance the _alignment_ between the frozen classifier and the output from the
feature extractor. For instance, FedBABU [2] employs various classifier
initialization techniques, keeping it fixed during the training of the feature
extractor. The methods proposed in [4, 5, 6, 7, 8] utilize more robust
initialization, the Equiangular Tight Frame (ETF) classifier [9], to replace
traditional random initialization and improve the local alignment strategy.

Figure 1: Overview of the proposed method, FedDr+ trained with
$\mathcal{L}_{\text{Dr+}}$. To enhance the local alignment, we employ dot-
regression loss $\mathcal{L}_{\text{DR}}$, which discards the pushing term of
cross-entropy loss, and propose a feature distillation
$\mathcal{L}_{\text{FD}}$ to preserve the knowledge imbued in the global
model. We describe $\mathcal{L}_{\text{DR}}$ in Section 2, and feature
distillation in Section 3 in detail.

A frozen classifier is also extensively explored in other research areas, such
as class imbalance [10] and class incremental learning [11], with a consistent
objective similar to aforementioned FL studies—enhancing alignment. Recently,
these fields have advanced by introducing and utilizing a novel type of loss,
called dot-regression loss $\mathcal{L}_{\text{DR}}$, which aims to achieve
alignment rapidly. In summary, $\mathcal{L}_{\text{DR}}$ originates from the
decomposition analysis of cross-entropy (CE) loss, which includes _pulling_
and _pushing_ components. As suggested in [10], the _pulling_ component is a
force that attracts features to the target class, whereas the _pushing_
component is a force that drives features away from other non-target classes.
$\mathcal{L}_{\text{DR}}$ discards the _pushing_ component, as it slows down
convergence to the desired alignment (refer to Figure 1).

Following the advancement of leveraging the frozen classifier with dot-
regression loss, we investigate the application of this loss to FL. However,
our findings indicate that dot-regression loss does not necessarily lead to
sufficient performance improvement of the aggregated server-side model,
although it enhances _local alignment_ as intended. We observe that this
drawback stems from the handling of unseen class samples. Specifically, while
alignment improves for the classes in the local training dataset, it
significantly deteriorates for unseen classes. This observation highlights the
need to preserve the representation of unobserved classes during local
training. To address this issue, we propose a training mechanism, termed
FedDr+, that employs dot-regression loss alongside feature distillation that
reduces the distance between feature vectors of local and global models.

Contributions. Our main contributions are summarized as follows:

  * •

We find that dot-regression loss is not easily compatible with FL, although it
can enhance the alignment of seen classes. The drawback comes from a
significant loss of information on unseen classes, which is vital in the
global model perspective. Therefore, we aim to preserve information of unseen
classes within the FL system.

  * •

To preserve global knowledge, including unseen class information while
maintaining the advantages of $\mathcal{L}_{\text{DR}}$, we propose FedDr+,
which utilizes a feature distillation when training local models. This
regularizer prevents the model from focusing solely on the local alignment.

  * •

We verify that the proposed method surpasses the conventional algorithms in
both global and personalized FL under various datasets and non-iid settings.

##  2 Preliminaries

###  2.1 Basic Setup of Conventional FedAvg Pipeline

Basic FL setup. Let $[N]=\\{1,\ldots,N\\}$ denote the indices of clients, each
with a unique training dataset
$D_{\text{train}}^{i}=\\{(x_{m},y_{m})\\}_{m=1}^{|D_{\text{train}}^{i}|}$,
where $(x_{m},y_{m})\sim\mathcal{D}^{i}$ for the $i^{\text{th}}$ client,
$x_{m}$ is the input data, and $y_{m}\in[C]$ is the corresponding label among
$C$ classes. Importantly, FL studies predominantly address the scenario where
the data distributions are heterogeneous, _i.e.,_ $\mathcal{D}^{i}$ varies
across clients. Knowledge distributed among clients is collected over $R$
communication rounds. The general objective of FL is to train a model fit to
the aggregated knowledge, $\bigcup_{i\in[N]}\mathcal{D}^{i}$. This objective
can be seen as solving the optimization problem:

$\min_{\bm{\Theta}=(\bm{\theta},\bm{V})}\sum_{i\in[N]}\frac{|D_{\text{train}}^{i}|}{\sum_{j\in[N]}|D_{\text{train}}^{j}|}\,\scalebox{1.45}{$\displaystyle\mathop{\mathbb{E}}$}_{(x,y)\sim\mathcal{D}^{i}}\Big{[}\mathcal{L}(x,y;\bm{\theta},\bm{V})\Big{]}\,,$
(1)

where $\mathcal{L}$ is the instance-wise loss function, $\bm{\theta}$ is the
weight parameter for the feature extractor, and
$\bm{V}=[v_{1},\ldots,v_{C}]\in\mathbb{R}^{d\times C}$ is the classifier
weight matrix. We use the notation $\bm{\Theta}$ to denote the entire set of
model parameters.

At the beginning of each round $r\in[R]$, the server has access to only a
subset of clients $\mathcal{S}_{r}\subset[N]$ participating in the
$r^{\text{th}}$ round. At each round $r$, the server transmits the global
model parameters $\bm{\Theta}_{r-1}^{g}$ to the participating clients. Each
client then updates the parameters with their private data
$D_{\text{train}}^{i}$ and uploads $\bm{\Theta}_{r}^{i}$ to the global server.
By incorporating the locally trained weights, the server then updates the
global model parameters to $\bm{\Theta}_{r}^{g}$.

FedAvg pipeline. Our study follows the conventional FedAvg [1] framework to
address the FL problem. FedAvg updates the global model parameters from
locally trained parameters by aggregating these local models into
$\bm{\Theta}_{r}^{g}=\sum_{i\in S_{r}}w_{r}^{i}\bm{\Theta}_{r}^{i}$, where
$w_{r}^{i}=|D_{\text{train}}^{i}|\,/\,{\sum_{j\in
S_{r}}|D_{\text{train}}^{j}|}$ is the importance weight of the $i^{\text{th}}$
client.

###  2.2 Dot-Regression Loss for Faster Feature Alignment

Dot-regression loss $\mathcal{L}_{\text{DR}}$. This loss [10] facilitates a
faster alignment of feature vectors (penultimate layer outputs)
$f(x;\bm{\theta})\in\mathbb{R}^{d}$ to the true class direction of $v_{y}$,
reducing the cosine angle as follows:

$\mathcal{L}_{\text{DR}}(x,y;\bm{\theta},\bm{V})=\frac{1}{2}\Big{(}\cos\big{(}f(x;\bm{\theta}),v_{y}\big{)}-1\Big{)}^{2}$
(2)

where $\cos(\mathrm{vec}_{1},\mathrm{vec}_{2})$ denotes the cosine of the
angle between two vectors $\angle(\mathrm{vec}_{1},\mathrm{vec}_{2})$.

The main motivation is that the gradient of the cross-entropy (CE) loss for
the feature vector can be decomposed into a _pulling_ and _pushing_ gradient,
and recent work indicates that we can achieve better convergence by removing
the pushing effect [10, 12]. The _pulling_ gradient aligns $f(x;\bm{\theta})$
with $v_{y}$, while the _pushing_ gradient ensures $f(x;\bm{\theta})$ does not
align with $v_{c}$ for all $c\neq y$ (Appendix B details the exact form of
pulling and pushing gradients). Since $\mathcal{L}_{\text{DR}}$ directly
attracts features to the true-class classifier, it drops the _pushing_
gradient, thereby increasing the convergence speed for maximizing
$\cos(f(x;\bm{\theta}),v_{y})$.

Frozen ETF classifier. Since $\mathcal{L}_{\text{DR}}$ focuses on aligning
feature vectors with the true-class classifier, the classifier is not required
to be trained. Instead, we construct the classifier to satisfy the simplex
Equiangular Tight Frame (ETF) condition, a constructive way to achieve maximum
angular separation between class vectors [10, 11]. Concretely, we initialize
the classifier weight $\bm{V}$ as follows and freeze it throughout training:

$\bm{V}\longleftarrow\sqrt{\frac{C}{C-1}}\bm{U}\left(\bm{I}_{C}-\frac{1}{C}\bm{1}_{C}\bm{1}_{C}^{\top}\right),$
(3)

where $\bm{U}\in\mathbb{R}^{d\times C}$ is a randomly initialized orthogonal
matrix. Note that each $v_{i}$ in the classifier weight $\bm{V}$ satisfies
$\cos(v_{i},v_{j})=-\frac{1}{C-1}$ for all $i\neq j\in[C]$111This relation for
cosines holds if the $v_{i}$’s are symmetrically distributed such that
$\bar{v}=\frac{1}{C}\sum_{i\in[C]}v_{i}=0$, and $\cos(v_{i},v_{j})$ are all
the same for $i\neq j$ ..

##  3 When Dot-Regression Loss Meets Federated Learning

Given our focus on applying $\mathcal{L}_{\text{DR}}$ to FL, we first examine
its impact on FL models compared to the CE loss $\mathcal{L}_{\text{CE}}$. In
summary, we find that while $\mathcal{L}_{\text{DR}}$ improves alignment and
performance on observed class labels, it faces challenge with unobserved
classes222While we use the term “unobserved” in this context, it also applies
to “rarely” existing classes., which are essential for the generalization
objective. To address this issue, we propose FedDr+, which integrates
$\mathcal{L}_{\text{DR}}$ 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. In this section, we conduct experiments on
CIFAR-100 [13] with a shard non-iid setting ($s$=10), where each client
contains at most 10 classes. We additionally employ LDA setting ($\alpha$=0.1)
in Section 3.4. Refer to Section 4 for more details on 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
$160^{\text{th}}$ and $240^{\text{th}}$ rounds.

(a) Local alignment

(b) Local accuracy

Figure 2: Comparison of (a) feature-classifier alignment and (b) accuracy on
the observed and unobserved classes test data for $\bm{\theta}_{r}^{i}$
trained with $\mathcal{L}_{\text{CE}}$ and $\mathcal{L}_{\text{DR}}$.

(c) Alignment gap

(d) Accuracy gap

Figure 3: Comparison of (a) feature-classifier alignment gap and (b) accuracy
gap on the observed and unobserved classes test data for $\bm{\theta}_{r}^{i}$
trained with $\mathcal{L}_{\text{CE}}$ and $\mathcal{L}_{\text{DR}}$.

###  3.1 Impact of Dot-Regression Loss on Local and Global Models

We investigate the performance of local models on average when trained with
$\mathcal{L}_{\text{DR}}$ compared to $\mathcal{L}_{\text{CE}}$. In Figure
3–3, we calculate the statistics on two datasets: the observed class set
$\mathcal{O}^{i}$, which includes classes present in each client’s training
data $D_{\text{train}}^{i}$, and the unobserved class set $\mathcal{U}^{i}$,
consisting of classes unseen during training. This partition highlights the
challenges associated with generalizing to unseen classes in FL.

First, we evaluate the feature-classifier alignment
$\cos(f(x;\bm{\theta}_{r}^{i}),v_{y})$ and accuracy of each local model on the
test data (Figure 3). We then observe the amount of change from the given
global model to each local model in every communication round (Figure 3). For
instance, the alignment gap is denoted by
$\cos(f(x;\bm{\theta}_{r}^{i}),v_{y})-\cos(f(x;\bm{\theta}_{r-1}^{g}),v_{y})$.

Performance analysis of local models. Figure 3 shows that
$\mathcal{L}_{\text{DR}}$, by focusing its pulling effects exclusively on
observed classes within a client’s dataset, effectively enhances alignment and
accuracy for these classes. However, this specificity leads to poor
generalization on unobserved classes, resulting in significantly weaker
performance than models trained with $\mathcal{L}_{\text{CE}}$. Figure 3
displays the different impacts on $\mathcal{O}^{i}$ and $\mathcal{U}^{i}$
during updates from the global model to local models.
$\mathcal{L}_{\text{DR}}$ significantly boosts alignment and accuracy for
$\mathcal{O}^{i}$ but causes significant reductions for $\mathcal{U}^{i}$.

Global model accuracy result. We confirm that $\mathcal{L}_{\text{DR}}$ shows
superior accuracy for $\mathcal{O}^{i}$ compared to $\mathcal{L}_{\text{CE}}$
but is less effective at generalizing to $\mathcal{U}^{i}$. In the shard
setting ($s=10$)—where each client has access to at most 10 out of 100
classes—this shortcoming significantly reduces the global model’s overall
accuracy ($\mathcal{L}_{\text{DR}}$: 42.52$\%$ vs. $\mathcal{L}_{\text{CE}}$:
46.38$\%$). Thus, it is crucial to develop methods that retain the strengths
of $\mathcal{L}_{\text{DR}}$, _i.e.,_ alignment of observed classes, while
improving generalization for unobserved classes, highlighting the need for
more adaptive loss functions in FL.

###  3.2 FedDr+: Dot-Regression and Feature Distillation for Federated
Learning

We propose FedDr+ to mitigate forgetting unobserved classes while retaining
the strengths of dot-regression loss in aligning features of observed classes.
Using $\mathcal{L}_{\text{DR}}$ with the frozen classifier $\bm{V}$, FedDr+
includes a regularizer that fully distills the global model’s feature vectors
$f(x;\bm{\theta}^{g})\in\mathbb{R}^{d}$ to the client features
$f(x;\bm{\theta})$, to enhance generalization across all classes. The proposed
loss function $\mathcal{L}_{\text{Dr+}}$, shown in Eq. (missing) 4, combines
$\mathcal{L}_{\text{DR}}$ with a regularizer
$\mathcal{L}_{\text{FD}}(x;\bm{\theta},\bm{\theta}^{g})=\frac{1}{d}\|f(x;\bm{\theta})-f(x;\bm{\theta}^{g})\|_{2}^{2}$.
Unless specified, we use a scaling parameter $\beta=0.9$ throughout the paper.

$\mathcal{L}_{\text{Dr+}}(x,y;\bm{\theta},\bm{\theta}^{g},\bm{V})=\beta\cdot\mathcal{L}_{\text{DR}}(x,y;\bm{\theta},\bm{V})+(1-\beta)\cdot\mathcal{L}_{\text{FD}}(x;\bm{\theta},\bm{\theta}^{g})$
(4)

Why feature distillation? To address data heterogeneity in FL, various
distillation methods have been explored, including model parameters [2, 14,
15, 16], logit-related measurement [16, 17, 18, 19, 20, 21, 22], and co-
distillation [23, 24]. In contrast, we utilize the feature distillation [25]
technique because the feature directly concerns alignment. On the other hand,
logits lose information from features when projected onto a frozen ETF
classifier [25, 26, 27, 28]. By distilling features, we leverage the global,
differentiated knowledge for each data input $x$. This approach aims to
minimize blind drift towards observed classes, and hence, we expect it to
enhance overall generalization.

(a) Feature distance

(b) Feature angle distance

(c) Feature norm difference

Figure 4: We present (a) feature distance, (b) feature angle distance, (c) and
feature norm difference from $\bm{\theta}_{r-1}^{g}$ to $\bm{\theta}_{r}^{i}$
for observed and unobserved classes by training with $\mathcal{L}_{\text{DR}}$
and $\mathcal{L}_{\text{Dr+}}$.

###  3.3 Effect of Feature Distillation

Our findings from Section 3.1 indicate that $\mathcal{L}_{\text{DR}}$ is
unsuitable for the heterogeneous FL environment. This is primarily because
there is a notable gap in how features align with the fixed classifier between
$\mathcal{O}^{i}$ and $\mathcal{U}^{i}$. To assess the effect of feature
distillation ($\mathcal{L}_{\text{FD}}$), which imposes a constraint on the
feature distance $\|f(x;\bm{\theta}_{r}^{i})-f(x;\bm{\theta}_{r-1}^{g})\|_{2}$
for $x\in\mathcal{O}^{i}$, we measure this distance for both $\mathcal{O}^{i}$
and $\mathcal{U}^{i}$ from the models trained with $\mathcal{L}_{\text{DR}}$
and $\mathcal{L}_{\text{Dr+}}$. We additionally analyze the angle distance,
$\angle(f(x;\bm{\theta}_{r}^{i}),f(x;\bm{\theta}_{r-1}^{g}))$, and feature
norm difference,
$\|f(x;\bm{\theta}_{r}^{i})\|_{2}-\|f(x;\bm{\theta}_{r-1}^{g})\|_{2}$, as
these factors influence the feature distance. These values are averaged over
the selected client set $\mathcal{S}_{r}$.

Feature distillation stabilizes the feature dynamics. By adding
$\mathcal{L}_{\text{FD}}$, as revealed in 4(a), the local model trained with
$\mathcal{L}_{\text{Dr+}}$ shows a reduction in feature distance for observed
classes, compared to the model trained with $\mathcal{L}_{\text{DR}}$. This
reduction happens even for unobserved classes. As demonstrated in 4(b) and
4(c), reduction of feature distance originates from reducing the feature angle
distance and feature norm difference for both class sets. In both local models
trained with $\mathcal{L}_{\text{DR}}$ and $\mathcal{L}_{\text{Dr+}}$, there
is a trend where the angle is significantly larger for $\mathcal{U}^{i}$ than
for $\mathcal{O}^{i}$ (4(b)), while the norm difference is smaller for
$\mathcal{U}^{i}$ than for $\mathcal{O}^{i}$ (4(c)). This large angle distance
of $\mathcal{U}^{i}$ leads to the degradation of the feature-classifier
alignment. By minimizing the angle distance via feature distillation, the
global model’s accuracy improved substantially, rising from $42.52\%$ with
$\mathcal{L}_{\text{DR}}$ to $48.69\%$with $\mathcal{L}_{\text{Dr+}}$.

Stabilized features enhance alignment and accuracy.

We confirm that feature distillation term $\mathcal{L}_{\text{FD}}$ stabilizes
feature dynamics for both $\mathcal{O}^{i}$ and $\mathcal{U}^{i}$, enhancing
the global model’s capabilities. While the feature difference is stabilized
via $\mathcal{L}_{\text{FD}}$, it is essential to verify whether this leads to
improved alignment and accuracy.

Figure 5: Comparison of alignment/accuracy on the observed and unobserved
classes test data for $\bm{\theta}_{r}^{i}$ trained with
$\mathcal{L}_{\text{DR}}$ and $\mathcal{L}_{\text{Dr+}}$.

In Figure 5, we examine in both aspects and illustrate the training curve. Our
proposed algorithm, _i.e.,_ $\mathcal{L}_{\text{Dr+}}$, demonstrates superior
performance for both $\mathcal{O}^{i}$ and $\mathcal{U}^{i}$ in terms of
alignment and accuracy. Notably, even with the addition of a term to the dot-
regression loss, alignment is improved. We attribute this improvement to the
enhanced knowledge of the global model which is preserved by preventing the
forgetting of previously trained knowledge. Even though the proposed
regularizer demonstrates a reasonable regularizing effect, one question
remains: “ _Is it superior to other previously used regularizers?_ ”

Table 1: Synergy of various FL algorithms and regularizers. Baseline indicates
training FL models without a regularizer. FD denotes feature distillation,
which is the regularizer we use in FedDr+.

|  Sharding ($s=10$) |  LDA ($\alpha=0.1$)
---|---|---
Algorithm | Baseline |  +Prox [14] |  +KD [29] |  +NTD [17] |  +LD [30] | +FD | Baseline |  +Prox [14] |  +KD [29] |  +NTD [17] |  +LD [30] | +FD
FedAvg [1] | 37.22 | 30.27 | 35.14 | 35.56 | 34.83 | 37.82 | 42.52 | 36.09 | 41.48 | 41.34 | 43.36 | 43.10
FedBABU [2] | 46.20 | 36.71 | 45.50 | 45.09 | 45.81 | 45.31 | 47.37 | 39.04 | 45.58 | 45.56 | 46.46 | 44.77
SphereFed [4] | 43.90 | 1.36 | 41.01 | 43.47 | 41.73 | 45.21 | 46.98 | 1.46 | 45.22 | 46.25 | 43.84 | 48.61
FedETF [5] | 32.42 | 25.18 | 32.76 | 31.98 | 32.25 | 32.77 | 46.27 | 34.92 | 44.94 | 45.77 | 44.36 | 45.92
Dot-Regression | 42.52 | 5.42 | 46.60 | 45.78 | 47.52 | 48.69 | 42.72 | 7.47 | 48.19 | 33.08 | 49.09 | 50.79

###  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, as suggested in Eq. (missing) 4. Our study includes
FedAvg [1] without classifier freezing and other advanced frameworks such as
FedBABU [2], SphereFed [4], FedETF [5], and dot-regression, all of which
update local models while freezing the classifier. In addition to the FD
regularizer, we consider regularizers such as Prox [14] to constrain the
distance between local and global model parameters, and several logit-based
regularizers—KD [29], NTD [17, 31], and LD [30]—to keep logit-related
measurement of local models from deviating significantly from that of the
global model. Specifically, KD applies the softened softmax probability from
the logit vector, NTD does the same but excludes the true class dimension, and
LD distills the entire logit vector.

Table 1 demonstrates that FedDr+ (dot-regression + FD) achieves the best
performance. Generally, Prox tends to be less effective than logit-based
regularizers, which are often outperformed by FD across most algorithms. This
is because, as noted in Section 3.2, with the frozen classifier, features are
expected to have rich information to mitigate the drift. Prox uniformly
regularizes all data instances, whereas logit and feature regularizers adapt
to both model parameters and data instances, offering more refined control.
Specifically, FD regularizer, with its higher dimensionality, captures the
global model’s information more precisely than logit-based ones, resulting in
better synergy.

##  4 Experiments and Results

In this section, we present the experimental results of FedDr+, encompassing
both global federated learning (GFL) and personalized federated learning
(PFL). Additionally, we perform various hyper-parameter sensitivity analyses,
exploring the impact of varying local epochs and the client sampling ratio on
performance, as well as the effect of different $\beta$ values in FedDr+.

###  4.1 Experimental Setup

Dataset and models. To simulate a realistic FL scenario involving 100 clients,
we conduct extensive studies on two widely used datasets: CIFAR-10 and
CIFAR-100 [13]. For CIFAR-10, we employ VGG11 [32], while for CIFAR-100,
MobileNet [33] is used. The training data is distributed among 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
$\frac{|D_{\text{train}}|}{100\times s}$ and
$\frac{|D_{\text{test}}|}{100\times s}$ samples per class, where $s$ denotes
the number of shards per client. On the other hand, LDA involves sampling a
probability vector from Dirichlet distribution,
$p_{c}=(p_{c,1},p_{c,2},\cdots,p_{c,100})\sim\text{Dir}(\alpha)$, and
allocating a proportion $p_{c,k}$ of instances of class $c\in[C]$ to each
client $k\in[100]$. Smaller values of $s$ and $\alpha$ increase the level of
data heterogeneity.

Implementation details. In each round of communication, a fraction of clients
equal to 0.1 is randomly selected to participate in the training process. The
total number of communication rounds is 320. The initial learning rate and the
number of local epochs for CIFAR-10 and CIFAR-100 are determined through grid
searches, with the detailed process and results provided in Appendix C. The
learning rate $\eta$ is decayed by a factor of 0.1 at the 160th and 240th
communication rounds. The number of local epochs is set to 10 for CIFAR-10 and
3 for CIFAR-100 in the main experiments.333Table 2 is constructed using the
code structure from https://github.com/Lee-Gihun/FedNTD, while the rest of the
implementation is based on https://github.com/jhoon-oh/FedBABU.

###  4.2 Global Federated Learning Results

Table 2: 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$\pm$Y. The
best performance in each case is highlighted in bold.

NIID Partition Strategy: Sharding
---
| MobileNet on CIFAR-100 | VGG on CIFAR-10
Algorithm |  $s$=10 |  $s$=20 |  $s$=50 |  $s$=100 |  $s$=2 |  $s$=5 |  $s$=10
FedAvg [1] |  36.63$\pm$ 0.22 |  42.25$\pm$ 1.42 |  45.57$\pm$ 0.22 |  48.20$\pm$ 1.36 |  72.08$\pm$ 0.67 |  81.53$\pm$ 0.35 |  82.38$\pm$ 0.40
SCAFFOLD [34] ($\times 3$) |  46.08$\pm$ 0.37 |  48.15$\pm$ 1.21 |  49.31$\pm$ 0.62 |  50.73$\pm$ 0.42 |  75.49$\pm$ 0.42 |  84.14$\pm$ 0.13 |  85.11$\pm$ 0.29
FedNTD [17] |  34.05$\pm$ 1.19 |  41.78$\pm$ 0.31 |  46.42$\pm$ 0.63 |  47.17$\pm$ 0.32 |  72.21$\pm$ 0.59 |  69.96$\pm$ 17.10 |  81.99$\pm$ 0.42
FedExP [35] |  36.85$\pm$ 0.11 |  42.49$\pm$ 1.22 |  45.07$\pm$ 0.92 |  48.09$\pm$ 1.00 |  72.31$\pm$ 0.60 |  81.41$\pm$ 0.19 |  82.47$\pm$ 0.16
FedBABU [2] |  45.97$\pm$ 0.48 |  45.53$\pm$ 0.79 |  46.52$\pm$ 0.51 |  46.02$\pm$ 0.28 |  71.99$\pm$ 0.52 |  81.07$\pm$ 0.60 |  82.32$\pm$ 0.06
SphereFed [4] |  42.71$\pm$ 0.65 |  48.63$\pm$ 0.90 |  52.16$\pm$ 0.22 |  53.41$\pm$ 0.19 |  76.33$\pm$ 0.33 |  83.67$\pm$ 0.18 |  84.36$\pm$ 0.30
FedETF [5] |  31.37$\pm$ 0.72 |  42.22$\pm$ 0.77 |  47.47$\pm$ 0.67 |  49.00$\pm$ 0.74 |  67.81$\pm$ 0.94 |  80.78$\pm$ 0.68 |  82.60$\pm$ 0.46
FedDr+ (Ours) |  48.21$\pm$ 0.56 |  50.77$\pm$ 0.14 |  52.15$\pm$ 0.03 |  52.41$\pm$ 0.81 |  76.57$\pm$ 0.51 |  83.22$\pm$ 0.34 |  84.14$\pm$ 0.27
NIID Partition Strategy: LDA
| MobileNet on CIFAR-100 | VGG on CIFAR-10
Algorithm |  $\alpha$=0.05 |  $\alpha$=0.1 |  $\alpha$=0.2 |  $\alpha$=0.3 |  $\alpha$=0.1 |  $\alpha$=0.2 |  $\alpha$=0.3
FedAvg [1] |  35.58$\pm$ 1.35 |  42.10$\pm$ 0.60 |  44.78$\pm$ 0.72 |  45.73$\pm$ 0.88 |  68.71$\pm$ 1.82 |  77.75$\pm$ 0.26 |  80.76$\pm$ 0.51
SCAFFOLD [34] ($\times 3$) |  40.54$\pm$ 0.48 |  46.14$\pm$ 0.70 |  47.98$\pm$ 0.93 |  48.06$\pm$ 1.08 | (Failed) |  80.15$\pm$ 0.29 |  82.63$\pm$ 0.23
FedNTD [17] |  31.78$\pm$ 3.14 |  40.41$\pm$ 0.96 |  43.10$\pm$ 2.03 |  43.04$\pm$ 0.82 |  70.22$\pm$ 0.40 |  77.16$\pm$ 0.20 |  79.50$\pm$ 0.56
FedExP [35] |  34.39$\pm$ 1.77 |  40.85$\pm$ 1.32 |  44.47$\pm$ 0.28 |  45.44$\pm$ 0.14 |  70.14$\pm$ 0.53 |  78.09$\pm$ 0.21 |  80.40$\pm$ 0.54
FedBABU [2] |  41.97$\pm$ 1.01 |  45.77$\pm$ 0.28 |  44.28$\pm$ 0.45 |  44.80$\pm$ 0.63 |  65.15$\pm$ 3.66 |  77.03$\pm$ 0.25 |  79.91$\pm$ 0.13
SphereFed [4] |  39.56$\pm$ 0.48 |  46.54$\pm$ 0.58 |  49.41$\pm$ 0.78 |  49.22$\pm$ 0.86 |  67.49$\pm$ 3.49 |  80.05$\pm$ 0.40 |  82.62$\pm$ 0.66
FedETF [5] |  40.71$\pm$ 0.90 |  45.63$\pm$ 0.33 |  46.28$\pm$ 1.05 |  46.69$\pm$ 0.87 |  70.75$\pm$ 0.36 |  77.86$\pm$ 0.46 |  79.95$\pm$ 0.34
FedDr+ (Ours) |  45.12$\pm$ 1.00 |  49.48$\pm$ 0.50 |  50.67$\pm$ 0.88 |  51.15$\pm$ 0.65 |  72.07$\pm$ 2.26 |  80.90$\pm$ 0.02 |  82.42$\pm$ 0.10

We compare FedDr+ with a range of GFL algorithms, considering both non-
freezing and freezing classifier approaches. Among non-freezing classifiers,
FedDr+ competes with FedAvg [1], SCAFFOLD [34], FedNTD [17], and FedExP [35].
FedDr+ is also evaluated against freezing classifier algorithms such as
FedBABU [2], SphereFed [4], and FedETF [5]. Among the baseline algorithms,
SCAFFOLD incurs a communication cost three times higher per round, denoted as
($\times 3$). Our experiments encompass heterogeneous settings involving
sharding and LDA non-IID environments.

Table 2 summarizes the accuracy comparison between the state-of-the-art GFL
methods and FedAvg under various conditions. While specific methods
demonstrated effectiveness in particular scenarios, some of these frequently
underperformed relative to the robustness of FedAvg. For example, SCAFFOLD
shown strong performance in the less heterogeneous sharding setting on
CIFAR-10; however, it failed in model training under the highly heterogeneous
LDA condition with $\alpha=0.1$. Notably, FedDr+ consistently exceeded all
baseline methods in performance across diverse experimental conditions and
often achieved state-of-the-art results. FedDr+ demonstrated exceptional
performance in highly heterogeneous FL environments, particularly excelling in
the CIFAR-100 LDA configuration with $\alpha=0.05$, achieving a notable 3.15%
improvement over all baseline models.

###  4.3 Personalized Federated Learning Results

Table 3: PFL accuracy comparison with MobileNet on CIFAR-100. For PFL, we
denote the entries in the form of X$\pm$(Y), representing the mean and
standard deviation of personalized accuracies across all clients derived from
a single seed.

Algorithm |  $s$=10 |  $s$=20 |  $s$=100 |  $\alpha$=0.05 |  $\alpha$=0.1 |  $\alpha$=0.3
---|---|---|---|---|---|---
Local only ($\mathcal{L}_{\text{CE}}$) |  58.05$\pm$(8.11) |  42.45$\pm$(6.44) |  18.69$\pm$(3.28) |  55.39$\pm$(8.79) |  43.76$\pm$(7.46) |  27.75$\pm$(5.32)
Local only ($\mathcal{L}_{\text{CE}}$+ETF) |  58.01$\pm$(7.34) |  41.62$\pm$(5.91) |  18.92$\pm$(3.00) |  55.34$\pm$(9.13) |  43.37$\pm$(7.12) |  27.87$\pm$(5.34)
Local only ($\mathcal{L}_{\text{DR}}$) |  60.68$\pm$(7.77) |  44.61$\pm$(6.61) |  20.98$\pm$(3.49) |  58.56$\pm$(9.16) |  46.72$\pm$(7.29) |  30.88$\pm$(5.33)
FedPer [36] |  70.67$\pm$(7.19) |  57.27$\pm$6.66 |  24.30$\pm$(4.34) |  62.67$\pm$(7.65) |  53.43$\pm$(6.60) |  35.68$\pm$(4.82)
Per-FedAvg [37] |  32.13$\pm$(10.90) |  36.66$\pm$(8.86) |  41.27$\pm$(7.43) |  28.81$\pm$(8.68) |  35.56$\pm$(6.56) |  42.80$\pm$(4.76)
FedRep [38] |  63.14$\pm$(7.63) |  51.69$\pm$(6.50) |  26.31$\pm$(4.74) |  57.53$\pm$(8.05) |  49.60$\pm$(6.25) |  37.00$\pm$(4.82)
Ditto [39] |  39.26$\pm$(14.43) |  38.18$\pm$(9.96) |  44.53$\pm$(5.08) |  35.81$\pm$(14.83) |  37.81$\pm$(11.80) |  43.72$\pm$(5.12)
FedAvg-FT [1] |  69.81$\pm$(6.78) |  56.13$\pm$(5.77) |  47.66$\pm$(5.20) |  63.37$\pm$(9.28) |  56.79$\pm$(5.96) |  50.12$\pm$(3.67)
FedBABU-FT [2] |  80.14$\pm$(6.25) |  70.89$\pm$(5.60) |  52.14$\pm$(5.09) |  75.50$\pm$(6.40) |  70.83$\pm$(5.06) |  56.91$\pm$(3.74)
SphereFed-FT [4] |  81.90$\pm$(5.86) |  71.56$\pm$(5.78) |  55.83$\pm$(4.67) |  73.21$\pm$(7.08) |  70.00$\pm$(5.09) |  60.03$\pm$(3.99)
FedETF-FT [5] |  53.75$\pm$(7.35) |  52.94$\pm$(5.71) |  51.69$\pm$(5.03) |  52.96$\pm$(8.01) |  53.97$\pm$(5.40) |  51.67$\pm$(3.83)
FedDr+ FT (ours) |  84.10$\pm$(5.43) |  75.42$\pm$(4.80) |  56.76$\pm$(4.91) |  78.55$\pm$(6.16) |  74.75$\pm$(4.75) |  62.16$\pm$(3.73)

We introduce FedDr+ FT, inspired by prior work [2, 4, 5, 40], which enhances
personalization by leveraging local data to fine-tune the global federated
learning (GFL) model. We fine-tune the FedDr+ GFL model using
$\mathcal{L}_{\text{Dr+}}$ to create FedDr+ FT, _i.e.,_ 2-step approach. For a
comprehensive analysis, we compare FedDr+ FT with existing personalized
federated learning (PFL) methods, including 1-step approaches, _i.e.,_
creating PFL models from scratch, such as FedPer [36], Per-FedAvg [37], FedRep
[38], and Ditto [39], as well as 2-step methods such as FedAVG-FT, FedBABU-FT
[2], SphereFed-FT [4], and FedETF-FT [5]. Additionally, we compare these
methods with various simple local models that have not undergone federated
learning: (1) Local only ($\mathcal{L}_{\text{CE}}$), trained with
$\mathcal{L}_{\text{CE}}$, (2) Local only ($\mathcal{L}_{\text{CE}}$ \+ ETF),
trained with $\mathcal{L}_{\text{CE}}$ and initializing the classifier with an
ETF classifier, and (3) Local only ($\mathcal{L}_{\text{DR}}$), trained using
$\mathcal{L}_{\text{DR}}$.

In Table 3, we first compare the performance of simple local models in PFL by
examining $\mathcal{L}_{\text{DR}}$ and $\mathcal{L}_{\text{CE}}$. While
methods using $\mathcal{L}_{\text{CE}}$ show no significant differences,
utilizing $\mathcal{L}_{\text{DR}}$ leads to substantial performance
improvements in PFL across all settings. The “Local only
($\mathcal{L}_{\text{CE}}$)” and “Local only ($\mathcal{L}_{\text{CE}}$ \+
ETF)” methods exhibit similar performance due to the nearly classwise
orthogonal nature of randomly initialized classifiers [2, 41, 42, 43, 44].
With a 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 consistently
demonstrates superior performance. This aligns with previous research [45, 46]
suggesting that fine-tuning from a well-initialized model yields better PFL
performance. Additionally, compared with 1-step algorithms, FedDr+ FT
continues to show superiority, outperforming all baseline methods across all
settings.

###  4.4 Sensitivity Analysis

(a) Client sampling ratio

(b) Local epochs

(c) $\beta$ sensitivity

Figure 6: Performance of baselines and FedDr+ on CIFAR-100 ($\alpha$=0.1 and
$s$=10) with various analyses: (a) client sampling ratio, (b) the number of
local epochs, and (c) sensitivity to $\beta$.

We explore the impact of varying client sampling ratio and local epochs on
performance, as well as the effect of different $\beta$ values in FedDr+, as
detailed in Figure 6. All experiments are conducted on MobileNet using the
CIFAR-100 dataset with sharding ($s$=10) and LDA ($\alpha$=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), FedBABU and SphereFed (all with classifier freezing)
show performance improvements with increasing fraction ratios, but FedDr+
consistently outperforms the baselines. The number of local epochs is crucial
in FL; too few epochs result in underfitting, while too many cause client
drift, degrading 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 or highly competitive.

Weight ratio $\beta$ analysis. We analyze the effect of scaling parameter in
FedDr+ by varying $\beta$ while keeping other hyperparameters constant. The
performance is evaluated for $\beta\in\\{0,0.1,0.3,0.5,0.7,0.9,1.0\\}$. When
$\beta=0$, only feature distillation is applied, and when $\beta=1$, only dot-
regression is used. $\beta\in\\{0,1\\}$ are generally less effective, whereas
$\beta\in\\{0.3,0.5,0.7,0.9\\}$ show consistently good performance, indicating
a balanced approach is beneficial.

##  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 [1, 2, 14, 47]. This approach faces challenges due to
data heterogeneity across clients, causing instability in the learning process
[34, 48]. To address this problem, strategies like classifier variance
reduction in FedPVR [49] and virtual features in CCVR [48] have been proposed.
Additionally, it is essential to distinguish between Global Federated Learning
(GFL) and Personalized Federated Learning (PFL), as these are crucial concepts
in FL. GFL aims to improve a single global model’s performance across clients
by addressing data heterogeneity through methods like client drift mitigation
[14, 34, 35], enhanced aggregation schemes [50, 51], and data sharing
techniques using public or synthesized datasets [20, 48, 52]. Otherwise, PFL
focuses on creating personalized models for individual clients by decoupling
feature extractors and classifiers for unique updates [2, 36, 38], modifying
local loss functions [37, 39], and using prototype communication techniques
[53, 54].

Frozen classifier in FL. By focusing on alignment, previous studies have
attempted to mitigate data heterogeneity by freezing the classifier [2, 4, 5].
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 federated learning to
achieve a 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 alignment performance
compared to previous approaches. However, optimizing the dot-regression loss
to align feature vectors with a fixed classifier caused the local model to
lose information on unobserved classes, thereby degrading global model
performance. To address these issues, FedLoGe [8] employing realignment
techniques to ensure the well-aligned local model’s performance translated to
the global model. Additionally, in FedGELA [6], the classifier is globally
fixed as a simplex ETF while being locally adapted to personal distributions.
Also, FedPAC [54] addressed these challenges by leveraging global semantic
knowledge for explicit local-global feature alignment. Besides alignment-
focused methods, there have been various attempts to maintain good local model
performance in the global model [55, 56, 57].

Knowledge distillation in FL. Multiple studies have explored knowledge
transfer techniques [18, 19, 58, 59, 60]. Specifically, knowledge distillation
(KD) has been widely studied in FL settings, such as in FedMD [16] and FedDF
[20], where a pretrained teacher model transfers knowledge to a student model.
Additional distillation-based methods, such as FedFed [61] and co-distillation
framework for PFL [23, 24], have also been explored. In contrast to existing
methods, we propose a loss function incorporating feature 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 PFL.

##  6 Conclusion

Motivated by the recent FL methods enhancing feature alignment with a 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-
classifier alignment, we find it improves alignment and accuracy in local
models but degrades the performance of the global model. This happens because
local clients trained with dot-regression tend to forget classes that have not
been observed. To address this, we propose FedDr+, combining dot-regression
with a 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).

## References

  * [1] Brendan McMahan, Eider Moore, Daniel Ramage, Seth Hampson, and Blaise Aguera y Arcas.  Communication-efficient learning of deep networks from decentralized data.  In Artificial intelligence and statistics, pages 1273–1282. PMLR, 2017.
  * [2] Jaehoon Oh, SangMook Kim, and Se-Young Yun.  FedBABU: Toward enhanced representation for federated image classification.  In International Conference on Learning Representations, 2022.
  * [3] 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 arXiv:2007.13518, 2020.
  * [4] Xin Dong, Sai Qian Zhang, Ang Li, and HT Kung.  Spherefed: Hyperspherical federated learning.  In European Conference on Computer Vision, pages 165–184. Springer, 2022.
  * [5] 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/CVF International Conference on Computer Vision, pages 5319–5329, 2023.
  * [6] Ziqing Fan, Jiangchao Yao, Bo Han, Ya Zhang, Yanfeng Wang, et al.  Federated learning with bilateral curation for partially class-disjoint data.  Advances in Neural Information Processing Systems, 36, 2024.
  * [7] Chenxi Huang, Liang Xie, Yibo Yang, Wenxiao Wang, Binbin Lin, and Deng Cai.  Neural collapse inspired federated learning with non-iid data, 2023.
  * [8] 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.
  * [9] Vardan Papyan, XY Han, and David L Donoho.  Prevalence of neural collapse during the terminal phase of deep learning training.  Proceedings of the National Academy of Sciences, 117(40):24652–24663, 2020.
  * [10] 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.
  * [11] Yibo Yang, Haobo Yuan, Xiangtai Li, Zhouchen Lin, Philip Torr, and Dacheng Tao.  Neural collapse inspired feature-classifier alignment for few-shot class incremental learning.  arXiv preprint arXiv:2302.03004, 2023.
  * [12] 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, pages 995–1005, 2021.
  * [13] Alex Krizhevsky, Vinod Nair, and Geoffrey Hinton.  Cifar-10 and cifar-100 datasets.  URl: https://www. cs. toronto. edu/kriz/cifar. html, 6(1):1, 2009.
  * [14] 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.
  * [15] Chaoyang He, Murali Annavaram, and Salman Avestimehr.  Group knowledge transfer: Federated learning of large cnns at the edge.  Advances in Neural Information Processing Systems, 33:14068–14080, 2020.
  * [16] Daliang Li and Junpu Wang.  Fedmd: Heterogenous federated learning via model distillation.  arXiv preprint arXiv:1910.03581, 2019.
  * [17] Gihun Lee, Minchan Jeong, Yongjin Shin, Sangmin Bae, and Se-Young Yun.  Preservation of the global knowledge by not-true distillation in federated learning.  Advances in Neural Information Processing Systems, 35:38461–38474, 2022.
  * [18] Sohei Itahara, Takayuki Nishio, Yusuke Koda, Masahiro Morikura, and Koji Yamamoto.  Distillation-based semi-supervised federated learning for communication-efficient collaborative training with non-iid private data.  IEEE Transactions on Mobile Computing, 22(1):191–205, 2021.
  * [19] 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.
  * [20] Tao Lin, Lingjing Kong, Sebastian U Stich, and Martin Jaggi.  Ensemble distillation for robust model fusion in federated learning.  Advances in Neural Information Processing Systems, 33:2351–2363, 2020.
  * [21] Wei-Chun Chen, Chia-Che Chang, and Che-Rung Lee.  Knowledge distillation with feature maps for image classification.  In Computer Vision–ACCV 2018: 14th Asian Conference on Computer Vision, Perth, Australia, December 2–6, 2018, Revised Selected Papers, Part III 14, pages 200–215. Springer, 2019.
  * [22] Biao Qian, Yang Wang, Hongzhi Yin, Richang Hong, and Meng Wang.  Switchable online knowledge distillation.  In European Conference on Computer Vision, pages 449–466. Springer, 2022.
  * [23] 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.
  * [24] Yae Jee Cho, Jianyu Wang, Tarun Chirvolu, and Gauri Joshi.  Communication-efficient and model-heterogeneous personalized federated learning via clustered knowledge transfer.  IEEE Journal of Selected Topics in Signal Processing, 17(1):234–247, 2023.
  * [25] 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, pages 1921–1930, 2019.
  * [26] Quanquan Li, Shengying Jin, and Junjie Yan.  Mimicking very efficient network for object detection.  In Proceedings of the ieee conference on computer vision and pattern recognition, pages 6356–6364, 2017.
  * [27] Jingzhi Li, Zidong Guo, Hui Li, Seungju Han, Ji-won Baek, Min Yang, Ran Yang, and Sungjoo Suh.  Rethinking feature-based knowledge distillation for face recognition.  In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 20156–20165, 2023.
  * [28] Emanuel Ben-Baruch, Matan Karklinsky, Yossi Biton, Avi Ben-Cohen, Hussam Lawen, and Nadav Zamir.  It’s all in the head: Representation knowledge distillation through classifier sharing.  arXiv preprint arXiv:2201.06945, 2022.
  * [29] Geoffrey Hinton, Oriol Vinyals, and Jeff Dean.  Distilling the knowledge in a neural network.  arXiv preprint arXiv:1503.02531, 2015.
  * [30] Taehyeon Kim, Jaehoon Oh, NakYil Kim, Sangwook Cho, and Se-Young Yun.  Comparing kullback-leibler divergence and mean squared error loss in knowledge distillation.  arXiv preprint arXiv:2105.08919, 2021.
  * [31] 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, pages 11953–11962, 2022.
  * [32] Karen Simonyan and Andrew Zisserman.  Very deep convolutional networks for large-scale image recognition.  arXiv preprint arXiv:1409.1556, 2014.
  * [33] 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 arXiv:1704.04861, 2017.
  * [34] Sai Praneeth Karimireddy, Satyen Kale, Mehryar Mohri, Sashank Reddi, Sebastian Stich, and Ananda Theertha Suresh.  Scaffold: Stochastic controlled averaging for federated learning.  In International conference on machine learning, pages 5132–5143. PMLR, 2020.
  * [35] Divyansh Jhunjhunwala, Shiqiang Wang, and Gauri Joshi.  Fedexp: Speeding up federated averaging via extrapolation.  In The Eleventh International Conference on Learning Representations, 2023.
  * [36] Manoj Ghuhan Arivazhagan, Vinay Aggarwal, Aaditya Kumar Singh, and Sunav Choudhary.  Federated learning with personalization layers.  arXiv preprint arXiv:1912.00818, 2019.
  * [37] Alireza Fallah, Aryan Mokhtari, and Asuman Ozdaglar.  Personalized federated learning with theoretical guarantees: A model-agnostic meta-learning approach.  Advances in Neural Information Processing Systems, 33:3557–3568, 2020.
  * [38] Liam Collins, Hamed Hassani, Aryan Mokhtari, and Sanjay Shakkottai.  Exploiting shared representations for personalized federated learning.  In International conference on machine learning, pages 2089–2099. PMLR, 2021.
  * [39] Tian Li, Shengyuan Hu, Ahmad Beirami, and Virginia Smith.  Ditto: Fair and robust federated learning through personalization.  In International Conference on Machine Learning, pages 6357–6368. PMLR, 2021.
  * [40] Seongyoon Kim, Gihun Lee, Jaehoon Oh, and Se-Young Yun.  Fedfn: Feature normalization for alleviating data heterogeneity problem in federated learning.  arXiv preprint arXiv:2311.13267, 2023.
  * [41] Andrew M Saxe, James L McClelland, and Surya Ganguli.  Exact solutions to the nonlinear dynamics of learning in deep linear neural networks.  arXiv preprint arXiv:1312.6120, 2013.
  * [42] Xavier Glorot and Yoshua Bengio.  Understanding the difficulty of training deep feedforward neural networks.  In Proceedings of the thirteenth international conference on artificial intelligence and statistics, pages 249–256. JMLR Workshop and Conference Proceedings, 2010.
  * [43] 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, pages 1026–1034, 2015.
  * [44] 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, pages 8109–8118, 2018.
  * [45] 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.  arXiv preprint arXiv:2206.15387, 2022.
  * [46] 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 Eleventh International Conference on Learning Representations, 2023.
  * [47] Durmus Alp Emre Acar, Yue Zhao, Ramon Matas, Matthew Mattina, Paul Whatmough, and Venkatesh Saligrama.  Federated learning based on dynamic regularization.  In International Conference on Learning Representations, 2021.
  * [48] Mi Luo, Fei Chen, Dapeng Hu, Yifan Zhang, Jian Liang, and Jiashi Feng.  No fear of heterogeneity: Classifier calibration for federated learning with non-iid data.  Advances in Neural Information Processing Systems, 34:5972–5984, 2021.
  * [49] Bo Li, Mikkel N Schmidt, Tommy S Alstrøm, and Sebastian U Stich.  On the effectiveness of partial variance reduction in federated learning with heterogeneous data.  arXiv preprint arXiv:2212.02191, 2022.
  * [50] Hongyi Wang, Mikhail Yurochkin, Yuekai Sun, Dimitris Papailiopoulos, and Yasaman Khazaeni.  Federated learning with matched averaging.  In International Conference on Learning Representations, 2020.
  * [51] Jianyu Wang, Qinghua Liu, Hao Liang, Gauri Joshi, and H Vincent Poor.  Tackling the objective inconsistency problem in heterogeneous federated optimization.  Advances in neural information processing systems, 33:7611–7623, 2020.
  * [52] Yue Zhao, Meng Li, Liangzhen Lai, Naveen Suda, Damon Civin, and Vikas Chandra.  Federated learning with non-iid data.  arXiv preprint arXiv:1806.00582, 2018.
  * [53] 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 Artificial Intelligence, pages 8432–8440, 2022.
  * [54] 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.
  * [55] 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 Representations, 2023.
  * [56] Xuming An, Li Shen, Han Hu, and Yong Luo.  Federated learning with manifold regularization and normalized update reaggregation.  Advances in Neural Information Processing Systems, 36, 2024.
  * [57] Hong-You Chen and Wei-Lun Chao.  On bridging generic and personalized federated learning for image classification.  In International Conference on Learning Representations, 2022.
  * [58] Yae Jee Cho, Andre Manoel, Gauri Joshi, Robert Sim, and Dimitrios Dimitriadis.  Heterogeneous ensemble knowledge transfer for training large models in federated learning.  arXiv preprint arXiv:2204.12703, 2022.
  * [59] Shuai Wang, Yexuan Fu, Xiang Li, Yunshi Lan, Ming Gao, et al.  Dfrd: Data-free robustness distillation for heterogeneous federated learning.  Advances in Neural Information Processing Systems, 36, 2024.
  * [60] Zhongyi Cai, Ye Shi, Wei Huang, and Jingya Wang.  Fed-co $\\_${$2$}: Cooperation of online and offline models for severe data heterogeneity in federated learning.  Advances in Neural Information Processing Systems, 36, 2024.
  * [61] 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.

\- Appendix -

FedDr+: Stabilizing Dot-regression with Global Feature Distillation for
Federated Learning

We organized notations at Appendix A. In Appendix B, we show the pulling and
pushing gradients of the CE loss in detail. Then, we elucidate the
experimental setup in Appendix C, encompassing dataset description, model
specifications, NIID partition, and hyperparameter search. In Appendix D, we
present additional experiment results of PFL and elapsed time measurement.

##  Appendix A Notations

Table 4: Notations used throughout the paper.

Indices |
---|---
$c\in[C]$ | Index for a class
$r\in[R]$ | Index for FL round
$i\in[N]$ | Index for a client
Dataset |
$D_{\text{train}}^{i}$ | Training dataset for client $i$
$D_{\text{test}}^{i}$ | Test dataset for client $i$
$(x,y)\in D_{\text{train,test}}^{i}\,;(x,y)\sim\mathcal{D}^{i}$ | Data on client $i$ sampled from distribution $\mathcal{D}^{i}$
| ($x$: input data, $y$: class label)
$\mathcal{O}^{i}$ | Dataset consists of observed classes in client $i$
$\mathcal{U}^{i}$ | Dataset consists of unobserved classes in client $i$
Parameters |
$\bm{\theta}$ | Feature extractor weight parameters
$\bm{V}=[v_{1},\ldots,v_{C}]\in\mathbb{R}^{C\times d}$ | Classifier weight parameters (frozen during training)
$v_{c},c\in[C]$ |  $c$-th row vector of $\bm{V}$
$\bm{\Theta}=(\bm{\theta},\bm{V})$ | All model parameters
$\bm{\Theta}_{r}^{g}=(\bm{\theta}_{r}^{g},\bm{V})$ | Aggregated global model parameters at round $r$
$\bm{\Theta}_{r}^{i}=(\bm{\theta}_{r}^{i},\bm{V})$ | Trained model parameters on client $i$ at round $r$
Model Forward |
$p(x;\bm{\theta})\in\mathbb{R}^{C}$ | Softmax probability of input $x$
$p_{c}(x;\bm{\theta}),c\in[C]$ |  $c$-th element of $p(x;\bm{\theta})$
$\mathcal{L}_{\text{CE}}(x;\theta)=-\log p_{y}(x;\bm{\theta})$ | Cross-entropy loss of input $x$
$f(x;\bm{\theta})\in\mathbb{R}^{d}$ | Feature vector of input $x$
$z(x;\bm{\theta})=f(x;\bm{\theta})\bm{V}^{\top}\in\mathbb{R}^{C}$ | Logit vector of input $x$
$z_{c}(x;\bm{\theta}),c\in[C]$ |  $c$-th element of $z(x;\bm{\theta})$

##  Appendix B Preliminaries: Pulling and Pushing Feature Gradients in CE

In this section, we first calculate the classifier gradient for features and
introduce the pulling and pushing effects of the cross-entropy objective.

###  B.1 Feature Gradient of $\mathcal{L}_{\text{CE}}$

We first provide two lemmas supporting Proposition 1, explaining the behavior
of pulling and pushing feature gradients in the cross-entropy (CE) loss.

######  Lemma 1.

For all $c,c^{\prime}\in[C],\frac{\partial
p_{c^{\prime}}(x;\bm{\theta})}{\partial
z_{c}(x;\bm{\theta})}=\begin{cases}p_{c}(x;\bm{\theta})\cdot(1-p_{c}(x;\bm{\theta}))&\text{if
}c=c^{\prime}\\\ -p_{c}(x;\bm{\theta})\cdot
p_{c^{\prime}}(x;\bm{\theta})&\text{else }\end{cases}.$

###### Proof.

Note that
$p(x;\bm{\theta})=\bigg{[}\frac{\exp(z_{j}(x;\bm{\theta}))}{\sum_{i=1}^{C}\exp(z_{i}(x;\bm{\theta}))}\bigg{]}_{j=1}^{C}\in\mathbb{R}^{C}$.
Then,

  1. (i)

$c=c^{\prime}$ case:

$\displaystyle\frac{\partial p_{c}(x;\bm{\theta})}{\partial
z_{c}(x;\bm{\theta})}$ $\displaystyle=\frac{\partial}{\partial
z_{c}(x;\bm{\theta})}\left\\{\frac{\exp(z_{c}(x;\bm{\theta}))}{\sum_{i=1}^{C}\exp(z_{i}(x;\bm{\theta}))}\right\\}$
$\displaystyle=\frac{\exp(z_{c}(x;\bm{\theta}))\left(\sum_{i=1}^{C}\exp(z_{i}(x;\bm{\theta}))\right)-{\exp(z_{c}(x;\bm{\theta}))}^{2}}{\left(\sum_{i=1}^{C}\exp(z_{i}(x;\bm{\theta}))\right)^{2}}$
$\displaystyle=p_{c}(x;\bm{\theta})-{p_{c}(x;\bm{\theta})}^{2}=p_{c}(x;\bm{\theta})(1-p_{c}(x;\bm{\theta})).$

  2. (ii)

$c\neq c^{\prime}$ case:

$\displaystyle\frac{\partial p_{c^{\prime}}(x;\bm{\theta})}{\partial
z_{c}(x;\bm{\theta})}$ $\displaystyle=\frac{\partial}{\partial
z_{c}(x;\bm{\theta})}\left\\{\frac{\exp(z_{c^{\prime}}(x;\bm{\theta}))}{\sum_{i=1}^{C}\exp(z_{i}(x;\bm{\theta}))}\right\\}=\frac{-\exp(z_{c}(x;\bm{\theta}))\exp(z_{c^{\prime}}(x;\bm{\theta}))}{\left(\sum_{i=1}^{C}\exp(z_{i}(x;\bm{\theta}))\right)^{2}}$
$\displaystyle=-p_{c}(x;\bm{\theta})p_{c^{\prime}}(x;\bm{\theta}).$

∎

######  Lemma 2.

$\nabla_{z(x;\bm{\theta})}\mathcal{L}_{\text{CE}}(x,y;\bm{\theta})=p(x;\bm{\theta})-\mathbf{e}_{y}$,
where $\mathbf{e}_{y}\in\mathbb{R}^{C}$ is the unit vector with its $y$-th
element as 1.

###### Proof.

$\displaystyle\frac{\partial\mathcal{L}_{\text{CE}}(x,y;\bm{\theta})}{\partial
z_{c}(x;\bm{\theta})}$ $\displaystyle=-\frac{\partial}{\partial
z_{c}(x;\bm{\theta})}\log
p_{y}(x;\bm{\theta})=-\frac{1}{p_{y}(x;\bm{\theta})}\frac{\partial
p_{y}(x;\bm{\theta})}{\partial z_{c}(x;\bm{\theta})}$
$\displaystyle\stackrel{{\scriptstyle(\bigstar)}}{{=}}\begin{cases}p_{c}(x;\bm{\theta})-1&\text{if
}c=y\\\ p_{c}(x;\bm{\theta})&\text{else
}\end{cases}=p_{c}(x;\bm{\theta})-\mathds{1}\\{c=y\\}.$

Note that $(\bigstar)$ holds by the Lemma 1. Therefore, the desired result is
satisfied. ∎

######  Proposition 1.

Given $(x,y)$, the gradient of the $\mathcal{L}_{\text{CE}}$ with respect to
$f(x;\bm{\theta})$ is given by:

$\nabla_{f(x;\bm{\theta})}\mathcal{L}_{\text{CE}}(x,y;\bm{\theta})=-(1-p_{y}(x;\bm{\theta}))v_{y}+\sum_{c\in[C]\setminus\\{y\\}}p_{c}(x;\bm{\theta})v_{c}.$
(5)

###### Proof.

$\displaystyle\nabla_{f(x;\bm{\theta})}\mathcal{L}_{\text{CE}}(x,y;\bm{\theta})$
$\displaystyle\stackrel{{\scriptstyle(\clubsuit)}}{{=}}\left[\nabla_{f(x;\bm{\theta})}z_{1}(x;\bm{\theta})|\cdots|\nabla_{f(x;\bm{\theta})}z_{C}(x;\bm{\theta})\right]\nabla_{z(x;\bm{\theta})}\mathcal{L}_{\text{CE}}(x,y;\bm{\theta})$
$\displaystyle=\sum_{c=1}^{C}\frac{\partial\mathcal{L}_{\text{CE}}(x,y;\bm{\theta})}{\partial
z_{c}(x;\bm{\theta})}\nabla_{f(x;\bm{\theta})}z_{c}(x;\bm{\theta})$
$\displaystyle=\frac{\partial\mathcal{L}_{\text{CE}}(x,y;\bm{\theta})}{\partial
z_{y}(x;\bm{\theta})}\nabla_{f(x;\bm{\theta})}z_{y}(x;\bm{\theta})+\sum_{c\in[C]\setminus\\{y\\}}\frac{\partial\mathcal{L}_{\text{CE}}(x,y;\bm{\theta})}{\partial
z_{c}(x;\bm{\theta})}\nabla_{f(x;\bm{\theta})}z_{c}(x;\bm{\theta})$
$\displaystyle=\frac{\partial\mathcal{L}_{\text{CE}}(x,y;\bm{\theta})}{\partial
z_{y}(x;\bm{\theta})}v_{y}+\sum_{c\in[C]\setminus\\{y\\}}\frac{\partial\mathcal{L}_{\text{CE}}(x,y;\bm{\theta})}{\partial
z_{c}(x;\bm{\theta})}v_{c}$
$\displaystyle\stackrel{{\scriptstyle(\spadesuit)}}{{=}}-(1-p_{y}(x;\bm{\theta}))v_{y}+\sum_{c\in[C]\setminus\\{y\\}}p_{c}(x;\bm{\theta})v_{c}.$

Employing the chain rule for $(\clubsuit)$ and invoking Lemma 2 for
$(\spadesuit)$ confirms the result.

∎

###  B.2 Physical Meaning of
$\nabla_{f(x;{\theta})}\mathcal{L}_{\text{CE}}(x,y;{\theta})$

Note that $\nabla_{f(x;\bm{\theta})}\mathcal{L}_{\text{CE}}(x,y;\bm{\theta})$
has two components: $\mathbf{F}_{\text{Pull}}=(1-p_{y}(x;\bm{\theta}))v_{y}$
and
$\mathbf{F}_{\text{Push}}=-\sum_{c\in[C]\setminus\\{y\\}}p_{c}(x;\bm{\theta})v_{c}$.
$\mathbf{F}_{\text{Pull}}$ adjusts the feature vector in the positive
direction of the actual class index’s classifier vector $v_{y}$, guiding
alignment towards $v_{y}$. Conversely, $\mathbf{F}_{\text{Push}}$ adjusts the
feature vector in the negative direction of the vectors in the not-true class
set $[C]\setminus\\{y\\}$, inducing misalignment towards $v_{c}$ for
$c\in[C]\setminus\\{y\\}$.

##  Appendix C Experimental Setup

###  C.1 Code Implementation

Our implementations are conducted using the PyTorch framework. Specifically,
the experiments presented in Table 2 are executed on a single NVIDIA RTX 3090
GPU, based on the code structure from the following repository:
https://github.com/Lee-Gihun/FedNTD. The other parts of our study are carried
out on a single NVIDIA A5000 GPU, utilizing the code framework from
https://github.com/jhoon-oh/FedBABU.

###  C.2 Datasets, Model, and Optimizer

To simulate a realistic FL scenario, we conduct extensive studies on two
widely used datasets: CIFAR-10 and CIFAR-100 [13]. A momentum optimizer is
utilized for all experiments. Unless otherwise noted, the basic setting of our
experiments follows the dataset statistics, FL scenario specifications, and
optimizer hyperparameters summarized in Table 5.

Table 5: Summary of Dataset, Model, FL System, and Optimizer Specifications Datasets | $C$ | $|D_{\text{train}}|$ | $|D_{\text{test}}|$ | $N$ | $R$ | $r$ | $E$ | $B$ | $m$ | $\lambda$
---|---|---|---|---|---|---|---|---|---|---
CIFAR-10 | 10 | 50000 | 10000 | 100 | 320 | 0.1 | 10 | 50 | 0.9 | 1e-5
CIFAR-100 | 100 | 50000 | 10000 | 100 | 320 | 0.1 | 3 | 50 | 0.9 | 1e-5

Note: In terms of dataset information, $C$ represents the number of classes in
the dataset, with $|D_{\text{train}}|$ and $|D_{\text{test}}|$ 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,
$r$ 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 $B$ is the batch size, and $m$ and $\lambda$ represent the momentum and
weight decay parameters, respectively. The initial learning rate $\eta$ is
decayed by a factor of 0.1 at the 160th and 240th communication rounds. The
initial learning rate $\eta$ and batch size $B$ were determined via extensive
grid search for each algorithm, details outlined in Appendix C.4.

###  C.3 Non-IID Partition Strategies

To induce heterogeneity in each client’s training and test data
($D_{\text{train}}^{i},D_{\text{test}}^{i}$), we distribute the entire class-
balanced datasets, $D_{\text{train}}$ and $D_{\text{test}}$, among 100 clients
using both sharding and Latent Dirichlet Allocation (LDA) partitioning
strategies:

  * •

Sharding [1, 2]: We organize the $D_{\text{train}}$ and $D_{\text{test}}$ by
label and divide them into non-overlapping shards of equal size. Each shard
encompasses $\frac{|D_{\text{train}}|}{100\times s}$ and
$\frac{|D_{\text{test}}|}{100\times s}$ samples of the same class, where $s$
denotes the number of shards per client. This sharding technique is used to
create $D_{\text{train}}^{i}$ and $D_{\text{test}}^{i}$, 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 a
result, each client has access to a maximum of $s$ different classes.
Decreasing the number of shards per user $s$ increases the level of data
heterogeneity among clients.

  * •

Latent Dirichlet Allocation (LDA) [48, 50]: We utilize the LDA technique to
create $D_{\text{train}}^{i}$ from $D_{\text{train}}$. This involves sampling
a probability vector
$p_{c}=(p_{c,1},p_{c,2},\cdots,p_{c,100})\sim\text{Dir}(\alpha)$ and
allocating a proportion $p_{c,k}$ of instances of class $c\in[C]$ to each
client $k\in[100]$. Here, $Dir(\alpha)$ represents the Dirichlet distribution
with the concentration parameter $\alpha$. The parameter $\alpha$ controls the
strength of data heterogeneity, with smaller values leading to stronger
heterogeneity among clients. For $D_{\text{test}}^{i}$, we randomly sample
from $D_{\text{test}}$ to match the class frequency of $D_{\text{train}}^{i}$
and distribute it to each client $i$.

Table 6: Hyperparameters for VGG11 training on CIFAR-10.

| Feature un-normalized algorithms | Feature normalized algorithms
---|---|---
Hyperparameters | FedAvg | FedBABU | SCAFFOLD | FedNTD | FedExP | FedETF | SphereFed |  FedDr+ (Ours)
$\eta$ | 0.01 | 0.01 | 0.01 | 0.01 | 0.01 | 0.05 | 0.55 | 0.35
Additional | None | None | None |  $(\beta,\tau)$=(1,3) |  $\epsilon$=0.001 |  $(\beta,\tau)$=(1,1) | None |  $\beta$=0.9

Table 7: Hyperparameters for MobileNet training on CIFAR-100.

| Feature un-normalized algorithms | Feature normalized algorithms
---|---|---
Hyperparameters | FedAvg | FedBABU | SCAFFOLD | FedNTD | FedExP | FedETF | SphereFed |  FedDr+ (Ours)
$\eta$ | 0.1 | 0.1 | 0.1 | 0.1 | 0.1 | 0.5 | 6.5 | 5.0
Additional | None | None | None |  $(\beta,\tau)$=(1,3) |  $\epsilon$=0.001 |  $(\beta,\tau)$=(1,1) | None |  $\beta$=0.9

###  C.4 Hyperparameter Search for $\eta$ and $E$

To optimize the initial learning rate ($\eta$) and the number of local epochs
($E$) for our algorithm, we conduct a grid search on the CIFAR-10 and
CIFAR-100 datasets. The process and reasoning are outlined below.

####  C.4.1 Rationale for Varying Initial Learning Rate ($\eta$)

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(x;\bm{\theta})$ are
normalized, the resulting gradient is scaled by
$\frac{1}{\|f(x;\bm{\theta})\|_{2}}$. This scaling effect necessitates a grid
search across various learning rates to account for the differences in
learning behavior.

####  C.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 is essential by
exploring a range of values.

####  C.4.3 Grid Search Process and Results

Considering the above reasons, we perform grid search for $\eta$ and $E$ on
CIFAR-10 and CIFAR-100 datasets. The grid search for CIFAR-10 uses a shard
size of 2, while for CIFAR-100, a shard size of 10 is used. The detailed
procedures for each dataset are provided below. These optimal settings have
also been confirmed to yield good performance in less heterogeneous settings.

CIFAR-10. We examine $\eta$ 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}.
The optimal learning rates vary by algorithm, and the results are summarized
in Table 6. Table 6 also includes the additional hyperparameters used for each
algorithm. The notation for these additional hyperparameters follows the
conventions used throughout this paper. The optimal number of local epochs is
found to be 10 for every algorithm.

CIFAR-100. We examine $\eta$ values from {0.1, 0.3, 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}. A 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 7. Table 7 also includes the
additional hyperparameters used for each algorithm. The notation for these
additional hyperparameters follows the conventions used throughout this paper.
The optimal number of local epochs is found to be 3 for every algorithm.

##  Appendix D Additional Experiment Results

Table 8: PFL accuracy comparison with MobileNet on CIFAR-100. For PFL, we
denote the entries in the form of X$\pm$(Y), representing the mean and
standard deviation of personalized accuracies across all clients derived from
a single seed.

Algorithm |  $s$=10 |  $s$=20 |  $s$=100 |  $\alpha$=0.05 |  $\alpha$=0.1 |  $\alpha$=0.3
---|---|---|---|---|---|---
Dot-Regression | 42.52 | 49.02 | 52.86 |  30.31$\pm$7.95 |  37.52$\pm$5.60 |  47.08$\pm$3.69
Dot-Regression FT ($\mathcal{L}_{\text{DR}}$) |  80.84$\pm$(5.99) |  74.18$\pm$(5.78) |  56.84$\pm$(5.04) |  72.02$\pm$(6.80) |  66.96$\pm$(5.36) |  60.34$\pm$(3.66)
Dot-Regression FT ($\mathcal{L}_{\text{Dr+}}$) |  80.82$\pm$(6.12) |  73.73$\pm$(5.75) |  56.69$\pm$(4.95) |  71.85$\pm$(7.03) |  66.59$\pm$(5.32) |  59.87$\pm$(3.65)
FedDr+ (ours) | 48.69 | 51.00 | 53.23 |  39.63$\pm$9.12 |  45.83$\pm$6.18 |  48.04$\pm$3.44
FedDr+ FT ($\mathcal{L}_{\text{DR}}$) (ours) |  84.23$\pm$(5.44) |  75.73$\pm$(4.79) |  56.90$\pm$(4.85) |  78.65$\pm$(6.17) |  74.86$\pm$(4.77) |  62.47$\pm$(3.72)
FedDr+ FT ($\mathcal{L}_{\text{Dr+}}$) (ours) |  84.10$\pm$(5.43) |  75.42$\pm$(4.80) |  56.76$\pm$(4.91) |  78.55$\pm$(6.16) |  74.75$\pm$(4.75) |  62.16$\pm$(3.73)

###  D.1 Personalized Federated Learning Results

We introduce FedDr+ FT and dot-regression FT, inspired by prior work [2, 4, 5,
40]. These methods enhance personalization by leveraging local data to fine-
tune the GFL model. We investigate the impact of fine-tuning using
$\mathcal{L}_{\text{Dr+}}$ and $\mathcal{L}_{\text{DR}}$ loss for each GFL
model to assess their effectiveness on personalized accuracy. Performance
metrics without standard deviations indicate results on $D_{\text{test}}$,
obtained from the GFL model after the initial step in the 2-step method. Our
experiments involve heterogeneous settings with sharding and LDA non-IID
environments, using MobileNet on CIFAR-100 datasets. We set $s$ as 10, 20, and
100, and the LDA concentration parameter ($\alpha$) as 0.05, 0.1, and 0.3.
Table 8 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 PFL model, again using $\mathcal{L}_{\text{DR}}$ or
$\mathcal{L}_{\text{Dr+}}$. This results in four combinations: Dot-Regression
FT ($\mathcal{L}_{\text{DR}}$), Dot-Regression FT
($\mathcal{L}_{\text{Dr+}}$), FedDr+ FT ($\mathcal{L}_{\text{DR}}$), and
FedDr+ FT ($\mathcal{L}_{\text{Dr+}}$). When the GFL model is fixed, using
$\mathcal{L}_{\text{DR}}$ for fine-tuning consistently outperforms
$\mathcal{L}_{\text{Dr+}}$ across all settings, because dot-regression focuses
on local alignment which advantages personalized fine-tuning. Conversely, when
the fine-tuning method is fixed, employing $\mathcal{L}_{\text{Dr+}}$ for the
GFL model consistently outperforms $\mathcal{L}_{\text{DR}}$ across all
settings. This aligns with previous research [45, 46] suggesting that fine-
tuning from a well-initialized model yields better PFL performance.

###  D.2 Elapsed Time Results

We compare FedDr+ with various GFL algorithms for the elapsed time per
communication round on CIFAR-100 ($s$=10). The results, detailed in Table 9,
show that FedDr+ exhibits a similar but slightly longer elapsed time than the
other algorithms.

Table 9: Elapsed time per round (in seconds) for various GFL algorithms.

| Non-feature normalized algorithms | Feature normalized algorithms
---|---|---
| FedAvg | FedBABU | SCAFFOLD | FedNTD | FedExP | FedETF | SphereFed |  FedDr+ (Ours)
Elapsed time | 21.3 | 20.9 | 22.3 | 22.9 | 20.3 | 22.2 | 22.3 | 24.4