Overcoming Resource Constraints in Federated Learning: Large Models Can Be Trained with only Weak Clients
Authors: Yue Niu
Saurav Prakash
Souvik Kundu
Sunwoo Lee
Salman Avestimehr
Venue: Accepted by TMLR ()

Abstract: Federated Learning (FL) is emerging as a popular, promising decentralized learning framework that enables collaborative training among clients, with no need to share private data between them or to a centralized server. However, considering many edge clients do not have sufficient computing, memory, or communication capabilities, federated learning of large models still faces significant bottlenecks. To keep such weak but crucial clients in the loop, prior works either consider a heterogeneous-client setting where clients train models with different sizes; or offload training to the server. However, the heterogeneous-client setting requires some clients to train full model, which is not aligned with the resource-constrained setting; while the latter ones break privacy promises in FL when sharing intermediate representations or labels with the server. To overcome these limitations, in this work, we formulate a realistic, but much less explored, cross-device FL setting in which no client can train a full large model nor is willing to share any intermediate information with the remote server. Under such a formulation, we develop a principal sub-model (PriSM) training methodology to collaboratively train a full large model, while assigning each client a small sub-model that is a probabilistic low-rank approximation to the full server model. When creating sub-models, PriSM first performs a principal kernel analysis in the orthogonal kernel space to obtain importance of each kernel. Then, PriSM adopts a novel importance-aware sampling process to select a subset of kernels (i.e., a kernel with high importance is assigned with a higher sampling probability). This sampling process ensures each sub-model is still a low-rank approximation to the full model, while all sub-models together achieve nearly full coverage on the principal kernels. To further improve memory efficiency while still preserving accuracy, PriSM also exploits low-rank structure in intermediate representations and allows each sub-model to learn only a subset of them. Our evaluations on various datasets and models (CNNs, LSTMs, Transformers) under different resource-constrained settings demonstrate that PriSM yields an accuracy improvement of up to $10\%$ compared to existing works. More importantly, PriSM does not incur significant accuracy degradation compared to full-model training (e.g., only $\sim 2\%$ accuracy drops for ResNet-18/CIFAR-10 when clients train only $0.2\times$ sub-models).

Published in Transactions on Machine Learning Research (10/2023)
Overcoming Resource Constraints in Federated Learning: Large Models Can Be Trained with only Weak Clients
Yue Niu, Saurav Prakash Deportment of Electrical and Computer Engineering University of Southern California
ynerdt,,aurmaprruucc..dd
Souvik Kundu Intel Labs, San Diego, USA
kowvkkk.unddaainn wwwtund@@nnee..co com
Sunwoo Lee
sumwool@knnaa.cm ac.kr
Department of Computer Science and Engineering, Inha University
Salman Avestimehr
avestime Grscc. edu
Deportmend of Electricul and Compuler Engineering
Universily of Southern California
Reviewed on OpenReview: https: // openreview. net/ forum? id= tziWnkL999
Abstract
Federated Learning (FL) emerging as a popular, promising decentralized learning frame- work that enables collaborative training among clients, with no need to share private data between them or to a centralized server. However, considering many edge clients do not have sufficient computing, memory, communication capabilities, federated learning of large models still faces significant bottlenecks. To keep such weak but crucial clients in the loop, prior works eir consider heterogeneous-client setting where clients train models with different sizes; or offload training to the server. However, the heterogeneous-clies client setting requires some clients lo train full model, which is not aligned with the resource- constrained setting; while the latter ones break privacy promises in FL when sharing in- termediate representations or labels with the server. To overcome these limitations, this work, we formulate realistic, but much less explored, cross-device FL setting in which no client can train a full large model nor is willing to share any intermediate informa- develop principal sub-model
tion with the remote server. Under such formulation, (PriSM) training methodology to collaboratively train a full large model, while assigning each client a small sub-model that is a probabilistic low-rank approximation to the full server model. When creating sub-models, PriSM first performs a principal kernel analysis in the orthogonal kernel space to obtain importance of each kernel. Then, PriSM adopts a novel importance-aware sampling process to select subset of kernels (i.e., kernel with high importance assigned with higher sampling probability). This sampling process ensures each sub-model is still low-rank approrimation to the full model, while all sub-models together achieve nearly full coverage on the principal kernels. To further improve memory efficiency while still preserving accuracy, PriSM also exploits low-rank structure in inter- modiate representations and allows each sub-model to learn only a subset of them. Our evaluations on various datasets and models (CNNs, LSTMs, Transformers) under different resource-constriine settings demonstrate that PriSM yields an accuracy improvement of up to 10% compared to existing works. More importantly, PriSM does not incur signifi- cant accuracy degradation compared to full-model training (c.g., only ~ 2% accuracy drops for ResNet-18/CIIARR- when clients train only 0.2x sub-models). Code is available at
* Authors have equal contribution.

---

machine Published in Transactions on Machine Learning Research (10/2023)
1 Introduction
Federated Learning (FL)) is emerging as a popular paradigm for decentralized and
learning as it allows clients to perform ML optimization jointly without directly sharing local data (McMahan et al. 2017 Kairouz et al. p-2TD- 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 data to participate in joint training without 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 FT. (Pichai 2019).
Despite significant progress in FL in the recent past, several crucial challenges still remain when moving to the edge. In particular, limited computation, memory, and communication capacities prevent clients from learning large models for leveraging vast amounts of local data at the clients. This problem is getting
increasing attention in current literature (Diao et al. 2021 Horvath et al. 2021 Yao et al. 2021: Hong et al.,
2020). For example, recent. works propose a sub- model training methodology by assigning clients with different subsets of the full model depending on their available resources (Diao et al. 2021 Horvath et al.
2022 Mei el. al., 2022 Vepakomma el al. 2018: He el al.
Yao et al. 2021 Li ct al. 2021 However, these works have an underlying assumption that some of the clients have sulficient resources to train a nearly full large model. In particular, methods like FedHM (Yao et al.] 2021 that adapt low-rank compression to FL, incur more memory footprint for intermediate representations, even for small sub-models compared to the full model training. As a result, server model size is limited by the clients with maximum computation, memory, and communication capacities. To overcome resource constraints on clients, other works (Vepakomma et al., 2018 Thapa et al. 2022 He et al., 2020) 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 representations and/or labels with the server, which directly leaks input information and potentially compromises privacy promises of FL (Erdogan et al.) 2022 Zhang et al. 2020
2021
Our Contributions. Unlike prior works, this work considers even more constrained and realistic settings at the edge, in which no client is capable of training large model nor is willing to share any intermediate data and/or labels with the server. To this end, we propose novel Principal methodology. At a high level, PriSM allows each client to only train small sub-model, while still allaining a full model on the server and achieving comparable accuracy as the standard training.
(PriSM) training
The cornerstone of PriSM is the models' inherent low-rank structure, which commonly used reducing compute costs (Khodak al., 2021 Denton et al.] 2014). However, naive low-rank approximation in FL (Yao et al.) 2021), where
60
BlackI-1
#RRoquired principal koomells
Block1-2
all clients only train top-k kernels, incurs notable accuracy drop, especially in very constrained scttings. In Figure we delve into the matter by showing the nunber of principal ker-
50-
Block2-1
Block2-2
nels required in the orthogonal space to accurately approximate each convolution layer in the first two ResBlocks ResNet-18
40
(He et al. 2016) during FL training]] 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 (See Sec 4.4.. Therefore, to avoid the reduction in server model capacity, it is essential to ensure that all server-side prin- cipal kernels are collaboratively trained on clients, especially when each client can only train a very small sub-model (c.g.: < 50% of the server model).
50% kemels
32-
30-
100
200
300
400
Communication round
Figure #principal kernels in the orthogonal space required to accurately approximate convolu- tion layers in the first two ResBlocks in ReaNet-18 during PL training. Blocki-j indicates j-th con- volution layer in i-th ResBlock. The server model gradually attains a low-rank structure. However, even in the final model, half of the principal ker- uels are still required to approxinate most layers.
Based on the above observations, we develop an effective prob- abilistic strategy to select subset of kernels and create sub-
1 See Sec 4.4 for further details, esperially for calculating the required number of principal kernels.
2

---

Published in Transactions on Machine Learning Research (10/2023)
model for each client as shown in Figure More specifically, PriSM first converts the model into orthogonal space where original convolution kernels are decomposed into principal kernels using singular value decom- position (SVD). To approximate the original server model, PriSM utilizes novel sampling process, where a principal kernel with a larger singular value has a higher sampling probability. The probabilistic process ensures that all sub-models can logether provide nearly full coverage of the principal kernels, thus reaching the near full-model training performance with significantly reduced costs on local computation and commu- nication during sub-model aggregation. PriSM further improves memory efficieney by exploiting low-rank structure in intermediate activations and allows each client to learn only subset of these representations while still preserving training performance. Thus, computation, memory, and communication bottlenecks at the edge are effectively resolved.
singuiar valuss:
(
sampling prob.:
Noum/ReLL
Norm//RLLU
(21-22.01))
Conv Saav:Linne Linear
SVD
Sampling
Conv/Linea
Norm/ R
Normy/Rel.l
original kernels
principal (orthogonal)
sampled kernels
kernels
Original Model
PriSM
Sub-Model
Figure 2: Creating clients' sub-models. PriSM samples subset of principal keruels create client's sub-model with loss computation and communication overhead. The importanco-aware sampling scheme (basod singular values) cusures every sub-model approximates the full large model, and all sub-models together provide nearly full coverage of the principal kernels. PriSM further improves memory efficiency by allowing each sub-model learn only subset of intermediate representations.
In the following, we summarize our contributions as follows:
We propose novel algorithm PriSM to overcome resource constraints of large model training in federated settings. PriSM uses a novel probabilistic sampling scheme to select a subset of the kernel and create a sub-model for each resource-constrainee client.
We present comprchensive empirical evidence showing that kernel sampling is crucial to preserve models' performance, rather than simply fixed kernel dropout.
We conduct exlensive evaluations of PriSM on vision and language tasks under resourced-constraed settings where no client is capable of training the large full model. In particular, we consider both resource constraints and heterogeneity in system capacitics as well as data distribution.
We further provide detailed insights into the performance gains attained by PriSM via 1) analyz- ing server model's rank structure during training; 2) investigating different sampling methods; 3) profiling the kernel sampling process; 4) breaking down costs in the system.
Our results demonstrate that PriSM delivers consistently better performance compared to prior works, especially when participating clients have very limited capacitios. For instance, on ResNet-18/ ResNatt18/CIFAR-10 CIFAR-10 (scc Figure 5). we show that PriSM only incurs around marginal accuracy drop for i.id and non-i.id datasets under constrained setting where clients train sub-models with only 20% of the principal kernels, accounting for 5% of the full server model. Compared to other solutions, PriSM improves the accuracy by up to 10%.
2 Related Works
Factorized Models. Training neural networks with layer factorization has been extensively studied in prior
literature DDetton et al. 2014 Khodak et. al.] 2021 Jaderberg et al. 2014 Novikov et al. 2015 Toannou
3

---

Published in Transactions on Machine Learning Research (10/2023)
et al. 2016 Wang et al. 2023 Niu et al. 2017 Mci et al. 2017). These works are based on the observation that well-trained neural networks have inherently low-rank structure and exhibit large correlations across kernels. A model's rank can be pre-defined or learned and adapted with layers during training (Wang et al, 2023). Hence, one can down-size the model with a low-rank approximation to provide a significant reduction in compulationn ions thus speeding up training. Furthermore, this can make model training more affordable for resource-contrranned devices. In addition models' low-rank structure, works such as (Niu al. 2022 also exploit low-rank structure in intermediate representations to reduce computation overhead.
Resource- Constrained Federated Learning. While federated learning opens the door for collaborative model training over edge users having rich (but private) data, the computation, memory, and communication footprint prohibits training of large models at the resource-constrranne clients. To address these resource limitations in federated learning, a number of works have been proposed in the literature (Diao et al.
2021
Horvath et al.) 2021: Yao et al.) 2021: Diao et al. 2021 Horvath et al.) [202T; Vepakomma et al.
2018
Poirot ct al., 2019 Chopra al. 2021 He et al. 2020 Prakash ct al.) 2020 Elkordy et al.) 2022).
al. 2018 Poirol et al. 2019 Chopra et al. 2021). 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. He ct al. (2020) proposes FedGKT that combines the model splitting approach with a bi-directional knowledge transfer technique between server and clients to achieve resource- constrained FL with much fewer communications than split learning. However, these works require sharing of intermediate activations (and in many cases, logits as well as labels) with the server, directly leaking input information and potentially costing the privacy promise of FL (Erdgaan et al. 2022 Zhang et al. 2020
Particularly, in split. learning (Vepakomiaa
The works closely related to ours HeteroFI. (Diao et al.) 2021 ) FjORD Horvath et al. 2021 SplitMix (Hong et al. 2022 FLANC (Mei et al. 2022) and FedHIM Yao et al. 2021 that aim to enable participation of a eesourceeconstrainee client by letting it train smaller sub-model based on its capabilities. In particular, HeteroFL and FjORD create sub-models for clients by selecting a certain fixed number of kernels of the server model based on each client's capacity. During model aggregation, the server aggregates cach kernel and recovers the full model (if all kernels are trained). Similarly, SplitMix also splits a large model into several sub-models. However, it trains and aggregates each sub-model separately, leading to server model as an ensemble of sub-models. On the other hand, FedHM and FLANC create sub-models 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, memory, and communication capacities, sacrificing the model performance. In particular, methods like FedHMM incur more memory footprints for intermediate representations, even for the small sub-models. This becomes even more critical in the realistic, cross-device FL setting wherein no client has the capacity to train a large model.
Structured pruning methods such as FedMask (Li et al., 2021) are also aimed at reducing training costs by pruning channels. However, FedMask requires warm-up stage with full-model training, which does not fit into our setting that clients can train a full model. Besides these sub-model training methods for reducing computations, FedPara (Hyeon-Woo et al. 2021), proposes low-rank factorized model training to reduce communication costs. However computational footprint still remains prohibitive as cach client conducts full- model training. FeDepth (Zhang el al., ) targels memory efficiency on edge clients. Instead of training a full model, it splits a full model into blocks and sequentially trains each block. As FeDepth still needs to train all blocks and upload models with the same size as the full model, computation, and communicatio costs are not reduced.
Therefore, to overcome the aforementioned challenges, we propose new training method to effectively address the computation, memory, and communication bottleneck at the edge, while still preserving the privacy promises of FL.
3 Method
In this section, we first motivate our proposal, Principal random Sub-Model training (PriSM), with an observation of orthogonality in neural network layers. Then, we describe the details of PriSM.
Notations- pl Frobenius norm. oj i-th singular value in a matrix. @: convolution. 4 matrix multipli- cation. (::): sum of element-wise multiplication or inner product. tr(A), trace of matrix.
4

---

Published in Transactions on Machine Learning Research (10/2023)
3.1 Motivation: An Observation Orthogonality
Without loss of generality, consider convolution layer with kernels W RNxMxxxx and imput X E I Mxxxw where N and M denote the number of output channels and input channels, & is kernel size, and h x w is the size of the input image along each channel. Based on a common technique im2col (Chellapilaa et al. 2006), the convolution layer can be converted to matrix multiplication as Y W. X, where W E R N MM* and X E M For kernel decorrelation, we apply singular value decomposition (SVD) to map kernels into orthogonal space as: W EE 0; u, off, where (u))= {wx)== sets of orthogonal vector: The convolution can be decomposed as
Σ v
Σ
o3 ui 0 XX.
(1)
For Vi + j, it is easy to verify that =4 tr(X'= U; u. Uj .03 X) =0, namely the output features Y; and Yj are orthogonal. Therefore, if regard W; o; u, of principal kernel, different principal kernels create orthogonal output features. To illustrate this, Figure 3 shows an 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 revealed in (iie et al. 2017 Balestriero et al. 2018 Wang et al. 2020), imposing orthogonality on kernels leads to better Training performance. This molivates us to initiate the training with set orthogonal kernels. Furthermore, to preserve kernel orthogonality during training, it is critical to constantly refresh the orthogonal space through ro-decompasition. The above intuitions based the obscrvation orthogonality play key role in PriSM, which is described in the following section.
Inputs
Kerne
Kerne 2
Kernel
Figure 3: Orthogonal features generated principal kernels. Different principal kernels capture different features (Kernel 1: outline of the cat, Kernel and textures but distinct regions).
Remark 3.1. Besides the convolution layer discussed above, the layer decomposition using SVD also applics to general linear layers (e.g., MLPs, LSTMs, Transformers), where weights W E NxM In particular, we can directly apply SVD and obtain the principal components.
3.2 PriSM: Principal Probabilistic Sub-Model Training
Motivated by the observation that output features orthogonal given orthogonal kernels, we propose PriSM, a new sub-model training method that directly trains orthogonal kernels in clients. Particularly, PriSM introduces two key components to ensure training performance with sub-models under very con- strained settings. First, considering different orthogonal kernels and their contributions to outputs are weighed by the corresponding singular values, PriSM devises a novel importanco-aware sampling schome to create client sub-models to achieve computation and communication efficiency. particular, the sampling scheme brings two benefits: i) each sub-model is a low-rank approximation of the server model: ii) the con- glomerate of the sampled sub-models enables nearly full coverage of orthogonal kernels. Second, to further improve the memory efficiency while maintaining training performance, PriSM exploits low-rank structure in activations and allows each client to learn only subset of intermediate representation alons in each layer. We start with unfolding sub-model creation, and then describe the training procedure in PrISM.
We assume w.l.o.g W matrix.
5

---

Published in Transactions on Machine Learning Research (10/2023)
3.2.1 Sub-Model Creation
This section explains how PriSM creates different sub-models for local clients.
Computation and communication efficiency via sub-model sampling. For a convolution layer with principal kernels { [W& 1 We)men i=1" and the corresponding singular values {ar)== randomly sample principal
N
kernels denoted by W with sampling probability for i-th kernel as follows:
P: E==1o3
(2)
Here, for a given layer, r is decided by e-th client's resource budget, and is smooth factor controlling the probability distribution in sampling. Indices of selected kernels are denoted as Iz.
The convolution output calculated as
Y= ΣY. Eo, oo u, of XX iEI.
iEI.
By performing the sampling process for every convolution layer, create random low-rank model for each client. Important kernels with large singular values are more likely to be chosen, and all sub-models can together provide nearly full coverage of principal kernels. The resulting sub-model hence has convolution layers with fewer kernels, reducing required computations and communication overhead (when aggregating models). Other element-wise layers, such as ReLU and batch normalization, remain the same.
Memory efficiency via learning a subset of features. We further exploit low-rank structure in inter- mediate representations. For sub-model above, each convolution layer is decomposed into two sublayers: Convy and Convy with kernels {ui}iez, and (vi respectively. Without further optimization, output from Convy consumes the same memory as the original layer, causing high memory pressure at the edge, However, inputs to Comny, X E RR' and kernels W E Nxx are both low-rank. Therefore, based on the rank incquality of matrix multiplication (Banerjoe & Roy 2014 we obtain
Rank(Y) S min (Rank X Rank(W)) s..
(4)
Therefore, the output of Convy is also low-rank. Such a low-rank structure reflects correlation among output channels, indicating channel redundancy outputs when N. Based on the observation, we further allow Convy to compute only subset of output. channels, therefore reducing memory footprints of activations while still preserving necessary information. Typically, given input to Convy with input channels, PriSM only computes r output channels. During implementation, when computing output features from Comvy. For simplicity, we will still use in the rest of the paper to denote the subset (1: r). The unselected channels are cached in ú; and are used in model aggregation.
replace the select u, with its subset u,(1 r)
3.2.2 Training
This section details the training procedure in PriSM. We describe each component below.
Local training. On each client, during local training, the sub-model with parameter {a,, u, v2 Weee, are updated. The sub-model also consists of trainable layers such as BatchNorm but with fewer channels denoted in I.. PriSM allows d; to be trained in local training, thus ensuring changes regarding each principal kernel's importance are captured in local clients. In addition, will be merged into u. as u, Volu, V. /o,0 to reduce memory footprints.
Sub-model aggregation. On the server side, with sub-models obtained from clients, we aggregate i-th principal kernel as follows:
W, = (I affuis a. 111) ( Easois) T
(5)
Loec
cee
6

---

Published in Transactions on Machine Learning Research (10/2023)
where C denotes the subset of active clients, a; is the aggregation coefficient for i-th kernel. We propose weighted averaging scheme: if i-th kernel is selected and trained by C₁ clients, then of = 1/C; Further- more, the unselected output channels ú; are concatenated with the aggregated u, before reconstructing the orthogonal kernel to preserve model capacity. W; was not selected by any client, it remained unchanged in the server. The full model in the original space is constructed by converting each 2-dimensional W; the original dimension M and combining them.
Orthogonal space refresh. After model aggregation, perform SVD the updated kernels W to preserve orthogonality among the principal kernels. Thus, in the next communicatio round, the importance- aware sampling can still create low-rank sub-models for different clients..
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 statisties; namely, the normalization always uses current. batch statistics in the training and evaluation phases. Each client applies normalization separately with no sharing of statisties during model aggregation. Such an adaptation is effective in ensuring consistent outputs between different sub-models and avoids potential privacy Icakage through the running statistics (Andreux et al. 2020).
Regularization. When learning a factorized model on a client, applying weight decay to u, and v; separately results in poor final accuracy. Inspired by Khodak al. (2020), for training on client e, we add regularization the subset of kernels as follows:
reg
A
Σ u,
'T
(6)
2
:EZ.
where À is the regularization factor, Ie denotes the subset of principal kernels on client
Algorithm D presents an overall description PriSM. We only show the procodure on a single convolution layer with kernels W for the sake of simplifying notations.
Algorithm PriSM: Principal Probabilistic Sub-Model Training
Input: layer parameters W, client capacities.
1: Decompose W into orthogonal kernel using SVD (W.)
2: for communication round t = 1.... do
3: Choose subset of clients + C.
4: for each client с E C do Compute the sub-model size for client c - [Tal-
Sub-model: Obtain sub-model following the procedure in Sec 3.2.1 Local training: Perform LocalTrain Ic, W
Ics W.
end for Sub-model aggregation: Aggregate parameters based on Eq.
5
W
+-
W
EEc
10: Orthogonal space refresh: Perform SVD on W.
11: end for
12: LocalTrain ++ Ic, W
13: for local iteration & = 1, ... K do
14: Sample an input batch from the local datasct + D
15: Perform the forward and backward pass + Dx. W
16: Update the local sub-model using SGD > W
17: end for
In the following remarks, we differentiate PriSM from Dropout and Low-Rank compression.
7

---

Published in Transactions on Machine Learning Research (10/2023)
Remark 3.2. PriSM us Dropout. PriSM shares some computation similarity with model training using regular dropoul (Diao el al. 2021; Horvath et al. 2021).. However, methods with regular dropoul simply select a fixed subset of kernels with size depending on a client's resource constraints, and suffer severe performance degradation, esp. with a high dropout probability. In contrast, PriSM performs importance-qware sampling in the orthogonal space. Each sub-model upprorsmates the full model, and different sub-models do not create significant inconsisteney. Importantly, PriSM obtains a final global model with full capacity, even though only allows clients to train sub-models.
Remark 3.3. PriSM es Low-Rank Compression. PriSM is not low-rank compression method. Low-rank compression methods such as FedHM (Yao et al. 2021) aim to construct a smaller server model by completely discarding some kernels even though they can still contribute to training performance (See Figure 7 in Section 7 and Figure 3 in Section 4-4). PriSM randomly select sub-models so that every kernel is possible to be learned. Furthermore, PriSM achieves memory efficiency by exploiting low-rank properties in intermediate represenlations, which is not seen in other low-rank methods.
4 Experiments
In this section, present numerical analyses of PriSM. First, we study the necessity of probabilistic kernel sampling in preserving models' capacity and accuracy via centralized training. Then, we evaluate PriSM under resource-constrained settings where no clients can train the large full model. Furthermore, we consider both homogeneous and heterogeneous client settings. Specifically, in homogeneous settings, all clients have the same limited computation, memory, and communication capacity, while in heterogeneous settings, clients' capacitics might vary. We also compare PriSM with three other basclines: ordered dropout orthogonal space (OrthDrop), ordered dropout in original space (OrigDrop), and SplitMix (Hong et al. 2022). At a high level, our results demonstrate that PriSM achieves comparable accuracy to full-model training even when only training very small sub-models on all clients. Additionally, we provide more insights into the superior performance of PriSM by 1) analyzing the server model's rank structure; 2) profiling the kernel sampling process during training; 3) investigating cost breakdowns in a federated system. Finally, we conduct ablation studies including analyzing several potential sampling strategies and impacts of training hyperparameters.
Baselines. Prior methods such as FjORD (Horvath et al. 2021 and HeteroFL (Diao al. 2021 select sub-models from the original kernel space, for which we denote as OrigDrop. In implementation, we follow the procedure in HeteroFL. On the other hand, we use OrthDrop to denote selecting fixed top-k kernels from the orthogonal space such as in FedHM (Yao et al.] 2021). For SplitMix, we follow the procedure in (Hong et al. 2022 to create atom sub-models and evaluate the ensemble of all sub-models on the server side.
Models and Datascts. We train ResNet-18 on CIFAR-10 (Krizhevsky et al., 2009), CNN on FEMNIST (Caldas et al., 2018) and LSTM model on IMDB (Mas el al.] 2011). We also train a transformer model DeiT (ouvrrn et al. 2021 on CIFAR-10. Detailed architectures are provided in Appendix A.1
Data Distribution. For CIFAR-10 and IMDB, we uniformly sample an equal number of images for each client when creating i.i.d datasets. For non-i.i.d datasets, we first use Dirichlet func- tion Dir(a) (Reddi et al.) (020) 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 a = and a = 0.1, where a smaller a indicates a higher degree of non-i.idd For FEMNIST, we directly use the dataset without any additional preprocessing.
fill model
ISSM
OriDoop
OrigDoop
07% 8.5
eb:
model size
FL Sctting. We simulate an FL setting with 100 clients with 20 random clients active in each communication round. Each client trains its model for 2 local epochs in each round. We use SGD with momentum during training. The learn- ing rate initially 0.1 for ResNet-18 and DeiT and 0.01 for CNN/FEMNTSS and decayed by a cosine annealing schedule.
Figure 4: Accuracy of VGG11 on CIFAR-10 with different sub-models. PriSM delivers much better performance under very constrained scttings (o.g., training with only 0.2x sub-models) compared to OrthDrop and OrigDrop.
Further details are in Appendix A.2
8

---

Published in Transactions on Machine Learning Research (10/2023)
4.1 Probabilistic Lub-Model Training in Centralized Settings
We first demonstrate that probabilistic kernel sampling is essential to preserve models' capacity and maintain accuracy in centralized settings. While centralized training is not our goal, it provides direct evidence that PriSM reduces training complexities without significantly sacrificing accuracy.
We use VGG11 as an example and train the model on CIFAR-10. During training, in each iteration, we sample subset of kernels and create sub-model as described in PriSM. In different iterations, different subsets of kernels will be sampled and trained. We adopt the following hyperparameters during training( SGD, batch size=128, weight decay=1e-4, epochs=150). To further demonstrate the benefits of kernel sam- pling, we also compared the accuracy with OrigDrop and OrthDrop, where we respectively extract fixed subset of kernels and create sub-models in the original and orthogonal kernel space.
Figure f shows the accuracy of VGG11 on CIFAR-10 with different sub-models. We observe that PriSM achieves much higher accuracy compared the other two basclines, especially under very constrained settings (o.g. 0.2x sub-models). Hence, compared to these compression methods that reduce models' capacity, PriSM preserves models' capacily via the probabilistic kernel sampling, therefore maintaining model accuracy.
4.2 Performance on Constrained Homogeneous Clients
Table 1: Model size, MACs, activation memory for sub-models in PriSM (batch size: 32). Unlike conventional low-rank methods, PriSM greatly reduces model size, computations, and activation memory when training with very small sub-models.
Model
Full
0.8x
0.6x
0.4x
0.2x
PriSM
Params
11 M
7.9 M (72%)
4.5 M (41%)
M (18%)
.5 M (4.5%)
MACs
35
25.5 G (73%)
14.7 G (42%)
6.87 G (20%) 18.9 M (60%)
G (5.6%)
Mem
31.5 M
37 M (115%)
28.3 M (90%)
9.5 M (30%)
ResNet-18
OrthDrop
(FedHM (Yao al.
2021
Params
(90%)
(67%)
M (44%)
2.5 M (22%)
MACs
35 G
28.8 G (82%)
22.4 G (64%)
16 G (46%)
7.4 G (23%)
Mem
31.5
44.9 (143%)
40.5 (129%)
36.1 M (115%)
31.7 M (101%)
In this setting, we assume that all clients have the same limited capacity. We vary the client sub-model size from 0.2 to 0.8 of the full server model, where O.r indicates that only O.r subset of the principal kernels are sampled in each convolution layer from the server model (denoted keep ratio). Accordingly, for given layer, r 0.z NN, where N is the number of output channels in the layer. In Table I memory and communication footprints for sub- models of ResNet-18. DefT/ CATT/CIAA-10,, CNN/FEMNIST and LSTM/IMDB also enjoy similar cost. reductions. Tt is worth noting that PriSM incurs much smaller costs compared to the deterministic compression methods(FedHMM) thanks to the memory-efficiense design.
we list the computation,
full model: 84.96
-9
b0c.
full model: 92.46
Jull model: 92
PrisM
Origttop
OrtiDrop
SplaMix
de
moddl size
(a) Lid
(b) non-i..d (a
(c) non-i.id (a 0.1)
Figure 5: Accuracy drops ResNet1s/CIFAB-100 homogeneous clients compared to full-model training. PriSM constantly delivers better performance compared OrigDrop, OrthDrop, and SplitMix, significantly outperforming them under very constrained realistic edge settings. (bar: mean; line: std)
ResNet-18 CIIAR-110. Figure 5 shows final validation accuracy drops of ResNet-18 with different sub- models i.i.d and non-i.i.d datasets compared to full-model training. We note that PriSM constantly

---

Published in Transactions on Machine Learning Research (10/2023)
delivers better performance than the other three baselines. The performance gap is more striking under very constrained settings. For instance, when only 0.2x sub-models are supported on clients, PriSM attains comparable accuracy as full-model training, and achieves up to 10%/14% 13% performance improvement compared to OrthDrop, OrigDrop and SplitMix on non-i.id dataset with a = 0.1. We also make two key observations. First, training with sub-models the orthogonal space provides better performance than the original space, which aligns with our intuition Section 3.1 Second, our importanco-aware sampling strategy for creating sub-models indispensable as demonstrated by the notable performance gap between PrISM, OrthDrop and MixitMix.
UNN/FEMNIST. We adopt CNN model as in (Horvath et al. 2021 which consists of two convolution layens (See Table S in Appendix for details). During training, we simulate 100 clients, in which each client contains 10 users' data from FEMNIST (Caldas et al. 2018). Figure shows accuracy drops of sub-model training with different sizes. Similar as in the ResNet experiments, PriSM incurs very small accuracy drops even under very constrained scttings. particular, 0.2x sub-model training only results in less than 1% accuracy drops, greatly outperforming the other three baselines.
LSTM/IMDB, The LSTM model used in FL training is de- tailed in Table g During training, we simulate 100 clients, in which each client, is assigned 375 training samples. We create local datasets with two distributions using the same method as CIFAR-10: i.id and non-i.idd (a = 0.1). Table 12 (in Appendix) lists detailed training hyperparameters. Figure shows accuracy drops of sub-model training IMDB in ho- mogeneous client settings. While the task IMDB is just binary classification problem, PriSM still achieves the best final server model accuracy on both i.id and non-i.i.d datasets. On i.i.d datasots, training using sub-models achioves comparable accuracy as full-model training, even only using very small sub- models such as 0.2x. On the other hand, on non-i.id dataset, OrigDrop suffers notable accuracy drops compared to PriSM.
900 dtop
haseline: 81. 86
I1 converge
quraix
oi
model
Figure 6: Accuracy drops FEMNIST ho- mogeneous clients compared to full-model training. PriSM incurs smaller accuracy drops under constrained scttings than the other three baselines.
-355
-375
baseline: 83.23
baseling: 84.99
PASM
Diop
04
d6
model
(a) i.i.d
(b) non-i.id (a 0.1)
Figure 7: Accuracy drops on IMDB on homogeneous clients compared to full-model training.
Transformer Ceansforee/CIIFAAE /CIFAR-10. PriSM also applies to the transformer architecture. To that end, we test PriSM DeiT Touvron et al. 2021), an efficient image transformer model. It consists of 12 attention layer, and each layer has 3 attention heads. Detailed architecture is provided in Table 6 (in Appendix A.I). We use a pro-trained model on ImageNet and adopt the same hyperparameter settings as in CocNNe--1/CIFAR-II Table 2 shows the final global model performance with different sub-model training. Similar to ResNet- 18/CIFAR-10 10, PriSM maintains comparable accuracy on both i.id and non-i.id datasets even under very constrained settings such when clients can only train 0.2x or 0.4x sub-models. Therefore, thanks to the probabilistic sampling process and memory-efficieng design, PriSM preserves the server model's capacity while significantly reducing training costs.
Table 2: Training performance of DelT/CIFAR resource-constrained clients.
Distribution
Full Model (Baseline)
0.8x
0.6x
0.4x
0.2x
L.d
92.04
91.92 + 0.18
91.36 0.26
90.12 + 0.31
88.34 0.28
non-i.i.d (a 1)
91.14
90.84 + 0.15
90.50 0.18
89.72 + 0.21
87.87 + 0.32
10

---

Published in Transactions on Machine Learning Research (10/2023)
4.3 Performance Constrained Heterogeneous Clients
To simulate clients with varying limited capacity, we simulate the following setting: 40% clients train 0.4x sub-models, and 60% clients train 0.2x sub-models. No participating client trains the full model. For baseline methods, we follow the same strategy in Section
Table 3 lists the final accuracy achieved by different methods under the heterogeneous setting, PriSM greatly outperforms the bascline methods even when 0.4x sub-models are supported on a small fraction of clients. Furthermore, similar to the results in Section 4.2 the benefits of training in the orthogonal space and importance-aware sampling strategy are also observed heterogeneous client settings.
Table 3: Training performance ResNet-18/CIPAR--D on heterogeneous clients.
Distribution
Full Model (Baseline)
OrigDrop
OrthDrop
PriSM
i.d non-i.i.d (1)
92.46
88.86 0.17
89.57 1 0.16
90.63 0.14
92
86.91 + 0.24
88.78 # 0.22
89.89 + 0.2
non-i.i.d (0.1)
84.96
76.38 0.92
78.37 + 0.99
82.58 0.51
120
Rock3
Blokd
-091
Phoppiee peineded bemmet
220-
4401
PRoppird prieaied
100
Rockd-2
Slodd5-2
60
100
Blol3-11
320
50% barral
Blo-6--2
1.0
431
103
2011
240
100
200
360
400
Cosnsanicatiss poca
Conrannncniia conzal
(a) ResBlock 3, (128 kernels). (b) ResBlock 5, 6 (256 kernels). (c) ResBlock 7, 8 (512 kernels).
Figure 8: The number of principal kernels required to accurately approximate each convolution layer in ResBlocks 3-8 in ResNet-18 (Results of ResBlocks and 2 are discussed in Figure The server model gradually attains low-rank structure. However, even in the final model, half of the principal kernels are still required to approximate most layers.
4.4 Insights into PriSM
We now focus on providing further insights into PriSM by analyzing some of its key aspects. To this end, wo first examine the low-rank structure of models during training, and pinpoint the cause behind the accuracy gap between fixed and random kernel dropout strategies in the orthogonal space. Thereafter, we study the sampling process and cost breakdown in the FL. system
4.4.1 Model's rank during training.
To analyze the server model's rank structure, we adopt a similar method as in Alter et al. (2000) to calculate the required number of principal kernols to accurately approximate each layer as 2 log Σ, p. logr.) Hore, p. is calculated as in Eq. (2) with x = 1. Figure g shows the number of required kernels for each layer in ResNet-18 during full-model FL (Blocki-j: j-th convolution layer in i-th ResBlock). We observe that a randomly initialized model is low-rank, rather the server model attains a low-rank structure gradually, Therefore, sub-models with fixed top-k principal kernels inevitably cause reductions in the server model capacity. Furthermore, even at the end of the training, around half the 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.
4.4.2 Kernel sampling profiling.
11

---

Published in Transactions on Machine Learning Research (10/2023)
Rocclll
Block4.1
Blaa66-2
Figure shows the average number of clients assigned for each orthogonal channel in one communication round. Each client trains a 0.2x sub-model of ResNet-18 on CIFAR-10 with ii.d distributions as in Sec 4.2 We observe that each kernel is se- lected by least one client each round, indicating every ker- nel will be activated and trained on clients each round. Fur- thermore, orthogonal kernels with larger a get more chances to be chosen, which ensures sub-models on all clients consistently approximate the full model.
Block4-2
Bla47-4
raskgged clees
Bloe244
Bleck5-1
Block--2
Bleck
Bloote-1 -
Borihh1
Blockó-1
Blod8 2
D
64
12a
256
s/2
pincipal loazel index
Figure 9: Average number of clients assigned for each orth. kernel during training. Kernels with large singular values get more chances to be chosen.
4.4.3 Runtime breakdown.
In this section, we investigate the relative cost of SVD by breaking down the model training time into four stages: sub-model creation, local training, model aggregation, and obtaining orthogonal kernels (SVD). We choose a large model, ResNet-18, as the larget model. We adopt hyperparameters as in Table [O] We run the server and elient process on NVIDIA RTX 5000. Table | lists each stage's average time in one communication round. We observe that the time of SVD is nearly negligible compared to the time spent on local training. Therefore, even for large models, the relative overhead of SVD is still very small.
Table 4: Training time breakdown for R8sNet-18 on CIFAR-10.
stage
sub-model create
local train
aggregate
SVD)
time
3.27 s
36.82 s
0.11 s
0.96 s
4.5
Ablation Study: Different Sampling Methods
We also conducted an additional study on the effects of different sampling strategies on the final serve model accuracy. In addition to the sampling method in PriSM (Eq (2), ; we also try two other methods: uni form sampling and softmaxbbased sampling. In uniform sampling, we regard each channel as equally important with a uniform sampling probability. On the other hand, the sof teatttaabaast sampling assigns cach channel probability that is calculated by applying the softmax function to the singular values. Hence, it also assigns high sampling probabilities for channels with large singular values.
Table 5 shows the final validation accuracy of the global model. We observe that training with uniform sampling incurs significant accuracy drops in both iid and non-i.id data distributions. It indicates that importance-aware sampling crucial to avoid a large variation between sub-models and the full model. On the other hand, while the boftmax-based sampling also considers each kernel's importance, the sampling method in PriSM performs better in choosing the right sub-model and covering the whole kernel space.
We also study the effects of two hyperparameters. the number of local epochs and active clients. Each of these two parameters affects the sampling process when creating sub-models for clients. Specifically, given a fixed number of total iterations, FL training with a small number local epochs per communication round performs a more frequent sampling process, thus making more orthogonal kernels to be selected and trained. Similarly, the FL training with a large number of active clients per round can also activate more orthogonal kernels. Due to space constraints, the results are deferred to Appendix A.3 At high level, fewer local epochs (given fixed total iterations) more active clients in each round result in higher sampling probability for each kernel. Therefore, training performance is improved shown in the results.
Table 5: Global model accuracy with different sampling strategies.
sampling
i.dd
non-i.i.d (a = 1)
non-i.i.d (a =
uniform
69.54 + 0.21
67.28 1 0.51
57.5 + 0.67
softmax-based
84.74 + 0.49
83.19 + 0.43
72.69 + 0.62
PriSM
90.57 + 0.25
89.36 + 0.21
80.72 + 0.59
12

---

Published in Transactions on Machine Learning Research (10/2023)
5 Conclusion
We have considered a practical yet under-explored problem of federated learning in resource-constrained 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 resource-limited clients by enabling them to train smaller sub-models, while still allowing the server to reconstruct the full model. Importantly, PriSM utilizes novel sampling approach to obtain sub-models for the clients, all of which together provide near-full model coverage. PriSM further improves memory efficiency by exploiting low-rank structure in intermediate activations. Our extensive empirical results diverse models and datasets demonstrate that PriSM performs significantly better than the prior baselines, especially when cach client can train only a very small sub-model. For instance, compared to full-model training, PriSM only incurs ~ 2% accuracy drop for 1esNNet-18/CFAA 10 when elients train only 0.2x sub-models, while yielding an accuracy improvement of up to 10% compared to existing works. We further present insights into performance gains of PriSM compared to other baseline methods and demonstrate the necessity of importance-awei kernel sampling in sub-model training.
Acknowledgement
This material is based upon work supported by ONR grant N00014-23-122191, ARO grant W911NF- 22-1-0165, Defense Advanced Research Projects Agency (DARPA) under Contract No. FASTNICS HIR001120COU0 and HR00112000160, and gifts from Intel and Qualcomm. The views, opinions, and/or findings expressed are those of the author(s) and should not be interpreted as representing the official views or policies of the Department of Defense or the U.S. Government.
References
Orly Alter, Patrick O Brown, and David Botstein. Singular value decomposition for genome-wide expression data processing and modeling. Proceedings of the National Academy of Sciences, 91(11)::0011-10006 2000.
Mathicu Andreux, Jean Ogier du Terrail, Constance Beguier, and Eric W Tramel. Siloed federated learn- ing for multi-centric histopathology datasets. In Domain Adaptation and Representation Transfer, and Distributed and Collaborative Learning, pp. 129 139. Springer, Lima, Peru, 2020.
Randall Balestriero et al. A spline theory of deep learning. In International Conference on Machine Learning, pp. 374 383. PMLR, 2018.
Sudipto Banerjee and Anindya Roy. Linear algebra and matrix analysis for statistics, volume 181. Cre Press Boca Raton, 2014.
Sebastian Caldas, Sai Meher Karthik Duddu, Peter Wu, Tian Li, Jakub Koneōný, H Brendan McMa- han, Virginia Smith, and Ameet Talwalkar. Leaf: A benchmark for federated settings. arXiv preprint arXiv:1812001097, 2018.
Kumar Chellapilla, Sidd Puri, and Patrice Simard. High performance convolutional neural networks for document processing. In Tenth international workshop on frontiers in handuriting recognation. Suvisoft, 2006.
Ayush Chopra, Surya Kant Sahu, Abhishek Singh, Abhinav Java, Praneeth Vepakomma, Vivek Sharma, and Ramesh Raskar. Adasplit: Adaptive trade-offs for resourco-constrained distributed deep learning. arXiv preprint xrX:::2112.01637, 2021.
Emily L Denton, Wojciech Zaremba, Joan Bruna, Yann LeCun, and Rob Fergus. Exploiting linear structure within convolutional networks for efficient evaluation. Advances in neural information processing systems, 27, 2014.
Enmao Diao, Jie Ding, and Vahid Tarokh. HeteroFL: Computation and communication efficient federated learning for heterogeneous clients. International Conference Learning Representations, 2021.
13

---

Published in Transactions on Machine Learning Research (10/2023)
Ahmed Roushdy Elkordy, Saurav Prakash, and Salman Avestimehr. Basil: A fast and byzantine-reslleen approach for decentralized training. IEEE Journal on Selected Areas in Commuuiccations, 40(9):2694 2716, 2022.
Ege Erdogan, Alptekin Küpçü, and A Ercüment Çiçek. Unsplit: Data-oblivious model inversion, model stealing, and label inference attacks against split learning. In Proceedings of the 21st Workshop on Privacy in the Electronic Society, pp. 115-124, 2022.
Chaoyang He, Murali Annavaram, and Salman Avestimehr. Group knowledge transfer: Federated learning large CNNs at the edge. Advonces in Neural Information Processing Systems, 33:14068 -40660 2020.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. Tn Proceedings of the IEEE conference computer vision and pattern recognition, pp. 770 778, 2016.
Junyuan Hong, Haotao Wang, Zhangyang Wang, and Jiayu Zhou. Efficient split-mix fedorated learning for on-demand and in-situ customization. In International Conference on Learning Representations, 2022.
Samuel Horvath, Stefanos Laskaridis, Mario Almeida, Ilias Leontiadis, Stylianos Venieris, and Nicholas Lane. FjORD: Fair and accurate federated learning under heterogeneous targets with ordered dropout. Advances in Neural Information Processing Systems, 34, 2021.
Nam Hyeon- Woo, Moon Ye-Bin, and Tac-Hyun Oh. Fedpara: Low-rank hadamard product for communication-elfiiientt federated learning, arXiv preprint arXiv:2108.06095 2021.
Yani Ioannou, Duncan Robertson, Jamic Shotton, Roberto Cipolla, and Antonio Criminisi. Training CNNs with low-rank filters for efficient image classification. International Conference on Learning Representa- tions (ICLR), 2016.
Max Jaderberg, Andrea Vedaldi, and Andrew Zisserman. Speeding up convolutional neural networks with low rank expansions. British Machine Vision Conference (BMVC), 2014.
Peter Kairouz, H Brendan McMahan, Brendan Avent, Aurélien Bellet, Mehdi Bennis, Arjun Nitin Bhagoji, Kallista Bonawitz, Zachary Charles, Graham Cormode, Rachel Cummings, et al. Advances and open problems in federated learning. Foundations and Trends0 in Machine Learning, 14(1 11::::1 2):1-210, 2021.
Mikhail Khodak, Neil A Tenenholtz, Lester Mackey, and Nicolo Fusi. Initialization and regularization of factorized neural layers. In International Conference on Learning Representations, 2020.
Mikhail Khodak, Neil Tenenholtz, Lester Mackey, and Nicolo Fusi. Initialization and regularization of factorized neural layers. International Conference on Learning Representations (ICLR), 2021.
Alex Krizhevsky, Geoffrey Hinton, et al. Learning multiple layers of features from tiny images. Citeseer, 2009.
Ang Li, Jingwei Sun, Xiao Zeng, Mi Zhang, Hai Li, and Yiran Chen. Fedmask: Joint computation and communication-effiienn personalized federated learning via heterogeneous masking. In Proceedings of the 19th ACM Conference Embedded Networked Sensor Systems, pp. 42 42-55, 2021.
Andrew Maas, Raymond E Daly, Peter T Pham, Dan Huang, Andrew Y Ng, and Christopher Potts. Learning word vectors for sentiment analysis. In Proceedings of the 49th annual meeting of the association for computational linguisties: Human language lechnologies, pp. 142 150, 2011.
Brendan McMahan, Eider Moore, Daniel Ramage, Seth Hampson, and Blaise Aguera y Arcas. Communication-effcieedd learning of deep networks from decentralized data. Tn Artificial intelligence and statistics, pp. 1273-1282. PMLR, 2017.
Chunsheng Mci, Zhenyu Liu, Yue Niu, Xiangyang Ji, Wei Zhou, and Dongsheng Wang. A 200mhz 202.4 gftops@ 10.8 vgg16 accelerator in xilixx vx690t. In 2017 IEEE Global Conference on Signal and Infor- mation Processing (GlobalSIP), pp. 784-788 IEFE, 2017.
14

---

Published in Transactions on Machine Learning Research (10/2023)
Yiqun Mei, Pengfei Guo, Mo Zhou, and Vishal Patel. Resource-adaptive federated learning with all-in-one neural composition. In Advances in Neural Informalion. Processing Systems, 2022.
Yue Niu, Zhenyu Liu, Chunsheng Mei, Xiangyang Ji, Wei Zhou, and Dongsheng Wang. Sensitivity- acceleration and compression algorithm for convolution neural network. In 2017 IEEE Global Conference on Signal and Information Processing (GlobalSIP), pp. 769 773. IEEE, 2017.
Yue Niu, Ramy E Ali, and Salman Avestimehr. 3legrace: Privacy-preserving dnn training over tees and gpus. Proceedings on Privacy Enhancing Technologies, 4:183-200 203, 2022.
Alexander Novikov, Dmitri Podoprikhin, Anton Osokin, and Dmitry P Vetrov. Tensorizing neural nelworks. Advances in neurol information processing systems, 28, 2015.
Sundar Pichai. Google's Sundar Pichai: Privacy Should Not Be Luxury Good. In New York Times, 2019.
Maarten G Poirot, Praneeth Vepakomma, Ken Chang, Jayashree Crapar,y-Cramer, Rajiv Gupta, and Ramesh Raskar. Split learning for collaborative deep learning in healthcare. arXiv preprint arXiv:1912. 12115, 2019.
Saurav Prakash, Sagar Dhakal, Mustafa Riza Akdeniz, Yair Yona, Shilpa Talwar, Salman Avestimehr, and Nageen Himayat. Coded computing for low-latency federated learning over wireless edge networks. IEEE Journal on Selected Areas in Communications,, 39(1):233 250, 2020.
Sashank Reddi, Zachary Charles, Manzil Zaheer, Zachary Garrett, Keith Rush, Jakub Koncčný, Sanjiv Kumar, and H Brendan McMahan. Adaptive federated optimization. arXiv preprint wxXi:2003300295, 2020.
Chandra Thapa, Pathum Chamikara Mahawaga Arachchige, Seyit Camtepe, and Lichao Sun. Splitfed: When federated learning moets split learning. In Proceedings of the AAAT Conference on Artificial Intelligence, volume 36, pp. 8485 8493, 2022.
Hugo Touvron, Matthieu Cord, Matthijs Douze, Francisco Massa, Alexandre Sablayrolles, and Hervé Jégou. Training data-efficient image transformers & distillation through attention. In International conference on machine learning, pp. 10347-10357. PMLR, 2021.
Praneeth Vepakomma, Otkrist Gupta, Tristan Swedish, and Ramesh Raskar. Split learning for health: Distributed deep learning without sharing raw patient data, arXiv preprint X81i::1812 200664 2018.
Hongyi Wang, Saurabh Agarwal, Yoshiki Tanaka, Eric Xing, Dimitris Papailiopoulos, et al. Cuttlefish: Low-rank model training without all the luning. Proceedings of Muchine Learning and Systems, 5, 2023.
Jiayun Wang, Yubei Chen, Rudrasis Chakraborty, and Stella X Yu. Orthogonal convolutional neural net- works. In Proceedings of the IEFE/CVF conference on computer vision and pattern recognition, pp.
11505- 2020.
Di Xie, Jiang Xiong, and Shiliang Pu. All you need is beyond good init: Exploring better solution for training extremely deep convolutional neural networks with orthonormality and modulation. In Proceedings of the IEEE Conference on Computer Vision and Pottern Recognition, pp. 6176-6185, 6185, 2017.
Dezhong Yao, Wanning Pan, Yao Wan, Hai Jin, and Lichao Sun. FedHM: Efficient Federated Learning for Heterogeneous Models via Low-rank Factorization. arXiv preprint arXiv:2111.14655, 2021.
Kai Zhang, Yutong Dai, Hongyi Wang. Eric Xing, Xun Chen, and Lichao Sun. Memory-adaptive depth-wise heterogenous federated learning, arXiv preprint ar arXiv:2303.0 04887, 2023.
Yuheng Zhang, Ruoxi Jia, Hengzhi Pei, Wenxiao Wang, Bo Li, and Dawn Song. The secret revcaler: Gener- ative model-inversion attacks against deep neural networks. In Proceedings of the IEEE/CVE Conference on Computer Vision and Pattern Recognition, pp. 253-261. 2020.
15

---

Published in Transactions on Machine Learning Research (10/2023)
A 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 code upon acceptance of the paper.
A.1 Models
RcsNet-18 SEANett1/CCIAA-O0, CIIAR-11.
We ResNet-18 optimized for CIFAR-10, in which kernel size in the first convolution layer is changed from 7 x to 3 x 3. Details are shown in Table g
DeiT/CIFAR-10. We adopt DeiT--tiy as in Touvron et al. 2021 The detailed model is shown in Table
Table 6: DeiT/CIFAR-10
Module
kernel size
#kernel
#layers
#heads
Embedding Attention
192
12
3
CNN CNN/FEMNISST, We use similar architecture as in FjORD Horvath et al. is shown in Table 7
2021
The detailed model
Table 7: CNN/FEMNIST
Module
#kernels
size
stride
ReLU
Convl
64
5
1
/
Pooling1
2
2
X
Conv12
64
3
1
/
Pooling2
2
X
Classification
10
X
LSTM/IMDB. We common LSTM model as shown in Table g
Table 8: LSTM/IMDB
Module
input size
output size
hidden size
#layers
Embedding
1001
64
LSTM
64
256
256
FC
256
1
A.2 Training Hyperparameters
homogeneous clients. We simulate 100 clients during FL training, in which each client is assigned 500 training samples for both i..d and non-i..d datasets. In each communication round, each client performs local training for 2 epochs using the local data, then uploads paramcters to the server for aggregation. Table 10 lists detailed hyperparameters during FL training with ResNet-18. For PriSM and baseline methods (OrigDrop and OrthDrop), we tune the initial learning rate from 0.01 to 0.2, with different learning rate schedulers. A cosine anncaling schoduler with an initial learning rate of 0.1 gives the best performance. We keep other parameters fixed (e.g., #clients, #local epochs, #rounds, batch size) to ensure they conducted in the same fair setting.
16

---

Published in Transactions on Machine Learning Research (10/2023)
Table 9: 1eeNet-118CIFARR
Module
#kernels
size
stride
Batch Norm
ReLU
Downsample
Conv1
64
3
X
Block1-1
64
3
ResBlock
x
Block1-2
3
Block2-1
64
3
1
/
ResBlock
X
Block2-2
64
3
/
/
Block3-1
128
3
/
ResBlock
/
Block3-2
128
3
/
Block4-1
128
3
ResBlock
Block4-2
128
3
Block5-1
256
3
ResBlock
/
Block5-2
256
3
Block6-1
256
3
ResBlock 6
X
Block6-2
256
3
Block7-1
512
3
ResBlock
Block7-2
512
3
Block8-1
512
3
ResBlock
X
Block8-2
512
3
Classification
10
x
X
X
Table 10: Hyperparameters for Ressee--B..Diit/CIAAR-- on homogeneous clients
#clients
#samples
distribution
augmentation
Datascts
100
500
i..d, non-i.i.d (a = 1,01)
flip, random crop
# #Rounds
#local epochs
batch size
#active clients
smooth factor
Training
1000/400 (DeiT)
2
32
20
2.5/4 (0.2 sub-model)
Optimizer
Momentum
wd
initial Ir
scheduler
Opt
SGD
0.9
0.0002
0.1
cosine anncaling
CNN /FEMNIST homogencous 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 II lists detailed hyperparameters during FL training with CNN.
Table 11: Hyperparameters for CNN/FEMNIST on homogeneous clients
#clients
#users/client
distribution
augmentation
Datasets
100
10
natural non-i.i.d
None
# #Rounds
#local epochs
batch size
#active clients
smooth factor
Training
100
2
32
20
2.5
Optimizer
Momentum
wd
initial ir
scheduler
Opt
SGD
0.9
0.0002
0.01
cosine annealing
17

---

Published in Transactions on Machine Learning Research (10/2023)
ResNet-18 CisNee--8/CIFARRE /CFAAR-1 on heterogeneous clients. We adopt the same setting as in Table I0 except the lact that clients might vary in computation and communication capacity. Therefore different model might train sub-models with different sizes (See Sec 4.3 in the main paper).
LSTM/IMDB on homogeneous clients The LSTM model used in FL training is detailed in Table g During training, we simulate 100 clients, in which each client assigned 375 training samples. We create local datasets with two different distributions using the same method in CIFAR-10: iidd and non-iid (a 0.1). Table 12 list detailed hyperparameters for training LSTM/IMDB.
Table 12: Hyperparameters for LSTM/IMDB on homogeneous clients
#clients
#samples
distribution
augmentation
Datasets
100
375
i.dd, non-iidd (a 0.1)
None
# Rounds
#local epochs
batch size
#active clients
smooth factor A
Training
300
2
32
20
Optimizer
Momentum
wd
initial Ir
scheduler
Opt
SGD
0.9
0.0002
0.1
cosine annealing
A.3
More Ablation Study
A.3.1 Effects of the number of local epochs.
To investigate the effects of the nunber of local epochs the final server model accuracy, train ResNet-18 in homogencous client settings. The sub-model trained on clients varies from 0.2x to 0.6x. We also train full model as baseline. The common training hyperparameters are the same as in Table I0 We fix the number of total iterations as 2000, namely, Round x Local epochs = 2000. Figure 10 shows final server model accuracy on i.i.d and non-i.i.d datasets with = 0.1. We observe that while final server model accuracy decreases as the number of local epochs increases, the accuracy gap between sub-model and full-model training also slightly increases. One potential reason is that the total number of sampling decreases in FL training with more local epochs. As result, orthogonal kernels are under-trained. Such observation also aligus with the intuition discussed in the main paper that all orthogonal kernels should be trained.
66
Accuracy (%)
Accuracy (%)
ocal epnchs
ooall epochs
(a) L..d datasets.
(b) non-i.id datasets (a= = 0.1).
Figure 10: Effects of the number of local epochs on training performance (RCSNet-IB/CIFAR-1)))
A.3.2 Effects of the number of active clients.
We follow the same settings as in Table [0 except that the number of active clients varies from 10 to 30 in each communication round. Similar above, we vary the sub-model trained on clients from 0.2x 0.6x. Figure [1 shows the final server model accuracy on i.id and non-i.id datasets with a 0.1. On i.i.d datasets, FT. training with different. number of active clients does not significantly change the training
18

---

Published in Transactions on Machine Learning Research (10/2023)
Accuracy (%)
Accuracy (%)
12
#Active clients
#Aclive clients
(a) i.id dalasets.
(b) non-i.id datasets (ar = 0.1).
Figure 11: Effects of the number of active clients on training performance (ResNet-18) CresNet-18/CIFAN /CIFAR-10).
performance. Furthermore, the accuracy gap between sub-model and full-model training is also almost the same, even though a less frequent sampling process is performed training with fewer clients. However, on non-i.i.d datasets, the final accuracy increases with an increasing number of participating clients. Further, the performance gap between sub-model and full-model training also shrinks as more sampling processes are performed in each round. In practical edge settings, with plenty of devices such smart-home devices connected, the performance gap can be possibly further shrunk.
19


Reviews:
[{"review_id": "e91927V87f", "reviewer": "Reviewer_o159", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n- The proposed PriSM method is easy to follow, and the presentation of this paper is concise and clear.\n- Adapting FL methods to hardware-heterogeneous environments is of great importance and can potentially attain a lot of attention.\n- The final results (especially for the end models' accuracy) are promising.\n\nWeaknesses:\n- My major concern is that it's still not clear to me why PriSM can significantly outperform the ordered dropout method, e.g., FjORD even after reading the explanations in Section 4.4. Overall, frankly speaking, PriSM is really similar to FjORD (except for the sampling scheme). So I wonder, after more efforts on careful hyperparameter sweeping and optimization, would the performance gap between PriSM and FjORD become smaller?\n- Most of the experimental results are focusing on CNNs running for vision tasks. I wonder if more language experiments can be performed, e.g., BERT/GPT fine-tuning on generative tasks.\n- It looks like there is a smarter way to adaptively select r during the entire model training, e.g., gradually using smaller rs during training. So I wonder if the method proposed in automatic low-rank training can be used in PriSM [1].\n- I wonder how the proposed PriSM method compare to the depth-wise model splitting approach proposed in [2].\n\n[1] https://proceedings.mlsys.org/paper_files/paper/2023/hash/b38f833dd45718cc414bd9b2a1d86c9a-Abstract-mlsys2023.html  \n[2] https://arxiv.org/abs/2303.04887\n", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper presents PriSM, a novel method aimed at enhancing the efficiency of heterogeneous federated learning algorithms. The central premise of PriSM involves using the singular values of each layer's model weight to guide model sampling, which could range from hidden neurons to convolutional kernels. The methodology was put to the test through numerous federated tasks and models, with the results underscoring the potency of PriSM. In comparison to prevalent benchmark methods that utilize low-rank approximation and structured/ordered dropout, PriSM demonstrated superior performance.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "- Results of more efforts careful hyperparameter sweeping and optimization on both PriSM and FjORD are expected to be shown in the paper.\n- More results on larger-scale language models, e.g., BERT and GPT are expected to be presented.\n- Some important works on low-rank model training (though not directly related to FL) shall be discussed.\n\n[1] https://proceedings.mlsys.org/paper_files/paper/2023/hash/b38f833dd45718cc414bd9b2a1d86c9a-Abstract-mlsys2023.html  \n[2] https://proceedings.mlsys.org/paper_files/paper/2021/hash/94cb28874a503f34b3c4a41bddcea2bd-Abstract.html", "broader_impact_concerns": "There is no concern about the ethical implications of the work.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "HNBDlFqlIY", "reviewer": "Reviewer_N2vX", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "**STRENGTHS.**\n\n* The paper is written in a concise style, with clear statements that are easily identifiable in the abstract.\n\n* The idea of decomposing the parameters using SVD (Singular Value Decomposition) appears promising for mitigating the memory and computing bottlenecks faced by edge clients.\n\n* The paper discusses the choice of hyperparameters, including insights on when to recalculate the orthogonal decomposition during training.\n\n**WEAKNESSES.**\n\n* Random assignment of orthogonal vectors may not be suitable to every client. Is it possible to have local parameters that diverge rapidly during training? Unless the parameter matrice is well approximated be few vectors (the spectrum decays rapidly), does this method work with very small sub-models?\n", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The authors focus on the cross-device federated learning setting. They propose to decompose parameters of neural networks into sub-models, this method facilitates the participation of weak clients by reducing their model size. To achieve this, the global model is decomposed into a sum of rank-1 matrices, and each part is randomly assigned to individual clients. Additionally, the authors present numerical illustrations of their algorithm, providing practical insights on its performance.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "**Clarifications:**\n\n* In the paragraph below the ``Motivation'' section, there seems to be overlapping definitions between the dimensionality of variable $X$, which is $M \\times H \\times W$, and the network parameter $W$.\n\n* The step \"Sub-model aggregation\" and equation (5) are little unclear. I don't understand the definition of the brackets in  $[\\text{Sum}(c\\in \\mathcal{C}) \\alpha_i^c u_i^c; \\hat{u}_i]$.\n\n* In Algorithm 1, I think the ``orthogonal space refresh'' step is written twice (step 2 and step 10).\n", "broader_impact_concerns": "I don't see any major concerns, the paper explores a new methodology for federated learning under resource constraints.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "kcnpf1BhOi", "reviewer": "Reviewer_phjK", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Pros:\n1. The paper is clearly well-written and easy to follow. The background and method have been explained clearly in detail.\n2. The paper conducts many experiments in different settings and has some empirical analysis.\n3. The experimental results show the proposed method could significantly achieve less performance drop compared with other baselines.\n\nCons:\n1. The proposed method's novelty is not very significant. As mentioned by the related work, using the principle component in federated learning has been used in federated learning several times. The main improvement is using an important sampling method other than the dropout.\n2. The proposed method mainly focuses on experimental evidence to show its effectiveness. It lacks a theoretical analysis of the convergence and how well the main principle component could approximate the original weight.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The paper proposes a new method in the federated learning setting where clients have limited computation resources and don't share the intermediate result or labels with the server. Specifically, the method includes first decomposing the model parameter using SVD and then ranking the singular value. Then every client just samples according to the rank based on its computation budget and only updates that main component. The experimental results on several datasets, and domains in both the homogeneous and heterogeneous settings show the proposed method could achieve less performance drop compared with other baselines. It has also included several ablation studies to show the importance of sampling, different sampling strategies, and runtime. ", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "1. Have some basic theoretical analysis to show the convergence with certain approximation errors. ", "broader_impact_concerns": "I don't see any ethical concerns. ", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: This work considers decomposing large models in cross-device federated settings into smaller sub-models, which are low-rank approximations of the large model on the server. 

All reviewers found the empirical results of the proposed method (PriSM) compelling, particularly when considering highly resource-constrained regimes where each client can only maintain a low capacity model. Reviewers phjK and 0159 mention concerns around a lack of theoretical guarantees and novelty of the approach relative to prior work that similarly explores low-rank training in both federated and non-federated settings. 

While I tend to agree with the reviewers that the techniques themselves are somewhat limited in novelty, I also agree with the authors regarding the significance of the proposed techniques: The authors have carefully compared to prior work in their experiments, and the proposed method seems to clearly lead to substantial empirical improvements, particularly in resource-constrained settings. I believe the work is thus worth showcasing at TMLR even in light of the mentioned limitations around novelty and lack of theoretical guarantees. 

However, I would encourage the authors to carefully incorporate the discussion from the rebuttal period around differences between this method and prior work (particularly HeteroFL and FjORD) in their revision, as this is helpful for more clearly putting the work in context and understanding the methodological modifications that lead to improved empirical performance.
Decision: Accept (poster)