Federated Learning of Large Models at the Edge via Principal Sub-Model Training
Abstract
Limited compute and communication capabilities of edge users create a significant bottleneck for federated learning (FL) of large models.
We consider a realistic, but much less explored, cross-device FL setting in which no client has the capacity to train a full large model nor is willing to share any intermediate activations with the server.
To this end, we present Principal Sub-Model (PriSM) training methodology, which leverages models’ low-rank structure and kernel orthogonality to train sub-models in the orthogonal kernel space.
More specifically, by applying singular value decomposition (SVD) to original kernels in the server model, PriSM first obtains a set of principal orthogonal kernels in which each one is weighed by its singular value.
Thereafter, PriSM utilizes a novel sampling strategy that selects different subsets of the principal kernels independently to create sub-models for clients.
Importantly, a kernel with a large singular value is assigned with a high sampling probability.
Thus, each sub-model is a low-rank approximation of the full large model, and all clients together achieve the near full-model training.
Our extensive evaluations on multiple datasets in resource-constrained settings show that PriSM can yield an improved performance of up to 10% compared to existing alternatives, with only around 20% sub-model training.
Introduction
Federated Learning (FL) is emerging as a popular paradigm for distributed and privacy-preserving machine learning as it allows local clients to perform ML optimization jointly without directly sharing local data [21, 15, 19].
Thus, it enables privacy protection on local data, and leverages distributed local training to attain a better global model.
This creates opportunities for many edge devices rich in data to participate in the joint training without direct data sharing.
For example, resource-limited smart home devices can train local vision or language models using private data and achieve a server model that generalizes well to all users via federated learning [23].
Despite significant progress in FL in the recent past, several crucial challenges still remain when moving to the edge.
In particular, limited computation and communication capacities prevent clients from learning large models for leveraging vast amounts of local data at the clients.
This problem has attracted a lot of attention [8, 11, 30, 27, 9].
For example, [8, 11, 30] propose to assign clients with different subsets of server model depending on their available resources.
However, these works have an underlying assumption that some of the clients have sufficient resources to train a nearly full large model.
As a result, server model size is limited by the clients with maximum computation and communication capacities.
To overcome resource constraints on clients, prior works such as [27, 9] change the training paradigm by splitting a model onto server and clients.
The computational burden on the clients is therefore relieved as the dominant part of the burden is offloaded to the server.
However, such a methodology requires sharing of intermediate activations and/or labels with the server, which directly leaks input information and potentially compromises privacy promises of FL.
To overcome the aforementioned limitations in prior works, we focus on an even more constrained and realistic setting at the edge, in which no client is capable of training a large model nor is willing to share any intermediate data and/or labels with the server.
To this end, we propose Principal Sub-Model (PriSM) training to allow each client to only train a small sub-model, while still enabling the server model to achieve comparable accuracy as the full-model training .
We exploit low-rank structure in models during the training, which is commonly used in reducing compute costs [16, 7].
However, naive lowrank approximation in FL [30], where all clients only train top-k kernels based on their capacities, incurs a notable accuracy drop, especially in very constrained settings.
In
Figure 1, we delve into this issue by showing the number of principal kernels required in the orthogonal space to accurately approximate each convolution layer in the first two ResBlocks in ResNet-18 [10] during FL training 2 .
We observe that even at the end of the FL training, around half of the principal kernels are still needed to sufficiently approximate each convolution layer.
We have similar findings for the remaining convolution layers.
Therefore, to avoid the reduction in server model capacity, it is essential to ensure that all server-side principal kernels are collaboratively trained on clients, especially when each client can only train a very small sub-model (e.g., <50% of the server model).
Based on our above observations, PriSM employs a novel probabilistic strategy to select a subset of kernels and create a sub-model for each client as shown in Figure 2.
More specifically, PriSM first converts the model into orthogonal space where original convolution kernels are decomposed into principal kernels using singular value decomposition (SVD).
To approximate the original server model, PriSM utilizes our novel sampling process, that is based on the singular values, such that a principal kernel with a larger singular value has a higher sampling probability.
Furthermore, the probabilistic process ensures that all sub-models can together provide a near server model coverage, thus leading to the near full-model training performance.
[Figure]
We conduct extensive evaluations for PriSM on vision and language tasks under resourced-constrained settings where no client is capable of training the large full model.
Our results demonstrate that PriSM delivers consistently better performance compared to other prior works, especially when participating clients have very limited capacities.
For instance, on ResNet-18/CIFAR-10, we show that PriSM only incurs around 2% and 3% accuracy drop for i.i.d and highly non-i.i.d datasets under a very constrained setting where all clients can only train around 20% of the server model.
Compared to other solutions, PriSM improves the accuracy by up to 10% .
Method
In this section, we first motivate our proposal, Principal random Sub-Model training (PriSM), with an observation of orthogonality in convolution layers.
Then, we describe the details of PriSM.
Notations.
We denote the Frobenius norm as ∥·∥ F , and σ i as i -th singular value in matrix A .
⊛ indicates convolution operation, and · indicates simple matrix multiplication.
⟨· , ·⟩ denotes sum of element-wise multiplication or inner product.
tr ( A ) is the trace of matrix A .
2.1 Motivation: An Observation on Orthogonality
We consider a convolution layer with kernels W ∈ R N × M × k × k and input X ∈ R M × H × W , where N and M denote the number of output channels and the number of input channels, k is kernel size, and
H × W is the size of the input image along each channel.
Based on a common technique im2col [5], the convolution layer can be converted to matrix multiplication as Y = W · X , where W ∈ R N × Mkand X ∈ R Mk 2 × HW .
For kernel decorrelation, we apply singular value decomposition (SVD) to map kernels into orthogonal space as: W = (cid:80) Ni =1 σ i · u i · v Ti , where { u i } Ni =1 , { v i } Ni =1 are two sets of orthogonal vectors 3 .
The convolution can be decomposed as follows:
[Equation]
For ∀ i ̸
= j , it is easy to verify that (cid:10) Y i , Y j (cid:11)
=
σ i · σ j · tr ( X T · v i · u Ti · u j · v Tj · X ) = 0 , namely the output features Y i and Y j are orthogonal.
Therefore, if we regard W i = σ i · u i · v Ti as a principal kernel, different principal kernels create orthogonal output features.
To illustrate this, Figure 3 shows a input image (left) and the outputs (right three) generated by principal kernels.
We can observe that principal kernels captures different features and serve different purposes.
As [29, 3, 28] reveal, imposing orthogonality leads to better training performance.
This motivates us to initiate the training with a set of orthogonal kernels.
Furthermore, to preserve kernel orthogonality during training, it is critical to constantly refresh the orthogonal space through re-decomposition.
The above intuitions play an important role in PriSM, which is described in the following section.
2.2 PriSM: Principal Random Sub-Model Training
In the realistic setting considered in our paper, participating clients are very resource-limited and incapable of performing full large model training.
Therefore, to train a large server model in FL, it is essential to distribute the training workloads among clients.
One way toward this goal is to conduct sub-model training on clients.
Based on the motivation in Section 2.1, we select sub-models and sample from the orthogonal space.
Further, as we observe in Section 2.1, different principal kernels capture distinct features in the orthogonal space, and their contributions are further weighed by their corresponding singular values as shown in Eq.
Therefore, we propose a novel importance-aware sampling strategy to create sub-models for clients.
This provides two key benefits: 1) each sub-model is a low-rank approximation to the server model; 2) the conglomerate of the sampled sub-models enables a near-full server model coverage.
Each participating client trains its sub-model and uploads it to the server.
The server then aggregates the updated sub-models, obtain the full model in the original space, and then re-decompose it to refresh the orthogonal space for the next round.
Input: layer parameters W , client capacities.
1: for communication round t = 1 , · · · , T do 2:
Decompose W into orthogonal kernel using SVD → (cid:8) W i (cid:9) N i =1 .
Choose a subset of clients → C .
4: for each client c ∈ C do 5:
Compute the sub-model size for client c → |I c | .
Obtain a sub-model using random sampling based on Eq.
(2) → I c , W c .
Perform LocalTrain ↔ I c , W c . //
Local training 8: end for 9:
Aggregate parameters based on Eq.
(3) W ← (cid:110) W c (cid:111) c ∈C . //
Sub-model aggregation 10:
Reconstruct W from W . //
Orthogonal space refresh 11: end for |——————————————————————————————— 12: LocalTrain ↔ I c , W c 13: for local iteration k = 1 , · · · , K do 14:
Sample an input batch from the local dataset → D k .
Perform the forward and backward pass ← D k , W c .
Compute additional gradients from regularization based on Eq.
(4) ← reg .
Update the local sub-model using SGD → W c .
18: end for |———————————————————————————————
We describe each components of PriSM below.
Sub-model sampling.
Given a resource budget in the client, suppose it can at most process r principal kernels for a convolution layer.
For a convolution layer with principal kernels (cid:8) W i (cid:9) N i =1 , and the corresponding singular values { σ i } Ni =1 , we randomly sample r principal kernels denoted by
I c without replacement with sampling probability for i -th kernel as follows:
[Equation]
Here, κ in Eq.
(2) is a smooth factor that controls the probability distribution for principal kernels to be chosen.
Therefore, with our proposed stochastic sampling strategy, important kernels with large singular values are more likely to be chosen, and all sub-models together provide a near-full model coverage.
Other element-wise layers such as ReLU and batch normalization remain the same.
Local training.
On each client, when performing optimization, the selected { u i , v i } i ∈I c are optimized, together with trainable parameters in other layers.
While in PriSM, singular values { σ i } i ∈I c are not updated during local training.
This is because each singular value indicates importance of each principal kernel.
Thus, freezing singular values across clients helps maintain consistency regarding importance for kernels and is useful in the aggregation of the sub-models as described next.
Sub-model aggregation.
On the server side, with sub-models obtained from clients, we first aggregate i -th principal kernel as follows:
[Equation]
where C denotes the subset of active clients, α i is the aggregation coefficient for i -th kernel.
We propose a weighted averaging scheme: if i -th kernel is selected and trained by C i clients, then α ci = 1 /C i .
Furthermore, for principal kernels that are not sampled during creation of sub-models, they remain intact during aggregation.
Orthogonal space refresh.
The full model in the original space is constructed by converting each 2-dimensional W i to the original dimension R M × k × k and combining them.
For the next communication round, the orthogonal space is refreshed by decomposing the updated W using SVD.
We further use two additional techniques to improve learning efficiency in the orthogonal space: activation normalization, and regularization on orthogonal kernels.
Activation normalization.
We apply batch normalization without tracking running statistics; namely, the normalization always uses current batch statistics in the training and evaluation phases.
Each client applies normalization separately with no sharing of statistics during model aggregation.
Such an adaptation is effective in ensuring consistent outputs between different sub-models and avoids potential privacy leakage through the running statistics [2].
Hence, it has been used in several sub-model FL training schemes [8, 30].
Regularization.
When learning a sampled subset of principal kernels on a client, naively applying weight decay to u i and v i separately results in poor final accuracy.
Inspired by [17], for training on client c , we add regularization to the subset of kernels as follows:
[Equation]
where λ is the regularization factor, I c denotes the subset of principal kernels on client c .
Algorithm 1 presents an overall description of PriSM.
We only show the procedure on a single convolution layer with kernels W for the sake of simplifying notations.
In the following remarks, we differentiate PriSM from Dropout and Low-Rank compression.
Remark 2.1 .
PriSM vs Dropout.
PriSM shares some computation similarity with model training using regular dropout [26].
However, regular dropout on the original kernels leads to significant convergence instability, especially with a high dropout probability [11].
In contrast, PriSM performs importance-aware sampling based on singular values.
Therefore, each sub-model is an approximation to the full model, and training different sub-models on clients does not create significant inconsistency.
Remark 2.2 .
PriSM vs Low-Rank Compression.
While PriSM exploits low-rank properties in models, it is not a low-rank compression method.
Low-rank compression methods aim to construct a smaller server model by completely discarding some kernels even though they can still contribute to model performance.
PriSM randomly select sub-models so that every kernel is possible to be learned.
Experiments
We evaluate PriSM under resourced-constrained settings where no clients can train the large full model.
Specifically, we consider homogeneous settings with all clients having the same limited compute and communication capacity.
We also compare PriSM with two other baselines: ordered kernel dropout in orthogonal space (OrthDrop), such as in [30]; and ordered kernel dropout in original space (OrigDrop), such as in [11, 8].
At a high level, our results demonstrate that PriSM achieves comparable server model accuracy even when only training very small sub-models on clients.
Baselines.
Prior methods such as FjORD [11] and HeteroFL [8] select sub-models from the original kernel space, for which we denote as OrigDrop.
On the other hand, we use OrthDrop to denote selecting fixed top-k principal kernels from the orthogonal space such as in FedHM [30].
Models and Datasets.
We train ResNet-18 on CIFAR-10 [18], CNN on FEMNIST [4] and LSTM model on IMDB [20] 4 .
ResNet-18 is optimized for CIFAR-10, where kernel size in the first convolution layer is reduced to 3 × 3 .
CNN is a small model with two convolution layers.
Data Distribution.
For CIFAR-10 and IMDB, we create balanced datasets during training with FL.
Given the total number of samples and participating clients, we uniformly sample the equal number of training images for each client when creating i.i.d datasets.
For non-i.i.d datasets, we first use
Dirichlet function Dir ( α ) [25] to create sampling probability for each client and then sample an equal number of training images for clients.
We create two different non-i.i.d datasets with α = 1 and α = 0 .
1 , where a smaller α denotes a higher degree of non-i.i.d.
For FEMNIST [4], we directly use the dataset without any additional preprocessing.
FL Setting.
We simulate an FL setting with 100 clients, and 20 clients are chosen uniformly at random in each communication round.
Each client trains its model for 2 local epochs and then uploads it to the aggregation server.
We use SGD with momentum in the local training.
Further details are provided in Appendix A.1.2.
Performance on homogeneous clients .
In this setting, we assume that all clients have the same limited compute and communication capacity.
We vary the client sub-model size from 0 .
2 to 0 .
8 of the full server model, where 0 .
× indicates only a 0 .
× subset of the principal kernels are sampled in each convolution layer from the server model (denoted as keep ratio in the results).
In Table 1, we list the computational and communication footprints for different sub-models of ResNet-18.
For
OrthDrop, we follow the strategy in [30] and select the top 0 .
× principal kernels for all clients.
For
OrigDrop, we select the first 0 .
× original kernels as in [8, 11].
[Table]
Figure 4 shows final validation accuracy of ResNet-18 with different sampled sub-model sizes on i.i.d and non-i.i.d ( α = 1 , 0 . 1 ) local datasets.
We note that PriSM constantly delivers better performance than the other two baselines.
The performance gap is even more striking under very constrained settings.
For instance, when only 0 .
2 × sub-models are supported on clients, PriSM attains comparable accuracy as full-model training, and achieves up to 10% performance improvement compared to
OrthDrop on non-i.i.d dataset with α = 0 .
Furthermore, we make two key observations.
First, training with sub-models in the orthogonal space (OrthDrop) provides better performance than in the original space (OrigDrop), which aligns with our intuition in Section 2.1.
Second, our importanceaware sampling strategy for creating sub-models is indispensable as demonstrated by the notable performance gap between PriSM and OrthDrop.
Conclusion
We have considered the practical, yet under-explored, problem of federated learning in a resourceconstrained edge setting, where no participating client has the capacity to train a large model.
As our main contribution, we propose the PriSM training methodology, that empowers the resource-limited clients by enabling them to train smaller sub-models.
At the same time, PriSM utilizes a novel sampling approach to obtain sub-models for the clients, all of which together ensure that the server model achieves close to the full-model performance.
Our extensive empirical results demonstrate that
PriSM performs significantly better than the prior baselines, especially when each client can train only a very small sub-model.
In particular, when each client is required to train a sub-model that is only around 20% in size of the server model, we demonstrate that PriSM achieves a performance advantage of up to 10% over the prior baselines.
Acknowledgments and Disclosure of Funding
This material is based upon work supported by Defense Advanced Research Projects Agency (DARPA) under Contracts No.
FASTNICS HR001120C0088 and HR001120C0160, NSF grants
CCF-1763673, CNS-2002874, ARO grant W911NF-22-1- and gifts from Intel, Qualcomm and
Cisco.
Checklist
For all authors... (a) Do the main claims made in the abstract and introduction accurately reflect the paper’s contributions and scope?
[Yes] (b) Did you describe the limitations of your work?
[No] .
(c) Did you discuss any potential negative societal impacts of your work?
[No] .
(d) Have you read the ethics review guidelines and ensured that your paper conforms to them?
[Yes] 2.
If you are including theoretical results... (a) Did you state the full set of assumptions of all theoretical results?
[N/A] (b) Did you include complete proofs of all theoretical results?
[N/A] 3.
If you ran experiments... (a) Did you include the code, data, and instructions needed to reproduce the main experimental results (either in the supplemental material or as a URL)?
[No] We will open-source our code upon acceptance of the paper.
(b) Did you specify all the training details (e g., data splits, hyperparameters, how they were chosen)?
[Yes] The training hyperparameters are provided in the Appendix.
(c) Did you report error bars (e g., with respect to the random seed after running experiments multiple times)?
[Yes] (d) Did you include the total amount of compute and the type of resources used (e.g., type of GPUs, internal cluster, or cloud provider)?
[Yes] Detailed information is provided in the Appendix. 4.
If you are using existing assets (e.g., code, data, models) or curating/releasing new assets... (a) If your work uses existing assets, did you cite the creators?
[Yes] (b) Did you mention the license of the assets?
[N/A] (c) Did you include any new assets either in the supplemental material or as a URL?
[Yes] (d) Did you discuss whether and how consent was obtained from people whose data you’re using/curating?
[Yes] (e) Did you discuss whether the data you are using/curating contains personally identifiable information or offensive content?
[No] All datasets used are well-known open datasets.
If you used crowdsourcing or conducted research with human subjects... (a) Did you include the full text of instructions given to participants and screenshots, if applicable?
[N/A] (b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, if applicable?
[N/A] (c) Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation?
[N/A]
A Appendix A.1 Models and Hyperparameters
In this section, we provide detailed information about models and hyperparameter settings for the results presented in the paper.
We will open-source our code upon acceptance of the paper.
A.1.1 Models
ResNet-18/CIFAR-10.
We use a ResNet-18 optimized for CIFAR-10, in which kernel size in the first convolution layer is changed from 7 × 7 to 3 × 3 .
Details are shown in Table 2. [Table]
Classification
[Table]
CNN/FEMNIST.
We use a similar architecture as in FjORD [11].
The detailed model is shown in
Table 3.
[Table]
LSTM/IMDB.
We use a common LSTM model as shown in Table 4.
[Table]
A.1.2 Training Hyperparameters
ResNet-18/CIFAR-10 on homogeneous clients.
We simulate 100 clients during FL training, in which each client is assigned 500 training samples for both i.i.d and non-i.i.d datasets.
In each communication round, each client performs local training for 2 epochs using the local data, then uploads parameters to the server for aggregation.
Table 5 lists detailed hyperparameters during FL training with ResNet-18. [Table]
CNN/FEMNIST on homogeneous clients.
We simulate 100 clients during FL training, in which each client is assigned 10 users’ data from the original training dataset.
We use the whole validation dataset to compute the validation accuracy.
Table 6 lists detailed hyperparameters during FL training with CNN. [Table]
LSTM/IMDB on homogeneous clients.
During training, we simulate 100 clients, in which each client is assigned 375 training samples.
We create local datasets with two different distributions using the same method as in CIFAR-10: i.i.d and non-i.i.d ( α = 0 . Table 7 list detailed hyperparameters for training LSTM/IMDB
[Table] [Table]
i.i.d ( α
[Table]
[Table] None
[Table]
A.2 Experiments on CNN/FEMNIST and LSTM/IMDB
Figure 5 and 6 show the training performance of PriSM on CNN/FEMNIST and LSTM/IMDB.
Similar as in ResNet/CIFAR-10, PriSM consistently delivers the best final server model accuracy on both i.i.d and non-i.i.d datasets.
B Model’s Rank during Training
To analyze the server model’s low-rank structure, we adopt a similar method as in [1] to calculate the required number of principal kernels to accurately approximate each layer as 2 − log ( (cid:80)
i p
i log p i ) .
Here, p i is calculated as in Eq.
(2) with κ = 2 .
Figure 7 shows the number of kernels required to approximate each layer with 3 × 3 kernels in ResNet-18 during FL with full models, where Block i - j indicates j -th convolution layer in i -th ResBlock.
First, we observe that while the server model attains a low-rank structure during training, a randomly initialized model does not exhibit a low-rank structure.
Therefore, selecting a fixed set of top-k principal kernels for the sub-models inevitably causes reductions in the server model capacity.
Furthermore, even at the end of the FL training, around half principal kernels are still required to approximate most layers.
In fact, some layers require even more principal kernels.
Therefore, our probabilistic sampling scheme is essential in preserving the server model capacity during FL training with sub-models.
[Figure]
C Related Works
Factorized Models .
Training neural networks with layer factorization has been extensively studied in prior literature [7, 16, 14, 22, 13].
Specifically, these works are based on the observation that well-trained neural networks have inherently low-rank structure and exhibit large-correlations across kernels.
Hence, one can potentially down-size the model with a low-rank approximation to provide significant reduction in computations thus speeding up training.
Furthermore, this can make model training more affordable for resource-constrained devices.
Resource-Constrained Federated Learning .
While federated learning opens the door for collaborative model training over edge users having rich (but private) data, the computational and communication footprint prohibits training of large high-performance models at the resource-constrained clients.
To address these resource limitations in federated learning, a number of works have been proposed in the literature [8, 11, 30, 8, 11, 30, 27, 24, 6, 9].
Particularly, in split learning [27, 24, 6], the model is partitioned into two parts, one (small) part is assigned to clients for local training, while the other (large) part is outsourced to the server.
[9] proposes FedGKT that combines the model splitting approach with a novel bi-directional knowledge transfer technique between server and clients to achieve resource-constrained FL with much fewer communications than split learning.
However, works such as split learning and FedGKT require sharing of intermediate activations (and in many cases, logits as well as labels) with the server, directly leaking input information and potentially compromising privacy promises of FL [31].
The works closely related to ours are HeteroFL [8], FjORD [11] and FedHM [30], that aim to enable participation of a resource-constrained client by letting it train a smaller sub-model based on its capabilities.
In particular, HeteroFL and FjORD create sub-models for clients by selecting certain fixed number of original kernels of the server model.
On the other hand, FedHM creates submodels using fixed subsets of factorized principal kernels.
However, in these works, the size of the server model gets limited by the clients with maximum computation and communication capacities, sacrificing the model performance.
This becomes even more critical in a realistic, cross-device FL setting wherein no client has the capacity to train a large model.
While another work, FedPara [12], proposes a low-rank factorized model training to reduce communication costs, computational footprint still remains prohibitive as every client is required to perform full-model training.