# Flow: Per-Instance Personalized Federated Learning Through Dynamic Routing

Kunjal Panchal,1 Sunav Choudhary,2 Hui Guan1

###### Abstract

Personalization in Federated Learning (FL) aims to modify a collaboratively
trained global model according to each client. Current approaches to
personalization in FL are at a coarse granularity, i.e. all the input
instances of a client use the same personalized model. This ignores the fact
that some instances are more accurately handled by the global model due to
better generalizability. To address this challenge, this work proposes Flow, a
fine-grained stateless personalized FL approach. Flow creates dynamic
personalized models by learning a routing mechanism that determines whether an
input instance prefers the local parameters or its global counterpart. Thus,
Flow introduces per-instance routing in addition to leveraging per-client
personalization to improve accuracies at each client. Further, Flow is
stateless which makes it unnecessary for a client to retain its personalized
state across FL rounds. This makes Flow practical for large-scale FL settings
and friendly to newly joined clients. Evaluations on Stackoverflow, Reddit,
and EMNIST datasets demonstrate the superiority in prediction accuracy of Flow
over state-of-the-art non-personalized and only per-client personalized
approaches to FL.

##  1 Introduction

Federated Learning (FL) is a distributed machine learning paradigm where
clients with resource-constrained edge devices collaboratively train a global
machine learning model. In FL, each client updates a local copy of the global
model weights using its local data, and sends updated weights back to a
server. The server aggregates the weights and updates the global model. FL is
beneficial for clients who do not have computational resources or the amount
of data needed to train a model locally. However, it suffers from a well-known
statistical heterogeneity issue: clients whose data distribution is different
from the global data distribution wouldn’t have performance incentives to
participate in FL training due to the arbitrarily bad prediction performance
of the global model on their data (Kairouz et al. 2021; Tan et al. 2022;
Fallah, Mokhtari, and Ozdaglar 2020).

(a) No Personalization

(b) Per-Instance Personalization

(c) Per-Client
Personalization

(d) Per-Client Per-Instance Personalization (Proposed)

Figure 1: Illustration of different granularity of personalization using
machine learning models’ computation graph. Yellow and green nodes are
operators in a model that are shared across clients. Edges represent tensors.
Stripes and checks show different values of model parameters in two clients.
Input-output pairs are colored in the same way as their activated execution
path in a model.

Personalization has been studied (Tan et al. 2022; Fallah, Mokhtari, and
Ozdaglar 2020; Arivazhagan et al. 2019) to incentivize statistically
heterogeneous clients through improved prediction performance. Typically, a
client trains the global model $\boldsymbol{w}^{g}$ and a personalized model
$\boldsymbol{w}^{p}$, sending $\boldsymbol{w}^{g}$ back to the server (Deng,
Kamani, and Mahdavi 2020; Liang et al. 2019; Collins et al. 2021; Pillutla et
al. 2022). Clients use the personalized model $\boldsymbol{w}^{p}$ instead of
the global model $\boldsymbol{w}^{g}$ for prediction. Recent literature
consistently demonstrate that personalized models can achieve higher
prediction performance than the global model aggregated over clients (Fallah,
Mokhtari, and Ozdaglar 2020; Deng, Kamani, and Mahdavi 2020; Liang et al.
2019; Collins et al. 2021; Pillutla et al. 2022; Li et al. 2020). Since these
approaches create a personalized model specific to each client’s data
distribution, we refer to them as per-client personalization.

Although a personalized model can achieve statistically higher accuracy than
the global model, there are still some input instances of a client that prefer
to use the global model $\boldsymbol{w}^{g}$ instead of the personalized model
$\boldsymbol{w}^{p}$. We use per-client approaches, including APFL (Deng,
Kamani, and Mahdavi 2020), Ditto (Li et al. 2021), LGFedAvg (Liang et al.
2019), FedRecon (Singhal et al. 2021), to train global and personalized models
with 100K number of clients on Stackoverflow (Federated 2022b) dataset. Our
experiments show that personalized models achieve 0.5% to 3.3% higher
prediction performance than a global model averaged over all clients. However,
the global model still outperforms the personalized model on 21% to 43% of
clients. Of the remaining clients whose personalized model outperforms the
global model, an average of 22.60% of the instances were correctly classified
by the global model, but incorrectly classified by the personalized model. The
observation reveals a major drawback of existing per-client personalization
approaches – that is, each client’s personalized model is a static network
that cannot accommodate the heterogeneity of the local client’s data
instances. Every data instance on a client is forced to use the identical
personalized model for prediction despite the fact that some instances could
benefit from the better generalizability of the global model.

In this paper, we propose a fine-grained personalized FL framework Flow that
leverages both per-instance and per-client personalization to further improve
the prediction performance of personalized models. The basic idea of Flow is
to create dynamic personalized models that are adaptive to not only each
client’s data distribution (per-client) but also each client’s individual data
instance (per-instance).

Figure 1 illustrates the difference between per-client and per-instance
personalization. Without personalization (Figure 1(a)), all the clients and
all the instances on a client execute the same model for prediction. With per-
instance personalization only (Figure 1(b)), all clients share the same global
model but each input instance could activate a different part of the model for
a prediction. An issue with this granularity is that the same input on
different clients would always produce the same output instead of a client-
specific prediction. Per-client personalization (Figure 1(c)) has distinct
personalized models for each client but all the instances on the same client
use the same model. Our proposed per-client per-instance personalization
(Figure 1(d)) allows instances on a client to flexibly choose which part of
the personalized model to activate depending on a client-specific instance.

The major challenge of per-client per-instance personalization is to design a
dynamic personalized model that is capable to characterize client-specific
data distribution while preserving the generalizability of the global model.
We address this challenge using the Dynamic Routing mechanism for both the
sequential natual language-based tasks (Liu and Deng 2018) and non-sequential
image-related tasks. In each FL round, a client first finetunes the global
parameters to generate local parameters. Then the client trains a dynamic
personalized model that integrates the global parameters, local parameters,
and a routing policy. The routing policy will determine whether to use the
local parameters or its global counterpart for each input instance.

Flow has two major advantages over many existing personalization approaches.
First, Flow leverages instance-level personalization to reconcile the
conflicting goal between the clients and the server: an individual client
wants improved prediction accuracy on its own dataset while the server wants
improved generalizability of the global model (measured usually by averaged
prediction accuracy across many clients). The dynamic personalized model of
Flow can achieve a better prediction performance on each client and the global
model offers a good starting point for any new client to personalize from.
Second, Flow is designed to be stateless because a client does not have to
carry its personalized model across FL rounds (Wu et al. 2022; Singhal et al.
2021). Being stateless brings two benefits. (1) Clients who don’t participate
in training frequently won’t suffer from stale states and thus poor accuracy.
Being stateless makes it easy to scale FL training to more than thousands of
clients. (2) Newly joined clients who did not participate in training before
could still pull the global model, personalize it, and perform inference.

Our evaluation on cross-device language and vision tasks shows that Flow
outperforms both per-client stateful and stateless personalization approaches.
For the newly joined clients, Flow achieves 0.90%, 2.38%, 1.26%
(Stackoverflow), 0.47%, 0.77%, 0.51% (Reddit) and 1.08%, 1.54%, 2.67% (EMNIST)
better accuracy against its best performing stateless, stateful, and a server-
side adaptive optimization baselines respectively.

Overall, this paper makes the following contributions:

  * •

It proposes a per-client per-instance personalization framework to improve
clients’ accuracy in FL. It is the first work on fine-grained personalization.

  * •

It uses dynamic routing to achieve the finer-grained personalization, which
strikes a balance between performance of the personalized model and the
generalizability of the global model.

  * •

Our evaluation on language and vision tasks using dynamic routing shows lower
instance misclassification rate for Flow over state-of-the-art personalization
and general FL approaches. Compared to stateful personalization approaches, we
consume fewer system resources.

##  2 Related Work

##### Per-Client Personalization

Personalization in Federated Learning (FL) has been explored for both vision
and language tasks, on the client-level granularity. APFL (Deng, Kamani, and
Mahdavi 2020) interpolates a client-specific personalized model state with the
global model sent by the server. Regularization is another popular way of
creating a personalized model. It encourages a personalized model of each
client to be close to the global model (Li et al. 2021; Dinh, Tran, and Nguyen
2020). Fine-tuning partial or entire global model to get the personalized
model has been studied in (Collins et al. 2021; Pillutla et al. 2022; Yu,
Bagdasaryan, and Shmatikov 2020). Multi-task learning has been utilized in
(Yu, Bagdasaryan, and Shmatikov 2020; Marfoq et al. 2021; Smith et al. 2017;
Dinh et al. 2021). Knowledge Distillation that enforces the similarity between
logits of the global and the personalized model is investigated in (Yu,
Bagdasaryan, and Shmatikov 2020; Zhang et al. 2021). Learning personalized
representations of the global model has been explored in (Liang et al. 2019),
and (Shamsian et al. 2021) trains a central hypernetwork model which is
trained to generate a personalized model for each client. While Ditto (Li et
al. 2021), FedRep (Collins et al. 2021), Local Adaptation (Yu, Bagdasaryan,
and Shmatikov 2020) have experiments on language tasks, most of the works we
discussed above focus more on vision tasks. Our work gives a specific design
for language tasks, along with a more general design for feed-forward
networks. Besides, none of the above work has explored instance-level
personalization. Since Flow is a client-side personalization approach. It can
work with any other server-side optimization methods like FedAvg and FedYogi.
For other personalization methods, composability would depend on whether a
method can derive a local model which can be incorporated in the dynamic
personalized model description provided in section C.

##### Stateless Personalization

Personalization approaches fall into two categories: stateless and stateful.
In stateful personalization, clients need to carry personalized states across
FL rounds (Deng, Kamani, and Mahdavi 2020; Collins et al. 2021; Li et al.
2021; Karimireddy et al. 2020). It is not practical for large-scale FL since a
client cannot participate in every FL round to keep their personalization
states from getting stale. This motivates several recent works to design
stateless personalization approaches such as FedRecon (Singhal et al. 2021),
and LGFedAvg (Liang et al. 2019). Specifically, FedRecon (Singhal et al. 2021)
creates personalized models by personalizing out-of-vocabulary tokens for each
client from the global model. LGFedAvg (Liang et al. 2019) divides a model
into two halves, where the feature representation half is unique for each
client, and the classifier half is shared among all the clients. Flow is also
designed to be stateless but features per-instance personalization to improve
clients’ accuracy.

##### Dynamic Routing

Motivated by saving computational effort for “easy to predict” instances,
instance-level dynamic routing has been a matter discussion in works related
to early exiting (Teerapittayanon, McDanel, and Kung 2017; Bolukbasi et al.
2017; Huang et al. 2018), layer skipping (Wang et al. 2018), and multi-
branching (Dai, Kong, and Guo 2020; Liu and Deng 2018). SkimRNN (Seo et al.
2018) explored temporal dynamic routing where depending on the input instance,
a trained policy can determine whether to skim through the instance with a
smaller RNN layer (which just updates the hidden states partially), or read
the instance with a standard (bigger) RNN layer (updating the entire hidden
state). Our work is motivated by the question of which route to pick,
depending on its utility?, which can be achieved by using a routing policy to
dynamically pick from two versions of encoders, which are equivalent in terms
of computational cost but different in terms of the data they are trained on.

##  3 Flow

We introduce Flow, a per-client per-instance stateless personalization method
that dynamically determines when to use a client’s local parameters and when
to use the global parameters based on the input instance. Flow works as
follows (shown in Algorithm 1).

During each FL round, upon receiving the global model $\boldsymbol{w}^{g}$,
each participating client personalizes and trains $\boldsymbol{w}^{g}$ in
three stages (shown in the procedure ClientTrain): (1) Split the train dataset
into two halves: $\mathcal{D}_{\ell}$ and $\mathcal{D}_{g}$, (2) derive the
local parameters $\boldsymbol{w}^{\ell}$ by finetuning the global parameters
$\boldsymbol{w}^{g}_{\mathsf{enc}}$ for $e^{\ell}$ epochs with
$\mathcal{D}_{\ell}$ (Lines 16 and 20), (3) construct a dynamic personalized
model by integrating the global parameters $\boldsymbol{w}^{g}$, the local
parameters $\boldsymbol{w}^{\ell}$, and randomly initialized dynamic routing
policy parameters $\theta$ (Line 17); train the routing policy $\theta$ and
$\boldsymbol{w}^{g}$ alternatively for $e_{g}$ epochs (Lines 23-24) with
$\mathcal{D}_{g}$, and (4) send the global model parameters
$\boldsymbol{w}^{g}$ back to the server for aggregation. The routing policy
$\theta$ determines whether the execution path of an instance uses
$\boldsymbol{w}^{\ell}$ or $\boldsymbol{w}^{g}$. The local parameters
$\boldsymbol{w}^{\ell}$ can be discarded since they will be recomputed
whenever a client next participates, in order to incorporate the latest global
model parameters.

The server aggregates all the clients’ updated global model parameters (shown
in the procedure ServerAggregate). Aggregation strategies are orthogonal to
this work and we adopt FedAvg (McMahan et al. 2016).

Figure 2: Illustration of the recurrent dynamic personalized model in Flow. A
general design for feed-foward architectures is shown in Appendix A.
Algorithm 1 Server Aggregation and Client Training in Flow

1:Global model $\boldsymbol{w}^{g}$, Routing policy $\theta$, Number of
clients per round $k$, Number of total rounds $r$, Train dataset of $c^{th}$
client $\mathcal{D}_{c}$, Epochs for the local model training $e_{\ell}$,
Epochs for the routing policy and the global model training $e_{g}$, Learning
rate for the local model $\eta_{\ell}$, Learning rate for the global model and
routing policy $\eta_{g}$

2:Global model $\boldsymbol{w}^{g}$

3:procedure ServerAggregate($k$, $r$)

4: $\boldsymbol{w}^{g}\leftarrow\text{Init
}\\{\boldsymbol{w}^{g}_{\mathsf{enc}},\boldsymbol{w}^{g}_{\mathsf{dec}}\\}$

5: for each round $i\in[1,r]$ do

6: $\mathcal{C}_{i}\leftarrow$ pick $k$ clients randomly

7: for each client $c\in\mathcal{C}_{i}$ do

8: $n^{c}_{i},\boldsymbol{w}^{c}_{i}\leftarrow$
ClientTrain($c,\boldsymbol{w}^{g}$)

9: end for

10: $n_{i}\leftarrow\sum_{c\in\mathcal{C}_{i}}n^{c}_{i}$

11:
$\boldsymbol{w}^{g}\leftarrow\frac{1}{n_{i}}\sum_{c\in\mathcal{C}_{i}}n^{c}_{i}\boldsymbol{w}^{c}_{i}$
//FedAvg

12: end for

13:end procedure

14:procedure ClientTrain($c,\boldsymbol{w}^{g}$)  15: initialize $\theta$
randomly  16:
$\boldsymbol{w}^{\ell}\leftarrow\\{\boldsymbol{w}^{g}_{\mathsf{enc}}\\}$
//local model 17:
$\boldsymbol{w}^{p}\leftarrow\\{\boldsymbol{w}^{g},\boldsymbol{w}^{\ell},\theta\\}$
//dynamic pers. model 18: $\mathcal{D}_{\ell},\mathcal{D}_{g}\leftarrow$
SplitDataset$(\mathcal{D}_{c})$ 19: for local epoch $e_{i}\in[1,e_{\ell}]$ do
20:
$\boldsymbol{w}^{\ell}\leftarrow\boldsymbol{w}^{\ell}-\eta_{\ell}\nabla_{\boldsymbol{w}^{\ell}}f(\boldsymbol{w}^{p},\mathcal{D}_{\ell})$
21: end for 22: for global epoch $e_{i}\in[1,e_{g}]$ do 23:
$\theta\leftarrow\theta-\eta_{g}\nabla_{\theta}f(\boldsymbol{w}^{p},\mathcal{D}_{g})$
24:
$\boldsymbol{w}^{g}\leftarrow\boldsymbol{w}^{g}-\eta_{g}\nabla_{\boldsymbol{w}^{g}}f(\boldsymbol{w}^{p},\mathcal{D}_{g})$
25: end for 26: return $|\mathcal{D}_{c}|,\boldsymbol{w}^{g}$ 27:end procedure

Next we discuss the design of the dynamic personalized model in detail
(recurrent design is illustrated in Figure 2). For recurrent networks, because
an instance is not fully available at the beginning, the dynamic routing needs
to take the gradually available tokens as input for making routing decisions.
The routing policy component takes as input both the hidden state from the
last timestep and the token available at the current timestep to make a
routing decision.

### Stage 1: Per-Client Personalization

Per-client personalization aims to derive local parameters that capture
client-specific patterns on each client’s local data. These local parameters
will be used to initialized a dynamic personalized model in the second stage.
Flow derives local parameters at the beginning of each FL round and discarding
them if they are no longer needed for inference in that round. This ensures
that Flow is stateless.

Flow derives local parameters
$\boldsymbol{w}^{\ell}\leftarrow\\{\boldsymbol{w}^{g}_{\mathsf{enc}}\\}$ by
finetuning only the encoder parameters $\boldsymbol{w}^{g}_{\mathsf{enc}}$ on
each client’s local data. The rest of the global model parameters is kept
frozen. The rationale is to finetune relatively small amount of parameters to
reduce training time. Previous work on per-client personalization (Liang et
al. 2019; Collins et al. 2021) shows that personalizing only encoder can
already improve the performance of the personalized model. Specifically, if
the architecture is recurrent, we select the RNN Cell (or its variants LSTM,
GRU) to finetune. The “encoder”, $\boldsymbol{w}^{g}_{\mathsf{enc}}$ refers to
the RNN layers and “decoder”, $\boldsymbol{w}^{g}_{\mathsf{dec}}$ refers to
the fully-connected layers that take hidden states as input. For convolution
architectures, we select the feature extractors to finetune.
$\boldsymbol{w}^{g}_{\mathsf{enc}}$ refers to feature extractors, while
$\boldsymbol{w}^{g}_{\mathsf{dec}}$ is the classifier head containing fully
connected layers.

### Stage 2: Per-Instance Personalization

Per-instance personalization allows an instance to flexibly select whether to
use global parameters or local parameters to maximize prediction performance.
Thus, a dynamic personalized model includes both local parameters
$\boldsymbol{w}^{\ell}$, the global parameters $\boldsymbol{w}^{g}$, and also
a routing policy component that determines the forward pass branch for an
input.

The challenge in training such a dynamic personalized model on each client is
to reconcile the conflicting goal of an individual client and the server. Some
clients would prefer to always use local parameters since local parameters are
already fine-tuned locally and could yield better accuracy statistically. On
the contrary, the server would prefer participating clients to collaboratively
update the global model parameters to improve its generalizability.

We use a two-pronged approach to address the issue. On one hand, we adopt the
temperature softmax function (Hinton et al. 2015) in the routing policy
component to allow soft decisions on which route to take, ensuring that global
parameters always get involved in the training. On the other hand, we add to
the loss function a regularization term that encourages instances to select
global parameters. Below we discuss the design of dynamic routing for
recurrent architectures and non-recurrent ones separately because the
availability of an input instance are different for the two architectures.

##### Dynamic Routing for Recurrent Networks.

An input instance for sequential tasks such as next word prediction is only
available gradually. So a new routing decision has to be made in every
timestep when a new token in the input instance is available, in order to
predict the next token.

Mathematically, our routing policy is parameterized with $\theta$, where
$\theta_{\boldsymbol{w}}\in\mathbb{R}^{2\times n}$ are weights and
$\theta_{\mathbf{b}}\in\mathbb{R}^{2}$ are biases. Let an input be
$\mathbf{x}=\mathbf{x}_{1},\cdots,\mathbf{x}_{T}$, where each
$\mathbf{x}_{t}\in\mathbb{R}^{m}$ represents a token available at timestep $t$
and $m$ is the embedding dimension. Let previous timestep’s hidden state be
$\mathbf{h}_{t-1}\in\mathbb{R}^{d}$, where $d$ is the hidden state dimension.

The routing policy has to pick one of the two routes
($\boldsymbol{w}^{\ell}_{\mathsf{enc}}$ and
$\boldsymbol{w}^{g}_{\mathsf{enc}}$). The probability distribution of picking
these routes for the recurrent networks is,

$\mathbf{r}_{t}=\text{softmax}_{\tau}(\theta_{\boldsymbol{w}}[\mathbf{x}_{t};\mathbf{h}_{t-1}]+\theta_{\mathbf{b}})\in\mathbb{R}^{2},$
(1)

where $\tau$ is the temperature parameter of the softmax function. Lower
$\tau$ (below 1) means the policy is more confident in its choice between
$\boldsymbol{w}^{\ell}_{\mathsf{enc}}$ and
$\boldsymbol{w}^{g}_{\mathsf{enc}}$, higher $\tau$ (above 1) means the policy
is less confident in its choice. We tune $\tau$ as a hyperparamter.

The hidden state after the timestep $t$ is the weighted average of the output
from the encoder that uses global parameters and the encoder that uses local
parameters:

$\mathbf{h}_{t}=r^{1}_{t}\times{f_{\boldsymbol{w}^{g}_{\mathsf{enc}}}}(\mathbf{h}_{t-1},\mathbf{x}_{t})+r^{2}_{t}\times{f_{\boldsymbol{w}^{\ell}_{\mathsf{enc}}}}(\mathbf{h}_{t-1},\mathbf{x}_{t}),$
(2)

where $f_{\boldsymbol{w}^{g}_{\mathsf{enc}}}$ and
$f_{\boldsymbol{w}^{\ell}_{\mathsf{enc}}}$ are global and local encoders
respectively and $\mathbf{r}_{t}=[r^{1}_{t},r^{2}_{t}]$. The objectives of
training the routing policy parameters $\theta$, and the global parameters
$\boldsymbol{w}^{g}$ are:

$\begin{split}min_{\\{\theta{},\boldsymbol{w}^{g}\\}}\frac{1}{NT}\sum_{i=1}^{N}\sum_{j=1}^{T}\left[\mathcal{L}(f_{\boldsymbol{w}^{p}}(x_{ij}),y_{ij})-\gamma\log(r^{1}_{ij})\right],\end{split}$
(3)

where $N$ is the number of instances of a particular client, and $T$ is the
maximum sequence length, and $\gamma$ is a regularizing hyperparameter. Note
that $i$ indexes an instance, and $j$ indexes a token at timestep $j$ of that
instance.

##### Dynamic Routing for Non-Recurrent Networks.

As an input instance for non-recurrent networks is typically available before
the execution of the models, the dynamic routing can be simplifed to be a one-
time decision making process. Specifically, for the routing policy parameters
$\theta_{\boldsymbol{w}}\in\mathbb{R}^{2\times n}$, $n$ is the dimension of
the flattened input vector, $\mathbf{x}\in\mathbb{R}^{n}$. The global and
local route probabilities will be calculated as follows,

$\mathbf{r}=\text{softmax}_{\tau}(\theta_{\boldsymbol{w}}[\mathbf{x}]+\theta_{\mathbf{b}})\in\mathbb{R}^{q}.$
(4)

The encoder outputs are the weighted average:

$\mathbf{x}_{\text{intermediate}}=r^{1}\times{f_{\boldsymbol{w}^{g}_{\mathsf{enc}}}}(\mathbf{x})+r^{2}\times{f_{\boldsymbol{w}^{\ell}_{\mathsf{enc}}}}(\mathbf{x}).$
(5)

The intermediate input $\mathbf{x}_{\text{intermediate}}$ will be fed to the
decoder $\boldsymbol{w}^{g}_{\mathsf{dec}}$ to generate the predicted outputs
for each instance. The routing policy and global objectives are similar to
Equation 3 with $T$ being always 1, as there’s no concept of sequence length
for non-sequential tasks.

### Pre-Inference Processing and Inference

A client has to create its dynamic personalized model before inference. It is
achieved by procedure ClientPreInference which is identical to ClientTrain in
Algorithm 1, except that we drop the $\boldsymbol{w}^{g}$ update in Line 24.
It trains the local parameters $\boldsymbol{w}^{\ell}$ and the routing policy
$\theta$ following similar steps as in training, with its training dataset.
Specifically, a client (1) pulls the latest global model $\boldsymbol{w}^{g}$,
(2) initializes and trains
$\boldsymbol{w}^{\ell}\leftarrow\\{\boldsymbol{w}^{g}_{\mathsf{enc}}\\}$ with
$\mathcal{D}_{\ell}$ which is a split of the training dataset, and (3)
initializes and trains the routing policy $\theta$ based on
$\boldsymbol{w}^{p}\leftarrow\\{\boldsymbol{w}^{g},\boldsymbol{w}^{\ell},\theta\\}$
with $\mathcal{D}_{g}$, the other split.

After pre-inference processing, a client uses $\boldsymbol{w}^{p}$ for
inference. To make inference as efficient as the per-client personalization
algorithms, the dynamic personalized model will make “hard” decisions during
inference – only the updates from one of the local and global routes are used
for an input instance of a timestep.

For recurrent networks, the hidden state update for timestep $t$ would be as
follows,

$\mathbf{h}_{t}=\begin{cases}{f_{\boldsymbol{w}^{g}_{\mathsf{enc}}}}(\mathbf{h}_{t-1},\mathbf{x}_{t})&\text{if
}r^{1}_{t}>0.5,\\\
{f_{\boldsymbol{w}^{\ell}_{\mathsf{enc}}}}(\mathbf{h}_{t-1},\mathbf{x}_{t})&\text{otherwise.}\end{cases}$
(6)

For non-recurrent networks, the intermediate input update for would be as
follows,

$\mathbf{x}_{\text{intermediate}}=\begin{cases}{f_{\boldsymbol{w}^{g}_{\mathsf{enc}}}}(\mathbf{x})&\text{if
}r^{1}>0.5,\\\
{f_{\boldsymbol{w}^{\ell}_{\mathsf{enc}}}}(\mathbf{x})&\text{otherwise.}\end{cases}$
(7)

In Flow, each client only needs to send (fetch) the global model parameters
back to (from) the server. The amount of communication is
$O(\boldsymbol{w}^{g})$. There’s no additional communication overhead compared
to FedAvg (McMahan et al. 2017). More details on resource usage are in
Appendix B.

##  4 Evaluations

We validate Flow by analyzing its empirical performance against various per-
client personalization approaches, both stateful and stateless, as well as
general FL approaches without client-side personalization.

### Experiment Settings

##### Models and Datasets

We target two next-word prediction tasks using datasets, Stackoverflow
(Federated 2022b) and Reddit (BigQuery 2022), and one image classification
task using EMNIST (Federated 2022a). Details on dataset heterogeneity are
available in Appendix C. For Stackoverflow, we use the same LSTM model (1
layer LSTM followed by 2 fully connected layers) as in (Reddi et al. 2021).
For Reddit, we use the LSTM model (2 stacked LSTM layers followed by a fully
connected layer) as in (McMahan et al. 2018). For EMNIST, we used the
recommended 2 layer CNN model in (Reddi et al. 2021). The detailed
hyperparameter settings are available in Appendix C.

##### Baselines

We compared Flow against the following baselines: three no personalization
approaches (FedAvg (McMahan et al. 2016), FedProx (Li et al. 2020), and
FedYogi (Reddi et al. 2021)), two stateful per-client personalization
approaches (APFL (Deng, Kamani, and Mahdavi 2020) and Ditto (Li et al. 2021)),
two stateless per-client personalization approaches (LGFedAvg (Liang et al.
2019), and FedRecon (Singhal et al. 2021)). Descriptions of each of these
baselines are given in Section 2. We also compare Flow with Dynamic Routing
(DR) where the global model itself is a dynamic network and clients finetune
it to create dynamic personalized models. The goal of this DR baseline is to
show whether Flow has benefits over a simple finetuning of
$\boldsymbol{w}^{g}$. Note that APFL interpolates an entire global model with
a personalized model by a client-specific interpolation factor $\alpha$, which
is similar to our $\mathbf{r}$ with the exception that $\mathbf{r}$ is
instance-specific, while $\alpha$ is the same for all instances.

All the baselines and Flow use the same train/validation/test data split. For
the train split, only FedRecon and Flow further separate the “train split”
into two sets: (1) $\mathcal{D}_{\ell}$, which will be the former half of all
the input sequences, (2) $\mathcal{D}_{g}$, which will be the latter half of
the same input sequences. This is in line with how datasets are split in
(Singhal et al. 2021). Vertical splits like this ensures that all the input
sequences contribute to both local and global training. Since there’s no
notion of sequentiality in image classification tasks, we split the training
dataset in two equal halves to derive $\mathcal{D}_{\ell}$ and
$\mathcal{D}_{g}$. Flow uses $\mathcal{D}_{l}$ to train local models and
$\mathcal{D}_{g}$ to train the routing policy and global model. FedRecon uses
$\mathcal{D}_{l}$ to reconstruct some local parameters and $\mathcal{D}_{g}$
to train global models. The validation and test dataset splits are the exact
same for all the baselines.

##### Evaluation Metrics

Prediction performance is evaluated by both validation accuracy and test
accuracy. Validation accuracy reported in Figure 3 is calculated by averaging
per-round accuracies over the range of 100 rounds. Test accuracy, as shown in
Table 1, is calculated at the end of the federated training for additional 100
rounds, picking test clients at random, and then averaging accuracies of all
100 rounds. For both validation and testing, we have used 100 clients for
Stackoverflow, 10 clients for both Reddit, and EMNIST (this count is the same
for training client count per round). The clients are randomly picked from a
cohort of 204,088 (Stackoverflow), 157,572 (Reddit), and 3400 (EMNIST) clients
available for testing.

Since many wouldn’t have participated once or more than once during the
training, we consider those as newly joined clients during inference. Both for
validation and test accuracies, clients first have to derive
$\boldsymbol{w}^{p}$ based on ClientPreInference procedure, then use that
$\boldsymbol{w}^{p}$ for inference. Since the inference is on
$\boldsymbol{w}^{p}$, we are mainly reporting the personalized accuracies for
our baseline approaches, Flow , and its variants discussed in the baselines
section above.

We implement Flow and its baselines in the Flower framework (Beutel et al.
2020). All the experiments have been repeated 3 times under different random
seeds, the results presented in the next section are the average of the 3 runs
for each approach and each experiment. Appendix C describes other
implementation details.

### Results and Analysis

(a) Stackoverflow

(b) Reddit

(c) EMNIST

Figure 3: Validation accuracy versus number of federated rounds of Flow and its baselines. Methods |  Stack-Overflow | Reddit | EMNIST
---|---|---|---
FedAvg ♢ |  24.98 $\pm$ 0.04% |  17.30$\pm$0.13% |  86.71$\pm$0.19%
FedYogi ♢ |  27.90$\pm$0.05% |  18.85$\pm$0.10% |  84.12$\pm$0.18%
FedProx ♢ |  25.38$\pm$0.04% |  17.58$\pm$0.14% |  89.01$\pm$0.10%
APFL ${}^{\dagger\;\S}$ |  25.48$\pm$0.07% |  18.48$\pm$0.08% |  90.14$\pm$0.14%
Ditto ${}^{\dagger\;\S}$ |  25.55$\pm$0.09% |  18.74$\pm$0.12% |  87.17$\pm$0.18%
LG-FedAvg ${}^{\ddagger\;\S}$ |  26.78$\pm$0.16% |  18.59$\pm$0.18% |  90.60$\pm$0.24%
FedRecon ${}^{\star\;\ddagger\;\S}$ |  28.26$\pm$0.20% |  18.89$\pm$0.23% | N/A
DR ${}^{\ddagger\;\S\;\flat}$ |  26.54$\pm 0.15$% |  18.70$\pm$0.16% |  89.78$\pm$0.20%
Flow ${}^{\ddagger\;\S\;\flat}$ |  29.16$\pm$ 0.21% |  19.36$\pm$ 0.19% |  91.68$\pm$ 0.25%
Table 1: Test accuracies for Stackoverflow, Reddit, and EMNIST (FT =
Finetuning, DR = Dynamic Routing, $\diamondsuit$ = Server-side Optimization,
$\dagger$ = Stateful Personalization, $\ddagger$ = Stateless Personalization,
$\star$ = with 500 OOV (Out of Vocabulary tokens), $\S$ = Per-Client
Personalization, $\flat$ = Per-Instance Personalization.)

##### Prediction Performance

Figure 3 shows the validation accuracy of Flow and its baselines on
Stackoverflow, Reddit, and EMNIST respectively. Table 1 reports the test
accuracies. We make the following observations. (1) Dynamic Routing (DR)
outperforms per-client personalization APFL, Ditto, and LGFedAvg, proving that
instances of the same client do prefer different parts of the model. (2) As a
stateless personalization approach, FedRecon (using FedYogi as aggregation)
outperforms stateful approaches since the latter require training of
personalized models over multiple rounds, which are impractical when only
$\sim$5-6% of total clients would participate during FL training.

Overall, Flow outperforms all the baselines. Flow outperforms FedRecon, the
best of all the baselines, by 0.90% on Stackoverflow and 0.47% on Reddit in
terms of test accuracy. It is because Flow is finer-grained and dynamic in its
choice of client-specific local parameters and global parameters. The results
demonstrate that per-instance personalization further improves the performance
of a per-client personalized model.

| Stackoverflow | Reddit | EMNIST
---|---|---|---
GR |  25.35$\pm$0.06% |  17.62$\pm$0.12% |  87.34$\pm$0.22%
LR |  28.56$\pm$0.16% |  19.01$\pm$0.16% |  89.94$\pm$0.27%
Flow |  29.16$\pm$0.21% |  19.36$\pm$0.19% |  91.68$\pm$0.25%
Table 2: Test accuracies of Flow and its two variants where all the clients
either just take the global or the local route. GR = Global Route only, LR =
Local Route only.

##### Importance of Per-Instance Personalization

In order to show the efficacy of dynamic routing approached used by Flow, we
compare Flow with two of its variants as follows: (a) Global route only: a
modification of Flow where $\mathbf{r}$ ($\mathbf{r}_{t}\;\forall t\in[1,T]$
in case of recurrent networks) is always 1.0 for all the instances, (b) Local
route only: a modification of Flow where $\mathbf{r}$
($\mathbf{r}_{t}\;\forall t\in[1,T]$ in case of recurrent networks) is always
0.0 for all the instances. The first variant indicates that the policy in
$\boldsymbol{w}^{p}$ is forced to take the global route for all the clients
and their instances, while the second baseline indicates that the policy in
$\boldsymbol{w}^{p}$ would always take the local route regardless of the
clients and their instances.

Test accuracy results in Table 2 shows that per-instance routing is superior
to routing designs which are not based on instances. Besides, the “local route
only” variant outperforms the “global route only”, which echoes the intuition
that finetuning a global model can increase clients’ accuracy. Note that the
test accuracies reported in Table 2 are based on evaluation done after the
global model is trained according to Algorithm 1. The increase in “local route
only“ variant displays better generalizability of the global model as compared
to rest of the baselines reported in Table 1.

Table 3 shows the reduction in the number of instances which are misclassified
by the personalized model, but correctly classified by the local model. This
indicates that using the global model for a subset of instances would be
indeed more beneficial.

##### Analysis on Per-Instance Routing Decisions

| Stackoverflow | Reddit | EMNIST
---|---|---|---
| C (%) | (I %) | C (%) | (I %) | C (%) | (I %)
APFL | 43.83 | 23.13 | 38.20 | 13.05 | 64.21 | 5.75
Ditto | 45.83 | 22.58 | 41.30 | 13.69 | 49.47 | 7.24
LGFedAvg | 61.73 | 23.10 | 64.61 | 14.35 | 68.42 | 4.86
FedRecon | 69.15 | 21.98 | 65.75 | 13.38 | N/A | N/A
Flow | 73.74 | 20.45 | 67.60 | 12.31 | 72.63 | 3.20
Table 3: For Flow and its baselines for each dataset - (i) C (%) is percent of
clients preferring the personalized model over the local model, (ii) I (%) is
percent of instances that are correctly classified by the local model, but not
by the personalized model (lower is better).

The analysis aims to show the importance of per-instance personalization by
examining the behavior of routing decisions ($\mathbf{r}$). The behavior of
routing decisions is closely related to the notion of heterogeneity of a
certain client. To understand it, we need a measure of the level of
heterogeneity within a client, with respect to the global data distribution.
We use Kullback-Leibler (KL) divergence as the metric. Definitions are in
Appendix E.

For RNNs, we analyze the behavior of $\mathbf{r}_{t}$ with respect to
timesteps $t$. We picked 100 test clients with their KL divergence having mean
of 5.67, median of 3.00, and the range of 0.05 to 24.63. Figure 4 shows the
average and the standard deviation of the local route probabilities taken by
all the instances of two clients on the RNN using stackoverflow. Lower
probability of local route means that a client prefers the global route, and
higher probability of local route means that a client prefers the local route.
Client #1 has relatively lower KL divergence value, indicating that most of
its instances perform well just with the global model. But there would still
be few instances (in this case, there are 63 such instances) who would have
preferred their local models. An opposite behavior is shown for Client #2
where, due to higher divergence, more instances (390 instances) prefer the
local model. But again, some instances (83 instances) which fit well within
the global data distribution would prefer the global model and hence take the
global route in our dynamic personalized model.

For the first few time-steps, the policy is less confident about its choice
between routes as it has only seen a few tokens. As it sees more tokens, its
confidence about picking a certain route increases. Figure 4 shows similar
behavior for both the lower and higher KL divergence clients, it shows that
the routing policy is independent of client heterogeneity.

(a) Client #1 with KLD 1.23

(b) Client #2 with KLD 7.11

Figure 4: Case study of routing decisions made across timesteps for
Stackoverflow.

We observe similar behavior of routing decisions on CNNs using EMNIST. Figure
5 shows the distribution of routing probabilities for two clients with
distinctive KL divergence. More instances on client #1 prefers global routes
than client #2 because client #1’s data distribution is closer to the global
distribution than client #2. However, there are still decent amount of
instances on client #1 that would choose the local routes since the local
routes fit these instances better.

(a) Client #1 with KLD 1.86
#Instances routed locally = 29
#Instances routed globally = 110

(b) Client #2 with KLD 6.14
#Instances routed locally = 141
#Instances routed globally = 59

Figure 5: Case study of routing probabilities for EMNIST.

##### Soft Decisions During Inference

We didn’t use soft decisions during inference for two reasons. First, it does
not improve the accuracy of Flow much. The test accuracies for Flow on
Stackoverflow and Reddit with hard decisions are 29.16% ± 0.21%, and 19.36% ±
0.19%, while with soft decisions, we observed 29.22% ± 0.24%, and 19.46% ±
0.15%. Second, using soft decisions means that both global routes and local
routes have to be executed. It doubles the computation and increases the
inference overhead.

##### Local Epoches and Regularization

We notice that for language tasks, lesser number of epochs are sufficient to
derive the dynamic personalized models, while image-related tasks require more
epochs, as discussed further in Appendix F. Appendix F also shows that adding
a regularization term improves the global model generalizability, which in
turn produces more accurate dynamic personalized models.

##  5 Conclusion

We presented Flow, a fine-grained personalized federated learning (FL)
approach to improve clients’ accuracy. For each client, Flow creates a dynamic
personalized model that leverages a dynamic routing mechanism to determine
whether to use local parameters or global parameters for inference depending
on the input instance. Our empirical evaluation shows that Flow achieves a
better personalized performance for newly joined clients.

## References

  * Arivazhagan et al. (2019) Arivazhagan, M. G.; Aggarwal, V.; Singh, A. K.; and Choudhary, S. 2019.  Federated learning with personalization layers.  _arXiv preprint arXiv:1912.00818_.
  * Beutel et al. (2020) Beutel, D. J.; Topal, T.; Mathur, A.; Qiu, X.; Parcollet, T.; and Lane, N. D. 2020\.  Flower: A Friendly Federated Learning Research Framework.  _arXiv preprint arXiv:2007.14390_.
  * BigQuery (2022) BigQuery, R. D. G. 2022.  Google BigQuery Reddit Dataset.  https://bigquery.cloud.google.com/dataset/fh-bigquery:reddit˙comments.  [Online; accessed 02-May-2022].
  * Bolukbasi et al. (2017) Bolukbasi, T.; Wang, J.; Dekel, O.; and Saligrama, V. 2017.  Adaptive Neural Networks for Efficient Inference.  In _Proceedings of the 34th International Conference on Machine Learning_.
  * Collins et al. (2021) Collins, L.; Hassani, H.; Mokhtari, A.; and Shakkottai, S. 2021.  Exploiting Shared Representations for Personalized Federated Learning.  _arXiv preprint 2102.07078_.
  * Dai, Kong, and Guo (2020) Dai, X.; Kong, X.; and Guo, T. 2020.  EPNet: Learning to Exit with Flexible Multi-Branch Network.  In _Proceedings of the 29th ACM International Conference on Information and Knowledge Management_. Association for Computing Machinery.
  * Deng, Kamani, and Mahdavi (2020) Deng, Y.; Kamani, M. M.; and Mahdavi, M. 2020.  Adaptive personalized federated learning.  _arXiv preprint 2003.13461_.
  * Dinh, Tran, and Nguyen (2020) Dinh, C. T.; Tran, N. H.; and Nguyen, T. D. 2020.  Personalized Federated Learning with Moreau Envelopes.  In _Proceedings of the 34th International Conference on Neural Information Processing Systems_.
  * Dinh et al. (2021) Dinh, C. T.; Vu, T. T.; Tran, N. H.; Dao, M. N.; and Zhang, H. 2021.  FedU: A Unified Framework for Federated Multi-Task Learning with Laplacian Regularization.  _arXiv pre-print 2102.07148_.
  * Fallah, Mokhtari, and Ozdaglar (2020) Fallah, A.; Mokhtari, A.; and Ozdaglar, A. 2020.  Personalized Federated Learning with Theoretical Guarantees: A Model-Agnostic Meta-Learning Approach.  In _Advances in Neural Information Processing Systems_.
  * Federated (2022a) Federated, E. D. T. 2022a.  TensorFlow Federated EMNIST Dataset.  https://www.tensorflow.org/federated/api˙docs/python/tff/simulation/datasets/emnist.  [Online; accessed 11-Sept-2022].
  * Federated (2022b) Federated, S. D. T. 2022b.  TensorFlow Federated Stackoverflow Datasets.  https://www.tensorflow.org/federated/api˙docs/python/tff/simulation/datasets.  [Online; accessed 02-May-2022].
  * FedML (2022) FedML, S. D. 2022.  FedML Federated Stackoverflow Heterogenity .  https://github.com/FedML-AI/FedML/blob/ecd2d81222301d315ca3a84be5a5ce4f33d6181c/doc/en/simulation/benchmark/BENCHMARK˙MPI.md#lightweight-and-shallow-neural-network-models.  [Online; accessed 14-Sep-2022].
  * Hinton et al. (2015) Hinton, G.; Vinyals, O.; Dean, J.; et al. 2015.  Distilling the knowledge in a neural network.  _Advances in Neural Information Processing Systems Deep Learning Workshop_.
  * Huang et al. (2018) Huang, G.; Chen, D.; Li, T.; Wu, F.; van der Maaten, L.; and Weinberger, K. 2018\.  Multi-Scale Dense Networks for Resource Efficient Image Classification.  In _International Conference on Learning Representations_.
  * Kairouz et al. (2021) Kairouz, P.; McMahan, H. B.; Avent, B.; Bellet, A.; Bennis, M.; Bhagoji, A. N.; Bonawitz, K. A.; Charles, Z.; Cormode, G.; Cummings, R.; D’Oliveira, R. G. L.; Eichner, H.; Rouayheb, S. E.; Evans, D.; Gardner, J.; Garrett, Z.; Gascón, A.; Ghazi, B.; Gibbons, P. B.; Gruteser, M.; Harchaoui, Z.; He, C.; He, L.; Huo, Z.; Hutchinson, B.; Hsu, J.; Jaggi, M.; Javidi, T.; Joshi, G.; Khodak, M.; Konecný, J.; Korolova, A.; Koushanfar, F.; Koyejo, S.; Lepoint, T.; Liu, Y.; Mittal, P.; Mohri, M.; Nock, R.; Özgür, A.; Pagh, R.; Qi, H.; Ramage, D.; Raskar, R.; Raykova, M.; Song, D.; Song, W.; Stich, S. U.; Sun, Z.; Suresh, A. T.; Tramèr, F.; Vepakomma, P.; Wang, J.; Xiong, L.; Xu, Z.; Yang, Q.; Yu, F. X.; Yu, H.; and Zhao, S. 2021.  Advances and Open Problems in Federated Learning.  _Foundations and Trends in Machine Learning_.
  * Karimireddy et al. (2020) Karimireddy, S. P.; Kale, S.; Mohri, M.; Reddi, S.; Stich, S.; and Suresh, A. T. 2020.  SCAFFOLD: Stochastic Controlled Averaging for Federated Learning.  In _Proceedings of the 37th International Conference on Machine Learning_.
  * Li et al. (2021) Li, T.; Hu, S.; Beirami, A.; and Smith, V. 2021.  Ditto: Fair and Robust Federated Learning Through Personalization.  _International Conference on Machine Learning_.
  * Li et al. (2020) Li, T.; Sahu, A. K.; Zaheer, M.; Sanjabi, M.; Talwalkar, A.; and Smith, V. 2020\.  Federated Optimization in Heterogeneous Networks.  In _Proceedings of Machine Learning and Systems_.
  * Liang et al. (2019) Liang, P. P.; Liu, T.; Ziyin, L.; Salakhutdinov, R.; and Morency, L.-P. 2019.  Think locally, act globally: Federated learning with local and global representations.  _Advances in Neural Information Processing Systems Workshop on Federated Learning_.
  * Liu and Deng (2018) Liu, L.; and Deng, J. 2018.  Dynamic Deep Neural Networks: Optimizing Accuracy-Efficiency Trade-offs by Selective Execution.  In _Association for the Advancement of Artificial Intelligence_.
  * Marfoq et al. (2021) Marfoq, O.; Neglia, G.; Bellet, A.; Kameni, L.; and Vidal, R. 2021.  Federated Multi-Task Learning under a Mixture of Distributions.  In _Advances in Neural Information Processing Systems_.
  * McMahan et al. (2017) McMahan, B.; Moore, E.; Ramage, D.; Hampson, S.; and Arcas, B. A. y. 2017.  Communication-Efficient Learning of Deep Networks from Decentralized Data.  In _Proceedings of the 20th International Conference on Artificial Intelligence and Statistics_.
  * McMahan et al. (2016) McMahan, H. B.; Moore, E.; Ramage, D.; and y Arcas, B. A. 2016.  Federated Learning of Deep Networks using Model Averaging.  _Proceedings of the 20th International Conference on Artificial Intelligence and Statistics_.
  * McMahan et al. (2018) McMahan, H. B.; Ramage, D.; Talwar, K.; and Zhang, L. 2018.  Learning Differentially Private Recurrent Language Models.  In _International Conference on Learning Representations_.
  * Pillutla et al. (2022) Pillutla, K.; Malik, K.; Mohamed, A.; Rabbat, M.; Sanjabi, M.; and Xiao, L. 2022\.  Federated Learning with Partial Model Personalization.  _arXiv preprint 2204.03809_.
  * Reddi et al. (2021) Reddi, S. J.; Charles, Z.; Zaheer, M.; Garrett, Z.; Rush, K.; Konečný, J.; Kumar, S.; and McMahan, H. B. 2021.  Adaptive Federated Optimization.  In _International Conference on Learning Representations_.
  * Seo et al. (2018) Seo, M.; Min, S.; Farhadi, A.; and Hajishirzi, H. 2018.  Neural Speed Reading via Skim-RNN.  In _International Conference on Learning Representations_.
  * Shamsian et al. (2021) Shamsian, A.; Navon, A.; Fetaya, E.; and Chechik, G. 2021.  Personalized Federated Learning using Hypernetworks.  In _Proceedings of the 38th International Conference on Machine Learning_.
  * Singhal et al. (2021) Singhal, K.; Sidahmed, H.; Garrett, Z.; Wu, S.; Rush, J. K.; and Prakash, S. 2021\.  Federated Reconstruction: Partially Local Federated Learning.  In _Advances in Neural Information Processing Systems_.
  * Smith et al. (2017) Smith, V.; Chiang, C.-K.; Sanjabi, M.; and Talwalkar, A. S. 2017.  Federated Multi-Task Learning.  In _Advances in Neural Information Processing Systems_.
  * Tan et al. (2022) Tan, A. Z.; Yu, H.; Cui, L.; and Yang, Q. 2022.  Towards Personalized Federated Learning.  _IEEE Transactions on Neural Networks and Learning Systems_.
  * Teerapittayanon, McDanel, and Kung (2017) Teerapittayanon, S.; McDanel, B.; and Kung, H. T. 2017.  BranchyNet: Fast Inference via Early Exiting from Deep Neural Networks.  _arXiv preprint 1709.01686_.
  * Wang et al. (2018) Wang, X.; Yu, F.; Dou, Z.-Y.; Darrell, T.; and Gonzalez, J. E. 2018.  SkipNet: Learning Dynamic Routing in Convolutional Networks.  In _The European Conference on Computer Vision_.
  * Wu et al. (2022) Wu, S.; Li, T.; Charles, Z.; Xiao, Y.; Liu, Z.; Xu, Z.; and Smith, V. 2022.  Motley: Benchmarking Heterogeneity and Personalization in Federated Learning.  _arXiv preprint arXiv:2206.09262_.
  * Yu, Bagdasaryan, and Shmatikov (2020) Yu, T.; Bagdasaryan, E.; and Shmatikov, V. 2020.  Salvaging Federated Learning by Local Adaptation.  _arXiv preprint 2002.04758_.
  * Zhang et al. (2021) Zhang, J.; Guo, S.; Ma, X.; Wang, H.; Xu, W.; and Wu, F. 2021.  Parameterized Knowledge Transfer for Personalized Federated Learning.  In _Advances in Neural Information Processing Systems_.

##  Appendix A Design of Flow for Feed Forward Networks

Figure 6 illustrates the dynamic personalized model in Flow for feed forward
architectures. Since the intermediate results of the layers of an encoder can
be too large for the routing network, we only use the initial input as an
input to the routing network. After incorporating the routing probabilities
with the outputs of both the global and local encoder, we derive an
intermediate output which can be fed into the shared decoder.

Figure 6: Illustration of a general dynamic personalized model in Flow for
feed-forward architectures.

##  Appendix B Resource Usage of Flow

We comment on the resource usages in communication, computation, and memory
requirements below.

##### Communication requirement:

In Flow, each client only needs to send (fetch) the global model parameters
back to (from) the server. Local parameters and dynamic personalized models
are derived locally on each client after it fetches the global model. There’s
no additional communication overhead compared to the standard FedAvg.

##### Computation requirement:

In one round of federated training, each client first derives the local
parameters by fine-tuning the global model for one epoch, and then trains the
dynamicpersonalized model. Compared to the standard FedAvg, the additional
computation overhead comes from two parts.

The first part is the finetuning of the global model to derive local
parameters. To reduce computation costs, Flow personalizes only the feature
extractors (Convolution layers or Recurrent cells) instead of the entire model
(which also consists of embeddings, and fully-connected layers), which makes
it more efficient compared to personalization approaches such as APFL and
Ditto. Our empirical evaluation shows that fine-tuning only the encoder is
1.5x and 1.12x faster compared to fine-tuning the entire global model, for RNN
used for Stackoverflow and CNN used for EMNIST respectively.

The second part is the training of additional policy modules. However, the
additional computation is a simple vector product following a softmax (see Eq.
2), which is often small compared to the other parts of the RNN model. The
FLOPS of the policy modules is only $308K$, 4.96% of that of the Stackoverflow
RNN model ($6.20M$), only $367.2K$, 5.14% of that of the Reddit RNN model
($7.14M$), and only $2.17M$, 0.886% of that of the EMNIST CNN model
($245.23M$).

##### Memory requirement:

Compared to FedAvg, the additional memory consumption is the local parameters
(encoder parameters) and the parameters from the policy module. Specifically,
the routing policy parameters counts are $|\theta|=153.8K$, $183.6K$, and
$217.54k$ for Stackoverflow, Reddit, and EMNIST respectively. The local
parameters counts are $2.06M$, $1.99M$, and $1.21M$ respectively. The total
additional parameters is only 49% (Stackoverflow), 35.73% (Reddit), and 45.88%
(EMNIST) of the global model parameters. However, personalization approaches
such as Ditto and APFL introduce double the extra parameters (100%), which is
more memory intensive than Flow.

##  Appendix C Datasets and Hyperparameters

### Dataset Splits

| Training | Validation | Test
---|---|---|---
| #Clients | #Total Instances | #Clients | #Total Instances | #Clients | #Total Instances
Stackoverflow | 342,477 | 135,818,730 | 38,758 | 16,491,230 | 204,088 | 16,586,035
Reddit | 157,572 | 56,725,920 | 157,572 | 7,090,740 | 157,572 | 7,090,740
EMNIST | 3400 | 671,585 | 3400 | 77,483 | 3400 | 77,483
Table 4: Train/Valid/Test datasets client and instance counts.

In Stackoverflow dataset, there’s an overlap between the train and test
clients. Since we need a training dataset for the pre-inference stage before
evaluating a personalized model’s performance, we use 195,303 clients which
have both the training and testing datasets. Reddit dataset consists of
157,572 unique users (November 2017 data from (BigQuery 2022)). Since it
didn’t have train/valid/test splits, we used 8:1:1 ratio for each of the
clients’ datasets. Similar with EMNIST, each client’s data is split in 8:1:1
ratio for train/valid/test datasets as well.

### Dataset Heterogeneity

Both Stackoverflow and Reddit datasets are naturally heterogeneous (FedML
2022). The two datasets have been used in the literature for studying data
heterogeneity (Wu et al. 2022; Yu, Bagdasaryan, and Shmatikov 2020). Each user
of Stackoverflow and Reddit is a client and his/her posts form a dataset for
the client. The dataset is heterogeneous in two ways: First, users have
different writing styles and thus clients’ datasets are not i.i.d. Second, the
total number of posts from each user is also different, leading to different
sizes of datasets per client.

The heterogeneity of EMNIST clients stems from the individual writing style of
each client (where a client is a person), as discussed in Appendix C.2 of
(Reddi et al. 2021).

### Flow Hyperparameters

#### Stackoverflow

Similar to AFO (Reddi et al. 2021), we used the sequence length $T=20$,
embedding size $m=96$, vocabulary size = 10,000, out of vocabulary bin count =
1, batch size $N=16$. We train on Stackoverflow for $r=2000$ rounds, with 100
clients per round. The local $\eta_{\ell}$ and global $\eta_{g}$ learning
rates are set to 0.3, we experimented on softmax temperature
$\tau\in\\{0.25,0.5,0.75,1.00,1.25\\}$ and local epochs
$e_{\ell}=\\{1,3,5\\}$. The combination of $\tau=0.75$ and $e_{\ell}=1$ worked
best for us. The regularization hyperparameter $\gamma$ is set to $10^{-2}$.

#### Reddit

We kept the Reddit vocabulary limited to 10,000 tokens also, out of vocabulary
bin count = 1, sequence length $T=10$, embedding size $m=200$, batch size
$N=5$. We train on Reddit for $r=1000$ rounds, with $10$ clients per round.
The local $\eta_{\ell}$ and global $\eta_{g}$ learning rates are set to 1. The
combination of $\tau=1.00$ and $e_{\ell}=1$ worked best for us. The
hyperparameter $\gamma$ is $10^{-2}$.

The regularization parameter $\gamma$ for both Stackoverflow and Reddit, we
tried $\\{0.1,0.01,0.001\\}$ and found that $\gamma=0.01$ gives the best
performance for both.

#### EMNIST

We trained each of the baselines for $r=1500$ rounds, with the batch size of
$N=20$ instances, 10 clients per round. We ran Flow for $e_{\ell}=5$ with
learning rate of $\eta_{\ell}=0.05$, regularization hyperparameter
$\gamma=0.05$, softmax temperature $\tau=0.75$.

### Baseline Hyperparameters

This sections talks about hyperparameters for all the baselines.

Starting with Stackoverflow, we tried local epochs $e_{\ell}\in[1,3,5,7]$ for
all the baselines, but due to each client having separate instances count and
different level of heterogeneity, $e_{\ell}=1$ works the best. All the
baselines pick 100 clients per round, from the same 195,303 clients available
to all of them. $\alpha\in[0,1]$ used for APFL (Deng, Kamani, and Mahdavi
2020) is based on each client for all the datasets. The default learning rate
is 0.3 (recommended in (Reddi et al. 2021)) for all but Ditto and LGFedAvg.
They require the learning rate to be 0.1. Flow and its baselines all train
Stackoverflow dataset for 2000 rounds. We use 10,000 vocab size.

For Reddit, Flow and its baselines pick 10 clients each round, for 1000
rounds. Once again, we see similar trend of detrimental accuracy upon
increasing epochs, so we keep local epochs for all the baselines to 1. We use
10,000 vocab size for Reddit as well. Learning rate of 1 was recommended in
(Yu, Bagdasaryan, and Shmatikov 2020), so keep that for all the baselines.
FedYogi has default $\beta_{1}=0.9$ and $\beta_{2}=0.999$ for all the
datasets.

In EMNIST, once again we use 10 clients per round, with 1500 rounds in total.
With the default learning rate of 0.05. Ditto uses 0.01 learning rate and 5
epochs. APFL uses 4 local epochs. The preprocessing methods employed here are
taken from (Reddi et al. 2021).

### Implementation

We implement Flow and its baselines in the Flower framework (Beutel et al.
2020). Experiments are conducted on two NVIDIA 2080TI GPUs, in stand-alone
simulation mode which mimics the real-world federated learning setup. Similar
to (Reddi et al. 2021), clients are randomly picked (because servers cannot
pick any particular client of its choice in the real-world deployment) without
replacement for each round and with replacement across each round.

##  Appendix D Pre-Inference

The pre-inference stage described in Section C is described in Algorithm 2.

Algorithm 2 Client Pre-Inference in Flow

1:Client ID $c$, Global model $\boldsymbol{w}^{g}$, Routing policy $\theta$,
Train dataset for $c^{th}$ client $\mathcal{D}_{c}$, Epochs for the local
model training $e_{\ell}$, Epochs for the routing policy training $e_{g}$,
Learning rate for the local model $\eta_{\ell}$, Learning rate for the routing
policy $\eta_{g}$

2:Dynamic Personalized model $\boldsymbol{w}^{p}$

3:procedure ClientPreInference($c,\boldsymbol{w}^{g}$)

4: $\boldsymbol{w}^{\ell}\leftarrow\\{\boldsymbol{w}^{g}_{\mathsf{enc}}\\}$
//local model

5:
$\boldsymbol{w}^{p}\leftarrow\\{\boldsymbol{w}^{g},\boldsymbol{w}^{\ell},\theta\\}$
//dynamic pers. model

6: $\mathcal{D}_{\ell},\mathcal{D}_{g}\leftarrow$
SplitDataset$(\mathcal{D}_{c})$

7: for local epoch $e_{i}\in[1,e_{\ell}]$ do

8:
$\boldsymbol{w}^{\ell}\leftarrow\boldsymbol{w}^{\ell}-\eta_{\ell}\nabla_{\boldsymbol{w}^{\ell}}f(\boldsymbol{w}^{p},\mathcal{D}_{\ell})$

9: end for

10: for global epoch $e_{i}\in[1,e_{g}]$ do

11:
$\theta\leftarrow\theta-\eta_{g}\nabla_{\theta}f(\boldsymbol{w}^{p},\mathcal{D}_{g})$

12: end for

13: return $\boldsymbol{w}^{p}$

14:end procedure

##  Appendix E Behavior of $\mathbf{r}_{t}$ for Sequential Tasks

The behavior of our routing policy ($\mathbf{r}_{t}$) is closely related to
the notion of heterogeneity of a certain client. To understand it, we need a
measure of the level of heterogeneity within a client, with respect to the
global data distribution. We use Kullback-Leibler (KL) divergence as the
metric in Eq. (6). In an ideal case where we have access to each client’s data
distribution, we could have compared a client’s local data distribution with
the distribution of all the data the global model has been trained on so far
(we call it global distribution). In the absence of the true local (and hence
global) distributions, we estimate those distributions at client side using
$\boldsymbol{w}^{\ell}$ and $\boldsymbol{w}^{g}$ as follows,

$\displaystyle D_{KL}(\mathcal{P}c\;||\;\mathcal{P}g)$
$\displaystyle=\sum_{(x,y)\in\mathcal{D}_{c}}\mathcal{P}c(x)\log\left(\frac{\mathcal{P}c(x)}{\mathcal{P}g(x)}\right)$
(8)
$\displaystyle\approx\sum_{(x,y)\in\mathcal{D}_{c}}\Pr[\boldsymbol{w}^{\ell}(x)=y]\log\left(\frac{\Pr[\boldsymbol{w}^{\ell}(x)=y]}{\Pr[\boldsymbol{w}^{g}(x)=y]}\right).$
(9)

Now we will discuss the distribution of $\mathbf{r}_{t}$ with respect to
timesteps $t$. Note that $\mathbf{r}_{t}^{1}$ represents the probability of
taking the global route, while $\mathbf{r}_{t}^{2}$ represents the probability
of taking the local route. Figure 4 shows the average and the standard
deviation of the local route probabilities taken by all the instances of a
certain client. Lower probability of local route means that a client prefers
the global route, and higher probability of local route means that a client
prefers the local route.

We picked 100 test clients with their KL divergence having mean of 5.67,
median of 3.00, and the range of 0.05 to 24.63. Client #1 from Figure 4(a) has
relatively lower KL divergence value, indicating that most of its instances
perform well just with the global model. But there would still be few
instances (in this case, there are 63 such instances) who would have preferred
their local models. An opposite behavior is shown for Client #3 in Figure 4(b)
where due to higher divergence, more instances (390 instances) prefer the
local model. But again, some instances (83 instances) which fit well within
the global data distribution would prefer the global model and hence take the
global route in our dynamic personalized model.

For the first several time-steps, the policy is less confident about its
choice between routes as it has only seen a few tokens. As policy sees more
tokens, its confidence about picking a certain route increases. We see this
behavior in Figure 4 from around $11^{th}$ time-step. The key takeaway is that
on per-client granularity, a client with high divergence would have just
preferred its local route for all the instances. But now that we can make per-
instance decisions with Flow, some instances with low divergence to the
estimated global distribution can pick the global route. Same intuition
applies to a client with low divergence still preferring the local model for
few instances.

(a) Stackoverflow

(b) Reddit

(c) EMNIST

Figure 7: Validation accuracy versus number of federated rounds of Flow and
its baselines.

##  Appendix F Ablation Studies

### Effects of Regularization

(a) Stackoverflow

(b) Reddit

(c) EMNIST

Figure 8: Effect of regularization

Figures 8 show the impact of the regularization terms in policy and global
objectives as stated in Equation 3. In order to keep the global parameters
generalizable for a newly joined client to adapt its personalized model from,
similar to Skim-RNN (Seo et al. 2018), we had jointly minimized the negative
mean of the log probability of the policy picking the global route. As
compared to no regularization, the regularized objective produces a more
generalizable global model $\boldsymbol{w}^{g}$. This is indicated from the
increment in accuracy for the validating clients, which haven’t taken part in
the training rounds.

##### Effect of Local Epochs

(a) Stackoverflow

(b) Reddit

(c) EMNIST

Figure 9: Effect of local epochs

Figures 9 show the effect of local epochs $e_{\ell}$. Recall that Flow
requires finetuning $\boldsymbol{w}^{g}_{\mathsf{enc}}$ for $e_{\ell}$ epochs
to create local parameters $\boldsymbol{w}^{\ell}_{\mathsf{enc}}$. In line
with FedRecon (Singhal et al. 2021) and MAML (Fallah, Mokhtari, and Ozdaglar
2020), having more rounds of training does not improve the trained model’s
performance. In fact, for language tasks, training for more rounds in Flow
results in overfitting and hence, degraded performance. Since the number of
training instances for each client differ widely in this statistically
heterogeneous setup, for most of the clients having 10 to 100 instances, a
single epoch proves to be sufficient to train the model.

Images require more number of epochs (5 in our case) to train the policy. We
note that in the language tasks, the policy is trained for each timestep in an
instance sequence. But for the image classification task, there’s no notion of
timesteps and a policy output is only dependent on an instance.