# FedMap: Iterative Magnitude-Based Pruning for Communication-Efficient
Federated Learning

Alexander Herzog12 *  — Robbie Southam2 *  — Ioannis Mavromatis2 *  — and
Aftab Khan2 *  —

Email<EMAIL_ADDRESS>
2Toshiba Europe Ltd., Bristol Research & Innovation Laboratory, UK

###### Abstract

Federated Learning (FL) is a distributed machine learning approach that
enables training on decentralized data while preserving privacy. However, FL
systems often involve resource-constrained client devices with limited
computational power, memory, storage, and bandwidth. This paper introduces
FedMap, a novel method that aims to enhance the communication efficiency of FL
deployments by collaboratively learning an increasingly sparse global model
through iterative, unstructured pruning. Importantly, FedMap trains a global
model from scratch, unlike other methods reported in the literature, making it
ideal for privacy-critical use cases such as in the medical and finance
domains, where suitable pre-training data is often limited. FedMap adapts
iterative magnitude-based pruning to the FL setting, ensuring all clients
prune and refine the same subset of the global model parameters, therefore
gradually reducing the global model size and communication overhead. The
iterative nature of FedMap, forming subsequent models as subsets of
predecessors, avoids parameter reactivation issues seen in prior work,
resulting in stable performance. In this paper we provide an extensive
evaluation of FedMap across diverse settings, datasets, model architectures,
and hyperparameters, assessing performance in both IID and non-IID
environments. Comparative analysis against the baseline approach demonstrates
FedMap’s ability to achieve more stable client model performance. For IID
scenarios, FedMap achieves over $90$% pruning without significant performance
degradation. In non-IID settings, it achieves at least $~{}80$% pruning while
maintaining accuracy. FedMap offers a promising solution to alleviate
communication bottlenecks in FL systems while retaining model accuracy.

###### Index Terms:

Federated Learning, Deep Learning, Internet of Things, Communication
Efficiency, Pruning

##  I Introduction

Deep Neural Networks (DNNs) often require vast amounts of training data,
particularly for intricate tasks such as image classification and language
modelling. In traditional, _centralised_ , Machine Learning (ML) approaches,
data is typically brought together at a central location for training. This
approach, however, encounters substantial challenges, particularly concerning
data privacy regulations and GDPR issues (with associated fines for non-
compliance) as well as the trustworthiness of the AI models [1].

Federated Learning (FL) [2] was proposed to address some of these challenges.
It offers a framework for distributed training, primarily focusing on neural
networks. Distinct from centralized models, FL ensures that raw data remains
exclusively with the clients and is never transferred to the central server.
In an archetypal FL system, data is harvested by client devices situated at
the network’s edge. The training regimen comprises local model updates using
each client’s distinct data and subsequently fusing all client models, usually
via a central server or parameter server (PS). By this mechanism, not only is
data privacy safeguarded – a vital aspect in settings with privacy
implications [3] – but the model also offers an efficient solution to
potential bandwidth constraints prevalent in large-scale use-cases, such as an
Internet of Things (IoT) setting. Transmitting model parameters instead of raw
data can significantly curtail data communication requirements [4].

However, it’s crucial to note that client devices in an FL system frequently
face resource limitations. Compared to the high-performance infrastructure in
data centres, these client devices often exhibit reduced computational
capabilities [5], limited memory [6, 7], smaller storage capacities [8], and
narrower communication bandwidth [9, 10, 11, 12]. Our work particularly
focuses on reducing communication bandwidth by collaboratively learning an
increasingly sparse global model. To achieve this, we strictly adapt the
iterative, magnitude-based pruning regime [13] and apply it to the FL
mechanism. Our method, named FedMap, ensures that all clients prune and refine
using the same pruning mask, representing the same subset of the global model.
As a result, once the client models are sufficiently refined, we proceed with
further pruning to remove even more parameters.

Unlike prior approaches reported in the literature (e.g., PruneFL [14] and
FederatedPruning [15]), a key advantage of FedMap is its ability to train a
global model from scratch, making it well-suited for use cases with heightened
privacy requirements. In the medical domain, increased sensitivity of
healthcare data and stricter regulations make it difficult to leverage pre-
trained models [16], underscoring the value of our proposed method aiming to
optimally train the global model without centralized data access. Similarly,
in the finance sector, financial institutions are often – and rightly so –
reluctant to share customer data due to privacy concerns, creating barriers to
developing AI models that rely on pre-trained weights [17]. These challenges
in such domains, where suitable pre-training data is often limited, position
FedMap as a particularly advantageous solution compared to other methods
reported in the literature.

Broadly, pruning techniques can be categorised into structured and
unstructured [18]. While structured pruning [19] eliminates entire units like
neurons or channels, unstructured pruning removes individual weights from the
model – FedMap is an unstructured pruning solution. As shown in the literature
[20, 21], unstructured pruning has proven to achieve higher pruning ratios
compared to structured pruning making it more suitable for communication-
constrained FL deployments. Moreover, as all clients prune the same model and
follow the same hyperparameters, masking information does not need to be
exchanged. That means that the PS can reliably reconstruct the models and more
importantly, that the models exchanged get ever-smaller after each round.

In this paper, we comprehensively explore FedMap’s performance across a varied
set of settings and pruning schedules, emphasising the versatility of our
approach. Moreover, we compare FedMap with FederatedPruning [15]. A
distinguishing feature of FedMap is its iterative approach, where subsequent
models are consistently formed as subsets of their predecessors. This design
precludes the parameter ‘reactivation’ issue observed in FederatedPruning;
aiming for a more stable performance, particularly at higher pruning regimes.
Through our experimental study, we show that FedMap contributes to a more
stable client model performance, delivering superior results in both IID and
non-IID settings.

Our main contributions in this paper are summarised below:

  * •

We propose a new method (FedMap) targeting the communication efficiency of FL
deployments

  * •

We extensively analyse our approach through a rigorous evaluation, with
diverse settings (including different datasets, model architectures,
hyperparameters), as well as evaluation under IID and non-IID scenarios for a
realistic performance assessment.

  * •

We also benchmark FedMap against the FederatedPruning [15] approach and
successfully demonstrate the effectiveness of our approach.

This paper is organized as follows: Section II covers preliminaries and
related work, including Federated Learning, Magnitude-Based Pruning, and
relevant literature. Section III introduces our method, FedMap, and its
extension for heterogeneous data distributions. Section IV details our
experiments, including various pruning schedules, local training influences,
non-IID data effects, and benchmarking. Section V discusses the results,
focusing on model performance, stability, and the integration of FedDR for the
deployment of FedMap to non-IID settings. Section VI concludes with a summary
of findings, limitations, and future research directions. Additional results
are provided in the appendices.

##  II Preliminaries and Related Work

Before further explaining FedMap’s functionality, we present integral details
of FL and magnitude-based pruning, setting a solid foundation for further
analysis. We also briefly compare FedMap with other similar solutions found in
the literature, paving the way for our experimental comparison in the
following sections.

###  II-A Federated Learning

Federated learning (FL) is a collaborative, privacy-preserving method for
training machine learning models in a distributed manner. During each
iteration of training/FL round $t\in\\{1,...,T\\}$ , a client
$n\in\\{1,...,N\\}$, with $T,N\in\mathbb{N}^{*}$ utilise a subset of their
local data $\mathcal{D}_{n,t}$, i.e., ${B_{n,t}\subset\mathcal{D}_{n,t}}$
(where ${B_{n,t}}$ denotes a batch of training examples), to calculate
gradients with respect to the global model’s parameters,
$\theta_{t}\in\mathbb{R}^{d}$:

$\nabla_{n,t,B_{n}}=\frac{\partial\mathcal{L}(\theta_{t},B_{n,t})}{\partial\theta_{t}},$
(1)

where $\mathcal{L}$ denotes the loss function. The choice of consensus
mechanism within FL influences whether gradients revert to the PS for
aggregation, e.g., Federated Stochastic Gradient Descent (FedSGD) method of
aggregation [2] or Federated Averaging [22].

Federated Averaging (FedAvg) [22], one of the most prevalent aggregation
approaches, employs locally computed gradients across multiple batches
contained in the client $n$’s training set at FL iteration $t$;
${\mathcal{D}_{n,t}=\\{B_{n,t}^{1},...,B_{n,t}^{\tau}\\}}$ with
$\tau=|D_{n,t}|$ and ${\mathcal{D}_{n,t}\subseteq\mathcal{D}_{n}}$.

Subsequently, the deviation between the updated local model and the global
model parameters, defined as $\Delta\theta_{n,t}=\theta_{t}-\theta_{n,t}$ is
communicated to the PS. By employing multiple batches and facilitating local
updates, FedAvg optimizes communication rounds between the PS and clients,
effectively reducing the communication frequency. The strategy the PS adopts
for aggregation determines the formation of the subsequent global model. In
the process of FedAvg, the changes in model parameters from multiple client
models are aggregated in an element-wise manner to compute a global model
update:

$\theta_{t+1}=\theta_{t}+\frac{1}{N}\cdot\sum^{N}_{n=1}\Delta{\theta}_{n,t},$
(2)

It is crucial to recognise that variations in data collection across clients
and temporal shifts as during FL training rounds may impact model performance.
In cases of independently and identically distributed client data (IID), the
model parameters exhibit significant adjustments in the initial FL rounds,
which progressively decrease. However, in practical scenarios, data leans
toward non-IID [23], implying various data-distribution shifts. In this work,
we consider both IID and non-IID scenarios, evaluating our proposed FedMap
methodology considering feature-, label- and sample-heterogeneity. See [24]
for a more detailed discussion of statistical heterogeneity in FL settings.

###  II-B Magnitude-Based Pruning

Achieving a high level of compression for neural networks can be accomplished
by eliminating parameters that have a minimal effect on model performance. In
unstructured pruning, the deletion of parameters is achieved by setting the
values of the parameters and the gradients during back-propagation to $0$
[18]. The importance of parameters can be assessed according to the impact of
removing a parameter on validation loss
$\mathcal{L}(\theta_{t},\mathcal{D}_{Val})$. Parameters whose elimination
incurs small perturbations to the validation loss are prioritised for removal.
Although there is a rich body of work on pruning at initialisation [25, 26,
27, 28, 29], iterative pruning has a proven track record of learning highly
sparse neural networks while maintaining a performance comparable to its
densely parameterised counterparts. The iterative pruning routine was first
described as early as 1989 in the work of LeCun et al. [13] proposing Optimal
Brain Damage (OBD).

Apart from the amendments (italic font) as result of later findings [30, 31,
20] iterative pruning in its current form has largely been described in the
OBD procedure:

  1. 1.

Choose a suitable network architecture

  2. 2.

Train the network until a reasonable solution is obtained.

  3. 3.

Determine the layerwise sparsity levels $p_{i}$ according to a pre-selected
pruning paradigm.

  4. 4.

Sort the parameters according to their magnitude and in every layer $i$ delete
the smallest $1-p_{i}$ (according to 3) parameters by magnitude.

  5. 5.

Iterate to step 2

Iterative magnitude-based pruning with layerwise pruning ratios is described
in Algorithm 1.

Algorithm 1 Iterative Layerwise Magnitude-Based Pruning

1:The untrained neural network, randomly initialised with parameters $\theta$;
The global pruning rate $p_{G}$; The number of iterations $T$; The number of
fine-tuning iterations $\mathbf{s}$

2:for $t=1$ to $T$ do

3: Train the network for $L$ epochs.

4: if t % $\mathbf{s}$ == 0 then

5: Compute the layerwise pruning ratios $p_{i}$ according

6: to layerwise, magnitude-based pruning scheme,

7: s.t. $\sum_{i=1}^{L}p_{i}|W^{i}|=p_{G}\cdot d,\quad\forall W^{i}\in\theta$.

8: Sort by the magnitude and set all but the top

9: $(1-p_{i})$ % of $W_{i}$ to 0. return The pruned network $\tilde{\theta}$.

It is important to note that the selection of parameters $s$ and $L$
significantly affects the model’s convergence behaviour, especially during the
first iteration of the algorithm. Therefore, careful selection of these
parameters is important to ensure that the model is trained towards a good
initial solution (good validation set performance) before proceeding with
pruning.

Since the inception of iterative pruning in [13], many methods have been
proposed to determine the importance of parameters. More recent work hints at
the effectiveness of simple magnitude-based pruning (MaP) [32], where the
magnitude of the parameters is used as a proxy to determine importance. The
removal of parameters smallest in magnitude is assumed to have the smallest
impact on $\mathcal{L}(\theta_{t},\mathcal{D}_{Test})$.

Using parameter magnitude as a proxy for importance can be explained by
viewing MaP as relaxed layerwise $l_{2}$ distortion minimization [20]. Without
considering the activation function $\sigma$, NNs can be viewed as linear
operators in layers $W^{i}\in\mathbb{R}^{m\times n\times[p]\times[q]}$, where
$W^{i}$ is the weight tensor of layer $i$, acting as an operator on input
vector $x$. Given $W^{i}\in\mathbb{R}^{m\times n}$ and $x\in\mathbb{R}^{n}$,
the induced matrix $l$ norm measures how much $W^{i}$ affects the length of
$x$, denoted as:

$||W^{i}||_{l}=\sup_{x\neq 0}{\frac{||W^{i}x||_{l}}{||x||_{l}}}.$ (3)

Viewing neural network layers as nested linear operators have been the basis
of recent pruning methodologies such as in the layer-adaptive magnitude-based
pruning (LAMP) scheme [20]. The LAMP score is a re-scaled version of the
magnitude that incorporates the model-level $l_{2}$-distortion incurred by
pruning. Given a global pruning target ratio $p_{G}$, the LAMP score
determines layerwise pruning levels, $p_{i}$, while the magnitude of the
parameter determines the inclusion or exclusion in the set of prunables.
Compared to other common selection schemes for MP such as the Erdős–Rényi-
Kernel methodology [31, 33], layerwise- ($p_{i}=p_{G}$) or global (select
$p_{G}$ % of parameters uniformly across all $L$ layers), the LAMP score marks
the state of the art for iterative magnitude-based pruning.

###  II-C Related Work

Communication Efficiency: Efficient communication is a significant bottleneck
in FL. Standard solutions to prevent excessive communication overheads include
reducing the update frequency between the server and the clients or size of
the model updates [4]. Many methods exist that precisely target communication
efficiency, ranging from quantisation to standard compression techniques (e.g.
finite source coding) or a combination of both. Traditional sparsification
methods such as Top-K sparsification-based methods achieve excellent
performance even under strict compression regimes. Additionally, their
convergence has theoretically been proven [34], especially in combination with
error-accumulation [35]. In Top-K, parameters are selected via the magnitude
of the difference between the global- and the local client model. The selected
value-index pairs are then sent to the server. Typically, the parameter’s
floating value is compressed by quantisation, and the integer encoding of the
index can be compressed via Golomb coding [36]. There is a rich body of work
on reducing or eliminating the need to send positional information. The
SmartIdx method by Wu et al. [37], is one example. The authors propose a
structured pruning scheme, where only whole structures, such as convolutional
filters and their module index, are sent, reducing the sending indices from
1:1 to n:1, where $n$ is the number of parameters in the substructure. Another
prominent example is the TCS method by Ozfatura et al. [38]. The authors
extend the canonical Top-K method with error accumulation and reduce the
number of non-zero entries in the binary index mask by reusing the previously
obtained global mask to select the client parameters plus a small percentage
of additional parameters.

Sparsification of Neural Networks: Sparsification or pruning (i.e. the removal
of connections by setting the respective weights to zero) is a broad term to
describe the process of finding subnetworks with similar- or (in some cases)
better performance than their dense counterparts. The most common are dense-
to-sparse methodologies, where a dense network is gradually pruned throughout
the training process. Recent discoveries on neural network pruning reveal that
with carefully chosen layerwise sparsity, a simple magnitude-based can achieve
a state-of-the-art tradeoff between sparsity and performance. Many methods
resort to handcrafted heuristics [31], [33] such as keeping parts of the model
dense, including the first or last layers [39], [40], [41]. Kusupati et al.
[42] provides an overview of the different pruning methods.

PruneFL [14] method of pruning under a FL framework outlines a two-stage
pruning process. In the initial stage, a powerful and trusted client pretrains
the model while also adaptively pruning the model until the model size
stabilises. This approach ensures a reduced computation and communication time
for each federated learning round, especially in heterogeneous setups.
However, the initial pruning might not be the most efficient since it draws
from the data of just one client. To circumvent this, further adaptive pruning
is introduced during FL, employing the standard FedAvg procedure and engaging
data from all participating clients. Their adaptive pruning involves removing
and adding parameters, which the authors call ‘reconfiguration’. Notably, this
reconfiguration transpires during multiple FL rounds at the selected client or
the server post-client updates. A significant element of their method is the
communication benefit, where bitmaps efficiently communicate client pruning
details in the uplink. In contrast, our proposed FedMap method aligns more
with Iterative Magnitude-Based Pruning. Instead of following a pre-training
(and simultaneous pruning) procedure as in PruneFL, we begin with an untrained
model, aiming to collaboratively learning a sparse global model. We emphasise
client-side pruning, avoiding mask transmissions and ensuring predictable
communication overheads. We introduce flexibility in pruning schedules without
broadcasting bitmaps during the uplink.

Similarlily, Federated Pruning [15] begins each round with the server crafting
a set of variable masks and dispatching the pruned model to clients. This
strategy shares a common thread – pruning at the round’s onset. However, while
the pruning masks are crafted server-side, in our proposed setup, clients
prune models on their end, eliminating the need for transmitting a sparse
model. To ensure reproducibility, our clients and the Parameter Server (PS)
initialise their models using the same random seed. While Federated Pruning
resorts to structured pruning, our approach, reminiscent of PruneFL, opts for
unstructured pruning.

In summary, our proposed approach, FedMap, offers several innovative
advantages over existing federated learning and pruning methods like PruneFL
and Federated Pruning. It features an analogous implementation to iterative
magnitude-based pruning, collaboratively learning a sparse global model from
dense client models without pretraining. Critically, it involves clients in
the pruning process before training, generating model masks locally to negate
transmission of masking information or sparse models, ensuring predictable
communication costs. Unlike prior work that deflates models before
aggregation, FedMap performs aggregation directly on dense client updates when
following the same pruning schedule. It offers flexible, client-driven pruning
schedules without incurring overhead from bitmap transmissions.

##  III Iterative Magnitude-based pruning for compression in FL

The Lottery Ticket Hypothesis (LTH) [43] states that given a large neural
network and a specific training task; there exists a sub-network within it
(“the winning ticket”), indicated by the nonzero elements of a binary mask
$M\in\mathbb{R}^{d}$, that when trained from its initialisation, can achieve
comparable or even better performance than the full network, using a fraction
of the parameters, $||M||_{0}<<d$. This observation has led researchers to
investigate techniques for finding such sub-networks, aiming to reduce the
computational resources required for training and inference. In the context of
FL, Lottery Ticket Networks (LTN) are commonly determined by pruning the
global model using client data. Most prominent in this context is the
LotteryFL method [44], where each client learns a personalised LTN (i.e. a
sub-network of the base model). Therefore, client-server communication can be
drastically reduced due to the compact size of the lottery networks. Since the
LTNs are personalised, both the values and their positional information are
communicated back to the server.

The observation of a similar test accuracy between the pruned and unpruned
network in a FL setup suggests the existence of a sparse, lottery ticket-like
sub-network, which can be uncovered during training. The insights gained from
this observation offer two avenues to reduce communication overheads:

  1. 1.

When clients prune after training on their local datasets $\mathcal{D}_{n,t}$,
clients need to transmit $||M_{n,t}\odot\theta_{n,t}||_{0}$ to the PS in
addition to $|M_{n,t}|=d$ bits to signal the pruned locations.

  2. 2.

Clients prune the global model before they begin training on their individual
datasets. Given that the parameter server already has access to the global
model, the clients are only required to send the parameters described by
$||M_{t}\odot\theta_{n,t}||_{0}$. Here, $M_{t}$ is the pruning mask determined
by the selected pruning technique, applied to the global model parameters
before a client’s training – which is the same on all client devices.

Our study is primarily focused on the latter approach, which offers some
distinct benefits. In the context of approach 1, each client produces a unique
pruning mask, this results in a collection of distinct pruning masks
$\mathbf{M}=\\{M_{1},M_{2},...,M_{n}\\}$ for the $N$ clients post-training.
For any two clients, $i$ and $j$ from the set of all clients, the support of
$M_{i}$ may or may not be a subset of the support of $M_{j}$. Consequently,
the union of supports of all masks is always greater than or equal to
$(1-p_{G})\cdot d$. Such a scenario can be interpreted as a form of
personalisation, similar to what has been demonstrated in [45]. On the PS,
client updates are aggregated according to their overlapping pruning masks to
produce a new global model. This demands sending more parameters and
additionally requires transmitting the merged client mask, resulting in an
enlarged downlink payload.

Contrarily, with approach 2, all clients share a common support, leading to a
uniformity in the masks; $supp(M_{i})=supp(M_{j})$, hence $M_{i,t}=M_{j,t}$,
ensuring the union of mask supports is precisely
$|\bigcup^{n}_{i=1}supp(M_{i})|=(1-p_{G})\cdot d$, as long as all clients
employ the same pruning rate $p_{G}$. We choose our pruning schedules s.t.
clients train unpruned models during the initial phase of FL training, similar
to what is demanded as initial condition in Algorithm 1. Post this phase,
pruning is executed before clients retrain their models, analogous to lines 2,
3 and 4 in Algorithm 1.

###  III-A FedMap Design principle

The main design principle behind FedMap revolves around collaboratively
learning and sparsifying the global model. In accordance with the iterative
pruning methodology (see Section II-B), clients follow a pruning schedule,
which determines how (we utilise the pruning methodology outlined in [20]) and
when (the size of the pruning interval $\mathbf{s}$, measured in FL rounds
$t$) to prune global model parameters. After clients locally update the global
models for $L$ local epochs on their respective datasets $D_{n,t}$; clients
compress the model by only keeping all nonzero values of the sparse client-
model, denoted as
$\Delta\check{\theta}_{n,t}\in\mathbb{R}^{\lfloor(1-p_{G})\cdot d\rceil}$,
which is then communicated to the PS, where aggregation is performed as
follows:

$\Delta\tilde{\theta}_{t+1}=\frac{\sum_{n=1}^{N}\Delta\tilde{\theta}_{n,t}\cdot\mathbb{1}(\Delta\tilde{\theta}_{n,t}\neq
0)}{\sum_{n=1}^{N}\mathbb{1}(\Delta\tilde{\theta}_{n,t}\neq 0)}.$ (4)

Producing the new global model is achieved by expanding the aggregation method
proposed in [22] (see Equation 2), to account for the client-subsets.
Specifically, for each client $n\in N$, the algorithm accumulates the changes
in parameters $\Delta\tilde{\theta}_{n,t}$ only when they are non-zero,
denoted via the indicator function $\mathbb{1}(\Delta\tilde{\theta}_{n,t}\neq
0)$. The summation of all these non-zero changes is then normalised by the
total number of clients that had non-zero changes at these specific locations
in their weights. In essence, the algorithm computes a weighted average of
parameter changes across all clients, depending on the support of individual
client masks.

At the beginning of each round $t$, each client $n$ receives the global update
$\Delta\check{\theta}_{t}$ from the PS (see Algorithm 2). Since all clients
follow the same pruning schedule, the sparse version of the update can be
recovered using the pruning mask from the previous round, $M_{t-1}$, to obtain
$\Delta\tilde{\theta}_{t}$. This operation is abbreviated with the function
$RFM(\cdot)$ (Recover From Mask), which can be formalised as follows:

$\tilde{\theta}=\left\\{\begin{matrix}\check{\theta}_{f(i)}&\text{ if
}M_{i}=1\\\ 0&\text{otherwise,}\end{matrix}\right.$ (5)

where $f(i)$ gives the index in $\check{\theta}$ corresponding to the $i$-th
non-zero entry in $M$.

After recovering the global model from the recovered model-delta
$\Delta\tilde{\theta}_{t}$ (see line 6 in Algorithm 2), the client computes
the size of the new set of nonzero parameters $K_{t}\in\mathbb{R}^{0:d}$ as a
function of the current FL round $t$ and model-size $d$ according to the
pruning schedule, denoted as $Schedule(\cdot)$. For a more thorough
description of the pruning schedules, see Section IV-A. The function
$Prune(\cdot)$ prunes the model to $K_{t}$ non-zero parameters; returning the
pruned model $\theta_{t}$ and the corresponding pruning mask $M_{t}$. We use
the aforementioned LAMP method [20] in the context of this work. After
training the sparse global model on $\mathcal{D}_{n,t}$ to obtain
$\tilde{\theta}_{n,t}$, the residuals are compressed via the $RWZ(\cdot)$
(Remove Where Zero) function. This function eliminates all sparse locations in
$\Delta\tilde{\theta}_{n,t}$ and returns
$\Delta\check{\theta}_{n,t}\in\mathbb{R}^{K_{t}}$.

The communication benefit with FedMap stems from only ever communicating
$K_{t}$ parameters during the uplink (see lines 12 and 17 in Algorithm 2),
without the necessity of sending additional masking information. In FedMap,
the PS acts solely as a vendor of aggregated client updates (see lines 16 and
17 in Algorithm 2), which has the same communication advantage for the
downlink.

Algorithm 2 The FedMap Algorithm

1:for $t=1,...,T$ do

2: Client Side:

3: for $n=1,...,N$ do

4: Receive $\Delta{\check{\theta}_{t}}$ from the PS

5: $\Delta\tilde{\theta}_{t}=RFM\left(\Delta\check{\theta}_{t},M_{t-1}\right)$

6: Update Global Model:
$\tilde{\theta}_{t}=\tilde{\theta}_{t-1}-\Delta\tilde{\theta}_{t}$

7: $K_{t}=Schedule(d,t)$

8: $\tilde{\theta}_{t},M_{t}\leftarrow Prune(\tilde{\theta}_{t},K_{t})$

9: Train $\tilde{\theta}_{t}$ on $\mathcal{D}_{n,t}$ for $L$ epochs to obtain
$\tilde{\theta}_{n,t}$.

10: $\Delta\tilde{\theta}_{n,t}=\tilde{\theta}_{t}-\tilde{\theta}_{n,t}$

11: Compress:
$\Delta\check{\theta}_{n,t}=RWZ\left(\Delta\tilde{\theta}_{n,t}\right)$

12: Send $\Delta\check{\theta}_{n,t}$ to the PS

13: Server Side (PS):

14: for $n=1,...,N$ do

15: Receive $\Delta{\check{\theta}_{n,t}}$ from client $n$

16: Global Averaging:
$\Delta\check{\theta}_{t+1}=\frac{1}{N}\cdot\sum_{n=1}^{N}\Delta{\check{\theta}_{n,t}}$

17: Send $\Delta\check{\theta}_{t+1}$ to clients

###  III-B Extending FedMap for Heterogeneous Data Distributions

Given recent developments in addressing challenges associated with learning
from non-IID data, we combine our FedMap method with the FedDR[46] method to
handle statistical heterogeneity. By using the Douglas-Rachford splitting
technique, FedDR aims to solve the common composite objective of the original
FL optimisation problem, which involves the addition of some regularizer
$\mathcal{L}(\theta)+g(\theta)$, where the authors propose $L1$
regularisation. Across our experimental evaluation, we find that $L1$
regularisation led to sub-optimal results, hence we omit the composite
setting. Additionally, we do not employ the randomised block-coordinate
technique proposed in [46], since this would require client sub-sampling which
is beyond the scope of this work.

Without applying $g(\theta)$, applying FedDR involves the following three
steps:

1) Intermediate variable update: After the client reconstructs the global
model, $\tilde{\theta}_{t}$, each client updates an intermediate variable,
$\tilde{\theta}^{y}_{n,t}$, influenced by both the global model, and the
client’s prior model $\tilde{\theta}_{n,t-1}^{y}$. The update is given as:

$\tilde{\theta}_{n,t}^{y}:=M_{t}\cdot\tilde{\theta}_{n,t-1}^{y}+\alpha(\tilde{\theta}_{t}-M_{t}\cdot\tilde{\theta}_{n,t-1})$
(6)

accumulating information from both the local and global models. Initially,
$\tilde{\theta}_{n,0}^{y}=\tilde{\theta}_{t=1}$, the global model at
initialisation. The gradual update helps address the non-convexity of the
problem by preventing drastic swings in model parameters. The step size,
$\alpha$, determines how quickly new knowledge is incorporated. We incorporate
updating of the intermediate variable in FedMap after line 8 in Algorithm 2.

2) Proximal operator application: A quadratic penalty is added to each
clients’ loss for each batch, forming the proximal operator. This is then
applied to the intermediate variable to obtain local model updates; augmenting
the local loss function:

$\displaystyle\tilde{\theta}_{n,t}$ $\displaystyle:=\text{prox}_{\eta
f_{c}}(\tilde{\theta}_{n,t}^{y})$ (7)
$\displaystyle:=\underset{\theta}{\text{arg
min}}\quad\mathcal{L}_{n}(\tilde{\theta}_{n,t})+\frac{1}{2\eta}\|\tilde{\theta}_{n,t}-\tilde{\theta}_{n,t}^{y}\|^{2}$

The quadratic penalty term is inversely modulated by $\eta$ where smaller
values for $\eta$ align the local model more with $\tilde{\theta}^{y}_{n,t}$,
balancing local adaptivity and global coherence.

3) Model reflection: To promote exploration, $\tilde{\theta}_{n,t}$ and
$\tilde{\theta}^{y}_{n,t}$ are combined in the reflection step:

$\tilde{\theta}^{x}_{n,t}=2\cdot\tilde{\theta}_{n,t}-\tilde{\theta}^{y}_{n,t}$
(8)

Finally, the difference
$\Delta\tilde{\theta}^{x}_{n,t}=\tilde{\theta}^{x}_{n,t}-M_{t}\cdot\tilde{\theta}^{x}_{n,t-1}$
, is sent back to the server for aggregation, where
$\Delta\tilde{\theta}^{x}_{n,t}$ replaces $\Delta\tilde{\theta}_{n,t}$ in line
10 in Algorithm 2.

##  IV Experiments

The experiments are organized into four categories: examining the impact of 1)
pruning schedules and 2) FedMap’s hyperparameters, 3) evaluating FedMap’s
performance in a non-IID deployment scenario, and 4) benchmarking against
_FederatedPruning_ [15].

###  IV-A Exp. A: Step-wise and Continuous Pruning Schedules

In our first experiment we study the impact of two different pruning schedules
under different pruning cadences (varying $\mathbf{s}$, see Fig. 1) on global
model performance. Central to the experiment was an exploration of two
distinct pruning schedules. The first schedule, known as the step-wise
schedule, involves a stepping function that determines when the model should
be pruned after a certain number of FL rounds. The step width of the model,
denoted as $\mathbf{s}$, follows the procedure outlined in Algorithm 1.

Secondly, we adopt a continuous pruning schedule, offering a smoother and more
gradual pruning of the model parameters. We utilize Bézier interpolation to
approximate the step-wise pruning schedule over the $T$ FL-iterations in our
experiments. With this continuous schedule, we closely match the cumulative
impact of the pruning process over $T$, effectively reducing the model at a
comparable pace, ensuring that at any given time, the model retains a similar
level of unpruned parameters. This smooth schedule aims for a more uniform
transition in pruning intensity.

Figure 1: The stepwise and continuous pruning schedules exemplified. We design
the schedules, s.t. the model can be trained to a good initial solution at the
start. After $\mathbf{s}$ FL rounds, clients prune the global model further,
s.t. $(1-p_{g})\cdot d$ parameters remain. We unanimously use an exponential
series, where $p_{g}=0.25$, progressively removing 25% of the remaining model
parameters.

In this first experiment, we trained 10 clients on IID splits of the CIFAR-10
dataset [47]. Each client participates in a set of $L=4$ epochs of local
training at every round $t$, before contributing to the global model update.
Furthermore, we compared the influence of pruning schedules on ResNet-56 [48]
and MobileNetV2 [49] network architectures111
https://github.com/chenyaofo/pytorch-cifar-models. The models were trained
using Stochastic Gradient Descent (SGD) with a learning rate of $\eta=0.01$,
weight decay rate of $5e^{-4}$ and batch size of $128$.

We experimented with the following pruning cadences;
$\mathbf{s}=\\{25,30,35,45,90\\}$, to assess the impact on global model
convergence in terms of overall performance and maximally achievable pruning
ratios of the models. Additionally, we fix the pruning proportion to
$p_{G}=0.25$, which means that over the step width $s$, an additional 25% of
the model parameters are pruned with respect to the unpruned parameters. We
also took precautions to set a defined pruning limit. Precisely, models were
pruned up to a maximum of 1% (100x compression) of their total parameters
accepting potential performance implications.

###  IV-B Exp. B: Influence of local training epochs and pruning cadence

Our second experiment aimed to explore the impact of modulating $L$ alongside
different pruning cadences. Specifically, we tested for $L=\\{2,4,8,16\\}$
while simultaneously adjusting $s=\\{45,90,135\\}$, using the CIFAR-10 dataset
as well as IID splits of the CIFAR-100 [47] and SVHN [50] datasets to expand
our empirical scope. To reduce the number of FL iterations, we re-calibrated
the pruning magnitude of the models to reach 5%, since beyond this point the
drop in performance is catastrophic. Based on the outcomes of the first
experiment, we favoured the step-wise pruning methodology over its continuous
counterpart. Finally, to make the setup more representative, we increased the
number of clients from 10 to 30.

###  IV-C Exp. C: Non-IID Experiments

In this set of experiments, we examine the performance of FedMap on two common
non-IID (Non-Independently and Identically Distributed) benchmarks, which
remark a more realistic scenario within real-world deployments. For this, we
use two datasets i.e., _i)_ the FEMNIST (Federated Extended MNIST), and _ii)_
the Shakespeare dataset [51]. FEMNIST is a modified version of the EMNIST [52]
dataset. It expands the $10$ class digit problem of the MNIST [53] dataset by
encompassing $62$ classes, including uppercase and lowercase characters. We
follow the partitioning scheme proposed in [51], which involves grouping
characters and digits by their respective author. In the context of Federated
Learning (FL), each author’s data corresponds to the data on an FL client.
This embodies a real-world scenario in which data is inherently partitioned
between different devices.

Our second dataset is the Shakespeare dataset, which includes all of William
Shakespeare’s works. Similarly to the FEMNIST dataset, the Shakespeare dataset
is organized based on a logical entity: the speaking roles in the plays [51].
Each speaking role is assigned to a different FL client, which simulates a
situation where each client device has data of a distinct categorical nature.
Table I provides a comprehensive summary of the dataset statistics. It is
imperative to note that the “Number of Slices” in the table corresponds to the
character writers in the FEMNIST dataset and individual speakers in the
Shakespeare dataset. This representation aids in understanding the
distribution and structure of the datasets in a manner conducive to FL
experimentation. The values for mean and standard deviation suggest that
FEMNIST presents a more uniformly distributed dataset across clients, offering
a benchmark for assessing our methodologies across varying degrees of non-IID
data.

TABLE I: Statistics of datasets in LEAF [51] Datasets | # Slices | Total Slices | Samples/Slice
---|---|---|---
|  |  | Mean | Std.
FEMNIST | 3,550 | 805,263 | 226.83 | 88.94
Shakespeare | 1,129 | 4,226,158 | 3,743.28 | 6,212.26

We subsample $150$ slices from either of the datasets. Due to the significant
volume of unused data, we employ a 60/40 train/test split, which translates to
90 FL clients for training and 60 evaluation slices for testing. Performance
assessment is carried out by averaging the results across all evaluation-
reserved slices, providing a robust evaluation of all FL configurations.

Models: For the FEMNIST dataset, we use the proposed model in LEAF [51], with
two convolutional layers followed by two fully connected layers. For the
Shakespeare dataset, characters are mapped to an eight-dimensional embedding,
then processed by a two-layer Long Short-Term Memory network (LSTM) [54] with
a sequence length of $80$. The LSTM outputs are mapped back to the vocabulary
for next-character prediction. Unlike in the original implementation, we added
batch-normalisation [55] for added stability and faster convergence, as well
as dropout [56] for regularisation to our setup.

Hyperparameters: For all experiments, data is batched with a size of $10$ to
align with the settings proposed in [46] and we use Stochastic Gradient
Descent (SGD) for optimisation. A learning rate of $3e^{-4}$ and $3e^{-3}$ is
used for the two datasets, respectively, and we set the dropout probability
for the LSTM with $0.5$. For the FEMNIST experiments we chose $L=20$ to align
with FedDR, and based on LEAF [51] we opt for $L=1$ for the Shakespeare
dataset.

The performance of models trained on the Shakespeare and FEMNIST datasets
responded significantly to changes in the hyperparameters for FedDR and FedMap
(see Section III-B). We observe that smaller step sizes such as the ones used
in experiment A, B, and C were too aggressive to maintain performance and was
further exacerbated for the FEMNIST dataset. To ensure sufficient convergence
and comparison, the step size has been titrated for both datasets such that
the models are able to learn enough before the first pruning round to prevent
an immediate and unrecoverable drop in performance. For Shakespeare this was
$\mathbf{s}=135$ and $\mathbf{s}=270$ for FEMNIST.

As described previously, the two main hyperparameters for FedDR are $\eta$ and
$\alpha$. We select these based on a small grid search around the parameters
of the original implementation for FEMNNIST in [46], that is,
$\eta=1000,\alpha=0.95$. Surprisingly, the original parameters were largely
unstable in our setup for FEMNIST, but more suitable for Shakespeare. FEMNIST
benefitted from a much lower $\eta$ at $10$, while Shakespeare performed best
at the proposed $\eta=1000$.

For $\alpha$ a value of $1$ and $1.2$ was found to be the best for Shakespeare
and FEMNIST respectively. Increasing $\alpha$ for FEMNNIST leads to higher
convergence but far less stability, while there is no significant difference
with Shakespeare.

After a certain number of pruning iterations, subsequent rounds in which
pruning takes place are markedly characterized by a drop in performance and
decreased stability. In FEMNIST, this occurs at the first pruning step, for
Shakespeare this is the second. We hypothesize that the performance of the
parameters chosen for FedDR may become unfaithful at these rounds and may
prove to be detrimental. Therefore, we propose several hybrid methods to alter
the training dynamics at these rounds:

  1. 1.

Switching from FedDR to FedAvg.

  2. 2.

Modifying the values of $\eta$ and $\alpha$.

For models trained on the Shakespeare dataset, adjusting $\alpha$ had a minor
impact on convergence, whereas varying $\eta$ had a much larger impact. Here
smaller values for $\eta$ seemed beneficial, leading to better stability,
faster convergence and better maximal performance. Performance improved as
$\eta$ increased. Specifically, the best results were obtained when we
completely removed the effect of the proximal operator and therefore
regularisation; just leaving the alterations in training dynamics via the
model reflection step. In contrast, FEMNIST experienced minimal effects from
the alteration of regularization (via adjustments $\eta$, and the best results
were found without change, i.e. $\eta=10$. Furthermore, increasing $\alpha$
surprisingly promoted stability, with the best value being $1.75$.

As a concluding extension, we investigated weighting local updates by each
client’s dataset size, aiming to equalise the influence of clients on the
model, regardless of their data volume. Specifically, this extension is
applied to Shakespeare, given its significantly larger sample skew; aiming to
prevent overfitting to predominant speaking-roles.

All of the above configurations are summarised in Table II. Additionally,
since the proximal term (see Eq. 7) represents the $L2$ difference between the
previous local model and the global model, this is skewed if the model is
pruned in the respective round. This is because the pruned model has fewer
trainable parameters. We took this into account and masked the parameters of
the local model of the previous round with the same pruning mask, preventing
any differences in pruned weights to accumulate. However, there was no effect
on training.

TABLE II: Specific FedDR configurations. Dynamic configurations change at the first pruning iteration for FEMNIST, and the second pruning iteration for Shakespeare. Name | Configuration Details
---|---
FedDR |  Basic FedDR implementation, no pruning.
FedMap-FedDR |  Basic FedDR with FedMap pruning.
FedMap-FedDR-C1 |  FedDR with pruning, changes to FedAvg.
FedMap-FedDR-C2 |  FedDR with pruning, changes $\alpha$ and $\eta$ to optimal values.
FedMap-FedDR-C3 |  FedMap-FedDR-C2, with updates weighted by local dataset size.

Figure 2: Comparison of pruning schedules trained on CIFAR10. Top row is with
Resnet56 and the bottom row is with MobileNet.

###  IV-D Exp. D: Benchmarking against FederatedPruning

We also performed experiments comparing against the FederatedPruning method
proposed by Lin et al. [15]. The key differences between FederatedPruning and
FedMap are as follows:

  1. 1.

FederatedPruning suggests pre-training and inital pruning on a centralised
server before beginning the FL procedure, akin to what [14] propose, which may
require settings where data exists sufficiently.

  2. 2.

FederatedPruning prunes the global model after global aggregation on the PS,
whereas in FedMap, pruning happens on the clients, eliminating the need to
communicate pruning masks.

  3. 3.

FederatedPruning proposes a novel structured ranking method to determine the
in- or exclusion from the set of prunables according to the $p_{G}$ in every
round; whereas in FedMap we rely on the state-of-the-art in unstructured
pruning, namely the LAMP method proposed in [20].

  4. 4.

FederatedPruning allows for the re-activation of previously pruned locations,
whereas in FedMap all pruning masks from progressive pruning are subsets of
preceding pruning masks.

To allow for a fair comparison, we operate FederatedPruning with LAMP instead
of their proposed structured pruning methodology, which we denote as
FederatedPruning-UnstructuredLamp in our experiments. Additionally, we chose
their structured pruning method targeting half-columns, which we denote as
FederatedPruning-StructuredColumnHalf, and which the authors chose within
their experimental settings.

We conducted two experiments, for two different models. First, we train the
MobileNetV2 model using the SVHN dataset, representing a relatively simple
task with an architecture optimised for mobile deployments [49]. This
combination of a lightweight model and a less complex dataset emulates
scenarios that may arise in real-world mobile deployments. Secondly, we train
the ResNet56 model on CIFAR100 dataset, combining a larger model architecture
with increased task-complexity. We follow the same pruning schedules and
training-parameters as outlined in our previous experiments.

Resnet56

MobilenetV2

Figure 3: Accuracy vs. rounds across variable number of local epochs ($L$) and
step-widths ($s$) for CIFAR10.

##  V Results and Discussion

###  V-A Investigating the Impact of Pruning Schedules on Model Performance
and Stability

In our first experiment, we assessed the influence of the type of pruning
schedule on the model’s ability to retain performance while progressing in
pruning. To this end, we conducted experiments using two image classification
models: Resnet56 [48] and MobileNetV2 [49]. Our findings were consistent
across both models, revealing some key insights. Expectedly, increased pruning
induced a noticeable decline in model performance. We discovered better
overall performance by extending the intervals between pruning iterations
(larger $\mathbf{s}$). This positive effect was observed even at higher
pruning rates, implying that less frequent pruning leads to better performance
at even higher pruning rates.

The two models produced varying results when comparing model sparsity and
accuracy. On average (across all interval sizes s), the step-wise pruning
schedule outperformed the continuous schedule. Leading to faster convergence
and similar performance (see MobileNetV2, bottom row in Figure 2) or increased
average performance (see ResNet56, top row in Figure 2). This trend is
especially evident at high pruning rates ($p_{G}\leq 20\%$).

A critical observation was made that both models maintained their performance
effectively up to a sparsity level of 80%, meaning that they retained only 20%
of their original parameters. After crossing a certain threshold, we observed
a significant reduction in the model’s performance. Despite this, we proceeded
with pruning beyond the 80% sparsity level, allowing us to gain insights into
how varying the step-width $\mathbf{s}$ affects the model’s performance under
extreme sparsity conditions, and to understand better the impact of different
pruning schedules on the model’s performance. After discovering that larger
$\mathbf{s}$ allow for more substantial pruning (retaining higher performance
longer), we investigate the underlying reasons for this effect.

Less frequent but more substantial pruning can improve model performance by
providing a stable learning and adaptation environment. When pruning occurs
less frequently, the models have more time to adjust and optimize their
remaining parameters in response to the reduced network complexity. This
uninterrupted learning phase is crucial for allowing the models to more
effectively assimilate the information from the training data, leading to a
more robust learning process [57]. As a result of reducing the frequency of
pruning, the final performance of the models can be improved. Frequent pruning
can cause disruption and hinder performance optimization within each pruned
state. By allowing models to settle and optimize their performance within each
state, they can maintain high accuracy levels, allowing for higher pruning
rates.

$\mathbf{L}$

Remaining
Parameters
(%)

(a) Resnet56

$\mathbf{L}$

Remaining
Parameters
(%)

(d) MobileNetV2

Figure 4: Average performance per pruning step ($s$) across variable different
values of local epochs ($L$) for CIFAR10.

Resnet56

Mobilenetv2_x0_5

Figure 5: Communicated bits (Mb) per round of FL for exp. B. on CIFAR10.

###  V-B Balancing Pruning Frequency and Local Training Regimes

Figure 3 (combined with Figure 8 in the appendix) illustrate our findings
indicating a significant relationship between task complexity, convergence,
and performance retention under high pruning rates. According to our analysis,
models trained on the CIFAR100 dataset may experience early performance
degradation (see Figure 8(a)). However, models trained on the CIFAR-10 and
SVHN datasets achieve peak performance more quickly and maintain superior
performance for a longer period (see Figures 8 and 8(b)). Models trained on
the latter can even withstand substantial pruning without significant loss of
performance. This phenomenon highlights the importance of model resilience in
the face of pruning, indicating considerable variation across different tasks.

Furthermore, the number of local training epochs emerges as a pivotal factor;
increasing the number of local training epochs, $\mathbf{L}$, enhances the
model’s ability to preserve high performance over time. However, prolonged
local training can increase the demand for client resources, such as energy
and compute, or increase the risk of model divergence. Excessive training on
local datasets can cause local models to deviate from the global model,
potentially reducing the overall performance across client devices and leading
to a degradation of the global model performance. While strategies to mitigate
this issue of model divergence have been proposed, such as those in [58] and
[59], they were not the focus of this work. Careful tuning of the local
training epochs $\mathbf{L}$ is therefore crucial to strike a balance between
preserving model performance under pruning and minimizing the associated
resource demands and risks of model divergence.

Our research has found that there is a certain point at which model
performance experiences a drastic decline due to excessive pruning. Before
this threshold is reached, the models can withstand pruning without
significant performance loss; however, after reaching this critical point,
model performance is irreversibly impaired. While the issue of performance
degradation following pruning is well-known in model pruning (see [20]),
deploying our FedMap methodology requires the inclusion of fallback mechanisms
to mitigate performance loss and ensure consistent delivery of expected
performance levels. This critical challenge must be addressed when deploying
FedMap in real-world scenarios.

In Figure 4, we demonstrate the delicate balance between pruning ratios,
pruning intervals $\mathbf{s}$ and local training epochs $\mathbf{L}$ on the
performance of Resnet56 and MobileNetV2 models trained on CIFAR10. Generally,
a larger $\mathbf{s}$ leads to better convergence at the beginning and to
models retaining higher overall performance for up to $p_{G}\geq 20\%$. Beyond
this point, and especially for pruning ratios where more than 90% of the
parameters are pruned, smaller $\mathbf{s}\in\\{45,90\\}$ lead to retaining
better performance. This phenomenon suggests that prolonged training may
unintentionally lead to over-fitting and, in turn, diminish the model’s
effectiveness at high pruning levels.

Additionally, we also performed the same experiment on the SVHN and CIFAR100
datasets (see Figures 9 and 11 in the appendix). For instance, a ResNet56
model trained on the SVHN task can withstand pruning down to 7.5% of its
original parameter count with minimal performance loss. Similarly, a
MobileNetV2 model can be aggressively pruned to 5% of its original size (20x
compression), accepting an 8.5% decrease in performance. When evaluating
models designed for the CIFAR100 dataset, it is evident that pruning beyond
75% leads to severe performance degradation. This threshold highlights a
crucial point at which reducing model parameters becomes counterproductive.
Furthermore, our investigation reveals that extending the number of FL
iterations (larger $T$) before pruning ($\mathbf{s}=135$) offers marginal
benefits in performance; for higher pruning ratios, a shorter pruning interval
($\mathbf{s}$) proves more advantageous, maintaining better performance under
such rigid compression conditions.

Moreover, when we increased the local training epochs ($\mathbf{L}$) for a
model, it became better equipped to handle higher pruning rates. This
highlights the importance of on-device training, as it helps to improve the
model’s ability to withstand significant parameter reductions. On the
contrary, larger pruning intervals $\mathbf{s}$ may lead to overall higher
top-performance. When considering very large pruning ratios, smaller
$\mathbf{s}$, lead to higher performance at higher pruning ratios.
Furthermore, we can see that the complexity of the task at hand also affects
how tolerant a model is to pruning.

In Figure 5, we present a quantitative analysis of the bytes communicated for
varying values of the step width, $\mathbf{s}$. Note, the CIFAR10 and SVHN
datasets are both problems with 10 classes each, hence the number of
parameters is equivalent (for results on the CIFAR100 and SVHN datasets, see
Figure 10). A lower value for $\mathbf{s}$ leads to faster pruning and less
data transmitted. This efficiency underscores the importance of optimising
$\mathbf{s}$ to reduce communication overhead in FL settings. Consistent with
our prior findings, increasing the number of local training epochs (denoted by
a large $\mathbf{L}$) is instrumental. Extended local training allows clients
to refine the model more thoroughly before aggregation, enhancing the model’s
resilience to pruning-induced performance decline. However, it is imperative
to recognise the potential drawbacks of excessively elongating the local
training phase. Overly extensive local training periods can inadvertently
impair the global model’s performance, a phenomenon well-documented in
existing literature [60, 61, 62, 58, 59]. This delicate balance emphasises the
need for strategic selection of $\mathbf{L}$, ensuring it is sufficient to
support effective pruning while avoiding the counterproductive effects of
over-training on the model’s aggregate performance.

Shakespeare

FEMNIST

Figure 6: Accuracy vs. rounds for the Non-IID experiments. Refer to table II
for FedDR configurations.

###  V-C Integrating FedDR into FedMap for Enhanced Non-IID Pruning

Figure 6 depicts our non-IID results for Shakespeare (left) and FEMNIST
(right) datasets. For both, as previously mentioned, the basic FedMap-FedDR
approach sharply drops in accuracy after the first few pruning points.
Configuration FedDR-C1, that switches to FedAvg after these pruning points,
was able to significantly stabilise at later pruning rounds and outperformed
FedMap-FedDR for Shakespeare. For FEMNIST, while it was able to stabilise
until the subsequent pruning point, it then experienced more rapid decline,
more than FedMap-FedDR. However, this decline was more consistent and the
final result was marginally higher than FedMap-FedDR.

Our approach to further adapting the hyperparameters, in configurations
FedMap-FedDR-C2 – FEMNIST, and FedMap-FedDR-C3 – Shakespeare, resulted in more
desirable results. For both datasets, these methods were able to provide large
improvements in stability and preserving significantly more accuracy than
FedMap-FedDR. For the Shakespeare dataset using FedMap-FedDR-C3, we were able
to achieve similar accuracy to FedAvg, showcasing the potential synergy
between FedMap and FedDR. For FEMNIST, we achieve at least 80% pruning (after
5 pruning steps i.e., by gradually removing 25% of the weights in the previous
step).

Although these are promising results for preserving accuracy in the context of
pruning, the optimal hyperparameter and performance seem highly dataset
dependent, as evidenced by their response to different settings. Currently, it
is very challenging to achieve a generalisable setup for maintaining
performance, but our experiments suggest numerous opportunities for further
tailoring these methods in the future.

MobileNetV2, SVHN

Resnet56, CIFAR100

(a) Accuracy vs. Rounds

(b) Total bits communicated (Mb) vs. Rounds

Figure 7: Comparison of accuracy and communication cost for Federated Pruning
benchmark (Exp. D).

###  V-D Evaluating FedMap Against Unstructured FederatedPruning

To ensure appropriate evaluation, we modified the FederatedPruning [15]
framework to incorporate the same unstructured pruning technique used in our
FedMap methodology. This adaptation, which we designate as FederatedPruning-
UnstructuredLAMP, employs the LAMP method [20]. Our comparative analysis,
depicted in Figure 7, yields several insights: The iterative update of pruning
masks at each specified interval $\mathbf{s}$ induces a notable drawdown
effect in the performance of the baseline FederatedPruning method, which is
proposed utilising structured pruning. This phenomenon can be attributed to
significant portions of the model being alternately activated or deactivated
across subsequent FL training cycles, necessitating a period of adjustment for
the global model to realign with the overarching training objectives.

This initial performance setback is significantly mitigated when
FederatedPruning is paired with our unstructured LAMP modification, and the
drawdown effect is predominantly observed mostly at higher pruning thresholds.
In more aggressive pruning scenarios, FederatedPruning, using the unstructured
approach, demonstrably surpasses FedMap in performance.

Conversely, FedMap is largely immune to the performance fluctuations
associated with rapid changes in pruning masks. This stability is due to its
pruning strategy, wherein the indices of parameters retained after each
pruning event are strictly subsets of those selected in the preceding
interval. This ensures consistent performance, with notable drawbacks only
manifesting at exceptionally high pruning ratios. For instance, ResNet56
trained on the CIFAR100 dataset illustrates that while performance may decline
at extreme pruning levels, it stabilises without further deterioration, unlike
the fluctuating recovery observed in FederatedPruning. Across all phases of
the FL training process, FedMap consistently demonstrates superior efficiency
in reducing communication overhead compared to FederatedPruning.

In summary, our FedMap methodology presents a formidable alternative to
conventional approaches like FederatedPruning [15]. It strikes a delicate
balance between optimising communication efficiency and maximising the
performance benefits derived from prolonged model training. This comparison
underscores a fundamental trade-off between minimising communication demands
and leveraging extended training periods to enhance model performance without
incurring adverse effects.

  * •

In centralised, i.i.d settings, Iterative Magnitude-based pruning achieves
very high sparsity ratios while still retaining good performance. FedMap
achieved convergence on par with FedAvg, but the sparsity ratios observed with
non-federated (centralised) iterative pruning are not achievable, even in iid
settings.

  * •

While FederatedPruning retains better performance than FedMap for very high
sparsity, it never converges to the same performance as does FedMap.

  * •

Another drawback of FederatedPruning is the high variance in model
performance. While in FedMap, the client-models achieve high performance even
at higher pruning ratios, the client-models in FederatedPruning wildly vary in
performance, not guaranteeing stable performance.

  * •

One additional benefit of FedMap is the relatively smaller transmission
overhead, compared to FederatedPruning. While with FedMap, no masking
information is ever transmitted, FederatedPruning transmits masking data every
FL round during the pruning- and fine-tuning phase of its lifecycle.

##  VI Conclusions, Limitations, and Future Work

In this paper we introduced FedMap, a novel communication-efficient pruning
technique tailored for federated learning deployments. Federated learning,
while enabling privacy-preserving distributed training on decentralized data,
often faces challenges due to the resource constraints of client devices.
FedMap addresses this issue by collaboratively learning an increasingly sparse
global model through iterative, unstructured pruning. A key strength of FedMap
is its ability to train a global model from scratch, making it well-suited for
privacy-critical domains where suitable pretraining data are limited. By
adapting iterative magnitude-based pruning to the federated setting, FedMap
ensures that all clients prune and refine the same subset of the global model
parameters, gradually reducing the model size and communication overhead. The
iterative nature of FedMap, where subsequent models are formed as subsets of
predecessors, avoids parameter reactivation issues encountered in prior work,
resulting in stable performance. Through extensive evaluation across diverse
settings, datasets, model architectures, and hyperparameters, assessing
performance in both IID and non-IID environments, we demonstrated the
effectiveness of FedMap. Comparative analysis against baseline approaches
highlighted FedMap’s ability to achieve more stable client model performance
and superior results.

During our evaluation, some limitations were identified, which, if addressed,
could further enhance the capabilities and applicability of FedMap. One
limitation lies in the current approach where all clients follow the same
pruning schedule, using the same parameters $s$ and $p_{G}$. While this
uniform approach simplifies the implementation, it may not be optimal for
scenarios with heterogeneous client devices and varying resource constraints.
To address this limitation, future work could explore adaptive pruning
schedules tailored to individual clients. By allowing clients to prune to
varying degrees based on their specific bandwidth limitations or hardware
capabilities, FedMap could better accommodate diverse client environments.
Clients with limited bandwidth could adopt more aggressive pruning schedules
to reduce communication overhead, while clients with more powerful hardware
could maintain higher model complexities to maximize performance.

While methods like FederatedPruning promise hardware-acceleration and thus
increasing training speeds while reducing inference-latency, we focus on
communication-reduction in this work and show that not only does iterative
pruning lead to more stable model behaviour but also allows to retain higher
model-performance at higher pruning rates. Moreover, although the present
study concentrates on the federated learning scenario, investigating the
potential application of FedMap in other distributed learning paradigms, such
as split learning or collaborative learning, could further broaden its impact
and utility. Overall, FedMap presents a promising solution to alleviate
communication bottlenecks in federated learning systems while retaining model
accuracy, offering a valuable contribution to the field of privacy-preserving
distributed machine learning.

## References

  * [1] R. Xu, N. Baracaldo, and J. Joshi, “Privacy-preserving machine learning: Methods, challenges and directions,” _CoRR_ , vol. abs/2108.04417, 2021. [Online]. Available: https://arxiv.org/abs/2108.04417
  * [2] J. Konecný, H. B. McMahan, and D. Ramage, “Federated optimization: Distributed optimization beyond the datacenter,” _ArXiv_ , 2015.
  * [3] X. Yin, Y. Zhu, and J. Hu, “A comprehensive survey of privacy-preserving federated learning: A taxonomy, review, and future directions,” _ACM Comput. Surv._ , vol. 54, no. 6, jul 2021. [Online]. Available: https://doi.org/10.1145/3460427
  * [4] A. Herzog _et al._ , “Selective updates and adaptive masking for communication-efficient federated learning,” _IEEE Transactions on Green Communications and Networking_ , pp. 1–1, 2024.
  * [5] K. Pfeiffer _et al._ , “Federated learning for computationally constrained heterogeneous devices: A survey,” _ACM Comput. Surv._ , vol. 55, no. 14s, jul 2023. [Online]. Available: https://doi.org/10.1145/3596907
  * [6] A. Imteaj _et al._ , “A survey on federated learning for resource-constrained iot devices,” _IEEE Internet of Things Journal_ , vol. PP, pp. 1–1, 07 2021.
  * [7] I. K. andAdwaita Janardhan Jadhav, “Federated learning in iot: a survey from a resource-constrained perspective,” 2023.
  * [8] V. K. Prasad _et al._ , “Federated learning for the internet-of-medical-things: A survey,” _Mathematics_ , vol. 11, no. 1, 2023. [Online]. Available: https://www.mdpi.com/2227-7390/11/1/151
  * [9] O. R. A. Almanifi _et al._ , “Communication and computation efficiency in federated learning: A survey,” _Internet of Things_ , vol. 22, p. 100742, 2023. [Online]. Available: https://www.sciencedirect.com/science/article/pii/S2542660523000653
  * [10] T. Tuor _et al._ , “Overcoming noisy and irrelevant data in federated learning,” in _2020 25th International Conference on Pattern Recognition (ICPR)_ , 2021, pp. 5020–5027.
  * [11] M. Asad _et al._ , “Limitations and future aspects of communication costs in federated learning: A survey,” _Sensors_ , vol. 23, no. 17, 2023. [Online]. Available: https://www.mdpi.com/1424-8220/23/17/7358
  * [12] J. Wen _et al._ , “A survey on federated learning: challenges and applications,” _International Journal of Machine Learning and Cybernetics_ , vol. 14, no. 2, pp. 513–535, Feb 2023. [Online]. Available: https://doi.org/10.1007/s13042-022-01647-y
  * [13] Y. LeCun, J. Denker, and S. Solla, “Optimal brain damage,” in _Advances in Neural Information Processing Systems_ , D. Touretzky, Ed., vol. 2. Morgan-Kaufmann, 1989. [Online]. Available: https://proceedings.neurips.cc/paper_files/paper/1989/file/6c9882bbac1c7093bd25041881277658-Paper.pdf
  * [14] Y. Jiang _et al._ , “Model pruning enables efficient federated learning on edge devices,” _IEEE Transactions on Neural Networks and Learning Systems_ , pp. 1–13, 2022.
  * [15] R. Lin _et al._ , “Federated Pruning: Improving Neural Network Efficiency with Federated Learning,” in _Proc. Interspeech 2022_ , 2022, pp. 1701–1705.
  * [16] N. Rieke _et al._ , “The future of digital health with federated learning,” _npj Digital Medicine_ , vol. 3, no. 1, p. 119, Sep 2020. [Online]. Available: https://doi.org/10.1038/s41746-020-00323-1
  * [17] C. M. Lee _et al._ , “Federated learning for credit risk assessment.” in _HICSS_ , 2023, pp. 386–395.
  * [18] Z. Yang and H. Zhang, “Comparative analysis of structured pruning and unstructured pruning,” in _Frontier Computing_ , J. C. Hung, N. Y. Yen, and J.-W. Chang, Eds. Singapore: Springer Nature Singapore, 2022, pp. 882–889.
  * [19] Y. He and L. Xiao, “Structured pruning for deep convolutional neural networks: A survey,” 2023.
  * [20] J. Lee _et al._ , “Layer-adaptive sparsity for the magnitude-based pruning,” in _International Conference on Learning Representations_ , 2021.
  * [21] S. A. Bhalgaonkar, M. V. Munot, and A. D. Anuse, “Pruning for compression of visual pattern recognition networks: A survey from deep neural networks perspective,” in _Pattern Recognition and Data Analysis with Applications_ , D. Gupta _et al._ , Eds. Singapore: Springer Nature Singapore, 2022, pp. 675–687.
  * [22] H. B. McMahan _et al._ , “Communication-efficient learning of deep networks from decentralized data,” in _AISTATS_ , 2017.
  * [23] K. Hsieh _et al._ , “The non-iid data quagmire of decentralized machine learning,” in _Proceedings of the 37th International Conference on Machine Learning_ , ser. ICML’20. JMLR.org, 2020.
  * [24] X. Ma _et al._ , “A state-of-the-art survey on solving non-iid data in federated learning,” _Future Generation Computer Systems_ , vol. 135, pp. 244–258, 2022. [Online]. Available: https://www.sciencedirect.com/science/article/pii/S0167739X22001686
  * [25] N. Lee, T. Ajanthan, and P. Torr, “Snip: single-shot network pruning based on connection sensitivity.” Open Review, 2019.
  * [26] P. de Jorge _et al._ , “Progressive skeletonization: Trimming more fat from a network at initialization,” in _International Conference on Learning Representations_ , 2021. [Online]. Available: https://openreview.net/forum?id=9GsFOUyUPi
  * [27] H. Tanaka _et al._ , “Pruning neural networks without any data by iteratively conserving synaptic flow,” in _Proceedings of the 34th International Conference on Neural Information Processing Systems_ , ser. NIPS’20. Red Hook, NY, USA: Curran Associates Inc., 2020.
  * [28] C. Wang, G. Zhang, and R. Grosse, “Picking winning tickets before training by preserving gradient flow,” in _International Conference on Learning Representations_ , 2020. [Online]. Available: https://openreview.net/forum?id=SkgsACVKPH
  * [29] S. M. Patil and C. Dovrolis, “Phew : Constructing sparse networks that learn fast and generalize well without training data,” in _Proceedings of the 38th International Conference on Machine Learning_ , ser. Proceedings of Machine Learning Research, M. Meila and T. Zhang, Eds., vol. 139. PMLR, 18–24 Jul 2021, pp. 8432–8442. [Online]. Available: https://proceedings.mlr.press/v139/patil21a.html
  * [30] G. Castellano, A. M. Fanelli, and M. Pelillo, “An iterative pruning algorithm for feedforward neural networks,” _IEEE transactions on neural networks_ , vol. 8 3, pp. 519–31, 1997. [Online]. Available: https://api.semanticscholar.org/CorpusID:6014155
  * [31] D. C. Mocanu _et al._ , “Scalable training of artificial neural networks with adaptive sparse connectivity inspired by network science,” _Nature Communications_ , vol. 9, no. 1, Jun 2018.
  * [32] T. Gale, E. Elsen, and S. Hooker, “The state of sparsity in deep neural networks,” _CoRR_ , vol. abs/1902.09574, 2019. [Online]. Available: http://arxiv.org/abs/1902.09574
  * [33] U. Evci _et al._ , “Rigging the lottery: Making all tickets winners,” in _International Conference of Machine Learning_ , 2020.
  * [34] D. Alistarh _et al._ , “The convergence of sparsified gradient methods,” _CoRR_ , vol. abs/1809.10505, 2018. [Online]. Available: http://arxiv.org/abs/1809.10505
  * [35] S. U. Stich, J. Cordonnier, and M. Jaggi, “Sparsified SGD with memory,” _CoRR_ , vol. abs/1809.07599, 2018. [Online]. Available: http://arxiv.org/abs/1809.07599
  * [36] F. Sattler _et al._ , “Sparse binary compression: Towards distributed deep learning with minimal communication,” in _2019 International Joint Conference on Neural Networks (IJCNN)_ , 2019, pp. 1–8.
  * [37] D. Wu _et al._ , “Smartidx: Reducing communication cost in federated learning by exploiting the cnns structures,” in _Thirty-Sixth AAAI Conference on Artificial Intelligence, AAAI 2022, Thirty-Fourth Conference on Innovative Applications of Artificial Intelligence, IAAI 2022, The Twelveth Symposium on Educational Advances in Artificial Intelligence, EAAI 2022 Virtual Event, February 22 - March 1, 2022_. AAAI Press, 2022, pp. 4254–4262. [Online]. Available: https://ojs.aaai.org/index.php/AAAI/article/view/20345
  * [38] E. Ozfatura, K. Ozfatura, and D. Gündüz, “Time-correlated sparsification for communication-efficient federated learning,” _2021 IEEE International Symposium on Information Theory (ISIT)_ , pp. 461–466, 2021.
  * [39] T. Lin _et al._ , “Dynamic model pruning with feedback,” 2020. [Online]. Available: https://arxiv.org/abs/2006.07253
  * [40] H. Mostafa and X. Wang, “Parameter efficient training of deep convolutional neural networks by dynamic sparse reparameterization,” _CoRR_ , vol. abs/1902.05967, 2019. [Online]. Available: http://arxiv.org/abs/1902.05967
  * [41] M. Zhu and S. Gupta, “To prune, or not to prune: Exploring the efficacy of pruning for model compression,” in _6th International Conference on Learning Representations, ICLR 2018, Vancouver, BC, Canada, April 30 - May 3, 2018, Workshop Track Proceedings_. OpenReview.net, 2018. [Online]. Available: https://openreview.net/forum?id=Sy1iIDkPM
  * [42] A. Kusupati _et al._ , “Soft threshold weight reparameterization for learnable sparsity,” _CoRR_ , vol. abs/2002.03231, 2020. [Online]. Available: https://arxiv.org/abs/2002.03231
  * [43] J. Frankle and M. Carbin, “The lottery ticket hypothesis: Finding sparse, trainable neural networks,” in _International Conference on Learning Representations_ , 2019. [Online]. Available: https://openreview.net/forum?id=rJl-b3RcF7
  * [44] A. Li _et al._ , “Lotteryfl: Empower edge intelligence with personalized and communication-efficient federated learning,” in _2021 IEEE/ACM Symposium on Edge Computing (SEC)_ , 2021, pp. 68–79.
  * [45] ——, “Fedmask: Joint computation and communication-efficient personalized federated learning via heterogeneous masking,” in _Proceedings of the 19th ACM Conference on Embedded Networked Sensor Systems_ , ser. SenSys ’21. New York, NY, USA: Association for Computing Machinery, 2021, p. 42–55. [Online]. Available: https://doi.org/10.1145/3485730.3485929
  * [46] Q. Tran-Dinh _et al._ , “Feddr – randomized douglas-rachford splitting algorithms for nonconvex federated composite optimization,” 2021.
  * [47] A. Krizhevsky, “Learning multiple layers of features from tiny images,” pp. 32–33, 2009. [Online]. Available: https://www.cs.toronto.edu/~kriz/learning-features-2009-TR.pdf
  * [48] K. He _et al._ , “Deep residual learning for image recognition,” in _2016 IEEE Conference on Computer Vision and Pattern Recognition (CVPR)_ , 2016, pp. 770–778.
  * [49] M. Sandler _et al._ , “Inverted residuals and linear bottlenecks: Mobile networks for classification, detection and segmentation,” _CoRR_ , vol. abs/1801.04381, 2018. [Online]. Available: http://arxiv.org/abs/1801.04381
  * [50] Y. Netzer _et al._ , “Reading digits in natural images with unsupervised feature learning,” in _NIPS Workshop on Deep Learning and Unsupervised Feature Learning 2011_ , 2011. [Online]. Available: http://ufldl.stanford.edu/housenumbers/nips2011_housenumbers.pdf
  * [51] S. Caldas _et al._ , “LEAF: A benchmark for federated settings,” _CoRR_ , vol. abs/1812.01097, 2018. [Online]. Available: http://arxiv.org/abs/1812.01097
  * [52] G. Cohen _et al._ , “Emnist: Extending mnist to handwritten letters,” _2017 International Joint Conference on Neural Networks (IJCNN)_ , pp. 2921–2926, 2017. [Online]. Available: https://api.semanticscholar.org/CorpusID:30587588
  * [53] L. Deng, “The mnist database of handwritten digit images for machine learning research [best of the web],” _IEEE Signal Processing Magazine_ , vol. 29, no. 6, pp. 141–142, 2012.
  * [54] S. Hochreiter and J. Schmidhuber, “Long short-term memory,” _Neural Comput._ , vol. 9, no. 8, p. 1735–1780, nov 1997. [Online]. Available: https://doi.org/10.1162/neco.1997.9.8.1735
  * [55] S. Ioffe and C. Szegedy, “Batch normalization: Accelerating deep network training by reducing internal covariate shift,” in _Proceedings of the 32nd International Conference on Machine Learning_ , ser. Proceedings of Machine Learning Research, F. Bach and D. Blei, Eds., vol. 37. Lille, France: PMLR, 07–09 Jul 2015, pp. 448–456. [Online]. Available: https://proceedings.mlr.press/v37/ioffe15.html
  * [56] N. Srivastava _et al._ , “Dropout: A simple way to prevent neural networks from overfitting,” _Journal of Machine Learning Research_ , vol. 15, no. 56, pp. 1929–1958, 2014. [Online]. Available: http://jmlr.org/papers/v15/srivastava14a.html
  * [57] D. M. L. Barbato and O. Kinouchi, “Optimal pruning in neural networks,” _Phys. Rev. E_ , vol. 62, pp. 8387–8394, Dec 2000. [Online]. Available: https://link.aps.org/doi/10.1103/PhysRevE.62.8387
  * [58] C. Ma _et al._ , “Federated learning with unreliable clients: Performance analysis and mechanism design,” _IEEE Internet of Things Journal_ , vol. 8, pp. 17 308–17 319, 2021. [Online]. Available: https://api.semanticscholar.org/CorpusID:234482857
  * [59] R. Xu, X. Feng, and H. Zheng, “Robust model aggregation for federated learning with heterogeneous clients,” _2021 7th International Conference on Computer and Communications (ICCC)_ , pp. 1606–1610, 2021. [Online]. Available: https://api.semanticscholar.org/CorpusID:246037503
  * [60] J. Zhang _et al._ , “Adaptive federated learning on non-iid data with resource constraint,” _IEEE Transactions on Computers_ , vol. 71, pp. 1655–1667, 2021. [Online]. Available: https://api.semanticscholar.org/CorpusID:237858014
  * [61] Y. Yan _et al._ , “On-device learning for model personalization with large-scale cloud-coordinated domain adaption,” _Proceedings of the 28th ACM SIGKDD Conference on Knowledge Discovery and Data Mining_ , 2022. [Online]. Available: https://api.semanticscholar.org/CorpusID:251518212
  * [62] Q. Wu, L. Liu, and S. Xue, “Global update guided federated learning,” _2022 41st Chinese Control Conference (CCC)_ , pp. 2434–2439, 2022. [Online]. Available: https://api.semanticscholar.org/CorpusID:248069338

##  Appendix A Additional Experimental Results

###  A-A FedMap Experimental Results

Resnet56

MobilenetV2

(a) CIFAR100

(b) SVHN

Figure 8: Accuracy vs. rounds across variable number of local epochs ($L$) and
step-widths ($s$). Experiments were performed with the following sets of
hyperparameters: $L\in\\{2,4,8\\},s\in\\{45,90,135\\}$

$\mathbf{L}$

Remaining
Parameters
(%)

(a) CIFAR100

$\mathbf{L}$

Remaining
Parameters
(%)

(b) SVHN

Figure 9: Average performance per pruning step ($s$) across variable different
values of local epochs ($L$) for Resnet56.

(a) CIFAR100

(b) SVHN

Figure 10: Communicated bits (Mb) per round of FL for exp. B. for Resnet56.

$\mathbf{L}$

Remaining
Parameters
(%)

(a) CIFAR100

$\mathbf{L}$

Remaining
Parameters
(%)

(b) SVHN

Figure 11: Average performance per pruning step (s) across variable different
values of local epochs (L) for MobileNetV2.

(a) CIFAR100

(b) SVHN

Figure 12: Communicated bits (Mb) per round of FL for exp. B. on MobileNetV2.