PerFedMask: Personalized Federated Learning with Optimized Masking Vectors
Authors: Mehdi Setayesh
Xiaoxiao Li
Vincent W.S. Wong
Venue: ICLR 2023 poster (2023)
Keywords: Computational capability
Data heterogeneity
Masking vectors
Personalized federated learning

Abstract: Recently, various personalized federated learning (FL) algorithms have been proposed to tackle data heterogeneity. To mitigate device heterogeneity, a common approach is to use masking.  In this paper, we first show that using random masking can lead to a bias in the obtained solution of the learning model. To this end, we propose a personalized FL algorithm with optimized masking vectors called PerFedMask. In particular, PerFedMask facilitates each device to obtain its optimized masking vector based on its computational capability before training.  Fine-tuning is performed after training. PerFedMask is a generalization of a recently proposed personalized FL algorithm, FedBABU (Oh et al., 2022). PerFedMask can be combined with other FL algorithms including HeteroFL (Diao et al., 2021) and Split-Mix FL (Hong et al., 2022). Results based on CIFAR-10 and CIFAR-100 datasets show that the proposed PerFedMask algorithm provides a higher test accuracy after fine-tuning and lower average number of trainable parameters when compared with six existing state-of-the-art FL algorithms in the literature. The codes are available at https://github.com/MehdiSet/PerFedMask.

Published as a conference paper at ICLR 2023
PERFEDMASK: PERSONALIZED FEDERATED LEARN- ING WITH OPTIMIZED MASKING VECTORS
Mehdi Setayesh, Xiaoxiao Li, and Vincent W.S. Wong" Department of Electrical and Computer Engineering
The University of British Columbia (setayeshm, x xiaoxiao 1i, vincentx)eece. cub ,
ABSTRACT
Recently, various personalized federated learning (FL) algorithms have been pro- posed to tackle data heterogeneity. To mitigate device heterogeneity, a common approach is to use masking. In this paper, first show that using random mask- ing can lead to a bias in the obtained solution of the learning model. To this end, we propose personalized FL algorithm with optimized masking vectors called PerFedMask. In particular, PerFedMask facilitates each device to obtain its opti- mized masking vector based its computational capability before training. Fine- tuning is performed after training. PerFedMask is generalization of recently proposed personalized FL algorithm, FedBABU (Oh et al., 2022). PerFedMask can be combined with other FI. algorithms including HeteroFL. (Diao et al. 2021) and Split-Mix FI. (Hong et al. 2022). Results based on CIFAR 10 and CIFAR- 100 datasets show that the proposed PerFedMask algorithm provides a higher test ac- curacy after fine-tuning and lower average number of trainable parameters when compared with six existing state-of-the-art FL algorithms in the literature. The codes are available at bettps:/ggitubb com/ Setehiiieee/Peredaaaaas
INTRODUCTION
Federated learning (FL) is distributed artificial intelligence (AD) framework, which allows multiple edge devices to train a single model collaboratively (Konečný et al. 2015 McMahan et al. (017)) The model is trained under the orchestration of central server. In typical FL algorithm, each communication round includes the following steps: (1) the edge devices download the latest model from the server be used their local model; (2) each device performs multiple local update iterations for updating the local model based on its local dataset; (3) the devices upload their updated local models to the server; (4) the server computes the new model by aggregating the local models. In practical systems, the devices may have diverse and limited computation, communication, and storage capabilities. Moreover, the local datasets available to the devices may be different in size, and contain non-independent and identically distributed (non-III) data samples across the devices.
Under these heterogeneous settings, the performance of the conventional FL algorithms can de- grade (Wang et al.) 2020; Li et al.] (2221). To handle the case when the data is non-IlI, some al., (020) have introduced new optimization frameworks to obtain a more stable global model for the devices. Another approach to address the data het- erogeneity issue is by designing a personalized model for cach device (Arivazhagan et al. 2019 Fallah al. 2020 Collins et al. 202 Oh et al., (022). In personalized FL algorithms, instead of obtaining a single model for all the devices, an initial model is obtained. This initial model can then be personalized for each device using its local data samples.
works (Li al., [020aa] Karimireddy
To overcome the computation limitation of the heterogeneous devices, one common approach is masking vectors. Masking vectors can be used train only sub-network of the learn- ing model for each device based on the computational capability of that device. Masking vectors can be combined with pruning and freezing methods. Pruning methods utilize masking vectors to keep the important parameters of the learning model and remove those which are unimportant from
"Corresponding author: Vincent W.S. Wong

---

Published as a conference paper at ICLR 2023
Server aggregation
Local head model parameters
Trainable porameters global during local upcanes Frozen perameters in the global model during local updaccs
Edge device
Edge dovine
Edge devie N
Figure I: Illustration of FL system using PerFedMask. The model is decoupled into global model and local head model. The local head model remains unchanged during training. The devices collaboratively train the global model. Some parts of the global model can be frozen for the devices during local updates using the optimized masking vectors. After training, a personalized model is obtained for each device by fine-tuning.
the model architecture. However, leveraging pruning in FL may incur additional communication overhead (Babakniya et al. 2022, Bibikar et al. 2022). Moreover, it results in different model ar- chitectures for the devices (Guo et al. 2016)) This may lead to accuracy loss, particularly when data heterogeneity exists in the system (Hong et al. 2022). In the freezing methods, the masking vectors used freeze some parts of the learning model for each device. Unlike pruning, the masked parameters are not removed but are frozen during local updates. Hence, a more stable FL algorithm obtained without changing the learning model architecture. [idahmeed et al. (202T) and Pfeiffer it al. (2022) have shown that freezing methods can reduce the computational and communication resources required for training the learning model in FL. However, the aforementioned works use heuristic approaches for designing the masking vectors, and do not provide theoretical analysis for their choice. Also, the aforementioned works do not address the data heterogeneity issue in their proposed algorithms. In this work, we aim to answer the following question: By exploiting freezing method in FL, what is a systematic approach to determine the masking vectors which can improve the final test accuracy in setting with data and device heterogeneities?
We first show that using the masking vectors to freeze the model parameters for the devices may lead to bias in the convergence bound. This bias hinder the success of employing masking vectors to tackle the device heterogeneity issue in FL. Using the insights from our analysis, we pro- pose PerFedMask, Personalized Federated Learning with Optimized Masking Vectors (see Fig. [ Specifically, by decoupling the learning model into a global model and local head model, we first freeze the local head model for all the devices. Then, each device based on its computational capability. In our work, the masking vectors are determined before training through minimizing the bias term in the convergence bound. The bias can be miti- gated by this approach. However, it may not be eliminated completely. Thus, after training of the global model, the frozen parameters of the local head model can assist fine- tue-tune the entire learn- ing model for each device. We demonstrate empirically the effectiveness of PerFedMask under the heterogeneous settings when compared with six existing state-of-the-art FL algorithms.
freeze a portion of the global model for
PerFedMask has several distinct advantages: (1) PerFedMask generalizes the recently proposed personalized FL algorithm, FedBABU (Oh et al. 2022). particular, FedBABU special case of PerFedMask when all the devices have the same computational capability. (2) PerFedMask is flexible. Since PerFedMask does not change the model architecture, can be combined with other FL algorithms such HeteroFL. (Diao al. 202T) and Split-Mix FL (Hong et al.) 2022) further improve the performance. (3) PerFedMask can address the objective inconsistency problem, which arises due to different number of local update iterations. Unlike FedNova (Wang al. 2020) which requires the modification of device optimizers to tackle this issue, in PerFedMask, we consider the same number of local update iterations for all the devices, while adjusting the required number of computations for those devices with lower computational capabilities.
2 RELATED WORK
FL Algorithms with non-IID Data. FedAvg (McMahan et al. 2017) is the most popular FL algo- rithm, which aims to find single model for all the devices. However, the local data samples the
2

---

Published as a conference paper at ICLR 2023
devices are usually IIn-IID, which can lead data heterogeneity among the devices. To address this problem. Li et al. (020a) proposed FedProx by introducing a proximal term in the objective function. However, this method does not attempt to find personalized model for each device. To tackle this issue, different personalized FL. algorithms have been proposed (Arivazhagan et al.) 2019 Fallah et al.] 2020; Collins et al. 202T Oh al. (022) and most of them utilize fine-tuning to obtain a higher accuracy, In particular, after training, each device uses its local data samples to fine-tune the obtained model by performing several iterations of stochastic gradient descent (SGD). Personalization by fine-tuning may be limited if we focus on increasing the performance of a single model across all devices (Deng et al., 2020; Oh et al.) 2022 Oh et al. (2022) proposed FedBABU, where after decoupling the learning model into a global model and local head model, the same randomly initialized local head model is considered for all the devices. In FedBABU, the local head model is not updated during training. Nevertheless, the aforementioned methods are not designed specifically to tackle the device heterogeneity issue.
FL Algorithms with Heterogeneous Devices. Since different types of devices (e.g., smartphones, tablets, laptops) may participate in FL, the devices' capabilities vary from one another. Pruning method is one approach to address the device heterogeneity issue. Diang et al. (2022); Babakniya et al.] (2022); Bibikar et al. (2022) proposed different algorithms for consistent mask adjustment at the devices and mask aggregation at the server during training. Some works have considered static pruning at initialization (i.e., before training). Diao et al. (2021) proposed HeteroFI., where the local model parameters for each device are considered be subset of the learning model parameters. Hong et al. (2022) proposed Split-Mix FL, where a large model is split into several small base sub- networks according the model width. Each device selects some of the base models for training. In general, the efficiency of different pruning techniques for obtaining an appropriate masking vector even in the centralized machine learning setting is still under investigation (Liu et al., 2022). Since the pruned connections are no longer available in the model architecture, pruning at initialization may inevitably degrade the performance of FL algorithms. Freezing method is another approach to address the device heterogeneity issue. [Sidahmed et al. (2021); Chen et al., (2021): PPeiiffer et al. (2022) proposed different algorithms, where the masking vectors are used to freeze a portion of the learning model for each device in FL. Hence, while the required computation and communication resources can be reduced for each device, a more stable FL algorithm is obtained without changing the learning model architecture. In this paper, we apply freezing method in FL.. We leverage the optimized masking vectors and fine-tuning to address both data and device heterogeneity issues. Appendix A contains full discussion of the related work.
PROBLEM SETTING
We consider one server and N edge devices. Each device n E [N] = f1, N} has its own set of local data samples Dn, In supervised learning setting, each device aims to find a learning model 0,, RR for predicting the true label yn given the input Fn, (æn: Yn) E Da,, where do denotes the dimension of the learning model. Let (f. represent the expected loss over the data distribution of device n. We have fn (0.) E E(... ((0,) Fn, Vn), where L(0,; Rns Bn) the loss function that measures the prediction error of 0n over data samples (In, 3n) E Dn, and pa. is the distribution over D,,. Formally, the optimization problem min N I fz. is solved in the
conventional FL. All our results can easily be extended to the weighted averaging case, where the devices have different size of data samples. In this work, we study both data and device heterogeneity issues using model personalization and masking vectors.
Model Personalization. In the heterogeneous data setting, the probability distribution Pr. varies across the devices. Unlike conventional FL problems, to obtain a more personalized solution for each device, the learning models 0n, n E [N]. are not equal to each other. Similar to FedBABU (Oh) ct al. 2022). we decouple each learning model 0n into a global representation model w E R and a device-specific head model Φn. E Rdd, where du and dg denote the dimensions of the global model and local head model, respectively. We have ddg. = do. To further improve accuracy after training, we freeze the local head model parameters during training. In particular, before training, Φ is initialized randomly by the server to be used for initialization of all the devices' local head models. We have on = Sg. n E [N]. 6,, remains unchanged during training until the convergence is reached for the global model Wg. The global model is obtained by minimizing the objective
3

---

Published as a conference paper at ICLR 2023
throughout all the communication rounds, where the operator {..-} denotes the concatenation of the learning model parameters. Since Øy and Φn, are never updated during training, with some abuse of notation,
function F(fWg. En (n ox.))
consider the objective function as F(w,)) - - Ex=1 fn(wy), After convergence of wg each device n obtains its personalized local head model Φn by fine-tuning the learning model 0n using its local data samples.
Masking Vectors. In the heterogeneous device setting, devices vary in their computational and communication capabilities. We consider that in each communication round, the devices perform local update iterations. When deploying the same local models on all the devices, some devices with limited computational capability are not able to complete 7 local update iterations and send their final local models to the server for aggregation in a timely manner. To address the device heterogeneity issue, masking vectors are used to freeze a part of the local model parameters by customization. A masking vector m, E {0, 1jdw is selected for each device n E [N] based on its computational capability. During local update iterations, each device only updates those parameters in the global model that correspond to non-zero values the masking vector m,, Other parameters, which correspond to the elements of m., with zero values, are frozen during local updates. Note that in FedBABU (Oh et al. 2022). all the elements of vector m,, are equal to one for all devices n E [M].
Let w (t) denote the local model of device at the beginning of local update iteration i in com- munication round t. At initialization (i.e., i = 1), we set w (t) (t). At local update iteration i 1, the local model of device is updated using SGD follows:
(I) where @ denotes the element-wise product, vn(t) is the learning rate, and (t) is the local batch sample chosen uniformly at random from the local dataset Dn.- After performing local update iterations, each device sends its final local model, i.e., w))+(() to the server. We have
w (t) + w - m/()m,, @ (w,,, (t), (t)),
i=
+-+ +1) 20 m))m,, m. u n
Algorithm in Appendix B describes the DeviceLocalU function based on ① and (2). In the aggregation step, we consider that the server aggregates the received final local models by taking the masking vectors of the devices into account. The global model for the next communication round can thus be determined through stable aggregation of unfrozen parameters, as follows:
Wy (y (1)
Σ kn "w""""""" +t).
(E)(())
where (k,)): = En' (m..)) denotes the 1-th element of vector kn. (()) denotes the set of devices participating in training in communication round 6. Specifically, the server selects a fraction e of N devices at random as participating devices in each communication round. For each device n, vector k, obtained as normalized masking vector using the vectors m. Nt)... Using k. in 3 indicates that the server only aggregates the updated parameters from the participating devices.
4 THEORETICAL RESULTS ON THE CONVERGENCE RATE OF THE GLOBAL MODEL
In this section, analyze the convergence rate of the global model when masking vectors are used by the devices. Without loss generality, we focus on non-convex loss functions. We also present the convergence results for strongly convex loss functions in Appendix C In both cases, assume that the loss functions are smooth. For simplicity. we obtain our convergence results for the full device participation scenarid") (i.c., following assumptions, which are commonly used for obtaining the convergence rate of different
1, (t) [N] for all t). The analysis relies on the
FL algorithms in the literature et 20206] [Reddi et al. 2021| Amiri al. 2022)
Assumption 1. The function (w)) n E [NV, is L-smooth and satisfies:
(w (1) < 2L (f, (t)) fD), n (N), i i= T, Vt, where f denotes the minimum value of fn (w).
Ussing the techniques presented in Ln al. 2020h Karimireddy et al. 2020 the extension to the general
case (i.c., C. 1) would he straightforward.
4

---

Published as a conference paper at ICLR 2023
Assumption 2. (u,, ((, is unbiased stochastic gradient of function fn(w). The vari- ance of the masked stochastic gradients is bounded for each device nt E [N]. We have
E k Sn (wF, (t); (t)) kn @ V.n (t))l S.. E [N], i= 1=.,...,7 Vt.
(5)
Assumption 3. The expected squared 1--orrm of the masked stochastic gradients for all the devices is uniformly bounded. We have
E m, O
Vin (1, (O S n n [N], i= T VL.
(6)
When the masking vectors determined based the computational capability of the devices, define the term 7/- = max; (ka) ) to quantify the degree of device heterogeneity in the network. Note that the full device participation scenario. N S Tn S I,
[AV] In addition, In inversely
proportional to the minimum non-zero element of the vector Σn =p'=1 Map @ M.y Hence, larger implies higher degree of device heterogeneity for device n E [N.. We first present the following quantify lemma, which is derived using Theorem in Fang al. (1994). From Lemma I
the impact of freezing the parameters by masking vectors on the convergence bound.
The following inequality holds for any vectors and E Rd,, for which there exists Q > 0 satisfying Imin (x @ x)11 S Q, and for any vector y E Rd:
Lemma 1.
(a. y ④ 2) max(y)) (z, z) + Q (dmax(y))->
-10--(001 I(w) 1=1
where (..-) denotes the inner product operator TR...
We use Lemma D to prove the following theorem concerning the device heterogeneity effect on the Fl. convergence bound. Devices with lower computational capability partially train the global model due to the zero-valued elements in their masking vectors. We show that employing the masking vectors to address the device heterogeneity issue in FI. leads to a bias term in the convergence bound. However, it does not affect the convergence rate.
Theorem 1. Under Assumptions 7 3 and for smooth and non-convex loss functions, if the total number of communication rounds T is pre-defined and the learning rate 1i(t) is small enough such
that (t) = r S LNI7. we have
T
2
E V Wy (t))1
nTT (F(w,(1)) - -F*) LNT
22
Sn
dbar
1.221 (T -1)(2r -1
24 diw
(kmn)
(8)
6
where VI is constant satisfying max; ( (n(ww(()) ④ V (F(W((())) S v for all E [N], i=
1,...,T, T. and t = =,..., F F(w"), where w is the global optimal point. L, 2 and G are
constants defined in Assumptions
The proof for Lemma I and Theorem can be found in Appendices D and respectively.
By employing the masking vectors in FI, the term En dou'tr Eim (kmn)) ap- pears on the right-hand hand side of (8. Since this term does not seale with the number of communication rounds T, it is considered as a bias term, which remains as residual in the convergence bound. Hence, those FL algorithmss, which use masking vectors to reduce the computational and communi- cation costs for the devices, may converge to a local minimum of the objective function F(wng). In
Remark 1.
PerFedMask, we design the masking vectors by minimizing Σn ddu En (kn)) to miti- gate the performance degradation due to this bias term. Appendix F we present a simple example gain insight regarding the selection masking vectors by minimizing the bias term.
5 OUR PROPOSED ALGORITHM
In this section, we propose a novel algorithm called PerFedMask, which aims to mitigate the perfor- mance degradation caused by bias described in Remark through: (1) systematically designing the
5

---

Published as a conference paper at ICLR 2023
Algorithm 1:
Training Procedure of PerFedMask
1: Input: Local datasets D.. maximum number of trainable parameters Vn,, v E [N]] the number of local
epochs E; the number of local batches B: participation ratio c.
2: Initialize the learning rate v7() and initialize randomly (1))): = (,, Φg). Model initialization
3: "Server Operation"
4: Select masking vector m. for each device by solving prsss Optimized masking using (a (9d)
for each communication round E (I....:7]) do
6: (1)) <- Random subset of max (cN, devices.
Selection of participating devices
7: for each device E (t)) in parallel do
8: w +1 (tt) peaieeLocallpdatee w (t), og mm, Su, D.. 17(1). # Local updates using eqn.
end
10: wa(( + 1) = Enerray An O w (1).
Aggregation at the server using eqn.
Update 7t().
12: end for
13: "Client Operation"
14: for each device n [MV] in parallel do
Fine-tune learning model 0, (w,(T + 1), Φn) using training data samples Da. # Fine-tuning
16: end for
masking vectors via an optimization framework, and (2) fine-tuning the local head models. First, each device determines the maximum number of parameters which can be trained during local up- date iterations based on its computational capability 2 These values are sent to the server. The server then determines the masking vector for each device before training.
Let thn denote the maximum number of parameters that can be trained by device E [N], where Un S dw-- Given th,, the server determines the masking vector m,, for each device n before training by minimizing the bias term. Here, we present the formulation of layer-wise masking) which decides whether or not to freeze all the parameters in each layer of the global model. We adopt layer- wise masking since most of the current machine learning frameworks such as PyTorch (PyTorch] 2022) and TensorFlow (Abadi et al. (0166) run at the granularity a full tensor, with no APIs which can provide parameter freezing masking, the number of optimization variables can be reduced. Thus, the complexity of obtaining the masking vectors layers in the global model, respectively. Let N and |my] denote the set and number of parameters in layer E A, respectively. We define ma E {0, 1 as the layer-wise masking vector. If (mn); = then all the elements (m,),, 1 optimization problem is solved by the server to obtain the masking vectors for the devices:
finer granularity. Also, by considering layer- laser-wise
also be reduced accordingly. Let A and |A] denote the set and number of
π, are equal to one. In PerFedMask, the following
punusk
minimize mn. Cas [N(N]
Σ
du max (kn))
-E
(kkn)) +
n=1
J'EA
(mn))
subject to
(kn)s = Σm==1 (m,)),
j j A, nE
(9a)
|T (m,)) =thn Gns n E [N].
(9b)
(mn)) E {0, 1}, jEA, [E]N)
(9c)
en 2 0, n [N],
(9d)
where en is slack variable, which prevents to train more than vn parameters for each device the layer-wise masking. Problem pmnak is mixed-integer nonlinear program, which NP-hard
and difficult to solve. In Appendix [ we show how to obtain close-to-optimal solution by using
successive convex approximation (Shen et al.) 2016) and MOSEK solver (MOSEK ApS 2022 Algorithm summarizes the training procedure of PerFedMask.
UUsing the curves similar to those in Fig. 3 Appendix each device can determine its maximum number of trainable parameters.
RRandom layer-wise masking has heen considered in Sidahmed et al. 2021 Pfeiffer et al. 2022

---

Published as a conference paper at ICLR 2023
6 EXPERIMENTAL RESULTS
6.1 EXPERIMENT SETUP
Datasets and Model Architectures. We conduct our experiments on CIFAR-10 and CIFAR-100 image classification tasks" Our experiments are performed with ResNet (PreResNet188 (ee et al. 2016) for CIFAR-10, and with MobileNet (Howard et al. (017) for CIFAR-100 In both cases, we set the number of devices 100. The data samples uniformly divided among the devices. Each device has 450 training data samples, 50 validation data samples, and 100 test data samples. The batch size is set to 50. To enable non-IlI data partitioning among the devices, we distribute 3 and 10 classes per device for CIFAR-10 and CIFAR-100 datasets, respectively. The same classes are considered in the training, validation, and test datasets.
Implementation Details. We denote the fraction of devices with the maximum computational capa- bility by v. That is, represents the ratio of devices which can completely update the entire global model during the local update iterations. Those devices are able to train x parameters of the global model. Since the remaining devices have lower computational capability, they should mask some parts of the global model during the local update iterations based on their capabilities. Unless stated otherwise, we set / = 0.5. For all the experiments, the learning rate starts with 0.1 and is decaved by factor of 0.1 in communication round E {5 ДT). Similar to FedBABU (Oh et al. 2022 we fix the product of the local epochs and the maximum number of communication rounds to 320. After T communication rounds, choose the model with the maximum validation accuracy and perform 5 local update iterations for fine-tuning the learning model using each device's local training data samples. We perform the experiments using Py library (PyTorch 2022) in Python 3.7. We apply layer-wise masking in our experiments.
Baselines. We compare the performance our proposed algorithm. PerFedMask, with the follow- ing FL algorithms: FedBABU (Oh et al. (2022) and FedProx (Li et al. 2020a), which have been proposed to tackle data heterogeneity; FedNova (Wang al., 2020), which has been proposed ad- dress objective inconsistency problem; HeteroFL (Diao et al. [2021] and Split-Mix FL (Hong et al., 2022]. which have been proposed to tackle device heterogeneity in the non-IID data settings.
Performance Metries. We consider the test accuracy as one of the performance metrics. Fine- tuning steps are performed for PerFedMask and FedBABU algorithms to personalize the learning model and obtain the local head model for each device. For fair comparison, the obtained learning model has been also fine-tuned for each device in the other FL algorithms. Thus, we report the test accuracy before and after fine-tuning. We also assess the average number of floating-poit operations (FLOPs) in each communication round for both the forward and backward propagation paths to show the required computation for the algorithms. We use PYPAPI (PyPAPI (017) to obtain the number of FLOPs. Since the average number of trainable parameters is equal to the average number of parameters transmitted from the devices the server in each communication round, we report this number to show the communication cost of each FL algorithm.
6.2 BENCHMARK EXPERIMENTS
We consider a heterogeneous device setting, where half of the devices (i.e., devices with maximum computational capability) perform four local epochs (i.e., E 4). Due the limited computational capability, the remaining devices perform two local epochs for updating all the parameters. Using PerFedMask, instead of considering different local epochs for the devices, we address the device heterogeneity issue by freezing some parts of the global model for those devices with lower compu- tational capability. For fair comparison, the number of frozen parameters are selected in a way such that the considered algorithms have the same number of FLOPs for each device.
Performance Comparison with the Baselines. Table T shows the obtained test accuracy after fine-tuning and the number of trainable parameters for CIFAR-10 and CIFAR- 100 datasets Our observations are as follows: (1) our proposed algorithm, PerFedMask, has comparable performance
'We also AlexNet on DomainNet dataset Li al. 2021 and provide the results in Appendix to show the performance under feature non-III configuration. We show some of the training curves in Appendix j Also, in Appendix K present additional statistics for some of the results in Table

---

Published as a conference paper at ICLR 2023
Table 1: Test accuracy after tine-tuning and number of trainable parameters of PerFedMask and the baseline algorithms for CIFAR-I0 and CIFAR- datasets
Test accuracy after fine-tuning
Dataset
PerFedMask (Ours)
FedBABU
FedProx
FodNova
HctcroFL
Split-Mix FL
FedAvg
88.43
88,20
84,96
84.26
87.33
85.56
84,99
CIFAR-10
0.1
83.60
84.27
74.55
71,88
73.34
77.76
71.19
1
72.40
69.01
64.63
65.24
68.65
65.95
65.27
CIFAR-100
0.1
67.47
66.32
59.36
60.42
65.87
62.35
59.12
Number trainable parameters
Dataset
PorFodMask (Ours)
FedBABU
FedProx
FodNova
HeteroFL
Split-Mix FL
FodAvg
CIFAR-I0
6.138M
11.167M
11.172M
11.172M
5.674M
0.793M
11.172M
CIFAR-100
1.803M
3.207M
3.309M
3.309M
1.774M
0.223M
3.309
to FedBABU and outperforms the other baselines in terms of test accuracy after fine-tuning, (2) By increasing the number of devices participating in FL (i.e., by increasing c), higher test accuracy can be achieved. (3) PerFedMask, HeteroFL, and Split-Mix FL can provide lower number of trainable parameters. Split-Mix FL has the lowest number of trainable parameters because it trains several low-width base models instead of the original learning model. (4) Different number of local update iterations for the devices may lead to the objective inconsistency problem. PerFedMask, HeteroFI., and Split-Mix FL. algorithms can address this problem by decreasing the number of FLOPs for the less capable devices. Thus, the same number local update iterations can be considered for all the devices. Also, this problem is tackled in FedNova through modifying the optimizer. However, other algorithms suffer from the objective inconsistency problem, which may degrade their performance.
Combining PerFedMask with HeteroFL and Split-Mix FL. Table Table[]showe showw the performance re- sults for PerFedMask and its combination with HeteroFL and Split-Mix FL algorithms. Table 2 also shows the performance results for FedBABU, HeteroF1., and Split-Mix FI. algorithms. The number of FLOPss for forward and backward indicates the number of required computations in the forward and backward propagation paths, respectively. In general, backpropagation dominates the computa- tional cost during training of a learning model (Xu et al. (022). In PerFedMask, by using masking vectors, there is no need to compute the partial derivative of the objective function with respect to the frozen parameters. Although PerFedMask has reduced the number trainable parameters and the backward FLOPs, it can still achieve higher test accuracy compared to other algorithms including FedBABU. Since PerFedMask does not change the architecture of the learning model, it can easily be combined with other FL algorithms. Table [ shows that combining PerFedMask with Split-Mix FL and HeteroFL algorithms can further reduce the number of FLOPs in the backward path and the number of trainable parameters. This combination provides a higher test accuracy after fine-tuning than Split-Mix FL and HeteroFL algorithms.
Table 2: Performance comparison on CIFAR-10 dataset when
1. Results for CIFAR-100 dataset can be
found in Appendix
Algorithm
Test accuracy
# of trainable parameters
FLOPs
Before fine-tuning
After fine-tuning
Forward
Backward
Per Sphi-Mix FL
51.88
87.74
0.69IM
0.178G
0.514G
ParkedMask HeteroFL
69.44
87.79
5.473M
1.111G
1.721G
ParFedMask
70.14
88.
6.138M
2.182G
2.697G
Split-Mix FL
57.96
85.56
0.793M
0.178G
0.541G
HeteroFL
62.58
87.33
5.674M
1.111G
1.749G
FedBABU
69.27
88.20
11.167M
2.182G
3.406G
6.3 ABLATION STUDIES
Effect of Fine-Tuning Steps. We investigate the impact of the number of fine-tuning steps on the final test accuracy of PerFedMask and FedBABU Fine-tuning steps equal to zero means the test accuracy is obtained before fine-tuning. Also, since the number of batches for each device's training dataset is equal to 9, increasing the fine-tuning steps by one leads to 9 more local update iterations. Results from Table 3 show that similar to FedBABU, PerkedMask can achieve better accuracy with
We also investigate the effect of freezing the local head models in Appendix M
8

---

Published as a conference paper at ICLR 2023
small number of fine-tuning steps. This characteristic is important when fine-tuning is restricted or costly for the devices.
Table Performance according to fine-tuning steps when 1
Fime-tuning sleps
Dataset
Algorithm
8
10
PerkedMaak
70.14
88.78
58.69
88.35
88.25
88.30
CIFAR-10
FedBABU
69.27
88.63
88.38
88
88.12
87.92
PerFedMask
32.04
72.62
72.84
72.48
72.11
72.40
CIFAR-100
FedBABU
29.70
69.03
69.15
68.86
68.83
68.81
Effect of Increasing the Number of Devices with Maximum Computational Capability. We investigate the impact of increasing v on the performance of PerFedMask in Table More devices in the network are able to train the entire global model as y increases. This leads an increase in the number of trainable parameters and number of backward FLOPs in PerFedMask. We can observe that by increasing v, the test accuracy before fine-tuning is improved. Note that PerFedMask can provide a comparable test accuracy after fine-tuning even for = 0.2, when compared with the case which all devices have the maximum computational capability (i.e., 1).
Table 4: Results of increasing for CIFAR-100 dataset when 1.
Algoridum
Test accuracy
# of trainable parameters # of backward FLOPs
Before line-luning
Afier line-tuning
29.29
72.07
0.91IM
0.617G
0.4
32.31
74.33
1.518M
0.675G
PerFedMask
0.6
32.79
72.82
2.095M
0.741G
0.8
33.59
72.64
2.647M
0.803G
34.73
73.76
3.207M
0,863G
Effeet of Masking Vectors Design. In T we consider three masking approaches: in sequential masking, the layers are masked sequentially; random masking. the layers are masked randomly; and in optimized masking, the layers are masked by solving problem 'pmask Optimized masking minimizes the bias described in Remark [ As shown Table 5 optimized masking can provide lower training loss and higher training and test accuracies. These results are compatible with our theoretical analysis. As shown in [3 PerFedMask enhances FL performance before fine-tuning by employing optimal masking vectors. The final test accuracy then be improved by fine-tuning.
Table 5: Results of different approaches for masking vectors design for CIFAR- 100 datasct when and 0.2.
Design upproach
Bias
Training loss
Training accuracy
Test accuracy before fine-tuning
Sequential masking
3.364M
2.324
50.11
27.02
Random masking
0.609M
1.706
64.40
28.88
Optimized masking
0.204M
1.648
66 63
29.80
7 CONCLUSION
In this work, we proposed a flexible and easy to implement personalized FL algorithm called PerFed- Mask. We provided theoretical and empirical grounds to justify the utility of PerFedMask in het- erogeneous data and device settings. In particular, PerFedMask employs (1) optimized masking vectors obtained by minimizing the bias term in the convergence bound, and (2) fine-tuning. The masking vectors are exploited to freeze some parts of the global model for each device based on its computational capability. Fine-tuning performed by each device after training to improve the final test accuracy. When compared with some existing state-of-the-art FL algorithms, PerFedMask can achieve higher test accuracy. can also decrease the average number of trainable parameters and the average number of FLOPs in each communication round without changing the learning model architecture. A future direction to consider freezing priority for different layers in the neural net- work architecture. For example, in Frankle al. (2021). has been shown that batch normalization layers in convolutional networks are important to be considered as the trainable parameters. Also, the approach of masking vectors design in this work can be emulated in pruning methods.

---

Published as a conference paper at ICLR 2023
ACKNOWLEDGMENTS
This work was supported in part by Rogers Communications Canada Inc., Natural Sciences and Engineering Research Council of Canada (NSERC), and Public Safety Canada (NS-5001-22170),
REFERENCES
Martín Abadi, Paul Barham, Jianmin Chen, Zhifeng Chen, Andy Davis, Jeffrey Dean, Matthieu Devin, Sanjay Ghemawat, Geoffrey Irving, Michael Isard, et al. TensorFlow: system for large- scale machine learning. In Proc. of USENIX Symp. on Operating Systems Design and Implemen- lution (OSDI), Savannah, GA, Nov. 2016.
Andrei Afonin and Sai Praneeth Karimireddy. Towards model agnostic federated learning using knowledge distillation. In Proc. of Int'l Conf. Learning Representations (ICLR), Apr. 2022.
Mohammad Mohammadi Amiri, Deniz Gündüüz, Sanjeev R. Kulkarni, and H. Vincent Poor. Con- vergence of federated learning over a noisy downlink. IEEE Trans. Wireless Commun., 21(3): Mar. 2022.
Manoj Ghuhan Arivazhagan, Vinay Aggarwal, Aaditya Kumar Singh, and Sunav Choudhary. Fed- erated learning with personalization layers. arXiv preprint arXiv:1912.00818 Dec. 2019.
Sara Babakniya, Souvik Kundu, Saurav Prakash, Yue Niu, and Salman Avestimehr. Federated sparse training: Lottery aware model compression for resource constrained edge. In Proc. of NeurlPS Workshop Federated Learning: Recent Advances and New Challenges, New Orleans, LA, Dec. 2022.
Sameer Bibikar, Haris Vikalo, Zhangyang Wang, and Xiaohan Chen. Federated dynamic sparse training: Computing less, communicating less, yet learning better. In Proc. of AAAI Conf. on Artif. Intell., Feb. 2022.
Chen Chen, Hong Xu, Wei Wang, Baochun Li, Bo Li, Li Chen, and Gong Zhang. Communication- efficient federated learning with adaptive parameter freezing. In Proc. of Int'l I Conf. on Distributed Computing Systems (ICDCS), Jul. 2021.
Liam Collins, Hamed Hassani, Aryan Mokhtari, and Sanjay Shakkottai. Exploiting shared repre- sentations for personalized federated learning. In Proc. of Int't II'I Conf. Mach. Learn. (ICML), Jul. 2021.
Yuyang Deng, Mohammad Mahdi Kamani, and Mehrdad Mahdavi. Adaptive personalized federated learning. arXiv preprint arXiv:2003.13461 2020.
Enmao Diao, Jie Ding, and Vahid Tarokh. HeteroFL: Computation and communication efficient federated learning for heterogeneous clients. In Proc. of Int't I Conf. on Learning Representations (ICLR), May 2021.
Alireza Fallah, Aryan Mokhtari, and Asuman E. Ozdaglar. Personalized federated learning with theoretical guarantees: A model-agnostic meta-learning approach. In Proc. of Conf. Neural Inf. Process. Syst. (NeurIPS), Vancouver, Canada, Dec. 2020.
Yuguang Fang, Kenneth A. Loparo, and Xiangbo Feng. Inequalities for the trace matrix product. IEEE Trans. Autom. Control, (9(1))24489-22 -2:20,9--2490, Dec. 1994.
Jonathan Frankle, David J. Schwab, and Ari S. Morcos. Training BatchNorm and only Normh On the expressive power of random features in CNNs. In Proc. of Int'l Conf. on Learning Repre- sentations (ICLR), May 2021.
Yiwen Guo, Anbang Yao, and Yurong Chen. Dynamic network surgery for efficient DNNs. Proc. of Conf. Neural Inf. Process. Syst. (NeurIPS), Barcelona, Spain, Dec. 2016.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recog- nition. In Proc. of IEEE Conf. on Comput. Vis. Pattern Recog. (CVPR), Las Vegas, NV, Jun. 2016.
10

---

Published as a conference paper at ICLR 2023
Junyuan Hong, Haotao Wang, Zhangyang Wang, and Jiayu Zhou. Efficient Split-Mix federated learning for on-demand and in-situ customization. In Proc. of Int'l Conf. on Learning Represen- tations (ICLR), Apr. 2022.
Samuel Horvath, Stefanos Laskaridis, Mario Almeida, Ilias Leontiadis, Stylianos I. Venieris, and Nicholas D. Lane. FjORD: Fair and accurate federated learning under heterogeneous targets with ordered dropout. In Proc. of Conf. Neural Inf. Process. Syst. (NeurIPS), Dec. 2021.
Andrew G. Howard, Menglong Zhu, Bo Chen, Dmitry Kalenichenko, Weijun Wang, Tobias Weyand, Marco Andrectto, and Hartwig Adam. MobileNets: Efficient convolutional neural networks for mobile vision applications. arXiv preprint arXiv: 704.04861, Apr. 2017.
Yuang Jiang, Shiqiang Wang, Victor Valls, Bong Jun Ko, Wei-Han Lee, Kin K. Leung, and Leandros Tassiulas. Model pruning enables efficient federated learning on edge devices. IEEE Trans. on Neural Networks and Learning Systems, Apr. 2022.
Sai Praneeth Karimireddy, Satyen Kale, Mehryar Mohri, Sashank Reddi, Sebastian Stich, and Ananda Theertha Suresh. SCAFFOLD: Stochastic controlled averaging for federated learning. In Proc. of Int'l Conf. on Machine Learning (ICML), Jul. 2020.
Jakub Konečný, Brendan McMahan, and Daniel Ramage. Federated optimization: Distributed opti- mization beyond the datacenter. arXiv preprint FrXiv:1511.03555. Nov. 2015.
Tian Li, Anit Kumar Sahu, Manzil Zaheer, Maziar Sanjabi, Ameet Talwalkar, and Virginia Smith. Federated optimization in heterogeneous networks. Proc. of Machine Learning and Systems (MLSys), Austin, TX, Mar. 2020a.
Xiang Li, Kaixuan Huang, Wenhao Yang. Shusen Wang, and Zhihua Zhang. On the convergence of FedAvg on non-IID data. In Proc. of Int'l Conf. on Learning Representations (ICLR), Apr. 2020b.
Xiaoxiao Li, Meirui Jiang, Xiaofei Zhang, Michael Kamp, and Qi Dou. FedBN: Federated learning on non-IID features via local batch normalization. In Proc. of Int'l Conf. on Learning Represen- tations (ICLR), May 2021.
Tao Lin, Lingjing Kong. Sebastian U. Stich, and Martin Jaggi. Ensemble distillation for robust model fusion in federated learning. In Proc. of Conf. Neural Inf. Process. Syst. (NeurIPS), Vancouver, Canada, Dec. 2020.
Shiwei Liu, Tianlong Chen, Xiaohan Chen, Li Shen, Decebal Constantin Mocanu, Zhangyang Wang, and Mykola Pechenizkiy. The unreasonable effectiveness of random pruning: Return of the most naive baseline for sparse training. In Proc. of Ini'l Conf. on Learning Representations (ICLR), Apr. 2022.
Yishay Mansour, Mehryar Mohri, Jae Ro, and Ananda Theertha Suresh. Three approaches for personalization with applications to federated learning. arXiv preprint 10619,, Jul. 2020.
Othmane Marfoq, Giovanni Neglia, Aurélien Bellet, Laetitia Kameni, and Richard Vidal. Federated multi-task learning under a mixture of distributions. In Proc. of Conf. Neural Inf. Process. Syst. (NeurIPS), Dec. 2021.
Brendan McMahan, Eider Moore, Daniel Ramage, Seth Hampson, and Blaise Aguera y Arcas. Communication-effiieenn learning of deep networks from decentralized data. In Proc. of Ini' Conf. Artificial Intelligence and Statistics (AISTATS), Ft. Lauderdale, FL, Apr. 2017.
MOSEK ApS. MOSEK modeling cookbook.
Available: https://docs.mosee com/ Aug. 2022.
McSEKMode www.oeeeeoonoooooooooooooee
pdf
Jachoon Oh, Sangmook Kim, and Se- Young Yun. FedBABU: Towards enhanced representation for federated image classification. In Proc. of Int'l Conf. on Learning Representations (ICLR), Apr. 2022.
11

---

Published as a conference paper at ICLR 2023
Kilian Pfeiffer, Martin Rapp, Ramin Khalili, and Jörg Henkel.
CoCo-FL: Communication-an and
computation-awai federated learning via partial NN freezing and quantization. arXiv preprint arXiv:2203.05466, Aug. 2022.
PyPAPI. PyPAPI's documentation [Online].
Available:
atts://ilozz.ithuub.b
pypapi/ 2017.
PyTorch. PyTorch documentation [Online]. Available:
htt.s://pptorcchorrg//oom
stable/ htle:/ndexxiiecc html 2022.
Sashank Reddi, 7achary Charles, Manzil Zaheer, Zachary Garrett, Keith Rush, Jakub Konečný, Sanjiv Kumar, and H. Brendan McMahan. Adaptive federated optimization. In Proc. of Ini' Conf. on Learning Representations (ICLR), May 2021.
Xinyue Shen, Steven Diamond, Yuantao Gu, and Stephen Boyd. Disciplined convex-concave pro- gramming. In Proc. of IEEE Conf. on Decision and Control (CDC), Las Vegas, NV, Dec. 2016.
Hakim Sidahmed, Zheng Xu, Ankush Garg, Yuan Cao, and Mingqing Chen. Efficient and pri- vate federated learning with partially trainable networks. In Proc. of NeurlPS Workshop New Frontiers in Federated Learning, Dec. 2021.
Jianyu Wang, Qinghua Liu, Hao Liang, Gauri Joshi, and H. Vincent Poor. Tackling the objective inconsistency problem in heterogeneous federated optimization. Proc. of Conf. Neural Inf. Process. Syst. (NeurIPS), Vancouver, Canada, Dec. 2020.
Daliang Xu. Mengwei Xu, Qipeng Wang. Shangguang Wang, Yun Ma, Kang Huang, Gang Huang, Xin Jin, and Xuanzhe Liu. Mandheling: Mixed dixee-precision on-device DNN training with DSP offloading. In Proc. of ACM Int'l Conf. Mobile Comput. Netw. (MobiCom), Sydney, Australia, Oct. 2022.
Tien-Ju Yang, Dhruv Guliani, Françoise Beaufays, and Giovanni Motta. Partial variable training for efficient on-device federated learning. In Proc. of IEEE Int'l Conf. Acoust. Speech Signal Process. (ICASSP), Singapore, Singapore. May 2022.
12

---

Published as a conference paper at ICLR 2023
A RELATED WORK
FL Algorithms with non-lID Data. In FedAvg (McMahan et al. 2017). the edge devices perform multiple local update iterations their local models before sending them the server for aggre- gation. In this way, the communication cost can be reduced. When there is data heterogeneity, conventional FL algorithms such FedAvg may have slow and unstable convergence (Karimireddy let al.] 2020). To handle the case when the data is non-llD, some works (Li et al.) 2020a Karim- ireddy et al.) 2020) have introduced new optimization frameworks to obtain a more stable model for the devices. Although the obtained model may still perform well on average across the devices, some of the devices may not be satisfied with the obtained final accuracy (Collins et al., 202T]. To address this issue, more personalized models are required for the devices. The personalized models can be obtained by using different methods such as meta-learning (Fallah et al.] (020), multi-task learning (Marfoq al.) 2021), clustering (Mansour et al. 2020), and decomposition. In the decom- position method, the learning model is decomposed into a global model and a device-specific head model (Arivazhagan et al. 2019 Collins al., 2021). In each communication round, the global model is updated by the devices and is aggregated by the server, whereas the head model is updated the server. Ohh al. (2022) proposed BedBABU using the
by each device but is not transmitted decomposition method. To improve the personalization ability, in FedBABU, the local head model parameters are frozen during training. After convergence local head model parameters during the fine-tuning steps by using its local data samples.
global model, each device updates the
FL Algorithms with Heterogeneous Devices. Given the disparities in devices' hardware, it is cru- cial to address device heterogeneity in FL. In general, this problem can be tackled by reducing the computational complexity of model training based on the devices' hardware capabilities. Lin et al. (220)) Afonin & Karimireddy (2022) proposed FL algorithms using knowledge distillation, which aim to generate compact models by transferring knowledge of a large model to smaller ones. Prun- ing method, where compact model for each device can be obtained by removing the parameters with little impact on the performance of the original learning model, is another approach to accom- modate device heterogeneity in FL. Due to the dynamic sparse training, pruning method may suffer from instability in convergence for finding the sparse masking vectors, which are used for training the sparse sub- networsss at the devices. This issue can be resolved by consistent mask adjustment procedure at the devices, at the expense of additional communication overhead (Jiang et al.) [022] Babakniya al. 2022 Bibikar al.] (0022). The ordered dropout technique has been proposed in Horvath et al. (2021) to dynamically adapt the model size used by each device based on its capabil- ities. Another pruning method that does not require additional communication cost is use fixed masking vector for each device before training based on its computational capability. This would result in training of the heterogeneous local models for the devices. In this regard, HeteroFL. (Diao et al. (021) aims to facilitate efficient training and stable aggregation of devices' local models. However, the available local data samples at the devices may not be used efficiently in HeteroFL.
In particular, HeteroFL considers different model architectures as the local models for the devices. Thereby, the available data samples at each device can be used only for training a specific model architecture. In Spxit-Mix FL (Hong et al. 2022), based on its computational capability, each device randomly selects some of the base models in each communication round to train them. After train- ing, by using ensemble learning approach, each device mixes the selected base models construct its desired model size. The potential drawback of this approach is that none of the devices train the original learning model.
Another line of research is related to the objective inconsistency problem in FL. Due to the differ- ent size of the devices' local dataset and their computational capability, some of the devices may finish their local update iterations faster. To prevent the devices from being idle, Wang et al. (2020) proposed let those faster devices continue their local update iterations until the slowest device finishes its local update. FedNova is proposed by Wang et al. (2020) to resolve the objective incon- sistency problem due to different number of local update iterations. Another approach to consider the same number of local update iterations for all the devices. However, to enable all the devices to finish their local update iterations in a timely manner, the required number of computations should be adjusted for those devices with lower computational capabilities. Our proposed algorithm as well as other FL algorithms which aim to adjust the required computations for each device based on its computational capability can be employed to address the objective inconsistency problem.
13

---

Published as a conference paper at ICLR 2023
B LOCAL UPDATES AT DEVICES
Algorithm 2: Local Update Function
I: function DeviceLocalUpdaanl
9y, m. f. D, n)
# Local update iterations for cach device
2: B - Split local data samples D into B local batches.
::= and w' + Wy.
each local epoch e E (1...., F}
The total number of local update iterations T = E x B
5: for each batch b E B do
6:
w <- w 17m & V f(w"", b).
7:
1
end for end for
9:
10: Return w²
11:
end function
C CONVERGENCE ANALYSIS FOR STRONGLY CONVEX LOSS FUNCTIONS
To show the convergence of MerkedMask for smooth and strongly convex loss functions, in addition Assumptions [ make the following assumption:
Assumption 4. (ww) [NV. is p-strongly convex and satisfies:
fv( fww (v w'V + || wll?, Vo, w, [N].
(10)
To quantify the degree data heterogeneity at each device [N], the term r, fn(w") f is defined (i et al. 2020b). Let ((t) = IE] [y 11 We first prove the following useful lemma.
Lemma 2. Under Assumptions 7 4 if the learning rate is small enough, i.e., n7() S L(NT+T) for all L = 1, I....,Ty T, have
5(t + 1) S (1 go n(()) 8(t) 21 mlt) 42 72
(11)
where
N
40
1LT
ins
(12)
N
N
du
A 2-8
in
1
+ 2TΩ
Σ
dn Σ (ka)))
(13)
N
n=1
n=1
1ml
G T(T
12 - 1)
N
92
(2 (1) Σ
Yn + 2Lr
In (2 + NTr) P. + N7
Sn
(14)
==1
and and should satisfy 1,...,7, 1=1,... T and minne[]] (Ja(w,(1) (1)) - fn (ww)) < Y, 1- 1,... T, respectively.
mmin; (ww* - - (1)) @ (w (t)) $ Q, for all n E [N],
i=
Proof. See Appendix N
Using Lemma 2 can state the following theorem for the convergence rate of smooth and strongly convex loss functions:
Theorem 2. Given Assumptions i if we choose k
2L (Nt + 1) and the learning rate Ty(t) = (ol under the full device participation scenario (i.e., 1), after T communication rounds, we have
2
EF(w, FF*
" 2 2 + k)
Lq
Aq 2
+ + ((++1)(1).
(15)
40
%
14

---

Published as a conference paper at ICLR 2023
Proof. See Appendix
Remark 2. The first term on the right-hand side of (15) appears in the convergence bound due
the device heterogeneity. In particular, Theorem 2 shows that for qi * O, the FL algorithm converges to a local optimal solution at the rate of O (/T)). This convergence rate is similar to the results presented Li et al. (2020b); Amiri et al. (2022), where device heterogeneity has not been considered. Hence, using the masking vectors in FL do not degrade the convergence rate of FL.
The result in Theorem 2 shows that when qı + 0, the FL algorithm converges to the global optimal solution for the smooth and strongly convex functions. Based on (I3), it is straight- forward to verify that without device heterogeneity in the network (i.e., when all the devices have the maximum computational capability and can update all the parameters of the global model), all the elements of vector kn. It [N] are equal to N. Hence, in this case, we have Tn = To and Q1 0. In general, based 473), one way to reduce the bias incurred by the device heterogeneity is to design the masking vectors by minimizing 91- For example, one can search for the masking vectors, which minimize Ex=1 dn
Remark 3.
Ei= (kmn)).
D PROOF OF LEMMA 1
Given vectors 2. y. and z, we form diagonal matrices X, Y, and Z, respectively. Note that we
can write (x, y @ z) the form of the trace of matrices X, Y. and product, i.e., (xx y + z)
Z). By using Theorem 3 in Fang et al. 1994), we have the following inequality:
Tr(XY Z) < (Y) Tr(X2) - Xx (Y) - Tr(Y)).
(16)
where x (Y) and Xa Z) are the largest eigenvalue of matrix Y and the smallest eigenvalue of matrix XZ, respectively. Since the considered matrices are diagonal, we have X (Y) = max; (y)) and (XZ) miny (x z),. Hence, we have
(n. ④ z) S max(y)r (a. 2) - min (a @ 2); (dmax(()) Σ (t))
(17)
1=1
Since d maxx( (y):
Et=1 (y): by considering min: (ær 3 2) -Q, Lemma
1 proved using
inequality (17
E PROOF OF THEOREM
Considering the smoothness of Jn (w), n E [N]. each communication round 2 1, we have
EF(w,(t (t + 1)) <
EF w (t)) + E (w + 1) w. (t))), (w (t)) + L 1g (t + 1) wt(t) (18)
We first find an upper bound for ||u + 1) w. (t)l as follows:
E I (t (++ 1) -w, (t)) (t)
Σ kn fi (t),
N
m (t)
kn V Vw (t),
(t), (t) k. @ f.
N
+ 72
kn c V fn (w)
(19)
A
where equality (a) results from 2 and
Equality (b) is obtained via basic equality E ||
E
Ez| 2 2 || for any random vector
15

---

Published as a conference paper at ICLR 2023
By using Assumption
2
can obtain upper bound of A1 follows:
N A1
=
kn
VJn (t), kn & Vfn(w(((t)
N
O (1, kn o (f)(w((())) 2
S Nr
kn
==1
NV
S N Sp
(20)
==1
By considering the convexity of ||-1|2 and by using Yo. max (kn)), we obtain an upper bound of A2 as follows:
N
A2 =
kn
< NT
V Sn [1)
n= i=1
NT Σ
(w (1))
(21)
By combining 19
20 and 21..
we have the following inequality:
E || (t 1) w. (t)l N 2 + NTi
(22)
Now, we aim to obtain an upper bound of E (wy (t -
(t),,, -0
(t)) We have
E (Wg(( (t + 1) (t),), (())(())
E
kn (t), (t)), (F(w)())))
N
7()EE
N
N (V VP(w, (1)) + 7()TV
du
7(()E
(-7m)
Σ
(ken)
R=1
x=1
1=1
duo
(tt)
E
V (w (t)), (F(w,(()) + (()r
(ken)
(23)
where equality (a) results from and Equality (b) follows from EV (w, (t), by (t) -ff( (t)). Inequality (c) holds by using Lemmal Inequality (d) follows from Yn N
To find an upper bound for -E t E.
(w (t)), VF(w, (t))).
we first represent it as fol-
lows:
V (t),, VF(w(())))
V (t)
V fn w, (t))
E
V (f
E IVF (VF(w (t)))
(24)
16

---

Published as a conference paper at ICLR 2023
(V |
Then, E Em (V
is bounded as follows:
I
- ((,()))))
V fn (w, (1)
#=
L
E (t(() - w (t)
(25)
N
where inequality (a) results from the convexity of 11.11- Inequality (b) results from Assumption 1 Now, we aim to bound E| (1() - (1]l for = ,,...,T. T. By using (I. we have
B wgygtt w (t)l
E
mt 3
fn (w) (t),
<1((1)i (
Σ
Bl|mm Sn (0)
/=1
(
(26)
where the last inequality results from Assumption 3 By combining (25 and 26 we have
E
Vin(w, (.,(,,
<13)
(27)
By combining (18 and 22
27
we have
dw
T EF(w, (t + 1)) EF (t))
() (t) Σ
& + n()T
Σ
dw in
Σ 1=1
(kin)
2
n=1
n((TT
1.2 (t) T(T
1)(2T -
E
BIFFF( (t))
12
ntt)
N
LNT
(fn(w((()))
(28)
N
Since n(t) 17 LN have
p(t
D El= ( - -LNT)a((() [i
N
0.
By rearranging the terms (28). we obtain
2 < (EF(w,(t) - EF(ws(t (t+ + LN Σ 22
E (F)
n=1
dew
(2
+ 24
duu In
+ 1,2 02 Σ (kin))
(29)
Finally, we multiply both sides of 29 by 1 and sum over 1...., T. Then, Theorem
concluded by considering that the first term on the right-hand side of (29) is a telescoping series. We have
2 W T
2
(EF(w,(t) -FFFEEF(W(( + 1))) (F (F(w,(1)) EF (T+ 1)))
2
(F(w (1))(1)) F*),
(30)
nTT
where the last inequality is obtained by considering that EF(w, +)) > Ft.
17

---

Published as a conference paper at ICLR 2023
Selection
Selection 2
Selection
Selection 2
Selection
171
111
m!1
n;
m
(d)
Selection
Selection
Selection
Selection
171
111
n
m
n
m!3
(b)
(c)
Selection
Selection 2
Selection
Selection
111
m
m
m
m
(c)
m
The masking element is set
m
The masking vector element is set
m,
(
Figure 2: Illustration of masking vectors selection in network with three devices. Device 3 has the maximum computational capability train x1 parameters of the learning model. We consider different scenarios for computational capability devices
(a) Devices and can train x0.25 parameters model. (b) Devices and can train x0.25 and x0.5 parameters of the learning model, respectively. (c) Devices 1 and 2 can train x0.25 and x0.75 parameters of the learning model, respectively. (d) Devices and 2 can train x0.5 parameters of the learning model. (e) Devices 1 and 2 train x0.5 and x0.78 75 parameters learning model, respectively. (f) Devices
learning
can train 0.75 parameters the learning model. (g) All three devices have maximum computational capability.
F SELECTION OF MASKING VECTORS IN A TOY EXAMPLE
a heterogeneous device setting, devices with lower computational capability partially train learning model. Those devices masking vectors to freeze portion of the learning model during training based on their computational capability. As shown in Theorem [ employing masking vectors lead to bias term in the convergence bound. In this section, demonstrate how the bias value impacted by the design of masking vectors using a simple example. We consider three devices. One device has the maximum computational capability. However, different scenarios considered for the computational capability of the other two devices. The devices aim to train model with four parameters. Hence, their masking vectors have four elements. Fig. shows the considered scenarios and the possible masking vectors selections, which may lead to different bias values.
For each scenario and for each possible selection, Table 6 shows vector kn, and the value of for each device as well as the obtained bias value. For this example, the results Table 0 illus- trate that the bias value is minimized by freezing the same parameters for two devices with lower computational capability, in case each parameter cannot trained at least by one of those devices (Figs. 2 a)) and b)). However, when each parameter can be trained at least by one of those less capable devices, the bias value is minimized by freezing distinct parameters for the devices (Figs. 2 f)). These results compatible with the empirical results obtained by Pfeiffer et al. (2022) and Yang et al. (2022), where random selection used by the devices to prevent freezing of the same parameters.
18

---

Published as a conference paper at ICLR 2023
Table 6: kn. Th. and the bias value for the possible masking vectors selections of the scenarios shown in Fig. 2 on Remark
'hee selection with the minimum bias value chosen based
Scenario Possible selections
k.
7m
Bias value
Selection
11
10
0,0,0
122
k3 1,1,1
73
k1 3.0,0,0
Selection
k2 00 0.4.0.0
4
R3 2.7.1.1
78
1,0,0,0
71
Selection
R2 0,0
72
3
Fig.
Rg 7.2.1.1
75
41 = 5,0,0,0
71
Seleetion
422 00_
72
4q 2' 2'
73
1 0,0
71
10
Selection
43 7
72
3
Fig.
78
0,0,0.0
71
Selection
72
2
73
Selection
78
7
71
Fig.
2
Selection 2
k.
72
4
70
21
71
Selection
22 0.
72
2
43
73
41
0.0
71
10
Selection
k2 3:
72
Fig.
ka
78
71
Selection
72
2
Selection
R3
R.
Fig.
Selection
k
0.
2
ly
Fig
1
R3
0
The results Table 6 also show that by increasing the computational capability of the devices, the bias value be decreased. In the extreme case that all three devices have the maximum computational capability (i.e., Fig. 2 g)). the bias value is zero.
G EFFECT OF MASKING RATE ON THE COMPUTATIONS
By increasing the masking rate (i.e., ))), each device can reduce the number of trainable parameters and the number of FLOPs based on its computational capability. Fig. 3 shows the reduction percentage]] which can be obtained for the number of trainable parameters and the number of FLOPs versus the masking rate. For the results Fig. 3 we have considered that a device performs each local update iteration batch containing 50 data samples of CIFAR-10 and CIFAR- 100 datasets using ResNet (PreResNNt18)) and MobileNet, respectively.
Reduction percentage the percentage change in value compared its maximum value.
19

---

Published as a conference paper at ICLR 2023
Bodentee hucetgg ger Laca tigtr flesson
90-
20
30
Masking Rat
Masking Ritt
(a)
(b)
Figure 3: Reduction percentage for the number of trainable parameters and number FLOPs versus masking
rate for (a) ResNet (PreResNet18) 18) and (b) MobileNet.
H SOLVING PROBLEM PMASK BY SCA
In this section, we first transform the non-convex constraints in optimization problem prask into convex functions or a difference of two convex functions. Then, by using successive convex approx- imation (SCA) method, we can obtain a suboptimal solution for problem pmask in polynomial time. We first rewrite constraint (9a in the form of the following inequalities:
(m)))
(kr)) S
jE A, n [N].
(31a)
2n (mn))s
(mn))
(kw)) 2 Σx (m,))
j jEA, [N]..
(31b)
Since (m,,) is binary variable, the non-convex constraint (3lb) can be expressed as the following convex constraint:
(mn))
(km)) Σ (m,)); JEA, [N].
(32)
The non-convex constraint (3la) can be rewritten as follows:
(kn), (mw)) S (mnnij, EEA, ne [N].
(33)
Equality (kw)) E (m,)); = ((Kn)) En= (m (km)) (EN=1 (m );
can be used express the left-hand side of (33) as a difference of two convex functions. We have
(kn))
(mn:))
(kr,)
(m)) < (m.)j, j A, n [N]. (34)
2
n'=1
n==l
Next, relax the binary constraint 9c in the form of the difference of two convex functions as follows:
mmn) Σ Σ mm.
50,
(35a)
n=AjEA
1=1
0 5 (m.))
< 1, j E A, [N].
(35b)
Finally, we define the following functions:
e(m,)), = =m,,),
(36a)
N
o((kn)j,, (k:m)) 3 + Σ
(m.))
(36b)
x'=1
20

---

Published as a conference paper at ICLR 2023
where vector M, =(m,)) n' E [N]).
Algorithm describes the SCA algorithm for solving problem Pinask Let i denote the iteration index. In Line 1, we initialize the maximum number of iterations gmm. In Line 2, the decision variables vin and k n E [N] are initialized in iteration i = 1 with a feasible solution of problem pousk In Line 4, the optimal solution of problem PPmus-SCAA (i.e., m, and kn n E [N] (i+1) determined. In Line 5, using r and k; update m..t1) and k( obtaining the first-order approximations (g))))) and ((k,), M,) of functions (;))))) and o ((ka)j Mj), respectively. The iteration index is updated Line 6. The steps within Lines 3 to 7 are repeated until the algorithm converges solution
[N] to be used for
reach imax, The convex problem
pmusk SCA, which is solved in each iteration i, is as follows:
SCA,
minimize
Σ
dw max (kn)s
b (ky:(( + on
mo. &n, NE[N]
J'EA
subject to constraints
9b
32
35b and (9d).
N
(kn)
+
mmn') )
0((k()) )j, ,²
$ {manj jEA, n e [N],
2
n'=1
mmn)
Σ
(())))))) SO.
n=1jEA
n=1jEA
Algorithm 3: SCA Algorithm for Solving Problem Prnask
1: Set := and initialize the maximum number of iterations max
2:
Initialize variables m.. and k. n E [.].
3: Repeat
4: Determine the optimal solution m., and ke, n [N] of problem SCA
Update variables rit, (f+1) == m,, and ((+1) == kinn n E [N].
6: Set i := i+
7: Until i 2ax or vit and Rn. v E [N] converge.
(6) and k = (())) E [N]. 8: Return ri..." := m..
I PERFORMANCE COMPARISON FOR DOMAINNET DATASET
Different from the considered class non-IllD configuration for CIFAR-10 and CIFAR-100 datasets, in this section, we evaluate the performance of our proposed algorithm under feature non-IID con- figuration. We perform our experiment with AlexNet on DomainNet dataset. The dataset contains images of six distinct domains including Clipart, Infograph, Painting, Quickdraw, Real, and Sketch. We consider 30 devices in the network, and split each domain among devices. We set the number of communication rounds and the learning rate to be 100 and 0.01, respectively. Table 7 shows the obtained test accuracy after fine-tuning and the number of trainable parameters for PerFedMask compared to other baseline algorithms. The results in Table 7 indicate that PerFedMask can achieve a higher test accuracy compared HeteroFL. and Split-Mix FL. algorithms, while the number of trainable parameters much less than FedBABU, FedProx, FedNova, and FedAvg algorithms.
Table 7:
Test accuracy after fine-tuning tuning and number of trainable parameters of PerFed- Mask and baseline algorithms for DomainNet dataset
Test accuracy after fine-tuning
PerFedMask (Ours)
FedBABU
FedProx
FedNova
HeteroFL
Split-Mix FL
FedAvg 72.24
70.68
72.29
71.95
72.85
67 68
68.44
Number of trainable parameters
PerFodMask (Ours)
FedBABU
FodProx
FedNova 57.063M
HeteroPL 98.88
Split- FL 4.058M
FodAvg 57.063M
31.221M
57.022M
37.063M
21

---

Published as a conference paper at ICLR 2023
J COMPARISON OF TRAINING CURVES
This section includes some of the training curves for our experiments. Fig. illustrates the train- ing loss over communication rounds for PerFedMask and the baseline algorithms CIFAR-100 and DomainNet datasets. Fig. 4 we have considered the full device participation scenario. The difference between the training loss of PerFedMask and FedBABU is the bias resulted by the de- vice heterogeneity. Although we have minimized this bias through solving Pmmask it has not been completely eliminated. Fine-tuning after training helps improve the performance of PerFedMask. When e 0.1, Fig. shows the evolution of the validation accuracy over communication rounds for PerFedMask, FedBABU, FedAvg, and Split-Mix FL algorithms on CIFAR-10 and CIFAR-100 datasets.
10
PerFodMack
3.5
PerFedMaakk
HecerofIL
Splii Mix FL
Split Mix FL.
TedNova
FalPhox
0.4
Blasehhr in desice holnogrnctt
FalProx
FedNma
FedAvs
FrdAvg FedBABU
0.2
HammFll
letengennctty
FedBABU
Thiining
Training
0.5
01
80
0
20
40
100
Communication Round
Communicarion Round
(a)
(b)
Figure 4: Training loss evolution over communication rounds for (a) CIFAR-100 and (b) DomainNet datasets.
is set to 1.
0.7
0.25
Validation Acuuney
Validation Accuracy
0.2
0.5
0.4
0.15
0.1
0.1
PaFulMakk
20
60
80
40
80
Communication Round
Communication Round
(a)
(b)
Figure 5: Validation accuracy evolution over communication rounds for (a) CIFAR-10 and (b) CIFAR-100 datascts. is set to 0.
K PERFORMANCE COMPARISON USING Box PLOTS
Tea Accerac atter Hote Thmne
(b)
Figure Test accuracy after fine-tuning for (a) CIFAR-10 (b) CIFAR-100 datascts, c set 0.1.
22

---

Published as a conference paper at ICLR 2023
L COMBINATION OF PERFEDMASK WITH HETEROFL AND SPLIT-MIX FL
Table 8: Performance comparison on CIFAR-I00 100 dataset when =
Algorithm
Test accuracy
of trainable parameters
of FLOPS
Before line tinetuning
Aller line-tuning
Torward
Backward
PerFedMask + Spx
29.66
70.78
0.204M
17.182M
0.122c
PerFedMask + HeteroFL.
30.94
71.56
1.611M
0.122G
0.464G
PerFedMask
32.04
72.40
1.803M
0.220G
0.711G
Split-Mix FL
23.09
65.95
0.223M
17.182M
0.127G
HeteroFL
30.13
68.65
1.774M
0.122G
0.473G
FedBABU
29.70
69.01
3.207M
0.220G
0.863G
M EFFECT OF FREEZING THE LOCAL HEAD MODELS DURING TRAINING
In this section, we investigate the impact of freezing the local head models on the test accuracy of compare the test accuracy of PerFedMask with (w/) and without
PerFedMask. In particular, (w/o) freezing the local head models. In the w/o head freezing scenario, the learning model is not decoupled into the global model and the local head model. Table 9 shows that the test accuracy of PerFedMask is improved by keeping the local head models frozen during training.
Table 9:
Test accuracy of PerFedMask with and without freezing the local head models
PerFedMask CIFAR- 10 dutaset
CIFAR- 100 dataset
0.1
w/ head freezing
83.60
88.43
67.47
72.40
w/o head freezing
71.56
85.76
64.85
67.89
N PROOF OF LEMMA 2
We define two auxiliary sequences g(t) ΣN
kn
EL
(n (t)
and g(t) DN
k 2 El Vn (wh, (t), (t)). We have
N
E
(yy(t + 1) - 201
Σ
kn
T+ x
- w²
n=1
N
N
Σ n=1
An C wg (1) (n)( Σ
ki
Sn (w (1), w
R=1
E (1w,((() (1))(()) w" 1²1
(I) (1))g(()) w² + ()t))(()) -
= E [1lw(()) (t) ~" -
2nn(t) (w, : - ((t g(t))
(37)
Since Et(() = g(),, we have EAA = 0. Next, we focus on bounding A
A [ty(t) =² (t)l || (t) w" 11 (t))) (g()) 2n(() (w w", g())
B
B2
(38)
23

---

Published as a conference paper at ICLR 2023
By considering the convexity of II-11" and by using Assumption 1
can bound B1 as follows:
N
B1 = ===10(()| =
kn
(fa(w)()))
N
Nr
(n (1))l
n=
i=1
N
N
mn
[ ( (1)]l
N
s 2L
Σ
( -7))-
(39)
== 1=1
Next, we aim to bound Bg. We have
B2 = - -2n(t) (w, (t) w", g()))
N
=- -22(t) wy (t) w"
kn 2 fzx
n=l
2m((t)
(wy(t) (t) - w (t) + w (t) w" kn @ fn (t))
n=1
N =t((
2) W (t) w k., (.n((((
N Σ
+ +n((l)
(w)) w. kn @ Vn (1))
(40)
n=
We first obtain an upper bound of Cj. We have
C S 2 w (t) - w (t), kn c (w)
(
< 2%
|| w0 (t) w (t) (t))
7n
||,(()) -w,(0) + 7(()mm (m),
(t)
< 7n ||0g(t) - (+)l + 2m (F,(w, fū),
(41)
(tt)
where inequality (a) results from triangle and Hölder's inequalities. Inequality (b) results from the inequality of arithmetic and geometric means (AM-GM) inequality. For inequality (c), we use Assumption T
Next, we obtain an upper bound of C2 by using Lemma 1 In particular, by considering a w - -w,, (t), y = kn and z = V (w,, (t)) in Lemma
we can bound C2 as follows:
dey
C2 S 7m (wo wi (1), (ln((((())))) + Ω
dun
Σ (kmn))
dea
5
- ((.() + 17n (- w/w (t)
20 w0'l²))
+5
do In
Σ 1=1
(kmn))
(42)
where the last inequality results from Assumption 4 Now, focus on bounding D. Considering
the inequality ||ær yll² ||a: 1 - ||y|| for any x. y TRd and by replacing with wi, (1) w
24

---

Published as a conference paper at ICLR 2023
and y with (t) - w (t) we have
D <
(y(() (x)]
1 2 || w²11
(43)
By combining 38
43
we obtain
wn(t)r
N
EAI S 1
Yn
B 2 - w='
2
12
N
+ (1 + pmn(t))
Yn E1| 1y (t) - -00 (1)]l
Bl
n=1 i=1
(l)
(N + 7m)))B(())) ((Jr((u, (7)
-2n(() Σ
In E (wx - Sn (w0)))
di
+ (t(()r
Σ ==1
dw dn in
Σ (kmn)e
(44)
I=1
By rearranging the terms in
44
we have
EA1 < 1 un(()r
N
(l)w,(t) - (1) 11
Yn
2
n=1
+ (1 + pn((t)
Yn E W (t) wr(((]]
n=1 1=1
N
2t)(t) Σ
Yn (1 - In(t) (Nr%, (f) fn("))
+ +2TTT12(1)
(Nrah + 7m) ((fx(")))) - {
N
du
+ (1)
(kin))
(45)
Now, we aim bound fw (t)) fn (w=) as follows:
fw (f)) - fa(w") = (fn(w' (t)) (wng(t))) + (fr (20 (t)) fn fa(w")
2 IV8
Vn (1)) m ()(()) + (Jr (1) fn(w"))
n(t)
1
|(,(() - w/(((]]
2n(()
+ (w, (t)) - fa(w"))
- -LL((t) (f((w,()) - -f:) 21y(() ||,,(() (1) - w (1)l
+ (fn(w, (t)) - fn(w"))
(46)
where inequality (a) results from the convexity of f.. (w). inequality (b) is obtained by using Cauchy- Schwarz and AM-GM inequalities, and inequality (c) is due Assumption 1
25

---

Published as a conference paper at ICLR 2023
By combining (26
(45
and (66),
we have
un(t)r
N
EA1 v
1
Elle0, (1)
-w" w=
2
n=
+ + TTT - (2T
1)
In (2 + pn(() Ln((t (NT% + 1))
6
n=1
N
-2n(1)7 (1 In(())
Σ
Ia (1 In((1) (NTY,. + 1)) 1)) Sn (nw)))
+ (L)
Σ
In (1+ (1 - Ln((t) (NrY, + ()) - -5))
N
due
+ 2n(TT
Σ
dw - -E (Kmn))
(47)
n=1
1=11
Lemma 3. For m]t) S T(NT+TJ for all t = 1...., T. we have
-2n(()r (1 Ln((t) Σ
In (1
+ +1) B(Jn(ung()) (nww))
N
1
Lo(()) (Th (NTT +1) 1) N (T f1)
(1 - (1-())(1-Lm(()) r
Σ
In
N
(48)
n=1
Proof. See Appendix P
Using Lemma 3 we can simplify (47 as follows:
EA1 S
1
2 UT
In n(t)
El Ellw,(()
w0112
+ C T(T 1)
N
(2 + 1)
I'
n=1
N + (n(t)r
Σ
In
N
n=1
+ 2LTT
In (2 + NTT.
n=1
da
+ (2)
dnn'n Σ (kmn))
(49)
Finally, we use 20 to bound EA, in 37 Lemma
2 is concluded by combining 20
37p and
49
PROOF OF THEOREM
First, through induction, we show that for diminishing stepsize r7t) qo(++a we have 8(t) < 14. where 3(t) = max 2 (t + x) + Ang (N + 1) 8(1)). Note that the considered 17(t) satisfies
(())
the mentioned condition in Lemma 2 Moreover, the definition of B(1) ensures that the inequality
26

---

Published as a conference paper at ICLR 2023
holds for 61). Now, assume that the inequality holds for t. We show that it also holds for + 1. we have
From Lemma 2
8(t + 1)
- (1
go (()))(
6(t) + 91 m7()) +
92 (²()
< (1
2
8(t)
241
4g2
+
1
1+ K 20 + k)
48 (t + x)²
t+ ((1)
+ 20091 (t + K)
+ 492 ((t)
R)+
/ (1 k)
2
(1 + k)²
1+R-1 8(t)
kt+
1+K
3(1)
(t (++R)
B(t)
(50)
+ K x
Finally, by the L-smoothness assumption for F, we have
(FF(w,(T) / - - ((T)) -
L B(T)
(51)
2
2 T
From the definition of ((t), we have B(T) - 2m1 (T' (T) + t + +K + 1) ((1. Combining this
with 51) completes the proof of Theorem 2
P PROOF OF LEMMA 3
First, we define two vectors æ and y E RN, where (æ),. =m (1-11114 Ln((t) (NT7a + 1)) and (y), Ia (m) n E [N], respectively. Let X and Y. respectively, denote the corresponding diagonal matrices of vectors æ and y. By using Theorem 3 in Fang et al. (1994), we have the following inequality:
Tr(XY) Z (x Tr(Y) + AY (Tr(X) (TNXX)NN(XXX))) NAw(X)),
(52)
where An (X) and (Y) are the smallest eigenvalue of matrices X and Y. respectively. Since X and Y are diagonal matrices, we have (x = min (z), and x (Y) = min (y)n- By considering ntt)== < E(N++Tj. all the elements of vector æ including (x are nonnegative. Also, (x)n is a quadratic function with respect to "/n For N I In I 1, by considering that (t)) S I(Nr+Ij. it can be verified that the minimum value of (æ)) is obtained at "in N. Thus, (x lower bounded as (x 2 & (1 Tn((t) (T + 1)). Hence, we have
-2(()) (1 - (n))t) Σ n=1
Yn (1 Ln((t) (NTr, + 1)) (f)(w,(t) - fnw"))
N 1 -n((t) (1 Lm(t)) (N(X)
(
Σ (w) fn(w")
==1
-2n(t)7 (1 Ln((t) min nE[M]
(fn (w,(w())) - fn((w)) -
(
5 2n((1) (1 In(()) AN ( F²)
2n)(t)T (1 Ln(()) T (Tr(X) - (X)(
2 2n(()) (1 - Ln((t) T Σ 7m (1 Lo(lt) (NTT + 1)) - -(1 -11 (T (++++) +1))
(53)
where for inequality (a),
We also the fact that -n)(()7 (1 In)(t) 0. Note that since 7tt)) S I(N++Ij- Vt. the quadratic function -2)(t)r (1 T.)(t)) is always nonpositive.
(52).
Inequality (b) results from the definition of F(w). For inequality (c), since F(w, (t)) > FF, use
the fact that -2n (t) (1 Ln((t) Ax (X)N (F(w,(t)) F*) S Finally, Lemma is concluded
by using the obtained lower bound for AN (X) and by rearranging the terms in 53
27


Reviews:
[{"review_id": "kgnOPbOfia", "reviewer": "Reviewer_Frgf", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The work studied masked federated learning (FL) for device-heterogeneous settings. The theorems on the convergence of masked FL revealed the limitation of existing masking strategies that induced biases to the model convergence. Therefore, the authors propose to optimize the per-client masks on the server. Empirical results demonstrate that the proposed method can improve multiple existing masked federated learning methods.", "summary_of_the_review": "The paper provides a novel insight into masked federated learning, which should be valuable for device-heterogeneous federated learning. Some issues on the clarity and experiments should be addressed during the rebuttal.", "strength_and_weaknesses": "Strength\n* The method is principled and motivated by the theoretic analysis of convergence. It is novel that the authors show the bias induced by heterogeneous masking. The finding reveals the critical upper bound of empirical studies of masked federated learning so far.\n* Because of the general non-convex and masking-agnostic assumptions, the method can be integrated with different masked federated learning methods, like HeteroFL, Split-Mix FL, etc. \n* The experiments also demonstrate that optimizing masks could improve different baselines. Therefore, the theorems and methods are generalizable and could be impactful in device-heterogeneous settings.\n\nWeakness\n* The method is motivated by the bias on the theoretic convergence bound. But the intuition behind the bias is not well-explained. The authors may elaborate on why the bias occurs. Instead of attributing the bias to the masking, it should be discussed how the bias is related to the heterogeneity. If all clients follow the same mask, how will the bias be? If all clients follow non-overlapped masks, will the bias be maximized? When the mask makes the bias minimized if can all clients afford the same number of parameters?\n* Authors provided two techniques for solving (10a-d). More details about how to choose one of the two techniques should be provided to ease the practice.\n* How does the method work on feature non-iid datasets? Since Split-Mix FL has shown that the HeteroFL would perform poorly on feature non-iid datasets, the Split-Mix FL based evaluation should be conducted on feature non-iid datasets.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "Yl3NMFTtB2", "reviewer": "Reviewer_pTHn", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This work proposes a specific masking strategy on each client that leads to two outcomes; 1) adhering to the computational constraints of each client and 2) improving the personalised model performance. Instead of prior works that “zero out” the masked parameters, this work instead “freezes” the masked parameters to the values communicated by the server. While strictly speaking this leads to less efficiency gains, the benefit is that each hardware constrained client can maintain the flexibility of the server model, while still being able to contribute in training (a part of) it. To optimise the masking assignment, the authors formalise it as a constraint satisfaction problem that targets to minimise the theoretical convergence bound bias coming from masking. As the latter is an NP-hard problem, the authors propose a specific approximation procedure that seems to empirically work. Finally, the improvement to the personalised model performance comes from also adopting the FedBABU strategy, i.e., splitting the learning of the body and head of the network; the head is frozen and kept local while training the (masked) global body and when the body has converged, the entire network is fine-tuned locally with the available data samples. The authors perform extensive experiments and ablation studies to show the performance of PerFedMask in various settings. ", "summary_of_the_review": "Based on all of the above, and primarily due to my novelty concerns, I cannot recommend acceptance. This work could benefit by better positioning in terms of works such as CoCo-FL and by making more clear the improvements that it brings to the table. \n", "strength_and_weaknesses": "Strengths\n- Extensive experimental evaluation\n- Theoretical analysis of convergence\n\nWeaknesses:\n- Server side optimization problem for masking does not consider FLOPs directly (which is a better proxy for computational cost, e.g., convolutional layers are more computationally expensive that linear ones, however the latter have more parameters) \n- Incremental; mostly a combination of FedBABU + freezing methods (such as CoCo-FL)\n- No indication of uncertainty / error bars in the results (so unclear whether the gains in some cases are significant) ", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "AoDBJ7Cq-CK", "reviewer": "Reviewer_nizq", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The author proposes to replace the random mask in FedBABU with a masked optimized for the device heterogeneity. The proposed masking is layer-wise so that all parameters in a layer will be frozen or not, based on the formulation Eq(10). Empirical evaluation shows better performance, less trainable parameters and FLOPS.", "summary_of_the_review": "The motivation of applying customized mask to FedBABU seems valid, but it is not surprising to see the empirical improvements. There are multiple aspects to further improve the paper as stated in weaknesses.\n", "strength_and_weaknesses": "Strengths:\n1. The motivation to let devices with more computational resources to do more training makes sense.\n2. There are empirical benefits of the proposed method regarding the accuracy and FLOPS.\n\nWeaknesses/Questions:\n1. How is \\phi_g updated?\n2. Can you show and compare the training curves?\n3. Can you confirm that among those devices with lower computational capacity, a higher computational capacity leads to more trainable parameters?\n4. In Eq (1), the \\theta_n in f_n(\\theta_n) is the concatenation of w and \\phi. However, in section 4 the arguments becomes w, contradicting the previous definition. What are the definitions of F_n and f_n exactly? Why theorem1 is not affected by data heterogeneity?\n5. In remark 1, minimizing upper bounds does not mean it is the best way.\n6. There is a lacking in the empirical validation for the benefits (e.g., time) of the proposed method when applying to heterogeneous devices.\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "1: The contributions are neither significant nor novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}]

Author rebuttal:

Decision: Accept: poster