SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Communication-cost is one of the major bottlenecks in Federated Learning due to slow connections between clients and the server. In Federated Learning, compression is commonly employed to mitigate communication costs, where clients compress the local updates before transferring to the server. A common compressor is choosing the Top-K coordinates, however, this approach is only good at reducing the upload communication cost. Since clients have different Top-K coordinates, the aggregated updates are often dense. More complicated approaches can save both upload and download communication cost, such as the Count-sketch based compressor [42]. In fact, we can develop a simple yet effective approach to reduce the communication cost based on FedSep.
In our FedSep framework, suppose we have the learning parameter  $\theta \in \mathbb{R}^d$ , and the communication parameter  $\omega \in \mathbb{R}^p$ . We choose  $p \ll d$  to have a communication efficient federated learning algorithm. More specifically, we consider the following formulation:
$$
\min  _ {\omega \in \mathbb {R} ^ {p}} \frac {1}{M} \sum_ {m = 1} ^ {M} \mathcal {L} \left(\theta_ {\omega} ^ {(m)}; \mathcal {D} _ {t r} ^ {(m)}\right) s. t. \theta_ {\omega} ^ {(m)} = \underset {\theta \in \mathbb {R} ^ {d}} {\arg \min } \frac {1}{2} \| S ^ {(m)} \theta - \omega \| _ {2} ^ {2} + \beta \| \theta \| _ {1} \tag {4}
$$
where  $\mathcal{D}_{tr}^{(m)}$  denotes the training distribution of the  $m_{th}$  client;  $\mathcal{L}(\cdot)$  denotes the loss function. In particular,  $S^{(m)}\in \mathbb{R}^{p\times d}$  ( $p\ll d$ ) is a random sketch matrix whose coordinates are sampled from Gaussian distribution. We choose the quadratic optimization problem (LASSO) as the decode function.
Eq. (4) is a special case of Eq. (1):  $\theta$  corresponds to the learning parameter  $y$ ,  $\omega$  corresponds to the communication parameter  $x$ ,  $\mathcal{L}(\theta; \mathcal{D}_{tr}^{(m)})$  corresponds to the learning problem  $f^{(m)}(y)$  and the decoding problem is a LASSO problem. We can solve it through Algorithm 1.
The LASSO problem of the decode operation involves a non-smooth  $L_{1}$  regularization term and we solve it with the standard proximal gradient method [3]. For the encode operation, we have:
$$
E n c ^ {(m)} \{\cdot \} = - \nabla_ {\omega} \theta_ {\omega} ^ {(m)} = \gamma S ^ {(m)} U \left(I - \left(I - \gamma (S ^ {(m)}) ^ {T} S ^ {(m)}\right) U\right) ^ {- 1} \tag {5}
$$
where  $I$  is the identity matrix and  $U = \text{Diag}\{\mathbb{I}_{\gamma \beta}(\theta_{\omega}^{(m)} + \gamma (S^{(m)})^T (\omega - S^{(m)}\theta_{\omega}^{(m)}))\}$  where  $\mathbb{I}_{\gamma \beta}(\cdot)$  is an indicator operator, which outputs 1 if the absolute value of the input is greater than  $\gamma \beta$ . Please refer to Appendix A.1 for more details of the proximal gradient method and the encode operator.
Remark 2. If we set  $\beta = 0$  in Eq. (4), i.e. we remove the sparsity constraints, both encode and decode operators have explicit solutions. More specifically, we have the decode operator  $Dec^{(m)}\{\cdot\} \coloneqq ((S^{(m)})^T S^{(m)})^{-1}(S^{(m)})^T$  and the encode operator  $Enc^{(m)}\{\cdot\} \coloneqq S^{(m)}((S^{(m)})^T S^{(m)})^{-1}$ . If we choose  $S^{(m)} = S$  for all  $m \in [M]$ , we get a linear compressor.
4.2 Model-Heterogeneous Federated Learning
In practical Federated Learning applications, the scale of the model is inherently limited by the on-device resources of the participating clients, which often exhibit significant diversity. As a result, we can choose different scale of models based on the available resources of each individual client. One widely-used approach is the sub-model extraction, i.e. each client selects a part of the server model to perform local training. Different strategies can be used to do extraction: fixed [11, 2], random [6] and roll [1]. In contrast, we provide a data-dependent approach to select the sub-models. As shown by the following formulation:
$$
\min  _ {\omega \in \mathbb {R} ^ {p}} \frac {1}{M} \sum_ {m = 1} ^ {M} \mathcal {L} \left(\theta_ {\omega} ^ {(m)}; \mathcal {D} _ {t r} ^ {(m)}\right) \tag {6}
$$
$$
s.t. \theta_ {\omega} ^ {(m)} = a _ {\omega} ^ {(m)} \odot \omega , a _ {\omega} ^ {(m)} = \underset {a \in \{0, 1 \} ^ {p}} {\arg \min } \mathcal {L} (a \odot \omega ; \mathcal {D} _ {v a l} ^ {(m)}) + \beta \mathcal {R} (T (a), p ^ {(m)} T _ {t o l})
$$
where  $\mathcal{D}_{tr}^{(m)}$  denotes the training distribution;  $\mathcal{L}(\cdot)$  denotes the loss function and  $\theta$  and  $\omega$  are the learning parameter. Note that  $\omega$  denotes the parameter of the full model and  $\theta$  is the parameter of the sub-model. In particular, we denote a mask vector  $a_{\omega}^{(m)}$ , whose value is from  $\{0,1\}$ .  $\odot$  represents the coordinate-wise multiplication.
For the decode operation, we have  $\theta_{\omega}^{(m)} = a_{\omega}^{(m)} \odot \omega$ . So we first need to find an optimal mask  $a_{\omega}^{(m)}$  instead of solving the complicated integer programming as in Eq. (6), we solve the following relaxed continuous problem:
$$
a _ {\omega} ^ {(m)} = \underset {a \in [ 0, 1 ] ^ {p}} {\arg \min } \mathcal {L} (a \odot \omega ; \mathcal {D} _ {v a l} ^ {(m)}) + \beta \mathcal {R} (T (a), p ^ {(m)} T _ {t o l}) \tag {7}
$$
Eq. (7) includes two parts.  $\mathcal{L}(a\odot \omega ;\mathcal{D}_{val}^{(m)})$  measures the loss of the extracted sub-model over  $\mathcal{D}_{val}^{(m)}$ . The second part  $\mathcal{R}(T(a),p^{(m)}T_{tol})$  is called resource loss [15].  $T_{tol}$  denotes FLOPS of
Figure 2: Test Accuracy w.r.t Communication Rate for FedSep and other baseline methods for MNIST Dataset. The first two plots show results under the I.I.D case, and the last two plots show results for the Non-I.I.D case.  $p$  in the second and the fourth plot is the dimension of the communication parameter. The local learning steps are set as  $I = 5$ .
the server model and  $T(a)$  denotes FLOPS of the sub-model. Therefore,  $p^{(m)} \in [0,1]$  controls the size of the sub-model. Following [15], we choose  $\mathcal{R}(\cdot)$  as:  $\mathcal{R}(T(a),p^{(m)}T_{tol}) = \log (max(T(a),p^{(m)}T_{tol}) / p^{(m)}T_{tol})$ . In summary, we choose the 'best' sub-model which keeps the loss value small and meanwhile, satisfies the FLOPS constraint. Note that this is in contrast to the data-independent sub-model selection methods [11, 6, 1].
Next for the encode operator  $Enc\{\cdot\}$ , we follow the definition of Eq. (2) to have:
$$
E n c ^ {(m)} \{\cdot \} = - \nabla_ {\omega} \theta_ {\omega} ^ {(m)} = - \nabla_ {\omega} \left(a _ {\omega} ^ {(m)} \odot \omega\right) = - D i a g \left\{a _ {\omega} ^ {(m)} \right\} - D i a g \left\{\omega \right\} \nabla_ {\omega} a _ {\omega} ^ {(m)} \tag {8}
$$
Where the last equality follows from the chain rule. Finally, combine Eq. (7) and Eq. (8), we can solve Eq. (6) with Algorithm 1.
Remark 3. For the encode operator Eq. (8), we have  $Enc^{(m)}(\cdot) \approx -Diag\{a_{\omega}^{(m)}\}$  by omitting the second term in Eq. (8), where  $\Delta \theta$  denotes the updates to the sub-model  $\theta$ . As shown in Line 21 of Algorithm 1, we have the update of the server model be  $\Delta \omega = \frac{1}{|\mathcal{M}_t|}\sum_{m \in \mathcal{M}_t}\eta_gDiag\{a_{\omega}^{(m)}\} \Delta \theta$ . In other words, the server simply aggregates the updates of the sub-model, this is the aggregation rule widely used in the partial training literature [11, 6, 1].
In this section, we perform numerical experiments to validate the efficacy of FedSep in solving communication-efficient FL and model-heterogeneous FL as discussed in Section 4. The code is written with Pytorch [39], and the Federated Learning environment is simulated via Pytorch.Distributed Package. We used servers with AMD EPYC 7763 64-core CPU and 8 NVIDIA V100 GPUs to run our experiments.
5.1 Experiments for Communication-Efficient Federated Learning
In this section, we provide numerical experiments for the communication-efficient FL task discussed in Section 4.1. We first introduce the experimental settings.
Dataset. We consider MNIST [29] and CIFAR-10 [27] in our experiments. We create the federated version of these datasets by evenly distributing the training samples among 10 clients. We consider an I.I.D setting and a Non-I.I.D setting. For the I.I.D setting, each client has data samples from all 10 classes, while for the Non-I.I.D setting, each client only has data samples from 2 classes.
Model. For experiments related to MNIST, we use a four-layer convolutional neural network. The first three layers contain 32 kernels of size  $3 \times 3$  and the fourth layer contains 32 kernels of size  $2 \times 2$ . For the experiments related to CIFAR-10, we increase the number of kernels to 64. The total number of parameters of the model is around  $10^{5}$ . For the sketch matrix  $S^{(m)}$  of our FedSep, we let all clients choose the same sketch matrix, and we test the dimension of the communication parameter  $p \in \{10, 100, 500, 1000, 2000, 5000, 10000, 20000\}$ .
Baselines and Hyper-parameters. For baselines, we first compare with the uncompressed baseline (the communication parameter and the learning parameter are in the same dimension); Next we compare with the local Top-K compressor (with error feedback [44]) and the Count-Sketch compressor [42]. Please refer to the Appendix for the hyper-parameter selection of FedSep and baselines.
Table 1: Test accuracy comparison between FedSep with other model-heterogeneous FL baseline methods. High data heterogeneity represents  $K = 2$  for CIFAR-10 and  $K = 20$  for CIFAR-100; Lower data heterogeneity represents  $K = 5$  for CIFAR-10 and  $K = 50$  for CIFAR-100.
Method	High Data Heterogeneity	Low Data Heterogeneity
CIFAR-10	CIFAR-100	CIFAR-10	CIFAR-100
KD-based	FedDF [34]	73.81 (± 0.42)	31.87 (± 0.46)	76.55 (± 0.32)	37.87 (± 0.31)
DS-FL [22]	65.27 (± 0.53)	29.12 (± 0.51)	68.44 (± 0.47)	33.56 (± 0.55)
Fed-ET [9]	78.66 (± 0.31)	35.78 (± 0.45)	81.13 (± 0.28)	41.58 (± 0.36)
PT-based	HeteroFL [11]	63.90 (± 2.74)	52.38 (± 0.80)	73.19 (± 1.71)	57.44 (± 0.42)
Federated Dropout [6]	46.64 (± 3.05)	45.07 (± 0.07)	76.20 (± 2.53)	46.40 (± 0.21)
ZeroFL [41]	64.61 (± 2.18)	51.39 (± 0.45)	83.31 (± 0.78)	53.62 (± 0.51)
FedDST [5]	67.65 (± 1.27)	54.21 (± 0.34)	84.57 (± 0.28)	54.97 (± 0.44)
Flash [2]	67.08 (± 1.46)	54.92 (± 0.29)	84.61 (± 0.37)	55.04 (± 0.32)
FedRolex [1]	69.44 (± 1.50)	56.57 (± 0.15)	84.45 (± 0.36)	58.73 (± 0.33)
FedSep (Ours)	71.13 (± 0.94)	58.16 (± 0.25)	84.61 (± 0.37)	61.41 (± 0.29)
Homogeneous (smallest)	38.82 (± 0.88)	12.69 (± 0.50)	46.86 (± 0.54)	19.70 (± 0.34)
Homogeneous (largest)	75.74 (± 0.42)	60.89 (± 0.60)	84.48 (± 0.58)	62.51 (± 0.20)
The experimental results are summarized in Figure 2 for MNIST (Figure 3 in the Appendix for CIFAR-10). As shown by the figures, FedSep outperforms other baselines. In particular, our FedSep achieves much better performance in the high-compression-rate regime and can get similar performance as other baselines in the low-compression-rate regime. Please refer to the Appendix for more ablation studies.
5.2 Experiments for Model-Heterogeneous Federated Learning
In this section, we provide numerical experiments for the task of a model-heterogeneous FL task discussed in Section 4.2. First, we introduce the experimental settings.
Dataset. We consider CIFAR-10 and CIFAR-100 [27] in our experiments. We create Federated datasets by evenly distributing images among clients; we have 100 clients in our experiments. We create data heterogeneity by controlling the number of classes  $K$  a client can have [1]. For CIFAR-10, we test  $K \in \{2,4,5,8\}$ , while for CIFAR-100, we test  $K \in \{5,10,20,50\}$ . Note that smaller values of  $K$  mean higher heterogeneity of the data. For FedSep, the validation and training set are the same.
Model. We choose ResNet-18 [17] in the experiments. Following the setting in [1, 11], we replace the batch normalization layer with static batch normalization and add a scalar module after each convolution layer. Instead of using the coordinate mask as defined in Eq. (6), we select kernels in each convolutional layer of ResNet-18; furthermore, we parameterize the mask  $a$  with a neural network  $HN(\phi)$  [15], i.e.  $a = HN(\phi)$ ; finally, we reuse the masks and only update the masks every two epochs, this stabilizes the training. This reduces the dimension of  $a$  that we need to optimize. In experiments, the size of the submodel  $p$  (Eq. (6)) of a client is randomly chosen from  $\{1,0.5,0.25,0.125,0.0625\}$ , where the ratio  $p$  is w.r.t the server model. For our FedSep, after optimizing Eq. (7), we select the top  $p$  percent kernels of each convolutional layer by the value of the leaned mask  $a$ , and then we only use the selected kernels to perform the training. In the encode operation, we only evaluate the first term of Eq. (8) as stated in Remark 3.
Baselines and Hyper-parameters. We compare with both state-of-the-art Knowledge Distillation-based methods: FedDF [34], DS-FL [22] and Fed-ET [9], and Partial Training Based methods: HeteroFL [11], ZeroFL [41], Federated Dropout [6], FedDST [5], Flash [2] and FedRolex [1]. Fjord [19] gets similar performance as HeteroFL, so we do not include it in the comparison. For ZeroFL, FedDST and Flash, we consider their heterogeneous-model version by varying the compression rate among clients.
The experimental results are summarized in Table 1. As shown in the table, our FedSep gets comparable performance with the Knowledge Distillation based methods, which needs additional public data, furthermore, FedSep outperforms the partial training based methods, including the state-of-the-art FedRolex [1] method. This result shows the effectiveness of the adaptive sub-model extraction strategy used by our FedSep. Our method chooses the most important part of the model
at each step for clients, thus our method converges faster than other data-independent methods. For more experimental results, please refer to the Appendix.
In this work, we propose a novel federated learning framework, i.e. FedSep with separated communication and learning components. Based on the observation that communication and learning set opposite requirements to the parameter, we let clients to have two layers: a communication layer and a learning layer. We formulate the decode operation from the communication layer to the learning layer as solving a minimization problem, therefore, FedSep has a bilevel structure. We propose an efficient algorithm to solve this bilevel problem and also prove that the algorithms converge to a stationary point with rate  $O(\epsilon^{-2})$ . Finally, we apply FedSep to solve the communication-efficient FL task and the model heterogeneous FL task. Numerical experiments show the superior performance of our FedSep over various baselines.
[1] S. Alam, L. Liu, M. Yan, and M. Zhang. Fedrolex: Model-heterogeneous federated learning with rolling sub-model extraction. arXiv preprint arXiv:2212.01548, 2022.
[2] S. Babakniya, S. Kundu, S. Prakash, Y. Niu, and S. Avestimehr. Federated sparse training: Lottery aware model compression for resource constrained edge. arXiv preprint arXiv:2208.13092, 2022.
[3] A. Beck and M. Teboulle. A fast iterative shrinkage-thresholding algorithm for linear inverse problems. SIAM journal on imaging sciences, 2(1):183-202, 2009.
[4] Y. Bengio. Gradient-based optimization of hyperparameters. Neural computation, 12(8):1889-1900, 2000.
[5] S. Bibikar, H. Vikalo, Z. Wang, and X. Chen. Federated dynamic sparse training: Computing less, communicating less, yet learning better. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 36, pages 6080-6088, 2022.
[6] S. Caldas, J. Konečny, H. B. McMahan, and A. Talwalkar. Expanding the reach of federated learning by reducing client resource requirements. arXiv preprint arXiv:1812.07210, 2018.
[7] D. Chen and M. T. Hagan. Optimal use of regularization and cross-validation in neural network modeling. In IJCNN'99. International Joint Conference on Neural Networks. Proceedings (Cat. No. 99CH36339), volume 2, pages 1275-1280. IEEE, 1999.
[8] T. Chen, Y. Sun, and W. Yin. A single-timescale stochastic bilevel optimization method. arXiv preprint arXiv:2102.04671, 2021.
[9] Y. J. Cho, A. Manoel, G. Joshi, R. Sim, and D. Dimitriadis. Heterogeneous ensemble knowledge transfer for training large models in federated learning. arXiv preprint arXiv:2204.12703, 2022.
[10] M. Dagréou, P. Ablin, S. Vaiter, and T. Moreau. A framework for bilevel optimization that enables stochastic and global variance reduction algorithms. arXiv preprint arXiv:2201.13409, 2022.
[11] E. Diao, J. Ding, and V. Tarokh. Heterofl: Computation and communication efficient federated learning for heterogeneous clients. arXiv preprint arXiv:2010.01264, 2020.
[12] C. B. Do, C.-S. Foo, and A. Y. Ng. Efficient multiple hyperparameter learning for log-linear models. In NIPS, volume 2007, pages 377-384. Citeseer, 2007.
[13] J. Domke. Generic methods for optimization-based modeling. In Artificial Intelligence and Statistics, pages 318-326. PMLR, 2012.
[14] L. Franceschi, M. Donini, P. Frasconi, and M. Pontil. Forward and reverse gradient-based hyperparameter optimization. In Proceedings of the 34th International Conference on Machine Learning-Volume 70, pages 1165–1173. JMLR.org, 2017.
[34] T. Lin, L. Kong, S. U. Stich, and M. Jaggi. Ensemble distillation for robust model fusion in federated learning. Advances in Neural Information Processing Systems, 33:2351-2363, 2020.
[35] Y. Lin, S. Han, H. Mao, Y. Wang, and W. J. Dally. Deep gradient compression: Reducing the communication bandwidth for distributed training. arXiv preprint arXiv:1712.01887, 2017.
[36] D. Maclaurin, D. Duvenaud, and R. Adams. Gradient-based hyperparameter optimization through reversible learning. In International Conference on Machine Learning, pages 2113-2122, 2015.
[37] B. McMahan, E. Moore, D. Ramage, S. Hampson, and B. A. y Arcas. Communication-efficient learning of deep networks from decentralized data. In Artificial Intelligence and Statistics, pages 1273-1282. PMLR, 2017.
[38] K. Nandakumar, N. Ratha, S. Pankanti, and S. Halevi. Towards deep neural network training on encrypted data. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition Workshops, pages 0–0, 2019.
[39] A. Paszke, S. Gross, F. Massa, A. Lerer, J. Bradbury, G. Chanan, T. Killeen, Z. Lin, N. Gimelshein, L. Antiga, et al. Pytorch: An imperative style, high-performance deep learning library. In Advances in Neural Information Processing Systems, pages 8024–8035, 2019.
[40] F. Pedregosa. Hyperparameter optimization with approximate gradient. arXiv preprint arXiv:1602.02355, 2016.
[41] X. Qiu, J. Fernandez-Marques, P. P. Gusmao, Y. Gao, T. Parcollet, and N. D. Lane. Zerof! Efficient on-device training for federated learning with local sparsity. arXiv preprint arXiv:2208.02507, 2022.
[42] D. Rothchild, A. Panda, E. Ullah, N. Ivkin, I. Stoica, V. Braverman, J. Gonzalez, and R. Arora. Fetchsgd: Communication-efficient federated learning with sketching. In International Conference on Machine Learning, pages 8253-8265. PMLR, 2020.
[43] A. Shaban, C.-A. Cheng, N. Hatch, and B. Boots. Truncated back-propagation for bilevel optimization. arXiv preprint arXiv:1810.10667, 2018.
[44] S. U. Stich. Local sgd converges fast and communicates little. arXiv preprint arXiv:1805.09767, 2018.
[45] S. U. Stich, J.-B. Cordonnier, and M. Jaggi. Sparsified sgd with memory. arXiv preprint arXiv:1809.07599, 2018.
[46] D. A. Tarzanagh, M. Li, C. Thrampoulidis, and S. Oymak. Fednest: Federated bilevel, minimax, and compositional optimization. arXiv preprint arXiv:2205.02215, 2022.
[47] A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, L. Kaiser, and I. Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017.
[48] W. Wen, C. Xu, F. Yan, C. Wu, Y. Wang, Y. Chen, and H. Li. Terngrad: Ternary gradients to reduce communication in distributed deep learning. arXiv preprint arXiv:1705.07878, 2017.
[49] R. A. Willoughby. Solutions of ill-posed problems (an tikhonov and vy arsenin). SIAM Review, 21(2):266, 1979.
[50] B. Woodworth. The minimax complexity of distributed optimization. arXiv preprint arXiv:2109.00534, 2021.
[51] J. Yang, K. Ji, and Y. Liang. Provably faster algorithms for bilevel optimization. arXiv preprint arXiv:2106.04692, 2021.