SlowGuess/NeurIPS_2023_OCR huggingface dataset:
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  $\mathbf{q}_{1,m}^2 G^2 / \mathbf{q}_{0,m}^2$ . 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  $\mathbf{q}_{1,m}^2 / \mathbf{q}_{0,m}^2$  would be higher for a heterogeneous client, since the client would get a higher probability of picking the local route ( $\mathbf{q}_{1,m}^2 \to 1$ ). The higher  $G$  and  $\mathbf{q}_{1,m}^2$  results in slower convergence. On the contrary, a homogeneous client would benefit from a low value of  $\mathbf{q}_{1,m}^2 / \mathbf{q}_{0,m}^2$ , 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,  $\delta_m^{\psi_g}$ , linearly affects the personalized model's convergence. Since the policy model is also globally aggregated, a heterogeneous client would have a high  $\delta_m^{\psi_g}$  and need more epochs per round to converge.
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).
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.
Table 2: Generalized  $(Acc_g)$  and Personalized  $(Acc_p)$  accuracy (the higher, the better) for Flow and baselines.
Datasets	Stackoverflow	Shakespeare	CIFAR10 (0.1)	CIFAR10 (0.6)
Baselines	Accg	Accp	Accg	Accp	Accg	Accp	Accg	Accp
LOCAL	-	15.93%	-	18.70%	-	49.78%	-	62.74%
FEDAVG	23.15%	-	52.00%	-	60.98%	-	67.50%	-
FEDAVGFT	23.83%	24.41%	52.12%	53.68%	61.23%	73.03%	68.19%	72.21%
KNNPER	23.16%	24.49%	51.87%	53.10%	59.62%	75.14%	69.22%	70.14%
PARTIALFED	-	-	-	-	62.57%	73.20%	66.93%	70.38%
APFL	22.96%	25.70%	52.38%	53.64%	62.87%	72.86%	69.53%	72.53%
DITTO	22.59%	24.36%	52.44%	53.95%	62.06%	72.06%	68.12%	70.31%
FEDREP	18.92%	21.04%	46.71%	50.09%	64.85%	68.62%	69.77%	63.61%
LGFEDAVG	22.61%	24.03%	51.08%	51.43%	56.63%	73.19%	67.48%	68.94%
HYPCLUSTER	23.75%	22.43%	51.92%	52.74%	63.64%	71.55%	65.44%	72.40%
Flow (Ours)	26.64%	29.49%	55.90%	56.20%	66.26%	76.47%	70.88%	77.11%
Overall, Flow achieves 1.11-  $3.46\%$  higher generalized accuracy and  $1.33 - 4.58\%$  higher personalized accuracy over the best performing baseline. In particular, Flow outperforms KN-NPER, another per-instance per-client personalization approach, by  $1.66 - 6.64\%$  and  $1.33 - 6.97\%$  in generalized and personalized accuracy metrics respectively. KN-NPER 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).
Table 3: % of clients for which  $Acc_p > Acc_g$  (the higher, the better).
Stackoverflow	Shakespeare	CIFAR10 (0.1)	CIFAR10 (0.6)
FEDAVGFT	79.26%	79.00%	97.18%	99.33%
KNNPER	82.73%	68.87%	90.00%	90.00%
PARTIALFED	-	-	88.30%	84.80%
APFL	69.66%	79.22%	87.48%	90.63%
DITTO	74.59%	73.74%	90.52%	89.61%
FEDREP	91.53%	79.78%	92.30%	84.64%
LGFEDAVG	83.47%	88.43%	88.41%	89.59%
HYPCLUSTER	80.46%	74.84%	95.11%	98.18%
Flow (Ours)	92.74%	89.77%	98.33%	99.62%
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.
Figure 2:  $w_{g}$  and  $w_{p}$  accuracies for Stackoverflow.
Figure 3: Behavior of the routing policy from  $\psi_{g}$  for all instances at each layer for Stackoverflow.
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  $\square$ ), (b) correctly classified by the personalized model but not the local model (noted as personalized-only, colored in blue  $\square$ ), and (c) correctly classified by both models (noted as both-correct, colored in green  $\square$  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,
(a) Ablation on the regularization term
(b) Ablation on per-instance personalization design
(c) Ablation on dynamic routing component
Figure 4: Ablation studies on Stackoverflow dataset.
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.
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\% \pm 0.23\%$  generalized accuracy with regularization, compared to  $24.16\% \pm 0.34\%$  without regularization, and  $29.49\% \pm 0.28\%$  personalized accuracy with regularization, compared to  $27.59\% \pm 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  $\mathbf{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  $\mathbf{q}_0 \in [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\% \pm 0.28\%$ , while with soft decisions, we observed  $29.57\% \pm 0.22\%$ . The rest of the datasets show a similar trend (see Appendix D, Table 8).
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.
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.
[1] 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.
[2] Tian Li, Anit Kumar Sahu, Ameet Talwalkar, and Virginia Smith. Federated learning: Challenges, methods, and future directions. IEEE Signal Processing Magazine, 2020.
[3] 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.
[4] Alysa Ziying Tan, Han Yu, Lizhen Cui, and Qiang Yang. Towards personalized federated learning. IEEE Transactions on Neural Networks and Learning Systems, 2022.
[5] 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.
[6] Yuyang Deng, Mohammad Mahdi Kamani, and Mehrdad Mahdavi. Adaptive personalized federated learning. arXiv preprint 2003.13461, 2020.
[7] 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.
[8] Liam Collins, Hamed Hassani, Aryan Mokhtari, and Sanjay Shakkottai. Exploiting shared representations for personalized federated learning. arXiv preprint 2102.07078, 2021.
[9] Krishna Pillutla, Kshitiz Malik, Abdelrahman Mohamed, Michael Rabbat, Maziar Sanjabi, and Lin Xiao. Federated learning with partial model personalization. arXiv preprint 2204.03809, 2022.
[10] 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.
[11] Sashank J. Reddi, Zachary Charles, Manzil Zaheer, Zachary Garrett, Keith Rush, Jakub Konečný, Sanjiv Kumar, and Hugh Brendan McMahan. Adaptive federated optimization. In International Conference on Learning Representations, 2021.
[12] Yishay Mansour, Mehryar Mohri, Jae Ro, and Ananda Theertha Suresh. Three approaches for personalization with applications to federated learning. arxiv preprint 2002.10619, 2020.
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  $\lambda = 0.5$ . For APFL, mixture hyperparameter  $\alpha$  is set to 0.25. DITTO has regularization hyperparameter  $\lambda = 0.1$ . There are 2 clusters by default for HYPCLUSTER. Flow and its variants were tested on the following choices of regularizing hyperparameters  $\gamma \in \{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.  $\lambda$  and  $\alpha$ , the mixture parameters, for KNNPER and APFL respectively, are set to 0.45 and 0.3. The regularization parameter  $\lambda$  for DITTO is set to 0.1. For Flow, the learning rate is set to 0.11 and the regularization parameter is picked from  $\gamma \in \{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  $\lambda = 0.4$ . For APFL, we have the default mixture parameter as  $\alpha = 0.25$ . DITTO has regularization hyperparameter as  $\lambda = 0.1$ . There are 2 clusters for the clustering