Paper: From Random to Relevant: Harnessing Salient Masks in Non-IID Federated Learning (ICLR-2024)

Reviewer comment (Theory):
The contribution of this paper is trivial. Although the proposed method seems efficient because the mask is found in the beginning, I don't think the proposed approach makes sense. Before the model training, the parameters $w_0$ are generated at random. Denote the gradient by $g_0$ for the initial model. Intuitively, when the model converges, the important parameters are with relatively large values, while the ignorable ones are close to 0. As the mask is merely generated in the beginning, I cannot see why Eq. (3) can find a reasonable mask. In other words, I cannot see the differences when the mask is generated arbitrarily.

Author response:
We thank the reviewer for his review and helpful suggestions. We appreciate the recognition of our approach's efficiency and the thoroughness of our literature review by the reviewer. Regarding the validity of our masking technique, we believe there might have been a miscommunication in our initial manuscript. The criterion for our masking process is firmly established in theory and on a comprehensive body of literature. 

To address this gap, we have expanded our explanation and provided additional evidence to substantiate the robustness and relevance of our method in the revised manuscript (Section 3.1 and Appendix A.1). This should clarify any uncertainties and demonstrate the validity of our approach clearly. We also address all the points brought up by the reviewer and make necessary changes in the updated manuscript (in purple).

### 1. Regarding the efficacy of the Masking criterion
> 1. "the proposed method seems efficient because the mask is found in the beginning, I don't think the proposed approach makes sense. Before the model training, the parameters are generated at random. .... In other words, I cannot see the differences when the mask is generated arbitrarily."

We acknowledge that our initial reliance on references for presenting the saliency criterion might have led to the reviewer's conclusions. To address this, we have thoroughly revised Section 3.1 of our manuscript, offering a comprehensive explanation of the saliency criterion, which is well-supported by extensive literature [1,2,3,4,5]. Furthermore, we present a detailed exposition of the process in Appendix A1. 

These revisions aim to clarify the mask generation process, explaining its methodical nature and demonstrating its effectiveness over random masking. Additionally, to substantiate our claims, we have included experiments comparing our method with random masking in the latter part of this response. Below we briefly provide an explanation of the masking process based on [1,2,3,4,5] here:

### Explanation of the validity of Connection Importance Criterion and Mask Generation
Since the goal is to measure the importance or saliency of each connection an auxiliary indicator variable $\mathbf{c} \in \\{ 0,1 \\}^d$ is introduced representing the degree of importance of parameter $\mathbf{w}$. Given a sparsity level $\lambda$, the objective can be re-written as follows:

\begin{align} 
    \mathcal{L} &= \frac{1}{M} \sum_i \ell ( \mathbf{c} \odot \mathbf{w}; \mathbf{x}_i, y_i) \quad \text{s.t.} \quad \mathbf{w} \in \mathbb{R}^d, \quad \mathbf{c} \in \\{ 0,1 \\}^d, \quad ||c||_0 \leq \lambda \notag
\end{align} 

The main goal of the above formulation is to use the $j$-th element $c_j$ of $\mathbf{c}$ as a metric for the importance of each connection $w_j$, where $w_j$ indicates the $j$-th connection, which enables us to determine the saliency of each connection by measuring its effect on the loss function.

The core idea of the importance criteria is to preserve the parameters that have maximum impact on the loss under perturbation. The importance criterion is then formulated as:
$$
    \mathbf{g}_j(\mathbf{w}; \mathcal{D}) 
    = \frac{\partial \mathcal{L} (\mathbf{w} \odot \mathbf{c})}{\partial \mathbf{c}}  \\
    = \frac{\partial \mathcal{L} (\mathbf{u})}{\partial \mathbf{u}} \cdot \frac{\partial \mathbf{u}}{\partial \mathbf{c}} \\
    = \frac{\partial \mathcal{L} (\mathbf{w})}{\partial \mathbf{w}} \odot \mathbf{w}
$$
Where we evaluate the above at $\mathbf{c}=1$ and use a substitution of variable $\mathbf{u} = \mathbf{w} \odot \mathbf{c}$ which implies, $\partial \mathbf{u}/ \partial \mathbf{c} = \mathbf{w}$ and $\frac{\partial \mathcal{L} (\mathbf{u})}{\partial \mathbf{u}} = \frac{\partial \mathcal{L} (\mathbf{w})}{\partial \mathbf{w}}$ at $\mathbf{c = 1}$ and relax the binary constraint on the indicator variables $\mathbf{c}$ to take on real values $\mathbf{c} \in [0,1]^d$. To achieve our primary objective of identifying significant connections towards the start of training, we utilize the magnitude of the gradient $|\mathbf{g}_j(\mathbf{w}; \mathcal{D})|$ as an indicator of connection importance (note that this is the gradient with respect to $\mathbf{c}$ and not $\mathbf{w}$). A higher value of this magnitude suggests that the corresponding connection $w_j$ associated with $c_j$ has a substantial influence on the loss function, implying its significance in the training process and that the connection should be preserved. Consequently, a saliency score for each connection can be calculated based on this principle as follows:

\begin{align}
    s_j = \big| g_j(\mathbf{w}; \mathcal{D}) \big| = \bigg| \frac{\partial \mathcal{L} (\mathbf{w})}{\partial \mathbf{w}} \odot \mathbf{w} \bigg|
\end{align}

We next extend this saliency criterion in section (3.1) in paper based on the data distribution at each site and generate a mask choosing the top-$\lambda$ ranked connections as: $\mathbf{m} = \mathbb{1}[s(w_j) \geq s(w_\lambda)]$

Reviewer comment (Novelty):
In addition to the mask initialization, the rest of the design is consistent with FedAvg. Following the first point, I cannot see the significance of this work.

Author response:
### 2. Significance of this work
> 2. In addition to the mask initialization, the rest of the design is consistent with FedAvg. Following the first point, I cannot see the significance of this work.
> 
In light of the above explanation on the validity of our mask generation process and its grounding in extensive literature we note the significance of our work here. 
- From our reading of the literature, the current leading works in terms of sparsity vs accuracy trade-offs, and hence communication efficiency in the realistic non-IID setting is DisPFL and FedPM [6, 7]. As FedPM trains a mask on the initial weights and never trains the weights, it loses accuracy in the process, similar to the works with supermask [10]. 
- DisPFL on the other hand finds a *random mask* at initialization and periodically updating the mask trains the model in a manner similar to FedAvg as well. One of the major goals of our work is to demonstrate that, in contrast to DisPFL's random initial masking and updating, finding a salient mask at initialization results in a **significantly** better sparsity vs accuracy trade-offs in the non-IID setting. 
- We also note that previous work [8,9] reported unsuccessful uses of the saliency criterion in the FL setting, whereas our proposed extensions for the saliency criterion in this work to the non-IID setting effectively helps the method to outperform competing similar methods.

Reviewer comment (Theory):
The authors mention the work based on decentralized FL. After reading the paper, I don't have an idea why it can work under peer-to-peer network architecture. Instead, it is solely workable with the client-server settings.

Author response:
### 3. SSFL in the Decentralized setting:
> 3. The authors mention the work based on decentralized FL. After reading the paper, I don't have an idea why it can work under peer-to-peer network architecture. Instead, it is solely workable with the client-server settings.

We appreciate the reviewer's interest in the applicability of our method in a decentralized setting. Since the reviewer does not mention *why this method would not work in a decentralized setting*, we provide a general explanation. In this work we do not assume the existence of a server, however, our proposed technique should be conceptually viable for both a server-client setup and a decentralized setup. We only require most client participation during initial weights alignment and transmitting and receiving saliency scores and sample count. These operations are required only once, at the start of training, for alignment purposes and would be pulled/gathered from all the clients at each client $k$. We make changes to clarify the exposition in Section 3.2.

Moreover, focusing on line-12 of Algorithm-1, we note that each client gathers the current masked weights from a subset of all other clients and trains the model. This is true for all the $K'$ selected clients, where they all gather weights among themselves. 

The work however, does not focus on the engineering details of deploying this in a real-world decentralized setting or go into exhaustive details as that is beyond the scope of our work and are deployment details. However, taking into reviewer's concern we have added modifications in the algorithm and updated the methods section 3.2 to help clarify. We are however happy to discuss if the reviewer has any specific concerns regarding SSFL's viability in the decentralized setting.

Reviewer comment (Experiments):
The experiments only show the result on a sparsity level of 50\%, which is not convincing. I would like to see the performance results when the sparsity level is at 10\% or smaller. According to Figure 4, I notice DisPFL and Ditto are yet to converge. I wonder about their final results when they converge.

Author response:
#### 4. Further Clarification on Initially Provided Performance Metrics Across Different Sparsity Levels 
> 4. The experiments only show the result on a sparsity level of 50%, which is not convincing. I would like to see the performance results when the sparsity level is at 10% or smaller.

We would like to point out that we have reported the performance of SSFL on varying sparsity levels of 50, 70, 80, 90 and 95 percent in Figure 2 of the initial paper submission and compared with DisPFL at the same sparsity levels as well as with the dense baseline FedAvg-FT. Thus, the results do include $<10\%$ density levels as asked by the reviewer. 
However, following reviewer's sugestion, we now also include results from two additional methods: SubFedAvg and FedPM. Please note that FedPM, a more recent method, automatically chooses the most optimal sparsity and does not allow sparsity sweeps. We present the results below:

| Methods               | Sps 0% | 50% | 70% | 80% | 90% | 95% |
|----------|----------|----------|----------|----------|----------|----------|
| **FedAvg-FT** (dense) | 88.02 | | | | | |
| **SSFL**              | | **88.29** | **87.57** | **87.43** | **84.58** | **84.18** |
| **DisPFL**            | | 85.82 | 83.39 | 82.93 | 76.17 | 75.00 |
| **SubFedAvg**         | | 0.7970 | 0.8489  | 0.7818 | 0.7813 | 76.77 |
| **FedPM**             | (optimal sps) 57.30 | N/A | N/A | N/A | N/A | N/A |

**Table:** *Accuracy at given sparsity levels. We note that FedPM optimally chooses its sparsity level and does not allow sparsity sweeps hence other sparsity levels are not available N/A.*

Reviewer comment (Experiments):
The authors should conduct an empirical study to show the performance differences when the mask is randomly drawn in the beginning.

Author response:
### SSFL masking compared with random masking
We appreciate the reviewer's suggestion about performing an empirical study comparing with random masking and acknowledge that this is indeed an interesting study to conduct. Taking into the reviewer's comment we report the comparison with random masking below on the Cifar10 non-IID split data:

| Methods              | 50% | 70% | 80% | 90% | 95% |
|----------|----------|----------|----------|----------|----------|
| **SSFL** | **88.29** | **87.57** | **87.43** | **84.58** | 84.18 |
| **Random** | 41.61 | 39.80 | 19.70 | 10.60 | - |

Since the performance of random masking falls around $10$% at $90$% sparsity, which is the accuracy of a random model on Cifar10, we do not conduct further experiments on higher sparsities during the discussion phase.

Reviewer comment (Theory):
While the approach is indeed innovative, it appears to be overly simplistic.

Author response:
1. While we respect the viewpoint that our approach may appear overly simplistic, we would like to pose an alternative perspective to consider the elegance of simplicity whilst maintaining effectiveness. The approach introduced by us is indeed intuitive, but it robustly builds upon the existing body of work in this field, extending the sparsity vs accuracy trade-off in the non-IID Federated Learning training regime. We would like to pinpoint that previous contributions tackling comparable problems [1] also hinge on the simplicity of initiating from a random mask and allow clients to update the mask using only the gradient as the update signal. Thus, we would not like to view simplicity as a quality that undermines the approach's value; instead, we view as an inherent strength of the design.
2. Furthermore, we would like to emphasize that our mask generation process, while designed in elegant simplicity, does contain nuanced elements in its execution. To provide a deeper understanding, we have expanded upon the process in section 3.1 of our manuscript, spotlighting these fine details. By laying bare the intricate facets of our seemingly simple approach, we hope to convince the reviewer of the depth of our work beneath its surface.

Reviewer comment (Experiments):
The experimental results are promising, but there is potential for improvement, particularly when tested on alternative models and datasets.

Author response:
We thank the reviewer for this suggestion.  To ensure a comprehensive
evaluation of our method in the initial submission, we have already tested on three distinct datasets, including the larger Tiny-Imagenet and have utilised modern Resnet models, distinctively different from what had been used in similar studies like FedPM [2]. This rigorous approach ensures the broad applicability of our method.

Responding to the invaluable feedback, we have taken one step further. We conducted additional experiments and have enriched our paper with new empirical results comparing our method with other established methods. These further reinforce the robust performance and major improvements our method exhibits across diverse conditions. Please refer to our detailed response in the opening common response.

We believe the additional experiments improved the part of the work that stands up under various testing conditions and comparison with diverse alternative models. We believe that these enhancements will significantly contribute to the reliability and novelty of our paper.

Reviewer comment (Experiments):
Can this approach be applied to Large Language Models (LLMs)? If so, incorporating empirical results would significantly enhance the paper.

Author response:
We certainly appreciate your suggestion about applying this approach to large language models. The idea is intriguing and could open up several interesting research areas. However, given the complexity involved in this process, it would necessitate extensive experimentation and research beyond the scope of our current work. It's an appealing future research direction, and we have added it to our discussion for potential future work in our revised manuscript. We hope that clarifying this in the paper helps make the potential of this study more clear.

Reviewer comment (Novelty):
Is your proposed method connected to the concept of core sets?

Author response:
We recognize that while the concepts are related, they are not identical as per our interpretation of this work. Core sets were originally developed as a technique to approximate a large set of data points with a smaller yet representative subset, preserving desired criteria.  This concept found its early applications in problems related to k-means and k-medians, where the goal was to represent a large dataset with a smaller, yet representative, subset of data points. In the field of sparsity and pruning, this idea was leveraged for pruning and compression in [3], where CNN model filters were grouped by their activation patterns and represented through their coreset characteristics followed by pruning.
In contrast, our work can be thought of as reliably finding axial subspaces within the model parameter space that encapsulates a good final solution and we don't believe this closely relates with the concept of core sets from our understanding.

Reviewer comment (Evaluation):
While the proposed method, SSFL, outperforms the selected baselines, overall the performance degrades considerably even for 50% sparsity. For standard gradient sparsification methods such as Top-k[1], Random-k[2], or rTop-k[3], there is no significant performance loss up to 90% sparsity with error feedback. I wonder if the authors have any explanation as to why SSFL and the selected baselines degrade the accuracy even with small sparsity ratios like 50%.

Author response:
We first note that the works [1,3] are in the distributed SGD and are not what we consider to be in FL setting, and while the work [3] is conducted in the FL setting, it works under the IID assumption for data distribution. We  summarize the biggest differences between the works mentioned by the reviewers and our work and the baselines below:

1. **Distributed SGD vs Federated Learning with real world assumptions:**
    - We would like to point out that some of the studies mentioned by the reviewer, specifically, works [1,3] focus on *gradient sparsification* techniques in a distributed, non-federated setting and are tested under IID assumptions, leading to divergent results. As a result, the sparsity vs accuracy trade-offs in those settings are close to standard multi-node training within a HPC cluster. 
    - In contrast, our method and the baselines that we compare against all use standard Federated Training technique with communications made only after $T$ steps of local training on local data. This arrangement makes the training task more challenging, simulates real-world scenarios more accurately, and substantially reduces communication requirements, leading to differing results. In essence, FedAvg-FT serves as the standard baseline for the non-IID setting we select.
     - We also note that SSFL outperforms our standard *dense* baseline FedAvg-FT (and FedAvg) in the non-IID setup at the lower sparsity level of $50\%$. This phenomena agrees with those reported in single node sparse training, where sparse models in lower sparsity range outperformed dense baselines [9].
    
2. **SSFL and selected baselines tested on the harder non-IID Assumption:** 
    - An even more significant difference is that all the methods [1,2,3] use the general CIFAR10 dataset assuming IID condition, while our work and all the baselines we report were tested on the CIFAR10 non-IID split dataset where the class allocation is done according to the Dirichlet distribution with $Dir(\alpha)$ [5]. Moreover, they also do not prune the whole network and leave the first and last layer intact heuristically.
    - Note, that this is a considerably harder task compared to the standard Cifar10 task by itself without sparsity considerations. As a further validation of this argument, we note that the last reference shared by the reviewer FedPM [4] which uses the non-IID Cifar10 split also report a accuracy drop in general and compared to us, a lower accuracy as well (57.30\% at optimal sparsity vs our 84.18% at 95% sparsity). Although, we believe there are more reasons for this drop in accuracy in FedPM that we explain later.

Reviewer comment (Experiments):
Is this because only a fixed sparse portion of the model is trained while the rest is kept at zero? If so, would it be possible to let each client find their own mask rather than aggregating them to have one global mask for all the clients. This way, different portions of the model could be trained by different clients at the cost of losing the communication gain from server to client communication -- which is typically not the main bottleneck compared to client to server communication especially given that the current approach results in considerable accuracy loss.

Author response:
- We thank the reviewer for this insightful suggestion. This is indeed a very interesting idea and is a natural progression from our work However, we note that we briefly explored local masking and training but it did not result in competitive performance ($70.22$% accuracy at $50$% sparsity level on non-IID cifar10). We hypothesize that this might be due to the complex way in how completely non-overlapping masks and weights interacted and updated with each other after aggregation along with random sampling of clients. It might be the case that the local models would drift far away in the restricted (sparse) solution space and their aggregation resulting in subpar models. This would indeed be a an interesting avenue for future research.

Reviewer comment (Experiments):
I also think comparing SSFL to these alternative sparse methods (Top-k[1], Random-k[2], or rTop-k[3]) would provide a more complete picture to the readers. Also, other methods for finding optimal global sparsity masks should be added as baselines as well such as FedPM [4].

Author response:
After noting the differences between the works suggested by the reviewer and SSFL and the baselines, taking into the reviewer's recommendation we conducted experiments with top-k, random-k and also compare FedPM (which used non-IID) We present the results from our experiments with FedPM, top-k and random masking. We note:
- We report comparison with FedPM on Resnet-18 trained on the Cifar10 non-IID dataset with $Dir(\alpha = 0.3)$.
- FedPM does not train the model weights and instead tries to find a super-mask and optimally determines its sparsity. This is effectively in line with works on super-masks [6]. Moreover, they use a heuristic for splitting the data in a non-IID manner.
- The top-k presented in the recommended works was implemented in relation to gradient sparsification in DSGD setting. For our training paradigm where the sparse weights are shared every $T$ local training steps, for the top-k experiment we resorted to choosing the top-k weights at the end of each communication round to share among the clients and aggregated them.
- We also note that at the end of training, SSFL and DisPFL ends up with sparse local models with potential for faster inference in contrast to top-k or random communication where the models remain dense.


    | Methods |  sparsity 50% | sparsity 70%  | sparsity 80% | sparsity 90% |
    |--------|--------|--------|--------|--------|
    | **SSFL**  | **88.29** | **87.57** | **87.43** | **84.58** |
    | **Random mask** | 41.61 | 39.80 | 19.70 | 10.60 |
    | **Top-k**  | 53.33 | 52.33 | - | - |
    | **FedPM**   | 57.30 (optimal sparsity) | N/A | N/A | N/A |
    
    **Table:** *Accuracy at given sparsity levels. We note that FedPM optimally chooses its sparsity level and does not allow sparsity sweeps hence other sparsity levels are not available N/A. At this stage we present Top-k for 50% and 70% as the runs are computationally demanding and are still computing, however we do not expect higher accuracy at higher sparsity levels*.