# Federated Learning under Partially Class-Disjoint Data via
Manifold Reshaping

Ziqing Fan, Jiangchao Yao🖂, Ruipeng Zhang {zqfan_knight, sunarker,
<EMAIL_ADDRESS>
Cooperative Medianet Innovation Center, Shanghai Jiao Tong University
Shanghai AI Laboratory

Lingjuan Lyu<EMAIL_ADDRESS>
Sony AI

Ya Zhang, Yanfeng Wang🖂 {ya_zhang<EMAIL_ADDRESS>
Cooperative Medianet Innovation Center, Shanghai Jiao Tong University
Shanghai AI Laboratory

###### Abstract

Statistical heterogeneity severely limits the performance of federated
learning (FL), motivating several explorations e.g., FedProx, MOON and FedDyn,
to alleviate this problem. Despite effectiveness, their considered scenario
generally requires samples from almost all classes during the local training
of each client, although some covariate shifts may exist among clients. In
fact, the natural case of partially class-disjoint data (PCDD), where each
client contributes a few classes (instead of all classes) of samples, is
practical yet underexplored. Specifically, the unique collapse and invasion
characteristics of PCDD can induce the biased optimization direction in local
training, which prevents the efficiency of federated learning. To address this
dilemma, we propose a manifold reshaping approach called FedMR to calibrate
the feature space of local training. Our FedMR adds two interplaying losses to
the vanilla federated learning: one is intra-class loss to decorrelate feature
dimensions for anti-collapse; and the other one is inter-class loss to
guarantee the proper margin among categories in the feature expansion. We
conduct extensive experiments on a range of datasets to demonstrate that our
FedMR achieves much higher accuracy and better communication efficiency.
Source code is available at: https://github.com/MediaBrain-SJTU/FedMR.

##  1 Introduction

Federated learning (McMahan et al. (2017); Li et al. (2020a); Yang et al.
(2019)) has drawn considerable attention due to the increasing requirements on
data protection (Shokri & Shmatikov (2015); Zhu & Han (2020); Hu et al.
(2021); Li et al. (2021c); Lyu et al. (2020)) in real-world applications like
medical image analysis (Guo et al. (2021); Park et al. (2021); Yin et al.
(2022); Dou et al. (2021); Jiang et al. (2022); Zhou et al. (2024; 2023)) and
autonomous driving (Liang et al. (2019); Pokhrel & Choi (2020)). Nevertheless,
the resulting challenge of data heterogeneity severely limits the application
of machine learning algorithms (Zhao et al. (2018)) in federated learning.
This motivations a plenty of explorations to address the statistical
heterogeneity issue and improve the efficiency (Kairouz et al. (2021); Wang et
al. (2020); Li et al. (2022)).

Existing approaches to address the statistical heterogeneity can be roughly
summarized into two categories. One line of research is to constrain the
parameter update in local clients or in the central server. For example,
FedProx (Li et al. (2020b)), FedDyn (Acar et al. (2020)) and FedDC (Gao et al.
(2022)) explore how to reduce the variance or calibrate the optimization by
adding the proximal regularization on parameters in FedAvg (McMahan et al.
(2017)). The other line of research focuses on constraining the representation
from the model to implicitly affect the update. FedProc (Mu et al. (2021)) and
FedProto (Tan et al. (2022)) introduce prototype learning to help local
training, and MOON (Li et al. (2021b)) utilizes contrastive learning to
minimize the distance between representations learned by local model and
global model, and maximize the distance between representations learned by
local model and previous local model. However, all these methods validate
their efficiency mostly under the support from all classes of samples in each
client while lacking a well justification on a natural scenario, namely
_partially class-disjoint data_ w.r.t. classes.

(a) Typical FL.

(b) FL under PCDD.

(c) Global / Collapsed / Reshaped (by FedMR) feature space.

Figure 1: Federated learning under partially class-disjoint data (PCDD).

As illustrated in Figure 1(a), in typical federated learning, each client
usually contains all classes of samples but under different covariate shifts,
and all clients work together to train a global model. However, in the case of
PCDD (Figure 1(b)), there are only a small subset of categories in each client
and all clients together provide information of all classes. Such a situation
is very common in real-world applications. For example, there are shared and
distinct Thyroid diseases in different hospitals due to regional diversity
(Gaitan et al. (1991)). Hospitals from different regions can construct a
federation to learn a comprehensive model for the diagnostic of Thyroid
diseases but suffer from the PCDD challenge. We conduct a toy study on a
simulated dataset (see details in the Appendix B), and visualize the feature
space under centralized training (the left panel of Figure 1(c)) and local
training under PCDD (the middle panel of Figure 1(c)) by projecting on the
unit sphere. As shown in Figure 1(c), PCDD induces a dimensional collapse onto
a narrow area due to the lack of support from all classes, and causes a space
invasion to the missing classes. Previous approaches such as FedProc and MOON
may implicitly constrain the space invasion by utilizing class prototypes or
global features generated from the global model, and methods like FedProx,
FedDyn, and FedNova could also help a bit from the view of optimization.
However, these methods are not oriented towards the PCDD issue, and they are
inefficient to avoid the collapse and invasion characteristics of PCDD and
achieve sub-optimal performance from the both view of experimental performance
shown in Table 1 and the feature variance of all methods shown in Table 7.

To address this dilemma, we propose a manifold-reshaping approach called FedMR
to properly prevent the degeneration caused by the locally class missing.
FedMR introduces two interplaying losses: one is intra-class loss to
decorrelate feature space for anti-collapse; and another one is the inter-
class loss to guarantee the proper margin among categories by means of global
class prototypes. The right panel of Figure 1(c) provides a rough
visualization of FedMR. Theoretically, we analyze the benefit from the
interaction of the intra-class loss and the inter-class loss under PCDD, and
empirically, we verify the effectiveness of FedMR compared with the current
state-of-the-art methods. Our contributions can be summarized as follows:

  * •

We are among the first attempts to study dimensional collapse and space
invasion challenges caused by PCDD in Generic FL that degenerates embedding
space and thus limits the model performance.

  * •

We introduce a approach termed as FedMR, which decorrelates the feature space
to avoid dimensional collapse and constructs a proper inter-class margin to
prevent space invasion. Our theoretical analysis confirms the rationality of
the designed losses and their benefits to address the dilemma of PCDD.

  * •

We conduct a range of experiments on multiple benchmark datasets under PCDD
and a real-world disease dataset to demonstrate the advantages of FedMR over
the state-of-the-art methods. We also develop several variants of FedMR to
consider the communication cost and privacy concerns.

##  2 Related Works

###  2.1 Federated Learning

There are extensive works to address the _statistical heterogeneity_ in
federated learning, which induces the bias of local training due to the
covariate shifts among clients (Zhao et al. (2018); Li et al. (2022); Zhang et
al. (2023a)). A line of research handles this problem by adding constraints
like normalization or regularization on model weights in the local training or
in the server aggregation. FedProx (Li et al. (2020b)) utilizes a proximal
term to limit the local updates so as to reduce the bias, and FedNova (Wang et
al. (2020)) introduces the normalization on total gradients to eliminate the
objective inconsistency. FedDyn (Acar et al. (2020)) makes the global model
and local models approximately aligned in the limit by proposing a dynamic
regularizer for each client at each round. FedDC (Gao et al. (2022)) reduces
the inconsistent optimization on the client-side by local drift decoupling and
correction. Another line of research focuses on constraining representations
from local models and the global model. MOON (Li et al. (2021b)) corrects
gradients in local training by contrasting the representation from local model
and that from global model. FedProc (Mu et al. (2021)) utilizes prototypes as
global information to help correct local representations. Our FedMR is also
conducted on representations of samples and classes but follows a totally
different problem and spirit.

###  2.2 Representation Learning

The collapse problem is also an inevitable concern in the area of
representation learning. In their research lines, there are several attempts
to prevent the potential collapse issues. For example, in self-supervised
learning, Barlow Twins, VICReg and Shuffled-DBN (Bardes et al. (2022); Zbontar
et al. (2021); Hua et al. (2021)) manipulate the rank of the (co-)variance
matrix to prevent the potential collapse. In contrastive learning, DirectCLR
(Jing et al. (2021)) directly optimizes the representation space without an
explicit trainable classifier to promote a larger feature diversity. In
incremental learning, CwD (Shi et al. (2022)) applies a similar technique to
prevent dimensional collapse in the initial phase. In our PCDD case, it is
more challenging, since we should not only avoid collapse but also avoid the
space invasion in the feature space.

###  2.3 Federated Prototype Learning

In many vision tasks (Snell et al. (2017); Yang et al. (2018); Deng et al.
(2021)), prototypes are the mean values of representations of a class and
contain information like feature structures and relationships of different
classes. Since prototypes are population-level statistics of features instead
of raw features, which are relatively safe to share, prototype learning thus
has been applied in federated learning. In Generic FL, FedProc (Mu et al.
(2021)) utilizes the class prototypes as global knowledge to help correct
local training. In Personalized FL, FedProto (Tan et al. (2022)) shares
prototypes instead of local gradients to reduce communication costs. We also
draw spirits from prototype learning to handle PCDD. Besides, to make fair
comparison to methods without prototypes and better reduce extreme privacy
concerns, we conduct a range of auxiliary experiments in Section 4.4.

##  3 The Proposed Method

###  3.1 Preliminary

##### PCDD Definition.

There are many nonnegligible real-world PCDD scenarios. ISIC2019 dataset
(Codella et al. (2018); Tschandl et al. (2018); Combalia et al. (2019)), a
region-driven subset of types of Thyroid diseases in the hospital systems, is
utilized in our experiments. In landmark detection (Weyand et al. (2020)) for
thousands of categories with data locally preserved, most contributors only
have a subset of categories of landmark photos where they live or traveled
before, which is also a scenario for the federated PCDD problem. To make it
clear, we first define some notations of the partially class-disjoint data
situation in federated learning. Let $C$ denote the collection of full classes
and P denote the set of all local clients. Considering the real-world
constraints like privacy or environmental limitation, each local client may
only own the samples of partial classes. Thus, for the $k$-th client $P_{k}$,
its corresponding local dataset $D_{k}$ can be expressed
as$D_{k}=\\{(x_{k,i},y_{k,i})|y_{k,i}=c\in
C_{k}\\},\text{~{}where~{}}C_{k}\subsetneq C.$ The number of samples of the
class $c$ ($c\in C_{k}$) in $P_{k}$ is $N^{c}_{k}$. We denote a local model
$f(\cdot;{\bf w}_{k})$ on all clients as two parts: a backbone network
$f_{1}(\cdot;{\bf w}_{k,1})$ and a linear classifier $f_{2}(\cdot;{\bf
w}_{k,2})$. The loss of the $k$-th client can be formulated as

$\small\ell^{\text{cls}}_{k}(D_{k};w_{k})=\frac{1}{N_{k}}\sum_{i=1}^{N_{k}}\ell(y_{k,i},f_{2}(z_{k,i};w_{k,2}))\big{|}_{z_{k,i}=f_{1}(x_{k,i};w_{k,1})},$

where $z_{k,i}$ is the feature representation of the input $x_{k,i}$ and
$\ell(\cdot,\cdot)$ is the loss measure. Under PCDD, we can empricially find
the dimensional collapse and the space invasion problems about representation.

##### FedAvg.

The vanilla federated learning via FedAvg consists of four steps (McMahan et
al. (2017)): 1) In round $t$, the server distributes the global model ${\bf
w}^{t}$ to clients that participate in the training; 2) Each local client
receives the model and continues to train the model, e.g., the $k$-th client
conducts the following,

$\small{\bf w}_{k}^{t}\leftarrow{\bf
w}_{k}^{t}-\eta\nabla\ell_{k}(b_{k}^{t};{\bf w}_{k}^{t}),$ (1)

where $\eta$ is the learning rate, and $b^{t}_{k}$ is a mini-batch of training
data sampled from the local dataset $\mathcal{D}_{k}$. After $E$ epochs, we
acquire a new local model ${\bf w}^{t}_{k}$; 3) The updated models are then
collected to the server as $\\{{\bf w}^{t}_{1},{\bf w}^{t}_{2},\dots,{\bf
w}^{t}_{K}\\}$; 4) The server performs the following aggregation to acquire a
new global model ${\bf w}^{t+1}$,

$\small{\bf w}^{t+1}\leftarrow\sum_{k=1}^{K}p_{k}{\bf w}^{t}_{k},$ (2)

where $p_{k}$ is the proportion of sample number of the $k$-th client to the
sample number of all the participants, i.e.,
$p_{k}=N_{k}/\sum_{k^{\prime}=1}^{K}N_{k^{\prime}}$. When the maximal round
$T$ reaches, we will have the final optimized model ${\bf w}^{T}$.

Figure 2: An illustration about the shift of the optimization direction under
PCDD. Here, we assume our client contains two classes $c_{1}$ and $c_{2}$ with
one missing class $c_{3}$. ${\bf w}^{*}$ is the optimal classifier direction
for $c_{2}$ (perpendicular to the plane $\alpha$) when all classes exist, and
${\bf w}_{1}^{*}$ is the learned classifier direction when $c_{3}$ is missing,
which can be inferred by the decision plane $\beta$ between $c_{1}$ and
$c_{2}$. As can be seen, PCDD leads to the angle shift $\theta$ in the
optimization.

###  3.2 Motivation

In Figure 2, we illustrate a low-dimensional example to characterize the
directional shift of the _local_ training under PCDD on the client side. In
the following, we use the parameter aggregation of a linear classification to
study the directional shift of _global_ model in the server, which further
clarifies the adverse effect of PCDD. Similar to Figure 2, let $c_{1}$,
$c_{2}$ and $c_{3}$ denote three classes of samples on a circular face
respectively centered at (1, 0), (0, $\sqrt{3}$) and (0, -$\sqrt{3}$) with
radius r=$\frac{1}{2}$ under a uniform distribution. Then, if there are
samples of all classes in each local client, we will get the optimal weight
for all categories as follows:

$\small w^{*}=\left[\begin{array}[]{cc}1&0\\\
-\frac{\sqrt{3}}{2}&\frac{1}{2}\\\ -\frac{\sqrt{3}}{2}&-\frac{1}{2}\\\
\end{array}\right].$

Note that, we omit the bias term in linear classification for simplicity.
Conversely, among total three participants, if each participant only has the
samples of two classes, e.g., ($c_{1}$, $c_{2}$), ($c_{1}$, $c_{3}$) and
($c_{2}$, $c_{3}$) respectively, then their learned weights can be inferred as
follows:

$\displaystyle\small\begin{split}w_{1}^{*},&~{}w_{2}^{*},~{}w_{3}^{*}=\left[\begin{array}[]{cc}\frac{1}{2}&-\frac{\sqrt{3}}{2}\\\
-\frac{1}{2}&\frac{\sqrt{3}}{2}\\\
0&0\end{array}\right],~{}\left[\begin{array}[]{cc}\frac{1}{2}&\frac{\sqrt{3}}{2}\\\
0&0\\\
-\frac{1}{2}&-\frac{\sqrt{3}}{2}\end{array}\right],~{}\left[\begin{array}[]{cc}0&0\\\
0&1\\\ 0&-1\end{array}\right].\end{split}$

After the server aggregation, we have the estimated weight

$\small\hat{w}^{*}=\left[\begin{array}[]{cc}\frac{1}{3}&0\\\
-\frac{1}{6}&\frac{\sqrt{3}+2}{6}\\\ -\frac{1}{6}&-\frac{\sqrt{3}+2}{6}\\\
\end{array}\right].$

Then, we can find that except the difference on amplitude between $w^{*}$ and
$\hat{w}^{*}$, a more important issue is the optimization direction for
$c_{2}$ (or $c_{3}$) shifts about 45∘ by computing the angle between vector
$(-\frac{\sqrt{3}}{2},\frac{1}{2})$ and vector
$(-\frac{1}{6},\frac{\sqrt{3}+2}{6})$ (or between vector
$(-\frac{\sqrt{3}}{2},-\frac{1}{2})$ and vector
$(-\frac{1}{6},-\frac{\sqrt{3}+2}{6})$). Actually, the angle shift can be
enlarged in some real-world applications, when the hard negative classes are
missing. However, if we can have the statistical centroid of the locally
missing class in the local client, namely $c_{3}$ in the case of Figure 2, it
is easy to find that the inferred optimal $\hat{w}^{*}$ is same to $w^{*}$ as
the decision plane can be normally characterized with the support of the
single point $c_{3}$. This inspires us to design the subsequent method111Note
that, we would like to point out that in the real-world case, we cannot
acquire such statistical centroid in advance but have to resort to the
training process along with the special technique design..

###  3.3 Manifold Reshaping

As the aforementioned analysis, PCDD in federated learning leads to the
directional shift of optimization both in the local models and in the global
model. An empirical explanation is that the feature representation of the
specific class that should support classification is totally missing, inducing
the feature representation of other observed classes arbitrarily distributes
as a greedy collapsed manifold, as shown in Figure 1(c). To address this
problem, we explore a manifold-reshaping method from both the intra-class
perspective and the inter-class perspective. In the following, we will present
two interplaying losses and our framework.

####  3.3.1 Intra-Class Loss

The general way to prevent the representation from collapsing into a low-
dimensional manifold, is to decorrelate dimensions for different patterns and
expand the intrinsic dimensionality of each category. Such a goal can be
implemented by manipulating the rank of the covariance matrix regarding
representation. Specifically, for each client, we can first compute the class-
level normalization for the representation $z^{c}_{k,i}\in\mathbb{R}^{d}$ as
$\hat{z}_{k,i}^{c}=\frac{z_{k,i}^{c}-\mu_{k}^{c}}{\sigma_{k}^{c}}$, where
$\mu^{c}_{k}$ and $\sigma_{k}^{c}$ are the mean and standard deviation of
features belonging to class $c$ and calculated as:
$\mu_{k}^{c}=\frac{1}{n_{k,c}}\sum_{i=1}^{n_{k,c}}z_{k,i}^{c}$ and
$\sigma_{k}^{c}=\sqrt{\frac{1}{n_{k,c}}\sum_{i=1}^{n_{k,c}}(z_{k,i}^{c}-\mu_{k}^{c})^{2}}$.
Then, we compute an intra-class covariance matrix based on the above
normalization for each observed class in the $k$-th client:

$\small
M_{k}^{c}=\frac{1}{N_{k}^{c}-1}\sum_{i=1}^{N_{k}^{c}}\left(\hat{z}_{k,i}^{c}\left(\hat{z}_{k,i}^{c}\right)^{\top}\right).$

Since each eigenvalue of $M_{k}^{c}\in\mathbb{R}^{d\times d}$ characterizes
the importance of a feature dimension within class, we can make them
distributed uniformly to prevent the dimensional collapse of each observed
class. However, considering the learnable pursuit of machine learning
algorithms, we actually cannot directly optimize eigenvalues to reach this
goal. Fortunately, it is possible to use an equivalent objective as an
alternative, which is clarified by the following lemma.

######  Lemma 1.

Assuming a covariance matrix $M\in\mathbf{R}^{d\times d}$ computed from the
feature of each sample with the standard normalization, and its eigenvalues
$\\{\lambda_{1},\lambda_{2},...,\lambda_{d}\\}$, we will have the following
equality that satisfied

$\small\sum_{i=1}^{d}{(\lambda_{i}-\frac{1}{d}\sum_{j=1}^{d}{\lambda_{j}})^{2}}=||M||^{2}_{F}-d.$

The complete proof is summarized in the Appendix A.2. From Lemma 1, we can see
that pursuing the uniformity of the eigenvalues for the covariance matrix can
transform into minimizing the Frobenius norm of the covariance matrix.
Therefore, our intra-class loss to prevent the undesired dimensional collapse
for observed classes is formulated as

$\small\ell_{k}^{\text{intra}}=\frac{1}{|C_{k}|}\sum_{c\in
C_{k}}{||M_{k}^{c}||^{2}_{F}}.$ (3)

####  3.3.2 Inter-Class Loss

Although the intra-class loss helps decorrelate the feature dimensions to
prevent collapse, the resulting space invasion for the missing classes can be
concomitantly exacerbated. Thus, it is important to guarantee the proper space
of the missing classes in the expansion as encouraged by equation 3. To
address this problem, we maintain a series of global class prototypes and
transmit them to local clients as support of the missing classes in the
feature space. Concretely, we first compute the class prototypes in the $k$-th
client as the average of feature representations (Snell et al. (2017); Yang et
al. (2018); Deng et al. (2021)):

$\small\left\\{g_{k}^{c}\bigg{|}g_{k}^{c}\leftarrow\frac{1}{N_{k}^{c}}\sum_{i=1}^{N_{k}^{c}}z^{c}_{k,i}\text{~{}for~{}}c\in
C_{k}\right\\}.$

Then, all client-level prototypes are submitted to the server along with local
models in federated learning. In the central server, the global prototypes for
all classes are updated as

$\small\left\\{g_{c}^{t}\bigg{|}g_{c}^{t}\leftarrow\sum_{k=1}^{K}p_{k}^{c}g_{k}^{c}\text{~{}for~{}}c\in
C\right\\},$

where $g_{c}^{t}$ is the global prototype of the $c$-th class in round $t$ and
$p_{k}^{c}=N_{k}^{c}/\sum_{k=1}^{K}N_{k}^{c}$. In the next round, the central
server distributes the global prototypes to all clients as the references to
avoid the space invasion. Formally, we construct the following margin loss by
contrasting the distances from prototypes to the representation of the sample.

$\small\ell_{k}^{\text{inter}}=\frac{1}{|C_{k}|(|C_{k}|-1)}\sum_{c_{i}\in
C_{k}}\sum_{c_{j}\in C_{k}\backslash c_{i}}{D_{c_{i},c_{j}}},\\\ $ (4)

where $D_{c_{i},c_{j}}$ is defined as:

$D_{c_{i},c_{j}}=\frac{1}{N_{k}^{c}}\sum_{n=1}^{N_{k}^{c}}\max\\{||z_{k,n}^{c_{i}}-g_{c_{i}}^{t}||-||z_{k,n}^{c_{i}}-g_{c_{j}}^{t})||,0\\}.$

In the following, we use a theorem to show how inter-class loss jointly with
intra-class loss makes the representation of the missing classes approach to
the optimal.

Figure 3: The framework of FedMR. On the client side, except the vanilla
training with the classification loss, the manifold-reshaping parts, i.e., the
intra-class loss and the inter-class loss, respectively help conduct the
feature decorrelation to avoid the dimensional collapse, and leverage the
global prototypes to construct the proper margin among classes to prevent the
space invasion. On the server side, except the model aggregation, the global
class prototypes are also the reference for missing classes participating in
the local training. Algorithm 1 FedMR

Input: a set of $K$ clients that participate in each round, the initial model
weights ${\bf w}^{0}$, the maximal round $T$, the learning rate $\eta$, the
local training epochs $E$.

for $t=0,1,\dots,T-1$ do

randomly sample $K$ clients

updates global model weights and global class prototypes (${\bf
w}^{t}\leftarrow\sum_{k=1}^{K}{p^{t}_{k}{\bf
w}^{t-1}_{k}}$$\forall{c},g^{t}_{c}\leftarrow\sum_{k=1}^{K}{p^{c}_{k}g^{c}_{k}}$).

distribute ${\bf w}^{t}$ and $G_{c}\\{g_{1}^{t},g_{2}^{t},...,g_{C}^{t}\\}$ to
the $K$ clients.

do in parallel for  $\forall{k\in K}$ clients

${\bf w}_{k}^{t}\leftarrow{\bf w}^{t}$.

for $\tau=0,1,...,E-1$ do

sample a mini-batch from local dataset and perform updates(${\bf
w}_{k}^{t}\leftarrow{\bf
w}_{k}^{t}-\eta\nabla\mathcal{L}_{k}(b_{k}^{t},G_{c};{\bf w}_{k}^{t})$ ).

end for

update local class prototypes
($\forall{c},~{}g_{k}^{c}\leftarrow\sum_{i=1}^{n_{k}^{c}}\frac{1}{n_{k}^{c}}z^{c}_{i}$),
and submit ${\bf w}^{t}_{k}$ and {$g_{k}^{1},g_{k}^{2},...,g_{k}^{C}$} to
server.

end in parallel

end for

######  Theorem 1.

Let the global optimal representation for class $c$ be
$g^{*}_{c}=[a_{c,1}^{*},...,a_{c,d}^{*}]$, and $z^{c,t}_{k}$ be the
representation of sample $x$ in the class $c$ of the $k$-th client. Assuming
that $\forall i$, both $|a^{*}_{c,i}|$ and $z^{c,t}_{k,i}$ are upper bounded
by $G$, and all dimensions are disentangled, in round t, the i-th dimension of
local representation $z^{c,t}_{k}$ satisfies

$\small|z_{k,i}^{c,t}-a_{c,i}^{*}|\leq
2(1-\hat{p}_{k}^{c}\Gamma)G+\delta\Gamma,$

where $\hat{p}_{k}^{c}$ is the accumulation regarding the $i$-th dimension of
the class-$c$ prototype, $\Gamma=\frac{1-(p_{k}^{c})^{t}}{1-p_{k}^{c}}$,
$(p_{k}^{c})^{t}$ refers to the $p_{k}^{c}$ raised to the power of t, and
$\delta$ is the maximum margin of the inter-loss term.

Note that, Theorem 1 shows five critical points: 1) The proof of the theorem
requires each dimension of the representation to be irrelevant to each other,
which is achieved by the intra-class loss. Although disentanglement of
dimensions might not be totally achieved in practical, empirically, we find
that the intra-class loss converges and maintains a relatively low value
easily, meaning that the model achieves good decorrelation. In the Appendix,
we show the training curve of intra-class loss during the federated training.
2) In the theorem, $\delta$ is a trade-off of training stability and
theoretical results determined by the margin. The larger the margin, the
larger the $\delta$, however the more stable the local training is. This is
because the global prototypes are not very accurate in the early stage of
local training and directly minimizing the distance of samples to their global
class prototypes can bring side effect to the feature diversity. When the
margin is removed
($D_{c_{i},c_{j}}=\frac{1}{N_{k}^{c}}\sum_{n=1}^{N_{k}^{c}}||z_{k,n}^{c_{i}}-g_{c_{i}}^{t}||$),
$\delta$ will be zero. 3) Without considering $\delta$, as $t$ increases,
$\Gamma$ is smaller and representation $z_{k}^{c}$ is closer to global optimal
prototype $g_{c}^{*}$, showing the promise of our method. 4) When t is large
enough, we can get an upper bound
$2\frac{1-p_{k}^{c}-\hat{p}_{k}^{c}}{1-p_{k}^{c}}G$, meaning more clients with
the specific dimensional information participating in the training, the
tighter the upper bound is. When all other clients can provide the support
($\hat{p}_{k}^{c}=1-p_{k}^{c}$), the error will be 0. 5) While
$\hat{p}_{k}^{c}$ denotes the proportions of a subset of clients that can
provide the support information for this dimension, the theoretical result
depends on the class distribution and overlap across the clients. The complete
proof is summarized in the Appendix A.1.

####  3.3.3 The Total Framework

After introducing the intra-class loss and the inter-class loss, we give the
total framework of FedMR. On the client side, local models are trained on
their partially class-disjoint datasets. Through manifold-reshaping loss,
dimensions of the representation are decorrelated and local class subspace is
corrected to prevent the space invasion, and gradually approach the global
space partition. The total local objective including the vanilla
classification loss can be written as

$\displaystyle\small\begin{split}\mathcal{L}_{k}=&\ell_{k}^{\text{cls}}+\underbrace{\left(\mu_{1}\ell_{k}^{\text{intra}}+\mu_{2}\ell_{k}^{\text{inter}}\right)}_{\text{manifold
reshaping}},\\\ \end{split}$ (5)

where $\mu_{1}$ and $\mu_{2}$ are the balancing hyperparameters and will be
discussed in the experimental part. In Figure 3, we illustrate the
corresponding structure of FedMR and formulate the training procedure in
Algorithm 1. In terms of the privacy concerns about the prototype transmission
and the communication cost, we will give the comprehensive analysis on FedMR
about these factors.

Table 1: Performance of FedMR and a range of state-of-the-art approaches on
four datasets under PCDD partitions. Datasets are divided into $\varrho$
clients and each client has $\varsigma$ classes (denoted as _$P\varrho
C\varsigma$_). We compute average accuracy of all partitions, highlight
results of FedMR, underline results of best baseline, and show the improvement
of FedMR to FedAvg (subscript of FedMR) and to the best baseline ($\Delta$).

Datasets | Split | FedAvg | FedProx | FedProc | FedNova | MOON | FedDyn | FedDC | FedMR | $\Delta$
---|---|---|---|---|---|---|---|---|---|---
FMNIST | P5C2 | 67.29 | 69.60 | 66.28 | 66.87 | 66.82 | $\underline{71.01}$ | 68.50 | $\textbf{75.51}_{8.22\%\uparrow}$ | +4.50
P10C2 | 67.33 | 67.76 | $\underline{69.08}$ | 48.44 | 67.93 | 67.16 | 67.36 | $\textbf{74.97}_{7.64\%\uparrow}$ | +5.89
P10C3 | 81.67 | 80.93 | 82.06 | 83.20 | $\underline{83.42}$ | 83.00 | 83.24 | $\textbf{83.55}_{1.88\%\uparrow}$ | +0.13
P10C5 | 88.53 | 89.22 | $\underline{89.23}$ | 88.86 | 88.98 | 88.38 | 89.22 | $\textbf{90.04}_{1.51\%\uparrow}$ | +0.81
IID | 91.93 | 91.95 | 92.06 | 91.84 | 92.12 | 91.76 | $\underline{92.15}$ | $\textbf{92.19}_{0.26\%\uparrow}$ | +0.04
avg | 79.35 | 79.89 | 79.74 | 75.84 | 79.87 | $\underline{80.26}$ | 80.09 | $\textbf{83.45}_{4.10\%\uparrow}$ | +3.19
SVHN | P5C2 | 81.85 | $\underline{81.83}$ | 79.54 | 81.11 | 81.60 | 79.89 | 81.63 | $\textbf{83.10}_{1.25\%\uparrow}$ | +1.27
P10C2 | 78.92 | 79.60 | 78.75 | 66.86 | $\underline{79.83}$ | 76.24 | 78.96 | $\textbf{82.47}_{3.55\%\uparrow}$ | +2.64
P10C3 | 87.70 | 87.40 | $\underline{88.13}$ | 87.50 | 87.83 | 87.27 | 88.05 | $\textbf{89.13}_{1.43\%\uparrow}$ | +1.00
P10C5 | 91.20 | 91.24 | 91.63 | $\underline{92.09}$ | 91.16 | 90.17 | 91.64 | $\textbf{92.18}_{0.98\%\uparrow}$ | +0.09
IID | 92.74 | 92.89 | $\underline{93.57}$ | 92.62 | 93.12 | 92.26 | 92.90 | $\textbf{93.04}_{0.30\%\uparrow}$ | -0.53
avg | 86.48 | 86.59 | 86.32 | 83.87 | $\underline{86.71}$ | 85.17 | 86.64 | $\textbf{87.98}_{1.50\%\uparrow}$ | +1.27
CIFAR10 | P5C2 | 67.68 | 68.18 | 69.27 | 67.57 | 66.86 | $\underline{69.64}$ | 69.18 | $\textbf{74.19}_{6.51\%\uparrow}$ | +4.55
P10C2 | 67.27 | $\underline{71.09}$ | 67.02 | 57.79 | 67.61 | 67.74 | 67.64 | $\textbf{73.32}_{2.23\%\uparrow}$ | +2.23
P10C3 | 77.82 | 77.89 | 77.87 | 77.22 | $\underline{78.42}$ | 77.99 | 77.94 | $\textbf{82.75}_{4.93\%\uparrow}$ | +4.33
P10C5 | 88.22 | 88.34 | 88.19 | 88.20 | 88.00 | $\underline{88.35}$ | 88.14 | $\textbf{89.06}_{0.84\%\uparrow}$ | +0.71
IID | 91.88 | 92.14 | 92.62 | 92.37 | 92.56 | 92.29 | $\underline{92.85}$ | $\textbf{93.06}_{1.18\%\uparrow}$ | +0.21
avg | 78.58 | $\underline{79.53}$ | 78.99 | 76.63 | 78.69 | 79.20 | 79.15 | $\textbf{82.48}_{3.90\%\uparrow}$ | +2.95
CIFAR100 | P10C10 | 54.31 | 54.79 | 54.69 | 54.45 | 54.98 | $\underline{55.94}$ | 54.73 | $\textbf{57.27}_{2.96\%\uparrow}$ | +1.33
P10C20 | 64.81 | 65.37 | 64.98 | $\underline{65.79}$ | 65.75 | 65.02 | 65.21 | $\textbf{65.81}_{1.00\%\uparrow}$ | +0.02
P10C30 | 69.35 | 69.75 | 69.64 | 69.55 | 69.51 | $\underline{69.84}$ | 69.38 | $\textbf{70.24}_{0.89\%\uparrow}$ | +0.40
P10C50 | 71.28 | 71.35 | $\underline{72.13}$ | 71.25 | 71.54 | 71.25 | 72.11 | $\textbf{72.17}_{0.89\%\uparrow}$ | +0.04
IID | 72.28 | 72.55 | $\underline{73.07}$ | 72.66 | 73.01 | 73.04 | 72.77 | $\textbf{72.79}_{0.51\%\uparrow}$ | -0.28
avg | 66.41 | 66.76 | 66.90 | 66.74 | 66.96 | $\underline{67.22}$ | 66.84 | $\textbf{67.66}_{1.25\%\uparrow}$ | +0.44

##  4 Experiment

###  4.1 Experimental Setup

##### Datasets.

We adopt four popular benchmark datasets SVHN (Netzer et al. (2011)), FMNIST
(Xiao et al. (2017)), CIFAR10 and CIFAR100 (LeCun et al. (1998)) in federated
learning and a real-world PCDD medical dataset ISIC2019 (Codella et al.
(2018); Tschandl et al. (2018); Combalia et al. (2019)) to conduct
experiments. Regarding the data setup, although Dirichelet Distribution is
popular to split data in FL, it usually generates diverse imbalance data
coupled with occasionally PCDD. In order to better study pure PCDD, for the
former four benchmarks, we split each dataset into $\varrho$ clients, each
with $\varsigma$ categories, abbreviated as _$P\varrho C\varsigma$_. For
example, P10C10 in CIFAR100 means that we split CIFAR100 into 10 clients, each
with 10 classes. Please refer to the detailed explanations and strategies in
the Appendix. ISIC2019 is a real-world federated application under the PCDD
situation and the data distribution among clients is shown in the Appendix
C.1.3. We follow the settings in Flamby benchmark (Terrail et al. (2022)).

##### Implementation.

We compare FedMR with FedAvg (McMahan et al. (2017)) and multiple state-of-
the-arts including FedProx (Li et al. (2020b)), FedProc (Karimireddy et al.
(2020)), FedNova (Li et al. (2021b)), MOON (Wang et al. (2020)), FedDyn (Acar
et al. (2020)) and FedDC (Gao et al. (2022)). To make a fair and comprehensive
comparison, we utilize the same model for all approaches and three model
structures for different datasets: ResNet18 (He et al. (2016)) (follow (Li et
al. (2021b; 2022))) for SVHN, FMNIST and CIFAR10, wide ResNet (Zagoruyko &
Komodakis (2016)) for CIFAR100 and EfficientNet (Tan & Le (2019)) for
ISIC2019. The optimizer is SGD with a learning rate 0.01, the weight decay
$10^{-5}$ and momentum 0.9. The batch size is set to 128 and the local updates
are set to 10 epochs for all approaches. The detailed information of the model
and training parameters are given in the Appendix.

###  4.2 Performance under PCDD

In this part, we compare FedMR with FedAvg and other methods on FMNIST, SVHN,
CIFAR10 and CIFAR100 datasets under partially class-disjoint situation. Note
that, in FedProx, MOON, FedDyn, FedProc, FedDC and our method, there are
parameters that need to set. We use grid search to choose the best parameters
for each method. See more concrete settings in the Appendix C.2 and Section
4.1.

As shown in Table 1, with the decreasing class number in local clients, the
performance of FedAvg and all other methods greatly drops. However, comparing
with all approaches, our method FedMR achieves far better improvement to
FedAvg, especially $7.64\%$ improvement vs. $1.75\%$ of FedProc for FEMNIST
(P10C2) and $6.51\%$ improvement vs. $1.96\%$ of FedDyn for CIFAR10 (P5C2).
Besides, FedMR also performs better under less PCDD, and on average of all
partitions listed in the table, our method outperforms the best baseline by
$3.19\%$ on FMNIST and $2.95\%$ on CIFAR10.

Table 2: Global test accuracy of methods on CIFAR10 and CIFAR100 under larger
scale of clients (P10, P50 and P100) and ISIC2019. _$P\varrho C\varsigma$_
denotes that the dataset is divided into $\varrho$ clients and each client has
$\varsigma$ classes of samples. We highlight results of FedMR, underline
results of best baseline, and show the improvement of FedMR to FedAvg (bottom
right corner of results) and to the best baseline ($\Delta$).

Datasets | Split | FedAvg | FedProx | FedProc | FedNova | MOON | FedDyn | FedDC | FedMR | $\Delta$
---|---|---|---|---|---|---|---|---|---|---
CIFAR10 | P10C3 | 77.82 | 77.89 | 77.87 | 77.22 | 78.42 | 77.99 | 77.94 | $\textbf{82.75}_{4.93\%\uparrow}$ | +4.33
P50C3 | 75.46 | 77.46 | 76.27 | 74.12 | 76.51 | 75.52 | 76.03 | $\textbf{79.58}_{4.12\%\uparrow}$ | +2.12
P100C3 | 71.65 | 72.37 | 72.44 | 70.46 | 72.46 | 71.85 | 73.95 | $\textbf{76.93}_{5.28\%\uparrow}$ | +2.98
CIFAR100 | P10C10 | 54.31 | 54.79 | 54.69 | 54.45 | 54.98 | 55.94 | 54.73 | $\textbf{57.27}_{2.96\%\uparrow}$ | +1.33
P50C10 | 49.84 | 51.17 | 51.94 | 50.22 | 52.19 | 50.53 | 51.17 | $\textbf{53.36}_{3.52\%\uparrow}$ | +1.17
P100C10 | 47.90 | 48.26 | 49.01 | 48.07 | 48.94 | 49.24 | 48.76 | $\textbf{49.60}_{1.70\%\uparrow}$ | +0.36
ISIC2019 | Real | 73.14 | 75.41 | 75.26 | 73.62 | 75.46 | 75.07 | 75.25 | $\textbf{76.55}_{3.41\%\uparrow}$ | +1.09

###  4.3 Scalability and Robustness

In the previous section, we validate the FedMR under 5 or 10 local clients
under partially class-disjoint data situations. In order to make a
comprehensive comparison, we increase the client numbers of CIFAR10 and
CIFAR100 and in each round, only 10 of clients participate in the federated
procedures. Besides, we also add one real federated application ISIC2019 with
multiple statistical heterogeneity problems including PCDD. The exact
parameters and communication rounds of all methods can be found in the
Appendix C.2.

Table 3: The communication cost of approaches with prototypes (i.e., FedProc
and FedMR) and without prototypes (i.e., FedAvg, FedProx, FedNova, MOON,
FedDyn and FedDC).

Method Type | FMNIST | SVHN | CIFAR10 | CIFAR100 | ISIC2019
---|---|---|---|---|---
w/o Prototypes | $11.182M$ | $11.184M$ | $11.184M$ | $36.565M$ | $4.875M$
w/ Prototypes | $11.187M$ | $11.189M$ | $11.189M$ | $36.629M$ | $4.880M$
Additional cost | $\color[rgb]{0,.5,.5}{0.044\%\uparrow}$ | $\color[rgb]{0,.5,.5}{0.044\%\uparrow}$ | $\color[rgb]{0,.5,.5}{0.044\%\uparrow}$ | $\color[rgb]{0,.5,.5}{0.175\%\uparrow}$ | $\color[rgb]{0,.5,.5}{0.102\%\uparrow}$

Table 4: Number of communication rounds and the speedup of communication when
reaching the best accuracy of FedAvg in federated learning on CIFAR100 under
three partition strategies.

Method | P10C10 | P50C10 | P100C10
---|---|---|---
(CIFAR100) | Commu. | Speedup | Commu. | Speedup | Commu. | Speedup
FedAvg | $400$ | $1\times$ | $400$ | $1\times$ | $400$ | $1\times$
FedProx | $352$ | $1.14\times$ | $370$ | $1.08\times$ | $384$ | $1.04\times$
FedProc | $357$ | $1.12\times$ | $381$ | $1.05\times$ | $389$ | $1.03\times$
FedNova | $290$ | $1.38\times$ | $385$ | $1.04\times$ | $\underline{382}$ | $\underline{1.05}\times$
MOON | $332$ | $1.20\times$ | $373$ | $1.07\times$ | $395$ | $1.01\times$
FedDyn | $\underline{178}$ | $\underline{2.25}\times$ | $\underline{366}$ | $\underline{1.09}\times$ | $384$ | $1.04\times$
FedDC | $275$ | $1.45\times$ | $393$ | $1.02\times$ | $387$ | $1.03\times$
FedMR | 149 | $\textbf{2.68}\times$ | 293 | $\textbf{1.37}\times$ | 368 | $\textbf{1.09}\times$

In Table 2, we divide CIFAR10 and CIFAR100 into 10, 50 and 100 clients and
keep the PCDD degree in the same level. As can be seen, with the number of
client increasing, the performance of all methods drops greatly. No matter in
the situations of fewer or more clients, our method achieves better
performance and outperforms best baseline by $2.98\%$ in CIFAR-10 and $0.95\%$
in CIFAR-100 on average. Besides, in Table 2, we verify FedMR with other
methods under a real federated applications: ISIC2019. As shown in the last
line of Table 2, our method achieves the best improvement of $3.41\%$ relative
to FedAvg and of $1.09\%$ relative to best baseline MOON, which means our
method is robust in complicated situations more than PCDD222Note that, in
Appendix, we provide two real-world datasets to further demonstrate the
efficiency of FedMR compared with baselines..

(a) Local burden.

(b) Computation times with global accuracy.

Figure 4: The average memory consuming, computation time of local training and
performance on all datasets of all baselines, FedMR and its light versions
(Lite 10 and Lite 50) for accelerating.

###  4.4 Further Analysis

Except performance under PCDD, we here discuss the communication cost between
clients and server, local burden of clients, and privacy, and conduct the
ablation study.

##### Communication Concern.

In terms of communication cost, our method needs to share the additional
prototypes. To show how much extra communication cost will be incurred, in
Table 3, we show the number of transmission parameters in each round to
compare the communication cost of methods with prototypes (FedProc and FedMR)
and without prototypes (FedAvg, FedProx, FedNova, MOON, FedDyn and FedDC).
From the results, the additional communication cost in single round is
negligible. Except for sharing class prototypes averaged from class
representations, we also tried to use 1-hot vectors to save computation and
communication, but the performance is unsatisfactory and even worse than the
FedAvg. This is because such discriminative structure might not fit the
optimal statistics as each class has different hardness to learn, and it is
not clear that arriving at such an optimization stationary from the given
initialization can be a good choice. In Table 4, we also provide the
communication rounds on CIFAR100 under three different partitions, where all
methods require to reach the best accuracy of FedAvg within 400 rounds. From
the table 4, we can see that FedMR uses less communication rounds (best
speedup) to reach the given accuracy, indicating that FedMR is a
communication-efficient approach.

##### Local Burden Concern.

In real-world federated applications, local clients might be mobile phones or
other small devices. Thus, the burden of local training can be the bottleneck
for clients. In Figure 4(a), we compute the number of parameters that needs to
be saved in local clients and the average local computation time per round. As
can be seen, FedDC, FedDyn and MOON require triple or even quadruple storing
memory than FedAvg, while FedProc and FedMR only need little space to
additionally store prototypes. In terms of local computation time, FedMR
requires more time to carefully reshape the feature space. To handle some
computing-restricted clients, we provide light versions of FedMR, namely Lite
10 and Lite 50, where local clients randomly select only 10 or 50 samples to
compute inter-class loss. From Figure 4(a), the training time of Lite 10 and
Lite 50 decreases sharply, while their performance is still competitive and
better than other baselines, as shown in Figure 4(b). Please refer to Appendix
C.6 for more details.

Table 5: Performance of FedMR on CIFAR10 and CIFAR100 when only 50%, 80% and 100% of clients are allowed to submit their class prototypes under different partitions. Datasets | Split | FedAvg | 50% | 80% | 100%
---|---|---|---|---|---
CIFAR10 | P10C3 | 77.82 | 80.42 | 80.27 | 82.75
P50C3 | 75.46 | 79.14 | 79.43 | 79.58
P100C3 | 71.65 | 76.73 | 76.87 | 76.93
CIFAR100 | P10C10 | 54.31 | 56.41 | 56.73 | 57.27
P50C10 | 50.22 | 51.06 | 52.57 | 53.36
P100C10 | 47.90 | 48.80 | 48.88 | 49.60
Table 6: The ablation study of FedMR. We illustrate average accuracy of FedMR on the four datasets without the inter-class loss or the intra-class loss or both. Results of all partitions are shown in Appendix C.7. Inter | Intra | FMNIST | SVHN | CIFAR10 | CIFAR100
---|---|---|---|---|---
- | - | 79.35 | 86.48 | 78.58 | 66.41
✓ | - | 80.18 | 87.46 | 78.86 | 66.62
- | ✓ | 81.03 | 87.15 | 80.89 | 67.11
✓ | ✓ | 83.45 | 87.98 | 82.48 | 67.66

##### Privacy Concern.

Although prototypes are population-level statistics of features instead of raw
features, which are relatively safe (Mu et al. (2021); Tan et al. (2022)), it
might be still hard for some clients with extreme privacy limitations. To deal
with this case, one possible compromise is allowing partial local clients not
to submit prototypes. In Table 5, we verify this idea for FedMR on CIFAR10 and
CIFAR100, where at the beginning, we only randomly pre-select 50% and 80%
clients to require prototypes. From Table 5, even under the prototypes of 50%
clients, FedMR still performs better than FedAvg, showing the elastic
potential of FedMR in the privacy-restricted scenarios. In the extreme case
where all prototypes of clients are disallowed to be submitted, we can remove
the inter-class that depends on prototypes from FedMR and use the vanilla
federated learning with the intra-class loss. As shown in Table 6, it can
achieve a promising improvement than that without the intra-class loss. Note
that, we cannot counteract the privacy concerns of federated learning itself,
and leave this in the future explorations.

Table 7: Variance of top 50 egienvalues of covariance matrices of all classes
within a mini-batch on CIFAR10.

Method | P5C2 | P10C2 | P10C3 | P10C5 | IID
---|---|---|---|---|---
FedAvg | 1201 | 1274 | 1055 | 814 | 640
FedProx | 931 | 977 | 874 | 727 | 582
FedProc | 1044 | 1074 | 955 | 739 | 522
FedNova | 1234 | 1277 | 977 | 755 | 572
MOON | 749 | 866 | 774 | 744 | 579
FedDyn | 854 | 906 | 844 | 759 | 599
FedDC | 1077 | 1104 | 784 | 766 | 572
FedMR (intra) | 478 | 437 | 372 | 407 | 538
FedMR (inter+intra) | 570 | 538 | 566 | 579 | 635

##### Decorrelating Analysis

Here, we empirically verify the effectiveness of FedGELA compared with all
methods on constraining feature spaces from the view of the variance of
eigenvalues of the covariance matrix M. In Table 7, we show the variance of
top-50 eigenvalues (sort from the largest to the smallest) of the covariance
matrix M within a mini-batch (batchsize is 128) after training 100 rounds on
CIFAR10, calculated
as:$\frac{1}{128}\sum_{i=1}^{50}{(\lambda_{i}-\frac{1}{50}\sum_{j=1}^{50}{\lambda_{j}})^{2}}.$
As can be seen, when the PCDD problem eased, the variance of FedAvg gradually
drops, which means in more uniform data distribution, the variance should be
relatively small. The slight and sharp decreasing variance of prior federated
methods and our FedMR indicate that they indeed help but still suffer from the
dimensional collapse problem caused by PCDD while FedMR successfully
decorrelates the dimensions. However the variance under the intra-class loss
is too small and far away from the values of FedAvg in the IID setting,
meaning it may enlarge the risk of the space invasion. In order to prevent
space invasion, our inter-class loss provide a margin for the feature space
expansion. In the table, we could see that the variance of FedMR (under the
intra-class loss and the inter-class loss) is a little larger compared to
FedMR (the intra-class loss) and approaches to the FedAvg under the IID
setting.

##### Ablation Study.

FedMR introduces two interplaying losses, the intra-class loss and the inter-
class loss, to vanilla FL. To verify the individual efficiency, we conduct an
ablation experiment in Table 6. As can be seen, the intra-class loss generally
plays a more important role in the performance improvement of FedMR, but their
combination complements each other and thus performs best than any of the
single loss, confirming our intuition to prevent the collapse and space
invasion under PCDD jointly. Besides, as FedMR and FedProc need local clients
additionally share class prototypes which might not fair for other baselines,
in Table 8, we properly configure all baselines with prototypes on FMNIST to
show the superiority of FedMR.

Table 8: The average accuracy of all methods adopted in Table 1 with or
without the aid of prototypes on FMNIST. Since FedProc is the prototype
version of FedAvg, here we don’t show them.

Method | FedProx | FedNova | MOON | FedDyn | FedDC | FedMR
---|---|---|---|---|---|---
w/o Prototypes | $79.89$ | $75.84$ | $79.87$ | $\underline{80.26}$ | $80.09$ | 81.03
w Prototypes | $80.29$ | $78.76$ | $80.12$ | $\underline{81.21}$ | $80.57$ | 83.45

##  5 Conclusion

In this work, we study the problem of _partially class-disjoint data_ (PCDD)
in federated learning, which is practical and challenging due to the unique
collapse and invasion problems, and propose a novel approach called FedMR to
address the dilemma of PCDD. Theoretically, we show how the proposed two
interplaying losses in FedMR to prevent the collapse and guarantee the proper
margin among classes. Extensive experiments show that FedMR achieves
significant improvements on FedAvg under the PCDD situations and outperforms a
range of state-of-the-art methods.

## Acknowledgements

The work is supported by the National Key R$\&$D Program of China (No.
2022ZD0160702), STCSM (No. 22511106101, No. 22511105700, No. 21DZ1100100), 111
plan (No. BP0719010) and National Natural Science Foundation of China (No.
62306178). Ziqing Fan and Ruipeng Zhang were partially supported by Wu Wen Jun
Honorary Doctoral Scholarship, AI Institute, Shanghai Jiao Tong University.

## References

  * Acar et al. (2020) 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_ , 2020.
  * Bardes et al. (2022) Adrien Bardes, Jean Ponce, and Yann Lecun.  Vicreg: Variance-invariance-covariance regularization for self-supervised learning.  In _ICLR 2022-10th International Conference on Learning Representations_ , 2022.
  * Borgli et al. (2020) Hanna Borgli, Vajira Thambawita, Pia H Smedsrud, Steven Hicks, Debesh Jha, Sigrun L Eskeland, Kristin Ranheim Randel, Konstantin Pogorelov, Mathias Lux, Duc Tien Dang Nguyen, et al.  Hyperkvasir, a comprehensive multi-class image and video dataset for gastrointestinal endoscopy.  _Scientific data_ , 7(1):283, 2020.
  * Codella et al. (2018) Noel CF Codella, David Gutman, M Emre Celebi, Brian Helba, Michael A Marchetti, Stephen W Dusza, Aadi Kalloo, Konstantinos Liopyris, Nabin Mishra, Harald Kittler, et al.  Skin lesion analysis toward melanoma detection: A challenge at the 2017 international symposium on biomedical imaging (isbi), hosted by the international skin imaging collaboration (isic).  In _2018 IEEE 15th international symposium on biomedical imaging (ISBI 2018)_ , pp. 168–172. IEEE, 2018.
  * Combalia et al. (2019) Marc Combalia, Noel CF Codella, Veronica Rotemberg, Brian Helba, Veronica Vilaplana, Ofer Reiter, Cristina Carrera, Alicia Barreiro, Allan C Halpern, Susana Puig, et al.  Bcn20000: Dermoscopic lesions in the wild.  _arXiv preprint arXiv:1908.02288_ , 2019.
  * Deng et al. (2021) Jiankang Deng, Jia Guo, Jing Yang, Alexandros Lattas, and Stefanos Zafeiriou.  Variational prototype learning for deep face recognition.  In _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition_ , pp. 11906–11915, 2021.
  * Dou et al. (2021) Qi Dou, Tiffany Y So, Meirui Jiang, Quande Liu, Varut Vardhanabhuti, Georgios Kaissis, Zeju Li, Weixin Si, Heather HC Lee, Kevin Yu, et al.  Federated deep learning for detecting covid-19 lung abnormalities in ct: a privacy-preserving multinational validation study.  _NPJ digital medicine_ , 4(1):1–11, 2021.
  * Fan et al. (2022) Ziqing Fan, Yanfeng Wang, Jiangchao Yao, Lingjuan Lyu, Ya Zhang, and Qi Tian.  Fedskip: Combatting statistical heterogeneity with federated skip aggregation.  In _2022 IEEE International Conference on Data Mining (ICDM)_ , pp. 131–140. IEEE, 2022.
  * Gaitan et al. (1991) Eduardo Gaitan, Norman C Nelson, and Galen V Poole.  Endemic goiter and endemic thyroid disorders.  _World journal of surgery_ , 15(2):205–215, 1991.
  * Gao et al. (2022) Liang Gao, Huazhu Fu, Li Li, Yingwen Chen, Ming Xu, and Cheng-Zhong Xu.  Feddc: Federated learning with non-iid data via local drift decoupling and correction.  In _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition_ , pp. 10112–10121, 2022.
  * Guo et al. (2021) Pengfei Guo, Puyang Wang, Jinyuan Zhou, Shanshan Jiang, and Vishal M Patel.  Multi-institutional collaborations for improving deep learning-based magnetic resonance image reconstruction using federated learning.  In _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition_ , pp. 2423–2432, 2021.
  * He et al. (2016) Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun.  Deep residual learning for image recognition.  In _Proceedings of the IEEE conference on computer vision and pattern recognition_ , pp. 770–778, 2016.
  * Hu et al. (2021) Rui Hu, Yanmin Gong, and Yuanxiong Guo.  Federated learning with sparsification-amplified privacy and adaptive optimization.  In _IJCAI_ , 2021.
  * Hua et al. (2021) Tianyu Hua, Wenxiao Wang, Zihui Xue, Sucheng Ren, Yue Wang, and Hang Zhao.  On feature decorrelation in self-supervised learning.  In _Proceedings of the IEEE/CVF International Conference on Computer Vision_ , pp. 9598–9608, 2021.
  * Jiang et al. (2022) Meirui Jiang, Zirui Wang, and Qi Dou.  Harmofl: Harmonizing local and global drifts in federated learning on heterogeneous medical images.  In _Proceedings of the AAAI Conference on Artificial Intelligence_ , volume 36, pp. 1087–1095, 2022.
  * Jing et al. (2021) Li Jing, Pascal Vincent, Yann LeCun, and Yuandong Tian.  Understanding dimensional collapse in contrastive self-supervised learning.  In _International Conference on Learning Representations_ , 2021.
  * Kairouz et al. (2021) Peter Kairouz, H Brendan McMahan, Brendan Avent, Aurélien Bellet, Mehdi Bennis, Arjun Nitin Bhagoji, Kallista Bonawitz, Zachary Charles, Graham Cormode, Rachel Cummings, et al.  Advances and open problems in federated learning.  _Foundations and Trends® in Machine Learning_ , 14(1–2):1–210, 2021.
  * Karimireddy et al. (2020) 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_ , pp. 5132–5143. PMLR, 2020.
  * LeCun et al. (1998) Yann LeCun, Léon Bottou, Yoshua Bengio, and Patrick Haffner.  Gradient-based learning applied to document recognition.  _Proceedings of the IEEE_ , 86(11):2278–2324, 1998.
  * Li et al. (2021a) Ning Li, Tao Li, Chunyu Hu, Kai Wang, and Hong Kang.  A benchmark of ocular disease intelligent recognition: One shot for multi-disease detection.  In _Benchmarking, Measuring, and Optimizing: Third BenchCouncil International Symposium, Bench 2020, Virtual Event, November 15–16, 2020, Revised Selected Papers 3_ , pp. 177–193. Springer, 2021a.
  * Li et al. (2021b) 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, 2021b.
  * Li et al. (2021c) Qinbin Li, Zeyi Wen, Zhaomin Wu, Sixu Hu, Naibo Wang, Yuan Li, Xu Liu, and Bingsheng He.  A survey on federated learning systems: vision, hype and reality for data privacy and protection.  _IEEE Transactions on Knowledge and Data Engineering_ , 2021c.
  * Li et al. (2022) Qinbin Li, Yiqun Diao, Quan Chen, and Bingsheng He.  Federated learning on non-iid data silos: An experimental study.  In _2022 IEEE 38th International Conference on Data Engineering (ICDE)_ , pp. 965–978. IEEE, 2022.
  * Li et al. (2020a) Tian Li, Anit Kumar Sahu, Ameet Talwalkar, and Virginia Smith.  Federated learning: Challenges, methods, and future directions.  _IEEE Signal Processing Magazine_ , 37(3):50–60, 2020a.
  * Li et al. (2020b) 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, 2020b.
  * Liang et al. (2019) Xinle Liang, Yang Liu, Tianjian Chen, Ming Liu, and Qiang Yang.  Federated transfer reinforcement learning for autonomous driving.  _arXiv preprint arXiv:1910.06001_ , 2019.
  * Lyu et al. (2020) Lingjuan Lyu, Han Yu, and Qiang Yang.  Threats to federated learning: A survey.  _arXiv preprint arXiv:2003.02133_ , 2020.
  * McMahan et al. (2017) 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_ , pp. 1273–1282. PMLR, 2017.
  * Mu et al. (2021) Xutong Mu, Yulong Shen, Ke Cheng, Xueli Geng, Jiaxuan Fu, Tao Zhang, and Zhiwei Zhang.  Fedproc: Prototypical contrastive federated learning on non-iid data.  _arXiv preprint arXiv:2109.12273_ , 2021.
  * Netzer et al. (2011) Yuval Netzer, Tao Wang, Adam Coates, Alessandro Bissacco, Bo Wu, and Andrew Y Ng.  Reading digits in natural images with unsupervised feature learning.  _NIPS Workshop on Deep Learning and Unsupervised Feature Learning_ , 2011.
  * Park et al. (2021) Sangjoon Park, Gwanghyun Kim, Jeongsol Kim, Boah Kim, and Jong Chul Ye.  Federated split task-agnostic vision transformer for covid-19 cxr diagnosis.  _Advances in Neural Information Processing Systems_ , 34, 2021.
  * Pokhrel & Choi (2020) Shiva Raj Pokhrel and Jinho Choi.  Federated learning with blockchain for autonomous vehicles: Analysis and design challenges.  _IEEE Transactions on Communications_ , 68(8):4734–4746, 2020.
  * Shakespeare et al. (1989) William Shakespeare et al.  _William Shakespeare: the complete works_.  Barnes & Noble Publishing, 1989.
  * Shi et al. (2022) Yujun Shi, Kuangqi Zhou, Jian Liang, Zihang Jiang, Jiashi Feng, Philip HS Torr, Song Bai, and Vincent YF Tan.  Mimicking the oracle: An initial phase decorrelation approach for class incremental learning.  In _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition_ , pp. 16722–16731, 2022.
  * Shokri & Shmatikov (2015) Reza Shokri and Vitaly Shmatikov.  Privacy-preserving deep learning.  In _Proceedings of the 22nd ACM SIGSAC conference on computer and communications security_ , pp. 1310–1321, 2015.
  * Snell et al. (2017) Jake Snell, Kevin Swersky, and Richard Zemel.  Prototypical networks for few-shot learning.  _Advances in neural information processing systems_ , 30, 2017.
  * Tan & Le (2019) Mingxing Tan and Quoc Le.  Efficientnet: Rethinking model scaling for convolutional neural networks.  In _International conference on machine learning_ , pp. 6105–6114. PMLR, 2019.
  * Tan et al. (2022) Yue Tan, Guodong Long, Lu Liu, Tianyi Zhou, Qinghua Lu, Jing Jiang, and Chengqi Zhang.  Fedproto: Federated prototype learning across heterogeneous clients.  In _AAAI Conference on Artificial Intelligence_ , volume 1, pp. 3, 2022.
  * Terrail et al. (2022) Jean Ogier du Terrail, Samy-Safwan Ayed, Edwige Cyffers, Felix Grimberg, Chaoyang He, Regis Loeb, Paul Mangold, Tanguy Marchand, Othmane Marfoq, Erum Mushtaq, et al.  Flamby: Datasets and benchmarks for cross-silo federated learning in realistic healthcare settings.  _arXiv preprint arXiv:2210.04620_ , 2022.
  * Tschandl et al. (2018) Philipp Tschandl, Cliff Rosendahl, and Harald Kittler.  The ham10000 dataset, a large collection of multi-source dermatoscopic images of common pigmented skin lesions.  _Scientific data_ , 5(1):1–9, 2018.
  * Wang et al. (2020) Jianyu Wang, Qinghua Liu, Hao Liang, Gauri Joshi, and H Vincent Poor.  Tackling the objective inconsistency problem in heterogeneous federated optimization.  _arXiv preprint arXiv:2007.07481_ , 2020.
  * Weyand et al. (2020) Tobias Weyand, Andre Araujo, Bingyi Cao, and Jack Sim.  Google landmarks dataset v2-a large-scale benchmark for instance-level recognition and retrieval.  In _Proceedings of the IEEE/CVF conference on computer vision and pattern recognition_ , pp. 2575–2584, 2020.
  * Xiao et al. (2017) Han Xiao, Kashif Rasul, and Roland Vollgraf.  Fashion-mnist: a novel image dataset for benchmarking machine learning algorithms.  _arXiv preprint arXiv:1708.07747_ , 2017.
  * Yang et al. (2018) Hong-Ming Yang, Xu-Yao Zhang, Fei Yin, and Cheng-Lin Liu.  Robust classification with convolutional prototype learning.  In _Proceedings of the IEEE conference on computer vision and pattern recognition_ , pp. 3474–3482, 2018.
  * Yang et al. (2019) Qiang Yang, Yang Liu, Tianjian Chen, and Yongxin Tong.  Federated machine learning: Concept and applications.  _ACM Transactions on Intelligent Systems and Technology (TIST)_ , 10(2):1–19, 2019.
  * Yao et al. (2022) Jiangchao Yao, Shengyu Zhang, Yang Yao, Feng Wang, Jianxin Ma, Jianwei Zhang, Yunfei Chu, Luo Ji, Kunyang Jia, Tao Shen, et al.  Edge-cloud polarization and collaboration: A comprehensive survey for ai.  _IEEE Transactions on Knowledge and Data Engineering_ , 35(7):6866–6886, 2022.
  * Ye et al. (2023a) Rui Ye, Zhenyang Ni, Fangzhao Wu, Siheng Chen, and Yanfeng Wang.  Personalized federated learning with inferred collaboration graphs.  2023a.
  * Ye et al. (2023b) Rui Ye, Mingkai Xu, Jianyu Wang, Chenxin Xu, Siheng Chen, and Yanfeng Wang.  Feddisco: Federated learning with discrepancy-aware collaboration.  _arXiv preprint arXiv:2305.19229_ , 2023b.
  * Yin et al. (2022) Youtan Yin, Hongzheng Yang, Quande Liu, Meirui Jiang, Cheng Chen, Qi Dou, and Pheng-Ann Heng.  Efficient federated tumor segmentation via normalized tensor aggregation and client pruning.  In _International MICCAI Brainlesion Workshop_ , pp. 433–443. Springer, 2022.
  * Zagoruyko & Komodakis (2016) Sergey Zagoruyko and Nikos Komodakis.  Wide residual networks.  In _British Machine Vision Conference 2016_. British Machine Vision Association, 2016.
  * Zbontar et al. (2021) Jure Zbontar, Li Jing, Ishan Misra, Yann LeCun, and Stéphane Deny.  Barlow twins: Self-supervised learning via redundancy reduction.  In _International Conference on Machine Learning_ , pp. 12310–12320. PMLR, 2021.
  * Zhang et al. (2023a) Ruipeng Zhang, Ziqing Fan, Qinwei Xu, Jiangchao Yao, Ya Zhang, and Yanfeng Wang.  Grace: A generalized and personalized federated learning method for medical imaging.  In _International Conference on Medical Image Computing and Computer-Assisted Intervention_ , pp. 14–24. Springer, 2023a.
  * Zhang et al. (2023b) Ruipeng Zhang, Qinwei Xu, Jiangchao Yao, Ya Zhang, Qi Tian, and Yanfeng Wang.  Federated domain generalization with generalization adjustment.  In _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition_ , pp. 3954–3963, 2023b.
  * Zhao et al. (2018) 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.
  * Zhou et al. (2021) Yuhang Zhou, Xiaoman Zhang, Shixiang Feng, Ya Zhang, et al.  Uncertainty-aware incremental learning for multi-organ segmentation.  _arXiv preprint arXiv:2103.05227_ , 2021.
  * Zhou et al. (2022) Yuhang Zhou, Shu-Wen Sun, Qiu-Ping Liu, Xun Xu, Ya Zhang, and Yu-Dong Zhang.  Ted: Two-stage expert-guided interpretable diagnosis framework for microvascular invasion in hepatocellular carcinoma.  _Medical Image Analysis_ , 82:102575, 2022.
  * Zhou et al. (2023) Yuhang Zhou, Jiangchao Yao, Feng Hong, Ya Zhang, and Yanfeng Wang.  Balanced destruction-reconstruction dynamics for memory-replay class incremental learning.  _arXiv preprint arXiv:2308.01698_ , 2023.
  * Zhou et al. (2024) Yuhang Zhou, Haolin Li, Siyuan Du, Jiangchao Yao, Ya Zhang, and Yanfeng Wang.  Low-rank knowledge decomposition for medical foundation models.  _arXiv preprint arXiv:2404.17184_ , 2024.
  * Zhu & Han (2020) Ligeng Zhu and Song Han.  Deep leakage from gradients.  In _Federated learning_ , pp. 17–31. Springer, 2020.

##  Appendix A Implementation of Method

###  A.1 Proof of Theorem 1

######  Theorem 2.

Let the global optimal representation for class $c$ denote
$g^{*}_{c}=[a_{c,1}^{*},...,a_{c,d}^{*}]$, and $z^{c,t}_{k}$ denote the
representation of sample $x$ in the class $c$ of the $k$-th client. Assuming
that $\forall i$, both $|a^{*}_{c,i}|$ and $|z^{c,t}_{k,i}|$ are upper bounded
by $G$, and all dimensions are disentangled. Then, in round t, the i-th
dimension of local representation $z^{c,t}_{k}$ satisfies

$\small|z_{k,i}^{c,t}-a_{c,i}^{*}|\leq
2(1-\hat{p}_{k}^{c}\Gamma)G+\delta\Gamma,$

where $\hat{p}_{k}^{c}$ is the accumulation regarding the $i$-th dimension of
the class-$c$ prototype, $\Gamma=\frac{1-(p_{k}^{c})^{t}}{1-p_{k}^{c}}$ and
$\delta$ is the maximum margin induced by the optimization of the inter-loss
term.

###### Proof.

If the global optimum satisfies:

$f_{1}(x_{c};w_{*,1})=g^{*}_{c}=[a_{c,1}^{*},...,a_{c,d}^{*}],$

we could define the local optimum of class c in client k:

$f_{1}(x_{c};w_{k,1}^{*})=z_{k}^{c}=[a_{c,1}^{*},...,a_{c,s}^{*},\sigma_{k,s+1}^{c},...,\sigma_{k,s+r}^{c}],$

where $\\{a_{c,1}^{*},...,a_{c,s}^{*}\\}$ denotes the dimensions relative to
classifying class c with other seen classes in client k. Since the intra-class
loss decorrelates each dimensions of the feature space,
$\\{\sigma_{k,s+1}^{c},...,\sigma_{k,s+r}^{c}\\}$ are free dimensions and
irrelevant to the seen class. By minimizing the inter-class loss

$D_{c_{i},c_{j}}=\frac{1}{N_{k}^{c}}\sum_{n=1}^{N_{k}^{c}}\max\\{||z_{k,n}^{c_{i}}-g_{c_{i}}^{t}||-||z_{k,n}^{c_{i}}-g_{c_{j}}^{t})||,0\\}.\\\
$

Since all dimensions are irrelevant (achieved by intra-class loss), we can
reach the optimum of each dimension of feature representations in the limit,
which has the following property:

$||(z_{k}^{c_{i}})_{m}-g_{c_{i},m}^{t}||\leq\min\\{||(z_{k}^{c_{i}})_{m}-g_{c_{j},m}^{t}||\\},~{}\forall{i}\neq
j,$

where m denote one of the dimensions. In terms of the representation space of
a certain class c, we have the dimensions $\\{a_{c,1}^{*},...,a_{c,s}^{*}\\}$
relevant to the seen classes in the local client reach the optimal. In this
case, we can rewrite the locality of
$\\{\sigma_{k,s+1}^{c},...,\sigma_{k,s+r}^{c}\\}$ by introducing a slack
variable as follows:

$\sigma_{k,s+j}^{c,t+1}-g_{c,s+j}^{t}=\xi^{t},$

where $\xi^{t}$ is the induced slack variable. The s+j -th element of class
prototypes of class c can be written as:

$g_{c,s+j}^{t}=\hat{p}_{k}^{c}a_{c,s+j}^{*}+p_{k}^{c}\sigma_{k,s+j}^{t}+\sum_{k^{\prime}}p_{k^{\prime}}^{c}\sigma_{k^{\prime},s+j}^{t},$

where $\hat{p}_{k}^{c}$ denotes the proportions of a subset of clients that
can provide the support information of these dimensions. Note that, such
support is related to class c to distinguish unseen classes in client k.
$p_{k^{\prime}}$ is the clients can not provide the support information except
for client k ($\hat{p}_{k}^{c}+p_{k^{\prime}}^{c}+p_{k}^{c}=1$). Putting all
things together, we could have the following equation:

$\sigma_{k,s+j}^{c,t+1}=\hat{p}_{k}^{c}a_{c,s+j}^{*}+p_{k}^{c}\sigma_{k,s+j}^{t}+\sum_{k^{\prime}}p_{k^{\prime}}^{c}\sigma_{k^{\prime},s+j}^{c,t}+\xi^{t}.$

Let $\sigma_{k,s+j}^{c,t+1}$ be $r_{t+1}$ for simplicity. We will have

$r_{t+1}=\hat{p}_{k}^{c}a_{k,s+j}^{*}+p_{k}^{c}r_{t}+\sum_{k^{\prime}}p_{k^{\prime}}^{c,t}+\xi^{t}.$

With the recursive iteration to $r_{0}$, we can compute that

$r_{t+1}=\frac{1-(p_{k}^{c})^{t}}{1-p_{k}^{c}}\hat{p}_{k}^{c}a_{k,s+j}^{*}+(p_{k}^{c})^{t}r_{0}+A+B,$

where A is defined as

$A=\sum_{k^{\prime}}p_{k^{\prime}}^{c,t}\sigma_{k^{\prime},s+j}^{c,t}+...+(p_{k}^{c})^{t}\sum_{k^{\prime}}p_{k^{\prime}}^{c,0}\sigma_{k^{\prime},s+j}^{c,0},$

and B is defined as

$B=\xi^{t}+p_{k}^{c}\xi^{t-1}+...+(p_{k}^{c})^{t}\xi^{0},$

If $|\xi^{t}|$ is bounded by $\delta$ and $\sigma_{k^{\prime},s+j}^{c,t}$ is
bounded by $G$ for all t, we have the inequality:

$|A|\leq\Gamma(1-\hat{p}_{k}^{c}-p_{k}^{c})G,$

and

$|B|\leq\Gamma\delta,$

where $\Gamma$ is defined as:

$\Gamma=\frac{1-(p_{k}^{c})^{t}}{1-p_{k}^{c}}.$

Therefore, we have the following bound

$\displaystyle|r_{t+1}-a_{c,s+j}^{*}|=$
$\displaystyle|\Gamma\hat{p}_{k}^{c}a_{k,s+j}^{*}-a_{c,s+j}^{*}+(p_{k}^{c})^{t}r_{0}+A+B|$
$\displaystyle\leq$
$\displaystyle(1-\hat{p}_{k}^{c}\Gamma)G+(p_{k}^{c})^{t}G+|A|+|B|$
$\displaystyle\leq$ $\displaystyle 2(1-\hat{p}_{k}^{c}\Gamma)G+\Gamma\delta.$

In the above first inequality, we use $|a+b+c|\leq|a|+|b|+|c|$. And in the
above second inequality, we combine the similar terms in A and B. As
$z_{k,i}^{c,t}=g_{c,i}^{*}$ for i=1,2,…,s, we universally have

$|z_{k,i}^{c,t}-a_{c,i}^{*}|\leq 2(1-\hat{p}_{k}^{c}\Gamma)G+\delta\Gamma,$

which completes the proof. ∎

###  A.2 Proof of Lemma 1

######  Lemma 2.

Assuming a covariance matrix $M\in\mathbf{R}^{d\times d}$ computed from the
feature of each sample with the standard normalization, and its eigenvalues
$\\{\lambda_{1},\lambda_{2},...,\lambda_{d}\\}$, we will have the following
equality that satisfied

$\small\sum_{i=1}^{d}{(\lambda_{i}-\frac{1}{d}\sum_{j=1}^{d}{\lambda_{j}})^{2}}=||M||^{2}_{F}-d.$

###### Proof.

On the right-hand side, we have

$\displaystyle||M||_{F}^{2}-d=$ $\displaystyle Tr((M)^{T}M-d)$
$\displaystyle=$ $\displaystyle Tr(U\sum U^{T}U\sum U^{T})-d$ $\displaystyle=$
$\displaystyle\sum_{i=1}^{d}\lambda_{i}^{2}-d.$

As we have applied the standard normalization on $M$, we have the the
characteristic of eigenvalues $\frac{1}{d}\sum_{j=1}^{d}\lambda_{j}=1$. Then
for the right-hand side, we can have the deduction

$\displaystyle\sum_{i=1}^{d}(\lambda_{i}-\frac{1}{d}\sum_{j=1}^{d}\lambda_{j}))^{2}=$
$\displaystyle\sum_{i=1}^{d}(\lambda_{i}-1)^{2}$ $\displaystyle=$
$\displaystyle\sum_{i=1}^{d}(\lambda_{i}^{2}-2\lambda_{i}+1)$ (6)
$\displaystyle=$ $\displaystyle\sum_{i=1}^{d}\lambda_{i}^{2}-d.$

This constructs the equality of the left-hand side and the right-hand side,
which completes the proof. ∎

##  Appendix B Simulation and Visualization

To show the dimensional collapse and verify our method, we simulate data for
four categories, and samples of each category are generated from a circle with
a different center and a radius of 0.5. We set centers (1,1), (1,-1), (-1,1)
and (-1,-1) for class 1 to 4, respectively. The input is the two-dimensional
position, and we adopt a three-layer MLP model with hidden size 128 and 3. We
visualize the outputs of the second layer (3 dimensions) and project them onto
an unit sphere. In the simulation, we adopt SGD with learning rate 0.1 and
generate 5000 samples for each category, and the number of iterations and
batch size are set to 50 and 128. In the paper, we train the model and show
the visualization of models trained on four classes of samples, two classes of
samples and two classes of samples with FedMR to show the global feature
space, collapsed feature space and reshaped feature space, respectively.

Table 9: Best hyper-parameters tuned from 0.000001 to 1 of FedMR and a range
of state-of-the-art approaches on four datasets under PCDD partitions.
Datasets are divided into $\varrho$ clients and each client has $\varsigma$
classes (denoted as _$P\varrho C\varsigma$_).

Datasets | Split | FedProx | FedProc | MOON | FedDyn | FedDC | FedMR ($\mu_{1}$) | FedMR ($\mu_{2}$)
---|---|---|---|---|---|---|---|---
FMNIST | P5C2 | 0.01 | 0.00001 | 0.001 | 0.0001 | 0.001 | 0.01 | 0.0001
P10C2 | 0.01 | 0.00001 | 0.0001 | 0.0001 | 0.001 | 0.01 | 0.0001
P10C3 | 0.01 | 0.00001 | 0.001 | 0.0001 | 0.0001 | 0.01 | 0.0001
P10C5 | 0.01 | 0.001 | 0.001 | 0.0001 | 0.001 | 0.01 | 0.0001
IID | 0.01 | 0.00001 | 0.01 | 0.0001 | 0.01 | 0.01 | 0.0001
SVHN | P5C2 | 0.001 | 0.0001 | 0.001 | 0.000001 | 0.0001 | 0.001 | 0.0001
P10C2 | 0.001 | 0.0001 | 0.001 | 0.000001 | 0.0001 | 0.001 | 0.0001
P10C3 | 0.001 | 0.0001 | 0.001 | 0.000001 | 0.0001 | 0.001 | 0.0001
P10C5 | 0.001 | 0.0001 | 0.1 | 0.000001 | 0.0001 | 0.001 | 0.0001
IID | 0.001 | 0.0001 | 0.1 | 0.00001 | 0.0001 | 0.001 | 0.0001
CIFAR10 | P5C2 | 0.01 | 0.001 | 0.01 | 0.0001 | 0.0001 | 0.1 | 0.001
P10C2 | 0.01 | 0.0001 | 0.0001 | 0.0001 | 0.0001 | 0.1 | 0.001
P10C3 | 0.01 | 0.001 | 0.0001 | 0.0001 | 0.0001 | 0.1 | 0.0001
P10C5 | 0.01 | 0.001 | 0.01 | 0.0001 | 0.00001 | 0.1 | 0.001
IID | 0.01 | 0.001 | 1 | 0.0001 | 0.001 | 0.1 | 0.0001
CIFAR100 | P10C10 | 0.01 | 0.001 | 0.1 | 0.0001 | 0.0001 | 0.1 | 0.0001
P10C20 | 0.01 | 0.001 | 0.1 | 0.0001 | 0.0001 | 0.1 | 0.0001
P10C30 | 0.01 | 0.001 | 0.1 | 0.0001 | 0.0001 | 0.1 | 0.0001
P10C50 | 0.01 | 0.001 | 0.1 | 0.0001 | 0.0001 | 0.1 | 0.0001
IID | 0.01 | 0.001 | 0.1 | 0.0001 | 0.0001 | 0.1 | 0.00001

Table 10: Best method-specific hyper-parameters (weights of proximal term in
FedProx, contrastive loss in MOON and so forth) tuned from 0.000001 to 1 of
FedProx, MOON, FedMR and so forth on CIFAR10 and CIFAR100 under larger scale
and a real-world application: ISIC2019. Datasets are divided into $\varrho$
clients and each client has $\varsigma$ classes (denoted as _$P\varrho
C\varsigma$_).

Datasets | Split | FedProx | FedProc | MOON | FedDyn | FedDC | FedMR ($\mu_{1}$) | FedMR ($\mu_{2}$)
---|---|---|---|---|---|---|---|---
CIFAR10 | P10C3 | 0.01 | 0.001 | 0.0001 | 0.0001 | 0.0001 | 0.1 | 0.0001
P50C3 | 0.0001 | 0.0001 | 0.1 | 0.0001 | 0.00001 | 0.01 | 0.0001
P100C3 | 0.0001 | 0.0001 | 0.1 | 0.0001 | 0.0001 | 0.01 | 0.0001
CIFAR100 | P10C10 | 0.01 | 0.001 | 0.1 | 0.0001 | 0.0001 | 0.1 | 0.0001
P50C10 | 0.0001 | 0.001 | 0.1 | 0.00001 | 00001 | 0.01 | 0.0001
P100C10 | 0.0001 | 0.0001 | 0.1 | 0.00001 | 0.0001 | 0.01 | 0.0001
ISIC2019 | Real | 0.001 | 0.0001 | 0.1 | 0.0001 | 0.0001 | 0.001 | 0.001

Table 11: Computation times and performance on P5C2 of FMNIST, SVHN and CIFAR10 and P10C10 of CIFAR100. Dataset | FedAvg | n=0 | n=10 | n=50 | n=128
---|---|---|---|---|---
FMNIST | 67.29/15.52s | 71.21/17.78s | 76.51/21.43s | 75.67/31.56s | 75.51/54.07s
SVHN | 81.85/22.15s | 82.74/24.11s | 83.51/28.31s | 83.48/57.25s | 83.10/88.20s
CIFAR10 | 67.68/15.74s | 72.26/17.70s | 73.19/24.56s | 73.56/47.26s | 74.19/60.07s
CIFAR100 | 54.31/35.12s | 55.38/42.14s | 55.57/52.18s | 56.35/105.34s | 57.27/195.12s
Table 12: Performance of FedAvg, inter-class loss, intra-class loss and both inter-class and inter-class losses (FedMR). Datasets | Split | FedAvg | Inter | Intra | FedMR
---|---|---|---|---|---
FMNIST | P5C2 | 0.6729 | 0.6717 | 0.7121 | 0.7497
P10C2 | 0.6733 | 0.6962 | 0.7028 | 0.7552
P10C3 | 0.8167 | 0.8241 | 0.8282 | 0.8324
P10C5 | 0.8953 | 0.8959 | 0.8899 | 0.9004
IID | 0.9193 | 0.9210 | 0.9185 | 0.9215
SVHN | P5C2 | 0.8185 | 0.8257 | 0.8274 | 0.8310
P10C2 | 0.7892 | 0.8138 | 0.8148 | 0.8247
P10C3 | 0.8770 | 0.8886 | 0.8807 | 0.8913
P10C5 | 0.9120 | 0.9164 | 0.9086 | 0.9209
IID | 0.9274 | 0.9283 | 0.9259 | 0.9304
CIFAR10 | P5C2 | 0.6768 | 0.6775 | 0.7226 | 0.7419
P10C2 | 0.6727 | 0.6766 | 0.7267 | 0.7332
P10C3 | 0.7782 | 0.7784 | 0.8051 | 0.8275
P10C5 | 0.8822 | 0.8851 | 0.8766 | 0.8906
IID | 0.9188 | 0.9253 | 0.9135 | 0.9306
CIFAR100 | P10C10 | 0.5431 | 0.5538 | 0.5615 | 0.5727
P10C20 | 0.6481 | 0.6486 | 0.6512 | 0.6581
P10C30 | 0.6951 | 0.6905 | 0.7021 | 0.7024
P10C50 | 0.7128 | 0.7129 | 0.7150 | 0.7217
IID | 0.7228 | 0.7252 | 0.7259 | 0.7279

##  Appendix C Implementation of Experiment

###  C.1 Models and Datasets

####  C.1.1 Models

For FMNIST, SVHN and CIFAR10, we adopt modified version of ResNet18 as
previous studies (Li et al. (2021b); Zhang et al. (2023b); Ye et al. (2023a);
Yao et al. (2022); Zhou et al. (2022; 2021)). For CIFAR100, we adopt wide
ResNet (https://github.com/meliketoy/wide-resnet.pytorch). For ISIC2019
dataset, we adopt EfficientNet b0 (Tan & Le (2019)) as the same of Flamby
benchmark (Terrail et al. (2022)).

####  C.1.2 Partition Strategies

Dirichlet distribution (Dir($\beta$) is not suitable to split data for pure
PCDD problems (some classes are partially missing). As an exemplar simulation
in Figure 6, Dirichlet allocation usually generates diverse imbalance data
coupled with occasionally PCDD. However, we do not focus on the locally-
imbalance of each client but the locally-balance of each client from limited
existing categories. This is the intuition difference from a Dirichlet
allocation. Therefore, taking an example of CIFAR100 (P10C30), to simulate
pure PCDD situation, we first divide all 100 classes to the client in order:
1-30 categories for client 1, 31-60 categories for client 2, 61-90 categories
for client 3, and 91-100 categories for client 4. Then the remain clients that
lacking categories (less than 30 classes) random choose categories from 1 to
100. After slicing the categories, the samples of each class are equally
divided into clients that have such class. Such partition strategy can ensure
the difference of class distributions among clients, all categories are
allocated and the sample numbers are roughly the same.

Figure 5: Heat map of the data distribution on ISIC2019 dataset.

####  C.1.3 ISIC2019 dataset

As shown in the Figure 5, we show the heat map of data distribution of
ISIC2019 dataset (Codella et al. (2018); Tschandl et al. (2018); Combalia et
al. (2019)). As can be seen, there are multiple statistical heterogeneity
problems including partially class-disjoint data (PCDD).

(a) $\beta=1$.

(b) $\beta=0.5$.

(c) $\beta=0.1$.

Figure 6: Heatmaps of data distribution of CIFAR10 generated by Dir($\beta$).

###  C.2 Parameters

As for model-common parameters like optimizer, lr and batchsize are all
aligned. We have verified $lr=0.01$ is stable and almost the best for all
methods in our settings. For method specific parameters like proximal term of
FedProx, reshaping loss of FedMR, contrastive loss of MOON, dynamic regulazer
of FedDyn and so forth are carefully searched and shown in the following.

####  C.2.1 Parameters in 4.2

We use grid search from 0.000001 to 1 (interval 10) to find best method-
specific parameters of all method on different datasets as shown in Table 9.
The total rounds are 100 for FMNIST, SVHN and CIFAR10 and 200 for CIFAR100.

####  C.2.2 Parameters in 4.3

We use grid search from 0.000001 to 1 (interval 10) to find best method-
specific parameters parameters of FedProx, MOON, FedMR and so forth on
different datasets as shown in Table 10. The total rounds are 100 for CIFAR10
(P10C3), 200 for CIFAR10 (P50C3) and 400 for CIFAR10 (P100C3) and CIFAR100.

Table 13: Performance of FedMR on two real-world datasets named HyperKvasir and ODIR under three PCDD situations. Datasets | Split | FedAvg | FedProx | MOON | FedDyn | FedProc | FedMR | $\Delta$
---|---|---|---|---|---|---|---|---
HyperKvasir | P10C2 | 69.67 | 68.43 | 69.33 | 70.02 | 69.54 | 70.79 | +0.75
P10C3 | 76.44 | 76.92 | 77.43 | 77.01 | 76.72 | 78.32 | +0.89
P10C5 | 89.48 | 89.23 | 89.13 | 89.34 | 89.27 | 89.54 | +0.06
ODIR | P3C3 | 59.89 | 60.24 | 59.85 | 60.78 | 60.88 | 61.79 | +0.91
P4C3 | 56.53 | 56.42 | 56.89 | 57.12 | 55.78 | 57.67 | +0.55
P5C5 | 54.51 | 54.07 | 54.11 | 54.17 | 53.73 | 54.62 | +0.21
Table 14: Results of FedMR on CIFAR10 when tuning $\mu_{1}$ of intra-class loss under the same $\mu_{2}$ of inter-class loss. Split | $\mu_{1}=1$ | $\mu_{1}=0.1$ | $\mu_{1}=0.01$ | $\mu_{1}=0.001$
---|---|---|---|---
P5C2 | 73.72 | 74.19 | 73.85 | 73.19
P10C2 | 72.99 | 73.32 | 73.00 | 73.12
P10C3 | 81.65 | 82.75 | 82.05 | 80.64
P10C5 | 88.76 | 89.06 | 88.96 | 88.64
P10C10 | 93.02 | 93.06 | 93.04 | 92.16

###  C.3 More results on real-world datasets

We additionally test our FedMR on two more datasets, named HyperKvasir (Borgli
et al. (2020)) and ODIR (Li et al. (2021a)) under three partitions. As shown
in Table 13, our method achieves the best average improvement of 1.02 and 1.05
relative to FedAvg and of 0.57 and 0.56 relative to best baseline on the
HyperKvasir and ODIR respectively.

###  C.4 More about intra-class loss

Since there are two additional loss in our method, it might raise heavy tuning
problem. As shown in Table 14, we record the results when tuning $\mu_{1}$
under the same $\mu_{2}$ on CIFAR10 and empirically observe that the
performance is good and stable for $\mu_{1}$ from a large range, which means
we only need to carefully tune $\mu_{2}$. What’s more, we also demonstrate the
value of intra-class loss to verify the effect of deccorelation during the
federated training. As shown in the Figure 7, we could see that intra-class
loss converges and maintains a relatively low value easily, supporting our
assumption in the Theorem 1 that the dimensions are decorrelated well by
intra-class loss.

Table 15: Performance on SVHN under partitions generated by dirichlet distribution. Dataset | Method | Full Participation(10 clients) | Partial Participation(50 clients)
---|---|---|---
#Partition | IID | $\beta=0.5$ | $\beta=0.1$ | IID | $\beta=0.5$ | $\beta=0.2$
SVHN | FedAvg | 92.74 | 91.24 | 75.24 | 91.29 | 89.29 | 84.70
Best Baseline | 93.50 | 92.46 | 76.26 | 91.67 | 91.27 | 87.78
FedMR(Ours) | 93.04 | 92.50 | 78.19 | 91.77 | 91.31 | 89.37

###  C.5 Performance on General Setting

To verify the effectiveness, we conduct the experiments on SHAKESPEARE
(Shakespeare et al. (1989)) and SVHN, whose data distributions follow the non-
PCDD setting. Specially, SHAKESPEARE is also commonly used as real-world data
heterogeneity challenge in federated learning. We select 50 clients of the
SHAKESPEARE into federated training and our method outperforms all methods and
achieves improvement of 3.86% to FedAvg and 1.10% to the best baseline. As for
SVHN, we split it by Dirichlet distribution as many previous FL works (Ye et
al. (2023b); Li et al. (2020b); Fan et al. (2022)). According to the results
in the Table 15, we can find FedMR still remains applicable and achieves
comparable performance.

Figure 7: Intra-class loss value of our FedMR during the federated training.

###  C.6 Lite Version

Here we introduce our light version for FedMR. Since our reshaping loss will
introduce additional computation times, we randomly select part of samples in
the mini-batch to compute inter-class loss for a computation friendly version.
In Table 11, we randomly select 0, 10, 50 and 128 samples in the mini-batch on
the four dataset (P10C10 for CIFAR100 and P5C2 for the others) to calculate
the inter-class loss and record the accuracy. As can be seen, the computation
time is reduced significantly and simultaneously maintains the competing
performance.

###  C.7 Ablation

In this part, we provide detailed results of ablation on FMNIST, SVHN, CIFAR10
and CIFAR100. As shown in the Table 12, the intra-class loss generally plays a
more important role in the performance improvement of FedMR on four datasets
under PCDD. Their combination complements each other and thus shows a best
improvement than any of the single loss, confirming our design from the joint
perspective to prevent the collapse under PCDD.