SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Kunjal Panchal
College of Information and Computer Sciences
University of Massachusetts, Amherst
kpanchal@umass.edu
Sunav Choudhary
schoudha@adobe.com
Nisarg Parikh
Khoury College of Computer Sciences
Boston, MA 02115
parikh.nis@northeastern.edu
Lijun Zhang, Hui Guan
{lijunzhang, huiguan}@umass.edu
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<sup>1</sup>.
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".
$^{1}$ https://github.com/Astuary/Flow
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.
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.
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  $\psi_{g}$  (Line 3), each participating client personalizes and trains  $w_{g}$  in five major stages: (1) Split the training dataset into two halves:  $\zeta_{m,\ell}$  and  $\zeta_{m,g}$  (Line 6).  $\zeta_{m,\ell}$  will be used to update the parameters in the local model parameters while  $\zeta_{m,g}$  will be used to train the parameters in the global model and a routing module. (2) Derive the local parameters  $w_{\ell}$  by finetuning the global parameters  $w_{g}$  for  $K_{1}$  epochs with  $\zeta_{m,\ell}$  (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  $\psi_{g,m}$ , and the local parameters  $w_{\ell}$ . Here, the routing policy  $\psi_{g}$  determines whether the execution
path of an instance should use  $w_{\ell}$  or  $w_{g}$ . (4) Train the routing policy  $\psi_{g,m}$  and the local version of the global parameters  $w_{g,m}$  alternatively for  $K_{2}$  epochs (Lines 9-10) with  $\zeta_{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  $\psi_{g,m}$  back to the server for aggregation (Line 14). Aggregation strategies are orthogonal to this work and we adopt FEDAVG [1].
Table 1: Notations
K	#Epochs for training
M	Total number of sampled clients
m	Client index ∈ [M]
M	Set of available clients
p	Client sampling rate
ηl	Local learning rate
wg	Global model
ψg	Policy module
wg,m	Local version of the global model
wℓ,m	Local model of mth client
wp,m	Personalized model of mth client
Dm	Data distribution of mth client
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_{\ell}$  and global model  $w_{g}$  and the routing module  $\psi_{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_{\ell,m}^{(r)}$  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_{g,m}^{(r)}$ , we finetune it for  $K$  epochs to get  $w_{\ell,m}^{(r)}$ . This local model would be reflective of the client's data distribution  $\mathcal{D}_m$ . Note that we use one half of the dataset to get  $w_{\ell,m}^{(r)}$  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_\ell$ . The update rule for  $w_{\ell,m}^{(r)}$  is:
$$
w _ {\ell , m} ^ {(r, K)} \leftarrow w _ {\ell , m} ^ {(r, 0)} - \eta_ {\ell} \sum_ {k = 1} ^ {K} \nabla f _ {m} \left(w _ {\ell , m} ^ {(r, 0)}; \zeta_ {m, \ell}\right), \text {w h e r e} w _ {\ell , m} ^ {(r, 0)} := w _ {g} ^ {(r)}. \tag {1}
$$
Algorithm 1: Pseudocode of Flow
1 Server randomly initializes the global model  $w_{g}$  and the policy module  $\psi_{g}$    
2 for each round do   
3 Send  $w_{g},\psi_{g}$  to sampled M clients   
4 for  $m\in [M]$  in parallel do   
5  $w_{g,m}\gets w_g;\psi_{g,m}\gets \psi_g;w_{\ell ,m}\gets w_{g,m}$    
6 Creating two mutually exclusive datasets   
7  $\zeta_{m,\ell},\zeta_{m,g}\gets \mathcal{D}_m$    
8 Train  $w_{\ell ,m}$  for  $K_{1}$  epochs   
9 for  $k\in [K_2]$  epochs do Update  $\psi_{g,m}$  according to Equation 4   
10 Update  $w_{g,m}$  according to Equation 5   
11 end   
12 Send back  $w_{g,m},\psi_{g,m},n_m:= |\zeta_{m,g}|$    
13 end   
14 Update  $w_{g}$  and  $\psi_{g}$  with weighted average of each client's  $w_{g,m}$  and  $\psi_{g,m}$
Figure 1: Illustration of the dynamic personalized model design proposed by Flow.
Routing Module. The routing module is a model with fully connected layers, shown as  $\psi_{\mathrm{base}}$  in Figure 1. After each layer  $\psi_{\mathrm{base}}$ , the model has early exits denoted as  $\psi_{\mathrm{exit}}$  which outputs a probability of choosing the layer in the global model  $w_{g}$  or the local model  $w_{\ell}$ . This probability at layer  $j \in [L]$  is computed as,
$$
\mathbf {q} ^ {(j)} = \left[ \mathbf {q} _ {0} ^ {(j)}, \mathbf {q} _ {1} ^ {(j)} \right] = \operatorname {s o f t m a x} \left(\psi_ {\text {e x i t}} ^ {(j)} \left(\psi_ {\text {b a s e}} ^ {(j)} (\hat {x})\right)\right) \in [ 0, 1 ] ^ {2} \text {w h e r e} \left\{ \begin{array}{l} \hat {x} \leftarrow x _ {m} \quad \text {i f} j = 0; \\ \hat {x} \leftarrow \psi_ {\text {b a s e}} ^ {(j - 1)} (\cdot) \text {o t h e r w i s e}, \end{array} \right. \tag {2}
$$
where  $\mathbf{q}_0^{(j)}$  and  $\mathbf{q}_1^{(j)}$  are the probability of picking the global parameter  $w_{g,m}^{(j)}$  and the local parameter  $w_{\ell,m}^{(j)}$  respectively.
In order to train the routing parameters, we compute the training loss based on the output of the personalized model  $w_{p,m}^{(r)}$ , which averages the global model's output and local model's output weighted by  $\mathbf{q}^{(j)}$ , for each layer  $j \in [L]$ :
$$
w _ {p, m} ^ {(j)} (\hat {x}) \leftarrow \mathbf {q} _ {0} ^ {(j)} \cdot w _ {g, m} ^ {(j)} (\hat {x}) + \mathbf {q} _ {1} ^ {(j)} \cdot w _ {\ell , m} ^ {(j)} (\hat {x}) \text {w h e r e} \left\{ \begin{array}{l l} \hat {x} \leftarrow x _ {m} & \text {i f} j = 0; \\ \hat {x} \leftarrow w _ {p, m} ^ {(j - 1)} (\cdot) & \text {o t h e r w i s e .} \end{array} \right. \tag {3}
$$
Let  $f_{m}(\cdot)$  denote the loss function on client  $m$ . Using the above personalized model  $w_{p,m}^{(r)}$ , Flow updates the policy parameters as follows,
$$
\psi_ {g, m} ^ {(r + 1)} \leftarrow \psi_ {g, m} ^ {(r)} - \eta_ {\ell} \nabla_ {\psi_ {g, m} ^ {(r)}} \left[ f _ {m} \left(w _ {p, m} ^ {(r)}; w _ {g} ^ {(r)}, \zeta_ {m, g}\right) - \frac {\gamma}{L} \sum_ {j \in [ L ]} \log \left(\mathbf {q} _ {0} ^ {(j)}\right) \right], \tag {4}
$$
where  $-\frac{\gamma}{L}\sum_{j\in [L]}\log (\mathbf{q}_0^{(j)})$  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  $\psi_{g}$ . Flow updates the global parameters as follows,
$$
w _ {g, m} ^ {(r + 1)} \leftarrow w _ {g, m} ^ {(r)} - \eta_ {\ell} \nabla_ {w _ {g, m} ^ {(r)}} f _ {m} \left(w _ {p, m} ^ {(r)}; \psi_ {g, m} ^ {(r + 1)}, \zeta_ {m, g}\right) \tag {5}
$$
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  $(\mathbf{q}^{(j)}) \in \{0,1\}^2$  for  $\mathbf{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.
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  $\sigma_{\ell}$ -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  $\beta$ -smoothness,  $\sigma_{\ell}$ -bounded local gradient variance,  $(G,B)$ -dissimilarity assumptions, using the learning rate  $\frac{1}{2\beta} \leq \eta_{\ell} \leq \frac{1}{2\sqrt{5}\beta BK^2\sqrt{R}}$  [for non-convex case] in Flow, then the following convergence holds:
$$
\begin{array}{l} \frac {1}{R} \sum_ {r = 1} ^ {R} \mathbb {E} \left\| \nabla F \left(w _ {g} ^ {(r)}\right) \right\| ^ {2} \leq \frac {2}{\eta_ {\ell} \mathbf {q} _ {0} ^ {2} R} \left[ \mathbb {E} \left[ F \left(w _ {g} ^ {(1)}\right) \right] - \mathbb {E} \left[ F \left(w _ {g} ^ {(R + 1)}\right) \right] \right] \\ + \frac {\sigma_ {\ell} ^ {2}}{2 \sqrt {5} B M \mathbf {q} _ {0} ^ {2} K \sqrt {R}} + \frac {2 \mathbf {q} _ {1} ^ {2} G ^ {2}}{\mathbf {q} _ {0} ^ {2} B ^ {2} R} + \frac {\sigma_ {\ell} ^ {2}}{B ^ {2} K R}. \\ \end{array}
$$
Discussion. Here,  $\mathbf{q}_0^2$  (and  $\mathbf{q}_1^2$ ) are the probability of picking the global (and local) weights averaged over all instances sampled from the global data distribution. Given a fixed  $\mathbf{q}_0^2$  and  $\mathbf{q}_1^2$ , 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  $\mathbf{q}_0^2$  in the bound.
One the one hand, when  $\mathbf{q}_0^2\rightarrow 1$ , Flow matches the linear (strong convex) and sub-linear (general convex and non-convex) convergence rates of FEDAVG [18]. Specifically,  $\mathbf{q}_0^2 = 1$  results in convergence rate of  $\mathcal{O}(1 / R^2)$  (strong convex),  $\mathcal{O}(1 / R^{2 / 3})$  (general convex), and  $\mathcal{O}(1 / R)$  (non-convex). When  $\mathbf{q}_0^2 = 1$ , which also indicates  $\mathbf{q}_1^2 = 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  $\mathbf{q}_0^2\rightarrow 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  $\delta_{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  $\beta$ -smoothness,  $\sigma_{\ell}$ -bounded local gradient variance,  $(G, B)$ -dissimilarity assumptions, and using the learning rate  $\eta_{\ell} \leq \frac{1}{K\sqrt{12\beta R}}$  [for non-convex case] in Flow, then the following convergence holds:
$$
\begin{array}{l} \frac {1}{R} \sum_ {r = 1} ^ {R} \mathbb {E} | | \nabla f _ {m} \left(w _ {p, m} ^ {(r, K)}\right) | | ^ {2} \leq \frac {2}{R} \left(\mathbb {E} \left[ f _ {m} \left(w _ {p, m} ^ {(1, K)}\right) \right] - \mathbb {E} \left[ f _ {m} \left(w _ {p, m} ^ {(R, K)}\right) \right]\right) \\ + \mathcal {O} \left(\frac {\beta^ {2}}{R ^ {2} K ^ {2}} \left(\sigma_ {\ell} ^ {2} + \left(\delta_ {m} ^ {\psi_ {g}} + \frac {\delta^ {\psi_ {g}}}{M}\right) K\right) \left(G ^ {2} + \frac {\mathbf {q} _ {1 , m} ^ {2} G ^ {2}}{\mathbf {q} _ {0 , m} ^ {2} R}\right)\right) \\ \end{array}
$$