Keywords
Efficient Federated Learning, Neuro-imaging, Sparsity
1 INTRODUCTION
Deep learning has transformed fields like computer vision, natural language processing, and is also starting
to transform the field of neuroimaging. As deep learning models grow, distributed and collaborative training
becomes essential, especially when sensitive data is spread across distant sites. Collaborative MRI data
analysis offers profound insights, allowing researchers to utilize data beyond a study’s original scope.
As MRI scans are often preserved, vast amounts of data accumulate across decentralized research sites.
Technological advancements have increased data complexity and reduced costs, prompting researchers to
harness data from various groups for larger sample sizes and revealing significant features while preserving
the privacy of the data.
Training models on more data, while preserving data privacy is thus crucial. Aggregating data from
different sources to a central server for training can however expose this sensitive information, raising
ethical concerns. Federated learning (FL) addresses this by allowing devices or organizations to train
models locally and share training details without sharing the actual data.
1
.CC-BY 4.0 International licensemade available under a
(which was not certified by peer review) is the author/funder, who has granted bioRxiv a license to display the preprint in perpetuity. It is
The copyright holder for this preprintthis version posted May 15, 2024. ; https://doi.org/10.1101/2024.05.14.594167doi: bioRxiv preprint
Bishal et al. Sparse FL for NeuroImaging
In federated learning (FL), a central server coordinates training, and client sites communicate only model
parameters, keeping local data private. In the decentralized setting, the server usually doesn’t exist and
clients train a model collaboratively among themselves. However, challenges arise due to data’s statistical
heterogeneity, limited communication bandwidth, and computational costs. In this work, we focus on
addressing the communication efficiency inherent in distributed federated neuro-imaging pipelines through
the training of sparse models at local sites.
1.1 Federated Learning
Federated Learning distinguishes itself from traditional distributed learning approaches, such as those
discussed by McMahan et al. (2017), through its unique characteristics:
1. Non-IID Data: The training data across different clients are not identically distributed, which means
that the data at each local site may not accurately represent the overall population distribution.
2. Unbalanced Data: The amount of data varies significantly across clients, leading to imbalances in
data representation.
3. Massive Distribution: Often, the number of clients exceeds the average number of samples per client,
illustrating the scale of distribution.
4. Limited Communication: Communication is infrequent, either among clients in a decentralized
setting or between clients and the server in a centralized setting, due to slow and expensive connections.
One of the main focuses of this work is to reduce the communication costs between the server and
clients in a centralized setting or among clients in a decentralized setting when dealing with non-IID and
unbalanced data. This is achieved by identifying a sub-network based on the data distributions at each
local site and transmitting only the parameters of this sub-network in each communication round r. In
each round, a fixed set of ˜K clients is sampled from all K clients, and federated training continues on the
selected sub-network of those clients. The general federated optimization problem encountered is detailed
next.
1.2 Federated Optimization Problem
In the general federated learning (FL) setting, a central server tries to find a global statistical model by
periodically communicating with a set of clients. The federated averaging algorithm proposed by McMahan
et al. (2017); Koneˇcn`y et al. (2016); Bonawitz et al. (2019) is applicable to any finite sum objective of the
form
min
w∈Rd
f(w), where f(w) = 1
n
nX
i=1
fi(w). (1)
In a typical machine learning problem, the objective function fi(w) = ℓ(xi, yi; w) is encountered, where
the ith term in the sum is the loss of the network prediction on a sample (xi, yi) made by a model with
parameter w. We assume that the data is partitioned over a total of K clients, with Pk denoting the set of
indices of the samples on client k, and nk = |Pk|. Thus, the objective in (1) can be re-written as
f(w) =
KX
k=1
nk
n Fk(w), where Fk(w) = 1
nk
X
i∈Pk
fi(w). (2)
Frontiers 2
.CC-BY 4.0 International licensemade available under a
(which was not certified by peer review) is the author/funder, who has granted bioRxiv a license to display the preprint in perpetuity. It is
The copyright holder for this preprintthis version posted May 15, 2024. ; https://doi.org/10.1101/2024.05.14.594167doi: bioRxiv preprint
Bishal et al. Sparse FL for NeuroImaging
In the typical distributed optimization setting, the IID assumption is made, which says the following: if
the partition Pk was created by distributing the training data over the set of clients uniformly at random,
then we would have EPk[Fk(w)] = f(w), where the expectation is over the set of examples assigned to a
fixed client k. In this work, we consider the non-IID setting where this does not hold and Fk could be an
arbitrarily bad approximation to f.
When designing an FL training paradigm, a set of core considerations have to be made to maintain data
privacy and address statistical or objective heterogeneity due to the differences in client data and resource
constraints at the client sites. A range of work tries to address the issue of heterogeneous non-IID data
McMahan et al. (2016); Kulkarni et al. (2020), however, some research also suggests that deterioration in
accuracy in the FL non-IID setting is almost inevitable Zhao et al. (2018).
1.3 Federated Learning in Neuro-imaging
Over the past decade, the field of neuroimaging has strongly embraced data sharing, open-source
software, and collaboration across multiple sites. This shift is largely attributed to the high costs and time
demands associated with neuroimaging data collection Landis et al. (2016); Rootes-Murdy et al. (2022).
By pooling data from different sources, researchers can explore findings that extend beyond the initial
scope of individual studies Poldrack et al. (2013). The practice of sharing data enhances the robustness
of research through larger sample sizes and the replication of results, offering significant benefits for
neuroimaging studies. Increasing the sample size not only makes predictions more reliable but also ensures
the reliability and validity of research findings, thereby preventing data manipulation and fabrication Ming
et al. (2017); Tenopir et al. (2011). Furthermore, aggregating data can lead to a more diverse sample by
combining otherwise similar datasets, thus reflecting a broader range of social health determinants for
more comprehensive results Laird (2021). Additionally, reusing data can significantly reduce research costs
Milham et al. (2018).
Federated Learning (FL) is increasingly recognized as a transformative approach in healthcare and
neuro-imaging. In the realm of biomedical imaging, FL has been applied to a variety of tasks. These
include whole-brain segmentation from MRI T1 scans Roy et al. (2019), segmentation of brain tumors
Sheller et al. (2019); Li et al. (2019), multi-site fMRI classification, and the identification of disease
biomarkers Li et al. (2020b). COINSTAC Plis et al. (2016a) offers a privacy-focused distributed data
processing framework specifically designed for brain imaging showcasing FL’s role in enhancing privacy
and efficiency in healthcare data analysis. Additionally, it has been utilized in discovering brain structural
relationships across various diseases and clinical cohorts through federated dimensionality reduction from
shape features Silva et al. (2019).
1.4 Efficiency in Federated Learning
The primary objective of model pruning is to identify subnetworks within larger architectures by
selectively removing connections. This technique holds considerable appeal for various reasons, particularly
for real-time applications on resource-constrained edge devices, which are prevalent in federated
learning (FL) and collaborative learning scenarios. Pruning large networks can significantly alleviate
the computational demands of inference Elsen et al. (2020) or hardware tailored to exploit sparsity Cerebras
(2019); Pool et al. (2021). More recently, the lottery ticket hypothesis has emerged Frankle and Carbin
(2019), suggesting the existence of subnetworks within densely connected networks. These subnetworks,
when trained independently from scratch, can attain comparable accuracy to fully trained dense networks
Frankle and Carbin (2019), revitalizing the field of sparse deep learning Renda et al. (2020); Chen et al.
Frontiers 3
.CC-BY 4.0 International licensemade available under a
(which was not certified by peer review) is the author/funder, who has granted bioRxiv a license to display the preprint in perpetuity. It is
The copyright holder for this preprintthis version posted May 15, 2024. ; https://doi.org/10.1101/2024.05.14.594167doi: bioRxiv preprint
Bishal et al. Sparse FL for NeuroImaging
(2020). This resurgence of interest has also extended into sparse reinforcement learning (RL) Arnob et al.
(2021); Sokar et al. (2021). Pruning techniques in deep learning can broadly be categorized into three
groups: methods that induce sparsity before training and during initialization Lee et al. (2018); Wang et al.
(2020); Tanaka et al. (2020), during training Zhu and Gupta (2018); Ma et al. (2019); Yang et al. (2019);
Ohib et al. (2022), and post-training Han et al. (2015); Frankle et al. (2021).
For pruning in the FL setting, using a Lottery Ticket like approach would result in immense inefficiency
in communication. Such methods Frankle and Carbin (2019); Bibikar et al. (2022) usually require costly
pruning and retraining cycles, often training and pruning multiple times to achieve the desired accuracy
vs sparsity trade-off. Relatively few research have leveraged pruning in the FL paradigm Li et al. (2020a,
2021); Jiang et al. (2022). In particular, with LotteryFL Li et al. (2020a) and PruneFL Jiang et al. (2022),
clients need to send the full model to the server regularly resulting in higher bandwidth usage. Moreover,
in Li et al. (2020a), each client trains a personalized mask to maximize the performance only on the local
data. A few recent works Bibikar et al. (2022); Huang et al. (2022); Qiu et al. (2022); Li et al. (2020a)
also attempted to leverage sparse training within the FL setting as well. In particular Li et al. (2020a)
implemented randomly initialized sparse mask, FedDST Bibikar et al. (2022) built on the idea of RigL Evci
et al. (2020) and mostly focussed on magnitude pruning on the server-side resulting in similar constraints
and Ohib et al. (2023) uses sparse gradients to efficiently train in a federated learning setting. In this work,
we try to alleviate these limitations which we discuss in the following section.
2 METHOD DESCRIPTION
In this section we present our proposed method. We first describe the process of discovering a sub-network
f(θ ⊙ m) within the full network f(θ), where θ, m ∈ Rd with ∥m∥0 < d . To discover a performant
sub-network an importance scoring metric is required, which we describe in section 2.1.1. Finally, we
delineate our proposed method in section 2.2
2.1 Sub-network discovery
Given a dataset D = {(xi, yi)}n
i=1 at a site k, the training of a neural network f parameterized by θ ∈ Rd
can be written as minimizing the following empirical risk:
argmin
θ
1
n
X
i
L(f(θ; xi), yi) s.t. θ ∈ H (3)
where θ ∈ Rd and L and H are the loss function and the constraint set respectively.
In general, in unconstrained (standard) training the set of possible hypotheses is considered to beH = Rd,
where d is the model dimension. The objective is to minimize the empirical risk L given a training set
{(xi, yi)}n
i=1 ∼ D
at the local client site k. Given access to the gradients of the empirical risk on a batch-
wise basis, an optimization algorithm such as Stochastic Gradient Descent (SGD) is typically employed to
achieve the specified objective. This process generates a series of parameter estimates, {θi}T
i=0
, where θ0
represents the initial parameters and θT the final optimal parameters. A sub-network within this network is
defined as a sparse version of this network with a mask m ∈ {0, 1}|θ| that results in a masked network
f(θ ⊙ m; xi). When aiming for a target sparsity level where k < d , the parameter pruning challenge
entails ensuring that the final optimal parameters, θT , have at most k non-zero elements, as denoted by
the constraint ∥θT ∥0 ≤ k. In many works, this sparsity constraint applies only to the final parameters and
Frontiers 4
.CC-BY 4.0 International licensemade available under a
(which was not certified by peer review) is the author/funder, who has granted bioRxiv a license to display the preprint in perpetuity. It is
The copyright holder for this preprintthis version posted May 15, 2024. ; https://doi.org/10.1101/2024.05.14.594167doi: bioRxiv preprint
Bishal et al. Sparse FL for NeuroImaging
not to any intermediate parameter estimates. However, in this work we maintain this sparsity constraint
throughout the entire training phase, that is throughout the entire evolution of θ from θ0 to θT .
The goal of discovering sub-networks at initialization introduces additional constraints to the previously
described framework by requiring that all parameter iterations fall within a predetermined subspace of
H. Specifically, the constraints seek to identify an initial set of parameters, θ0, that has no more than
k1 non-zero elements ( ∥θ0∥0 ≤ k1), and ensure that all intermediate parameter sets, θi, belong to a
subspace ¯H ⊂ H for all i in {1, . . . , T}, where ¯H is the subspace of Rd spanned by the natural basis
vectors {ej}j∈supp(θ0). Here, supp(θ0) represents the support of θ0, or the set of indices corresponding to
its non-zero entries. This approach not only specifies a sub-network at initialization with k parameters but
also maintains its structure consistently throughout the training.
2.1.1 Connection Importance Criterion
Lee et al. (2018) introduced a technique inspired by the saliency criterion originally proposed by
Mozer and Smolensky (1988). They contributed an important insight, demonstrating that this criterion is
remarkably effective in predicting the significance of each connection in a neural network at the initialization
phase. The core concept revolves around retaining those parameters that, when altered, would have the most
substantial effect on the loss function. This is operationalized by considering a binary vector c ∈ {0, 1}m
and utilizing the Hadamard product ⊙. Consequently, SNIP calculates the sensitivity of connections based
on this approach. For each connection θj in the network, the importance score for that connection is
calculated as following:
gj(θ; D) = ∂L(θ ⊙ c)
∂c
c=1
= ∂L(θ)
∂θ ⊙ θ (4)
After determining g(θ), the parameters associated with the highestk magnitudes of |gj(θ; D)| are retained.
Essentially, SNIP prioritizes weights that, regardless of their direction, are distant from the origin and yield
large gradient values. It’s noteworthy that the objective of SNIP can be reformulated as follows De Jorge
et al. (2020); Frankle et al. (2021):
max
c
S(θ, c) :=
X
i∈supp(c)
|θi∇L(θ)i| s.t. c ∈ {0, 1}m, ∥c∥0 = k. (5)
It is trivial to note that the optimal solution to the above problem can be obtained by selecting the indices
corresponding to the top-k values of |θi ∇L(θj)|.
2.2 Proposed Method
We propose a novel method for efficient distributed sub-network discovery for distributed neuroimaging
and propose a method for training such sparse models or subnetworks in a communication efficient
manner called Sparse Federated Learning for Neuro-Imaging or NeuroSFL with the goal of tackling
communication inefficiency during decentralized federated learning with non-IID data distribution in the
context of distributed neuro-imaging data. The proposed method initiates with the common initialization
θ0 at all the local client models. Next, importance scores sj are calculated for each model parameter in the
network based on the information from the imaging data available across all the client sites. At this stage,
each client has a unique set of importance scores for their parameters in the local network f based on the
local data available at that site similar to Lee et al. (2018); De Jorge et al. (2020). All the clients transmit
Frontiers 5
.CC-BY 4.0 International licensemade available under a
(which was not certified by peer review) is the author/funder, who has granted bioRxiv a license to display the preprint in perpetuity. It is
The copyright holder for this preprintthis version posted May 15, 2024. ; https://doi.org/10.1101/2024.05.14.594167doi: bioRxiv preprint
Bishal et al. Sparse FL for NeuroImaging
these scores to each other and a mask m is created corresponding to the top-k % of the aggregated saliency
scores:
m = Tk(
K−1X
k=0
sk) (6)
Where the Tk is the top-k operator. This mask is then used to train the model fk(θ ⊙ m; x) at site k on
their local data (x, y) ∼ D k.
For the federated training among a total of K clients, the clients are trained locally, and at the end of
local training they share their trained parameters which are then averaged; we call this a communication
round. At the start of this local training, each site k starts with the same initial model weights θ0 which
at each site k is denoted as θk,0 at training step t = 0 which are then masked with the generated saliency
mask m to produce the common masked initialization θm
k,0 as follows:
θm
k,0 = θk,0 ⊙ m
Next these models at each site k are trained on their local dataset (x, y) ∼ D k.
The masked models f(θk,0 ⊙ m) across all the sites are trained for a total of T communication rounds to
arrive at the final weights θk,T at each local site. In each communication round t, only a random subset
F ′ = {f1, f2, ..., fK′} of K′ clients where F ′ ⊆ F the set of all clients, and K′ ≤ K are trained on
their local data (a way to avoid the straggler effect in the real world, wherein a large client group the
update might be bottle-necked by the slowest, most resource-constrained client site). At the end of local
training, the selected subset F ′ of the clients contribute to training the updated model θm
k,t, which is the
starting weights for the next communication round. When sharing the updated weights only the weights
corresponding to the 1’s in the binary maskm are shared among the clients and with the server, as only
these weights are being trained and the rest of the weights are zero-ed out. This results in the gains in
communication efficiency. To efficiently share the model weights, the clients only share their sparse masked
weights θm
F ′ = θF ′ ⊙ m among the selected clients in F ′ using the compressed sparse row (CSR) encoding.
The algorithm for the training process is delineated in Algorithm 1.
2.3 Iterative Importance Score Calculation: IterativeSNIP
In this section, we test the effectiveness of iterative-SNIP De Jorge et al. (2020), which is an iterative
version of the application of saliency criterion in eq. (4). We briefly describe the iterative-SNIP next. We
assume k to be the number of parameters to be preserved post pruning. Given that we have some pruning
schedule (similar to learning rate schedule: linear, exponential etc.) to dividek into a set of natural numbers
{kt}T
t=1 such that kt > k t+1 and kT = k. Now, given the binary masking variable ct corresponding to
kt, the formulation of pruning from kt to kt+1 can be made using the connection sensitivity (4) similar to
De Jorge et al. (2020) as:
ct+1 = argmax
ˆθ,c
S( ¯θ, c) s.t. c ∈ {0, 1}m, ∥c∥0 = kt+1, c ⊙ ct = c, (7)
where ¯θ = θ ⊙ ct. The constraint c ⊙ ct = c is added to ensure that no previously pruning parameter
is re-activated. Assuming that the pruning schedule ensures a smooth transition from one topology to
Frontiers 6
.CC-BY 4.0 International licensemade available under a
(which was not certified by peer review) is the author/funder, who has granted bioRxiv a license to display the preprint in perpetuity. It is
The copyright holder for this preprintthis version posted May 15, 2024. ; https://doi.org/10.1101/2024.05.14.594167doi: bioRxiv preprint
Bishal et al. Sparse FL for NeuroImaging
another (∥ct∥0 ≈ ∥ct+1∥0) such that the gradient approximation ∂L( ¯θ)
∂ ¯θ
ct
≈ ∂L( ¯θ)
∂ ¯θ
ct+1
is valid, (7) can
be approximated as solving (5) at ¯θ.
In the scenario where the schedule parameter is set to T = 1 , the original SNIP saliency method is
recovered. This is basically employing a gradient approximation approach between the initial dense
network c0 = 1 and the resulting mask c. We conduct experiments with IterativeSNIP in the federated
neuroimaging setting and present our findings in section 4.2.
Algorithm 1 NeuroSFL
Input: Total number of clients K; Total communication rounds T
Output: Sparse local models ˆθC
m
1: Initialize local models with θ0 and transmit to all clients.
2: m = Tk(PK−1
k=0 Sk(θ0)) # generate global mask from importance scores
3: Transmit m to all the sites.
4: θm
k,0 ← θk,0 ⊙ m #apply the mask at all sites k = 1, 2, ..., K
5: for t = 0 to T − 1 do
6: {ci} ˜K
i=1 # Sample a set of ˜K clients uniformly from the set of all clients
7: for site k in parallel for all ˜K clients do
8: θm
k,t ← csr(θm
k,t); #Gather all masked weights θk,m where k ∈ {1, 2, 3, . . . ˜K}
9: ˆθm
k,t ←
1
˜K
P ˜K−1
k=0 θm
k
#combine the weights of models in the selected sites
10: for τ = 0 to N − 1 do
11: Sample a batch of data δk,t,τ from the local dataset.
12: gm
k,t,τ ← ∇ θL( ˆθm
k,t,τ ; δk,t,τ ) ⊙ m # calculate and mask gradients
13: ˆθm
k,t,τ +1 ← ˆθm
k,t,τ − ηgm
k,t,τ # take optimization step with masked gradients on masked weights
14: end for
15: transmit the non-zero elements of the updated model ˆθm
k,t,N −1 back to all clients.
16: end for
17: end for
3 EXPERIMENTS
3.1 Dataset and non-IID partition
We evaluated NeuroSFL on the ABCD dataset. ABCD study is the largest long-term study of brain
development and child health in the US. It recruited over 10 thousand children of 9 and 10 years old
from 21 sites and followed them for 10 years with annual behavioral and cognitive assessments and
biannual MRI scans Garavan et al. (2018). Along with multi-session brain MRI scans for structure and
function, the ABCD study also includes key demographic information including gender, racial information,
socio-economic backgrounds, cognitive development, and mental and physical health assessments of the
subjects. The ABCD open-source dataset can be found on the National Institute of Mental Health Data
Archive (NDA) https://nda.nih.gov/. In this study, we used data from the ABCD baseline, which
contain 11,875 participants aged 9-10 years.
T1-weighted MRI images were preprocessed using the Statistical Parametric Mapping 12 (SPM12)
software toolbox for registration, normalization, and tissue segmentation. Then the gray matter density maps
Frontiers 7
.CC-BY 4.0 International licensemade available under a
(which was not certified by peer review) is the author/funder, who has granted bioRxiv a license to display the preprint in perpetuity. It is
The copyright holder for this preprintthis version posted May 15, 2024. ; https://doi.org/10.1101/2024.05.14.594167doi: bioRxiv preprint
Bishal et al. Sparse FL for NeuroImaging
were smoothed by a 6 mm3 Gaussian kernel, creating images with the dimensionality of (121,145,121) of
voxels at Montreal Neuroimaging Institute (MNI) space with each voxel having dimensions of 1.5×1.5×1.5³
mm.
We simulated the heterogeneous data distributions across federated clients through the adoption of two
distinct data partitioning strategies. We outline these strategies for generating non-IID data partitions with
a comprehensive discussion in section 3.1.1.
3.1.1 Generating non-IID data partition with Dirichlet Distribution
In this section, we provide the necessary background on generating non-identical data distribution in the
client sites using the Dirichlet Distribution, specifically for the context of federated learning.
3.1.1.1 non-IID data in FL
Federated Learning (FL), as introduced by McMahan and Ramage (2017), is a framework designed
for training models on decentralized data while preserving privacy. It utilizes the Federated Averaging
(FedAvg) algorithm where each device, or client, receives a model from a central server, performs stochastic
gradient descent (SGD) on its local data, and sends the models back for aggregation. Unlike data-center
training where data batches are often IID (independent and identically distributed), FL typically deals with
non-IID data distributions across different clients. Hence, to evaluate federated learning it is crucial to not
make the IID assumption and instead generate non-IID data among clients for evaluation Hsu et al. (2019).
3.1.1.2 Generating non-IID data from Dirichlet Distribution
In this study, we assume that each client independently chooses training samples. These samples are
classified into N distinct classes, with the distribution of class labels governed by a probability vector
q, which is non-negative and whose components sum to 1, that is, qi > 0, i ∈ [1, N] and ∥q∥1 = 1. For
generating a group of non-identical clients, q ∼ Dir(αp) is drawn from the Dirichlet Distribution, with p
characterizing a prior distribution over the N classes and α controls the degree of identicality among the
existing clients and is known as the concentration parameter.
In this section, we generate a range of client data partitions from the Dirichlet distribution with a range
of values for the concentration parameter α for exposition. In fig. 1, we generate a group of 10 balanced
clients, each holding an equal number of total samples. Similar to Hsu et al. (2019) the prior distribution p
is assumed to be uniform across all classes. For each client, given a concentration parameter α, we sample
a q from Dir(α) and allocate the corresponding fraction of samples from each client to that client. fig. 1
illustrates the effect of the concentration parameter α on the class distribution drawn from the Dirichlet
distribution on different clients, for the CIFAR-10 dataset. When α → ∞, identical class distribution is
assigned to each classes. With decreasing α, more non-identicalness is introduced in the class distribution
among the client population. At the other extreme with α → 0, each class only consists of one particular
class.
3.2 Architecture, Hyperparameters and experimental details
Here we provide a comprehensive overview of the architecture, hyperparameters, and the experimental
setup we use to evaluate our proposed NeuroSFL method on the neuroimaging Adolescent Brain Cognitive
Development (ABCD) data. Our study focuses on the task of classifying a participant’s sex based on
MRI scans, by employing a 3D variant of the well-known AlexNet model Krizhevsky et al. (2012). The
Frontiers 8
.CC-BY 4.0 International licensemade available under a
(which was not certified by peer review) is the author/funder, who has granted bioRxiv a license to display the preprint in perpetuity. It is
The copyright holder for this preprintthis version posted May 15, 2024. ; https://doi.org/10.1101/2024.05.14.594167doi: bioRxiv preprint
Bishal et al. Sparse FL for NeuroImaging
(a) (b) (c) (d) (e)
Figure 1. Generating non-identical client data partitions using the Dirichlet Distribution for the Cifar10
dataset among 10 clients. Distribution among classes is represented using different colors. (a) Dirichlet,
α → ∞ results in identical clients (b-d) Client distributions generated from Dirichlet distributions with
different concentration parameters α (e) Dirichlet, α → 0.0 results in each client being assigned only 1
particular class.
3D variant was referenced from Abrol et al. (2021), which has a specific channel configuration for the
convolutional layers set as: 64C-128C-192C-192C-128C, where ’C’ denotes channels.
We optimized the learning rate for this task through an exhaustive search ranging from LR=1 × 10−3
to 1 × 10−6, achieving a delicate balance between rapid convergence and fine-tuning during training. We
employed a batch size of 32 and a learning rate decay factor of 0.998 was applied. We applied varying
sparsity levels, ranging from 0%, 50%, 80%, 90%, and 95% to assess the overall performance. A split of
80/20 was used for training and testing for each individual site. Our training consists of 5 epochs with 200
communication rounds.
3.3 Baselines
We compared our method with both centralized and decentralized baselines. Centralized baseline includes
FedAvg McMahan et al. (2017), FedAvg-FT Cheng et al. (2021) which are the standard dense baselines,
and for the decentralized FL setting, we take the sparse Dis-PFL Dai et al. (2022).
In FedAvg McMahan et al. (2017), each client trains its local model using its local data, and then these
local models are aggregated or averaged to update the global model. On the other hand, FedAvg-FT Cheng
et al. (2021) extends the FedAvg algorithm by incorporating fine-tuning or transfer learning. Specifically,
after the global model is trained using FedAvg, the global model is then fine-tuned or adapted using
additional data from a central server or other external sources. This fine-tuning step allows the global model
to adapt to new tasks or data distributions beyond what was initially learned from the federated learning
process. We also compare with DisPFL Dai et al. (2022) with varying sparsity levels. DisPFL is a new
sparse FL technique that randomly prunes each layer similar to Evci et al. (2020) and uses the prune and
regrow method from that work as well, resulting in a dynamically sparse method.
In exploring the impact of using unique local masks instead of a global mask on the performance of
FL, we established IndividualSNIP as a baseline, representing an approach where unique local masks are
devised from the saliency criterion, and local models are trained based on these masks. Moreover, to isolate
the impact of just using global masking, that is using the same random mask in all clients, instead of using
different unique random masks at different sites we compare our method and competing methods against
random global masking as well in fig. 2.
Additionally, with further exploration to investigate how different methods of model pruning and selection
impact the performance of our approach, we further experiment with other techniques named IterSNIP and
WeightedSNIP. IterSNIP builds upon the traditional SNIP method Lee et al. (2018) by incorporating multiple
minibatches during the training process of mask generation. This approach aggregates saliency scores from
Frontiers 9
.CC-BY 4.0 International licensemade available under a
(which was not certified by peer review) is the author/funder, who has granted bioRxiv a license to display the preprint in perpetuity. It is
The copyright holder for this preprintthis version posted May 15, 2024. ; https://doi.org/10.1101/2024.05.14.594167doi: bioRxiv preprint
Bishal et al. Sparse FL for NeuroImaging
these mini-batches to generate a comprehensive and robust pruning mask. Conversely, WeightedSNIP
adopts a different strategy, deriving a global mask through a weighted average of saliency scores based on
the frequency of data at each site, and assigning importance levels to individual sites based on the amount
of data at sites.
4 RESULTS AND DISCUSSION
4.1 Effect of varying sparsity levels
We first explore the effect of sparsity on IID data in section 4.1.1 and then explore the efficacy of
NeuroSFL on non-IID data in section 4.1.2.
4.1.1 Effect of varying sparsity levels on IID data
The performance of various methods across different sparsity levels was evaluated, as presented in
Table 1, and visually presented in Figure 2. Sparse baselines, including Ours (NeuroSFL), IndividualSNIP,
DisPFL Dai et al. (2022), and Global Random Mask, were compared against dense baselines such as
FedAvg-FT Cheng et al. (2021) and FedAvg McMahan et al. (2017). Notably, our proposed NeuroSFL,
exhibited robust performance across varying sparsity levels, achieving an accuracy of 92.52% at 0% sparsity
and maintaining high accuracy even at higher sparsity levels, with 71.18% accuracy at 95% sparsity. In
comparison, IndividualSNIP demonstrated decreasing accuracy as sparsity increased, with a significant
drop to 52.70% at 95% sparsity. This is in line with expectation as individual-SNIP only incorporates the
saliency scores from a single site at random and does not incorporate information from the datasets at all
the participating cites.
Moreover, in contrast to NeuroSFL, DisPFL, and Global Random Mask also showcased diminishing
accuracy with increasing sparsity, highlighting the effectiveness of our proposed approach in mitigating the
adverse effects of sparsity on model performance on neuroimaging data. Notably, Global Random Mask
outperformed DisPFL on lower sparsities, suggesting that in general global random masks might be more
suitable for federated applications compared to even targeted unique local masks which DisPFL employs.
Dense baselines, such as FedAvg-FT and FedAvg, even while being not sparse and using full
communication achieved comparable performances to NeuroSFL in the non-extreme sparsity region.
Figure 2. (a) Comparison of methods for gender classification using MRI Scans of ABCD dataset. (b)
Gender differences in each of the 21 ABCD sites along with the performance of the global model on each
site with 50% sparsity.
Frontiers 10
.CC-BY 4.0 International licensemade available under a
(which was not certified by peer review) is the author/funder, who has granted bioRxiv a license to display the preprint in perpetuity. It is
The copyright holder for this preprintthis version posted May 15, 2024. ; https://doi.org/10.1101/2024.05.14.594167doi: bioRxiv preprint
Bishal et al. Sparse FL for NeuroImaging
References
Abrol, A., Fu, Z., Salman, M., Silva, R., Du, Y ., Plis, S., et al. (2021). Deep learning encodes
robust discriminative neuroimaging representations to outperform standard machine learning. Nature
Communications 12. doi:10.1038/s41467-020-20655-6
Arnob, S. Y ., Ohib, R., Plis, S., and Precup, D. (2021). Single-shot pruning for offline reinforcement
learning. arXiv preprint arXiv:2112.15579
Bibikar, S., Vikalo, H., Wang, Z., and Chen, X. (2022). Federated dynamic sparse training: Computing
less, communicating less, yet learning better. In Proceedings of the AAAI Conference on Artificial
Intelligence. vol. 36, 6080–6088
Bonawitz, K., Eichner, H., Grieskamp, W., Huba, D., Ingerman, A., Ivanov, V ., et al. (2019). Towards
federated learning at scale: System design. Proceedings of machine learning and systems 1, 374–388
[Dataset] Cerebras (2019). Wafer Scale Engine: Why We Need Big Chips for Deep
Learning. https://cerebras.net/blog/cerebras-wafer-scale-engine-why-we-
need-big-chips-for-deep-learning/
Chen, T., Frankle, J., Chang, S., Liu, S., Zhang, Y ., Wang, Z., et al. (2020). The lottery ticket hypothesis
for pre-trained bert networks. Advances in neural information processing systems 33, 15834–15846
Cheng, G., Chadha, K., and Duchi, J. (2021). Fine-tuning is fine in federated learning. arXiv preprint
arXiv:2108.07313
Dai, R., Shen, L., He, F., Tian, X., and Tao, D. (2022). Dispfl: Towards communication-efficient
personalized federated learning via decentralized sparse training. arXiv preprint arXiv:2206.00187
De Jorge, P., Sanyal, A., Behl, H. S., Torr, P. H., Rogez, G., and Dokania, P. K. (2020). Progressive
skeletonization: Trimming more fat from a network at initialization. arXiv preprint arXiv:2006.09081
Elsen, E., Dukhan, M., Gale, T., and Simonyan, K. (2020). Fast sparse convnets. In Proceedings of the
IEEE/CVF conference on computer vision and pattern recognition. 14629–14638
Evci, U., Gale, T., Menick, J., Castro, P. S., and Elsen, E. (2020). Rigging the lottery: Making all tickets
winners. In International Conference on Machine Learning (PMLR), 2943–2952
Frankle, J. and Carbin, M. (2019). The lottery ticket hypothesis: Finding sparse, trainable neural networks.
In 7th International Conference on Learning Representations, ICLR 2019, New Orleans, LA, USA, May
6-9, 2019
Frankle, J., Dziugaite, G. K., Roy, D., and Carbin, M. (2021). Pruning neural networks at initialization:
Why are we missing the mark? In 9th International Conference on Learning Representations, ICLR
2021, Virtual Event, Austria, May 3-7, 2021
Garavan, H., Bartsch, H., Conway, K., Decastro, A., Goldstein, R., Heeringa, S., et al. (2018). Recruiting
the abcd sample: Design considerations and procedures. Developmental cognitive neuroscience 32,
16–22
Frontiers 15
.CC-BY 4.0 International licensemade available under a
(which was not certified by peer review) is the author/funder, who has granted bioRxiv a license to display the preprint in perpetuity. It is
The copyright holder for this preprintthis version posted May 15, 2024. ; https://doi.org/10.1101/2024.05.14.594167doi: bioRxiv preprint
Bishal et al. Sparse FL for NeuroImaging
Han, S., Mao, H., and Dally, W. J. (2015). Deep compression: Compressing deep neural networks with
pruning, trained quantization and huffman coding. arXiv preprint arXiv:1510.00149
Hsu, T.-M. H., Qi, H., and Brown, M. (2019). Measuring the effects of non-identical data distribution for
federated visual classification. arXiv preprint arXiv:1909.06335
Huang, T., Liu, S., Shen, L., He, F., Lin, W., and Tao, D. (2022). Achieving personalized federated learning
with sparse local models. arXiv preprint arXiv:2201.11380
Jiang, Y ., Wang, S., Valls, V ., Ko, B. J., Lee, W.-H., Leung, K. K., et al. (2022). Model pruning enables
efficient federated learning on edge devices. IEEE Transactions on Neural Networks and Learning
Systems
Koneˇcn`y, J., McMahan, H. B., Ramage, D., and Richt´arik, P. (2016). Federated optimization: Distributed
machine learning for on-device intelligence. arXiv preprint arXiv:1610.02527
Krizhevsky, A., Sutskever, I., and Hinton, G. E. (2012). Imagenet classification with deep convolutional
neural networks. In Advances in Neural Information Processing Systems, eds. F. Pereira, C. Burges,
L. Bottou, and K. Weinberger (Curran Associates, Inc.), vol. 25
Kulkarni, V ., Kulkarni, M., and Pant, A. (2020). Survey of personalization techniques for federated
learning. In 2020 Fourth World Conference on Smart Trends in Systems, Security and Sustainability
(WorldS4) (IEEE), 794–797
Laird, A. R. (2021). Large, open datasets for human connectomics research: Considerations for reproducible
and responsible data use. NeuroImage 244, 118579
Landis, D., Courtney, W., Dieringer, C., Kelly, R., King, M., Miller, B., et al. (2016). Coins data exchange:
An open platform for compiling, curating, and disseminating neuroimaging data. NeuroImage 124,
1084–1088
Lee, N., Ajanthan, T., and Torr, P. H. (2018). Snip: Single-shot network pruning based on connection
sensitivity. arXiv preprint arXiv:1810.02340
Li, A., Sun, J., Wang, B., Duan, L., Li, S., Chen, Y ., et al. (2020a). Lotteryfl: Personalized and
communication-efficient federated learning with lottery ticket hypothesis on non-iid datasets. arXiv
preprint arXiv:2008.03371
Li, A., Sun, J., Zeng, X., Zhang, M., Li, H., and Chen, Y . (2021). Fedmask: Joint computation and
communication-efficient personalized federated learning via heterogeneous masking. In Proceedings of
the 19th ACM Conference on Embedded Networked Sensor Systems. 42–55
Li, W., Milletar`ı, F., Xu, D., Rieke, N., Hancox, J., Zhu, W., et al. (2019). Privacy-preserving federated
brain tumour segmentation. In Machine Learning in Medical Imaging: 10th International Workshop,
MLMI 2019, Held in Conjunction with MICCAI 2019, Shenzhen, China, October 13, 2019, Proceedings
10 (Springer), 133–141
Li, X., Gu, Y ., Dvornek, N., Staib, L. H., Ventola, P., and Duncan, J. S. (2020b). Multi-site fmri analysis
using privacy-preserving federated learning and domain adaptation: Abide results. Medical Image
Analysis 65, 101765
Ma, R., Miao, J., Niu, L., and Zhang, P. (2019). Transformed ℓ1 regularization for learning sparse deep
neural networks. Neural Networks 119, 286–298
McMahan, B., Moore, E., Ramage, D., Hampson, S., and y Arcas, B. A. (2017). Communication-efficient
learning of deep networks from decentralized data. In Artificial intelligence and statistics (PMLR),
1273–1282
[Dataset] McMahan, B. and Ramage, D. (2017). Federated learning: Collaborative machine learning
without centralized training data
Frontiers 16
.CC-BY 4.0 International licensemade available under a
(which was not certified by peer review) is the author/funder, who has granted bioRxiv a license to display the preprint in perpetuity. It is
The copyright holder for this preprintthis version posted May 15, 2024. ; https://doi.org/10.1101/2024.05.14.594167doi: bioRxiv preprint
Bishal et al. Sparse FL for NeuroImaging
McMahan, H. B., Moore, E., Ramage, D., Hampson, S., and y Arcas, B. A. (2016). Communication-
efficient learning of deep networks from decentralized data. In International Conference on Artificial
Intelligence and Statistics
Milham, M. P., Craddock, R. C., Son, J. J., Fleischmann, M., Clucas, J., Xu, H., et al. (2018). Assessment
of the impact of shared brain imaging data on the scientific literature. Nature Communications 9, 2818
Ming, J., Verner, E., Sarwate, A., Kelly, R., Reed, C., Kahleck, T., et al. (2017). Coinstac: Decentralizing
the future of brain imaging analysis. F1000Research 6
Mozer, M. C. and Smolensky, P. (1988). Skeletonization: A technique for trimming the fat from a network
via relevance assessment. Advances in neural information processing systems 1
Ohib, R., Gillis, N., Dalmasso, N., Shah, S., Potluru, V . K., and Plis, S. (2022). Explicit group sparse
projection with applications to deep learning and NMF. Transactions on Machine Learning Research
Ohib, R., Thapaliya, B., Gaggenapalli, P., Liu, J., Calhoun, V ., and Plis, S. (2023). Salientgrads: Sparse
models for communication efficient and data aware distributed federated training. arXiv preprint
arXiv:2304.07488
Plis, S. M., Sarwate, A. D., Wood, D., Dieringer, C., Landis, D., Reed, C., et al. (2016a). Coinstac: a
privacy enabled model and prototype for leveraging and processing decentralized brain imaging data.
Frontiers in neuroscience10, 204805
Plis, S. M., Sarwate, A. D., Wood, D., Dieringer, C., Landis, D., Reed, C., et al. (2016b). COINSTAC: A
privacy enabled model and prototype for leveraging and processing decentralized brain imaging data.
Frontiers in Neuroscience10. doi:10.3389/fnins.2016.00365
Poldrack, R. A., Barch, D. M., Mitchell, J. P., Wager, T. D., Wagner, A. D., Devlin, J. T., et al. (2013).
Toward open sharing of task-based fmri data: the openfmri project. Frontiers in neuroinformatics7, 12
[Dataset] Pool, J., Sawarkar, A., and Rodge, J. (2021). Accelerating Inference with
Sparsity Using the NVIDIA Ampere Architecture and NVIDIA TensorRT. https:
//developer.nvidia.com/blog/accelerating-inference-with-sparsity-
using-ampere-and-tensorrt/
Qiu, X., Fernandez-Marques, J., Gusmao, P. P., Gao, Y ., Parcollet, T., and Lane, N. D. (2022). Zerofl:
Efficient on-device training for federated learning with local sparsity. arXiv preprint arXiv:2208.02507
Renda, A., Frankle, J., and Carbin, M. (2020). Comparing rewinding and fine-tuning in neural network
pruning. arXiv preprint arXiv:2003.02389
Rootes-Murdy, K., Gazula, H., Verner, E., Kelly, R., DeRamus, T., Plis, S., et al. (2022). Federated analysis
of neuroimaging data: a review of the field. Neuroinformatics , 1–14
Roy, A. G., Siddiqui, S., P¨olsterl, S., Navab, N., and Wachinger, C. (2019). Braintorrent: A peer-to-peer
environment for decentralized federated learning. arXiv preprint arXiv:1905.06731
Sheller, M. J., Reina, G. A., Edwards, B., Martin, J., and Bakas, S. (2019). Multi-institutional deep learning
modeling without sharing patient data: A feasibility study on brain tumor segmentation. In Brainlesion:
Glioma, Multiple Sclerosis, Stroke and Traumatic Brain Injuries: 4th International Workshop, BrainLes
2018, Held in Conjunction with MICCAI 2018, Granada, Spain, September 16, 2018, Revised Selected
Papers, Part I 4(Springer), 92–104
Silva, S., Gutman, B. A., Romero, E., Thompson, P. M., Altmann, A., and Lorenzi, M. (2019). Federated
learning in distributed medical databases: Meta-analysis of large-scale subcortical brain data. In 2019
IEEE 16th international symposium on biomedical imaging (ISBI 2019) (IEEE), 270–274
Sokar, G., Mocanu, E., Mocanu, D. C., Pechenizkiy, M., and Stone, P. (2021). Dynamic sparse training for
deep reinforcement learning. arXiv preprint arXiv:2106.04217
Frontiers 17
.CC-BY 4.0 International licensemade available under a
(which was not certified by peer review) is the author/funder, who has granted bioRxiv a license to display the preprint in perpetuity. It is
The copyright holder for this preprintthis version posted May 15, 2024. ; https://doi.org/10.1101/2024.05.14.594167doi: bioRxiv preprint
Bishal et al. Sparse FL for NeuroImaging
Tanaka, H., Kunin, D., Yamins, D. L., and Ganguli, S. (2020). Pruning neural networks without any data
by iteratively conserving synaptic flow. In Advances in Neural Information Processing Systems 33:
Annual Conference on Neural Information Processing Systems 2020, NeurIPS 2020, December 6-12,
2020, virtual, eds. H. Larochelle, M. Ranzato, R. Hadsell, M. Balcan, and H. Lin
Tenopir, C., Allard, S., Douglass, K., Aydinoglu, A. U., Wu, L., Read, E., et al. (2011). Data sharing by
scientists: practices and perceptions. PloS one 6, e21101
Wang, C., Zhang, G., and Grosse, R. B. (2020). Picking winning tickets before training by preserving
gradient flow. In 8th International Conference on Learning Representations, ICLR 2020, Addis Ababa,
Ethiopia, April 26-30, 2020
Yang, H., Wen, W., and Li, H. (2019). Deephoyer: Learning sparser neural network with differentiable
scale-invariant sparsity measures. arXiv preprint arXiv:1908.09979
Zhao, Y ., Li, M., Lai, L., Suda, N., Civin, D., and Chandra, V . (2018). Federated learning with non-iid data.
arXiv preprint arXiv:1806.00582
Zhu, M. and Gupta, S. (2018). 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
Frontiers 18
.CC-BY 4.0 International licensemade available under a
(which was not certified by peer review) is the author/funder, who has granted bioRxiv a license to display the preprint in perpetuity. It is
The copyright holder for this preprintthis version posted May 15, 2024. ; https://doi.org/10.1101/2024.05.14.594167doi: bioRxiv preprint