Title: Flow: Per-instance Personalized Federated Learning
Venue: NeurIPS 2023
Accepted: yes
Keywords: federated learning

personalization

statistical heterogeneity

dynamic routing
Abstract: Federated learning (FL) suffers from data heterogeneity, where the diverse data distributions across clients make it challenging to train a single global model effectively. Existing personalization approaches aim to address the data heterogeneity issue by creating a personalized model for each client from the global model that fits their local data distribution. However, these personalized models may achieve lower accuracy than the global model in some clients, resulting in limited performance improvement compared to that without personalization. To overcome this limitation, we propose a per-instance personalization FL algorithm Flow. Flow creates dynamic personalized models that are adaptive not only to each client’s data distributions but also to each client’s data instances. The personalized model allows each instance to dynamically determine whether it prefers the local parameters or its global counterpart to make correct predictions, thereby improving clients’
accuracy. We provide theoretical analysis on the convergence of Flow and empirically demonstrate the superiority of Flow in improving clients’ accuracy compared to state-of-the-art personalization approaches on both vision and language-based tasks.

Flow: Per-instance Personalized Federated Learning

Abstract

Federated learning (FL) suffers from data heterogeneity, where the diverse data distributions across clients make it challenging to train a single global model effectively. Existing personalization approaches aim to address the data heterogeneity issue by creating a personalized model for each client from the global model that fits their local data distribution. However, these personalized models may achieve lower accuracy than the global model in some clients, resulting in limited performance improvement compared to that without personalization. To overcome this limitation, we propose a per-instance personalization FL algorithm Flow . Flow creates dynamic personalized models that are adaptive not only to each client's data distributions but also to each client's data instances. The personalized model allows each instance to dynamically determine whether it prefers the local parameters or its global counterpart to make correct predictions, thereby improving clients' accuracy. We provide theoretical analysis on the convergence of Flow and empirically demonstrate the superiority of Flow in improving clients' accuracy compared to state-of-the-art personalization approaches on both vision and language-based tasks. The source code is available on GitHub 1 .

1 Introduction

Federated Learning (FL) is a distributed machine learning paradigm that enables edge devices, known as 'clients', which collaboratively train a machine learning model called a 'global model' [1]. However, because the server, which is the FL training orchestrator, does not have access to or knowledge of client data distributions, it poses a challenge of statistical heterogeneity [2, 3]. This heterogeneity hinders the server's ability to train an ML model on a large quantity and variety of data and also impacts the client's ability to benefit from a generalizable model without sharing any information about its data. To address this challenge, personalization has been studied [4, 5] to improve prediction performance. Recent literature consistently demonstrates that personalized models achieve higher prediction performance than the global model aggregated over clients [5-10]. These approaches typically create a personalized model specific to each client's data distribution, and thus we refer to them as 'per-client personalization'.

However, we have identified two factors that limit the performance improvement of existing per-client personalization approaches in terms of clients' accuracy. First, as reported in the evaluation in Sec. 5, we found that personalized models can achieve lower accuracy than the global model on up to 31% clients, causing limited improvement in accuracy averaged across all clients. Second, even if personalized models achieve higher accuracy on a client compared to the global model, they can still produce incorrect predictions on up to 11% of data instances on clients that could be correctly handled by the global model, causing limited accuracy improvement from personalization on that client. These observations reveal a significant drawback of existing per-client personalization approaches: each client's personalized model is a static network that cannot accommodate the heterogeneity of the local client's data instances. As a result, every data instance on a client is constrained to use its personalized model for prediction, even though some instances could benefit from the better generalizability of the global model.

To overcome the above limitation, this paper proposes a per-instance and per-client personalized FL algorithm Flow via dynamic routing to improve clients' accuracy. Flow creates dynamic personalized models that are adaptive to each client's individual data instances (per-instance) as well as their data distribution (per-client). In a FL round of Flow , each client has both the global model parameters that are shared across clients and the local model parameters that are adapted to the local data distribution of the client by fine-tuning the global model parameters. Flow creates a dynamic personalized model per client that consists of local and global model components and a dynamic routing module. The dynamic routing module allows each data instance on a client to determine whether it prefers the local model parameters or their global model counterparts to make correct predictions, thereby improving the personalized model's accuracy on the client compared to that of a local model or a global model. At the same time, through dynamic routing, Flow could identify instances in each client that agree with the global data distribution to further improve the performance of the global model, offering a good starting point for any new client to personalize from. Since Flow is a client-side personalization approach, it can work with server-side optimization methods like FEDYOGI [11].

We theoretically analyze how dynamic routing affects the convergence of the global model and personalized model in Flow and also empirically demonstrate the effectiveness of Flow in improving clients' accuracy compared to state-of-the-art personalization approaches on cross-device language and vision tasks. For the newly joined clients, the global model from Flow achieves 2.81% (Stackoverflow), 3.46% (Shakespeare), and 0.95%-1.41% (CIFAR10) better accuracy on the global model against its best performing baselines respectively. After personalization, the dynamic personalized model from Flow sees improvements of 3.79% (Stackoverflow), 2.25% (Shakespeare), and 3.28%-4.58% (CIFAR10) against the best performing baselines. Our in-depth analysis shows that Flow achieves the highest percentage of clients who benefit from personalization compared to all baselines and reduces the number of instances that are misclassified by the personalized model but correctly classified by the global model, contributing to the better personalized accuracy from Flow .

We summarize the contributions as follows:

We propose a per-instance and per-client personalization approach Flow that creates personalized models via dynamic routing, which improves both the performance of the personalized model and the generalizability of the global model.

We derive convergence analysis for both global and personalized models, showing how the routing policy influences convergence rates based on the across- and within- client heterogeneity.

We empirically evaluate the superiority of Flow in both generalization and personalized accuracy on various vision and language tasks in cross-device FL settings.

2 Related Work

Personalized Federated Learning. Personalization in FL has been explored primarily on client-level granularity. APFL [6] interpolates client-specific local model weights with the global model weights which are sent by the server. Meanwhile, Flow is an intermediate-output interpolation method and also includes a dynamic policy to interpolate at instance-level granularity. We note that DAPPER [12] interpolates the client dataset with a global dataset, which is impractical for the cross-device use cases FL focused on in this paper. 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 as explored

in DITTO [13] and PFEDME [14]. Finetuning partial or entire global model to get the personalized model has been studied in LGFEDAVG [7], which finetunes and personalizes the feature extractor and globally shares the classifier head, FEDBABU [15] which freezes the classifier head and only updates and aggregates the feature extractor, and FEDREP [8] which finetunes and personalizes the classifier head and globally shares the feature extractor. Learning personalized representations of the global model has been explored in FEDHN [16] where a central hypernetwork model is trained to generate a personalized model for each client. PARTIALFED [17] makes a layer-wise choice between global and local models to create a personalized model, but it sends the personalized models back to the server in lieu of a separate global model. This method has limitations in terms of training the base global model (on which the personalized model would be based) with FEDAVG, which has been observed to be insufficient against non-iid data [18]. Besides, the strategy does not create a dynamic personalized model during inference. None of the above work has explored instance-level personalization.

Personalized models from previous work KNNPER [19] exhibits per-instance personalization behavior. The personalized model of a client makes a prediction based on features extracted from the global model from an instance as well as features from the instance's nearest neighbors. However, KNNPER trains the global model parameters in the same way as FEDAVG, which has been shown to perform poorly in heterogeneous data settings [18]. In contrast, Flow 's dynamic routing mechanism results in a global model with better-generalized performance amidst heterogeneous instances, which ultimately leads to a higher boost in the performance of personalized models as well.

Dynamic Routing. Motivated by saving compute effort for 'easy to predict' instances, instance-level dynamic routing has been a matter of discussion in works related to early exiting [20-22], layer skipping [23], and multi-branching [24, 25]. SKIMRNN [26] 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 Depending on the models' utility and instances' heterogeneity, which route to pick? . This can be achieved by using a routing policy to dynamically pick from two versions of a model, which are equivalent in terms of computational cost but different in terms of the data they are trained on.

3 Our Approach

We introduce Flow , a per-instance and per-client 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. Table 1 summarizes the notation used in this paper.

Algorithm 1 describes the workflow of Flow . During each FL round, the server samples M participating clients with a sampling rate of p . Upon receiving the global model w g and policy parameters ψ g (Line 3), each participating client personalizes and trains w g in five major stages: (1) Split the training dataset into two halves: ζ m,ℓ and ζ m,g (Line 6). ζ m,ℓ will be used to update the parameters in the local model parameters while ζ m,g will be used to train the parameters in the global model and a routing module. (2) Derive the local parameters w ℓ by finetuning the global parameters w g for K 1 epochs with ζ m,ℓ (Line 7). (3) Construct a dynamic personalized model w p,m by integrating the local versions of the global parameters w g,m and policy parameters ψ g,m , and the local parameters w ℓ . Here, the routing policy ψ g determines whether the execution

path of an instance should use w ℓ or w g . (4) Train the routing policy ψ g,m and the local version of the global parameters w g,m alternatively for K 2 epochs (Lines 9-10) with ζ m,g . Although K 1 and K 2 can be different, we later use K to denote both K 1 and K 2 for ease of theoretical analysis. (5) Send the local version of the global model parameters w g,m and the policy parameters ψ g,m back to the server for aggregation (Line 14). Aggregation strategies are orthogonal to this work and we adopt FEDAVG [1].

Next, we discuss the design of the dynamic personalized model w p in detail. Figure 1 illustrates the design of w p . In Flow , w p is made of three components, the local model w ℓ and global model w g and the routing module ψ g that selects the execution of local or global model layers for each data instance.

Local Parameters. We use finetuning to get the local parameters w ( r ) ℓ,m in the r -th round since finetuning has proven to be a less complex yet very effective method of personalization [27, 28]. Given a global model of the previous round, w ( r ) g,m , we finetune it for K epochs to get w ( r ) ℓ,m . This local model would be reflective of the client's data distribution D m . Note that we use one half of the dataset to get w ( r ) ℓ,m and reserve the other half of the training dataset for updating the global model and the policy (i.e., the routing module) parameters. This is to make sure that the policy parameters, which would decide between local and global layers based on each input instance, are not overfitted in favor of the local parameters. In Figure 1, the local parameters are shaded green, denoted by w ℓ . The update rule for w ( r ) ℓ,m is:

Algorithm 1: Pseudocode of Flow

ψ

𝒘

g

𝒘

ℓ

𝒒

𝒘

g

𝒘

ℓ

𝑋

𝑌

inputs

outputs

𝒘 𝑝

trained

from

𝒘

g

probability of

picking a route

personalized model

global model

local model

policy

model

ψ base

ψ exit

✕

L

towards next

policy layer

from

previous

layer

𝑋

…

…

to next

layer

𝑌

Routing Module. The routing module is a model with fully connected layers, shown as ψ base in Figure 1. After each layer ψ base, the model has early exits denoted as ψ exit which outputs a probability of choosing the layer in the global model w g or the local model w ℓ . This probability at layer j ∈ [ L ] is computed as,

where q ( j ) 0 and q ( j ) 1 are the probability of picking the global parameter w ( j ) g,m and the local parameter w ( j ) ℓ,m respectively.

In order to train the routing parameters, we compute the training loss based on the output of the personalized model w ( r ) p,m , which averages the global model's output and local model's output weighted by q ( j ) , for each layer j ∈ [ L ] :

Let f m ( · ) denote the loss function on client m . Using the above personalized model w ( r ) p,m , Flow updates the policy parameters as follows,

where -γ L ∑ j ∈ [ L ] log( q ( j ) 0 ) is a regularization term that encourages the global model to learn on heterogeneous instances to improve the global model's generalizability.

Global Parameters. Global parameters w g are trained alternatively with the routing parameters ψ g . Flow updates the global parameters as follows,

The goal of the alternative training is for the instances that are characterized better by the global data distribution to get diverted to the global model and the rest of the instances that are captured better by the local data distribution routed to the local model. This improves the stability of the global model compared to approaches like APFL or KNNPER which still use the global model trained on all the instances, regardless of the level of heterogeneity. The effectiveness of the alternative training is validated empirically and discussed in Appendix D.2.

Soft versus Hard Policy. During training, we use soft policy where the probability in Equation 2 ranges over [0 , 1] to update the parameters in the global model and the routing module. But during inference, we use a discrete version of the policy round ( q ( j ) ) ∈ { 0 , 1 } 2 for q ( j ) in Eq. 3. The rationale behind this is twofold: (a) Using hard policy saves compute resources during inference by executing either the global or the local layer for an instance instead of both. (b) Our empirical evaluation shows negligible difference in performance of personalized models with soft and hard policies during the inference.

The dynamic nature of the personalized model in Flow introduces additional storage and computational overhead compared to the canonical method FEDAVG with Fine Tuning (called FEDAVGFT). However, compared to other state-of-the-art personalization methods such as DITTO and APFL, Flow requires similar or even less storage and computational overhead. Detailed analysis and comparison with baselines is in Appendix C.

4 Theoretical Analysis

In this section, we give convergence bounds for the global model w g and the personalized model w p of an arbitrary client m in smooth non-convex cases. The bounds for strong and general convex cases are available in Appendix E, Sections E.3 and E.5. To derive the bounds, we adopt two commonly used assumptions in FL convergence analysis from AFO (Assumption 2, [11]) and SCAFFOLD (Assumption 1, [18]): (1) We assume local variance between a client's expected and true gradient is σ ℓ -bounded. (2) We also assume that the dissimilarity between aggregated gradients of local expected risk and the true global gradient is ( G,B ) -bounded, where both G and B are constants capturing the extent of gradient dissimilarities. A detailed description is in Appendix Section E.2.

Now we present the bounds on the norm of expected global (and local) risks on global (and personalized) models respectively, at R -th (last) round. The proofs are in Appendix Sections E.3 and E.5.

Theorem 4.1 (Convergence of the Global Model) . If each client's objective function f m (and hence the global objective function F ) satisfies β -smoothness, σ ℓ -bounded local gradient variance, ( G,B ) -dissimilarity assumptions, using the learning rate 1 2 β ≤ η ℓ ≤ 1 2 √ 5 βBK 2 √ R [for non-convex case] in Flow, then the following convergence holds:

Discussion. Here, q 2 0 (and q 2 1 ) are the probability of picking the global (and local) weights averaged over all instances sampled from the global data distribution. Given a fixed q 2 0 and q 2 1 , a larger G increases the bound, indicating slower convergence of the global model due to higher heterogeneity. We made two additional observations on the convergence of Flow depending on q 2 0 in the bound.

One the one hand, when q 2 0 → 1 , Flow matches the linear (strong convex) and sub-linear (general convex and non-convex) convergence rates of FEDAVG [18]. Specifically, q 2 0 = 1 results in convergence rate of O (1 /R 2 ) (strong convex), O (1 /R 2 / 3 ) (general convex), and O (1 /R ) (non-convex). When q 2 0 = 1 , which also indicates q 2 1 = 0 , the local parameters do not influence global model updates since all instances on a client will be routed to the global model and their gradients solely depend on the global model parameters. Flow degrades to FEDAVG in this case in terms of the convergence of the global model.

On the other hand, when q 2 0 → 0 indicating all instances are likely to be routed to the local model, the bound goes to infinity. It is because the global model parameters won't get updated and thus won't be able to converge. This observation validates the necessity of the regularization term in Eq. 4 to encourage instances to pick global model parameters.

Now we present the convergence bounds of the personalized model for the m -th client. The bound uses a definition of gradient diversity (noted as δ m ) to quantify the diversity of a client's gradient with respect to the global aggregated gradient, following the prior work [29]. Higher diversity implies higher heterogeneity of the client.

Theorem 4.2 (Convergence of the Personalized Model) . If each client's objective function f m satisfies β -smoothness, σ ℓ -bounded local gradient variance, ( G,B ) -dissimilarity assumptions, and using the learning rate η ℓ ≤ 1 K √ 12 βR [for non-convex case] in Flow, then the following convergence holds:

Discussion. The theorem implies two main properties of the personalized models in Flow . First, for all convex and non-convex cases, the convergence rate of the personalized model in Flow is affected by the routing policy through the ratio q 2 1 ,m G 2 / q 2 0 ,m . We know that a higher value of the gradient dissimilarity constant G , indicates higher heterogeneity between the aggregated and expected global model. The ratio of q 2 1 ,m / q 2 0 ,m would be higher for a heterogeneous client, since the client would get a higher probability of picking the local route ( q 2 1 ,m → 1 ). The higher G and q 2 1 ,m results in slower convergence. On the contrary, a homogeneous client would benefit from a low value of q 2 1 ,m / q 2 0 ,m , which would offset the high value of G . Hence a homogeneous client's personalized model would converge faster than the one of a heterogeneous client. Second, we observe that gradient diversity of the policy model, δ ψ g m , linearly affects the personalized model's convergence. Since the policy model is also globally aggregated, a heterogeneous client would have a high δ ψ g m and need more epochs per round to converge.

5 Experiments and Results

We empirically evaluate the performance of Flow against various personalization approaches for five non-iid vision and language tasks in terms of clients' accuracy.

Datasets, Tasks, and Models. We have experiments on two language and three vision tasks. The first three datasets which are described below represent real-world heterogeneous data where each author or user is one client. (a) Stackoverflow [30] dataset is used for the next word prediction task, using a model with 1 LSTM and 2 subsequent fully connected layers. (b) Shakespeare [31] dataset is for the next character prediction task, using a 2 layer LSTM + 1 layer fully connected model. (c) EMNIST [32] dataset is for 62-class image classification, which uses 2 CNN layers followed by 2 fully-connected layers. The models for the above three datasets have been described in [11]. The next two datasets are federated and artificially heterogeneous versions of CIFAR10/100

datasets. (d-e) CIFAR10/100 [33] datasets are for 10- and 100-class image classification tasks with ResNet18 [34] model. Both CIFAR10/100 have two heterogeneous versions each: 0.1-Dirichlet is more heterogeneous, and 0.6-Dirichlet is less heterogeneous. Details about the datasets and the hyperparameters are in Appendix B.

Baselines and Metrics. We compare Flow with the following baselines: the classic FL algorithms FEDAVG [35], state-of-the-art personalized FL approaches including FEDAVGFT (FedAvg + Finetuning) [36], PARTIALFED [17], APFL [6], FEDREP [8], LGFEDAVG [7], DITTO [13], HYPCLUSTER [12], and KNNPER [19], and the LOCAL baseline which trains a local model on each client's dataset without any collaboration. We use the adaptive version of PARTIALFED as it shows better performance compared to the alternatives in [17]. We evaluate Flow and these baselines in terms of generalized accuracy and personalized accuracy , which correspond to the accuracy of the global model and the personalized model on clients' test data split.

We use Flower [37] library to implement Flow and all its baselines. We use an NVidia 2080ti GPU to run all the experiments with 3 runs for each. The random seeds used are 0, 44, and 56. We do not observe significant difference in results using other random seeds (see results in Appendix D.4).

5.1 Performance Comparison

Generalized and Personalized Accuracy. The performance of Flow and its baselines are reported in Table 2 for four datasets. Note that the LOCAL baseline has only personalized accuracy as it doesn't create a global model collaboratively. The FEDAVG has only generalized accuracy as it doesn't personalize the global model for each client. Since PARTIALFED is a stateful approach, we are unable to run it on the cross-device datasets (Stackoverflow, Shakespeare, EMNIST). Results for the rest of the datasets and their variance across 3 different runs are reported in Appendix D.

Overall, Flow achieves 1.113.46% higher generalized accuracy and 1.33-4.58% higher personalized accuracy over the best performing baseline. In particular, Flow outperforms KNNPER, another per-instance perclient personalization approach, by 1.66-6.64% and 1.33-6.97% in generalized and personalized accuracy metrics respectively. KNNPER allows each instance to personalize the prediction of the global model based on its k- nearest neighbors at inference time. However, the global model is trained through the classic FL method FEDAVG, which is agnostic to the heterogeneity of data instances. In contrast, Flow trains the global model differently via a parameterized dynamic routing module, which learns to put emphasis on data instances that are more aligned with the global data distribution to improve the performance of the global model.

Flow also outperforms the per-client personalization approaches including APFL, DITTO, HYPCLUSTER, and FEDAVGFT by 1.35-3.46 % (generalized accuracy) and 2.25-4.58% (personalized accuracy).

We see improvements in generalized accuracy of Flow because of the fact that the global model in Flow is trained based on the instances which align more with the global distribution. We see improvements in personalized accuracy due to the limitation of the per-client approaches where some instances being correctly classified by the global model are incorrect on the personalized model. We next give insights on why

Flow achieves better performance in person-

alized and generalized accuracy compared to these personalized baselines.

FedAvg + FT

APFL

Ditto

knnPer

FedRep

LGFedAvg

HypCluster

Flow

Baselines

0.00

0.05

0.10

0.15

0.20

0.25

0.30

% of Instances

both-correct

personalized-only

global-only

1

2

3

4

5

6

7

8

9

10

11

12

13

14

15

16

17

18

19

20

21

22

Layer indices

0.0

0.2

0.4

0.6

0.8

1.0

(Probability of picking the global route)

both-correct

global-only

personalized-only

Percentage of Clients Benefiting from Personalization. The goal of personalization is to achieve higher prediction accuracy in each client by creating a per-client personalized model from the global model. We can thus measure the effectiveness of personalization by computing the percentage of clients for which the personalized model achieves higher task accuracy than the global model. The higher the percentage is, the better (or more effective) the personalization approach is.

Table 3 reports the results. We observed that Flow achieves the highest percentage of clients who benefit from personalization compared to all personalization baselines, echoing the better personalized accuracy from Flow . The percentage of clients who prefer personalized models can be as low as 68.87% (KNNPER on Shakespeare), which means personalization hurts up to 31% of clients' accuracy, as mentioned in the introduction. As a contrast, Flow improves the percentage of clients benefiting from personalization to 89.77%-99.62% because each instance, in a client, has a choice between the global model parameters and the local model parameters and can choose the one that better fits it. Note that the comparison is in favor of the baselines since Flow also achieves better generalized accuracy, which makes it even harder for personalized models to further improve prediction accuracy.

Breakdown of Correctly Classified Instances. Figure 2 further shows the breakdown of the percentage of instances that are (a) correctly classified by the global model but not the personalized model (noted as global-only , colored in yellow ), (b) correctly classified by the personalized model but not the local model (noted as personalized-only , colored in blue ), and (c) correctly classified by both models (noted as both-correct , colored in green in Flow and baselines on Stackoverflow. The percentage of instances in y-axis is averaged over the test splits of all clients.

Overall, Flow increases the both-correct bars compared to all the baselines, which are the instances that contribute to the generalized performance of the global model. This explains the better generalized accuracy of Flow . Flow also increases the personalized-only bars and decreases the global-only bars, which correspond to the heterogeneous instances that prefer personalized models instead of the global model. This further explains the better personalized accuracy of Flow . Notably, for Stackoverflow dataset, existing personalization approaches still result in up to 4.74-7.93% of instances incorrectly classified by the personalized model but correctly by the global model. Flow reduces it to 1.12-2.42%. Similarly for the CIFAR10 (0.6) dataset, as mentioned in the introduction, we notice up to 11.4% of instances falling under the global-only category, which Flow reduces to 2.55%. It echoes the aforementioned effectiveness of personalization in Flow . The instance-wise accuracy breakdown for the rest of the datasets is detailed in Appendix D, Figure 10.

Analysis of Routing Decisions. We further analyze the behavior of routing decisions for instances of a client that fall in the above three cases, global-only , personalized-only , and both-correct . Figure 3 shows the per-layer routing policies of the dynamic personalized model from Flow on Stackoverflow. For instances that fall into each category, we average the policy value from Eq. 2 and report the statistics on the probability of picking the global parameters for each layer. The statistics for the rest of the datasets are detailed in Appendix D.

For instances that are correctly classified by w g but not by w p ( global-only ), we see a clear trend of the routing parameters getting more confident about picking the global parameters. As a contrast,

250

500

750

1000 1250 1500 1750 2000

Number of Rounds

0.20

0.22

0.24

0.26

0.28

0.30

Validation Accuracy (Personalized)

of Flow

Flow with Regularization

Flow without Regularization

(a) Ablation on the regularization

term

250

500

750

1000 1250 1500 1750 2000

Number of Rounds

0.18

0.20

0.22

0.24

0.26

0.28

0.30

Validation Accuracy (Personalized)

of Flow variants

Per-instance Per-client Flow

Per-instance Flow

Per-client Flow

(b) Ablation on per-instance

personalization design

250

500

750

1000 1250 1500 1750 2000

Number of Rounds

0.21

0.22

0.23

0.24

0.25

0.26

0.27

Validation Accuracy (Personalized)

of Flow variants

Dynamic Routing

Fixed

q

0 = 1.00

Fixed

q

0 = 0.75

Fixed

q

0 = 0.50

Fixed

q

0 = 0.25

(c) Ablation on dynamic routing

component

for instances that are correctly classified by w p but not by w g ( personalized-only ), we see the trend of routing policy being more confident in picking the local parameters. For instances that can be correctly classified by both models ( both-correct ), the routing policy still prefers the global parameters over local parameters. This is due to the regularization term in Eq. 4, which encourages instances to pick the global model over the local model in order to improve the generalizability of the global model. Our ablation study in the next section demonstrates the importance of regularization.

5.2 Ablation Studies

Here we highlight some results of three ablation studies on regularization, per-instance personalization, dynamic routing, and hard policies during inference. More results are in Appendix D.

Regularization. The regularization term in Eq. 4 promotes the global model layers whenever possible. It helps boost the generalized performance of the global model, which in turn also produces better personalized models. We use the results on the Stackoverflow dataset in Figure 4a to illustrate the importance of regularization. At the end of the training, we get 26.64% ± 0.23% generalized accuracy with regularization, compared to 24.16% ± 0.34% without regularization, and 29.49% ± 0.28% personalized accuracy with regularization, compared to 27.59% ± 0.36% without regularization. The importance of regularization in the policy update rule is also highlighted in Theorem 4.1, which states that only picking local route does not lead to global model convergence; the regularization term can encourage the global model to converge faster.

Per-instance Personalization. The dynamic personalized model in Flow allows each instance to choose between the local model and global model layers. To verify this per-instance personalization design, we create two variants of Flow , named 'Per-Instance' Flow (PI-FLOW) where both paths of a dynamic model are global models, and 'Per-Client' Flow (PC-FLOW) which is simply FEDAVGFT. Compared to the personalized accuracy of per instance and per client Flow (29.49% ± 0.28%), PI-FLOW and PC-FLOW achieve 26 . 31% ± 0 . 19% and 24 . 41% ± 0 . 26% respectively on Stackoverflow dataset, as shown in Figure 4b. The results demonstrate the effectiveness of the per-instance personalization design in Flow .

Dynamic Routing. This ablation study aims to learn whether dynamically interpolating global and local routes has any advantages over fixing the routing policy throughout the training. In Figure 4c, we compare the validation accuracy curves of dynamic routing in Flow and dynamic routing with instance-agnostic static routing during the training phase. We observe that (a) For the case of fixed policy of q 0 = 0 . 25 , the validation accuracy has bad performance due to the fixed policy only choosing the local route. This is due to using hard policy during inference, and (b) The cases of fixed policy q 0 ∈ [0 . 50 , 0 . 75 , 1 . 00] will only pick the global route during inference, which are also outperformed by the dynamic routing variant. With dynamic routing, the choice between local and global parameters depends on each instance during inference.

Soft versus Hard Decisions during Inference. We did not use soft decisions during inference since it only negligibly improves the accuracy of Flow . The test accuracies after personalization for Flow on Stackoverflow with hard decisions are 29.49% ± 0.28%, while with soft decisions, we observed 29.57% ± 0.22%. The rest of the datasets show a similar trend (see Appendix D, Table 8).

6 Conclusion

This paper proposed Flow , a per-instance and per-client personalization method to address the statistical heterogeneity issue in Federated Learning. Flow is motivated by the observation that the personalized models from existing personalized approaches achieve lower accuracy in a significant portion of clients compared to the global model. To overcome this limitation, Flow creates dynamic personalized models with a routing policy that allow instances on each client to choose between global and local parameters to improve clients' accuracy. We derived error bounds for global and personalized models of Flow , showing how the routing policy affects the rate of convergence. The theoretical analysis validates our empirical observations related to clients preferring either a global or a local route based on the heterogeneity of individual instances. Extensive evaluation on both vision and language-based prediction tasks demonstrates the effectiveness of Flow in improving both the generalized and personalized accuracy.

Acknowledgments and Disclosure of Funding

This material is based upon work supported by the National Science Foundation under Grant No. 2312396 , 2220211 , and 2224054 . Any opinions, findings, and conclusions or recommendations expressed in this material are those of the author(s) and do not necessarily reflect the views of the National Science Foundation. The work is also generously supported by Adobe Gift Fund.

References

H. Brendan McMahan, Eider Moore, Daniel Ramage, and Blaise Agüera y Arcas. Federated learning of deep networks using model averaging. Proceedings of the 20th International Conference on Artificial Intelligence and Statistics , 2016.

Tian Li, Anit Kumar Sahu, Ameet Talwalkar, and Virginia Smith. Federated learning: Challenges, methods, and future directions. IEEE Signal Processing Magazine , 2020.

Bing Luo, Wenli Xiao, Shiqiang Wang, Jianwei Huang, and Leandros Tassiulas. Tackling system and statistical heterogeneity for federated learning with adaptive client sampling. In IEEE INFOCOM 2022 IEEE Conference on Computer Communications , 2022.

Alysa Ziying Tan, Han Yu, Lizhen Cui, and Qiang Yang. Towards personalized federated learning. IEEE Transactions on Neural Networks and Learning Systems , 2022.

Alireza Fallah, Aryan Mokhtari, and Asuman Ozdaglar. Personalized federated learning with theoretical guarantees: A model-agnostic meta-learning approach. In Advances in Neural Information Processing Systems , 2020.

Yuyang Deng, Mohammad Mahdi Kamani, and Mehrdad Mahdavi. Adaptive personalized federated learning. arXiv preprint 2003.13461 , 2020.

Paul Pu Liang, Terrance Liu, Liu Ziyin, Ruslan Salakhutdinov, and Louis-Philippe Morency. Think locally, act globally: Federated learning with local and global representations. Advances in Neural Information Processing Systems Workshop on Federated Learning , 2019.

Liam Collins, Hamed Hassani, Aryan Mokhtari, and Sanjay Shakkottai. Exploiting shared representations for personalized federated learning. arXiv preprint 2102.07078 , 2021.

Krishna Pillutla, Kshitiz Malik, Abdelrahman Mohamed, Michael Rabbat, Maziar Sanjabi, and Lin Xiao. Federated learning with partial model personalization. arXiv preprint 2204.03809 , 2022.

Tian Li, Anit Kumar Sahu, Manzil Zaheer, Maziar Sanjabi, Ameet Talwalkar, and Virginia Smith. Federated optimization in heterogeneous networks. In Proceedings of Machine Learning and Systems , 2020.

Sashank J. Reddi, Zachary Charles, Manzil Zaheer, Zachary Garrett, Keith Rush, Jakub Koneˇ cný, Sanjiv Kumar, and Hugh Brendan McMahan. Adaptive federated optimization. In International Conference on Learning Representations , 2021.

Yishay Mansour, Mehryar Mohri, Jae Ro, and Ananda Theertha Suresh. Three approaches for personalization with applications to federated learning. arxiv preprint 2002.10619 , 2020.

Tian Li, Shengyuan Hu, Ahmad Beirami, and Virginia Smith. Ditto: Fair and robust federated learning through personalization. International Conference on Machine Learning , 2021.

Canh T. Dinh, Nguyen H. Tran, and Tuan Dung Nguyen. Personalized federated learning with moreau envelopes. In Proceedings of the 34th International Conference on Neural Information Processing Systems , 2020.

Jaehoon Oh, SangMook Kim, and Se-Young Yun. FedBABU: Toward enhanced representation for federated image classification. In International Conference on Learning Representations , 2022.

Aviv Shamsian, Aviv Navon, Ethan Fetaya, and Gal Chechik. Personalized federated learning using hypernetworks. In Proceedings of the 38th International Conference on Machine Learning , 2021.

Benyuan Sun, Hongxing Huo, Yi Yang, and Bo Bai. Partialfed: Cross-domain personalized federated learning via partial initialization. In Advances in Neural Information Processing Systems , 2021.

Sai Praneeth Karimireddy, Satyen Kale, Mehryar Mohri, Sashank Reddi, Sebastian Stich, and Ananda Theertha Suresh. SCAFFOLD: Stochastic controlled averaging for federated learning. In Proceedings of the 37th International Conference on Machine Learning , 2020.

Othmane Marfoq, Giovanni Neglia, Laetitia Kameni, and Richard Vidal. Personalized federated learning through local memorization. International Conference on Machine Learning , 2022.

Surat Teerapittayanon, Bradley McDanel, and H. T. Kung. Branchynet: Fast inference via early exiting from deep neural networks. arXiv preprint 1709.01686 , 2017.

Tolga Bolukbasi, Joseph Wang, Ofer Dekel, and Venkatesh Saligrama. Adaptive neural networks for efficient inference. In Proceedings of the 34th International Conference on Machine Learning , 2017.

Gao Huang, Danlu Chen, Tianhong Li, Felix Wu, Laurens van der Maaten, and Kilian Weinberger. Multiscale dense networks for resource efficient image classification. In International Conference on Learning Representations , 2018.

Xin Wang, Fisher Yu, Zi-Yi Dou, Trevor Darrell, and Joseph E. Gonzalez. Skipnet: Learning dynamic routing in convolutional networks. In The European Conference on Computer Vision , 2018.

Xin Dai, Xiangnan Kong, and Tian Guo. 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, 2020.

Lanlan Liu and Jia Deng. Dynamic deep neural networks: Optimizing accuracy-efficiency trade-offs by selective execution. In Association for the Advancement of Artificial Intelligence , 2018.

Minjoon Seo, Sewon Min, Ali Farhadi, and Hannaneh Hajishirzi. Neural speed reading via skim-RNN. In International Conference on Learning Representations , 2018.

Jian Xu, Xinyi Tong, and Shao-Lun Huang. Personalized federated learning with feature alignment and classifier collaboration. In The Eleventh International Conference on Learning Representations , 2023. URL https://openreview.net/forum?id=SXZr8aDKia .

Shanshan Wu, Tian Li, Zachary Charles, Yu Xiao, Ziyu Liu, Zheng Xu, and Virginia Smith. Motley: Benchmarking heterogeneity and personalization in federated learning. arXiv preprint 2206.09262 , 2022.

Blake E Woodworth, Kumar Kshitij Patel, and Nati Srebro. Minibatch vs local sgd for heterogeneous distributed learning. In Advances in Neural Information Processing Systems , 2020.

Stackoverflow. TensorFlow Federated Datasets. https://www.tensorflow.org/federated/api_ docs/python/tff/simulation/datasets , 2023. [Online; accessed 09-Jan-2023].

Shakespeare. TensorFlow Federated Datasets. https://www.tensorflow.org/federated/api_ docs/python/tff/simulation/datasets/shakespeare , 2023. [Online; accessed 09-Jan-2023].

Gregory Cohen, Saeed Afshar, Jonathan Tapson, and André van Schaik. Emnist: Extending mnist to handwritten letters. In 2017 International Joint Conference on Neural Networks (IJCNN) , 2017.

Alex Krizhevsky, Geoffrey Hinton, et al. Learning multiple layers of features from tiny images, 2009.

Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition , 2016.

Brendan McMahan, Eider Moore, Daniel Ramage, Seth Hampson, and Blaise Aguera y Arcas. Communication-Efficient Learning of Deep Networks from Decentralized Data. In Proceedings of the 20th International Conference on Artificial Intelligence and Statistics , 2017.

Liam Collins, Hamed Hassani, Aryan Mokhtari, and Sanjay Shakkottai. Fedavg with fine tuning: Local updates lead to representation learning. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho, editors, Advances in Neural Information Processing Systems , 2022. URL https://openreview. net/forum?id=G3fswMh9P8y .

Daniel J Beutel, Taner Topal, Akhil Mathur, Xinchi Qiu, Titouan Parcollet, and Nicholas D Lane. Flower: A friendly federated learning research framework. arXiv preprint 2007.14390 , 2020.

Wenqi Wei, Ling Liu, Margaret Loper, Ka Ho Chow, Mehmet Emre Gursoy, Stacey Truex, and Yanzhao Wu. A framework for evaluating gradient leakage attacks in federated learning. arxiv preprint 2004.10397 , 2020.

Xiao Jin, Pin-Yu Chen, Chia-Yi Hsu, Chia-Mu Yu, and Tianyi Chen. Catastrophic data leakage in vertical federated learning. Advances in Neural Information Processing Systems , 2021.

Wenqi Wei, Ling Liu, Yanzhao Wut, Gong Su, and Arun Iyengar. Gradient-leakage resilient federated learning. In IEEE 41st International Conference on Distributed Computing Systems , 2021.

Milad Nasr, Reza Shokri, and Amir Houmansadr. Comprehensive privacy analysis of deep learning: Passive and active white-box inference attacks against centralized and federated learning. In IEEE symposium on security and privacy , 2019.

Hongsheng Hu, Zoran Salcic, Lichao Sun, Gillian Dobbie, and Xuyun Zhang. Source inference attacks in federated learning. In IEEE International Conference on Data Mining , 2021.

Mohammad Naseri, Jamie Hayes, and Emiliano De Cristofaro. Toward robustness and privacy in federated learning: Experimenting with local and central differential privacy. Proceedings of the 29th Network and Distributed System Security Symposium , 2022.

Kang Wei, Jun Li, Ming Ding, Chuan Ma, Howard H. Yang, Farhad Farokhi, Shi Jin, Tony Q. S. Quek, and H. Vincent Poor. Federated learning with differential privacy: Algorithms and performance analysis. IEEE Transactions on Information Forensics and Security , 2020.

Stackoverflow Dataset FedML. FedML Federated Stackoverflow Heterogenity . https: //github.com/FedML-AI/FedML/blob/ecd2d81222301d315ca3a84be5a5ce4f33d6181c/doc/ en/simulation/benchmark/BENCHMARK_MPI.md , 2022. [Online; accessed 14-Sep-2022].

Jianyu Wang, Qinghua Liu, Hao Liang, Gauri Joshi, and H Vincent Poor. Tackling the objective inconsistency problem in heterogeneous federated optimization. Advances in neural information processing systems , 2020.

Ellango Jothimurugesan, Kevin Hsieh, Jianyu Wang, Gauri Joshi, and Phillip Gibbons. Federated learning under distributed concept drift. In NeurIPS 2022 Workshop on Distribution Shifts: Connecting Methods and Applications , 2022.

CIFAR100. TensorFlow Federated Datasets. https://www.tensorflow.org/federated/api_docs/ python/tff/simulation/datasets/cifar100/load_data , 2023. [Online; accessed 18-Jan-2023].

Hanxiao Liu, Karen Simonyan, and Yiming Yang. DARTS: Differentiable architecture search. In International Conference on Learning Representations , 2019.

Xuanyi Dong and Yi Yang. Searching for a robust neural architecture in four gpu hours. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition , 2019.

A Limitations, Future Work, and Broader Impact

Learning on naturally heterogeneous datasets can be challenging, as the true data distributions of individual clients are unknown, making it difficult to correlate the divergence between client data distribution and the global data distribution with routing policy decisions. In our approach, we estimate the distribution divergence by measuring the difference between inference losses on global and local models, which helps us reason about routing probabilities for global and local routes. To further improve our understanding of the model performance, we plan to propose a metric that quantifies the difference in performance when a particular dataset is included versus excluded.

Flow has shown the promise of per-instance personalization in improving clients' accuracy. This approach also holds the potential of preserving privacy by protecting against gradient leakage [38-40] and membership inference [41, 42] attacks that are easier to carry out in vanilla FL. Studying the relationship between personalization and privacy, and comparing our approach to traditional methods like Differential Privacy (DP) [43, 44] can reveal properties of personalization that go beyond improved accuracy.

B Datasets and Hyperparameters

Stackoverflow The Stackoverflow dataset [30] is comprised of separate clients designated for training, validation, and testing. The dataset contains a total of 342,477 train clients, whose combined sample count equals 135,818,730. Similarly, the dataset contains 38,758 validation and 204,088 test clients, whose combined sample counts equal 16,491,230 and 16,586,035, respectively. This dataset is naturally heterogeneous [45] since each user of Stackoverflow represents a client, with their posts forming the dataset for that client. The heterogeneity of the dataset arises from the fact that users have different writing styles, meaning the clients' datasets are not i.i.d., and the total number of posts from each user varies, leading to different dataset sizes per client.

We have trained Flow and its baselines on the Stackoverflow dataset for 2000 rounds. The one layer LSTM we have used has 96 as embedding dimension, 670 as hidden state size, and 20 as the maximum sequence length [11]. The batch size used for each client on each baseline is 16. The vocabulary of this language task is limited to 10k most frequent tokens. The default learning rate used is 0.1. The number of clients per round is set to 10, as is the common practice in [13, 46, 12, 10, 47]. For client-side training, the default epoch count is 3 for all the algorithms.

For KNNPER, we used 5 nearby neighbors, and the mixture parameter is λ = 0 . 5 . For APFL, mixture hyperparameter α is set to 0.25. DITTO has regularization hyperparameter λ = 0 . 1 . There are 2 clusters by default for HYPCLUSTER. Flow and its variants were tested on the following choices of regularizing hyperparameters γ ∈ { 1e-1, 1e-2, 1e-3, 1e-4 } , where 1e-3 gave the best personalized accuracy.

Shakespeare The Shakespeare dataset [31] consists of 715 distinct clients, each of which has its own training, validation, and test datasets. The combined training datasets of all clients contain a total of 12,854 instances, while the combined validation and test datasets contain 3,214 and 2,356 instances, respectively. The Shakespeare dataset is considered heterogeneous due to the fact that each client is a play written by William Shakespeare, and these plays have varying settings and characters.

All the baselines and Flow variants have been run for 1500 rounds, with 10 clients per round. The 2 layer LSTM used [11] has 8 as embedding size, vocabulary size of 90 most frequently used characters, and 256 as hidden size. The default epoch count is 5 for each client, for each algorithm. The batch size is 4 since bigger batch sizes resulted in the divergence of the global model across all the different runs. The default learning rate is 0.1.

Since each client has a sample count under 20, we have used 3 as the nearest neighbor sample count for KNNPER. λ and α , the mixture parameters, for KNNPER and APFL respectively, are set to 0.45 and 0.3. The regularization parameter λ for DITTO is set to 0.1. For Flow , the learning rate is set to 0.11 and the regularization parameter is picked from γ ∈ { 1e-1, 1e-2, 1e-3 , 1e-4 } similar to Stackoverflow.

EMNIST The EMNIST dataset [32] comprises 3400 distinct clients, each of which has its own training, validation, and test datasets. The combined total number of instances in the train datasets of all clients is 671,585, whereas the validation and test datasets of all clients combined contain 77,483 instances each. The heterogeneity of EMNIST clients is due to the individual writing styles of each client, with each client representing a single person. This is discussed in Appendix C.2 of [11].

The default round count for all the baselines and Flow variants is 1500, with 10 clients participating per round. Similar to AFO [11], we have used a shallow convolution neural network with 2 convolution layers. Each client uses 3 local epochs for on-device training. The default batch size is 20, and the default learning rate is 0.01.

For LOCAL only training, we have used 10 epochs per client with a learning rate of 0.05. The nearest sample count for KNNPER is 10 and the mixture parameter is λ = 0 . 4 . For APFL, we have the default mixture parameter as α = 0 . 25 . DITTO has regularization hyperparameter as λ = 0 . 1 . There are 2 clusters for the clustering

algorithm HYPCLUSTER. And for Flow , along with its variants, we have picked γ ∈ { 1e-1, 1e-2 , 1e-3, 1e-4 } as the regularizing hyperparameter.

CIFAR10 The CIFAR10 dataset is derived from the centralized version of the CIFAR10 dataset [33], which comprises 50,000 images. The federated CIFAR10 dataset consists of 500 unique clients, each of which has 100 training samples and 20 testing samples. The training and testing samples for each client are determined according to the Dirichlet distribution [11]. The heterogeneity of a client is determined by the Dirichlet distribution parameter α ∈ [0 , 1] , where a client is more heterogeneous than α → 0 . In this context, heterogeneity refers to the dissimilarity of the dataset instances sampled from a distribution. We conducted experiments on clients with α values of 0.1 and 0.6.

We ran all the experiments for 4000 rounds for the CIFAR10 dataset. ResNet18 [34] is used for all the algorithms. The default batch size is 20 and the default learning rate is 0.05. Each client individually trains their local versions of the global model for 3 epochs.

For LOCAL only training, 20 epochs per client were used. The learning rate was 0.1 for the same. The nearest sample count and the mixture hyperparameter for KNNPER are set to 5 and 0.5. PARTIALFED learning rate is set to 0.11, with the local epoch count is 5. APFL has mixture hyperparameter set as α = 0 . 2 . And DITTO has a regularization hyperparameter set as λ = 0 . 01 . Flow and its variants have their regularization hyperparameter as γ ∈ { 1e-1, 1e-2, 1e-3 , 1e-4 } .

CIFAR100 Like CIFAR10, the CIFAR100 dataset [48] is derived from the CIFAR100 dataset [33] consisting of 50,000 images. The number of clients and the count of training and testing images are identical to those of CIFAR10. Similarly, we also conducted experiments with the Dirichlet parameter set to α = 0 . 1 and α = 0 . 6 .

Similar to CIFAR10, we have a 4000 round count for all the algorithms ran on the CIFAR100 dataset. We have again used ResNet18 [34]. The default local epoch count is 3, and the default learning rate is 0.05. We have used 20 batch size for all the algorithms. For each round, 10 clients participate as is the norm stated in the Stackoverflow dataset description.

LOCAL only training has 20 epochs per client, and 0.1 learning rate. 5 nearest samples are used for KNNPER, while the mixture parameter λ is set to 0.4. PARTIALFED, just like in CIFAR10, has 0.11 learning rate and 5 local epochs per client. APFL has 0.25 as mixture parameter α . DITTO has 1e-2 as regularization parameter λ . For both CIFAR10 and CIFAR100, we have 2 as the default cluster count for HYPCLUSTER. Flow and its variants get {1e-1, 1e-2, 1e-3 , 1e-4} as the regularization hyperparameter γ .

C Computation, Communication, and Storage

While Flow introduces additional storage and computational overhead compared to the canonical method FEDAVG with Fine Tuning (called FEDAVGFT), compared to other state-of-the-art personalization methods such as DITTO and APFL, Flow requires similar or even less storage and computational overhead.

To illustrate, Table 4 compares the storage, computational overhead, and communication cost for personalized models from Flow and baselines using the CNN for EMNIST, and RNN for Stackoverflow.

Referring to Table 4, for Flow , ψ g is the policy module. For our experiments on Flow , the policy module has 1.23% for Stackoverflow RNN, 8.39% for Shakespeare RNN, 27.86% for EMNIST CNN, 35.51% for CIFAR10/100 parameters of w g . For EMNIST CNN, for one epoch, the dynamic routing takes 0 . 3 M FLOPs, while the rest of the model computations take 2 . 9 M FLOPs. Hence the overhead is 10.34%. For Stackoverflow RNN one epoch, the dynamic routing takes 0 . 89 M FLOPs, while the rest of the model computations take 12 . 34 M FLOPs, making the overhead 7.21%.

Below, we highlight the computational and storage analysis on Stackoverflow dataset for the best performing methods, APFL, and HYPCLUSTER. The computational overhead is calculated for one epoch of training, which is then multiplied with however many epochs the baseline needs for convergence. APFL needs to train two separate local models for more numbers of epochs than Flow , hence the higher computational overhead. In comparison with HYPCLUSTER, Flow introduces 7.23% overhead. With respect to both APFL and HYPCLUSTER, Flow introduces only 1.23% storage overhead.

Moreover, from Table 4, we can see that compared to some other state-of-the-art personalization techniques (e.g., FEDREP, APFL), Flow requires similar or even less computations, even with the dynamic routing operations.

Table 4: Computational and Communication costs, and Storage of Flow and its baselines. A = Local Storage of Personalized Model (unit: parameter count) for general case, B = Local Storage of Personalized Model (unit: parameter count) for Stackoverflow RNN case, C = Computational Overhead of Personalized Model of the RNN used for Stackoverflow (unit: FLOPs for training),

D = Communication Cost (unit: parameter count) for general case, and E = Communication Cost (unit: parameter count) for Stackoverflow RNN case.

D Additional Results

D.1 Generalized and Personalized Accuracy

Generalized (Personalized) accuracy is calculated based on the global (personalized) model, where each participating client's test dataset is used to compute accuracy of the global (personalized) model.

Generalized accuracy is formulated as

Personalized accuracy is formulated as

We have reported Generalized (Personalized) Accuracy Acc g ( Acc p ) of Flow , averaged across 1000 clients in Table 5, for all the datasets. Similarly, variance of accuracies across 3 different runs (based on seeds 0, 44, 56) is reported in Table 6. The learning curves for both Generalized and Personalized accuracies for all datasets for Flow and its baselines are in Figures 5 and 6.

Flow sees an improvement of 1.11-3.46% in Acc g and 1.33-4.58% in Acc p over the best performing baseline. Besides the main observations listed in Section 5, we discuss results on the CIFAR100 dataset here. For CIFAR100 (0.6), Flow (40.08% ± 0.27%) matches the personalized accuracy of the highest performing baseline, PARTIALFED (40.18% ± 0.19%), while achieving 1 . 98% point increase in generalized accuracy. And for CIFAR100 (0.1), Flow improves personalized accuracy by 1 . 78% points. For generalized accuracy, Flow ( 34 . 00% ± 0 . 32% ) reaches close to the best performing baseline, PARTIALFED ( 34 . 79% ± 0 . 29% ). The reason behind the on-par performance of Flow with PARTIALFED can be attributed to the statefulness of PARTIALFED.

0.7

200

400

600

800

1000

1200

1400

Number of Rounds

0.30

0.35

0.40

0.45

0.50

0.55

Validation Accuracy (Generalized)

FedAvg

FedAvg + FT

APFL

Ditto

knnPer

FedRep

LGFedAvg

HypCluster

Flow

200

400

600

800

1000

1200

1400

Number of Rounds

0.65

0.70

0.75

0.80

0.85

0.90

Validation Accuracy (Generalized)

FedAvg

FedAvg + FT

APFL

Ditto

knnPer

FedRep

LGFedAvg

HypCluster

Flow

(b) EMNIST

0

500

1000 1500 2000 2500 3000 3500 4000

Number of Rounds

0.1

0.2

0.3

0.4

0.5

0.6

Validation Accuracy (Generalized)

FedAvg

FedAvg + FT

APFL

Ditto

knnPer

PartialFed

FedRep

LGFedAvg

HypCluster

Flow

(c) CIFAR10 (0.1)

0

500

1000 1500 2000 2500 3000 3500 4000

Number of Rounds

0.3

0.4

0.5

0.6

0.7

Validation Accuracy (Generalized)

FedAvg

FedAvg + FT

APFL

Ditto

knnPer

PartialFed

FedRep

LGFedAvg

HypCluster

Flow

(d) CIFAR 10 (0.6)

0

500

1000 1500 2000 2500 3000 3500 4000

Number of Rounds

0.05

0.10

0.15

0.20

0.25

0.30

0.35

Validation Accuracy (Generalized)

FedAvg

FedAvg + FT

APFL

Ditto

knnPer

PartialFed

FedRep

LGFedAvg

HypCluster

Flow

(e) CIFAR100 (0.1)

0

500

1000 1500 2000 2500 3000 3500 4000

Number of Rounds

0.10

0.15

0.20

0.25

0.30

0.35

0.40

Validation Accuracy (Generalized)

FedAvg

FedAvg + FT

APFL

Ditto

knnPer

PartialFed

FedRep

LGFedAvg

HypCluster

Flow

(f) CIFAR100 (0.6)

250

500

750

1000 1250 1500 1750 2000

Number of Rounds

0.10

0.15

0.20

0.25

Validation Accuracy (Generalized)

FedAvg

FedAvg + FT

APFL

Ditto

knnPer

FedRep

LGFedAvg

HypCluster

Flow

(g) Stackoverflow

With the assumption of full device participation, PARTIALFED makes use of each client's previous state of the personalized model to further train its layer-wise model building policy. With Flow , both the assumptions of full device participation and statefulness of the personalized model are not necessary. Since the clients do not necessarily have to carry their personalized model states to the upcoming rounds, the personalized model recreated by Flow might be unable to compete against stateful approaches like PARTIALFED. Although because of the per-instance routing, Flow still manages to outperform PARTIALFED for the CIFAR10 (0.1/0.6) datasets, and gives comparable performance for the CIFAR100 (0.1/0.6) datasets.

200

400

600

800

1000

1200

1400

Number of Rounds

0.35

0.40

0.45

0.50

0.55

Validation Accuracy (Personalized)

FedAvg + FT

APFL

Ditto

knnPer

FedRep

LGFedAvg

HypCluster

Flow

200

400

600

800

1000

1200

1400

Number of Rounds

0.4

0.5

0.6

0.7

0.8

0.9

Validation Accuracy (Personalized)

FedAvg + FT

APFL

Ditto

knnPer

FedRep

LGFedAvg

HypCluster

Flow

(b) EMNIST

0

500

1000 1500 2000 2500 3000 3500 4000

Number of Rounds

0.625

0.650

0.675

0.700

0.725

0.750

0.775

Validation Accuracy (Personalized)

FedAvg + FT

APFL

Ditto

knnPer

PartialFed

FedRep

LGFedAvg

HypCluster

Flow

(c) CIFAR10 (0.1)

0

500

1000 1500 2000 2500 3000 3500 4000

Number of Rounds

0.45

0.50

0.55

0.60

0.65

0.70

0.75

Validation Accuracy (Personalized)

FedAvg + FT

APFL

Ditto

knnPer

PartialFed

FedRep

LGFedAvg

HypCluster

Flow

(d) CIFAR 10 (0.6)

0

500

1000 1500 2000 2500 3000 3500 4000

Number of Rounds

0.225

0.250

0.275

0.300

0.325

0.350

0.375

0.400

0.425

Validation Accuracy (Personalized)

FedAvg + FT

APFL

Ditto

knnPer

PartialFed

FedRep

LGFedAvg

HypCluster

Flow

(e) CIFAR100 (0.1)

0

500

1000 1500 2000 2500 3000 3500 4000

Number of Rounds

0.15

0.20

0.25

0.30

0.35

0.40

Validation Accuracy (Personalized)

FedAvg + FT

APFL

Ditto

knnPer

PartialFed

FedRep

LGFedAvg

HypCluster

Flow

(f) CIFAR100 (0.6)

250

500

750

1000 1250 1500 1750 2000

Number of Rounds

0.10

0.15

0.20

0.25

Validation Accuracy (Generalized)

FedAvg + FT

APFL

Ditto

knnPer

FedRep

LGFedAvg

HypCluster

Flow

(g) Stackoverflow

D.2 Effectiveness of the Alternative Training of the Global Model

We experimented with three modes of training for the global model: (a) Global model trained first, then the policy module, (b) Policy module trained first, then the global model, (c) Global model and Policy module trained alternatively. The results are shown in Figure 7. We see that the alternate training results in a more stable training compared to the other two modes of training. These results are in conformance with other works [49, 50] which have also used alternative training for policy and model weights training.

D.3 Dataset Split

We experimented with the Stackoverflow Next Word Prediction task on initial rounds for the local and global training dataset splits. The plot is shown in Figure 8. We observe that for the dataset split size of 0.75:0.25 for local and global datasets respectively, the global model does not get sufficient samples to converge, resulting in worse personalized model performance since the personalized model is based on the global model. While a split of 0.25:0.75 for local and global datasets has closer performance to that of a 0.50:0.50 split, lesser data (and hence fewer iterations) to the local model leads to local model weights being similar to that of global model, diminishing the impact of personalization.

D.4 More Seeds

We have re-run Stackoverflow experiment with seeds in {1, 2, 3, 4, 5 , 6, 7, 8, 9}, the results show no material difference. See Figure 9.

200

400

600

800

1000

1200

1400

Number of Rounds

0.65

0.70

0.75

0.80

0.85

0.90

Validation Accuracy (Generalized)

Policy Module First

Alternate

100

150

200

250

300

350

400

450

500

Number of Rounds

0.14

0.16

0.18

0.20

0.22

0.24

0.26

0.28

0.30

0.32

Validation Accuracy (Personalized)

|

m

,

| =

1

2

|

m

|

|

m

,

| =

1

4

|

m

|

|

m

,

| =

3

4

|

m

|

250

500

750

1000 1250 1500 1750 2000

Number of Rounds

0.14

0.16

0.18

0.20

0.22

0.24

0.26

0.28

0.30

0.32

Validation Accuracy (Personalized)

Seed = 1

Seed = 2

Seed = 3

Seed = 4

Seed = 5

Seed = 6

Seed = 7

Seed = 8

Seed = 9

D.5 Percentage of Clients Benefiting from Personalization

In this section we discuss the effect of personalization, by comparing each client's performance on their individual personalized models with their performance on the global model. The evaluation, just as in section D.1, is done on the test datasets of all the clients. The goal with any personalization method is to make each client's personalized model more beneficial (for us, in terms of accuracy) compared to the global model. Hence we want Acc p > Acc g , to incentivize personalization for each client. As shown in Table 7, compared to the best performing baseline, Flow improves the utility of personalization by up to 3.31% points.

D.6 Breakdown of Correctly Classified Instances

Here we show a detailed view of how instances (across all the clients) get classified correctly between global and personalized models for each of the baselines. For the plots in Figures 10, y -axis represent % of instances correctly classified by (a) Both the global and the personalized models ( both-correct ), (b) Only the global model ( global-only ), and (c) Only the personalized model ( personalized-only ). This % of instances metric is averaged across all clients, and is based on their test datasets. The goal here is to increase the % of instances for both-correct and personalized-only , and reduce the % of instances for global-only . We make the following observations for each of the datasets: Since Flow improves both the generalized and personalized accuracies, we see higher both-correct for Stackoverflow (by 2.75% points), Shakespeare (by 4.34% points), EMNIST (by 3.17% points), CIFAR10 (0.1) (by 5.24% points), CIFAR10 (0.6) (by 0.03% points), CIFAR100 (0.1) (by 0.63% points) and CIFAR100 (0.6) (by 2.78% points).

Due per-instance personalization, we see improvements in personalized accuracy, but those improvements are also included in the both-correct bars, so solely comparing personalized-only bar lengths is not a right comparison. Similarly, we see fewer instances in global-only bars due to the increase in instances which fall under both-correct .

D.7 Analysis of Routing Decisions

Now we show probability value analysis of the routing policy for CIFAR10/100 datasets. Here we have fixed the client as the client which had the highest loss difference between its global and personalized models for Flow . This analysis was done during the inference stage, on the test dataset of the above-mentioned client. The box plots show statistics on the probability of picking the global route for all the instances. Echoing the observations made in Section 5, in Figure 11, we see a trend in increasing probability for the global parameters for the instances which are correctly classified by only the global model. In the contrary, for the instances which

FedAvg + FT

APFL

Ditto

knnPer

FedRep

LGFedAvg

HypCluster

Flow

Baselines

0.0

0.1

0.2

0.3

0.4

0.5

0.6

% of Instances

both-correct

personalized-only

global-only

FedAvg + FT

APFL

Ditto

knnPer

FedRep

LGFedAvg

HypCluster

Flow

Baselines

0.0

0.2

0.4

0.6

0.8

1.0

% of Instances

both-correct

personalized-only

global-only

(b) EMNIST

FedAvg + FT

APFL

Ditto

knnPer

PartialFed

FedRep

LGFedAvg

HypCluster

Flow

Baselines

0.0

0.1

0.2

0.3

0.4

0.5

0.6

0.7

0.8

% of Instances

both-correct

personalized-only

global-only

(c) CIFAR10 (0.1)

FedAvg + FT

APFL

Ditto

knnPer

PartialFed

FedRep

LGFedAvg

HypCluster

Flow

Baselines

0.0

0.1

0.2

0.3

0.4

0.5

0.6

0.7

0.8

% of Instances

both-correct

personalized-only

global-only

(d) CIFAR 10 (0.6)

FedAvg + FT

APFL

Ditto

knnPer

PartialFed

FedRep

LGFedAvg

Flow

Baselines

0.0

0.1

0.2

0.3

0.4

% of Instances

both-correct

personalized-only

global-only

(e) CIFAR100 (0.1)

FedAvg + FT

APFL

knnPer

PartialFed

FedRep

LGFedAvg

HypCluster

Flow

Baselines

0.0

0.1

0.2

0.3

0.4

% of Instances

both-correct

personalized-only

global-only

(f) CIFAR100 (0.6)

can only be classified by the personalized model, the probability for taking the global route is lower as the input passes through more layers.

D.8 Ablation Study: Regularization

Figures 12 and 13 show the validation curves for generalized and personalized accuracy with and without the regularization term used in the policy learning objective as shown in Equation 4. With regularization, we see an improvement of 2.18% (Stackoverflow), 1.86% (Shakespeare), 3.98% (EMNIST), 2.55% (CIFAR10 0.1), 4.36% (CIFAR10 0.6), 0.91% (CIFAR100 0.1), 3.46% (CIFAR100 0.6) for the generalized accuracy. And for the personalized accuracy, we see an improvement of 1.92% (Stackoverflow), 2.02% (Shakespeare), 3.01% (EMNIST), 0.65% (CIFAR10 0.1), 3.98% (CIFAR10 0.6), 2.42% (CIFAR100 0.1), 2.19% (CIFAR100 0.6).

D.9 Ablation Study: Per-instance Personalization

Figure 14 show the validation curves for 3 Flow variants: (a) Per-instance Per-client Flow , which is the primary design proposed in this work, (b) Per-instance Flow , which makes choices between two global routes solely based on each client's instances, (c) Per-client Flow , which is simply FEDAVGFT where the personalization only depends on a client, and not on any specific instances.

With all the datasets, we see a trend of Per-instance Flow outperforming Per-client Flow by 1.88% (Stackoverflow), 0.82% (Shakespeare), 5.07% (EMNIST), 2.90% (CIFAR10 0.1), 2.41% (CIFAR10 0.6), 7.52% (CIFAR100 0.1), 1.09% (CIFAR100 0.6). We also see the trend of Per-instance Flow outperforming Per-Instance Per-Client Flow by 3.19% (Stackoverflow), 1.24% (Shakespeare), 0.94% (EMNIST), 0.55% (CIFAR10 0.1), 4.49% (CIFAR10 0.6), 3.88% (CIFAR100 0.1), 1.37% (CIFAR100 0.6).

D.10 Ablation Study: Soft versus Hard Policy

Table 8 shows the personalized accuracy of the test clients while using soft and hard policies during inference. We see that the accuracy difference between the two designs are statistically insignificant. Hence, using a hard

1

2

3

4

5

6

Layer indices

0.0

0.2

0.4

0.6

0.8

1.0

(Probability of picking the global route)

both-correct

global-only

personalized-only

(a) CIFAR10 (0.1)

1

2

3

4

5

6

Layer indices

0.2

0.4

0.6

0.8

1.0

(Probability of picking the global route)

both-correct

global-only

personalized-only

(b) CIFAR 10 (0.6)

1

2

3

4

5

6

Layer indices

0.3

0.4

0.5

0.6

0.7

0.8

0.9

1.0

(Probability of picking the global route)

both-correct

global-only

personalized-only

(c) CIFAR100 (0.1)

1

2

3

4

5

6

Layer indices

0.2

0.4

0.6

0.8

1.0

(Probability of picking the global route)

both-correct

global-only

personalized-only

(d) CIFAR100 (0.6)

policy for inference not only saves half the compute resources, but also doesn't affect the personalized model's performance.

D.11 Ablation Study: Dynamic Routing

The accuracy curves are given in Figure 15. The curves show that the phenomena of 'dynamic probabilities based on each instance works consistently better than the fixed probabilities for all the clients throughout the training' is indeed observable across all the datasets. The intuition behind that is discussed in Section 5.2, under 'Dynamic Routing'.

200

400

600

800

1000

1200

1400

Number of Rounds

0.48

0.50

0.52

0.54

0.56

Validation Accuracy (Generalized)

of Flow

Flow with Regularization

Flow without Regularization

200

400

600

800

1000

1200

1400

Number of Rounds

0.55

0.60

0.65

0.70

0.75

0.80

0.85

0.90

Validation Accuracy (Generalized)

of Flow

Flow with Regularization

Flow without Regularization

(b) EMNIST

0

500

1000 1500 2000 2500 3000 3500 4000

Number of Rounds

0.54

0.56

0.58

0.60

0.62

0.64

0.66

0.68

Validation Accuracy (Generalized)

of Flow

Flow with Regularization

Flow without Regularization

(c) CIFAR10 (0.1)

0

500

1000 1500 2000 2500 3000 3500 4000

Number of Rounds

0.45

0.50

0.55

0.60

0.65

0.70

Validation Accuracy (Generalized)

of Flow

Flow with Regularization

Flow without Regularization

(d) CIFAR 10 (0.6)

0

500

1000 1500 2000 2500 3000 3500 4000

Number of Rounds

0.26

0.28

0.30

0.32

0.34

Validation Accuracy (Generalized)

of Flow

Flow with Regularization

Flow without Regularization

(e) CIFAR100 (0.1)

0

500

1000 1500 2000 2500 3000 3500 4000

Number of Rounds

0.15

0.20

0.25

0.30

0.35

0.40

Validation Accuracy (Generalized)

of Flow

Flow with Regularization

Flow without Regularization

(f) CIFAR100 (0.6)

250

500

750

1000 1250 1500 1750 2000

Number of Rounds

0.20

0.21

0.22

0.23

0.24

0.25

0.26

0.27

Validation Accuracy (Generalized)

of Flow

Flow with Regularization

Flow without Regularization

(g) Stackoverflow

200

400

600

800

1000

1200

1400

Number of Rounds

0.48

0.50

0.52

0.54

0.56

Validation Accuracy (Personalized)

of Flow

Flow with Regularization

Flow without Regularization

200

400

600

800

1000

1200

1400

Number of Rounds

0.6

0.7

0.8

0.9

Validation Accuracy (Personalized)

of Flow

Flow with Regularization

Flow without Regularization

(b) EMNIST

0

500

1000 1500 2000 2500 3000 3500 4000

Number of Rounds

0.64

0.66

0.68

0.70

0.72

0.74

0.76

0.78

Validation Accuracy (Personalized)

of Flow

Flow with Regularization

Flow without Regularization

(c) CIFAR10 (0.1)

0

500

1000 1500 2000 2500 3000 3500 4000

Number of Rounds

0.650

0.675

0.700

0.725

0.750

0.775

0.800

Validation Accuracy (Personalized)

of Flow

Flow with Regularization

Flow without Regularization

(d) CIFAR 10 (0.6)

0

500

1000 1500 2000 2500 3000 3500 4000

Number of Rounds

0.39

0.40

0.41

0.42

0.43

0.44

0.45

Validation Accuracy (Personalized)

of Flow

Flow with Regularization

Flow without Regularization

(e) CIFAR100 (0.1)

0

500

1000 1500 2000 2500 3000 3500 4000

Number of Rounds

0.15

0.20

0.25

0.30

0.35

0.40

Validation Accuracy (Personalized)

of Flow

Flow with Regularization

Flow without Regularization

(f) CIFAR100 (0.6)

250

500

750

1000 1250 1500 1750 2000

Number of Rounds

0.20

0.22

0.24

0.26

0.28

0.30

Validation Accuracy (Personalized)

of Flow

Flow with Regularization

Flow without Regularization

(g) Stackoverflow

200

400

600

800

1000

1200

1400

Number of Rounds

0.48

0.50

0.52

0.54

0.56

Validation Accuracy (Personalized)

of Flow variants

Per-instance Per-client Flow

Per-instance Flow

Per-client Flow

200

400

600

800

1000

1200

1400

Number of Rounds

0.60

0.65

0.70

0.75

0.80

0.85

0.90

0.95

Validation Accuracy (Personalized)

of Flow variants

Per-instance Per-client Flow

Per-instance Flow

Per-client Flow

(b) EMNIST

0

500

1000 1500 2000 2500 3000 3500 4000

Number of Rounds

0.62

0.64

0.66

0.68

0.70

0.72

0.74

0.76

0.78

Validation Accuracy (Personalized)

of Flow variants

Per-instance Per-client Flow

Per-instance Flow

Per-client Flow

(c) CIFAR10 (0.1)

0

500

1000 1500 2000 2500 3000 3500 4000

Number of Rounds

0.55

0.60

0.65

0.70

0.75

0.80

Validation Accuracy (Personalized)

of Flow variants

Per-instance Per-client Flow

Per-instance Flow

Per-client Flow

(d) CIFAR 10 (0.6)

0

500

1000 1500 2000 2500 3000 3500 4000

Number of Rounds

0.32

0.34

0.36

0.38

0.40

0.42

0.44

0.46

Validation Accuracy (Personalized)

of Flow variants

Per-instance Per-client Flow

Per-instance Flow

Per-client Flow

(e) CIFAR100 (0.1)

0

500

1000 1500 2000 2500 3000 3500 4000

Number of Rounds

0.15

0.20

0.25

0.30

0.35

0.40

Validation Accuracy (Personalized)

of Flow variants

Per-instance Per-client Flow

Per-instance Flow

Per-client Flow

(f) CIFAR100 (0.6)

250

500

750

1000 1250 1500 1750 2000

Number of Rounds

0.18

0.20

0.22

0.24

0.26

0.28

0.30

Validation Accuracy (Personalized)

of Flow variants

Per-instance Per-client Flow

Per-instance Flow

Per-client Flow

(g) Stackoverflow

200

400

600

800

1000

1200

1400

Number of Rounds

0.35

0.40

0.45

0.50

0.55

Validation Accuracy (Personalized)

of Flow variants

Dynamic Routing

Fixed

q

0 = 1.00

Fixed

q

0 = 0.75

Fixed

q

0 = 0.50

Fixed

q

0 = 0.25

200

400

600

800

1000

1200

1400

Number of Rounds

0.60

0.65

0.70

0.75

0.80

0.85

0.90

0.95

Validation Accuracy (Personalized)

of Flow variants

Dynamic Routing

Fixed

q

0 = 1.00

Fixed

q

0 = 0.75

Fixed

q

0 = 0.50

Fixed

q

0 = 0.25

(b) EMNIST

0

500

1000 1500 2000 2500 3000 3500 4000

Number of Rounds

0.65

0.70

0.75

0.80

Validation Accuracy (Personalized)

of Flow variants

Dynamic Routing

Fixed

q

0 = 1.00

Fixed

q

0 = 0.75

Fixed

q

0 = 0.50

Fixed

q

0 = 0.25

(c) CIFAR10 (0.1)

0

500

1000 1500 2000 2500 3000 3500 4000

Number of Rounds

0.60

0.65

0.70

0.75

0.80

Validation Accuracy (Personalized)

of Flow variants

Dynamic Routing

Fixed

q

0 = 1.00

Fixed

q

0 = 0.75

Fixed

q

0 = 0.50

Fixed

q

0 = 0.25

(d) CIFAR 10 (0.6)

0

500

1000 1500 2000 2500 3000 3500 4000

Number of Rounds

0.10

0.15

0.20

0.25

0.30

0.35

0.40

Validation Accuracy (Personalized)

of Flow variants

Dynamic Routing

Fixed

q

0 = 1.00

Fixed

q

0 = 0.75

Fixed

q

0 = 0.50

Fixed

q

0 = 0.25

(e) CIFAR100 (0.1)

0

500

1000 1500 2000 2500 3000 3500 4000

Number of Rounds

0.10

0.15

0.20

0.25

0.30

0.35

0.40

Validation Accuracy (Personalized)

of Flow variants

Dynamic Routing

Fixed

q

0 = 1.00

Fixed

q

0 = 0.75

Fixed

q

0 = 0.50

Fixed

q

0 = 0.25

(f) CIFAR100 (0.6)

250

500

750

1000 1250 1500 1750 2000

Number of Rounds

0.21

0.22

0.23

0.24

0.25

0.26

0.27

Validation Accuracy (Personalized)

of Flow variants

Dynamic Routing

Fixed

q

0 = 1.00

Fixed

q

0 = 0.75

Fixed

q

0 = 0.50

Fixed

q

0 = 0.25

(g) Stackoverflow

E Proofs

E.1 Flow : Detailed

Here we give a detailed version of Flow (Algorithm 2) for proving its convergence properties. Here we are assuming that the global and local model output interpolation is model-wise (after the final layer), not layer-wise.

Algorithm 2: Flow

E.2 Basics

Weperform theoretic analysis of Flow based on the following setup: There are total M clients. A client is denoted by a unique integer m associated with it where m ∈ [ M ] . Each client m has a dataset S m = { ( x ( i ) m , y ( i ) m ); i ∈ [ n m ] } where ( x ( i ) m , y ( i ) m ) has been sampled from D m distribution of the m th client. n m = |S m | is the total sample count of the m th client. Total sample count across all the participating client is n = ∑ m ∈ [ M ] n m . The ratio of m th client's sample count to total sample count is α = n m n .

The global distribution is defined as D = ∑ m ∈ [ M ] q m D m where q m is the weight associated with m th client and ∑ m ∈ [ M ] q m = 1 .

Note that w p,m is a combination of outputs of w g,m (Global parameters) and w ℓ,m (Local parameters) on each layer. For tractability of analysis, we will assume that the combination is only after the last layer. Hence,

The local model update rule is,

where w ( r, 0) ℓ,m = w ( r, 0) g,m = w ( r ) g . Indices r ∈ [ R ] and k ∈ [ K ] are the global round and the local epoch indices.

The policy update rule is,

The global model update rule is,

We list out all the optimization problems relevant to Flow :

Local true risk of the personalized model

where f m is a loss function associated with the m th client.

Local empirical risk of the personalized model

Local true risk of the global model

Local empirical risk of the global model

Local minimizer of local empirical risk of the personalized model

Global true risk of the global model

Global empirical risk of the global model

Local minimizer of global empirical risk

We also use the following assumptions similar to [11, 18, 6]:

Assumption E.1 (Strong Convexity) . f m is µ -convex for µ ≥ 0 . Hence,

We also generalize our convergence analysis for µ = 0 , general convex cases.

Assumption E.2 (Smoothness) . The gradient of f m is β -Lipschitz,

Assumption E.3 (Bounded Local Variance) . h m ( w ) := ∇ f m ( w ( x m ) , y m ) is an unbiased stochastic gradient of f m with variance bounded by σ 2 ℓ .

Assumption E.4 (( G,B )-Bounded Gradient Dissimilarity) . There exists constants G ≥ 0 and B ≥ 1 such that

for a convex f m . And for a non-convex f m ,

The derivation is given in Section D.1 of Scaffold [18].

We also use a definition to quantify the diversity of a client's gradient with respect to the global gradient as defined in [29]:

Definition E.5 (Gradient Diversity) . The difference between gradients of the m th client's true risk and the global true risk based on the global model w is,

E.3 Convergence Proof for the Global Model: Convex (Strong and General) Cases

A client's local update for one local epoch on the global model, starting with w ( r, 0) g,m ← w ( r ) g , is

And a client's local update for K epochs on the global model, would be

In both the above cases, the gradient is with respect to w g parameters.

The global model update is,

We first start with a lemma which binds the deviation between the local model w ( r,K ) ℓ,m and the global model starting point w ( r ) g for it at round r .

Lemma E.6 (Local model progress) . If m th client's objective function f m satisfies Assumptions E.2, E.3, and condition η ℓ ≤ 1 β √ 2 K ( K -1) in Algorithm 2, the following is satisfied:

Proof.

Here we have used triangle inequality and variance separation.

Assuming η ℓ ≤ 1 β √ 2 K ( K -1) , we get

Unrolling the above recursion,

Now we move forward to a lemma which binds the deviation between the local version of the global model w ( r,k ) g,m and the global model starting point w ( r ) g for it round r .

Lemma E.7 (Local version of the global model progress) . If m th client's objective function f m satisfies Assumptions E.1, E.2, E.3, and condition η ℓ ≤ 1 β √ 2 k in Algorithm 2, the following is satisfied:

Proof. We start by expanding w ( r,k ) g,m in terms of its previous epoch iterate.

Using triangle inequality and separation of variance, we get,

Using the convex property of f m , we get

Unrolling the recursion,

Assuming that η ℓ ≤ 1 β √ 2 k we get kη 2 ℓ β 2 ≤ 1 ,

Lemma E.8 (Deviation of the personalized model from the global model) . If m th client's objective function f m satisfies Assumptions E.1, E.2, E.3, and condition η ℓ ≤ min ( 1 β √ 2 K ( K -1) , 1 β √ 2 K ) in Algorithm 2, the following is satisfied:

Proof.

Using lemmas E.6 and E.7,

Theorem E.9 (Convergence of the Global Model for Convex Cases) . If each client's objective function f m satisfies Assumptions E.1, E.2, E.3, E.4 using the learning rate 1 µR ≤ η ℓ ≤ min ( 1 4 √ 10 βBK 2 , 1 8 B 2 β ) in Algorithm 2, then the following convergence holds: (Strong Convex Case)

where q 2 0 / 1 are the probabilities of picking global/local routes averaged over all the instances sampled from the global distribution.

Proof. From the update rules stated in Equations 10 and 11:

Taking squared norm and expectation on both sides with respect to the choice of h m ,

Separating mean and variance according to Lemma 4 of Scaffold [18],

Bounding T 1

Using perturbed strong convexity lemma (Lemma 5) from [18], we get,

Next, using Assumption E.4,

Bounding T 2

Plugging in T 1 and T 2 bounds,

Rearranging the terms, and replacing E || ψ ( r ) g || 2 and E || 1 -ψ ( r ) g || 2 with q 2 0 (probability of picking global route averaged over the instances sampled from the global distribution) and q 2 1 respectively,

Assuming η ℓ K 2 ≥ 80 η 3 ℓ K 4 β 2 B 2 ( η ℓ β +1) = ⇒ η ℓ ≤ 1 4 √ 10 βBK 2 and η ℓ K 2 ≥ 4 η 2 ℓ KβB 2 = ⇒ η ℓ ≤ 1 8 B 2 β , we get

Moving E [ F ( w ( r ) g ) ] -F ( w ∗ g ) to the left-hand side, and rest of the terms on right-hand side,

Unrolling the recursion over R rounds and then using the linear convergence lemma (Lemma 1) for strong convex case from Scaffold [18],

Unrolling the recursion over R rounds and then using the sublinear convergence lemma (Lemma 2) for general convex case from Scaffold [18],

E.4 Convergence Proof for the Global Model: Non-convex Case

We start with a non-convex version of Lemmas E.7 and E.8,

Lemma E.10 (Local version of the global model progress) . If m th client's objective function f m satisfies Assumptions E.2, E.3, in Algorithm 2, the following is satisfied:

Proof. We start by expanding w ( r,k ) g,m in terms of its previous epoch iterate.

Using triangle inequality and separation of variance, we get,

(65)

(68)

Unrolling the recursion,

Lemma E.11 (Deviation of the personalized model from the global model) . If m th client's objective function f m satisfies Assumptions E.2, E.3, and condition η ℓ ≤ 1 2 √ 2 βK in Algorithm 2, the following is satisfied:

Proof.

Using lemmas E.6 and E.10,

Assuming 8 K 2 η 2 ℓ β 2 ≤ 1 = ⇒ η ≤ 1 2 √ 2 βK and unrolling the recursion over w ( r,i -1) p,m -w ( r ) g ,

Theorem E.12 (Convergence of the Global Model for Non-convex Case) . If each client's objective function f m satisfies Assumptions E.2, E.3, E.4, using the learning rate 1 2 β ≤ η ℓ ≤ min ( 1 2 √ 5 βBK 2 , 1 3 √ 40 K 4 β 3 B 2 ) in Algorithm 2, then the following convergence holds:

Proof. From the update rule stated in Equation 11, and β -smoothness of f m , we have

Taking expectation on both sides,

Using Equation 10 for second and third terms, and using the fact that the expectation is with respect to the choice of h m ,

where α m = n m n , which are the weights for weighted aggregation according to the sample count, as shown in Equation 11.

Separating mean and variance according to Assumption E.3,

Using ⟨ a, b ⟩ = -1 2 || a -b || 2 + 1 2 || a || 2 + 1 2 || b || 2 ,

Using Lemma E.11,

Using Assumption E.4 for non-convex case, we get,

Rearranging the terms to put E ∥ ∥ ∥ ∇ F ( w ( r ) g ) ∥ ∥ ∥ 2 on left-hand side,

Assuming 10 K 4 β 2 η 3 ℓ B 2 ≤ η ℓ 2 = ⇒ η ℓ ≤ 1 2 √ 5 βBK 2 and 20 K 4 β 3 η 4 ℓ B 2 ≤ η ℓ 2 = ⇒ η ℓ ≤ 1 3 √ 40 K 4 β 3 B 2 ,

Taking average over all the R rounds,

E.5 Convergence Proof for the Personalized Model: Convex (Strong and General) Cases

Lemma E.13 (Local progress of the personalized model) . If m th client's objective function f m satisfies Assumptions E.1, E.2, E.3, and E.4 and conditioning on η ℓ ≤ 1 β √ 6 K in Algorithm 2, the following are satisfied:

Proof.

Using the convexity of f m ,

and we get,

Using Lemma E.6 and E.7, and smoothness property,

Lemma E.14 (Deviation of local parameters from the aggregated global parameters) . If m th client's objective function f m satisfies Assumptions E.3, E.4, in Algorithm 2, the following is satisfied:

Proof. Stating the aggregate rule from Algorithm 2, Lines 12 ,19 and 20,

Using Lemma 8 from [29] and Lemma E.17,

Lemma E.15 (One epoch progress of the personalized model) . If m th client's objective function f m satisfies Assumptions E.1, E.2, E.3, and E.4 in Algorithm 2, the following are satisfied:

and hence,

Proof.

Using, and,

we get,

From Lemmas E.6 and E.7.

Summing over i = k to K ,

Theorem E.16 (Convergence of the Personalized Model for Convex (Strong and General) Cases) . If each client's objective function f m satisfies Assumptions E.2, E.3, E.4, using the learning rate 1 µR ≤ η ℓ ≤ 1 Kβ 2 in Algorithm 2, then the following convergence holds:

(Strong Convex Case)

where 1 R ∑ R r =1 E ||∇ F ( w ( r ) g ) || 2 is bounded as shown in Theorem E.12.

Proof. We restate the update rules of the personalized model in Algorithm 2,

And using Assumption E.4,

Rearranging the terms,

For strong convex ( µ > 0 ) case, using the linear convergence rate lemma from [18] (Lemma 1) and Definition E.5,

For general convex ( µ = 0 ) case, using the sublinear convergence rate lemma from [18] (Lemma 2), and conditioning on η 2 ℓ K 2 β 4 ≤ 1 = ⇒ η ℓ ≤ 1 Kβ 2 ,

E.6 Convergence Proof for the Personalized Model: Non-convex Case

Lemma E.17 (One round progress of the local model) . If m th client's objective function f m satisfies Assumptions E.3, E.4, in Algorithm 2, the following is satisfied:

Proof.

The last inequality follows from Assumption E.4.

We proceed with a lemma which binds the deviation of the personalized model w p of an arbitrary client m over one round, i.e., w ( r +1) p,m and w ( r ) p,m , for non-convex case.

LemmaE.18 (Local progress of personalized model) . If m th client's objective function f m satisfies Assumptions E.3, E.4, and η ℓ ≤ 1 K √ 12 β ( K -1) , in Algorithm 2, the following is satisfied:

Proof. We start with using the update rule stated for the personalized model at the beginning of Theorem E.16,

Expanding by one iterate,

The inequality was derived from the fact that E || -η ℓ ∇ ψ ( r,k ) g,m f m ( w ( r,k ) p,m ) || 2 = E || ψ ( r,k +1) g,m -ψ ( r,k ) g,m || 2 ≤ 1 . Unrolling the recursion across r ∈ [ R ] , then using Lemmas E.6 and E.10 and Assumption E.4,

Assuming 1 K -1 ≥ 12 K 2 η 2 ℓ β = ⇒ η ℓ ≤ 1 K √ 12( K -1) β ,

Lemma E.19 (One round progress of personalized model) . If m th client's objective function f m satisfies Assumptions E.3, E.4, in Algorithm 2, the following is satisfied:

Proof.

Using the Lemmas E.18 and E.14, we proceed as

Theorem E.20 (Convergence of the Personalized Model for Non-convex Cases) . If each client's objective function f m satisfies Assumptions E.2, E.3, E.4 using the learning rate η ℓ ≤ 1 K √ 12 β in Algorithm 2, then the following convergence holds:

Proof. According to the update rule of Equation 10 and β -smoothness of f m , we have,

Taking expectation on both sides,

Rearranging the terms to put 1 2 E ||∇ f m ( w ( r,K ) p,m ) || 2 at LHS,

Taking an average over all the rounds r ∈ [ R ] ,

Assuming 12 K 2 η 2 ℓ β ≤ 1 ≤ 1 = ⇒ η ℓ ≤ 1 K √ 12 β ,

Plugging in Theorem E.12 to get bounds on ∑ R r =1 E ||∇ F ( w ( r ) g ) || 2 would get us bounds on 1 R ∑ R r =1 E ||∇ f m ( w ( r,K ) p,m ) || 2 .

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper proposes a per-instance and per-client personalization approach Flow that creates personalized models via dynamic routing to improve both the performance of the personalized model and the generalizability of the global model. Convergence analysis for both global and personalized models are provided. Empirical evaluation shows that Flow can achieve better performance in terms of generalization and personalized accuracy on several vision and language tasks in cross-device FL setting.\n\nSTRENGTHS:\nS1: The paper provides two essential observations that limits the performance of existing personalized FL methods, which well motivates the proposal of training dynamic network to benefit from both the global model and the personalized model.\n\nS2: The basic idea behind Flow that improving the performance of the personalized model via dynamic routing is very intuitive. Theoretical analysis and experimental evaluation also show the effectiveness.\n\nS3: The experiments are well designed which not only show the superior performance of the proposed Flow against existing state-of-the-art personalized FL methods, but also give in-depth analysis of how the proposed dynamic routing strategy works.\n\nS4: The paper is well-written and very easy-to-follow.\n\nWEAKNESSES:\nW1: The routing module is not very intuitive to me. Why the model is designed as a fully connected network rather than a sequential model such as RNN? The authors should elaborate more on this issue.\n\nW2: The effectiveness of the alternative training for the global parameters (Equation 5) is not validated. More evaluations are needed.\n\nQUESTIONS:\nQ1: Why the routing model is designed as a fully connected network rather than a sequential model such as RNN? See W1.\n\nQ2: Whether the alternative training of the global parameters is effective? See W2.\n\nQ3: Theorem 4.1 and 4.2 are not the standard form of convergence because they do not show the global optima.\n\nLIMITATIONS:\nSee W1-W2.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-05"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper proposes a per-instance personalized federated learning algorithm called Flow. The algorithm creates dynamic personalized models that adapt to each client's individual data instances. The personalized models allow each instance to determine whether it prefers the local parameters or the global counterpart for making predictions, resulting in improved accuracy. The paper provides theoretical analysis and empirical results demonstrating the superiority of Flow compared to state-of-the-art personalization approaches.\n\nSTRENGTHS:\n1) The paper proposes a novel approach for per-instance personalized federated learning, which is not explored in previous work. \n\n2) The proposed problem “each clients’ data instance is limited to using its personalized model for prediction, while some instances could benefit from better generalization of the global model” is a real PFL problem need to be solved. \n\n3) Both CV and NLP datasets were used for the experiment evaluation. The logic of the paper is clear and easy to understand.\n\nWEAKNESSES:\n1) There are additional computational and communication overheads that are not analyzed inside the paper.\n\n2) some hyperparameters are not introudced: for example, the values of K1 and K2. \n\n3) There are no converge curves compared to other baseline algorithms.\n\nQUESTIONS:\n- What is split data used for?\n- Why is there no \"Ablation on Dynamic Routing Component\" like Figure 4(c)  about other datasets except for stackoverflow. I am very interested in the results, but I cannot find them in the appendix.\n- In Table 2, why did LGFedAvg's Acc_g achieve a relatively high score and even surpass FedAvg at the CIFAR10(0.6) setting? Didn't LGFedAvg only aggregate classification layers? In terms of my own experiments, it is reasonable that it should be worse than the fully aggregated FedAvg.\n- the seed selection \"0, 44, and 56\" seems deliberate to me. \n- Dynamic routing is done by layer-wise to select the layer of the local model or the layer of the global model. My concern is how to ensure that the input space and output space of the local or global model are matched.\n\nLIMITATIONS:\nInstead of using a real-world federal learning dataset, the division is performed on a general NLP/CV dataset. Routing module may cause the risk of leaking client privacy distribution.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 5, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-06"}

comments: SUMMARY:
This paper presents a novel approach to address the data heterogeneity issue in Federated Learning, resulting in improved accuracy for clients. The proposed approach, called Flow, is a per-instance and per-client personalized Federated Learning algorithm that uses dynamic routing to create adaptive personalized models for each client's individual data instances. The paper provides theoretical analysis on the convergence of Flow and empirically demonstrates its superiority in improving clients' accuracy compared to state-of-the-art personalization approaches on both vision and language-based tasks.

STRENGTHS:
The paper presents a novel approach to address the data heterogeneity issue in Federated Learning, which is a significant challenge in real-world applications.
The proposed approach, called Flow, is a per-instance and per-client personalized Federated Learning algorithm that uses dynamic routing to create adaptive personalized models for each client's individual data instances.
The paper provides theoretical analysis on the convergence of Flow, which adds to the credibility of the proposed approach.
The paper empirically demonstrates the superiority of Flow in improving clients' accuracy compared to state-of-the-art personalization approaches on both vision and language-based tasks.
The paper is well-structured and easy to follow, with clear explanations of the proposed approach and experimental results.

WEAKNESSES:
The motivation of this paper is not convincing enough for me, how can per-instance FL solve the challenges is not clear. 
The paper does not provide a detailed analysis of the computational and communication costs of Flow, which could be important factors in practical applications.

QUESTIONS:
see above

LIMITATIONS:
see above
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-08

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposed the Flow, which achieves per-instance and per-client personalized federated learning by dynamic routing between a global model and a personalized model. Flow creates a personalized model that adapts to each client's data instances and distribution. In each federated learning round, clients have both global and local model parameters, which are fine-tuned to fit the local data distribution. The dynamic routing module predicts for each instance using either local or global model parameters based on the instance preference, thereby improving the accuracy of the personalized model. The authors also improve the accuracy of the global model by identifying the instances on the clients that are in agreement with the global data distribution. The extensive experiments have been conducted to evaluate the advantage of the proposed Flow.\n\nSTRENGTHS:\n1) The authors proposed a new personalized FL algorithm that utilizes the dynamic routing strategy to generate dynamic inference models that best fit the local data, which is a new aspect of personalized FL.\n2) The experiments are comprehensive, and the code is included.\n3) The writing is clear.\n\nWEAKNESSES:\n1）The Flow proposed in this paper takes a different approach compared to general local model personalization methods in FL literature. In some typical personalized FL algorithms, a large global model is generated on the server. When the global model is deployed for the personalized requirements of clients, NAS or dynamic routing are adopted to get a smaller model which can better fits the local data distribution (e.g., OFA and AdaptiveNet[1-2]). In contrast, the authors suggest storing both the personalized local models and global models on the client, and dynamically selecting the inference path based on routing decisions, which can adapt to local data distribution and instances. However, this approach introduces additional local storage and computational overhead. Especially when the global model has a large parameter size, the proposed algorithm is limited.\n2) Some metrics proposed in the paper are heuristic. For example, in line 124, the authors divide the local dataset into two parts because they considered that fine-tuning the local model with the entire local data may lead to overfitting.\n3) The selection of baselines lacks some personalized FL algorithms based on NAS or NAS-inspired methods. Flow and these baselines are both algorithms based on path-searching, and it would be interesting to supplement the related comparative experimental results.\n\n4) Some claims are unclear. For example, in the abstract, it is a little confused that the personalized models could achieve lower accuracy than the global model. \nIt seems that the per-instance personalization FL is more fine-grainded than the current personalization FL. However, this would also introduce additional overhead. The authors need to discuss the overhead of this operation. \n\n5) The authors only consider two types of models for the clients to select: i.e., the persoanlized model and the global model. However, in conditional computing paradigm, such as mixture of expert (MoE) framework, each instance can be routed to different kinds of experts.  \n\nReferences:\n1) Once-for-All: Train One Network and Specialize it for Efficient Deployment\n2) AdaptiveNet: Post-deployment Neural Architecture Adaptation for Diverse Edge Environments\n\nQUESTIONS:\n1) Is there any experimental results to support the heuristic step in line 124? What does the proportion of data used for fine-tuning impact the final experimental performance?\n2) How does it compare to the existing flexible personalized FL algorithms, such as some algorithms based on NAS?\n3) What is the additional training time overhead for fine-tuning the local models?\n4) What is the additional time overhead for dynamic routing decisions during inference?\n5) Does the current reuslts can be extended to the more general dynamic routing framework, such as MoE?\n\nLIMITATIONS:\nPlease refer to the weakness discussed above.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-09"}