# Bridging Model Heterogeneity in Federated Learning via Uncertainty-based
Asymmetrical Reciprocity Learning

Jiaqi Wang  Chenxu Zhao  Lingjuan Lyu  Quanzeng You  Mengdi Huai  Fenglong Ma

###### Abstract

This paper presents FedType, a simple yet pioneering framework designed to
fill research gaps in heterogeneous model aggregation within federated
learning (FL). FedType introduces small identical proxy models for clients,
serving as agents for information exchange, ensuring model security, and
achieving efficient communication simultaneously. To transfer knowledge
between large private and small proxy models on clients, we propose a novel
uncertainty-based asymmetrical reciprocity learning method, eliminating the
need for any public data. Comprehensive experiments conducted on benchmark
datasets demonstrate the efficacy and generalization ability of FedType across
diverse settings. Our approach redefines federated learning paradigms by
bridging model heterogeneity, eliminating reliance on public data,
prioritizing client privacy, and reducing communication costs.

Federated Learning, Model Heterogeneity

##  1 Introduction

Federated Learning (FL) is designed to enable the collaborative training of a
machine learning model without the need to share clients’ data. Many prevalent
FL models, including FedAvg (McMahan et al., 2017) and FedProx (Li et al.,
2020), mandate that clients employ an identical model structure and target for
training a shared global model. However, clients may possess diverse model
structures, introducing model heterogeneity within the FL framework. The goal
of this challenging task is to learn personalized client models instead of a
powerful global model.

Recently, many studies have emerged to tackle the challenge of model
heterogeneity (Huang et al., 2022; Li & Wang, 2019; Yi et al., 2023; Lin et
al., 2020; Yu et al., 2022; Wang et al., 2023). These efforts can be
categorized into two groups based on the approach employed for information
exchange between clients and the server. The first category focuses on the
transmission of additional side information, such as logits (Huang et al.,
2022), class scores (Li & Wang, 2019), and label-wise representations (Yi et
al., 2023; Tan et al., 2022), which are derived from utilizing public data on
individual clients. Conversely, the second category involves the direct upload
of client models to the server for processes like distillation (Lin et al.,
2020; Yu et al., 2022) or model reassembly (Wang et al., 2023) with the help
of public data. While they successfully achieve the goal of heterogeneous
model aggregation, they encounter significant drawbacks:

Diminishing returns in public data incorporation – As previously mentioned,
prevailing approaches rely on extra public data, either at the client or
server side, for conducting heterogeneous model aggregation. The latest study
(Wang et al., 2023) underscores that the choice of public data plays a pivotal
role in influencing model performance. Furthermore, using extra data amplifies
the learning cost for models. To mitigate these challenges, a pertinent
research question arises: _Can heterogeneous models be successfully aggregated
without depending on any external public data?_

Disclosure risks raised by exchanging sensitive information – While exchanging
side information between clients and the server can alleviate communication
costs, this straightforward approach raises concerns about the potential
disclosure of sensitive client information (Lyu et al., 2022). Furthermore,
the complete upload of model structures and parameters poses substantial
security risks (Tolpegin et al., 2020) and privacy concerns (Bouacida &
Mohapatra, 2021), particularly for business corporations or entities operating
in sensitive domains. In light of these challenges, another pivotal research
question surfaces: _Can heterogeneous model aggregation be achieved while
exchanging only non-sensitive information, mitigating the risks associated
with data privacy and security?_

Figure 1: Overview of the proposed FedType framework. (a) demonstrates the
workflow of the proposed FedType to address the model-heterogeneous issue in
FL, (b) is the local update demonstration for a data sample $\mathbf{x}_{j}$
using the proposed uncertainty-based asymmetrical reciprocity learning, and
(c) is the illustration of backward knowledge distillation with the proposed
uncertainty-based behavior imitation learning.

Necessity of efficient communication – Efficient communication stands out as a
foundational challenge in federated learning. While numerous effective
solutions have been devised for homogeneous federated learning (Diao et al.,
2020; Yao et al., 2021; Cho et al., 2023), a notable research gap persists in
heterogeneous federated learning. While the exchange of side information
effectively reduces communication costs, as discussed previously, these
approaches introduce vulnerabilities to FL systems. Hence, an urgent
requirement arises for an efficient, secure solution that operates without
relying on public data in heterogeneous FL.

To resolve the issues mentioned previously, this paper introduces a novel,
powerful, yet straightforward framework, FedType, to bridge model
heterogeneity gaps in federated learning without relying on any public data
and ensuring efficient communication through uncertainty-based asymmetrical
reciprocity learning, as depicted in Figure 1. FedType achieves heterogeneous
model aggregation, communication efficiency, and system security
simultaneously.

FedType involves introducing small identical proxy models for clients, serving
as agents for information exchange between clients and the server to
facilitate heterogeneous model aggregation. In the local training phase, each
client employs a novel bidirectional knowledge distillation strategy, referred
to as asymmetrical reciprocity learning, to simultaneously update its large
private model and small proxy model. This strategy leverages the distinctive
characteristics of these two models, elaborated in Section 3.1. Additionally,
an _uncertainty-based behavior imitation learning_ method is developed to
enhance the guidance provided by the small proxy model in facilitating the
learning of the large client model, as detailed in Section 3.2. Ultimately,
each local client is mandated to upload only its proxy model to the server.
This facilitates model aggregation using existing approaches employed in
homogeneous FL, as outlined in Section 3.3. We perform comprehensive
experiments on multiple benchmark datasets, evaluating the FedType framework
in both heterogeneous and homogeneous scenarios against state-of-the-art
baselines. The experimental results consistently affirm the effectiveness of
the FedType framework, underscoring its practicality and robust performance in
real-world applications.

Remarkably, the FedType framework111The code is available at
https://github.com/JackqqWang/FedType. presents several notable advantages in
comparison to existing approaches. Firstly, it eliminates the necessity for
relying on any public data across clients and the server, effectively
mitigating the adverse effects associated with the initialization of public
data in model-heterogeneous federated learning. Secondly, FedType adopts an
identical yet compact proxy model to facilitate information exchange between
clients and the server, resulting in a substantial reduction in communication
costs. Thirdly, by utilizing proxy models, the framework ensures the privacy
and security of client models, offering protection against potential harmful
attacks. Fourthly, the proposed framework exhibits versatility and can be
applied to both heterogeneous and homogeneous federated learning settings.
Lastly, this flexible framework is compatible with any existing federated
learning model.

##  2 Related Work

###  2.1 Heterogeneous Federated Learning

Most existing federated learning studies (McMahan et al., 2017; T Dinh et al.,
2020; Zhang et al., 2022; Marfoq et al., 2022; Bao et al., 2023; Dennis et
al., 2021; Zhou et al., 2022) have concentrated on the homogeneous setting,
requiring all clients to adopt an identical model structure. Recent research
has begun exploring the heterogeneous setting, allowing for varying model
structures across clients. Within this domain, the approaches include sub-
model training (Alam et al., 2022), sparse model-adaption (Chen et al., 2023),
and hypernetworks (Shamsian et al., 2021), which impose constraints on the
relationship between the clients’ models and the global model, thus still
restricting the freedom of clients in using their preferred models.

To facilitate more flexible cooperation among heterogeneous models, several
studies have explored alternatives to averaging model parameters, such as
aggregating extra information like logits (Huang et al., 2022), class scores
(Li & Wang, 2019), and label-wise representations (Yi et al., 2023; Tan et
al., 2022). However, these methods potentially raise privacy concerns. In
response, recent research has focused on exchanging model parameters through
strategies like ensemble learning (Lin et al., 2020), mutual learning (Yu et
al., 2022; Shen et al., 2023), or model reassembly (Wang et al., 2023).

These approaches exhibit two prevalent limitations: (1) Their dependence on
public data. In real-world applications, obtaining access to public data may
not always be feasible, and selecting suitable public data without preliminary
knowledge of the client’s data poses a complex challenge. (2) All the
mentioned studies require sharing sensitive information, such as model
structure, model parameters, or data-related insights from the private local
models, with the server. This raises significant privacy concerns.

###  2.2 Conformal Prediction

Conformal prediction has seen significant popularity in recent years,
particularly due to its capacity to generate prediction sets with guaranteed
error rates under minimal assumptions. Initially introduced in (Vovk et al.,
1999) and further elaborated in (Shafer & Vovk, 2008) and (Balasubramanian et
al., 2014), conformal prediction offers a distribution-free uncertainty
quantification technique that has been effectively applied in various
applications (Angelopoulos & Bates, 2021; Bhatt et al., 2021; Fisch et al.,
2021; Sankaranarayanan et al., 2022). Several recent studies focus on adapting
conformal prediction within FL, specifically targeting the challenges
associated with label shift (Plassier et al., 2023) and the quantification of
uncertainty in distributed environments (Lu et al., 2023). To the best of our
knowledge, there is limited work utilizing conformal prediction to enhance the
performance of heterogeneous FL frameworks.

##  3 Methodology

As illustrated in Figure 1(a), FedType comprises two key components: local
update and server update. During the $t$-th communication round, FedType
initiates the training process for a proxy model $\mathbf{p}_{i}^{t}$ and a
client private model $\mathbf{w}_{i}^{t}$ using the data $\mathcal{D}_{i}$
from the $i$-th client. It is essential to note that in model-heterogeneous
federated learning, the structures of private models
$\\{\mathbf{w}_{1}^{t},\cdots,\mathbf{w}_{N}^{t}\\}$ differ from one another,
where $N$ represents the number of clients. Conversely, the structures of
proxy models $\\{\mathbf{p}_{1}^{t},\cdots,\mathbf{p}_{N}^{t}\\}$ are uniform.
Furthermore, it is worth mentioning that the parameter size of proxy model
$\mathbf{p}_{i}^{t}$ is typically much smaller than that of private model
$\mathbf{w}_{i}^{t}$.

Only the learned proxy models
$\\{\mathbf{p}_{1}^{t},\cdots,\mathbf{p}_{N}^{t}\\}$222In the cross-device
setting, we will randomly select $B\ll N$ clients at each communication round.
are transmitted to the server for model aggregation during the server update,
leading to a significant reduction in communication costs. The resulting
aggregated model $\mathbf{p}_{g}^{t}$ is then disseminated to each client in
the subsequent communication round as the initialized proxy model, denoted as
$\mathbf{p}_{i}^{t+1}=\mathbf{p}_{g}^{t}$ ($\forall i\in[1,\cdots,N]$). These
two updates are executed iteratively until FedType converges. The whole
algorithm flow can be found in Appendix A. Next, we provide the details of
each update.

###  3.1 Local Update via Uncertainty-based Asymmetrical Reciprocity Learning

The inherent challenge posed by model heterogeneity renders the direct
aggregation of uploaded client models unfeasible. Despite various proposed
approaches for heterogeneous model aggregation (Lin et al., 2020; Yu et al.,
2022; Huang et al., 2022; Li & Wang, 2019; Yi et al., 2023), as discussed in
Section 1, they still exhibit several limitations. In contrast to existing
work, we present a simple yet novel uncertainty-based asymmetrical reciprocity
learning (UARL) approach to tackle the challenges posed by model
heterogeneity. As shown in Figure 1 (c), UARL is a bidirectional knowledge
distillation (KD)-based model.

The forward KD (FKD) follows the traditional “teacher-student” knowledge
distillation approach. In FKD, a small proxy/student model
$\mathbf{p}_{i}^{t}$ is distilled from the large, private client/teacher model
$\mathbf{w}_{i}^{t}$ using the client data $\mathcal{D}_{i}$ for the $i$-th
client during the $t$-th communication round. The forward loss can be
formulated as follows:

$\overrightarrow{\mathcal{J}}_{i}^{t}=\sum_{j=1}^{|\mathcal{D}_{i}|}\text{KL}(\mathbf{w}_{i}^{t}(\mathbf{x}_{j})||\mathbf{p}_{i}^{t}(\mathbf{x}_{j})),$
(1)

where $|\mathcal{D}_{i}|$ denotes the number of data stored in the $i$-th
client, $\text{KL}(\cdot,\cdot)$ is the Kullback–Leibler divergence, and
$\mathbf{x}_{j}\in\mathcal{D}_{i}$ represents the input data.

The backward KD (BKD) poses a challenge, as the capability of the small proxy
model $\mathbf{p}_{i}^{t}$ is typically weaker than that of the large client
model $\mathbf{w}_{i}^{t}$. Directly applying traditional knowledge
distillation may lead to a degradation in the power of client models. To
overcome this asymmetrical reciprocity issue, we introduce a novel
uncertainty-based behavior imitation learning method to transfer diverse
knowledge from the proxy model $\mathbf{p}_{i}^{t}$ to the client model
$\mathbf{w}_{i}^{t}$. Further details can be found in Section 3.2.

Let $\overleftarrow{\mathcal{J}}_{i}^{t}$ denote the backward KD loss. We then
use the following loss function to train the $i$-th client at the $t$-th
communication round:

$\mathcal{J}_{i}^{t}=\sum_{j=1}^{|\mathcal{D}_{i}|}\text{CE}(\mathbf{w}_{i}^{t}(\mathbf{x}_{j}),\mathbf{y}_{j})+\overrightarrow{\mathcal{J}}_{i}^{t}+\overleftarrow{\mathcal{J}}_{i}^{t},$
(2)

where $\mathbf{y}_{j}$ denotes the vectorized ground truth obtained from the
real label set $\mathcal{Y}_{j}$ of $\mathbf{x}_{j}$. The algorithm flow can
be found in Algorithm 1.

Input: Client training data $\mathcal{D}_{i}$, private model
$\mathbf{w}_{i}^{t-1}$, proxy model initialized by $\mathbf{p}_{g}^{t-1}$,
validation data $\mathcal{D}_{i}^{\prime}$, local training epoch $R$,
hyperparameters

Divide the shuffled $\mathcal{D}_{i}$ into $R$ parts
$\\{\mathcal{D}_{i}^{1},\cdots,\mathcal{D}_{i}^{R}\\}$;

Initialize epoch-level models: $\mathbf{w}_{i}^{0}=\mathbf{w}_{i}^{t-1}$ and
$\mathbf{p}_{i}^{0}=\mathbf{p}_{g}^{t-1}$;

for  _each epoch $r=1,\cdots,R$_ do

2 Train the conformal prediction models $\mathbf{cw}_{i}^{r}$ and
$\mathbf{cp}_{i}^{r}$ using the learned epoch-level models
$\mathbf{w}_{i}^{r-1}$ and $\mathbf{p}_{i}^{r-1}$ with the validation data
$\mathcal{D}_{i}^{\prime}$ according to Eq. (5); Initialize the loss
$\mathcal{J}_{i}^{r}=0$; for  _each sample
$\mathbf{x}_{j}\in\mathcal{D}_{i}^{r}$_ do

3 Calculate the classification loss via
$\text{CE}(\mathbf{w}_{i}^{r-1}(\mathbf{x}_{j}),\mathbf{y}_{j})$; Calculate
the FKD loss $\overrightarrow{\mathcal{J}}_{i}^{r}(\mathbf{x}_{j})$ via Eq.
(1); Calcuate the performance change $\Delta^{r}$ on the proxy model on
$\mathcal{D}_{i}^{\prime}$; Obtain the prediction set $\mathcal{S}_{j}^{r}$
via Eq. (6) using $\mathbf{cp}_{i}^{r}$; Obtain the prediction set
$\mathcal{L}_{j}^{r}$ via Eq. (6) using $\mathbf{cw}_{i}^{r}$; Calculate
$\eta_{j}^{r}$ according to Eq. (8); Calculate the BKD loss
$\overleftarrow{\mathcal{J}}_{i}^{r}(\mathbf{x}_{j})$; $\mathcal{J}_{i}^{r}$
+=
$\text{CE}(\mathbf{w}_{i}^{r-1}(\mathbf{x}_{j}),\mathbf{y}_{j})+\overrightarrow{\mathcal{J}}_{i}^{r}(\mathbf{x}_{j})+\overleftarrow{\mathcal{J}}_{i}^{r}(\mathbf{x}_{j})$;

4 end for

5 Update the models $\mathbf{w}_{i}^{r}$ and $\mathbf{p}_{i}^{r}$ by
optimizing $\mathcal{J}_{i}^{r}$;

6 end for

Return: Trained models $\mathbf{w}_{i}^{t}$ and $\mathbf{p}_{i}^{t}$.

Algorithm 1 Epoch-Level Algorithm Flow of UARL.

###  3.2 Uncertainty-based Behavior Imitation Learning

The initialization of the proxy model $\mathbf{p}_{i}^{t}$ is carried out
using the aggregated global model $\mathbf{p}_{g}^{t-1}$ (refer to Section
3.3). This global model encapsulates diverse knowledge from other clients,
making it crucial to transfer this knowledge to the private client model.
However, due to the inherently weaker capability of the proxy model
$\mathbf{p}_{i}^{t}$, direct application of the traditional “teacher-student”
learning paradigm for knowledge transfer to the large model
$\mathbf{w}_{i}^{t}$ is impractical. Such a direct approach may introduce
additional noise to the large model, potentially impeding the overall training
efficiency of the entire framework.

####  3.2.1 Backward Knowledge Distillation Loss

To tackle this challenge, we introduce an uncertainty-based behavior imitation
learning approach. This method exclusively relies on the use of _partial
logits_ generated by the proxy model $\mathbf{p}_{i}^{t}$ with a high level of
confidence or certainty. The intention is to employ these confident
predictions to guide the learning process of the large model
$\mathbf{w}_{i}^{t}$. In essence, when dealing with a specific data sample
$\mathbf{x}_{j}\in\mathcal{D}_{i}$, if the proxy model $\mathbf{p}_{i}^{t}$
exhibits high confidence in predicting certain classes, transferring this
behavioral information—rather than the complete logits—can still be beneficial
for the training of the large model.

Let $\mathcal{S}_{j}^{t}$ represent the set of class labels predicted with
high confidence by $\mathbf{p}_{i}^{t}$. In our proposed behavior imitation
learning, the objective is to strengthen the probability/logit associated with
the labels in $\mathcal{S}_{j}^{t}$ as predicted by the large model
$\mathbf{w}_{i}^{t}$. Importantly, we refrain from imposing a strict
requirement for the large model $\mathbf{w}_{i}^{t}$ to prioritize ranking the
labels in $\mathcal{S}_{j}^{t}$ at the top positions. To address this, we
introduce a novel ranking-based behavior imitation learning loss, outlined as
follows:

$\begin{split}&\overleftarrow{\mathcal{J}}_{i}^{t}=\sum_{j=1}^{|\mathcal{D}_{i}|}\eta_{j}^{t}\sum_{k\in\mathcal{S}_{j}^{t}}\log(\frac{\exp(\mathbf{w}_{i}^{t}(\mathbf{x}_{j})[k])}{\Phi_{j}^{t}}),\\\
\Phi_{j}^{t}=&\sum_{s\in\mathcal{S}_{j}^{t}}\exp(\mathbf{w}_{i}^{t}(\mathbf{x}_{j})[s])+\sum_{v\in\mathcal{V}_{j}^{t}}\exp(\mathbf{w}_{i}^{t}(\mathbf{x}_{j})[v]),\end{split}$
(3)

where $\eta_{j}^{t}$ is the estimated consensus weight to determine the amount
of the transferred knowledge from the proxy model $\mathbf{p}_{i}^{t}$ to the
large model $\mathbf{w}_{i}^{t}$.
$\mathcal{V}_{j}^{t}=\mathcal{Y}-\mathcal{S}_{j}^{t}$ denotes the class labels
associated with low confidence, and $\mathcal{Y}$ represents the complete set
of class labels in $\mathcal{D}_{i}$. In Eq. (3), $\eta_{j}^{t}$ and
$\mathcal{S}_{j}^{t}$ are unknown variables. To proceed, we introduce a novel
approach to quantify an uncertainty set $\mathcal{S}_{j}^{t}$ for each data
$\mathbf{x}_{j}$ through dynamic conformal prediction, and the uncertainty
sets will be further used to estimate $\eta_{j}^{t}$.

####  3.2.2 Conformal Model Training

Conformal prediction (Angelopoulos & Bates, 2021) stands out as a reliable and
interpretable approach for quantifying uncertainty, providing prediction sets
accompanied by a designated level of confidence or probability. In
mathematical terms, conformal prediction involves an uncertainty set function
$f(\mathbf{p}_{j}^{t},\mathbf{x}_{j})$ that maps $\mathbf{x}_{j}$ to a subset
of $\mathcal{Y}$ (i.e.,
$f(\mathbf{p}_{j}^{t},\mathbf{x}_{j})=\mathcal{S}_{j}^{t}{\subseteq}\mathcal{Y}$),
satisfying the condition:

$P(\mathcal{Y}_{j}\in\mathcal{S}_{j}^{t}))\geq 1-\theta,$ (4)

where $\theta$ represents a predefined confidence level.

To estimate the prediction set $\mathcal{S}_{j}^{t}$, we need to train a
conformal model $\mathbf{cp}_{j}^{t}$ first using the validation dataset
$\mathcal{D}_{i}^{\prime}$ where
$\mathcal{D}_{i}\cap\mathcal{D}_{i}^{\prime}=\emptyset$ as follows:

$\mathbf{cp}_{i}^{t}=\text{Cmodel}(\mathbf{p}_{i}^{t},\mathcal{D}_{i}^{\prime}),$
(5)

where $\text{Cmodel}()$ is constructed using the split conformal prediction
framework, detailed in Appendix B.

Existing conformal prediction approaches are principally developed to quantify
uncertainty in static, well-trained models and are not tailored to address
dynamic scenarios. However, the FedType framework undergoes iterative
training, resulting in dynamic changes to the models $\mathbf{p}_{i}^{t}$ and
$\mathbf{w}_{i}^{t}$ at each communication round and even each epoch $r$.333In
our implementation, the parameters of two models change at each epoch $r$
during each communication round $t$ in the training stage. Here, we omit the
notation of epoch $r$ in the rest of this section for simplicity and
readability, which can be treated as $r=1$ in each communication round $t$.
Consequently, the uncertainty set $\mathcal{S}_{j}^{t}$ is predicted
dynamically. Next, we will use the trained conformal model
$\mathbf{cp}_{i}^{t}$ to generate the prediction set $\mathcal{S}_{j}^{t}$ for
each training data $\mathbf{x}_{j}\in\mathcal{D}_{i}$.

####  3.2.3 Uncertainty Set Prediction

A straightforward approach would involve directly applying existing conformal
prediction methods, such as regularized adaptive prediction sets (RAPS)
(Angelopoulos et al., 2020), to generate the uncertainty set
$\mathcal{S}_{j}^{t}$ for $\mathbf{x}_{j}$ at each communication round $t$.
However, as mentioned earlier, the capability of the proxy model
$\mathbf{p}_{i}^{t}$ is weak, especially at the initial stages of training.
Utilizing prediction sets from such unreliable models may introduce adverse
effects when training Eq. (3).

To address this challenge, we introduce a dynamic adjustment mechanism for the
size of the prediction set $\mathcal{S}_{j}^{t}$ based on the observed changes
in model performance on the training data $\mathcal{D}_{i}$. In essence, a
decrease in a model’s performance may be indicative of low-quality prediction
sets. Therefore, to enhance the informativeness of the prediction sets, it is
necessary to reduce their size by refining the conformal prediction algorithm.
Based on this intuition, we propose a dynamic conformal prediction for
federated learning training based on RAPS using the following uncertainty
prediction set generation:

$\begin{split}\mathcal{S}_{j}^{t}&:=\\{y:\pi_{j}(y)\cdot u+\rho_{j}(y)\\\
&+g(\Delta^{t},\lambda)\cdot(o_{j}(y)-\kappa_{reg})^{+}\leq\tau\\}.\end{split}$
(6)

$\pi_{j}(y)$ represents the probability assigned to the label $y$ as predicted
by the conformal model $\mathbf{cp}_{i}^{t}$ for the $j$-th data point
$\mathbf{x}_{j}$. The parameter $u$ is a predefined randomized factor
determining the value jump for each new label $y$. Furthermore,
$\rho_{j}(y)=\sum_{y^{\prime}=1}^{|\mathcal{Y}|}\pi_{j}(y^{\prime})\mathbbm{1}_{\\{\pi_{j}(y^{\prime})>\pi_{j}(y)\\}}$
denotes the total probability mass associated with the set of labels that are
more likely than the label $y$.

Besides, $g(\Delta^{t},\lambda)$ is a piecewise calibration function444Here,
we define a simple linear function to characterize the dynamic adjustment of
prediction sets. We can also use other monotonically decreasing functions as
alternative solutions., which is defined as follows:

$g(\Delta^{t},\lambda)=\begin{cases}\lambda\cdot\Delta^{t}-\Delta^{t}+\lambda,&\text{if
$\Delta^{t}<0$,}\\\ \lambda,&\text{otherwise}.\end{cases}$ (7)

$\Delta^{t}=\mathcal{A}(\mathbf{p}_{i}^{t},\mathcal{D}_{i}^{\prime})-\mathcal{A}(\mathbf{p}_{i}^{t-1},\mathcal{D}_{i}^{\prime})$,
where $\mathcal{A}(\mathbf{p}_{i}^{t},\mathcal{D}_{i}^{\prime})$ denotes the
validation accuracy on $\mathcal{D}_{i}^{\prime}$ using the trained model
$\mathbf{p}_{i}^{t}$. In the event of a performance drop, i.e.,
$\Delta^{t}<0$, we actively reduce the size of the prediction set by adjusting
the value of $g(\Delta^{t},\lambda)$. On the other hand, if $\Delta^{t}\geq
0$, we maintain $g(\Delta^{t},\lambda)=\lambda$, where $\lambda\geq 0$ serves
as a predefined regularization hyperparameter.

In addition, $(z)^{+}$ denotes the positive part of $z$, and $\kappa_{reg}$ is
another regularization hyperparameter. The variable
$o_{j}(y)=|y^{\prime}\in\mathcal{Y}:\pi_{j}(y^{\prime})\geq\pi_{j}(y)|$
represents the ranking of $y$ among the labels in $\mathcal{Y}$ based on the
prediction probability $\pi_{j}$. Finally, $\tau$ denotes the cumulative sum
of the sorted, penalized classifier scores. The details of the designed
dynamic conformal prediction can be found in Appendix B.

####  3.2.4 Consensus Weight Estimation

In Eq. (3), $\eta_{j}^{t}$ is a key factor to control the behavior of
knowledge transfer from the proxy model $\mathbf{p}_{i}^{t}$ to the large
model $\mathbf{w}_{i}^{t}$. Intuitively, if the proxy model
$\mathbf{p}_{i}^{t}$ is significantly confident on certain predicted labels on
data $\mathbf{x}_{j}$ i.e., the small size of $\mathcal{S}_{j}^{t}$, then it
is necessary to transfer such high-quality knowledge to the large model
$\mathbf{w}_{i}^{t}$ as much as possible, as the guidance of model training.
On the other hand, we need to reduce the amount of transferred knowledge with
low quality, i.e., the uncertainty set $\mathcal{S}_{j}^{t}$ with large size.

The quality of knowledge is a relative variable, which can be determined by
the size of uncertainty sets predicted by both models $\mathbf{p}_{i}^{t}$ and
$\mathbf{w}_{i}^{t}$ on $\mathbf{x}_{j}$ using Eq. (6) with the conformal
models $\mathbf{cp}_{i}^{t}$ and $\mathbf{cw}_{i}^{t}$ learned by Eq. (5),
denoted as $\mathcal{S}_{j}^{t}$ and $\mathcal{L}_{j}^{t}$, respectively.
Mathematically, we define the consensus weight $\eta_{j}^{t}$ as follows:

$\eta_{j}^{t}=\begin{cases}|\mathcal{S}_{j}^{t}\cap\mathcal{L}_{j}^{t}|/|\mathcal{S}_{j}^{t}\cup\mathcal{L}_{j}^{t}|,&\text{if
$|\mathcal{S}_{j}^{t}|\geq|\mathcal{L}_{j}^{t}|$},\\\
|\mathcal{S}_{j}^{t}\cap\mathcal{L}_{j}^{t}|/|\mathcal{S}_{j}^{t}|,&\text{if
$|\mathcal{S}_{j}^{t}|<|\mathcal{L}_{j}^{t}|$},\end{cases}$ (8)

where $|\cdot|$ denotes the size of the set. We can observe that such a design
encourages the proxy model to transfer confident knowledge to the large model.

###  3.3 Server Update

After training each client using Eq. (2), we will obtain both proxy and
private models. The proxy models
$\\{\mathbf{p}_{1}^{t},\cdots,\mathbf{p}_{N}^{t}\\}$ will be uploaded to the
server to conduct aggregation, resulting in a share proxy global model
$\mathbf{p}_{g}^{t}$ using any existing data-free aggregation approaches, such
as FedAvg (McMahan et al., 2017) and FedProx (Li et al., 2020). The global
model $\mathbf{p}_{g}^{t}$ will be distributed to each client at the next
communication round as the initialization of the proxy client model, i.e.,
$\mathbf{p}_{i}^{t+1}=\mathbf{p}_{g}^{t}$. The proposed FedType framework will
be iteratively executed for the client update and server update until
convergence.

Table 1: Performance (%) comparison under the heterogeneous cross-device
settings.

Agg. Method | Model | FMNIST | CIFAR-10 | CIFAR-100
---|---|---|---|---
$\alpha=1$ | $\alpha=0.5$ | $\alpha=0.1$ | $\alpha=1$ | $\alpha=0.5$ | $\alpha=0.1$ | $\alpha=1$ | $\alpha=0.5$ | $\alpha=0.1$
✗ | FedProto | 85.05 | 87.66 | 89.04 | 76.59 | 78.17 | 82.96 | 58.03 | 66.31 | 68.60
✗ | FML | 86.54 | 90.71 | 92.63 | 80.80 | 85.24 | 88.58 | 58.77 | 66.90 | 68.74
FedAvg | FedType${}_{\text{global}}$ | 84.11 | 83.93 | 81.32 | 66.40 | 63.39 | 58.17 | 38.36 | 38.17 | 35.45
FedType${}_{\text{proxy}}$ | 86.09 | 89.45 | 93.16 | 80.65 | 82.57 | 85.04 | 56.24 | 61.06 | 62.31
FedType${}_{\text{private}}$ | 87.26 | 91.22 | 94.77 | 82.56 | 86.83 | 91.90 | 57.33 | 65.69 | 68.14
FedProx | FedType${}_{\text{global}}$ | 86.96 | 86.44 | 84.29 | 68.26 | 65.86 | 63.75 | 41.88 | 39.31 | 36.53
FedType${}_{\text{proxy}}$ | 87.03 | 91.50 | 92.64 | 82.19 | 82.48 | 87.80 | 58.56 | 61.22 | 62.64
FedType${}_{\text{private}}$ | 87.65 | 93.84 | 94.98 | 83.69 | 86.92 | 92.03 | 59.18 | 65.45 | 68.37
pFedMe | FedType${}_{\text{global}}$ | 87.82 | 87.13 | 85.86 | 68.71 | 65.22 | 64.95 | 41.55 | 40.92 | 38.60
FedType${}_{\text{proxy}}$ | 88.63 | 92.05 | 93.38 | 82.64 | 83.00 | 88.14 | 59.04 | 62.68 | 64.89
FedType${}_{\text{private}}$ | 88.96 | 92.36 | 94.86 | 83.47 | 87.24 | 92.16 | 59.78 | 67.07 | 69.51
pFedBayes | FedType${}_{\text{global}}$ | 88.20 | 87.85 | 86.04 | 68.41 | 66.87 | 63.32 | 43.73 | 41.24 | 38.72
FedType${}_{\text{proxy}}$ | 89.69 | 92.11 | 93.29 | 83.33 | 84.49 | 89.10 | 59.47 | 62.96 | 63.51
FedType${}_{\text{private}}$ | 90.26 | 93.17 | 95.88 | 84.09 | 88.67 | 92.38 | 59.62 | 67.35 | 69.60

##  4 Experiments

###  4.1 Experimental Setups

Data preparation. We assess the effectiveness of the proposed FedType approach
through image classification tasks conducted in the cross-device scenario on
FMNIST, CIFAR-10, and CIFAR-100 datasets, and cross-silo scenario on Fed-
ISIC19 dataset (Ogier du Terrail et al., 2022). For the cross-device
experiments, we follow existing work (Yurochkin et al., 2019; Hsu et al.,
2019) to set heterogeneity degrees by adjusting the Dirichlet distribution’s
concentration parameter $\alpha$. We set $\alpha=1,0.5,0.1$, respectively in
our experiments. The details of the data partition can be found in Appendix C.

Baselines. In our scenario, clients employ distinct network structures, and we
abstain from using public data. FML (Shen et al., 2023) is the only baseline
with the same setting as ours, which uses bidirectional knowledge distillation
to learn the model. Besides, we employ FedProto (Tan et al., 2022) as another
baseline, which aggregates class prototypes instead of model parameters.

The outcomes of FedType consist of three models: a shared global model
(FedType${}_{\text{global}}=\mathbf{p}_{g}^{T}$), a set of proxy client models
(FedType${}_{\text{proxy}}=\\{\mathbf{p}_{1}^{T},\cdots,\mathbf{p}_{N}^{T}\\}$),
and a set of private client models
(FedType${}_{\text{private}}=\\{\mathbf{w}_{1}^{T},\cdots,\mathbf{w}_{N}^{T}\\}$),
where $T$ represents the total number of communication rounds. Subsequently,
we compare the average client accuracy across these three types of models.
Moreover, the inherent flexibility of model aggregation in FedType enables us
to assess various representative aggregation approaches, including FedAvg
(McMahan et al., 2017), FedProx (Li et al., 2020), pFedMe (T Dinh et al.,
2020), and pFedBayes (Zhang et al., 2022). More details of the baselines can
be found in Appendix D.

Impletation details. To replicate the model-heterogeneous scenario, we
assemble a model pool comprising private client models, encompassing
ResNet-18, ResNet-34, ResNet-50, RestNet-101, ResNet-152, VGG-11, VGG-13,
VGG-16, and VGG-19, which will be randomly assigned to a client. In our
primary experiments, we designate ResNet-18 as the small proxy model, as
presented in Table 1. Additionally, we delve into alternative options for
proxy model selection, detailed in Table 3. Our reported metric is the average
accuracy on _100 clients for the cross-device setting and 6 clients for the
cross-silo setting_. More details about the model and hyperparameters can be
found in Appendix E and F, respectively.

###  4.2 Results of the Heterogeneous Model Setting

####  4.2.1 Cross-device Evaluation

Table 1 presents the average client accuracy across three image datasets,
employing different aggregation methods under varying label heterogeneity
distributions ($\alpha$’s). Observing the results, it is evident that the
proposed FedType${}_{\text{private}}$ consistently outperforms both the shared
global model FedType${}_{\text{global}}$ and the learned proxy model
FedType${}_{\text{proxy}}$ across diverse aggregation methods.

Particularly noteworthy is the global model’s lower performance, especially in
challenging tasks. This aligns with our design, where the shared global model
primarily serves as an agent for information exchange in personalized
federated learning. Despite sharing the same network structure as the global
model, the proxy models, through uncertainty-based asymmetrical reciprocity
learning, acquire valuable knowledge, contributing to their enhanced
performance. Furthermore, comparing different aggregation methods reveals that
personalized approaches (pFedMe and pFedBayes) exhibit superior performance
compared to general methods (FedAvg and FedProx), aligning with our
expectations.

Notably, an increase in the value of $\alpha$ corresponds to an overall
performance improvement for FedType${}_{\text{proxy}}$ and
FedType${}_{\text{private}}$. This observation aligns with our data partition
method, where training and testing data follow the same distribution. A larger
$\alpha$ increases label categories for each client, rendering the
classification task more challenging. However, the performance of
FedType${}_{\text{global}}$ contradicts this trend. This is attributed to the
shared global model serving as an average representation of all proxy client
models, performing better in scenarios where data follows an independent and
identical distribution, such as a large $\alpha$.

####  4.2.2 Cross-silo Evaluation

An additional experiment is conducted to assess the effectiveness of the
proposed FedType under the cross-silo setting, utilizing the Fed-ISIC19
dataset. The results are illustrated in Figure 2, where $x$-axis denotes the
aggregation method, and $y$-axis is the average client accuracy. Similar
observations to those in Table 1 emerge, where the private models outperform
the proxy models, which, in turn, surpass the global models. This consistent
trend validates the efficacy of the proposed uncertainty-based asymmetrical
reciprocity learning in mutually enhancing the capabilities of proxy and
private models.

Figure 2: Results (%) of the cross-silo evaluation.

####  4.2.3 Ablation Study

We compare our model with the following variants of FedType (i.e.,
FedType${}_{\text{private}}$) in this ablation study: (1)
FedType${}_{\text{sym}}$: We replace the backward loss
$\overleftarrow{\mathcal{J}}_{i}^{t}$ with traditional knowledge distillation
loss in Eq. (2), which can be treated as symmetrical reciprocity learning. (2)
FedType${}_{\text{Top}K}$: We choose top $K$ labels based on the ranking of
the logits to construct the label set rather than using the proposed conformal
prediction set estimation in the Section 3.2.3. Here, we set $K=3$. (3)
FedTypeη=1: We simply set the consensus weight $\eta_{j}^{t}=1$ in Eq. (8).
(4) FedTypeg=0.5: We simply set the function
$g(\Delta^{t},\lambda)=\lambda=0.5$ in Eq. (6). The ablation study results are
shown in Table 2. The cross-device scenario is validated on the CIFAR-10
dataset by setting $\alpha=0.5$ and the number of clients as 100. The
aggregation method on the server is FegAvg for both cross-device and cross-
silo scenarios.

Observing Table 2 reveals that each component effectively enhances performance
but to varying degrees. Firstly, FedType${}_{\text{sym}}$ exhibits the least
favorable performance among all baselines, underscoring the imperative need
for considering asymmetrical knowledge distillation. Secondly, the outcomes of
FedType${}_{\text{Top}K}$ suggest the effectiveness of asymmetrical
reciprocity learning. However, opting for an arbitrarily fixed number of top
ranks proves suboptimal. Therefore, it becomes essential to dynamically
determine the ranks for each sample. Thirdly, the results of FedTypeη=1 and
FedTypeg=0.5 demonstrate that employing uncertainty set prediction contributes
to a performance increase compared with FedType${}_{\text{Top}K}$. In
conclusion, considering asymmetrical reciprocity learning, dynamically
adjusting the prediction set based on the model training performance, and
estimating the consensus weights of samples prove to be valuable strategies
for enhancing overall performance.

Table 2: Ablation study performance ($\%$) comparison.

Dataset | CIFAR-10 | Fed-ISIC19
---|---|---
FedType${}_{\text{sym}}$ | 82.24 | 49.89
FedType${}_{\text{Top}K}$ | 84.69 | 51.26
FedTypeη=1 | 85.12 | 52.40
FedTypeg=0.5 | 86.24 | 54.17
FedType | 86.83 | 55.95

Table 3: Proxy model study. The approximate model sizes are shown using the
model parameters (in millions).

Proxy | Parameter | FMNIST | CIFAR-10 | CIFAR-100 | Fed-ISIC19
---|---|---|---|---|---
Model | Size | $\alpha=1$ | $\alpha=0.5$ | $\alpha=0.1$ | $\alpha=1$ | $\alpha=0.5$ | $\alpha=0.1$ | $\alpha=1$ | $\alpha=0.5$ | $\alpha=0.1$ | ✗
ShuffleNet-V2 | 2.27M | 87.11 | 88.86 | 91.58 | 81.25 | 81.79 | 84.48 | 50.44 | 57.17 | 60.35 | 47.48
MobileNet-V1 | 3.21M | 87.59 | 88.93 | 91.13 | 81.04 | 82.95 | 85.06 | 51.31 | 59.30 | 62.76 | 50.32
EfficientNet-B0 | 5.29M | 88.43 | 89.55 | 92.14 | 82.13 | 83.76 | 87.02 | 54.12 | 61.89 | 63.61 | 53.13
ResNet-18 | 11.17M | 87.26 | 91.22 | 94.77 | 82.56 | 86.83 | 91.90 | 57.33 | 65.69 | 68.14 | 55.95

####  4.2.4 Dynamic Conformal Prediction

In Eq. (6), we introduce a calibration function $g(\Delta^{t},\lambda)$ to
regulate the size of prediction sets based on the model’s performance change.
In this experiment, our objective is to examine the impact of selecting the
function $g(\Delta^{t},\lambda)$ when $\Delta^{t}\in[-1,0]$. Let
$g_{1}=g(\Delta^{t},\lambda)=\lambda\cdot\Delta^{t}-\Delta^{t}+\lambda$ and
$g_{2}=g(\Delta^{t},\lambda)=\lambda$ ($\forall\Delta^{t}\in[-1,1]$). We also
consider two alternative quadratic functions:
$g_{3}=g(\Delta^{t},\lambda)=\lambda\cdot{\Delta^{t}}^{2}+\lambda$ and
$g_{4}=g(\Delta^{t},\lambda)=-\lambda\cdot{\Delta^{t}}^{2}-{\Delta^{t}}+\lambda$.
When $\Delta^{t}\in(0,1]$, $g(\Delta^{t},\lambda)=\lambda$ for all functions
with $\lambda=0.5$.

Figure 3 (a) and (b) illustrate the $g$ function’s representation and its
impact on private models respectively. Analyzing the experimental outcomes, we
note the following observations: (1) The performance with functions $g_{1}$,
$g_{3}$, or $g_{4}$ consistently surpasses that of $g_{2}$. This suggests that
adjusting the value of $\lambda$ in response to accuracy fluctuations
positively influences performance on both the Cifar-10 and Fed-ISIC19
datasets. (2) Among $g_{1}$, $g_{3}$, and $g_{4}$, $g_{3}$ exhibits superior
performance on these datasets. A potential explanation is its steeper slope
near a $\Delta^{t}$ value of -1, aiding $g(\Delta^{t},\lambda)$ in rapidly
declining to identify a more optimal prediction set for the subsequent
training epoch. (3) While outcomes with different $g$ function configurations
vary, they remain within a stable range, indicating the robustness and
adaptability of our proposed FedType to the selection of the $g$ function.

Figure 3: Calibration function $g(\Delta^{t},\lambda)$ study.

####  4.2.5 Alternative Proxy Model Selection

The proxy model used in all the previous experiments is ResNet-18. To examine
the impact of proxy model selection on performance, we select four widely-used
compact models as proxies, including MobileNet-V1 (Howard et al., 2017),
ShuffleNet-V2 (Ma et al., 2018), and EfficientNet-B0 (Tan & Le, 2019). Model
details can be found in Appendix E. The results are shown in Table 3, where
the aggregation method is still FedAvg. We can observe that the performance
variation with these different proxy models generally aligns with their
respective sizes and capabilities, which aligns with our expectations. This
study of proxy models underscores the resilience of our approach to variations
in proxy model choice and further affirms the generalizability of our proposed
framework.

####  4.2.6 Client Model Architecture Analysis

In our experiments, we use a mixed combination of client models. To explore
the relationship between the model architecture difference and the framework
performance, we conduct the following experiment with $\alpha=0.5$, and the
other settings are the same as the experiments shown in Table 1. In this
experiment, we fixed the proxy model, which is ResNet-18, but used different
private models. FedType-ResNet denotes all the private client models selected
from the ResNet family pool, including ResNet-18, ResNet-34, ResNet-50,
RestNet-101, and ResNet 152. FedType-VGG denotes all the private models
belonging to the VGG family, including VGG-11, VGG-13, VGG-16, and VGG-19.
FedType-Mix is used in the main experiments, using the mixed ResNet and VGG
client models.

We report the results in Table 4. We have several exciting observations.
First, the private model performs better than the proxy model, which is better
than the global model. This observation is the same as we discussed in the
main results shown in Table 1. Second, FedType-ResNet outperforms FedType-Mix,
which further outperforms FedType-VGG. The observation demonstrates that if
the proxy and private models have similar model architectures, the proposed
bidirectional knowledge distillation performs the best.

Table 4: Performamce (%) of different client model combinations, where all
clients’ personal models are from ResNet or VGG family.

Dataset | Model | FedAvg | FedProx | pFedMe | pFedBayes
---|---|---|---|---|---
global | proxy | private | global | proxy | private | global | proxy | private | global | proxy | private
FMNIST |  FedType-ResNet | 84.55 | 89.63 | 91.12 | 85.51 | 91.86 | 93.61 | 87.41 | 93.55 | 93.64 | 87.62 | 92.33 | 94.86
FedType-Mix | 83.93 | 89.45 | 91.22 | 86.44 | 91.50 | 93.84 | 87.13 | 92.05 | 92.36 | 87.85 | 92.11 | 93.17
FedType-VGG | 83.22 | 88.87 | 90.86 | 84.31 | 91.07 | 92.46 | 86.65 | 91.77 | 92.20 | 87.04 | 91.98 | 92.46
CIFAR-10 |  FedType-ResNet | 64.59 | 83.62 | 86.88 | 66.11 | 83.79 | 87.25 | 66.89 | 84.08 | 87.10 | 67.03 | 84.55 | 88.96
FedType-Mix | 63.39 | 82.57 | 86.83 | 65.86 | 82.48 | 86.92 | 65.22 | 83.00 | 87.24 | 66.87 | 84.49 | 88.67
FedType-VGG | 61.22 | 80.61 | 83.07 | 64.95 | 80.87 | 84.14 | 63.40 | 81.36 | 84.52 | 65.48 | 81.27 | 86.11
CIFAR-100 |  FedType-ResNet | 38.95 | 62.48 | 65.53 | 40.66 | 63.07 | 67.89 | 41.56 | 62.97 | 68.16 | 41.82 | 63.57 | 67.94
FedType-Mix | 38.17 | 61.06 | 65.69 | 39.31 | 61.22 | 65.45 | 40.92 | 62.68 | 67.07 | 41.24 | 62.96 | 67.35
FedType-VGG | 35.96 | 58.74 | 62.41 | 36.77 | 58.91 | 62.88 | 39.52 | 59.61 | 63.05 | 40.08 | 61.63 | 63.85
FedISIC-19 |  FedType-ResNet | 42.86 | 53.55 | 55.84 | 42.84 | 55.51 | 56.90 | 42.21 | 53.88 | 57.03 | 42.37 | 55.22 | 57.10
FedType-Mix | 42.30 | 53.48 | 55.95 | 41.77 | 54.30 | 56.86 | 42.35 | 53.85 | 56.68 | 41.82 | 54.04 | 57.21
FedType-VGG | 40.76 | 51.06 | 52.90 | 42.04 | 51.98 | 53.69 | 42.13 | 52.61 | 53.12 | 41.01 | 52.48 | 54.55

####  4.2.7 Upper-bound Performance Exploration

In this experiment, we aim to investigate the performance upper bound, which
will be obtained by training the federated learning framework in the
homogeneous setting, where each client uses the largest model (i.e., VGG-19 in
our experiments). Since the homogeneous setting aims to train a shared global
model, we report the average performance tested on each client with the
learned global model from Homo-VGG-19 in Table 5, compared with the
performance of private models obtained by our proposed FedType with
$\alpha=0.5$. To maximize the performance of the private model with our
setting, we also test another model, named VGG-19+ResNet-18, in which all
clients used VGG-19 as the private model and ResNet-18 as the proxy model.

Table 5: Performance upper-bound analysis.

Dataset | Model | FedAvg | FedProx | pFedMe | pFedBayes
---|---|---|---|---|---
FMNIST | Homo-VGG-19 | 92.86 | 94.05 | 94.26 | 94.87
VGG-19+ResNet-18 | 91.35 | 94.12 | 93.45 | 94.33
FedType | 91.22 | 93.84 | 92.36 | 93.17
CIFAR-10 | Homo-VGG-19 | 87.89 | 88.62 | 89.10 | 89.92
VGG-19+ResNet-18 | 87.04 | 87.80 | 87.98 | 89.55
FedType | 86.83 | 86.92 | 87.24 | 88.67
CIFAR-100 | Homo-VGG-19 | 65.91 | 67.52 | 69.81 | 69.98
VGG-19+ResNet-18 | 65.77 | 66.71 | 67.89 | 68.89
FedType | 65.69 | 65.45 | 67.07 | 67.35
FedISIC-19 | Homo-VGG-19 | 56.89 | 58.94 | 59.06 | 59.22
VGG-19+ResNet-18 | 56.62 | 57.12 | 57.59 | 57.96
FedType | 55.95 | 56.86 | 56.68 | 57.21

We can observe that the upper bound performance obtained by the Homo-VGG-19
model performs the best, and VGG-19+ResNet-18 is better than FedType but with
marginal improvements, which aligns with our expectations. However, our model
and VGG-19+ResNet-18 use ResNet-18 (with 11.17 million parameters) as the
proxy model to exchange information between clients and the server. However,
the Homo-VGG-19 model has 144 million parameters, which is 12.9X of our model.
Considering the communication costs, FedType is an effective solution to
address the model heterogeneity challenge in federated learning.

####  4.2.8 Scalability Analysis

In this experiment, we validate the scalability of the proposed FedType with a
different number of clients on the CIRFA-10 dataset with a sample rate 10% and
$\alpha=0.5$. The client model pool contains ResNet-18, ResNet-34, and
ResNet-50. The proxy model used in this experiment is MobileNet-V1, and the
aggregation method is FedAvg. Table 6 presents our experimental findings with
client counts of 50, 100, and 300. A noticeable trend is the marginal decline
in performance as the number of clients rises. This outcome is attributable to
the distribution of the same total volume of training data among a greater
number of clients, resulting in reduced data availability per client. Such a
scenario leads to a performance dip within a specific range. These results
effectively illustrate the scalability of our proposed FedType, highlighting
its adaptability to varying client numbers in federated learning environments.

###  4.3 Results of the Homogeneous Model Setting

While the primary focus of the proposed FedType is addressing the challenge of
model heterogeneity, it showcases versatility by extending its application to
the model homogeneous setting. To illustrate its generalization ability, we
homogenize all clients to utilize the same model structure, VGG-11, with
FedType employing ResNet-18 as the proxy model in this experiment. A
comparison is made with two representative federated learning approaches,
FedAvg and pFedMe, and the results are outlined in Figure 4. Notably, FedType,
with fewer parameters (only around 8% of baselines’ parameters), achieves
comparable performance with existing homogeneous federated learning models.
This underscores the effectiveness and advantages of FedType, even in a
homogeneous model setting. More experimental results can be found in Appendix
Sections G to M.

Table 6: scalability study in the cross-device setting. Setting | 50 | 100 | 300
---|---|---|---
FedType${}_{\text{global}}$ | 62.88 | 59.40 | 56.01
FedType${}_{\text{proxy}}$ | 84.33 | 81.46 | 77.90
FedType${}_{\text{private}}$ | 86.74 | 84.26 | 80.15
Figure 4: Homogenenous evaluation.

##  5 Conclusion

We have designed FedType to effectively tackle the challenges of model
heterogeneity without relying on public data. By implementing a novel
uncertainty-based asymmetrical reciprocity learning method, we have not only
demonstrated the feasibility but also showed the superiority of FedType in
handling diverse model structures while safeguarding client privacy and
minimizing communication costs. Our comprehensive experiments across multiple
benchmark datasets illustrate the effectiveness of FedType in different
scenarios. We believe that FedType significantly contributes to the
advancement of general federated learning and holds immense potential for
practical applications in real-world scenarios.

## Acknowledgements

This work is partially supported by the National Science Foundation under
Grant No. 2333790, 2212323, and 2238275 and the National Institutes of Health
under Grant No. R01AG077016.

## Impact Statement

The introduction of FedType represents a significant advancement in the field
of federated learning (FL), addressing crucial challenges associated with
heterogeneous model aggregation. By leveraging small identical proxy models,
FedType ensures secure and efficient information exchange among clients. The
implications of this research are far-reaching, promising to enhance the
application of federated learning in real-world scenarios where data privacy
and efficient communication are paramount. By eliminating the reliance on
public data and ensuring robust performance across diverse conditions, FedType
has the potential to revolutionize the deployment of FL in sensitive and
resource-constrained environments, including healthcare, finance, and beyond.

The proposed framework has two major limitations. First, selecting the proxy
model is non-trivial and further influences the framework performance. In
future work, we plan to design an automatic strategy to adaptively select the
proxy model according to the types of private models. Second, although the
proposed framework is effective and efficient, the designed uncertainty-based
asymmetrical reciprocity learning between the proxy and private model runs
slightly slower than the naive bidirectional knowledge distillation. Thus, we
need to develop a more efficient approach to accelerate this step. From our
perspective, there is no negative social impact on the designed framework.

## References

  * Alam et al. (2022) Alam, S., Liu, L., Yan, M., and Zhang, M.  Fedrolex: Model-heterogeneous federated learning with rolling sub-model extraction.  _Advances in Neural Information Processing Systems_ , 35:29677–29690, 2022.
  * Angelopoulos et al. (2020) Angelopoulos, A., Bates, S., Malik, J., and Jordan, M. I.  Uncertainty sets for image classifiers using conformal prediction.  _arXiv preprint arXiv:2009.14193_ , 2020.
  * Angelopoulos & Bates (2021) Angelopoulos, A. N. and Bates, S.  A gentle introduction to conformal prediction and distribution-free uncertainty quantification.  _arXiv preprint arXiv:2107.07511_ , 2021.
  * Angelopoulos et al. (2022) Angelopoulos, A. N., Bates, S., Fisch, A., Lei, L., and Schuster, T.  Conformal risk control.  _arXiv preprint arXiv:2208.02814_ , 2022.
  * Balasubramanian et al. (2014) Balasubramanian, V., Ho, S.-S., and Vovk, V.  _Conformal prediction for reliable machine learning: theory, adaptations and applications_.  Newnes, 2014.
  * Bao et al. (2023) Bao, W., Wang, H., Wu, J., and He, J.  Optimizing the collaboration structure in cross-silo federated learning.  _arXiv preprint arXiv:2306.06508_ , 2023.
  * Barber et al. (2021) Barber, R. F., Candes, E. J., Ramdas, A., and Tibshirani, R. J.  Predictive inference with the jackknife+.  2021\.
  * Barber et al. (2023) Barber, R. F., Candes, E. J., Ramdas, A., and Tibshirani, R. J.  Conformal prediction beyond exchangeability.  _The Annals of Statistics_ , 51(2):816–845, 2023.
  * Bhatt et al. (2021) Bhatt, U., Antorán, J., Zhang, Y., Liao, Q. V., Sattigeri, P., Fogliato, R., Melançon, G., Krishnan, R., Stanley, J., Tickoo, O., et al.  Uncertainty as a form of transparency: Measuring, communicating, and using uncertainty.  In _Proceedings of the 2021 AAAI/ACM Conference on AI, Ethics, and Society_ , pp. 401–413, 2021.
  * Bouacida & Mohapatra (2021) Bouacida, N. and Mohapatra, P.  Vulnerabilities in federated learning.  _IEEE Access_ , 9:63229–63249, 2021.
  * Chen et al. (2023) Chen, D., Yao, L., Gao, D., Ding, B., and Li, Y.  Efficient personalized federated learning via sparse model-adaptation.  _arXiv preprint arXiv:2305.02776_ , 2023.
  * Cho et al. (2023) Cho, Y. J., Wang, J., Chirvolu, T., and Joshi, G.  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.
  * Dennis et al. (2021) Dennis, D. K., Li, T., and Smith, V.  Heterogeneity for the win: One-shot federated clustering.  In _International Conference on Machine Learning_ , pp. 2611–2620. PMLR, 2021.
  * Diao et al. (2020) Diao, E., Ding, J., and Tarokh, V.  Heterofl: Computation and communication efficient federated learning for heterogeneous clients.  In _International Conference on Learning Representations_ , 2020.
  * Fisch et al. (2021) Fisch, A., Schuster, T., Jaakkola, T., and Barzilay, R.  Few-shot conformal prediction with auxiliary tasks.  In _International Conference on Machine Learning_ , pp. 3329–3339. PMLR, 2021.
  * Guo et al. (2017) Guo, C., Pleiss, G., Sun, Y., and Weinberger, K. Q.  On calibration of modern neural networks.  In _International conference on machine learning_ , pp. 1321–1330. PMLR, 2017.
  * He et al. (2016) He, K., Zhang, X., Ren, S., and Sun, J.  Deep residual learning for image recognition.  In _Proceedings of the IEEE conference on computer vision and pattern recognition_ , pp. 770–778, 2016.
  * Howard et al. (2017) Howard, A. G., Zhu, M., Chen, B., Kalenichenko, D., Wang, W., Weyand, T., Andreetto, M., and Adam, H.  Mobilenets: Efficient convolutional neural networks for mobile vision applications.  _arXiv preprint arXiv:1704.04861_ , 2017.
  * Hsu et al. (2019) Hsu, T.-M. H., Qi, H., and Brown, M.  Measuring the effects of non-identical data distribution for federated visual classification.  _arXiv preprint arXiv:1909.06335_ , 2019.
  * Huang et al. (2022) Huang, W., Ye, M., and Du, B.  Learn from others and be yourself in heterogeneous federated learning.  In _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition_ , pp. 10143–10153, 2022.
  * Kuleshov et al. (2018) Kuleshov, V., Fenner, N., and Ermon, S.  Accurate uncertainties for deep learning using calibrated regression.  In _International conference on machine learning_ , pp. 2796–2804. PMLR, 2018.
  * Li & Wang (2019) Li, D. and Wang, J.  Fedmd: Heterogenous federated learning via model distillation.  _arXiv preprint arXiv:1910.03581_ , 2019.
  * Li et al. (2020) Li, T., Sahu, A. K., Zaheer, M., Sanjabi, M., Talwalkar, A., and Smith, V.  Federated optimization in heterogeneous networks.  _Proceedings of Machine learning and systems_ , 2:429–450, 2020.
  * Lin et al. (2020) Lin, T., Kong, L., Stich, S. U., and Jaggi, M.  Ensemble distillation for robust model fusion in federated learning.  _Advances in Neural Information Processing Systems_ , 33:2351–2363, 2020.
  * Lu et al. (2023) Lu, C., Yu, Y., Karimireddy, S. P., Jordan, M., and Raskar, R.  Federated conformal predictors for distributed uncertainty quantification.  In _International Conference on Machine Learning_ , pp. 22942–22964. PMLR, 2023.
  * Lyu et al. (2022) Lyu, L., Yu, H., Ma, X., Chen, C., Sun, L., Zhao, J., Yang, Q., and Philip, S. Y.  Privacy and robustness in federated learning: Attacks and defenses.  _IEEE transactions on neural networks and learning systems_ , 2022.
  * Ma et al. (2018) Ma, N., Zhang, X., Zheng, H.-T., and Sun, J.  Shufflenet v2: Practical guidelines for efficient cnn architecture design.  In _Proceedings of the European conference on computer vision (ECCV)_ , pp. 116–131, 2018.
  * Maddox et al. (2019) Maddox, W. J., Izmailov, P., Garipov, T., Vetrov, D. P., and Wilson, A. G.  A simple baseline for bayesian uncertainty in deep learning.  _Advances in neural information processing systems_ , 32, 2019.
  * Marfoq et al. (2022) Marfoq, O., Neglia, G., Vidal, R., and Kameni, L.  Personalized federated learning through local memorization.  In _International Conference on Machine Learning_ , pp. 15070–15092. PMLR, 2022.
  * McMahan et al. (2017) McMahan, B., Moore, E., Ramage, D., Hampson, S., and y Arcas, B. A.  Communication-efficient learning of deep networks from decentralized data.  In _Artificial intelligence and statistics_ , pp. 1273–1282. PMLR, 2017.
  * Neal (2012) Neal, R. M.  _Bayesian learning for neural networks_ , volume 118.  Springer Science & Business Media, 2012.
  * Ogier du Terrail et al. (2022) Ogier du Terrail, J., Ayed, S.-S., Cyffers, E., Grimberg, F., He, C., Loeb, R., Mangold, P., Marchand, T., Marfoq, O., Mushtaq, E., et al.  Flamby: Datasets and benchmarks for cross-silo federated learning in realistic healthcare settings.  _Advances in Neural Information Processing Systems_ , 35:5315–5334, 2022.
  * Papadopoulos et al. (2002) Papadopoulos, H., Proedrou, K., Vovk, V., and Gammerman, A.  Inductive confidence machines for regression.  In _Machine Learning: ECML 2002: 13th European Conference on Machine Learning Helsinki, Finland, August 19–23, 2002 Proceedings 13_ , pp. 345–356. Springer, 2002.
  * Plassier et al. (2023) Plassier, V., Makni, M., Rubashevskii, A., Moulines, E., and Panov, M.  Conformal prediction for federated uncertainty quantification under label shift.  _arXiv preprint arXiv:2306.05131_ , 2023.
  * Platt et al. (1999) Platt, J. et al.  Probabilistic outputs for support vector machines and comparisons to regularized likelihood methods.  _Advances in large margin classifiers_ , 10(3):61–74, 1999.
  * Sankaranarayanan et al. (2022) Sankaranarayanan, S., Angelopoulos, A., Bates, S., Romano, Y., and Isola, P.  Semantic uncertainty intervals for disentangled latent spaces.  _Advances in Neural Information Processing Systems_ , 35:6250–6263, 2022.
  * Shafer & Vovk (2008) Shafer, G. and Vovk, V.  A tutorial on conformal prediction.  _Journal of Machine Learning Research_ , 9(3), 2008.
  * Shamsian et al. (2021) Shamsian, A., Navon, A., Fetaya, E., and Chechik, G.  Personalized federated learning using hypernetworks.  In _International Conference on Machine Learning_ , pp. 9489–9502. PMLR, 2021.
  * Sharma et al. (2023) Sharma, A., Veer, S., Hancock, A., Yang, H., Pavone, M., and Majumdar, A.  Pac-bayes generalization certificates for learned inductive conformal prediction.  _arXiv preprint arXiv:2312.04658_ , 2023.
  * Shen et al. (2023) Shen, T., Zhang, J., Jia, X., Zhang, F., Lv, Z., Kuang, K., Wu, C., and Wu, F.  Federated mutual learning: a collaborative machine learning method for heterogeneous data, models, and objectives.  _Frontiers of Information Technology & Electronic Engineering_, 24(10):1390–1402, 2023.
  * Simonyan & Zisserman (2014) Simonyan, K. and Zisserman, A.  Very deep convolutional networks for large-scale image recognition.  _arXiv preprint arXiv:1409.1556_ , 2014.
  * T Dinh et al. (2020) T Dinh, C., Tran, N., and Nguyen, J.  Personalized federated learning with moreau envelopes.  _Advances in Neural Information Processing Systems_ , 33:21394–21405, 2020.
  * Tan & Le (2019) Tan, M. and Le, Q.  Efficientnet: Rethinking model scaling for convolutional neural networks.  In _International conference on machine learning_ , pp. 6105–6114. PMLR, 2019.
  * Tan et al. (2022) Tan, Y., Long, G., Liu, L., Zhou, T., Lu, Q., Jiang, J., and Zhang, C.  Fedproto: Federated prototype learning across heterogeneous clients.  In _Proceedings of the AAAI Conference on Artificial Intelligence_ , volume 36, pp. 8432–8440, 2022.
  * Tolpegin et al. (2020) Tolpegin, V., Truex, S., Gursoy, M. E., and Liu, L.  Data poisoning attacks against federated learning systems.  In _Computer Security–ESORICS 2020: 25th European Symposium on Research in Computer Security, ESORICS 2020, Guildford, UK, September 14–18, 2020, Proceedings, Part I 25_ , pp. 480–501. Springer, 2020.
  * Vovk (2015) Vovk, V.  Cross-conformal predictors.  _Annals of Mathematics and Artificial Intelligence_ , 74:9–28, 2015.
  * Vovk et al. (1999) Vovk, V., Gammerman, A., and Saunders, C.  Machine-learning applications of algorithmic randomness.  1999\.
  * Vovk et al. (2005) Vovk, V., Gammerman, A., and Shafer, G.  _Algorithmic learning in a random world_ , volume 29.  Springer, 2005.
  * Wang et al. (2023) Wang, J., Yang, X., Cui, S., Che, L., Lyu, L., Xu, D., and Ma, F.  Towards personalized federated learning via heterogeneous model reassembly.  In _Thirty-seventh Conference on Neural Information Processing Systems_ , 2023.
  * Yao et al. (2021) Yao, D., Pan, W., O’Neill, M. J., Dai, Y., Wan, Y., Jin, H., and Sun, L.  Fedhm: Efficient federated learning for heterogeneous models via low-rank factorization.  _arXiv preprint arXiv:2111.14655_ , 2021.
  * Yi et al. (2023) Yi, L., Wang, G., Liu, X., Shi, Z., and Yu, H.  Fedgh: Heterogeneous federated learning with generalized global header.  _arXiv preprint arXiv:2303.13137_ , 2023.
  * Yu et al. (2022) Yu, S., Qian, W., and Jannesari, A.  Resource-aware federated learning using knowledge extraction and multi-model fusion.  _arXiv preprint arXiv:2208.07978_ , 2022.
  * Yurochkin et al. (2019) Yurochkin, M., Agarwal, M., Ghosh, S., Greenewald, K., Hoang, N., and Khazaeni, Y.  Bayesian nonparametric federated learning of neural networks.  In _International conference on machine learning_ , pp. 7252–7261. PMLR, 2019.
  * Zhang et al. (2018) Zhang, X., Zhou, X., Lin, M., and Sun, J.  Shufflenet: An extremely efficient convolutional neural network for mobile devices.  In _Proceedings of the IEEE conference on computer vision and pattern recognition_ , pp. 6848–6856, 2018.
  * Zhang et al. (2022) Zhang, X., Li, Y., Li, W., Guo, K., and Shao, Y.  Personalized federated learning via variational bayesian inference.  In _International Conference on Machine Learning_ , pp. 26293–26310. PMLR, 2022.
  * Zhou et al. (2022) Zhou, Y., Wu, J., Wang, H., and He, J.  Adversarial robustness through bias variance decomposition: A new perspective for federated learning.  In _Proceedings of the 31st ACM International Conference on Information & Knowledge Management, Atlanta, GA, USA, October 17-21, 2022_, pp. 2753–2762. ACM, 2022.

##  Appendix A Algorithm Flow

In this subsection, we show the whole algorithm flow of our proposed FedType
in Algorithm 2.

Input: Client training data ${\mathcal{D}_{1},\cdots,\mathcal{D}_{N}}$,
private models $\\{\mathbf{w}_{1},\cdots,\mathbf{w}_{N}$}, proxy models
$\\{\mathbf{p}_{1},\cdots,\mathbf{p}_{N}\\}$, validation data
$\\{\mathcal{D}_{1}^{\prime},\cdots,\mathcal{D}_{N}^{\prime}\\}$,
communication round T, local training epoch $R$, and hyperparameters.

for  _each communication round $t=1,2,\cdots$,T_ do

8 Client Update  __

9 for  _each epoch $r=1,\cdots,R$_ do

10 for  _active client $i\in[1,\cdots,B]$_ do

11 Conduct uncertainty-based asymmetrical reciprocity learning following the
Algorithm 1.

12 end for

13

14 end for

15 Obtain private models $\\{\mathbf{w}^{t}_{1},\cdots,\mathbf{w}^{t}_{N}\\}$
and proxy models $\\{\mathbf{p}^{t}_{1},\cdots,\mathbf{p}^{t}_{N}\\}$; Upload
proxy models $\\{\mathbf{p}^{t}_{1},\cdots,\mathbf{p}^{t}_{N}\\}$ to the
server;

16

17 Server Update  __

18 Obtain the aggregated model $\mathbf{p}^{t}_{g}$ via FedAvg or other
existing methods; Distribute the aggregated model $\mathbf{p}^{g}_{t}$ back to
clients.

19

20

21 end for

Algorithm 2 Algorithm Flow of FedType.

##  Appendix B Details of Conformal Prediction

###  B.1 Conformal Prediction

Conformal prediction is a promising statistical framework for describing the
prediction uncertainty (Vovk et al., 2005; Shafer & Vovk, 2008; Angelopoulos &
Bates, 2021; Angelopoulos et al., 2022; Barber et al., 2023; Sharma et al.,
2023). It guarantees finite-sample coverage under the mild assumption of
exchangeability. The main idea of conformal prediction involves establishing a
non-conformity score $S$ to assess the conformity between new test data and
existing data distribution. Compared with existing traditional methods for
estimating prediction uncertainty, such as Bayesian neural networks (Neal,
2012; Kuleshov et al., 2018; Maddox et al., 2019) and Platt scaling (Platt et
al., 1999; Guo et al., 2017), conformal prediction offers numerous advantages,
including its post-hoc nature, and being distribution-free and model-agnostic.

The original version of conformal prediction, also known as full conformal
prediction or transductive conformal prediction, is regarded for its
exceptional uncertainty estimation capabilities. However, its considerable
computational cost restricts its practical application. To address this,
several alternative methods have been developed to reduce the computational
demands of full conformal prediction while retaining its core utility, for
example, Split Conformal Prediction (Inductive CP) (Papadopoulos et al.,
2002), Cross-CP (Vovk, 2015), and jackknife+ (Barber et al., 2021). In
particular, Split Conformal Prediction has garnered significant attention due
to its ease of implementation. Our paper is based on split conformal
prediction, which we will introduce below.

###  B.2 Split Conformal Prediction Model Training

Next, we describe the process for deriving $\mathbf{cp}_{i}^{t}$ as outlined
in Eq. (5) following the methodology presented in the work (Angelopoulos et
al., 2020). We describe the process of obtaining $\mathbf{cp}_{i}^{t}$ in
Algorithm 3 below.

Input: A trained proxy model $\mathbf{p}_{i}^{t}$ using the training dataset
$\mathcal{D}_{i}$, validation data $\mathcal{D}_{i}^{\prime}$.

  * •

Step 1: Define the nonconformity score function $S$ based on the nonconformity
measure and given hyperparameters;

  * •

Step 2: Calculate the nonconformity score $s_{j}$ for each
$\mathbf{x}_{j}\in\mathcal{D}_{i}^{\prime}$, i.e., $s_{j}=S(\mathbf{x}_{j})$,
forming the nonconformity score set $\\{s_{j}\\}$;

  * •

Step 3: Obtain the score quantile $Q_{\theta}$ related to $\theta$ on the
nonconformity score set, i.e.,
$Q_{\theta}:=\frac{1}{n}\lceil(1-\theta)(n+1)\rceil$-th quantile
of $\\{s_{j}\\}_{j=1}^{n}$, where $n=|\mathcal{D}_{i}^{{}^{\prime}}|$;

  * •

Step 4: Integrate $\mathbf{p}_{i}^{t}$ with $Q_{\theta}$.
$\mathbf{cp}_{i}^{t}$ is first initialized by trained $\mathbf{p}_{i}^{t}$ and
then predicts and calculates the non-conformity score
on the test data coupled with the candidate label, which is then compared with
the $Q_{\theta}$.

Return: The trained conformal model $\mathbf{cp}_{i}^{t}$.

Algorithm 3 Algorithm Flow of Cmodel Training with Eq. (5).

###  B.3 Regularized Adaptive Prediction Sets Converge Guarantee

Note that our proposed dynamic adjustment mechanism retains all the
characteristics of traditional conformal prediction requirements.
Additionally, it ensures the isolation of $\mathcal{D}_{i}$ and
$\mathcal{D}^{{}^{\prime}}_{i}$ throughout the entire federated learning
process, including during local training phases. Thus, regularized adaptive
prediction sets converge can be guaranteed. More details about the upper bound
proof and lower bound proof can be found in Section D in (Angelopoulos &
Bates, 2021).

##  Appendix C Datasets

We assess the effectiveness of the proposed FedType approach through image
classification tasks conducted in both cross-device and cross-silo scenarios.
For the cross-device evaluation, we employ three image classification datasets
– FMNIST, CIFAR-10, and CIFAR-100. To introduce heterogeneity in federated
learning, a methodology inspired by existing work (Yurochkin et al., 2019; Hsu
et al., 2019) is followed. This involves manipulating the concentration
parameter $\alpha$ of the Dirichlet distribution to partition the datasets.
Here, $\alpha$ signifies the label heterogeneity across clients, with smaller
values concentrating labels on a few categories for a client (where $\alpha\to
0$ implies a client stores data with a single category). Conversely, larger
$\alpha$ values lead to clients holding a more diverse set of label categories
(where $\alpha\to+\infty$ indicates each client possesses data spanning all
categories). Our data distribution process involves initially allocating data
to clients, followed by a subsequent split into local model training, testing,
and conformal learning sets in a 7:2:1 ratio. _While the splitting method
guarantees an identical distribution for both sets, aligning effectively with
the personalized federated learning setting, it is noteworthy that the
increased value of $\alpha$ introduces heightened difficulty in the
classification tasks._

For the cross-silo setting, the Fed-ISIC19 dataset (Ogier du Terrail et al.,
2022) is employed, featuring 23,247 dermoscopy images encompassing eight
distinct types of melanoma. Following the data partition strategy of FLamby
(Ogier du Terrail et al., 2022), the number of training/testing data on six
clients is distributed as 9,930/2,483, 3,163/791, 2,690/673, 655/164, 351/88,
and 180/45, respectively. The cross-silo setting necessitates the involvement
of all clients in the training process at each communication round.

##  Appendix D Baseline Introduction

In our paper, though there is no previous heterogeneous FL work sharing the
exact setting with our work, we introduce the selected baselines used for the
homogeneous setting comparison.

  * •

FedAvg (McMahan et al., 2017): It is the vanilla baseline. In this approach,
active local clients train their models and send the parameters to a central
server. The server then computes the average of these local model parameters
and redistributes the aggregated global model to the active clients for
subsequent rounds of local training.

  * •

FedProx (Li et al., 2020): this work introduces a proximal term in the local
training of each client, quantifying the divergence between the local and
global models. This term acts as a constraint, ensuring that the local models’
personalized optimization does not deviate significantly from the global
model.

  * •

pFedMe (T Dinh et al., 2020): It uses regularized loss and decouples the
personalization problem into a bi-level optimization. pFedMe aims to develop
personalized models for each client. It does so by integrating Moreau
envelopes into the learning process to help regularize the local updates
during training;

  * •

pFedBayes (Zhang et al., 2022): It proposes an algorithm to take consideration
of the global distribution while conducting local model training. By
leveraging Bayesian inference, pFedBayes not only customizes models to fit
individual client needs better but also provides a robust framework for
managing the inherent uncertainties and variabilities in distributed datasets.

_Notably, these approaches are all used as model aggregation solutions in the
experiments due to the flexibility of our proposed framework._

##  Appendix E Model Details and Implementation

We introduce the models that we use in our experiment one by one. For all the
following client private models, we adjust the dimension of the linear layer
to fit the number of classes of the datasets accordingly.

  * •

ResNet family: ResNet (He et al., 2016) is a type of convolutional neural
network (CNN) architecture proposed in 2015. The key innovation of ResNet is
its use of residual blocks. These blocks allow the network to learn residual
functions with reference to the layer inputs, instead of learning unreferenced
functions. In our work, we use ResNet-18, 34, 50, 101, and 152. The
implementation is based on the Pytorch official
library555https://pytorch.org/vision/main/models/resnet.html.

  * •

VGG family: VGG (Simonyan & Zisserman, 2014), short for Visual Geometry Group,
refers to a deep CNN architecture. VGG was one of the first to demonstrate
that depth of the network (i.e., the number of layers) is a critical component
for achieving high performance in visual recognition tasks. In our work, we
use VGG-11, 13, 16, and 19. The implementation is based on the Pytorch
official library666https://pytorch.org/hub/pytorch_vision_vgg/.

Except for ResNet-18, we also test the following models as the proxy modes in
our experiments:

  * •

ShuffleNet: ShuffleNet(Zhang et al., 2018) is an efficient CNN designed
primarily for mobile and computing devices with limited computational
capacity. The key innovation in ShuffleNet is the use of pointwise group
convolutions and channel shuffle operations to significantly reduce the
computational cost and the number of parameters. In our work, we use
ShuffleNet V2 (Ma et al., 2018) and implement it using the Pytorch
library777https://pytorch.org/hub/pytorch_vision_shufflenet_v2/.

  * •

MobileNet: MobileNet (Howard et al., 2017) is a class of efficient CNN
architectures designed specifically for mobile and embedded vision
applications. The key innovation in MobileNet is the use of depthwise
separable convolutions, which significantly reduce the number of parameters
and the computational burden compared to standard convolutions used in more
traditional CNN architectures. In our experiment, we use MobileNet V1 (Howard
et al., 2017) and implement it based on the Github resource
888https://github.com/jmjeon94/MobileNet-Pytorch/blob/master/MobileNetV1.py.

  * •

EffcientNet: EfficientNet is an efficient CNN structure introduced in (Tan &
Le, 2019). The primary innovation of EfficientNet lies in its novel compound
scaling method, which uniformly scales the depth, width, and resolution of the
network with a set of fixed scaling coefficients. In our paper, we use
EfficientNet-B0 and implement via Pytorch
library999https://pytorch.org/vision/main/models/efficientnet.html.

##  Appendix F Hyperparameter Details

Our experimental setup involves 100 communication rounds, 100 clients, a 20%
sample ratio for the cross-device experiments, and five local training epochs.
All experiments are conducted on an NVIDIA A100 with CUDA version 12.0,
running on a Ubuntu 20.04.6 LTS server. All baselines and the proposed FedType
are implemented using PyTorch 2.0.1. For the local update, we set the learning
rate as 0.0001, the batch size is 16, and the optimizer used in the
optimization is Adam. Following the provided value in the work (Angelopoulos
et al., 2020), we set $\lambda=0.5$, $\kappa_{reg}=5$, and $\theta=0.1$ for
the local conformal model and proxy conformal model in Eq. (6), the conformal
learning batch size is 32. The learning rate in the Platt scaling process is
0.01, and the maximum iteration is 10 following the default setting.

##  Appendix G Empirical Study of Convergence

In this section, we present the convergence results of our proposed FedType
framework as illustrated in Figure 5 across different datasets: FMNIST,
CIFAR-10, CIFAR-100, and Fed-ISIC19, shown in panels (a), (b), (c), and (d),
respectively. These results are obtained under the setting of $\alpha=0.5$
with FedAvg employed as the aggregation method. The x-axis represents the
number of communication rounds, while the y-axis tracks the average accuracy
of clients. Observations from these results indicate that the private model,
the proxy model, and the global model within the FedType framework all achieve
convergence. This effectively demonstrates the robust convergence properties
of our proposed approach empirically.

Figure 5: The empirical convergence of our proposed FedType.

##  Appendix H Study of Consensus Weight $\eta$

As shown in Section 3.2.4 and Eq. (3), $\eta$ serves as a crucial metric in
our framework, quantifying the degree of consensus to regulate knowledge
transfer from the proxy model $\mathbf{p}_{i}^{t}$ to the larger model
$\mathbf{w}_{i}^{t}$. In this section, we analyze the behavior of $\eta$
across various datasets: FMNIST, CIFAR-10, CIFAR-100, and Fed-ISIC19, as
depicted in Figure 6 in panels (a), (b), (c), and (d), respectively, where
x-axis denotes the number of communication rounds, and y-axis represents the
average $\eta$ values of all training data for all active/selected clients at
round $t$. These observations are made under the setting of $\alpha=0.5$ for
FMNIST, CIFAR-10, and CIFAR-100, utilizing FedAvg as the aggregation method.

The results indicate a consistent increase in the value of $\eta$ across
communication rounds, suggesting that the large model and the proxy model are
achieving greater consensus as training progresses. Notably, an $\eta$ value
of 1 implies complete uniformity in the uncertainty set for input data between
the large and proxy models. Interestingly, for simpler datasets like FMNIST
and Fed-ISIC19, $\eta$ attains more instances of 1 and converges more rapidly
compared to the more complex CIFAR-10 and CIFAR-100 datasets. These
experimental findings highlight the dynamic nature of $\eta$ throughout the
iterative learning process and validate our strategy of assigning increased
weight to the knowledge transferred from the proxy model to the large model as
consensus grows.

Figure 6: The value of $\eta$ of our proposed FedType.

##  Appendix I Calibration Function Study

In Section 4.2.4, we detailed a dynamic calibration function designed to
adjust the prediction set size in response to changes in model performance.
This section delves deeper into the training process, specifically examining
changes quantified with $\alpha=0.5$ and FedAvg as the aggregation method. We
focus on the FMNIST, CIFAR-10, CIFAR-100 datasets under a cross-device
setting, and Fed-ISIC19 under a cross-silo setting. At the 50th communication
round, we highlight changes in active clients’ average proxy set size and
average client accuracy (last batch of the training epoch), as presented in
Table 7.

From these results, we draw several insights: (1) In scenarios where accuracy
improves across training epochs, such as with CIFAR-10 and Fed-ISIC19, the
prediction set size of the proxy model remains fairly consistent. This
stability is attributed to our specifically designed function
$g(\Delta^{t},\lambda)$. In instances where $\Delta^{t}$ falls within the
range of $(0,1]$, the function maintains $g(\Delta^{t},\lambda)=\lambda$, with
$\lambda$ consistently set at 0.5. (2) Conversely, where there’s a decrease in
accuracy, as observed in the CIFAR-100 dataset between the 1st and 3rd epochs
(accuracy dropping from 62.15% to 61.42%), the prediction set size
correspondingly reduces from 3.06 to 2.56. In contrast, between the 3rd and
5th epochs of the same dataset, as accuracy increases from 52.56% to 52.77%,
the set size also rises from 2.63% to 2.69%. Similar patterns are noted in the
FMNIST dataset. These findings underscore the efficacy of our dynamic
calibration function in adapting the prediction set size in tandem with
fluctuations in model accuracy during the training process.

Table 7: Calibration function quantitative analysis. $|\mathcal{S}^{50}|$
denotes the average size of the prediction sets by proxy models.

Dataset | FMNIST | CIFAR-10 | CIFAR-100 | Fed-ISIC19
---|---|---|---|---
Epoch | Accuracy | $|\mathcal{S}^{50}|$ | Accuracy | $|\mathcal{S}^{50}|$ | Accuracy | $|\mathcal{S}^{50}|$ | Accuracy | $|\mathcal{S}^{50}|$
1 | 87.85 (-) | 2.15 (-) | 83.26 (-) | 2.13 (-) | 62.15 (-) | 3.06 (-) | 52.10 (-) | 2.69 (-)
2 | 87.62 ($\downarrow$) | 1.75 ($\downarrow$) | 83.57 ($\uparrow$) | 2.11 ($\downarrow$) | 61.52 ($\downarrow$) | 2.75 ($\downarrow$) | 52.47 ($\uparrow$) | 2.56 ($\downarrow$)
3 | 88.21 ($\uparrow$) | 2.00 ($\uparrow$) | 83.92 ($\uparrow$) | 2.13 ($\uparrow$) | 61.42 ($\downarrow$) | 2.56 ($\downarrow$) | 52.56 ($\uparrow$) | 2.63 ($\uparrow$)
4 | 88.43 ($\uparrow$) | 2.19 ($\uparrow$) | 84.11 ($\uparrow$) | 2.13 (-) | 62.56 ($\uparrow$) | 2.62 ($\uparrow$) | 52.62 ($\uparrow$) | 2.56($\downarrow$)
5 | 88.56 ($\uparrow$) | 2.25 ($\uparrow$) | 84.21 ($\uparrow$) | 2.19 ($\uparrow$) | 62.87 ($\uparrow$) | 3.12 ($\uparrow$) | 52.77 ($\uparrow$) | 2.69 ($\uparrow$)

Figure 7: Prediction set consensus study.

##  Appendix J Conformal Prediction Set Visualization

In this section, our focus is on evaluating the performance of conformal
models for both private and proxy models on identical data points throughout
the iterative training process. This evaluation is aimed at understanding how
these models progress towards consensus on a certain data sample
$\mathbf{x}_{j}$. Utilizing the CIFAR-10 dataset and adhering to the settings
outlined in Table 1, we graphically represent the prediction set of one data
sample randomly selected from CIFAR-10 for both private and proxy models in
Figure 7. On this graph, the x-axis represents the communication round, while
the y-axis corresponds to label IDs (ranging from 1 to 10, corresponding to
the dataset’s labels).

From our analysis, it is evident that the labels within the prediction sets
generated by the conformal models of both private and proxy models
increasingly overlap as the communication rounds progress. Notably, in later
communication rounds, we observe a reduction in the size of both prediction
sets, with each set encapsulating the ground truth label 5. This convergence
of label predictions between the two models’ prediction sets suggests a
gradual move toward consensus within a specific range for the given input
data. This observation is crucial as it demonstrates the effectiveness of the
iterative training process in aligning the outputs of the private and proxy
models, contributing significantly to the overall model consensus and
performance.

##  Appendix K Experiment Results on Text Classification Task

In addition to our previous experiments, we also conducted tests on text
classification using the AG news dataset to further validate our FedType
framework. For these experiments, we created a heterogeneous model pool by
adding 3, 4, or 5 linear layers to the base of a pre-trained DistillBERT
model. Specifically, we used DistillBERT with 3 linear layers as the proxy
model in the heterogeneous setting and DistillBERT with 4 linear layers for
the homogeneous setting. The experimental setup included using $\alpha=0.5$
and FedAvg as the aggregation method, with a total of 50 communication rounds,
keeping all other settings consistent with those in Table 1. During training,
we only fine-tuned the appended linear layers, keeping the pre-trained
DistillBERT layers fixed. The results, presented in Table 8, affirm the
effectiveness of FedType in text classification tasks.

Interestingly, despite the utilization of a common pre-trained language model,
which led to similar accuracy levels as observed in our image classification
tasks, there were still noticeable differences in the performance of the
global, proxy, and private models. This aligns with the observations from our
image classification experiments, thereby reinforcing the efficacy of FedType
in handling text classification tasks. These results not only demonstrate the
versatility of our approach but also its adaptability across different types
of data.

Table 8: Performance (%) of the text classification task under the
heterogeneous cross-device settings .

Dataset | FedType${}_{\text{global}}$ | FedType${}_{\text{proxy}}$ | FedType${}_{\text{private}}$
---|---|---|---
Heterogenous | 91.21 | 91.86 | 92.16
Homogenous | 90.13 | 90.89 | 91.44

##  Appendix L Homogeneous Results on Fed-ISIC19 dataset

Figure 8: Homogeneous evaluation with Fed-ISIC19 dataset under the cross-silo
setting.

In this section, we report the experiment results on Fed-ISIC19 under the
homogeneous setting due to the page limit of the main paper. We maintain all
the other settings as the heterogeneous setting except replacing the private
model of clients with VGG-11, which is the same as the setting in Section 4.3.
The results are shown in figure 8. Based on the experiment results, we observe
that they generally align with the results under the heterogenous setting in
Figure 2. It further demonstrates the effectiveness and robustness of our
proposed FedType under both the heterogenous setting and the homogeneous
setting.

Table 9: Comparision of the average of 3 epoch training time (second) on
datasets

Setting | FMNIST | CIFAR-10 | CIFAR-100 | Fed-ISIC19
---|---|---|---|---
Private model training in FedType | 265.67 | 279.33 | 312.67 | 526.00
Private model training with symmetrical KD loss | 238.00 | 241.67 | 275.33 | 493.67
Proxy model training in FedType | 161.33 | 164.67 | 176.33 | 371.00
Proxy model training with only CE loss | 129.67 | 131.33 | 152.00 | 253.33

##  Appendix M Resource Usage Discussion

In this section, we evaluate the computational efficiency of our proposed
FedType framework. Our focus is on the duration of each training epoch,
particularly how it compares to scenarios without our specially designed
module. To this end, we compare the training time of the private model in
FedType with its counterpart that only uses symmetrical KD loss. Similarly, we
assess the training time of the proxy model in FedType against training with
only the CE loss. The average training time over three epochs is reported. The
results in Table 9 indicate that, while our approach does entail additional
training time within a certain range, this increase is justifiable considering
the performance improvements detailed in the ablation study (Table 2) and the
communication cost reduction observed in Figure 4. This aspect is also
addressed in the discussion on limitations, where we contemplate strategies to
further minimize computational costs in future work.

You can have as much text here as you want. The main body must be at most $8$
pages long. For the final version, one more page can be added. If you want,
you can use an appendix like this one.

The $\mathtt{\backslash onecolumn}$ command above can be kept in place if you
prefer a one-column appendix, or can be removed if you prefer a two-column
appendix. Apart from this possible change, the style (font size, spacing,
margins, page numbering, etc.) should be kept the same as the main body.