SlowGuess/NeurIPS_2023_OCR huggingface dataset:
algorithm HYPCLUSTER. And for Flow, along with its variants, we have picked  $\gamma \in \{1\mathbf{e} - 1,1\mathbf{e} - 2,1\mathbf{e} - 3,1\mathbf{e} - 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  $\alpha \in [0,1]$ , where a client is more heterogeneous than  $\alpha \rightarrow 0$ . In this context, heterogeneity refers to the dissimilarity of the dataset instances sampled from a distribution. We conducted experiments on clients with  $\alpha$  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  $\alpha = 0.2$ . And DITTO has a regularization hyperparameter set as  $\lambda = 0.01$ . Flow and its variants have their regularization hyperparameter as  $\gamma \in \{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  $\alpha = 0.1$  and  $\alpha = 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  $\lambda$  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  $\alpha$ . DITTO has 1e-2 as regularization parameter  $\lambda$ . 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  $\gamma$ .
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$ ,  $\psi_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.3M$  FLOPs, while the rest of the model computations take  $2.9M$  FLOPs. Hence the overhead is  $10.34\%$ . For Stackoverflow RNN one epoch, the dynamic routing takes  $0.89M$  FLOPs, while the rest of the model computations take  $12.34M$  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.
Baselines	A	B	C	D	E
FEDAVGFT	|wg|	72.38M	Not Applicable	|wg|	72.38M
KNNPER	|wg| + (#instances * |intermediate representation|)	72.42M	12.46M	|wg|	72.38M
PARTIALFED	|wg| + 2*# layers of wg	72.38M	36.9M	|wg|	72.38M
APFL	3|wg|	217.14M	73.8M	|wg|	72.38M
DITTO	2|wg|	144.76M	36.9M	|wg|	72.38M
FEDREP	|wg|	72.38M	51M	|wg(base)|	70.98M
LGFEDAVG	|wg|	72.38M	10.5M	|wg(head)|	1.39M
HYPCLUSTER	2|wg|	144.76M	36.9M	|wg|	72.38M
FLOW (ours)	2|wg| + |ψg|	145.651M	39.57M	|wg| + |ψg|	73.271M
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
$$
A c c _ {g} = \frac {1}{M} \sum_ {m \in [ M ]} \frac {\sum_ {(x , y) \in S _ {m} ^ {\text {t e s t}}} \mathbb {1} \{y = w _ {g} (x) \}}{S _ {m} ^ {\text {t e s t}}}. \tag {6}
$$
Personalized accuracy is formulated as
$$
A c c _ {p} = \frac {1}{M} \sum_ {m \in [ M ]} \frac {\sum_ {(x , y) \in S _ {m} ^ {\text {t e s t}}} \mathbb {1} \{y = w _ {p , m} (x) \}}{S _ {m} ^ {\text {t e s t}}}. \tag {7}
$$
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.
Table 5: Generalized  $(Acc_g)$  and Personalized  $(Acc_p)$  accuracy (the higher, the better) for Flow and baselines. Variance across different runs is reported in Appendix D, Table 6.
Datasets	Stackoverflow	Shakespeare	EMNIST	CIFAR10 (0.1)	CIFAR100 (0.1)	CIFAR10 (0.6)	CIFAR100 (0.6)
Baselines	Accg	Accp	Accg	Accp	Accg	Accp	Accg	Accp	Accg	Accp	Accg	Accp	Accg	Accp
LOCAL	-	15.93%	-	18.70%	-	28.18%	-	49.78%	-	36.19%	-	62.74%	-	21.31%
FEDAVG	23.15%	-	52.00%	-	85.10%	-	60.98%	-	28.11%	-	67.50%	-	30.33%	-
FEDAVGFT	23.83%	24.41%	52.12%	53.68%	89.57%	90.14%	61.23%	73.03%	29.60%	31.02%	68.19%	72.21%	31.15%	37.24%
KNNPER	23.16%	24.49%	51.87%	53.10%	85.20%	88.28%	59.62%	75.14%	28.08%	33.62%	69.22%	70.14%	30.66%	34.39%
PARTIALFED	-	-	-	-	-	-	62.57%	73.20%	34.79%	40.64%	66.93%	70.38%	37.72%	40.18%
APFL	22.96%	25.70%	52.38%	53.64%	88.40%	89.44%	62.87%	72.86%	31.05%	32.56%	69.53%	72.53%	36.37%	36.74%
DITTO	22.59%	24.36%	52.44%	53.95%	89.08%	91.30%	62.06%	72.06%	28.14%	35.45%	68.12%	70.31%	35.11%	36.07%
FEDREP	18.92%	21.04%	46.71%	50.09%	89.95%	89.77%	64.85%	68.62%	26.10%	33.72%	69.77%	63.61%	28.42%	31.02%
LGFEDAVG	22.61%	24.03%	51.08%	51.43%	87.43%	91.70%	56.63%	73.19%	31.65%	39.63%	67.48%	68.94%	35.01%	33.90%
HYPCLUSTER	23.75%	22.43%	51.92%	52.74%	89.47%	90.49%	63.64%	71.55%	31.57%	33.04%	65.44%	72.40%	34.76%	36.22%
Flow (Ours)	26.64%	29.49%	55.90%	56.20%	90.88%	94.18%	66.26%	76.47%	34.00%	42.42%	70.88%	77.11%	39.70%	40.08%
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\% \pm 0.27\%)$  matches the personalized accuracy of the highest performing baseline, PARTIALFED  $(40.18\% \pm 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\% \pm 0.32\%)$  reaches close to the best performing baseline, PARTIALFED  $(34.79\% \pm 0.29\%)$ . The reason behind the on-par performance of Flow with PARTIALFED can be attributed to the statefulness of PARTIALFED.
Table 6: Variance of generalized and personalized accuracies across 3 different runs (seeds = 0, 44, 56) for Flow and its baselines.
Datasets	SO NWP	Shakespeare	EMNIST	CIFAR10 (0.1)	CIFAR100 (0.1)	CIFAR10 (0.6)	CIFAR100 (0.6)
Baselines	Accg	Accp	Accg	Accp	Accg	Accp	Accg	Accp	Accg	Accp	Accg	Accp	Accg	Accp
LOCAL	-	0.25%	-	0.46%	-	1.14%	-	1.56%	-	0.43%	-	0.89%	-	0.25%
FEDAVG	0.07%	-	0.39%	-	1,32%	-	1.12%	-	0.31%	-	0.82%	-	0.15%	-
FEDAVGFT	0.09%	0.26%	0.51%	0.59%	1.16%	1.21%	0.99%	0.89%	0.46%	0.62%	1.10%	1.26%	0.33%	0.42%
KNNPER	0.16%	0.24%	0.36%	0.41%	0.95%	1.02%	1.41%	1.57%	0.34%	0.57%	0.91%	1.06%	0.24%	0.29%
PARTIALFED	-	-	-	-	-	-	1.36%	1.39%	0.29%	0.46%	0.96%	1.97%	0.09%	0.19%
APFL	0.19%	0.20%	0.41%	0.53%	1.41%	1.50%	1.24%	1.31%	0.36%	0.72%	0.70%	0.97%	0.42%	0.59%
DITTO	0.12%	0.15%	0.49%	0.56%	1.12%	1.22%	1.35%	1.41%	0.43%	0.69%	0.84%	0.87%	0.28%	0.34%
FEDREP	0.15%	0.29%	0.50%	0.65%	0.89%	0.94%	0.95%	1.02%	0.59%	0.79%	0.96%	1.14%	0.14%	0.10%
LGFEDAVG	0.08%	0.16%	0.32%	0.56%	1.10%	1.17%	1.21%	1.24%	0.47%	0.51%	0.82%	0.96%	0.23%	0.21%
HYPCLUSTER	0.20%	0.19%	0.56%	0.73%	0.90%	1.13%	1.43%	1.49%	0.39%	0.47%	0.98%	0.76%	0.35%	0.46%
Flow	0.23%	0.28%	0.40%	0.49%	1.16%	1.21%	1.23%	1.25%	0.32%	0.36%	0.78%	0.86%	0.21%	0.27%
(b) EMNIST
(c) CIFAR10 (0.1)
(a) Shakespeare
(d) CIFAR 10 (0.6)
(e) CIFAR100 (0.1)
(f) CIFAR100 (0.6)
(g) Stackoverflow
Figure 5: Learning curves on Generalized Accuracy Metric of Flow and its baselines.
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.
(c) CIFAR10 (0.1)
(f) CIFAR100 (0.6)
(e) CIFAR100 (0.1)
(g) Stackoverflow
Figure 6: Learning curves on Personalized Accuracy Metric of Flow and its baselines.
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.
Figure 7: Results for different modes of training for EMNIST
Figure 8: Earlier stages of Stackoverflow NWP training with different local:global dataset splits
Figure 9: Stackoverflow NWP training with seed choices of [1,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.
Table 7: % of clients for which  $Acc_p > Acc_q$  (the higher, the better).
Stackoverflow	EMNIST	Shakespeare	CIFAR10 (0.1)	CIFAR100 (0.1)	CIFAR10 (0.6)	CIFAR100 (0.6)
FEDAVGFT	79.26%	81.48%	79.00%	97.18%	91.74%	99.33%	88.54%
KNNPER	82.73%	89.97%	68.87%	90.00%	94.71%	90.00%	96.37%
PARTIALFED	-	-	-	88.30%	90.32%	84.80%	98.64%
APFL	69.66%	93.39%	79.22%	87.48%	86.18%	90.63%	92.03%
DITTO	74.59%	79.26%	73.74%	90.52%	91.45%	89.61%	97.45%
FEDREP	91.53%	82.20%	79.78%	92.30%	78.81%	84.64%	99.54%
LGFEDAVG	83.47%	66.16%	88.43%	88.41%	86.39%	89.59%	91.73%
HYPCLUSTER	80.46%	80.70%	74.84%	95.11%	93.70%	98.18%	99.73%
Flow (Ours)	92.74%	96.70%	89.77%	98.33%	97.29%	99.62%	99.75%
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
(d) CIFAR 10 (0.6)
(f) CIFAR100 (0.6)
Figure 10: Different combinations of  $w_{g}$  and  $w_{p}$  accuracies.
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
(a) CIFAR10 (0.1)
(c) CIFAR100 (0.1)
(b) CIFAR 10 (0.6)
(d) CIFAR100 (0.6)
Figure 11: Behavior of  $\psi_g$  for all instances with respect to each layer of a client with highest loss difference between personalized and global models.
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".
Table 8: Test (personalized) accuracy of two of the Flow variants: (a) Soft Policy variant where the probability  $\mathbf{q}$  is continuous in the range of  $\left\lbrack  {0,1}\right\rbrack$  during inference. (b) Hard Policy variant where the probability  $\mathbf{q}$  is discrete over the set  $\{ 0,1\}$  during inference.
Datasets	Stackoverflow	Shakespeare	EMNIST
Soft Policy	29.57% ± 0.22%	57.01% ± 0.53%	94.97% ± 1.06%
Hard Policy	29.49% ± 0.28%	56.20% ± 0.49%	94.18% ± 1.21%
Datasets	CIFAR 10 (0.1)	CIFAR 100 (0.1)	CIFAR 10 (0.6)	CIFAR 100 (0.6)
Soft Policy	77.24% ± 1.30%	42.75% ± 0.30%	77.02% ± 0.90%	39.74% ± 0.13%
Hard Policy	76.47% ± 1.25%	42.42% ± 0.36%	77.11% ± 0.86%	40.08% ± 0.27%
(e) CIFAR100 (0.1)
(g) Stackoverflow
Figure 12: Generalized accuracy of the ablation study on the regularization term used in the policy learning objective.
(e) CIFAR100 (0.1)
(g) Stackoverflow
Figure 13: Personalized accuracy of the ablation study on the regularization term used in the policy learning objective.
(e) CIFAR100 (0.1)
(g) Stackoverflow
Figure 14: Ablation of the dynamic routing component (Per-client Flow), and the local component (Per-instance Flow).
(g) Stackoverflow
Figure 15: Personalized Accuracy of the Ablation Study on the Dynamic Routing Component.
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
Input:  $R$  Total number of rounds,  $r\in [R]$  : Round index,  $M$  Total number of clients,  $m\in [M]$  : Client index,  $\mathcal{M}$  : Set of available clients,  $p$  : Client sampling rate,  $K$  Total local epoch count,  $k\in [K]$  : Epoch index,  $\eta_{\ell}$  Local learning rate,  $w_{g}^{(r)}$  Global model at  $r^{th}$  round,  $w_{g,m}^{(r,k)}$  : mth client's local update of the global model for  $r^{th}$  round and  $k^{th}$  epoch,  $w_{\ell,m}^{(r,k)}$  : mth client's local model for  $r^{th}$  round and  $k^{th}$  epoch,  $w_{p,m}^{(r,k)}$  : mth client's personalized model for  $r^{th}$  round and  $k^{th}$  epoch,  $\psi_{g}^{(r)}$  Global policy model at  $r^{th}$  round,  $\psi_{g,m}^{(r,k)}$  : mth client's routing policy for  $r^{th}$  round and  $k^{th}$  epoch,  $\mathcal{D}_{m}$  Data distribution of  $m^{th}$  client,  $\mathcal{S}_{m}$  Dataset of  $m^{th}$  client,  $\zeta_{m,\ell}$  Dataset used to train  $w_{\ell}$ ,  $\zeta_{m,g}$  Dataset used to train  $w_{g}$  and  $\psi_{g}$  Output:  $w_{g}^{(R + 1)}$  Global model at the end of the training   
1 Server randomly initializes  $w_{g}^{(1)}$    
2 for  $r\in [R]$  round do Sample M clients from  $\mathcal{M}$  with the rate of p Send  $w_{g}^{(r)},\psi_{g}^{(r)}$  to all the clients for  $m\in [M]$  in parallel do  $w_{g,m}^{(r,0)}\gets w_{g}^{(r)};\psi_{g,m}^{(r,0)}\gets \psi_{g}^{(r)};w_{\ell,m}^{(r,0)}\gets w_{g,m}^{(r,0)}$ $\zeta_{m,\ell},\zeta_{m,g}\gets S_m / *$  Creating two mutually exclusive datasets \*/   
9 for  $k\in [K_1]$  epochs do  $\begin{array}{rl}{w_{\ell,m}^{(r,k)}\gets w_{\ell,m}^{(r,k - 1)} - \eta_{\ell}\nabla f_m(w_{\ell,m}^{(r,k - 1)};\zeta_{m,\ell})} & {} \end{array}$    
10 end   
11 for  $k\in [K_2]$  epochs do   
12  $\forall (x_m,y_m)\sim \zeta_{m,g}$  , define  $\tilde{w}_{p,m}^{(r,k - 1)}(x_m)\gets \psi_{g,m}^{(r,k - 1)}(x_m)\cdot w_{g,m}^{(r,k - 1)}(x_m) + (1 - \psi_{g,m}^{(r,k - 1)}(x_m))\cdot w_{\ell,m}^{(r,K)}(x_m)$ $\psi_{g,m}^{(r,k)}\gets \psi_{g,m}^{(r,k - 1)} - \eta_{\ell}\nabla_{\psi_{g,m}^{(r,k - 1)}}\left[f_m(\tilde{w}_{p,m}^{(r,k - 1)};\zeta_{m,g})\right]$    
13  $\forall (x_m,y_m)\sim \zeta_{m,g}$  , define  $\begin{array}{rl}{w_{p,m}^{(r,k - 1)}(x_m)\gets \psi_{g,m}^{(r,k)}(x_m)\cdot w_{g,m}^{(r,k - 1)}(x_m) + (1 - \psi_{g,m}^{(r,k)}(x_m))\cdot w_{\ell,m}^{(r,K)}(x_m)} & {} \end{array}$    
15  $\begin{array}{rl}{w_{g,m}^{(r,k)}\gets w_{g,m}^{(r,k - 1)} - \eta_{\ell}\nabla_{w_{g,m}^{(r,k - 1)}}f_m(w_{p,m}^{(r,k - 1)};\zeta_{m,g})} & {} \end{array}$    
16   
17   
18   
19   
20   
21 end
E.2 Basics
We perform 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 \in [M]$ . Each client  $m$  has a dataset  $S_{m} = \{(x_{m}^{(i)}, y_{m}^{(i)}); i \in [n_{m}]\}$  where  $(x_{m}^{(i)}, y_{m}^{(i)})$  has been sampled from  $D_{m}$  distribution of the  $m^{th}$  client.  $n_{m} = |\mathcal{S}_{m}|$  is the total sample count of the  $m^{th}$  client. Total sample count across all the participating client is  $n = \sum_{m \in [M]} n_{m}$ . The ratio of  $m^{th}$  client's sample count to total sample count is  $\alpha = \frac{n_{m}}{n}$ .