FedPeWS: Personalized Warmup via Subnetworks for Enhanced Heterogeneous Federated Learning

## Abstract

Abstract Statistical data heterogeneity is a significant barrier to convergence in federated learning (FL). While prior work has advanced heterogeneous FL through better optimization objectives, these methods fall short when there is extreme data heterogeneity among collaborating participants. We hypothesize that convergence under extreme data heterogeneity is primarily hindered due to the aggregation of conflicting updates from the participants in the initial collaboration rounds. To overcome this problem, we propose a warmup phase where each participant learns a personalized mask and updates only a subnetwork of the full model. This personalized warmup allows the participants to focus initially on learning specific subnetworks tailored to the heterogeneity of their data. After the warmup phase, the participants revert to standard federated optimization, where all parameters are communicated. We empirically demonstrate that the proposed personalized warmup via subnetworks ( FedPeWS ) approach improves accuracy and convergence speed over standard federated optimization methods.

## 1 Introduction

Federated learning (FL) is a distributed learning paradigm where participants collaboratively train a global model by performing local training on their data and periodically sharing local updates with the server. The server, in turn, aggregates the local updates to obtain the global model, which is then transmitted to the participants for the next round of training []. While FL preserves data confidentiality by avoiding collating participant data at the server, statistical heterogeneity between local data distributions is a significant challenge in FL []. Several attempts have been made to tackle heterogeneity via federated optimization algorithms [, , , , , , , ], dropout [, ], and batch normalization [].

Consider the scenario where multiple hospitals collaborate to learn a medical image classification model that works across imaging modalities and organs, where the data from each hospital pertains to a different modality (e.g., histopathology, CT, X-ray, ultrasound, etc.) and/or organ (e.g., brain, kidney, colon, etc.). Most of the existing heterogeneous FL algorithms fail when there is such extreme data heterogeneity among collaborating participants, especially when the model is learned from scratch (with random initialization). The main reason for this failure is the high degree of conflicts between the local updates during the initial collaboration rounds. While enforcing a strong regularization constraint on the local updates ([]) can partially alleviate this problem, it dramatically slows down local learning and hence, convergence speed.

In this work, we explore an alternate approach to minimize the initial conflicts between heterogeneous participants by allowing participants in FL to initially train a partial subnetwork using only their local datasets. This warmup phase enables the participants to focus first on learning their local data well before engaging in broader collaboration. Thus, our proposed approach can be summarized as follows (see Figure 1). Initially, each participant uses a personalized binary mask tailored to their data distributions, allowing them to first learn their local data distributions and optimize their local (sparse) models. During this warmup phase, participants transmit only their masked updates to the server, and this process continues for a certain number of collaboration rounds. At the end of the warmup phase, the participants switch to standard federated optimization methods for subsequent collaboration. Our contributions are as follows:

Figure: Figure 1: Conceptual illustration of training personalized subnetworks in federated learning.
Refer to caption: x1.png

- 1.
We introduce a novel concept in federated learning, termed as personalized warmup via subnetworks (FedPeWS), which helps the global model to generalize to a better solution in fewer communication rounds. This is achieved through a neuron-level personalized masking strategy that is compatible with other FL optimization methods.
- 2.
We propose an algorithm to identify suitable subnetworks (subset of neurons) for each participant by simultaneously learning the personalized masks and parameter updates. The proposed algorithm does not make any assumptions regarding the data distributions and incorporates a mask diversity loss to improve the coverage of all neurons in the global model.
- 3.
For simple cases involving a small number of participants with known data distributions, we show that it is possible to skip the mask learning step and use fixed masks (that partition the network) determined by the server. We refer to this variant as FedPeWS-Fixed.
- 4.
We empirically demonstrate the efficacy of the FedPeWS approach under both extreme non-i.i.d. and i.i.d. data scenarios using three datasets: a custom synthetic dataset, a combination of MNIST and CIFAR-10 datasets, and a combination of three distinct medical datasets (PathMNIST, OCTMNIST and TissueMNIST).

## 2 Related Work

#### Collaborative Learning.

FL is a distributed learning paradigm that addresses data confidentiality concerns [], particularly in environments where data can not be centralized due to regulatory or practical reasons []. One of the seminal FL algorithms, FedAvg [], involves participants training models locally on their data and periodically transmitting their model parameters to a central server. The server averages these parameters to update the global model, which is then redistributed to the participants for further local refinement. FedAvg has inspired a plethora of variants and extensions aimed at enhancing performance [, , ], scalability [, ], communication efficiency [, , ], privacy/confidentiality [, , ], robustness [], and fairness [, , ]. For example, strategies such as weighted averaging or adaptive aggregation have been proposed to accommodate the non-i.i.d. nature of distributed data sources $-$ a scenario where data is not identically distributed across all participants, which can significantly hinder model performance [, , , , , ]. Specifically, FedProx [] addresses data heterogeneity by integrating a proximal term into the FedAvg framework. There is also a body of work that focuses on addressing the heterogeneity problem through personalization-based approaches, utilizing local-centric objectives [, , , , , ].

#### Independent Subnet Training.

Independent subnet training (IST) is a variant of distributed learning that focuses on enhancing model personalization and reducing communication overhead by training separate subnetworks for different participants []. IST distributes neurons of a fully connected neural network disjointly across different participants, forming a group of subnets. Then, each of these subnets is trained independently for one or more local SGD steps before synchronization. In every round, after broadcasting the server weights, each participant gets updated neurons to focus on, and the local subnet training continues. This approach led to different works along the line of using subnetwork training for efficiency [, , , , , , ] in FL. In our work, we adopt IST’s core principle of selecting neurons rather than focusing on weight values, which in turn narrows the search space. A key distinction between our method and IST lies in how the neurons are selected and the necessity of covering all neurons. Whilst IST typically involves random sampling of masks in each training round by the server and full coverage of neurons, we do not randomly sample neurons; instead, we use a learnable mask for each participant that is trained along with the parameters, and we relax the assumption of full coverage of neurons.

#### Finding Subnetworks in FL.

Another relevant idea is the Lottery Ticket Hypothesis (LTH) [], which attempts to identify subnetworks within a larger network. LTH is a model personalization technique, which focuses on sparsifying the network to create a smaller-scale version that improves per-round communication efficiency. In contrast to LTH, our method is directed towards training a shared global model and simultaneously improving convergence speed (reducing number of communication rounds). After LTH, there has been a growing interest in finding sparse and trainable networks at initialization [, , ]. Recently, in [], sparse networks were found inside the main model to increase communication efficiency in FL. The proposed FedPM method focuses on finding a subnetwork by freezing the model weights and training for masks on a weight level, in contrast to IST, which works on a neuron level. FedPM utilizes the sigmoid function to obtain probability values from unbounded mask scores and then uses Bernoulli sampling to obtain binary masks. We use a similar approach in our FedPeWS algorithm to learn the neuron-level personalized masks.

## 3 Preliminaries

Our goal is to minimize a sum-structured federated learning optimization objective:

$$ $\begin{split}x^{\star}\leftarrow\operatorname*{arg\,min}_{x\in\mathbb{R}^{d}} \Bigg{[}f(x)\coloneqq\frac{1}{N}\sum_{i=1}^{N}f_{i}(x)\Bigg{]},\end{split}$ (1) $$

where the components $f_{i}:\mathbb{R}^{d}\rightarrow\mathbb{R}$ are distributed among $N$ local participants and are expressed in a stochastic format as $f_{i}(x)\coloneqq\mathbb{E}_{\xi\sim\mathcal{D}_{i}}\big{[}F_{i}(x,\xi)\big{]}$. Here, $\mathcal{D}_{i}$ represents the distribution of $\xi$ at participant $i\in[N]\coloneqq\{1,\ldots,N\}$. This problem encapsulates standard empirical risk minimization as a particular case when each $\mathcal{D}_{i}$ is represented by a finite set of $n_{i}$ elements, i.e., $\xi_{i}=\{\xi_{i}^{1},\ldots,\xi_{i}^{n_{i}}\}$. In such cases, $f_{i}$ simplifies to $f_{i}(x,\xi_{i})=\frac{1}{n_{i}}\sum_{j=1}^{n_{i}}F_{i}(x,\xi_{i}^{j})$. Our approach does not impose restrictive assumptions on the data distribution $\mathcal{D}_{i}$. In fact, we specifically focus on the extreme heterogeneous (non-i.i.d.) setting, where $\mathcal{D}_{i}\neq\mathcal{D}_{i^{{}^{\prime}}},\forall~{}i\neq i^{{}^{\prime}}$ and the local optimal solution $x_{i}^{\star}\leftarrow\operatorname*{arg\,min}_{x\in\mathbb{R}^{d}}f_{i}(x)$ might significantly differ from the global minimizer of the objective function in Equation 1.

We are especially interested in the supervised classification task and let $\mathcal{M}_{x}:\mathcal{Z}\rightarrow\mathcal{Y}$ be a classifier parameterized by $x$. Here, $\mathcal{Z}\subseteq\mathbb{R}^{D}$ and $\mathcal{Y}=\{1,2,\ldots,M\}$ denote the input and label spaces, respectively, $D$ is the input dimensionality, $M$ is the number of classes, and $d$ represents the number of parameters in the model $\mathcal{M}$. We set $F_{i}(x,\xi_{i}^{j})=\mathcal{L}(\mathcal{M}_{x}(\mathbf{z}_{i}^{j}),y_{i}^{j})$, where $\mathcal{L}$ is an appropriate loss function and $\xi_{i}^{j}\coloneq(\mathbf{z}_{i}^{j},y_{i}^{j})$ is a labeled training sample such that $\mathbf{z}_{i}^{j}\in\mathcal{Z}$ and $y_{i}^{j}\in\mathcal{Y}$. Furthermore, we mainly focus on the cross-silo FL setting ($N$ is small).

#### Federated Averaging (FedAvg).

A common approach for solving Equation 1 in the distributed setting is FedAvg []. This algorithm involves the participants performing $K$ local steps of stochastic gradient descent (SGD) and communicating with the server over $T$ communication rounds. The server initializes the global model with $x_{g}^{0}$ and broadcasts it to all participants, which is then used to initialize the local models, i.e., $x_{i}^{1,0}=x_{g}^{0}$. In each communication round, the updates from the participants are averaged on the server and sent back to all participants. For a local step $k\in[K]$, communication round $t\in[T]$, and participant $i\in[N]$, the local and global iterates are updated as:

$$ $x_{i}^{t,k}=x_{i}^{t,k-1}-\eta_{\ell}\nabla f_{i}\big{(}x_{i}^{t,k-1},\xi_{i}^ {t,k-1}\big{)},\quad x_{i}^{t}=x_{i}^{t,K},\text{ and }x_{g}^{t}=x_{g}^{t-1}- \eta_{g}(x_{g}^{t-1}-\frac{1}{N}\sum_{i=1}^{N}x_{i}^{t}),$ (2) $$

where $\eta_{\ell}$ and $\eta_{g}$ are the local and global learning rates, respectively. The server then broadcasts the updated global model $x_{g}^{t}$ to all participants, which is then used to reinitialize the local models as $x_{i}^{t+1,0}=x_{g}^{t}$.

In the FedAvg algorithm, the number of communication rounds necessary to achieve a certain precision is directly proportional to the heterogeneity measure []. Notably, this relationship holds true in convex settings; however, in non-convex scenarios, the algorithm may either not converge or may converge to a suboptimal solution. Stemming from this observation, our objective is to reduce the number of requisite communication rounds to achieve convergence, while simultaneously achieving a better solution.

Figure: Algorithm 1 FedPeWS (For FedPeWS-Fixed variant, the steps highlighted in green are omitted and instead the server sets $m_{i}^{t}=m_{i}$, $\forall~{}t\in[W]$.)

## 4 Proposed FedPeWS Method

Figure: Figure 2: Illustration of the proposed FedPeWS algorithm for two participants, which aggregates partial subnetworks ($x_{i}^{t}\odot m_{i}^{t}$) during the warmup phase to obtain a shared global model $x_{g}^{t}$. Here, $x_{i}^{t}$ and $m_{i}^{t}$ denote the local model and personalized mask of the $i^{\text{th}}$ participant in the $t^{\text{th}}$ round.
Refer to caption: x2.png

The core idea of the proposed FedPeWS method is to allow participants to learn only a personalized subnetwork (a subset of parameters) instead of the entire network (all parameters) during the initial warmup phase. Let $m_{i}\in\{0,1\}^{d}$ be a binary mask vector denoting the set of parameters that are learned by participant $i$, $i\in[N]$. Note that $m_{i}(\ell)=1$ indicates that the $\ell^{\textrm{th}}$ element of $x_{i}$ is selected for learning (value 0 0 indicates non-selection), $\ell\in[d]$. Thus, during the warmup phase, the objective in FedPeWS is to learn the parameters $x$ along with the personalized masks $m_{i}$, i.e.,

$$ $\min_{x,\{m_{i}\}_{i\in[N]}}\frac{1}{N}\sum_{i=1}^{N}f_{i}(x\odot m_{i}),$ (3) $$

$\odot$ denotes element-wise multiplication. Note that $\mathcal{M}_{x\odot m_{i}}$ denotes the personalized subnetwork of participant $i$. When personalized masks are employed, the update rules can be modified as:

$$ $x_{i}^{t,k}=x_{i}^{t,k-1}-\eta_{\ell}\nabla f_{i}\big{(}x_{i}^{t,k-1}\odot m_{ i}^{t},\xi_{i}^{t,k-1}\big{)}\;\text{and}\;x_{g}^{t}=x_{g}^{t-1}-\eta_{g}(x_{g }^{t-1}-\frac{\sum_{i\in[N]}x_{i}^{t}\odot m_{i}^{t}}{\sum_{i\in[N]}m_{i}^{t}}).$ (4) $$

The obvious questions regarding the FedPeWS method are: (i) how to learn these personalized masks $m_{i}$? and (ii) what should be the length of the warmup period?

Identification of personalized subnetworks: It is not straightforward to directly optimize for the personalized binary (discrete) masks $m_{i}$ in Equation 3. Hence, we make the following design choices. Firstly, personalized masks are learned at the neuron-level and then expanded to the parameter-level. Following IST [], masks are specifically applied only to the hidden layer neurons, while the head and tail neurons remain unaffected. However, unlike IST, the neuron-level masks are not randomly selected in each collaboration round. Instead, we learn real-valued personalized neuron-level mask score vectors $s_{i}\in\mathbb{R}^{h}$, which in turn can be used to generate the binary masks. Here, $h$ denotes the number of hidden neurons in the classifier $\mathcal{M}$ and $h\ll d$. A higher value of element $s_{i}(\ell)$, $\ell\in[h]$, indicates that the $\ell^{\textrm{th}}$ neuron is more likely to be selected by participant $i$. Let $\mathcal{G}:\mathbb{R}^{h}\rightarrow\{0,1\}^{d}$ be the mask generation function that generates the binary parameter-level masks $m_{i}$ from neuron-level mask score vectors $s_{i}$, i.e., $m_{i}=\mathcal{G}(s_{i})$. $\mathcal{G}$ consists of three steps. Firstly, we convert $s_{i}$ into probabilities by applying a sigmoid function, i.e., $\theta_{i}=\sigma(s_{i})$, where $\theta_{i}\in[0,1]^{h}$ is the mask probability vector and $\sigma$ is the sigmoid function. Next, binary neuron masks $\tilde{m}_{i}$ are obtained by sampling from a Bernoulli distribution with parameter $\theta_{i}$, i.e., $\tilde{m}_{i}(\ell)\sim Bernoulli(\theta_{i}(\ell))$, $\forall~{}\ell\in[h]$. Finally, these binary neuron masks can be directly mapped to the binary parameter-level mask $m_{i}$, i.e., if a neuron is selected, all the weights associated with the selected neuron are also selected. Thus, Equation 3 can be reparameterized as:

$$ $\min_{x,\{s_{i}\}_{i\in[N]}}\frac{1}{N}\sum_{i=1}^{N}f_{i}(x\odot\mathcal{G}(s _{i})).$ (5) $$

The above equation can be optimized alternatively for the mask score vectors $s_{i}$ and the parameters $x$. The participants first optimize for the mask scores while the model parameters $x_{i}^{t,k}$ are frozen (Procedure I), and then switch to optimizing the model parameters while freezing the mask scores (Procedure II). In the mask training step (Procedure I), the optimization objective is defined as:

$$ $\mathcal{L}_{s}=f_{i}\Big{(}x_{i}^{t,k}\odot\mathcal{G}\left(s_{i}^{t,k}\right ),\xi_{i}^{t,k-1}\Big{)}-\lambda\|\sigma\left(s_{i}^{t,k}\right)-\theta_{g \backslash\{i\}}^{t}\|_{2}^{2};\qquad s_{i}^{t,k+1}\leftarrow s_{i}^{t,k}-\eta _{s}\nabla_{s}\mathcal{L}_{s},$ (6) $$

where $\nabla_{s}$ indicates that the gradient is w.r.t. mask score vector $s$, $\eta_{s}$ is the local learning rate for updating $s$, $\theta_{g}^{t}$ is the global mask probability at round $t$, $\theta_{g\backslash\{i\}}^{t}$ is the global mask probability excluding the probability mask of the current participant $i$, and $\lambda$ is the weight assigned to the mask diversity measure (second term). It is important to note that the personalized masks may not cover all neurons in the network. Maximizing the mask diversity measure encourages personalized masks to deviate as much as possible from the global mask, which facilitates better coverage of all the neurons in the global model. The diversity measure has an upper bound due to the sigmoid function:

$$ $\|\sigma\left(s_{i}^{t,k}\right)-\theta_{g\backslash\{i\}}^{t}\|_{2}^{2}\leq h.$ (7) $$

Given the difficulty in calculating $\nabla_{s}\mathcal{L}_{s}$ directly due to the discrete nature of Bernoulli sampling, we employ the straight-through estimator (STE) [] to approximate the gradients, which does not compute the gradient of the given function and passes on the incoming gradient as if the function was an identity function.

During Procedure II, the optimization function for the model weights is expressed as:

$$ $\mathcal{L}_{x}=f_{i}\Big{(}x_{i}^{t,k}\odot\mathcal{G}\left(s_{i}^{t,k}\right ),\xi_{i}^{t,k-1}\Big{)};\qquad x_{i}^{t,k+1}\leftarrow x_{i}^{t,k}-\eta_{\ell }\nabla_{x}\mathcal{L}_{x},$ (8) $$

where $\nabla_{x}$ indicates that the gradient is w.r.t. weights $x$. The FedPeWS algorithm alternates between these two procedures for $W$ rounds, where $W$ is the number of warmup rounds.
At this point, the warmup stops and the participants switch to standard training for $(T-W)$ collaboration rounds.
This approach ensures that each participant effectively contributes to the FL process while also tailoring the learning to their specific data distributions. The number of warmup rounds $W$ (or the proportion of warmup rounds $\tau=\frac{W}{T}$) is a key hyperparameter of the FedPeWS algorithm, along with the weight $\lambda$ assigned to the mask diversity loss. While it would be ideal to have a principled method to select these hyperparameters, we use a grid search to tune them, which is currently a limitation.

Use of fixed subnetworks: When the number of participants is small and the data distributions of the participants are known apriori, the server can partition the full model into subnetworks of the same depth and assign a fixed subnetwork to each participant, i.e., $m_{i}^{t}=m_{i}$, $\forall~{}t\in[W]$. Participants transmit only the masked updates back to the server during warmup, which then aggregates these masked parameters and redistributes them in their masked form. For the sake of utility, the server can design personalized masks such that the union of these masks covers all the neurons. This variant of FedPeWS is referred to as FedPeWS-Fixed and follows the same algorithm in Algorithm 1, except for the omission of the highlighted (green) steps.

## 5 Experiments and Results

### 5.1 Datasets and Network Architecture

#### Synthetic Dataset.

To effectively evaluate the performance of the proposed algorithm, we generated a custom synthetic dataset that simulates the extreme non i.i.d. scenario. This dataset encompasses four classes, each characterized by four 2D clusters determined by specific centers and covariance matrices. Note that the clusters from different classes interleave each other as shown in Figure 3. For this dataset, we utilize a neural network consisting of five fully-connected (FC) layers, each followed by ReLU activation functions, except the last layer. To enhance the dataset complexity and aid FC network learning, we transform these 2D points into 5D space using the transformation $[x,y,x^{2},y^{2},xy]$, based on their $(x,y)$ coordinates. We generate two versions of this dataset, Synthetic-32K and Synthetic-3.2K, depending on the number of data points in the training set. The former has $32000$ samples, with each class containing $8000$ data points, while the latter has ten times fewer data points.

Figure: Figure 3: Samples from the custom synthetic dataset.
Refer to caption: x3.png

#### CIFAR-MNIST.

We integrate two distinct datasets, CIFAR-10 [] and MNIST [], to explore how different clients might adapt when faced with disparate data sources. CIFAR-10 comprises of $32\times 32$ pixel images categorized into 10 object classes. MNIST, typically featuring $28\times 28$ pixel images across 10 digit classes, is upscaled to $32\times 32$ pixel to standardize dimensions with CIFAR-10. We compile a balanced dataset by randomly selecting 400 samples from each class for the training set and 200 samples for the test set from the combined pool of 20 classes. This setup aims to simulate a FL environment where multiple clients handle significantly varied data types. For this dataset, we employ a convolutional neural network comprising four convolutional layers, each having a kernel size of 3 and padding of 1, followed by max pooling. This is succeeded by a fully connected layer. This architecture was used because of its simplicity and widespread use in the literature [, ].

#### {Path-OCT-Tissue}MNIST.

We amalgamate three distinct medical datasets: PathMNIST, OCTMNIST, and TissueMNIST [], to develop a universal medical prognosis model capable of recognizing various tasks using a single model. The datasets contain 9, 4, and 8 classes, respectively, totaling to 21 classes. For this dataset, we utilized the same architecture and training details described in the CIFAR-MNIST dataset.

### 5.2 Experimental Setup

#### Dataset partitioning.

For scenarios with a smaller number of collaborators ($N={2,3,4}$), we manually partition the training dataset to tailor the data distribution to specific participants. In the $N=2$ scenario, we partition as follows: (i) For the Synthetic dataset, encompassing both Synthetic-32K and Synthetic-3.2K, even-numbered classes are assigned to participant 1, while odd-numbered classes are allocated to participant 2. (ii) For the CIFAR-MNIST combination, all CIFAR-10 samples are assigned to participant 1, with MNIST samples allocated to participant 2. In the $N=3$ scenario, the {Path-OCT-Tissue}MNIST dataset is partitioned into three splits corresponding to the individual datasets, with PathMNIST assigned to participant 1, OCTMNIST to participant 2, and TissueMNIST to participant 3. For the $N=4$ scenario, the synthetic dataset is divided so that each class is exclusively allocated to one of the four participants.

For scenarios with a larger number of participants ($N\geq 10$), we employ a Dirichlet distribution to partition the training set. This approach utilizes a concentration parameter $\alpha$ to simulate both homogeneous and heterogeneous data distributions [, , , ]. We experiment with various values of $\alpha$, specifically $\alpha\in\{0.1,0.5,1.0,2.0,5.0\}$, to explore the effects of dataset heterogeneity (lower $\alpha$ values) and homogeneity (higher $\alpha$ values) on the model performance. This methodological diversity allows us to comprehensively assess our approach under varying data conditions. Results for large $N$ ($>100$) are reported in the appendix.

#### Training details.

In federated optimization, we primarily benchmark against the FedAvg algorithm [], a standard approach in federated learning. However, our algorithm is designed to be versatile, functioning as a ‘plug-and-play’ solution that is compatible with various other optimizers. To demonstrate this adaptability, we also conduct experiments using FedProx [], showcasing our method’s capabilities across different optimization frameworks. For our experiments, we fix the local learning rate $\eta_{\ell}=0.001$ in the Synthetic-32K dataset case, and we set $\eta_{\ell}=0.01$ for other experiments. Also, the mask learning rate is fixed $\eta_{s}=0.1$. Furthermore, we vary the global learning rate $\eta_{g}\in\{0.1,0.25,0.5,1.0\}$ to observe the differences in optimization behavior between the baseline and our proposed methods. Additionally, we employ two distinct batch sizes $\{32,8\}$ for Synthetic-32K and Synthetic-3.2K, respectively. For experiments involving the CIFAR-MNIST and {Path-OCT-Tissue}MNIST datasets, we standardize the batch size to 64.
We conduct our experiments on NVIDIA RTX A6000 GPUs on an internal cluster server, with each run utilizing a single GPU. The execution time for each run is capped at less than an hour, which indicates the maximum execution time rather than the average. All results are averaged over three independent runs and the average accuracy is reported on the global test dataset.

**Table 1: The required number of collaboration rounds to reach target accuracy $\upsilon$ $\%$ and the final accuracy after $T$ rounds. The results are averaged over 3 seeds. $\times$ indicates that the algorithm cannot reach target accuracy $\upsilon$ within $T$ rounds and NA means that it reaches $\upsilon$ only in one random seed.**
| Dataset $/$ Batch size | Synthetic-32K, $32$ | Synthetic-3.2K, $8$ |  |  |  |
| --- | --- | --- | --- | --- | --- |
| Parameters $\{\eta_{g}/\lambda/\tau\}$ | $\{1.0/5.0/0.125\}$ | $\{0.5/2.0/0.2\}$ | $\{0.25/1.0/0.1875\}$ | $\{0.1/2.0/0.1\}$ |  |
| Target accuracy $\upsilon(\%)$ | $99$ | $90$ | $75$ | $99$ |  |
| No. of rounds to<br>reach target accuracy | FedAvg | $148\pm 3.79$ | $199\pm\text{NA}$ | $\times$ | $371\pm\text{NA}$ |
| FedPeWS | $\mathbf{115\pm 7.21}$ | $\mathbf{182\pm 6.81}$ | $\mathbf{286\pm 7.93}$ | $\mathbf{301\pm 10.59}$ |  |
| Final accuracy after<br>$T$ collaboration rounds | FedAvg | $99.94\pm 0.05$ | $91.40\pm 7.25$ | $67.64\pm 0.90$ | $97.33\pm 3.89$ |
| FedPeWS | $\mathbf{99.96\pm 0.01}$ | $\mathbf{99.49\pm 0.60}$ | $\mathbf{83.50\pm 3.52}$ | $\mathbf{99.66\pm 0.19}$ |  |

Figure: Figure 4: Results of the experiments on Synthetic-{32, 3.2}K datasets with batch sizes {32, 8}, with different global learning rates $\eta_{g}\in\{1.0,0.5,0.25,0.1\}$ and communication rounds $T\in\{200,250,400,500\}$. Refer to Table 1 for the corresponding numbers. In all the above scenarios, FedPeWS converges faster to a better solution compared to FedAvg.
Refer to caption: x4.png

### 5.3 Experimental Results

Figure: Figure 5: Visualization of validation accuracy and loss on the Synthetic-32K dataset with $N=4$ participants and a global learning rate $\eta_{g}=1.0$.
Refer to caption: x5.png

Our experimental analysis focuses on assessing the performance of our proposed FedPeWS algorithm within the FL framework. The key findings from our studies are as follows: (i) The FedPeWS approach demonstrates a significant reduction in the number of communication rounds required to achieve target accuracy while also enhancing the final accuracy post-convergence. (ii) The FedPeWS algorithm is robust across different levels of data heterogeneity. (iii) In scenarios where full knowledge of the participant data distributions is available, the server can employ the FedPeWS-Fixed method (Figure 6). While the FedPeWS-Fixed variant shows competitive effectiveness comparable to our primary FedPeWS algorithm, the latter offers broader applicability in real-world settings.

#### Improved communication efficiency and accuracy.

We initially report the required number of communication rounds to reach the target accuracy and the final accuracy after $T$ communication rounds for the synthetic dataset in Table 1. The results underscore that the incorporation of a personalized warmup phase in a federated setup significantly reduces the required number of communication rounds across all tested scenarios. Notably, in specific instances, such as with the Synthetic-32K dataset and $\eta_{g}=0.25$, the conventional FedAvg algorithm does not meet the target accuracy within the $T$ communication rounds. Conversely, in scenarios where $\eta_{g}\in\{0.25,0.1\}$, FedAvg only achieves the target accuracy in one of the seeds, exhibiting suboptimal performance in the other two runs. From Figure 4, it is evident that our proposed FedPeWS algorithm surpasses FedAvg in both communication efficiency and accuracy.

We also consider a more extreme data heterogeneity scenario with $N=4$ participants, depicted in Figure 5, where FedAvg completely fails by reaching only $58.4\pm 2.33\%$, whereas our FedPeWS approach reaches $91.13\pm 3.55\%$ accuracy by significantly outperforming the base optimizer (FedAvg) with a gain of $\mathbf{32.72\%}$. It is crucial to highlight that in this experiment, we set $\lambda=0.0$, effectively not enforcing diversity as outlined in Equation 6. This approach focuses solely on optimizing the masks using the first loss component, which depends only on the data distributions of each participant. This shows that, in specific scenarios, we can learn the personalized masks (Procedure I) without the need to adjust the $\lambda$ parameter, while still achieving a better performance than the base optimizer.

Figure: (a) CIFAR-MNIST dataset.
Refer to caption: x6.png

Figure: Figure 7: Top: illustration of number of samples per class allocated to each client, that is indicated by dot sizes, for different concentration $\alpha$ values. Bottom: visualization of the experiments on CIFAR-MNIST dataset with $N=10$ participants with different levels of heterogeneity.
Refer to caption: x8.png

#### Sensitivity to λ 𝜆 \lambda italic_λ and τ 𝜏 \tau italic_τ parameters.

Figure 6 showcases the results of experiments on the CIFAR-MNIST dataset with $N=2$ participants and {Path-OCT-Tissue}MNIST dataset with $N=3$ participants.
The left-side plots of Figures 6(a) and 6(b), which depict the performance of the global model (averaged over 3 runs), demonstrate that our method consistently achieves higher accuracy. The right side figures feature heatmap plots that annotate the global model accuracy obtained varying $\lambda\in\{0,0.1,0.3,0.5,1,2,5,10,100,1000\}$ and $\tau\in\{0.0,0.2,0.4,0.5,0.6,0.8\}$ parameters. An additional row labeled $(\lambda=-)$ represents the FedPeWS-Fixed approach, where user(server)-defined fixed masks are employed. In this method, we simply split the full network into $N$ partitions, with each partition assigned to a participant (for detailed instructions on setting masks, please see Section B.2 in the appendix). The results indicate that our approach has a low sensitivity to variations in $\lambda$ and $\tau$. For more detailed insights, please refer to Tables 6 and 7 in the appendix.

#### Varying degrees of heterogeneity.

Figure 7 demonstrates that our FedPeWS approach consistently outperforms FedAvg, with gains directly related to the degree of data heterogeneity. The figure clearly shows that the advantage of using our method is more pronounced under conditions of high data heterogeneity. As heterogeneity levels decrease, our method becomes comparable to FedAvg.

**Table 2: The required number of collaboration rounds to reach target accuracy $\upsilon$ $\%$ using FedProx algorithm and the final accuracy after $T$ rounds. The results are averaged over 3 seeds. $\times$ indicates that the algorithm cannot reach target accuracy $\upsilon$ within $T$ rounds.**
| Dataset $/$ Batch size | Synthetic-32K, $32$ | Synthetic-3.2K, $8$ |  |  |  |
| --- | --- | --- | --- | --- | --- |
| Parameters $\{\eta_{g}/\lambda/\tau\}$ | $\{1.0/0.1/0.125\}$ | $\{0.5/0.1/0.2\}$ | $\{0.25/1.0/0.1875\}$ | $\{0.1/1.0/0.1\}$ |  |
| Target accuracy $\upsilon(\%)$ | $99$ | $90$ | $75$ | $99$ |  |
| No. of rounds to<br>reach target accuracy | FedProx | $138\pm 13.22$ | $\times$ | $\times$ | $362\pm 20.00$ |
| FedPeWS | $\mathbf{115\pm 5.29}$ | $\mathbf{211\pm 16.52}$ | $\mathbf{314\pm 27.83}$ | $\mathbf{344\pm 27.30}$ |  |
| Final accuracy after<br>$T$ collaboration rounds | FedProx | $99.95\pm 0.02$ | $82.43\pm 7.98$ | $69.26\pm 6.03$ | $\mathbf{99.92\pm 0.06}$ |
| FedPeWS | $\mathbf{99.96\pm 0.01}$ | $\mathbf{98.40\pm 1.84}$ | $\mathbf{90.40\pm 3.91}$ | $\mathbf{99.92\pm 0.07}$ |  |

Figure: Figure 8: Comparison of our proposed method and FedProx [] on Synthetic-{32, 3.2}K datasets. Refer to Table 2 for the corresponding numbers.
Refer to caption: x9.png

#### FedProx.

We also present results using the FedProx optimizer on Synthetic-32K and Synthetic-3.2K datasets in Figure 8 and Table 2, employing global learning rates $\eta_{g}=\{1.0,0.5,0.25,0.1\}$. Note that we adapt Algorithm 1 to incorporate the FedProx algorithm as the base optimizer, instead of FedAvg.
We selected the best performing proximal term scaler $0.01$ after tuning and evaluating different values from a set of potential values $\{0.001,0.01,0.1,0.5\}$, based on the findings in []. The results demonstrate that FedPeWS outperforms FedProx in terms of both communication efficiency and final accuracy across the tested scenarios, except the last scenario (Synthetic-3.2K dataset with batch size 8 and $\eta_{g}=0.1$), where the performance of FedPeWS is comparable to that of FedProx.

## 6 Conclusion

In this work, we introduced a novel concept called personalized warmup via subnetworks for heterogeneous FL $-$ a strategy that enhances convergence speed and can seamlessly integrate with existing optimization techniques. Results demonstrate that the proposed FedPeWS approach and achieves higher accuracy than the relevant baselines, especially when there is extreme statistical heterogeneity. Limitations of FedPeWS include the need to tune two additional hyperparameters (no. of warmup rounds and mask diversity weight) and the lack of theoretical convergence analysis.

## Appendix A Related Work

Table 3 shows the comparison of our proposed approach to the existing literature.

**Table 3: Comparison of approaches for handling data heterogeneity in federated learning.**
|  | Shared<br>Global Model | Level of Mask<br>Personalization | Learnable<br>Mask | Learnable<br>Parameters |
| --- | --- | --- | --- | --- |
| FedPM [] | ✓ | parameter-level | ✓ | $\bm{\times}$ |
| IST [] | ✓ | neuron-level | $\bm{\times}$ (random) | ✓ |
| LTH [] | $\bm{\times}$ | parameter-level | $\bm{\times}$ (pruned) | ✓ |
| FedWeIT [] | $\bm{\times}$ | parameter-level | ✓ | ✓ |
| FjORD [] | ✓ | parameter-level | $\bm{\times}$ (slimmed) | ✓ |
| FedPeWS (ours) | ✓ | neuron-level | ✓ | ✓ |

## Appendix B Additional Experimental Details

### B.1 Network Architecture Details

The network for the synthetic dataset (detailed in Section 5.2) consists of five fully connected (FC) layers, each followed by ReLU activation functions, except for the last layer. We provide the details of this architecture in Table 4.

**Table 4: Architecture for synthetic dataset models used in the experiments.**
| Layer | Input | FC1 | FC2 | FC3 | FC4 | FC5 |
| --- | --- | --- | --- | --- | --- | --- |
| Dimensions | [5] | [5, 32] | [32, 64] | [64, 128] | [128, 32] | [32, 4] |

The network for the CIFAR-MNIST and {Path-OCT-Tissue}MNIST datasets includes three convolutional layers followed by max pooling, and a fully connected layer. The details of this architecture is provided in Table 5.

**Table 5: Architecture for CIFAR-MNIST dataset models. Every convolutional layer is followed by a max pooling layer with kernel size 2 and stride 2.**
| Layer | Input | Conv1 | Conv2 | Conv3 | Flatten | FC |
| --- | --- | --- | --- | --- | --- | --- |
| Dimensions | [3, 32, 32] | [3, 32, 3, 3] | [32, 64, 3, 3] | [64, 128, 3, 3] | [2048] | [2048, 20] |

### B.2 Fixed Mask Generation

Figure 9 illustrates how we design masks for FedPeWS-Fixed experiments in scenarios with $N=2$ participants. For cases involving $N=4$ participants, the full network $\mathcal{M}_{x}$ (classifier $\mathcal{M}$ parameterized with $x$) is divided into four subnetworks, vertically, with each subnetwork corresponding to one of the participants. As such, we vertically partition the hidden neurons in the network into $N$ groups (subnetworks) and design the mask to assign each group to one participant, ensuring no overlap. This design choice is based on the assumption that classes held by each participant are highly heterogeneous, thus preventing any intersection in the masks. This setting is specifically tailored for the FedPeWS-Fixed method and doesn’t necessitate performing optimization over the masks $m_{i}$, they are kept fixed.

Figure: Figure 9: Illustration of manual mask setting in the FedPeWS-Fixed method. The left figure illustrates the complete network with all neurons active and full connections. The middle figure represents subnetwork 1, utilizing only the left portion of the full network, where $m_{1}$ corresponds to this left side. Conversely, the right figure indicates the part of the network used for subnetwork 2. This setting is employed in all experiments involving $N=2$ participants.
Refer to caption: x10.png

## Appendix C Experimental Results

### C.1 Wall-clock Time vs. Accuracy

Figure 10 illustrates the wall-clock time versus accuracy results, which correspond to Figure 4 in the main paper. From this comparison, FedPeWS demonstrates a slightly improved performance over FedAvg in terms of wall-clock time in two of the four scenarios. However, it underperforms slightly in the remaining two scenarios, with only a marginal increase in time. This variance is attributed to the alternation between training masks and weights during the warm-up phase, impacting the time efficiency.

Figure: Figure 10: Wall-clock time vs. accuracy plot corresponding to Figure 4 of the main paper.
Refer to caption: x11.png

### C.2 FedPeWS-Fixed . Mask Length Study

In this section, we explore the impact of mask length on the performance of the FedPeWS-Fixed method with parameter $W=120(\tau=0.4)$ using the CIFAR-MNIST dataset. We examine two scenarios for splitting the network into two subnetworks:

- 1.
$|m_{1}|<|m_{2}|$: 75% of the mask is assigned to Participant 2, 25% to Participant 1.
- 2.
$|m_{1}|=|m_{2}|$: equal sized masks are assigned to each participant.

Figure: Figure 11: Mask length study using FedPeWS-Fixed method on CIFAR-MNIST dataset.
Refer to caption: x12.png

Figure 11 displays the validation accuracy over $T=300$ communication rounds for both scenarios. The leftmost plot shows the accuracy of the global model, while the middle and rightmost plots the accuracy for each of the participants. Both mask length scenarios converge to a comparable accuracy levels, with a marginal difference of $0.5\%$ higher accuracy in the scenario where $|m_{1}|<|m_{2}|$. This is likely due to the larger mask size, which aids in learning the more complex CIFAR-10 dataset held by Participant 2.

### C.3 Sensitivity Analysis

In this section, we detail the sensitivity analysis of the $\lambda$ and $\tau$ parameters conducted on the CIFAR-MNIST dataset (Table 6) and the {Path-OCT-Tissue}MNIST dataset (Table 7). This analysis particularly includes the standard deviation of the accuracy achieved by the tested algorithms after $T$ communication rounds and over three independent evaluations. Results with the best performance are highlighted in green.

**Table 6: Ablation study of the parameters $\lambda$ and $\tau$ on the CIFAR-MNIST dataset with $N=2$ participants over three independent runs. The first column ($\tau=0.0$) corresponds to the FedAvg algorithm. The last row $(-)$ presents results for the FedPeWS-Fixed algorithm.**
|  | Proportion of warmup rounds $\tau=W/T$ |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- |
| $\mathbf{\lambda}(\downarrow),\tau(\rightarrow)$ | $\tau=0.0$ (FedAvg) | $\tau=0.2$ | $\tau=0.4$ | $\tau=0.5$ | $\tau=0.6$ | $\tau=0.8$ |
| 0.0 | $71.23\pm 0.71$ | $68.01\pm 0.88$ | $66.01\pm 0.48$ | $65.96\pm 1.03$ | $65.40\pm 1.95$ | $65.77\pm 0.41$ |
| 0.1 | $68.77\pm 1.09$ | $70.39\pm 1.00$ | $70.61\pm 1.17$ | $69.48\pm 0.37$ | $70.36\pm 2.06$ |  |
| 0.3 | $70.86\pm 0.23$ | $73.00\pm 0.65$ | $73.91\pm 0.71$ | $73.26\pm 0.46$ | $73.02\pm 1.05$ |  |
| 0.5 | $71.43\pm 0.56$ | $74.17\pm 0.91$ | $73.84\pm 0.12$ | $73.66\pm 0.88$ | $75.05\pm 0.45$ |  |
| 1.0 | $72.26\pm 0.54$ | $74.46\pm 0.44$ | $74.54\pm 0.76$ | $74.91\pm 0.42$ | $73.81\pm 0.87$ |  |
| 2.0 | $72.61\pm 0.79$ | $73.68\pm 0.17$ | $\mathbf{75.35\pm 0.50}$ | $74.76\pm 0.54$ | $74.46\pm 0.56$ |  |
| 5.0 | $72.60\pm 0.45$ | $\mathbf{75.22\pm 0.33}$ | $\mathbf{75.00\pm 0.74}$ | $75.01\pm 0.71$ | $73.96\pm 1.60$ |  |
| 10.0 | $72.29\pm 0.48$ | $\mathbf{74.97\pm 0.65}$ | $74.31\pm 0.95$ | $74.03\pm 0.30$ | $71.91\pm 2.69$ |  |
| 100.0 | $71.64\pm 0.47$ | $72.92\pm 0.39$ | $73.96\pm 0.65$ | $73.13\pm 0.73$ | $72.43\pm 3.55$ |  |
| 1000.0 | $71.58\pm 0.53$ | $73.18\pm 0.73$ | $73.32\pm 1.70$ | $73.87\pm 1.16$ | $72.52\pm 1.92$ |  |
| $-$ | $72.72\pm 0.44$ | $\mathbf{75.22\pm 0.19}$ | $\mathbf{75.05\pm 0.42}$ | $72.51\pm 3.89$ | $73.77\pm 0.20$ |  |

**Table 7: Ablation study of the parameters $\lambda$ and $\tau$ on the combination of {Path-OCT-Tissue}MNIST datasets with $N=3$ participants over three independent runs. The first column ($\tau=0.0$) corresponds to the FedAvg algorithm. The last row $(-)$ presents results for the FedPeWS-Fixed algorithm.**
|  | Proportion of warmup rounds $\tau=W/T$ |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- |
| ${\lambda}(\downarrow),\tau(\rightarrow)$ | $\tau=0.0$ (FedAvg) | $\tau=0.2$ | $\tau=0.4$ | $\tau=0.5$ | $\tau=0.6$ | $\tau=0.8$ |
| 0.0 | $52.25\pm 0.57$ | $53.04\pm 2.08$ | $50.35\pm 0.61$ | $47.83\pm 1.04$ | $47.89\pm 1.17$ | $46.80\pm 1.76$ |
| 0.1 | $52.52\pm 1.79$ | $51.72\pm 0.63$ | $49.83\pm 1.83$ | $48.83\pm 2.51$ | $47.50\pm 1.07$ |  |
| 0.3 | $52.74\pm 1.75$ | $54.36\pm 0.80$ | $51.07\pm 0.25$ | $50.94\pm 2.22$ | $50.62\pm 1.91$ |  |
| 0.5 | $54.30\pm 2.08$ | $54.42\pm 1.43$ | $52.02\pm 2.53$ | $51.43\pm 1.60$ | $48.97\pm 0.80$ |  |
| 1.0 | $\mathbf{54.89\pm 0.72}$ | $53.59\pm 1.40$ | $51.49\pm 1.19$ | $50.29\pm 2.70$ | $52.21\pm 1.02$ |  |
| 2.0 | $\mathbf{54.75\pm 1.12}$ | $\mathbf{54.45\pm 1.78}$ | $53.02\pm 0.45$ | $52.26\pm 1.81$ | $52.49\pm 0.49$ |  |
| 5.0 | $\mathbf{54.91\pm 0.91}$ | $\mathbf{54.99\pm 0.90}$ | $52.41\pm 0.45$ | $52.95\pm 1.49$ | $52.53\pm 1.35$ |  |
| 10.0 | $\mathbf{55.12\pm 1.16}$ | $\mathbf{55.03\pm 1.39}$ | $52.85\pm 0.95$ | $50.82\pm 3.03$ | $52.27\pm 1.11$ |  |
| 100.0 | $54.22\pm 1.74$ | $52.31\pm 2.55$ | $52.10\pm 0.87$ | $49.52\pm 4.34$ | $51.46\pm 2.92$ |  |
| 1000.0 | $53.45\pm 1.65$ | $53.82\pm 2.16$ | $51.16\pm 1.92$ | $51.30\pm 2.02$ | $51.19\pm 1.05$ |  |
| $-$ | $53.69\pm 0.77$ | $51.78\pm 0.44$ | $50.24\pm 1.88$ | $51.12\pm 0.72$ | $49.87\pm 1.23$ |  |

For the CIFAR-MNIST dataset, the preferred values of $\lambda$ that yield optimal outcomes range within $\{2.0,5.0,10.0\}$, and for $\tau$, the values are $\{0.4,0.5\}$. A similar pattern is observed in the {Path-OCT-Tissue}MNIST dataset experiment, with a small difference, in which it shows a preference for fewer warmup rounds ($\tau\in\{0.2,0.4\}$) and demonstrates optimal performance with the same set of $\lambda$ values. This consistency across different datasets indicates robustness in the parameter settings for achieving high accuracy.

### C.4 Larger number of participants

Although our primary focus is on the cross-silo setting, we extend our study to include a large-scale scenario involving 200 participants on the CIFAR-MNIST dataset. We adopt a Dirichlet partition strategy with concentration parameter $\alpha=0.5$ and implement this scenario with a partial participation rate of $0.1$. The outcomes of this experiment, as depicted in Figure 12, indicate a superior performance compared to the conventional FedAvg algorithm, thereby further substantiating the validity and effectiveness of our proposed method. The parameters set for FedPeWS are: $\tau=25$ and $\lambda=0.5$.

Figure: Figure 12: Visualization of global model performance with $N=200$ participants with a partial participation rate of $0.1$. Smoothing is applied as a running average with a window size of 5. A learning rate scheduler is implemented at rounds $200$ and $400$ with a learning rate decay factor of $0.1$.
Refer to caption: x13.png

## Appendix D Neuron Activations

In this section, we examine the extent to which neurons in each layer are activated. Our study uses the Synthetic-32K dataset and the FedPeWS-Fixed method (with parameter $W=50$). The vertical dashed line ($W=50$) indicates the point at which participants switch to using full masks.

Figure: Figure 13: Neuron activation study on the Synthetic-32K dataset with a global learning rate $\eta_{g}=1.0$. The experiment uses the FedPeWS-Fixed method with parameter $W=50$, indicated by the vertical dashed line, marking the switch to full masks by each participant.
Refer to caption: x14.png

Figure 13 displays the neuron activations, measured as the sum of activations over a batch of samples randomly selected from each participant’s dataset, over $T=250$ communication rounds. The top row shows the outcomes for Participant 1, and the bottom row shows the activations for Participant 2. Each column corresponds to different fully connected layers (FC1 to FC4) in the network. Observations are as follows:

- 1.
Before switching $(t\leq W)$: for Participant $i$, subnetwork $i$ shows higher activation patterns in all given FC layers, $i\in[1,2]$, while the other subnetwork exhibits a minimal activation.
- 2.
After switching to full mask $(t>W)$: (i) there is a noticeable increase in activations for both participants upon switching to using full masks, (ii) Participant 1 with its originally initialized subnetwork 1, shows a substantial increase in activations compared to subnetwork 2 across all layers. The same pattern is obesrved for Participant 2 with subnetwork 2.

These findings suggest that the personalized warmup strategy helps the network learn which paths to follow when specific data points are fed into the network. This supports the superiority of our method and corroborates the claims made in the main paper.

## Peer Reviews

confidence: 4
contribution: 1
initial_rating: 3
presentation: 3
questions: Based on my understating, in the first phase each client learns its own subnetwork. Then in the second phase the proposed algorithm employs federated averaging. Can you explain why does the algorithm split the learning rounds into two parts? I am not clear about the intuition behind this and its helpfulness in personalized federated learning.
soundness: 2
strengths: The paper studies the problem of personalized federated learning from learning subnetworks point of view. The presentation of the paper is good.
summary: The paper proposes a personalized federated learning algorithm called FedPeWS. The paper proposes to split the learning rounds into two parts. The first $W$ rounds is called warmup phase where each client learns a personalized subnetwork of the model. The key parts of the proposed algorithm is to identify the subnetwork for each client and determining the warmup length.
weaknesses: - The proposed algorithm is too intuitive. And reading the paper I am not clear why the proposed algorithm can help personalized federated learning. Also reading the introduction, the motivation is not convincing to me.
- Since this work is based on intuitions, I believe the paper should improve the experimental study significantly. Specifically, I believe the paper should add more baselines to the paper with more analytical explanations about the result.

---

confidence: 4
contribution: 2
initial_rating: 3
presentation: 2
questions: 1. Since the research purpose of this paper is to solve the problem of extreme data heterogeneity, more comparisons can be made with some of the current state-of-the-art methods for solving data heterogeneity. I think that besides the FedProx method mentioned in this paper, the MOON [1] method is also a worthy comparison. And I think FedPeWS is essentially an initialization method, so it can also be combined with methods such as PFL [2] in the subsequent standard federated optimization to investigate whether this method, FedPeWS, can bring further enhancement to the existing state-of-the-art methods.
2. In the related experimental descriptions of Table 1 and Table 2, there is no mention of the number setting of clients N. If I understand correctly, the settings on synthetic datasets for both experiments are N=2. This implies that when comparing with FedProx, the comparison has been made only in the scenario of synthetic datasets (N=2). If possible, I'd like to see experimental comparisons with the FedProx method under more datasets (and at the same time larger N value).
3. In this paper, collaboration rounds are used in the experimental part of the main text to compare the convergence speeds of FedAvg and FedPeWS, but since FedPeWS contains a greater computational overhead for one collaboration round during the warm-up phase (but only need to transmit sub-network, which is theoretically less expensive in terms of communication overheads), simply comparing the number of rounds is not convincing enough. Although the experimental analysis for wall-clock time is also presented in the Appendix, it can also be seen that FedPeWS does take more time in the warm-up phase, and thus the balance between the parameter $\tau$ and the actual convergence speed is very important in practical applications. To summarize, I think how to determine the size of the warm-up rounds $W$, or the size of $\tau$, in the experiments is an issue worth discussing, which also involves the analysis of the overall computational and communication overheads. (By the way, if I understand correctly, $\tau =25$ in Appendix C.4 should be changed to $W=25$)

[1] Li, Qinbin, Bingsheng He, and Dawn Song. "Model-contrastive federated learning." Proceedings of the IEEE/CVF conference on computer vision and pattern recognition. 2021.
[2] Arivazhagan, Manoj Ghuhan, et al. "Federated learning with personalization layers." arXiv preprint arXiv:1912.00818 (2019).
soundness: 2
strengths: 1. This paper addresses the problem of data heterogeneity from a novel perspective. 
2. The authors have conducted extensive experiments under a variety of datasets and made a detailed analysis of the ablation experiments around the relevant parameters.
summary: To address the problem of extreme data heterogeneity in federated learning, this paper starts from the perspective of initialization. It argues that an appropriate personalized initial method can enable the clients to quickly learn their local data well before engaging in broader collaboration, and thus achieve faster convergence and a higher final accuracy. In this paper, the authors propose FedPeWS, a personalized warm-up method for the initial phase of federated learning training, which simultaneously optimize the personalized masks of the clients and update the corresponding sub-networks during the initial warm-up phase to enables the global network to adapt faster and better to extreme data heterogeneity scenarios. Experimental results under various scenarios and parameter settings validate the effectiveness of this method.
weaknesses: 1. As mentioned by the authors, this paper lacks proofs related to convergence, and therefore at the same time, it fails to give an analysis of the relationship between hyperparameters such as $\lambda$, $\tau$ and convergence to determine a better range of theoretical values. Therefore for different experimental scenarios, perhaps multiple experiments (e.g. grid search) are required to determine the most suitable parameter values, which has limitations in terms of generalizability and cost.
2. The experimental results show that the effectiveness of FedPeWS and FedPeWS-Fixed methods are comparable, and in some stages FedPeWS-Fixed even outperforms FedPeWS, which I think is an issue that should not be ignored. In conjunction with Appendix B.2, I think FedPeWS-Fixed is a rigid approach to divide the network evenly and fixedly, while the FedPeWS method dynamically optimizes the personalized masks to select the appropriate sub-networks. Intuitively this dynamic optimization method should be more effective than the simple fixed division method, but the experimental results are the opposite, does it mean that the dynamic optimization method of FedPeWS does not achieve much extra gain? I think this is a question worth analyzing.

---

confidence: 5
contribution: 1
initial_rating: 1
presentation: 2
questions: See comments above.
soundness: 1
strengths: The paper studies an important problem which is convergence and performance issues of FL under extreme data heterogeneity.
summary: The paper proposes "FedPeWS," a personalized warmup method for federated learning (FL) using subnetwork masking to improve model convergence under extreme data heterogeneity. It introduces a warmup phase where each participant trains a personalized subnetwork before switching to standard federated optimization, claiming enhanced accuracy and efficiency over baseline methods.
weaknesses: Overall, while this paper attempts to address challenges in federated learning under extreme data heterogeneity, there are several critical issues regarding its hypotheses, methodology, and experimental validation. I hope these comments provide constructive insights for improvement:

* The paper is built upon unproven hypotheses, especially as stated in the abstract (lines 15-17): "we hypothesize that ... rounds." Presenting such assumptions without experimental or theoretical proof weakens the foundation of the study. For any claim regarding the benefits of a personalized warmup phase, empirical or theoretical evidence is essential to establish credibility.
* In lines 44-45, the authors suggest that the primary cause of federated learning failure under extreme data heterogeneity lies in high conflict between local updates during initial rounds. This premise is misleading, as conflicts due to objective misalignment and heterogeneity persist throughout the entire FL process, not just in early rounds. A more nuanced discussion around ongoing client conflicts across rounds would provide a clearer picture of challenges in FL.
* Lines 51-53 mention that each participant uses a personalized binary mask to learn local data distributions and optimize local (sparse) models. However, there is no clear explanation of how the nonlinear parameter space of a model can be mapped directly to each client’s data distribution. To the best of my knowledge, there is no work which has been able to drive a specific subnetwork structure representing the data distribution in NN with non-linear activation.
* The paper’s methodology seems to extend existing techniques without substantial innovation. Additionally, recent work in FL initialization with pre-trained models has shown that warmup rounds may not be necessary, as initialization with pre-trained or meta-learned models can often yield better results. A thorough discussion of and comparison with these approaches would better position the contributions within the current literature. I also urge the authors to enhance their literature review. 
* The experimental setup lacks robust benchmarks, as it does not evaluate the method on deeper architectures or complex datasets like CIFAR-100. Additionally, comparisons with state-of-the-art FL methods (such as those involving advanced initialization or optimization techniques) would strengthen the empirical evaluation and actual benefits of the method. Right now, the actual benefits of method are questionable.

---

confidence: 4
contribution: 2
initial_rating: 3
presentation: 3
questions: See weaknesses
soundness: 2
strengths: - The idea of using independent subnet training in federated learning looks quite novel and reasonable to the reviewer.
- The paper is well written and easy to follow.
summary: This paper develops a warmup algorithm to tackle data heterogeneity in federated learning. The idea is to first allow clients to personalize their models in initial rounds, and then adopt conventional federated approach. Some experiments are provided to showcase the effectiveness of the algorithm. The paper is well written and easy to read.
weaknesses: While the reviewer appreciates the novelty of the approach, they are a set of concerns/comments/questions: 

- While the paper emphasizes that the primary benefit of the algorithm is improved convergence, there is a lack of theoretical analysis. 
- How does the proposed method differ from FedPM (Isik et al., 2023) which also uses mask training? 
- The biggest concern from the reviewer is that the numerical results are on a relatively simple setting:
   -- The datasets are too simple. The authors are suggested to try harder ones, e.g., CIFAR100, (Tiny) Imagenet
   -- The model used are too simple. The authors are suggested to try ResNet and ViT
   -- The comparison to SOTA is missing. Current comparison only involves FedAvg and FedProx and is far from enough. The authors are suggested to MOON, FedOpt, FedUV, etc.
   -- The number of clients is too small. While the authors tried N=200 in the appendix, the improvement is much less obvious than the simple setting in the main text. 
   -- Important ablation study is missing. For example, the authors might try replacing the current warmup strategy by random neuron sampling and other IST approaches. 
   -- Data heterogeneity with \alpha being 0.1 is not extreme. The authors are suggested to try smaller values, e.g., 0.01.