On Sampling Strategies for Spectral Model Sharding
Authors: Denis Korzhenkov
Christos Louizos
Venue: NeurIPS 2024 poster (2024)
Keywords: federated learning
singular vector decomposition
heterogeneous devices
Primary area: deep_learning_architectures
TL;DR: We propose new sampling strategies and techniques for local training of low-rank factorized neural networks for the scenario of federated learning on heterogeneous devices.

Abstract: The problem of heterogeneous clients in federated learning has recently drawn a lot of attention. Spectral model sharding, i.e., partitioning the model parameters into low-rank matrices based on the singular value decomposition, has been one of the proposed solutions for more efficient on-device training in such settings. In this work we present two sampling strategies for such sharding, obtained as solutions to specific optimization problems. The first produces unbiased estimators of the original weights, while the second aims to minimize the squared approximation error. We discuss how both of these estimators can be incorporated in the federated learning loop and practical considerations that arise during local training. Empirically, we demonstrate that both of these methods can lead to improved performance in various commonly used datasets.

On Sampling Strategies for Spectral Model Sharding
Denis Korzhenkov Qualcomm AI Research* Amsterdam, The Netherlands akorzhen@qti qual comm. com
Christos Louizos Qualcomm Al Research" Amsterdam, The Netherlands clouizosßqti. qual comm. com
Abstract
The problem of heterogeneous clients in federated learning has recently drawn a lot of attention. Spectral model sharding, i.e., partitioning the model parameters into low-rank matrices based on the singular value decomposition, has been one of the proposed solutions for more efficient on-device training in such settings. In this work, we present two sampling strategies for such sharding, obtained as solutions to specific optimization problems. The first produces unbiased estimators of the original weights, while the second aims to minimize the squared approximation error. We discuss how both of these estimators can be incorporated in the federated learning loop and practical considerations that arise during local training. Empiri- cally, we demonstrate that both of these methods can lead to improved performance on various commonly used datasets.
1 Introduction
Due to increasing concerns about user data privacy, federated approach becomes viable alternative for common centralized methods of machine learning [25]. Unlike centralized training on powerful servers, the federated setting faces a variety of user devices like personal computers, mobile phones, tablets, etc. which may have different constraints terms of memory footprint, spare CPU resources, battery levels, and many others [34]. Therefore, in practice, the scenario where each user trains the same model may be impossible.
To tackle this problem, several methods which carefully create smaller sub-models from a larger global model were proposed [6, 17, 1, 44]. The general recipe for each communication round is composed of three steps. Initially. the server receives from each client the information about its capabilities and sends a sub-model of the corresponding size to that client. Afterwards, each of the clients update their respective sub-models for specific number of steps. Finally, the server collects the fune-tuned models from each client and aggregates them into single global model.
In order to create sub-models, the recently proposed method called PriSM suggests to conduct singular value decomposition (SVD) of the weight matrices of affine layers and randomly sample the terms from this decomposition [32, 33]. The number of sampled terms is hyperparameter which depend on the capacity of the client. Despite the fact that PriSM has shown promising results on several datasets, its sampling design remains heuristic.
In this work, we show that by introducing certain first principles, better strategies can be obtained via solutions to specific optimization problems. More specifically, our aim is to obtain the inclusion probabilities for the SVD factors in way that minimizes the expected, under these probabilities, squared approximation error. isolation, this requirement leads to the deterministic top-n sampling. theorem implies [8]. However, we propose additional principles
the Yorkart-Youn-MMee
"Qualcomm Al Research, Qualcomm Technologies Netherlands B.V (Qualcomm Al Research an initiative of Qualcomm Technologies, Inc.).
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

which lead to two novel strategies, respectively. The first is inspired by the commonly used inverted implementation of dropout regularization [37. 24], and searches for the solution in the family of unbiased estimators. The second one takes into account all the (participating in a given round) client specific sub-models, and aims to reconstruct the full weight matrix from all these observations (instead of considering each client separately). Notably, the solutions for both strategies can be presented in closed form. Furthermore, since the sampling of sub-models takes place on the server side, the computation of the optimal probabilities in these settings does not require any additional calculation on potentially weak clients. Finally, we employ such procedures for sampling without replacement which preserve given inclusion probabilities [39].
2 Related Works
Heterogeneity in its different aspects is ubiquitous in real-life setups for federated learning and is therefore a long-standing topic of research [11, 45]. Probably the most well-studied issue is the non-i..d. distributed data samples, available on clients' devices [19]. However, the deployment of federated systems also faces the variety of devices which often have different computational constraints [34]. This makes training the same model for each client infeasible and motivates the search for new solutions.
Many popular approaches to tackle this kind of heterogeneity follow one of two directions. Methods based on knowledge distillation in most cases client models to train new predictor the server side [27. 14]. However, these techniques typically require a separate dataset on the server, which may be a serious limitation in practice. In contrast, partial training- solely on the client side. To fit the clients' constraints, they sample sub-models which have size that allows them be trained on the corresponding device. After local training, these sub-models aggregated into global model. Notably, this allows for final model size that is otherwise infcasible to train in any of the available clients.
approaches do all the training
Many such methods conduct the sub-sampling in the original model space, e.g., by selecting specific output dimensions each layer. While HeteroFL [6] simply selected the top-left corner of the weight matrix of the required size, methods such as Federated Dropout [4] and FjORD [17] argued for randomized selection. FedRolex [1] replaced randomness with rolling windows, with the aim of covering all possible sub-models in shorter period of time.
Recent advances in using factorized low-rank models, both for training from scratch [22, 21, 38, 43] and for fine-tuning [20]. found applications in the realm of federated learning [2, 46]. E.g. FLANC used this idea to address heterogeneous computational constraints by representing weight matrix as a product of two factors, one fixed for all clients and the other having a varying size based on the capacity of each client [30]. Such approaches for model sub-sampling can be generally useful, given the fact that the (stable) rank of the model weights tends, in some cases, to decrease during training [41, 43]. In a similar spirit, FedHM [44] applied truncated SVD to create the sub-models, while PriSM [32, 33] employed a random procedure for choosing the SVD factors. In our work, we propose novel strategies for the selection of sub-models and their training in such scenarios and demonstrate their advantages.
3 Method
3.1 Preliminaries
In the case of cross-device federated learning. low-end clients may not be able to perform the computations with the full weight matrix W in all the layers due to various constraints. A possible remedy for this problem can be to use an approximation for W. In this work we consider a particular case of such approximations, namely sub-sampling the terms from the singular decomposition (SVD) W = - Auuuo," where both sets of column vectors {u; E Ro ), and {u; E R ); are orthonormal, and the singular values are sorted in non-increasing order X 2 X2 2 > AN > 0. Due to the usage of singular decomposition, we name this type of sub-moded sampling as spectral model sharding.
More formally, let 0 <r S be the keep ratio which depends on the client's 's capabilities, then the goal is to sample the [Nr] terms from the full sum. Following prior works, keep ratio r is shared

---

from the server and performs local training with the factorized weight matrix ((o for some predefined number of epochs. During training (Ce) and ((e) updated while the diagonal matrix k(c) which has values of (e() on the diagonal is kept frozen. When the local training is across layers [33, 32, 30]. With sufficiently low value of r, this reduces both communication costs
and computational complexity. We introduce vector of indicators 2 (z1,..., =N) E {0, 1)" such that Σ 2, = n, and the corresponding sub-layer is created with weight matrix W,
W
Σ wanrAmii
Σ WI, UI,
where {wr E [R are scalar auxiliary multipliers, and I is the set of selected indices, I {i: 24 = -1.
If the sampling of indicators is random, we denote the corresponding marginal inclusion proba- bilities as Th = Pr(z; = 1). The joint distribution of indicator variables is interchangeably referred to as p(z) p(I). Wherever the expectation symbol E is used in the text, expectation w.r.t. p(2) is assumed. it is necessary to emphasize that the particular estimator, multiplier, vector, etc. used by the client e, this is shown with the client index in parentheses, e.g. W(())
Based the previous models of FedHM [44] and PriSM [32, 33], consider the following pipeline of sharded model training in the case of federated learning on heterogeneous devices:
1. In the beginning of each communication round, the server performs SVD for the weight matrices of the affine layers. Then, for each client e randomly samples a subset of indices I from the decomposition with keep ratio ,(t), W((c -j= Wyfor 1/ 7,00"
o/T The singular values are
absorbed into the vectors, namely u/ A,u,, of= VT,u Note that in both FedHM and PriSM, the auxiliary multipliers w/ (which can be used to 'compensate" for the dropped terms) were set to However, since they are actively used in our method, we include them in the description.
2. Each participating client receives the matrices ((e)
E R and
V = -(1/0)
RS as well the vector (()) (Wzjor..... (gz0))
E R"
done, the updated matrices ((e) and (e) are sent back the server.
The server aggregates each vector from the decomposition separately, i.e. u/ Ee ((()
_,, Dio where ( is the size of the local dataset of the client and u, u(()) is the
updated value of the i-th vector received from the client c. If the i-th term not selected by any client in the current round, it remains the same. Vectors {ur}, are aggregated in the same way. A Afterward, the updated full weight matrix is calculated W Σ.V Ei= uf and new communication round begins.
FedHM proposed to the terms corresponding to the largest singular values, deterministic selection instead of random sampling. In contrast, PriSM presented a randomized algorithm, which samples the terms according to the magnitudes of the singular values. It employs the NumPy [12] method numpy. random. choice with the option repl acplaceFaase, and the associated probabilities set proportional to X, where i is hyperparameter depending on the keep ratio Sampling in this manner results in a probability law similar to Wallenius' noncentral hypergeometric distribution [42, 5. 10, 9]. see Appendix C for details. Such sampling procedures make the analysis of the resulting marginal inclusion probabilities { complicated [40]. Thus, the statistical properties of the obtained matrix approximation are opaque.
3.2 Sampling Distribution
As opposed the prior works, we consider the search an optimal sampling distribution be the central task. We introduce two assumptions which we use to derive our two different strategies.
3

---

3.2.1 Unbiased Estimator
The bias and variance are inherent trade-offs for an estimator. To keep the bias in check, we propose an assumption of the unbiasedness of the sampled estimator W of the original weight matrix W, i.. EWW- W. where the expectation is taken over the sampling distribution p(Z) with the marginal inclusion probabilities (TI,..., NT) As shown in Appendix B.I, this requirement necessarily leads to the following values of auxiliary multipliers wi = n; This selection of multipliers known in literature a of inverted dropout [37, 24] which assumes dividing the non-zeroed activations by the inclusion probability training time.
estimator [16] and resembles the common implementation
With the bias of the estimator taken under control, we aim to reduce the approximation error of the resulting layer. For any input vector E R the following inequality holds
Wr Wa 2 - || W w/ [ | - [ |W- w -
(2)
where 1/-11/- is squared Frobenius norm which equals the sum of the squared singular values. Thus, we can control the upper-bound of the expected error by searching for sampling distribution p(T) that minimizes the Frobenius discrepancy
W wI
min P(Z)
Theorem 3.1. For unbiased estimator W of the rype specified in Eq. (1) and consisting of terms, the Frobenius discrepancy can be expressed in terms of the marginal inclusion probabilities
E
W
W
(4)
and the optimal set of inclusion probabilities has the following form
ifi S t, ifi>t
(n- Emmetr de
(5)
for = ,,...,N,, where t E 11.
Proof. See Appendix B.1.
As follows from Theorem 3.1, to find the true arguments the minima of Eq. (4), one needs to sweep over n possible values of evaluate the Frobenius discrepancy and select the with the minimal discrepancy. Note that this procedure takes place on the server side and therefore does not require extra computation on clients' devices.
3.2.2 Collective Estimator
Due to the well-known bias-variance trade-off, unbiased estimators in practice can have too large variance. This motivates us consider another perspective. Consider the simplifying assumption where C' clients participating in the current communication round share the same number of terms n in their respective estimators. We can target 'reconstructing' the full matrix W from the whole set of
i..dd observations [fr(e)] 11e5s The simplest way of obtaining such reconstruction is averaging,
ie. W = Wee)) and we refer to the matrix W as the collective estimator. To ensure that our reconstruction is accurate, we aim to find the optimal sampling distribution as well as the set of auxiliary multipliers,
min (wth
Ep(r)
||w - |
(6)
Theorem 3.2. For a collective estimator W, the average value of the Frobenius discrepancy can be expressed in terms of just the marginal inclusion probabilities
E W 11
Σ Xurts (-2
(22 C
Wi
WiiTi (C
N
X2
C
1=1

---

and the optimal set of inclusion probabilities and auxiliary weights has the following form in the case
of C
1,
if
C
(n-t+ ).,
W/ 1 T. (C
Ti
ift u,
(8)
AE
fiittu,
for = I.....,,, N, where and u are integer constants such that S S n, and 0 s u < - t. For
C' = 1 the optimal values are Wi I[i S n).
Proof. For C' the proof immediately follows from the Yoetart-ooug--Wae Mirskky theorem [8] which states that truncated SVD provides the best low-rank approximation of the original matrix in terms of Frobenius discrepancy. Since all the estimators defined by Eq. (1) are low-rank approximations, is impossible obtain average error which is strictly lower than the error of truncated SVD, For the case C 1, see Appendix B.2.
Note that for a group consisting of single client (i.e. C - 1) the optimal solution coincides with lop-n sampling used in FedHM method. For larger groups, similarly to the unbiased estimator, sweeping across all possible values of and is required on the server side. Nevertheless, this search can be performed in vectorized form.
In practice, to apply Theorem 3.2, the server can cluster the heterogeneous clients into homogeneous groups which share that same keep ratio r. and compute a specific sampling distribution for each group. Such clustering is often employed by methods targeting diverse clients, ef. [30].
3.2.3 Joint Distribution and Sampling
The strategies derived in Secs. 3,2.1 and 3.2.2 provide the values of marginal inclusion probabilities but do not specify the joint distribution, i.e. the co-occurrences of sampled terms. Tillé [39] presented a systematic survey of multiple sampling procedures which preserve the given inclusion probabilities. Notably, each of the methods, while keeping the mean vector E (21. , 2N) (#11-... Tx) fixed, has different covariance matrices. We follow the authors' recommendation and stick with the Conditional Poisson scheme (CPS) for our experiments. The joint distribution achieved with this method has the maximum entropy among all distributions over samples of a size of exactly n with the same marginals. In theory, such design should allow the model to better explore the interaction between the singular vectors during training. The reference implementation of this sampling algorithm provided in the accompanying package" for [39].
3.3 Local Training
While the derivations in Sec. 3.2 provide an optimal (in certain sense) estimator Wi(e) of the full weight matrix W for the client e before the local training starts on that client, they do not provide any guidance how optimize the received sub-model.
Auxiliary multipliers. As indicated in Sec. 3.1, we follow Khodak et al. [22] and Niu et al. [33] and absorb the square root of the singular value X; into both left and right singular vectors, "/ = -veu, v/ = VAe, During early stages of our experiments auxiliary multipliers Wi in the same manner. However, this led to unsatisfactory results. Note that, except for the corner cases n. {0, 1). the multiplier w; is in inverse proportion to X; and therefore the magnitude of their product is almost independent of i. Therefore, after the product vw; VX; is directly 'baked' into the sub-model weights, is harder for the clients to distinguish between important and uninformative terms.
explored the absorption of the
Instead, we treat the auxiliary multipliers as scaling factors, frozen for the current round of local training. Hence, the effective weight matrix of the sub-model is trained locally in the factorized form W = UNVT with learnable terms U and V. This could introduce an issue of stale multipliers: during the course of local training the original meaning, along with the guarantees on the desired estimator properties, of the auxiliary multipliers is gradually fading away. However, this was not
http::///raan..ppjjeee.e
5

---

something we observe in practice and any attempts to mitigate this (e.g., with gradual decrease of the scaling factors) did not provide consistent improvement for the model performance.
Effective learning rate. Nevertheless, the chosen way of incorporating the auxiliary multipliers into the model still has a drawback. Namely, the gradient of the scalar loss function w.r.t. the factors
AL is expressed OL OL VΩ,
UR. This means that the auxiliary multipliers W, aV aa
AL
at aw
also affect the effective learning rate of the corresponding column vectors vi, and v when first-order optimization methods like SGD are used. For the collective estimator it follows from Eq. (8) that the values of the multipliers are bounded by the size of the group. 1 S wi S C. However, for the unbiased estimators there is no such upper bound as Wi = म. Moreover, for small values of the keep ratio r there necessarily should exist relatively large values of प. This is because
E Ei= WI, = EE = = u,; NN
(9)
and one needs to approximate N with n 2 rN randomly chosen terms. This can lead to excessively high values of the effective learning rate. Based on that, we adjust the nominal learning rate for each vector or, equivalently, override the value of the gradient before conducting the optimization
BL 8L min 1, I for some threshold N 1. With this 8L aL min (1, ).). dr, de
step, Du, Bi, modification, the effective learning rate cannot be more than times higher than the nominal. In all reported experiments we use 7 10. Worth noting that equalizing the efficient learning rate with 7 = 1 turned out very detrimental for the performance, as we found early in our experiments. This coincides with the intuition behind some of adaptive optimization methods like AdaGrad [7] that perform larger updates for the parameters which are in charge of 'less frequent" features.
wi
Frobenius decay and momentum, As proposed in other works [22, 33, 44] we use Frobenius weight
decay (FD) during local training, i.e. we apply an additional loss function proportional to the || |
for each effective weight matrix. The weight of FD in the resulting loss function is set to 1 10 Additionally, confirming findings of Niu et al. [33], we found it necessary to use plain SGD with momentum weight of 0.9 during local optimization of the sub-model weights.
4 Experiments
4.1 Experimental Setup
Datasets and models. To evaluate the proposed strategies, we conducted experiments on several datasets which are usually employed test federated systems. In case of artificial simulation of clients, we follow [19] and sample the prior probabilities of classes for each client from the Dirichlet distribution Dir(ap) where p is the vector of class proportions in the original dataset and is hyperparameter that controls the amount of non-i.i.d. in the data splits. Note that this notation is different from the one used in some of the recent papers, e.g. [44, 32, 33].
Since in our work we do not aim to reach new state-of-the-art in the field and instead focus on analyzing the relative performance of the proposed strategies of spectral sharding, we have chosen to test all the methods in the presence of high data non-i.i.d.-nes. between clients.
For CIFAR- 10 [23] we split the data with a = 1 and conduct experiments with a ResNet- model [15] with the kernel size of the first convolutional block equal to 3 3 and the normalization layer replaced with GroupNorm [18]. For Tinymmageee [31] we use = 10 and a compact transformer (CCT) model [13], namely CCT-6/3x2. For CIFAR-1O0 [23] the two-staged Pachinko allocation method (PAM) [26] is used: for each client, at first parameters of the multinomial distribution over the twenty coarse labels are sampled with a = 1, and afterwards the distribution over the coarse-to- labels with a = 10. On this dataset we train both the ResNet and CCT models described above. We select Shakespeare [29] as an example of a dataset with a natural data split over clients. We train a tiny transformer model with three GPT-2 [35] blocks the task of next character prediction and report the performance in terms of accuracy in accordance with prior works [36].
For all the architectures we do not decompose the very first and very last layers of the model [22, 33). However, unlike PriSM, we decompose all the affine layers inside the attention blocks, not just the fully connected ones [43]. Also, we follow Niu et al. [33], Wang et al. [43] instead of Khodak

---

et al. [22] and decompose a convolutional layer to sequence of a regular convolution with output channels followed by 1 x convolution. For image datasets we employ per-image pixel value standardization similarly to [36] but do not apply random cropping.
Training. We train all models from scratch with cosine annealing learning rate scheduler [28]. The initial value for learning rate is 0.1 for CIFAR-IO, 0.05 for CIFAR-100 and Tinylmagenet, and 0.1 for the Shakespeare data. The client's batch size equals 32, 64, 128, and 10 respectively. All experiments are run with three random seeds which also affect data splitting between clients, if applicable. Standard deviation is reported in all tables and plots based on those runs.
Federated setup. While Shakespeare dataset naturally contains 715 clients, we simulate 100 clients for all other datasets. We randomly sample 10 clients for each communication rounds which results in participation ratio of 10% on image datasets and about 1.4% for Shakespeare. In each communication round all participating clients train their sub-models for two local epochs. The total number of local epochs (e.g., number of local epochs per round x number of communication rounds) equals 2,000 for CIFAR-10, 3,000 for CIFAR-100, 5,000 for TinyImagenet and 3,000 for Shakespeare.
Baselines. As stated above, we focus on exploring different strategies of sampling sub-models when a particular approach of training on weak and, possibly, heterogeneous devices was chosen, namely the approach described in Sec. 3.1. Therefore, we compare our presented strategies, denoted as Unbiased and Collective, with the Top-n sampling proposed by [44] and the PriSM method introduced in [32, 33]. We copy the value of the hyperparameter k required for sampling in PriSM from the original implementation. In detail, depends on the keep ratio r: k if S 0.2 and k = 2.5 otherwise.
To decouple the sampling strategies themselves from other training details discussed in Sec. 3 we additionally introduce modifications to the PriSM strategy: motivated from our unbiased method, we train it with auxiliary multipliers to allow for (approximate) unbiasedness of the estimators, i.e. wi = =,, (this strategy is dubbed as PriSM + Wallenius), and clip the effective learning rate (PriSM + Wallenius RlipLR). Exact computation of the mean vector for the Wallenius' distribution occurring in these two strategies is very time-consuming for the ranks N of weight matrices used in practice. For that reason we use the approximate algorithm² from [9].
Finally, we explore the simplest option of compensating for the missing terms of the estimator, namely, introducing a scaling factor that keeps the Frobenius norm of the estimator equal to that of
E Xg (c) Σ 2 the full matrix, i.e. w
for all i. This modification is marked as +Sealed.
All the considered methods require equal communication costs and equal amount of computation on the client side per round, since the overhead caused by the vector of auxiliary weights introduced by novel strategies is negligible.
4.2 Results
4.2.1
Main Results
We report the accuracy of different strategies in all the datasets considered in Tab. 1. In our experi- ments we mostly explore weak clients with low keep ratio. The results are provided for two setups with homogeneous clients (r 0.1 and 0.2) and one setup with two groups, namely 60% clients having a keep ratio of 0.2 and 40% of 0.4. Also, for reference purposes we train a conventional full model on all the clients without sharding, as well an wee-paarmeterized model with keep ratio r = 1. Similarly to what was reported by Khodak et al. [22], we observe that overparameterization can provide a better performance than usual training.
As the results demonstrate, the preservation of the Frobenius norm (+Scaled) the least effective and stable among the considered modifications. It is often detrimental for the Top-n strategy, and although sometimes improves the performance of PriSM, the PriSM modification based on unbiasedness (+ Cliip is generally more successful. In certain cases, it even achieves the best accuracy among all models, e.g. see the results of ResNet on CIFAR- 100. However, overall the 'unbiasedness' has an inconsistent effect on PriSM: it is usually beneficial in case of ResNet but is less helpful or architectures. As has already been noted, deeper analysis of such
even harmful for wwtpppeettppppaeeeeepppppppppppp BiasedUrn
7

---

Table 1: Accuracy achieved by different strategies under limited computational budget. Our Unbiased and Collective strategies outperform the vanilla Top-n and PriSM baselines on all datasets except for Shakespeare, although the gap is not significant for that dataset. In addition, the modifica- tions proposed for local training allow PriSM to significantly improve its performance for ResNet architecture, and sometimes even surpass other methods in the current setting.
STRATEGY
CIFAR-IO
REsNET
TINVMAGENTT
CCT
Rus
CIFAR-100
CCT
SHARESPPARE TRANSFORMTER
rutler
ToP-n
67.44
21
47.34g 10
Too- SCALED
23.01
22.29 m
13 m
33.96
36. 14g0
PRISM
71.15a3
35.81
18. 45
31.0020.00
46 46.91.1.35
PRISM + SCALED
76.03
55. 10
28 65
12. 1igess
45.34
PRISM WALLENIUS
77 77.53
37.48
37 1os
42.04.1.11
45
PRISM WALLENIUS CLIPLR
82.42
40 36
35.50.1.24
48 14.24
80.57
37.57
46.01
48.253a.
COLLECTIVE
59 00.27
38.72 go.0
37 87.85
49.98,1.7
48.6440.50
Keep ratio 0.2
TOP-n
78. 37 a0
35.97
61.70gi sa
51. 11 g0.50
TOP-n++ + SCALED
80,75a1
31, 00 g0
35 25
37.68 66.66.277
45.8 45.8220.3
PRISM
62.59
38 66
33 76
62.04g
49 40.632n
PRISM SCALED
81.6315
60.30):044
10.80 1..8
51.64:130
PRISM WALLENIUS
82. 82.15 14
11
18.20
45 45.3520.50
PRISM WALLENIUS CLrLR
81.0011 24
38.12 18
38.54
49.44.0.35
UNBIASED
84,3331114
41.28 10.51
10.34
50.1220.55
COLLECTIVE
85 20.00....
41.12
47 14
64.31
50.05go.1
60% client have keep nallo 0.2 and 40% clients have
01.4
Top-n
64 41.432mm
36.80
41 41:34g1
51 35
Ten-n-4 SCALED
12
12:24.aaa
47.00
30.1
82:37 10
31.25
23.04:22
20.4
29.75
PRISM SCALED
82 78.8 14
34.80 28.
28 22
12.46
49.69
PRISM+ WALLENIUS
78
34.80-1..o
34.051222
34.28-1100
38.21:2.2
PRISM WALLENIUS + CLPLRR
82 82.9510054
39 30.11
32 88
10.07
49, 49.4519.55
UNRIASED
09
40 78..1
44 61....
54. 16
51.0550.43
COLLECTIVE
39.30
45.74
02.02.0.0
50.741..1
Keep ratio
41.28
03.45.o.1
09.25
32.10.1.21
No sharding
87. 11
38.08.0.0
63.3520.05
60.05.00.7
52.89
Table 2: Longer training on CIFAR-100. Table 3: Influence of the clipped learning rate.
Despite our 'unbiased' modification of PriSM demonstrated the best accuracy in case of limited computation budget, the more explorative Collec- tive strategy closes the gap in performance if the number of communication rounds is increased.
In our experiments, we find that clipping the ef- fective learning rate is beneficial for the Unbiased strategy in case of all architectures and values of the keep ratio r. Without clipping the perfor- mance drops consistently.
STRATRGY
COMMUNICATIO
1,500
STRATEGY
RESNET
CCT
ToP-n
21. 21.13
92.22s1. 94
Keep ratio 0.1 UNBIASED
TOP-L + SCALED
14.18gam
39.688g in
12
45.01_1.07
PRISM
PRISM SCALED
26 1777
18-4fgr
57. $7.08
51.85g
UNBIASED w/o CLIPLR
30.14+1.88
32.46-0.71
PAIS WALLENIUS
37.70
19 19.8140
Keep ratio r = 0.2
PRISM + WALLENIUS CLIPLR
40.36gn-
60.02ga
UNBIASED
46.34 46.34go.22
53.30_0.81
COLLECTIVE
UNBIASED
85.12
87.83
59.641000
20.
UNBIASED w/o CLIPLR
43.11
46.67=0.85
behavior is complicated due to unclear statistical properties of the PriSM sampling procedure and we leave this fo future work. One possible reason is that the approximation given by Wallenius' distribution may be less precise in some cases.
Our strategies presented in Sec. 3 in most cases outperform the baselines for all datasets but Shake- speare where Top-n demonstrates superior results, although the difference in performance is not significant. As the results imply, the Collective strategy is more suitable if the keep ratio is small, while for larger r the Unbiased strategy seems more preferable. This may be explained by the negative influence of too large auxiliary multipliers which also amplify the impact of less informative terms during the forward pass; see Sec. 3.3.
Interestingly, in all experiments, the novel sampling strategies significantly improve the results of baselines if those had large accuracy gap compared the full model. For the Shakespeare dataset, in contrast, all the low-rank methods achieve performance close to the full model even for small values of the keep ratio. Therefore, seems that the performance on Shakespeare is determined more by the properties of the data split and the corresponding federated setup than by the training strategy.
4.2.2 Discussion
Explorative and exploitative strategies. To understand the behavior of the strategies better, we propose tool named average normalized marginal entropy (ANME). For each decomposed layer with
8

---

MSN
MSN
Waterta
13-
10
(a) Accuracy of the global ResNet model on the test (b) Average normalized marginal entropy (ANME) subset of 1IFAR-100.
shows how diverse cach strategy is.
Communication efficiency. The original PriSM method is too explorative (high ANME), while our 'unbiased' modification (+ (+Wallemius) makes it the most exploitative strategy and allows to achieve the best performance in some experiments with limited computational budget.
Figure 1:
inclusion probabilities (TI, (TI,...;; TN) calculate the mean entropy of the corresponding marginal Bernoulli distributions. N Ex: O<C,<<1 (-x, log ग. - -1 πς) log (1 - T.)). It is casy to show that the minimum value of this quantity equals and achieved by the deterministic Top-n strategy. The maximum value is achieved when all probabilities are equal, i.c. T; #. Therefore, we normalize the mean entropy from above by dividing by the mean entropy of the uniform inclusion probabilities. To finish the computation of ANME, we average this normalized entropy across all sharded layers in the network.
Intuitively, low values of ANME mean that the server tends to send the same terms to all the clients, and we name such strategies exploitative. On the contrary, high value of ANME shows that the strategy is explorative and the selection of terms for each client is 'more random'. We observe the same qualitative behaviour in all our experiments (see Fig. 1b): PriSM the most explorative strategy, while the modification based on unbiasedness turns it into the most exploitative among randomized methods. The Unbiased and Collective strategies are between these two extreme types of behaviour, and the latter is usually more exploitative than the former one. Based on this observation, we check if training with more communication rounds can help the explorative strategies to perform better and report the results in Tab. 2. Indeed, all methods improve with more training and the largest improvement is for PriSM, which the most explorative. We also observe that our collective strategy overtakes our modified PriSM baseline and achieves the best performance.
Influence of learning rate clipping. In all the experiments we found that clipping the effective learning rate is necessary for strategies. Tab. 3 illustrates this observation by using models trained on CIFAR-100. Notably, Tab. I demonstrates that, empirically, such clipping is not always beneficial for the modified PriSM method. This may be explained by the fact that 'unbiased' PriSM becomes extremely exploitative. In this case, terms which could produce too large auxiliary multipliers are very unlikely to be sampled in practice, resulting in the large variance of the left-hand side of Eq. (9).
More results. Please refer to the Appendix A for additional results on model convergence, post-client updates and influence data heterogeneity.
5 Conclusion
We presented two novel sampling strategies for spectral model sharding, grounded as solutions to specific optimization problems. Alongside them, we presented techniques that facilitate local training with such methods. As shown, in number of cases these techniques can also significantly improve the performance of the strategies proposed earlier in the literature. Nonetheless, our strategies demonstrate superior performance on number of commonly used datasets in the presence of high data heterogeneity between clients. As downside, in certain cases their learning curve is less steep than that of the baselines due to their more explorative nature. We leave the improvement of the convergence speed for future work.

---

References
[1] S. Alam, L. Liu, M. Yan, and M. Zhang. Fedrolex: Modee-hhetroogeneous federated learning with rolling sub-model extraction. In NeurlPS, 2022. 1,2
[2] Babakniya, A. Elkordy, Y. Ezzeldin, Q. Liu, K.-B. Song, M. EL-Khamy, and Avestimehr. A: Federated parameter efficient fine-tuning of language models. In NeurlPS Workshop. 2023.
[3] K. E. W. Brewer. A Simple Procedure for Sampling rpswort. Australian Journal of Statistics, 176):166-117 172, 1975. 14
[4] Caldas, J. Konečný. B. McMahan, and A. Talwalkar. Expanding the reach of federated learning by reducing client resource requirements, 2019. 2
[5] Chesson. A non-central multivariate hypergeometric distribution arising from biased sampling with application to selective predation. Journal of Applied Probability, 13(4):795-797, 1976. 3,
20
[6] E. Diao, Ding, and V. Tarokh. HeteroFL: Computation and communication efficient federated learning for heterogeneous clients. In /CLR, 2021. 1,2
[7] Duchi, Hazan, and Y. Singer. Adaptive Subgradient Methods for Online Learning and Stochastic Optimization. JMLR, 12661):2121-2159, 2011. 6
[8] C. Eckart and G. Young. The approximation of one matrix by another of lower rank. Psychome- trika, 133)2111-218, 1936. 5,5
[9] A. Fog. Calculation Methods for Wallenius' Noncentral Hypergcometrie Distribution. Commu- nications in Statistics Simulation and Computation, 37(2):258-273, 2008. 3, 7, 20
[10] A. Fog. Sampling Methods for Wallenius' and Fisher's Noncentral Hypergeometric Distribu- tions. Communications in Statistics Simulation and Computation, (7(2)::41-257,2000 3,
20
[11] D. Gao, X. Yao, and Q Yang. A survey on heterogeneous federated learning, 2022. 2
[12] C. R. Harris, K. J. Millman, S. J. van der Walt, R. Gommers, P. Virtanen, D. Cournapeau, E. Wieser, J. Taylor, S. Berg, N. J. Smith, R. Kern, M. Picus, S. Hoyer, M. H. van Kerkwijk, K. Sheppard, Reddy, W. Weckesser, H. Abbasi, C. Gohlke, and T. E. Oliphant. Array programming with NumPy. Nature, 585(7825):357-367 2020. 3, 20
M. Brett, A. Haldane, J. F. del Río, M. Wiebe, P. Peterson, P.
[13] A. Hassani, S. Walton, N. Shah, Abuduweili, J. Li, and H. Shi. Escaping the big data paradigm with compact transformers, 2022. 6
[14] C. He, M. Annavaram, and S. Avestimehr. Group knowledge transfer: federated learning of large CNNs at the edge. In NeurIPS, 2020. 2
[15] K. He, X. Zhang, S. Ren, and J. Sun. Deep residual learning for image recognition. In CVPR, 2016.
[16] D. G. Horvitz and D. J. Thompson. A generalization of sampling without replacement from a finite universe. Journal of the American Statistical Association, 47(260):663 685, 1952. 4
[17] Horváth, S. Laskaridis, M. Almeida, Leontiadis, S. Venieris, and N. Lane. FjORD: Fair and accurate federated learning under heterogeneous targets with ordered dropout. In NeurIPS, 2021. 1.2
[18] K. Hsieh, A. Phanishayee, O. Mutlu, and P. Gibbons. The non-IID data quagmire of decentral- ized machine learning. In ICML, 2020. 6
[19] H. Hsu, H. Qi, and M. Brown. Measuring the effects of non-identical data distribution for federated visual classification. In NeurIPS Workshop, 2019. 2, 6
10

---

[20] E. J. Hu, Y. Shen, P. Wallis, Z. Allen-Zhu Y. Li, S. Wang, L. Wang, and W. Chen. LoRA: Low-rank adaptation large language models. In ICLR, 2022.
[21] R. Kamalakara, A. Locatelli, B. Venkitesh, Ba, Y. Gal, and A. N. Gomez. Exploring Low Rank Training of Deep Neural Networks, 2022. 2
[22] M. Khodak, N. A. Tenenholtz, L. Mackey, and N. Fusi. Initialization and regularization of factorized neural layers. In ICLR, 2021. 2, 2, 6,
[23] A. Krizhevsky. Learning multiple layers of features from tiny images, 2009. 6, 21
[24] F-F. L. J. Johnson, and S. Yeung. CS231n: Convolutional neural networks for visual recogni- tion, 2017. URL https::tp//. //sp33nn git..b..... neural-networike 2/. 2,4
[25] Q. Li, Z. Wen, Z. Wu, S. Hu, N. Wang. Y. Li, X. Liu, and B. He. A survey on federated learning systems: Vision, hype and reality for data privacy and protection. IEEE Transactions on Knowledge & Data Engineering, 35(04):3347-3366 2023.
[26] W. Li and A. McCallum. Pachinko allocation: DAG-structured mixture models of topic correlations. In ICML, 2006.
[27] Lin, Kong, S. U. Stich, and M. Jaggi. Ensemble distillation for robust model fusion in federated learning. In NeurlPS, 2020. 2
[28] I. Loshchilov and F. Hutter. SGDR: Stochastic gradient descent with warm restarts. In ICLR,
2017.
[29] B. McMahan, E. Moore, D. Ramage, S. Hampson, and B. A. y. Areas. Communication-Efficennn Learning of Deep Networks from Decentralized Data. In AISTATS, 2017. 6, 21
[30] Y. Mei, P. Guo, M. Zhou, and Patel. adspuice-adaptive federated learning with all-in-one neural composition. In NeurlPS, 2022. 3,
[31] M. A. Mnmoustafa. Tiny ImageNet, 2017. URL https: com/ com/competitions/ tiny- imaggnee.. 6, 21
[32] Y. Niu, S. Prakash, S. Kundu, S. Lee, and S. Avestimehr. Federated learning of large models at the edge via principal sub-model training. In NeurlPS Workshop, 2022. 2, 3, 6,
[33] Y. Niu, Prakash, Kundu, S. Lee, and Avestimehr. Overcoming Resource Constraints in Federated Learning: Large Models Can Be Trained with only Weak Clients. TMLR, 2023, 1, 2. 3, 5.
[34] K. Pfeiffer, M. Rapp. Khalili, and Henkel. Federated learning for computationally constrained heterogeneous devices: A survey. ACM Compul. Surv. 55(14s) 2023. 1,2
[35] A. Radford, J. Wu, R. Child, D. Luan, D. Amodei, I. Sutskever, al. Language models unsupervised multitask learners, 2019. 6
[36] S. J. Reddi, Z. Charles, M. Zaheer, Z. Garrett, K. Rush, J. Konečný, S. Kumar, and H. B. McMahan. Adaptive federated optimization. In ICLR, 2021. 6,
[37] N. Srivastava, G. Hinton, Krizhevsky, Sutskever, and Salakhutdinov. Dropout: A simple way to prevent neural networks from overfitting. JMLR, 15(56): 1929-19999.9 2014. 2, 4
[38] Y. Sui, M. Yin, W. Yang, Y. Gong, J. Xiao, H. Phan, D. Ding, X. Xu, S. Liu, Z. Chen, and
B. Yuan. ELRT: Towards efficient low-rank training for compact neural networks, 2023. 2
[39] Y. Tillé. Sampling algorithms. Springer, 2006. ISBN 978-0-387-308144-2 2, 14, 17
[40] Y. Tillé, Remarks on some misconceptions about unequal probability sampling without replace-
ment. Computer Science Review, 47:100533, 2023.
[41] N. Timor, G. Vardi, and Shamir. Implicit Regularization Towards Rank Minimization in ReLU Networks. In ALT, 2023. 2

---

[42] K. T. Wallenius. Biased Sampling: the Noncentral Hypergeometric Probability Distribution. Stanford University, 1963. 20
[43] H. Wang, S. Agarwal, P. U-chupala, Y. Tanaka, E. Xing, and D. Papailiopoulos. Cuttlefish: Low-Rank Model Training without All the Tuning. In MLSys, 2023. 2,
[44] D. Yao, W. Pan, M. J. O'Neill, Y. Dai, Y. Wan, H. Jin, and L. Sun. FedHM: Efficient federated learning for heterogeneous models via low-rank factorization, 2022. Version: 2. 1, 6, 7
[45] M. Ye, X. Fang, B. Du, P. C. Yuen, and D. Tao. Heterogeneous federated learning: State-of-the- art and research challenges, ACM Comput. Surv., 56(3). 2023. 2
[46] Z. Zhang, Y. Yang. Y. Dai, Q. Wang, Y. Yu, Qu, and Z. Xu. FedPETuning: When federated learning meets the parameter-efficiett tuning methods of pre-trained language models. In ACL, 2023.
12

---

Unblased Collective
3.5
3.0
Train cross-entropy loss
2.5
2.0
1.5
1.0
0.5
0.0
2000
4000
6000
8000
10000
Communication rounds
Figure 2: Convergence analysis. When being trained longer, the proposed strategies demonstrate the decrease of the cross-entropy loss of the global model on the training set. Unbiased strategy reaches the train accuracy of 97.0%. and Collective strategy achieves 98.4%. This serves as empirical evidence of convergence for our method.
Tope
Tope
50
30
(a) High data heterogeneity a
(b) Low heterogeneity a 10.
Figure 3: Impact of data heterogeneity communication efficiency. When data distribution between clients is closer to i..dd, the most exploitative Top-n strategy demonstrates the best training speed in the beginning, however it overfits soon. For more heterogeneous data, strategies with exploration achieve much better performance.
A Additional Discussion
Convergence analysis. It is rather involved provide the theoretical proof of convergence for any of spectral sharding methods. In practice, we observe that the models trained in this manner successfully converge even the presence of high data heterogeneity between clients. As an example, Fig. demonstrates the cross-entropy loss on the train set for the ResNet model with keep ratio = 0.1 trained on CIFAR- 10 with non-i.i.d.-ness parameter a = The loss was computed at the end of each communication round, the number of rounds was intentionally increased up to 10,000 for this experiment. is easy see that the loss successfully decreases as the training progresses. Therefore, given these empirical observations, did not focus the formal convergence analysis. Having said
13

---

that, it is an interesting direction for future work, as it can potentially uncover further improvements to our approach.
Impact of data heterogeneity, Throughout our paper, evaluate all the strategies in the presence of a relatively high degree of data heterogeneity between clients. E.g., note that a = 1 for CIFAR-10 10 roughly corresponds to only one two distinct class labels present for each client, and TinyImagenet means that each client only has access to 20-30% of all class labels.
10 for
We demonstrate the effect of data heterogeneity on the communication efficiency of the best per- forming strategies on CIFAR-10 10 in Fig. 3. When clients' data distributions are closer to i.i.d., ie. = 10, Top-n strategy is the most efficient in the beginning, however it overfits soon. For more heterogeneous data, i.c. a = 1, the proposed modification of PriSM and the novel Collective strategy the most efficient and deliver the best accuracy.
Note that the novel strategies developed in our work do not use any knowledge of data distribution between clients. Neither they rely on the presence of a public dataset on the server side. Nevertheless, we believe that incorporating such information can be beneficial for the Collective strategy and leave this for future work.
Influence of the joint distribution. To test the sensitivity of the results to the choice of the joint distribution of the sampled indices while keeping the marginal distributions frozen, compare the selected CPS sampling against two other methods, Brewer's [3] and Minimum support [39]. The advantage of those methods is that they allow for significantly faster sam- pling than CPS. Perhaps surprisingly, we do not notice any significant difference in our ex- periments, as reported in Tab. We suggest that the effect of high entropy which is charac- teristic for CPS may be shown on more large-scale problems and leave this for future work.
Post-client updates. To analyze the post- gradient updates, we followed the optimization path {0:)²= of the server model under each strategy with T communication rounds and cal- culated the cosine similarity between the server model update Satardingsstratege Og under that strategy and FeddAgg under FedAvg method with full-size model on each client. Subse- quently, the model weights were updated ac- cording to the selected strategy: 0,+1 = 0, - strategy The results of these exper- iments are provided in Fig. 4. As shown, there a notable connection between the relative per- formance of each strategy and its "alignment" with FedAvg updates. The deterministic Top-n strategy lags behind its randomized counterparts in terms of accuracy in the considered setting (see Tab. 1), and similar observation is made for its update directions.
PMSM WWil CipLA
30c 450
Figure 4: Comparison with FedAvg weight up- dates. For ResNet model trained on CIFAR-10, up- dates provided by the Top-n strategy significantly deviate from those of FedAvg method. This cor- relates with the worse performance in this experi- ment.
Schatten norm. Throughout this paper we used Frobenius norm to upper-bound the operator norm, see Eq. (2). While this choice allowed us to derive closed-fone forem-formmoo solution which enables fast caleulation on the server side, this upper bound is loose. Better approximation can achieved by
using Schatten p-norm with > 2. Schatten norm is defined as ||W (Xf ....... X²) ) where Ag....,A AN are singular values of the matrix W. Frobenius norm is an example of Schatten p-norm with p = 2.
14

---

Table Influence of the joint distribution. We have not found the evidence of the advantage of maximum entropy distribution (CPS) over other methods for our tasks. On CIFAR- 100, all three sampling methods perform similarly provided that marginal inclusion probabilities are preserved.
Table Unbiased estimator based on Schatten norm. While Schatten norm approximates the largest singular value better than the Frobenius norm, it does not allow the closed-form deriva- tion of sampling strategies. In our experiments with numerical constrained optimization, we have not seen any significant improvement over the methods that are based the Frobenius norm.
Mopsr
RATIO
STRATREY
CPS
RREWER
MISSUPRORT
RESNEY
0.1
COLASETIVE
UNRIASSI
35. 12
37. 53g
35.67a:
37.50g
35, 194 37.89 37.89
0.2
16 COLASETIVE
UNrIASSI
17 17.7gg
06.942
45.58g114
40.13
45.71g
STRATEGY
CIFAR-10 RESNET
RESNET
CIFAR-100 100
CCT
CCT
0.2
0.1
COLLECTIVE
COLAEETIVE
UNBEASSI
UNREASEI
53. 23.30gm
54.51
15.01g11m in
10.00 10.00g1
55 28
45 15
55.74gm
09.51
52 22.Mgam
45.3tgr vo
52 18gg
50.65g1
UNBIASED
SCHATTEN C+ = 4) 80.72 + 0.67
COLLECTIVE
82.59 0.27
80.57 + 1.39
37.83
34.54 + 0.62
35.12 1.51
1.24
45.01 + 1.07
49.53
45.75 4:76::0444 0:44
4.0...1.. 1.77
For the unbiased matrix estimators of the kind specified in Eq. (1), the average Schatten norm can be
upper-bounded with the help of Jensen's inequality, since (-) is a concave function for p
F
Pp(()) W W/1 11, p(s)
Σ
X
TT
1/2
111)
s
Ep((a)
Σ
X
T -
2j
1
(10)
1=1
n. + me (T 1)]] 21
Σ
X 1
Unfortunately, were not able to derive closed-form solution for large values of p in the same way as we proved Theorem 3.1. Therefore, we opted for numerical constrained optimization on the server side to find the optimal values of inclusion probabilities. As demonstrated in Tab. 5, for p 4 we did not find any significant improvement over our strategies presented in Sec. 3. For larger p. the optimization procedure turned out unstable, and the results were unsatisfactory.
Computational resources. All our experiments were conducted with a single GPU and required not more than 10 Gb VRAM. We used simulation of the federated learning environment with all participating clients being trained in sequence on the same GPU.
B Optimal Sampling Distribution
In the main text we discussed estimators obtained from singular vector decomposition. However, we found that the obtained results are more general and apply to any orthogonal decomposition. Therefore, in this section we provide proof's in a slightly generalized form.
More specifically, we consider a vector space equipped with an inner product and the corresponding induced norm. Assume that some element of that space Y may be represented as a sum of mutually orthogonal terms Y = > X,, and the terms are enumerated in descending order w.r.. their norms, i.c. A; |x,ll, and Å 2 ... 2 Ay > 0. In the particular case of SVD, Y is equal to the matrix W. X; = A and the Frobenius inner product is considered.
B.1 Unbiased Estimator
Let Y be an estimator of the element Y of the following type,
- Σ w (I)Xz, Σ 2, (Z),,
(II)
where I C (1,..., N) and the size of I equals n, 2; is an indicator variable 2; = 2I) = T(i E 2), and Wi E IR are scalar values, potentially depending the whole subset I. Assuming there exists
15

---

distribution p(Z) over all subsets of size n. the expected value of the estimator can be calculated
N
(2) - x Σ
p(I)
Σ p(T) (Z), (12)
I
1=1 I
B.1.1 Horvitz-Thompson Estimator
The bias of the estimator above equals ( - r - p(z) (2, 1] and
(i) all the {X,}; are mutually orthogonal and (ii) the estimator Y is unbiased, the following constraint obtained
BR( (Z) = for i = I...,,N. N.
(13)
Taking into account that 2; is binary random variable, we can use that 2 2, and derive the following
1 = E [2));( = EE [2:] = E 22 [w 2i]]
(14)
Pr(z: 1) [uy 1] T,E [uy 23 1],
(15)
leading to E [w,]2, = 1] =
Due to the mutual orthogonality of the {Xi}:
P Y 8
Σ X [=w,; 1]²
(16)
Now it is possible to state the following optimization problem to find the unbiased estimator with lowest mean squared error
min p(2)
X EptT [2:2; 112
s.t.
(Z) [aain]]
=, 1=,,..,, N,
(17)
m n.
0 S S 1.
The last constraint takes place since the size the subset is set equal to n which implics Σ, 24 = n, and due to the linear property of expectation E Σ, 2 Σ, Ez, Σ n. Furthermore, the quadratic term under the expectation can be simplified thanks to the unbiasedness and lower-bounded with the help of Jensen's 's inequality
E 2:00 1]2
E [
2E [2:
[z,
TE
[w]] Iz; = 1|
(18)
my (EE [wy]2; 2; 11)²
mi (m, 1)²
म. 1,
and this lower bound can be achieved by setting wi a non-random value w, π,
B.1.2 Inclusion Probabilities
With the optimal values of auxiliary multipliers known, the optimization problem can be rewritten as
min m.
X ( (m-
0 ग 51 I,..,,N, N,
(19)
R =1
1=1
16

---

This is typical contrained optimization problem which can be solved with the help of Lagrangian
function [(( , 8. (,)), (d:) = a Σ; A.m; + ( _े - + 7 (- (- >, 6 (T3 - 1),
leading to the following equations for i ,,..,,
0x. = at?
0L
+ + - 7/2 + 6: =0
i
Σ T: = n.
7: = 0,
(20)
6, (T, 1) = 0
a, 76 6 2 0,
a² ++ 8² +
0.
Case a = In this case 9-- 8, = 3 const.
Vi % < 6, But 7/ and 6; cannot be both greater than zero simultaneously since this would imply that π.
ak 6% > + Th = 1 % = 0 0 3 3 < ~0
and T; = 1 at the same time. Therefore, the only way for the strict incquality to hold true is to set Vi n 0, 8, > 0 + Vi ग. 1. The last equality leads to contradiction in case n N, which is our main case of interest.
Therefore, Vi 5, 0 0 7/4 = 3 const. If B 0, this breaks the constraint of existence of at = 0, and this leads to contradiction least one non zero coefficient. Therefore, Vi % = 8 > 0 + Vi n, given that n > 0.
Case = 1. W.l.o.g. we can assume than A W
M
Ax 0. From our equations,
X2
Vi 7 3 8; %
If 3E > 0 3 Ti m0= \ 0, and this contradicts the assumptions above. Therefore, vi 7 0.
x²
and Vi पे
B + 82
X
Now let us assume that 3k 0% > 0 + T 1 + 3 = X 6x + Vj
X 5x + 83 Vj + k 13 - A < 8, 5g. Therefore, j < k. then X3 \ > 0. and consequently 5, 5g 0 which leads to 5 Summarizing, if nT& = 1, then TΙ
Based on the previous observation, we define
t= min (1 [ISiSN; T. < <1} 1,
(21)
which means that कi -.. कe 1, and Te+1,..., n 1 - 6, for -t+1 1, ..... N, and
.ी 3. After that, the sum of all the inclusion probabilities equals 1- >i = t ++ Σ 2i=+t vB
2 Xi
IN EEun n-t
3
(n t)
Since all the probabilities हı should be less than
2j=t+1
ΣN EN Ay n-1
1 for 1 > f, there is a natural constraint de+1
Thus, the search of the optimal set of inclusion probabilities can be done in time not worse than O(N)). See Algorithm for details. In practice, significantly increased with vectorization. Note that the solution always exists in Algorithm because
found that the computation speed can be
fort = n 1 the condition A++1 < بي holds true as obviously An >inn Av.
Remark B.I. The optimal set of probabilities (Iı,..., TT) specified by Eq. (5) is balanced [39] w.r.t. the variables (AI...., Ax), i.e. for any subset of indices I sampled according to the optimal distribution p (Z) the following equality holds
X
T
17

---

Therefore, in the case of spectral model sharding, before the local training starts, the sum of singular values of the estimator W is equal to the sum of singular values of the full matrix W.
B.2 Collective Estimator
Now assume a slightly modified version of the setting described above: There exist C' > 1 clients each of which samples an i.i.d. estimator y(e), 5 S C, of the target value Y using the same distribution p (I). In this section consider simpler not depend on the sampled subset of indices anymore, i.e. ÿr(c) = > 2
of weighting where coefficients wy do
wx Imagine another participant, 'server', who tries to reconstruct the target value Y by averaging the clients' estimators, Y = - P(c) t Σ w;XK; Ee==" (()) We aim to minimize the average squared error between the server's 's estimator and ground-true true value Y. As before, due to the mutual orthogonality
of (Xi), this error can be expressed solely in terms of the magnitudes,
N
2
E
Y
212
E
A
X
W1
C
1=1
c=1
N
C
2
X
1 C
Wi
(c)
c=1
C
X7
1
2wh
Σ
E2
(c)
Ez (c)
+2 2
Σ
E2
2
C
c=1
d<c
(22)
2wi C
u
C(C
X?
1
CTe
+
CT, +
m
C
2
w²
w (C 1)
X
2w
+
C TT
+
m
C
1=1
WiTT (C 1)
Σ
x +
-
C
C
Now we can formulate a new constrained optimization problem,
min (1
N
X WWi
-22 +
Wi +
WiTTg (C 1)
C C
(23)
s.t.
mi
n1
i=1
0
< T
<
1,
N.
The corresponding Lagrangian function equals
W (C
[(fm), (w.). , a, a, (m). (1.))
Σ
X X2 -
C
C
13
Σ Mi
(24)
Σ
V/ (-m,) + Σ
5g (iT -
18

---

and leads to the following conditions for all i == ......NNN
w;
-2+
Wi
2w (C (C-1) 1)
3
7: + 6, = 0,
0T.
C
C
OL
2w0
2 (C - 1)
= ad n.
-2 +
+
= 0.
Dw;
C
C
Σ n n
(25)
7 0, 8 (T. - 1) = O,
a. 76. 8 2 0,
a² + 8 + Σ
7 + Σ
80 > 0.
Case a 0. This case does not differ from the one of the setting of the unbiased estimator.
Case a = 1. Since Σ, गं = n, there exists such & that the corresponding inclusion probability
Th is greater than zero, 3k T& > + -2 + 2u + (C
+ Wkk (C-J- As we look for 0 < m 2 1, this implies S Wk C. Also, since Th > 0, then 7t =
C
Xe -22
2w, mu (C-1) + 6x
Rus (- wx + 0 B 8x 2
If Tk 1, then 8% = 0 Wk
VCP d
= Tx
C-1
1
c 9
1]. Furthermore, in this case
1 < wg C, then 2 < B < C
max
x
C
<
min X²C.
k0
k:0cm m <1
X Otherwise, if T. = 1, then wi = L, and B = -6 + C
and since oz 2 0, the following incquality
holds true, 3 < min 12
< min X This inequality k: m1=1
Combining these observations together, we obtain
max k:0-m. <1
can hold true only if, similarly to the case of unbiased estimator, several largest magnitudes correspond to inclusion probabilities equal to 1, and all the rest inclusion probabilities strictly less than 1.
Revisiting the optimization criterion ΣN F; = Σn XURT (-2+ +
wy + W.r.(C-1)). we can
C
now consider each separate term of the sum. Starting with the product of the inclusion probability and the weighting coefficient
if T = 1, if < प. < if =T 0,
N
3
1
VCB C A
C-1
VCB CC
(26)
0
we can write down the equation of each term
if = 1,
F.
X2
C-T
JC (-(C-] ) (-2+
X.v X.CC = -Tdo d; \i V@
+1
if < क. < 1,
if T. = 0.
(27)
As derived above, in case 0 n; we have 3 < XCC A; monotonically decreasing function of X; with negative values if n. <1. Since the goal is to minimize Σ. F,, this implies that it is beneficial to assign non-zero probabilities to greater values of \ Therefore, the optimal solution (Ti,..., nT) has the following form for some f 2 0, u 2 0,
V8- Therefore, F, is
Te 1.
19

---

1> T(11 2 ... 2..2 4tta > 0,
Tt4u+1 TM= = 0.
Since the sum of all inclusion probabilities should be equal to n,
VC Ex= VC i=++1
Σ
H+
u
A; (n - (C-
V ++u -i++++
C- = = -
The search for the best solution takes no longer then (N) time, see Algorithm 2 for details. As before, the calculations may be done in the vectorized form.
Remark B.2. Interestingly, for large values of C, the optimal values given by Eq. (8) are approaching some unbiased Soovitz--Thopsonn estimator of W, in accordance to the probability theory's law of large numbers. However, this estimator is not guaranteed to have the properties that are provided by Theorem 3.1. Therefore, for large values of C' it is more beneficial to use the Unbiased estimator instead of the Collective one.
Remark B.3. It is straightforward to show that if the constraint of unbiasedness is applied to the collective estimator, i.e. EW = W, then it necessarily implies that w. ., just as in the case of the unbiased estimator. Therefore, the unbiased collective estimator is the same as the unbiased estimator. Since Theorem 3.2 provides the estimator with the least Frobenius discrepancy among all the collective estimators, the achieved least discrepancy cannot be greater than the least possible discrepancy among unbiased collective estimators. This concludes the proof that the discrepancy of the collective estimator is not greater than the discrepancy of the unbiased estimator. However, as Remark B.2 implies, in the limit case -> oo these discrepancies become equal.
C Wallenius' Distribution and numpy random. choice
The official documentation" of NumPy [12] currently does not contain description of how exactly sampling without replacement with unequal probabilities is done if the function numpy random. choice is employed. Nevertheless, the source code suggests that sampling is conducted in the following way:
1. The remaining number of items is sampled with replacement from the multinomial distribu- tion with probabilities proportional to the current chances (in the beginning, those chances are equal to the probabilities provided by the user).
2. The unique items from the sampled subset are added the output of the function, and the chances of those items are set equal to zero
3. The steps above are repeated with the updated chances until the requested number of items is selected.
The analysis of the statistical properties of the output sample seems infeasible if the probabilities given by the user are far from being uniform. However, the case of non-uniform chances is exactly what we are interested in when working with singular values of weight matrices.
Interestingly, this algorithm from NumPy resembles the description of multivariate Wallenius' non- central hypergeometric distribution [42, 5, 9, 10] which is often defined through the um problem. In detail, if one draws items one by one, instead trying to get the maximum possible number each step, this would provide a sample from a special case of Wallenius' law.
Clearly, the Despite this, we found that in our simulations the mean vector of the Wallenius' distribution was quite good approximation for the empirical inclusion probabilities of the output of NumPy function.
procedures
different significantly enough produce different distributions.
'ntpp:://gitabb.cccommm/// http://www...i.ia/////..//////aa
20

---

D Licenses for Assets
1. CIFAR-IO [23]: unknown
2. CIFAR-100 [23]: unknown
3. Tinylmagenet [31]: unknown
4. Shakespeare [29]: Apache License, Version 2.0
Impact Statement
This paper presents work whose goal is advance the field of federated learning. The potential societal consequences include the increased abilities of training more powerful models on a federation of user devices in privacy-friendly manner. At the moment, we do not see any direct negative consequences of the conducted research.
21

---

Algorithm 1 Inclusion Probabilities for the Unbi- Algorithm 2 Inclusion Probabilities and Auxil- ased Strategy
iary Multipliers for the Collective Strategy
Require: Magnitudes A 2 A2
AN
Require: Magnitudes 2 X 2 AN 0, number of samples 0 < n N, number
0, number of samples () < N.
Ensure: Optimal
inclusion
probabilities
of clients C.
( [.1....,"N mT)
Ensure: Optimal inclusion probabilities
1: E* + 00
Best criterion value
T .y), optimal auxiliary multipliers
2: for = t
do
(w 1..... WN)
3: A, < Σ 21= A
# Consider top-n case first, i.e. n, u
4: if 1-11 A. then
for = 1 x-1
1: for 1 ..... n do
5:
to t do
6:
m f-1 end for
TT f Wi f
7:
end for
8:
for + N do (n-
5: for n 1,...,N. N do
9:
Tl f- end for
10:
11:
= Σ - X
7: Wi f
0
end for
12:
if E < E"
then
9: E
Σ i= X
E + E (Tf.....
10: (Ti..... .i..............................
N
.1..... TN)
14:
T
(T1 .......:.. TT))
11: (wi
WN) t [1...... WN
15:
end if
# Other cases
16: end if
12: for t = 0,
1 do
17:
end for
13: for
do
14:
W;
15:
Wi
16:
end for
17:
E.
Ei=1 X
18:
for 1,.... N do
VO
19:
VB t (n-t)(- 1)tu
20:
if Aet1
vB
& VB
VC
vC
A1tu V then
21:
for i = t l...., u do
AvC.
22:
T; T-
vB
CCVB
23:
w; A
24:
end for
25:
for i = 1 + + 1,....N
N do
26:
n;
27:
W/ + 0
28:
end for E 29:
E
t+u d; (>,
C-1
i=tt1
VC
30:
if E E then
31:
F <
32:
(T²
1. f- (m) ......... TN)
33:
(w] W.....
wx) t (Wi,... WN
34:
end if
35:
end if
36:
end for
37:
end for
22

---

NeurIPS Paper Checklist
1. Claims
Question:
Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification:
The strategy using unbiased estimators is introduced in Sec. 3.2.1. The
strategy minimizing the squared approximation error is derived in Sec. 3.2.2. The practical considerations are discussed in Sec. 3.2.3. The quantitative evaluation is provided in Sec. 4.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results be expected to generalize other settings.
is fine include aspirational goals as motivation long is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: The used assumptions and their limitations are explicitly discussed in Secs. 3.2.1, 3.2.2 and and Appendix A.
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect the scope of the claims made, e.g., if the approach was only tested on a few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
23

---

Answer: [Yes]
Justification: Due to the lack of space, the proof of Theorem 3.1 was put in Appendix B. The proof of Theorem 3.2 for the case of C = given in the main text, and the proof for C >> is provided in Appendix B.2.
Guidelines:
The answer NA means that the paper does not include theoretical results,
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided not)?
Question:
Answer: [Yes]
Justification: The step-by-step algorithms for the proposed strategies are presented in Algo- rithms 1 and 2. All the training details are provided in Sec.
Guidelines:
The answer NA means that the paper does not include experiments. If the paper includes experiments, No answer this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how replicate the results, access a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how reproduce that algorithm. (b) If the contribution is primarily
model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should possible for other researchers to have some path to reproducing verifying the results.

---

5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes] for data and [No] for code
Justification: All the datasets used are public, see the provided references in Sec. 4. The code cannot be released at the moment due to copyright procedures.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips cc/
for more details.
While we encourage the release of code and data, we understand that this might not be possible, "No" an acceptable answer. Papers cannot be rejected simply for including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https: / //niss. wwwpcccpppcc.aeeecc...ccc.ccetle httpiiigggi/////eeeeeeeeeesseeee httgg//dde/ooeeuubiiieeeeeeeeeee
for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental Setting/Details
Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Question:
Answer: [Yes]
Justification: All these details are provided in Sec.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment Statistical Significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [Yes]
Justification: All the reported results in both tables and plots include error bars obtained from multiple runs with different random seeds, as explained See. 4.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by emoo bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
25

---

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
It should be clear whether the error bar is the standard deviation the standard error of the mean.
It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not show in tables figures symmetric error bars that would yield results that are out range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Question:
Answer: [Yes]
Justification:
Due to the lack of space, the details are provided in Appendix A.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount compute required for each of the individual experimental runs as well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that NeurlPS Code of Ethics https: / hww.rpsu....p//. ctt.bliieeehiisccieeeeeeee.ee.e
didn't make into the paper).
9. Code Of Ethics
Question: Does the research conducted the paper conform, in every respect, with the
Answer: [Yes]
Justification:
Our research conforms with the NeurlPS Code of Ethics.
Guidelines:
The answer NA means that the authors have not reviewed the NeurlPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., there a special consid- cration due to laws or regulations their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification: Due to the lack of space, the impact statement was put to the supplementary material.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
26

---

If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), faimess considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g.. gated release of models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Question:
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Answer: [NA]
Justification: Since our work is mostly theoretical, we believe that our paper poses no such risk.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification: All the models, algorithms and datasets are properly cited in the text. We re- re-implemented all the models and algorithms ourselves. The linceses for datasets are listed in the supplementary material.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset used and, possible, include URL.
The name of the license (c.g., CC-BY 4.0) should be included for each asset.
27

---

For scraped data from a particular source (e.g., website). the copyright and terms of service of that souree should be provided.
assets are released, the license, copyright information, and terms use in the package should be provided. For popular datasets, paperswi thcode. com comddataees has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information is not available online, the authors are encouraged to reach out to the asset's s creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification:
No new assets released.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/mmodel as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification: Our paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According to the NeurlPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification:
Our paper does not involve crowdsoureing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with human subjects.
28

---

Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
29


Reviews:
[{"review_id": "FXSNV9sjQJ", "reviewer": "Reviewer_fYuS", "summary": "he study addresses the challenge of heterogeneous clients in federated learning by introducing two sampling strategies for spectral model sharding. These strategies, derived from specific optimization problems, either produce unbiased estimators of the original weights or minimize the squared approximation error. Empirical results demonstrate that these methods enhance performance on various datasets, despite sometimes slower convergence rates.", "questions": "$\\bf Q1.$ Considering the experiments, could one suggest reducing the model size based on the fact that trained layer using your algorithm is low rank? How do you justify using your method instead of considering a model with fewer parameters?\n\n$\\bf Q2.$ It would be beneficial to add the performance results of FedAvg, specifically examining the effect of local epochs, and local training (training on local data only). Additionally, comparing the communication costs of FedAvg and your algorithm can provide a fair measure of the reduced communication cost.\n\n$\\bf Q3.$ Have you tried factorizing different layers? Which layer's factorization can reduce the communication cost the most, and which one affects the performance the most?", "limitations": "Adding a convergence analysis of the algorithm will make the paper stronger.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The paper clearly introduces the problem and identifies a significant gap in the current research. The proposed method is well-presented, with experiments that support the main claims. Additionally, it includes a detailed discussion on the performance of the algorithms, impact of data heterogeneity, communication complexity, and the influence of hyperparameters.", "weaknesses": "Practitioners can benefit greatly from reading the paper. Adding a convergence analysis of the proposed algorithm would provide valuable insights for technical readers in the field of federated learning.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their time and effort in evaluating our work. We are happy to take the feedback into account and improve our presentation accordingly. Below, we address the questions raised in the review.\n\n**Q1 and Q2.** While using a model with fewer parameters is a straightforward solution for low-capacity clients, we believe that the capabilities of contemporary large vision and language models, as demonstrated in recent years, indicate that larger model size often significantly improves the performance. As a consequence, the challenge of training a large model on a set of low-end devices naturally emerges. One advantage of spectrally sharded models is the ability to train a server model that is significantly larger than the capacity of any individual participating client. \n\nIn the context of the aforementioned challenge, we compare the communication costs of the full-size ResNet model trained on CIFAR-100 against different spectral sharding strategies with a keep ratio of 0.1, see Fig. R1 in the attached pdf. These results correspond to those in Tab. 1 of the submitted paper (where FedAvg was referred to as “No Sharding”). In addition, we trained slim versions of various models each containing approximately the same number of parameters as the subsampled per-client models in the considered strategies.  \n\nAs shown, the full model trained with FedAvg is the least communication-efficient. The slim version of the ResNet model demonstrates the best efficiency in this experiment. However, as mentioned above, the ultimate goal of sharded training is to enable the user to train a model of size which surpasses the individual capabilities of clients. And slim models do not provide this ability.\n\n\n**Q3.** The low capacity of a participating client imposes a global constraint on the total computational budget. Distributing this budget among different layers of the network is a complicated and valuable problem, as demonstrated in recent works e.g. [1, 2]. Our study focuses on sampling strategies, which represent a subsequent step after the budget has been distributed. For simplicity, in our work, we spend our budget in proportion to the original layer widths, using the same keep ratio for each subsampled layer. \n\nThe communication cost and influence on the performance of different layers are intrinsic properties of each specific architecture. E.g., as mentioned in the submitted manuscript, various methods (including ours) do not decompose the very first and the very last layers which are crucial for model performance, and we confirmed this finding in our early experiments.  Speaking of the communication cost, for example, in ResNet architecture the number of parameters in a residual block grows with depth and, as noted in [3], “last 3 bottleneck residual blocks of ResNet-50 contain ∼60% of the entire model parameters”. At the same time, modern transformers typically consist of blocks of the same size.\n\n\n**Convergence analysis.** We agree that a formal proof of the convergence would strongly support not only our approach but the entire group of spectral sharding-based methods. However, due to the complexity of such an analysis,  in our work, we opted for an empirical check of convergence which is provided in Appendix A of the submitted paper. We aim to provide rigorous proof in future work.\n\n\nReferences.\n\n[1] Wang et al. Pufferfish: Communication-efficient Models At No Extra Cost. In MLSys, 2021.\n\n[2] Wang et al. Cuttlefish: Low-Rank Model Training without All the Tuning. In MLSys, 2023.\n\n[3] Yao et al. FedHM: Efficient Federated Learning for Heterogeneous Models via Low-rank Factorization. 2021.", "claims_and_evidence": "", "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": "", "extra_scores": {}}, {"review_id": "2KVCnRRMzc", "reviewer": "Reviewer_7vsp", "summary": "To handle system heterogeneity in federated learning, spectral model sharding partitions model parameters into low-rank matrices using singular value decomposition (SVD). As done in prior work, e.g., PRiSM, SVD of the weight matrices of affine layers are done and only randomly sampled terms from this decomposition are locally updated at each client (different for each client based on their system capabilities), boosting on-device training efficiency. This study introduces two sampling strategies for sharding: one produces unbiased estimators of the original weights, and the other minimizes squared approximation error. Both strategies can be incorporated into the federated learning loop and authors show improvements on various FL datasets. These solutions have closed-form optimals that can be computed on the server side, avoiding additional computations on client devices.", "questions": "- For the unbiased estimator, the sum of inclusion RVs is a Poisson Binomial distribution with variance that scales with $n$ for a matrix of size $n \\times n $. Does this not violate the computational constraints on each user, what if the user can only handle a fixed rank of $3$? Or, is the approach limited by the fact that the computational constraints for each user need to scale with model size?", "limitations": "The authors mention some limitations in Appendix A, though the limitations section can use some better sign posting.", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- Sampling methods proposed (e.g., sampling based on eigen value), make it hard to analyze the marginal distribution over the inclusion probabilities, and thus the approximation error of the resulting matrix, whereas the strategies explored in this work are more amenable to analysis.\n- The preliminaries section presents a crisp and clear overview of spectral model sharding, though maybe some notational overload (e.g., definition of $V^{(c)}$) could have been avoided.\n- The discussion on local training is particularly informative, and is useful for practitioners to understand how different choices of the sampling strategy and corresponding auxiliary weights can affect local training dynamics. The proposed solutions to control the effective learning rate also seem reasonable. This is also empirically verified in Table 3.\n- The framework presents a way to propose and analyze new variants of the marginal inclusion probabilities, which can easily be plugged into Conditional Poisson schemes that allow for more interactions between different singular vectors. \n- The empirical results in Table 1 suggest that when 60% clients have keep ratio $r = 0.2$ and 40% clients have $r = 0.4$, the unbiased and collective strategies significantly outperform PRiSM, top-n (and variants). Though, the gains are marginal for a fixed keep ratio. Nevertheless, its nice that the authors evaluate confidence intervals, which suggests that the gains are at least statistically significant. \n- The proposed modification of PriSM: PriSM + Wallenius + ClipLR, reduces exploration during training, improves exploitation, and thus empirically improves convergence of FedAvg training.", "weaknesses": "- The paper only analyzes the pre-client update approximation errors, i.e., the error induced by spectral sharding locally at each client. It would be nice if the authors can comment on how the post-client update average deviates from the weight matrix obtained by running one round of FedAvg without spectral sharding (for the proposed sampling strategies).\n- While the authors analyze the Frobenious norm of the discrepancy, a tighter norm to analyze would have been the spectral norm. It is unclear if and how the optimal solutions (and findings) in Frobenious norm transfer.\n- For the collective estimator, the paper only provides a closed form solution when all clients in the per-round average have the same constraint $n^{(c)} = n$. Can the authors elaborate on the case where there is heterogeneity within each round of the update? This seems limiting in practical settings. \n- It is unclear why the collective estimator is a natural class to analyze in terms of estimators that optimally tradeoff bias and variance?\n\nNote: On the outset, this work seemed like a small modification of PRiSM, but reading it more carefully, I feel the paper clearly presents that the choice of sampling in spectral model sharding can significantly affect the convergence and performance in FL. While the theoretical results on the approximations are nice, there is definitely more room for analyzing convergence, post-gradient update performance etc., for different types of sampling strategies. At the same time, the empirical results would also be more impactful if extended to larger models (> 1B parameters, e.g., in LLMs) and settings with both system and statistical heterogeneity. Thus, I am giving a score of 5 for now, but will definitely increase, if the authors can respond on some of my concerns/clarifications in the rebuttal period.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the time and effort invested in our work. We are happy to take the feedback into account and improve our presentation accordingly. Below we address the questions raised in the review.\n\n**Post-client update.** We appreciate this valuable suggestion! To analyze the post-gradient updates, we followed the optimization path $ \\\\{ \\\\theta_t \\\\}\\_{t=1}^{T} $ of the server model under each strategy with $T$ communication rounds and calculated the cosine similarity between the server model update $\\\\Delta\\_{sharding\\\\;strategy} \\\\theta_t$ under that strategy and $\\\\Delta_{FedAvg} \\\\theta_t$ under FedAvg method with full-size model on each client. Subsequently, the model weights were updated according to the selected strategy: $\\\\theta_{t+1} = \\\\theta_t + \\\\Delta_{sharding\\\\;strategy} \\\\theta_t. $ The results of these experiments are provided in Fig. R2 in the attached pdf. As shown, there is a notable connection between the relative performance of each strategy and its “alignment” with FedAvg updates. The deterministic Top-n strategy lags behind its randomized counterparts in terms of accuracy in the considered setting (see Tab. 1 of the paper), and a similar observation is made for its update directions.\n\n\n**Tighter error bound.** While the Frobenius norm is convenient for analysis and allows for closed-form solutions, it indeed provides a relatively loose bound for the spectral norm. For this reason, in Appendix A of the submitted paper, we present the results from experiments using Schatten norms which offer a tighter bound. However, as we demonstrated, the obtained results did not significantly differ from those described in the main text.\n\n**System heterogeneity.** When heterogeneous clients participate in the same round, we follow prior approaches [1] and cluster the clients based on their capacity. After that, the sampling distribution of singular values is computed within each cluster independently. We used this methodology to conduct experiments for the third, “heterogeneous” block of Tab. 1 in the submitted paper. We agree that this approach is ad-hoc, and aim to develop a more grounded solution in future work.\n\n\n**Motivation of the collective estimator.** We agree with the reviewer that the collective estimator may not be the most natural choice from the perspective of the bias-variance trade-off. In the paper, we mentioned this trade-off  (i) to motivate the unbiased strategy and (ii) to explain why the unbiased strategy might have limitations such as large variance. Our motivation for the collective estimator was different. We aimed to generalize the well-known property of truncated SVD (Eckart–Young–Mirsky theorem) by incorporating the knowledge about the number of participating clients. Therefore, we formulated the optimization problem to find the estimator with the least expected Frobenius discrepancy.\n\n**Larger models.** We admit that demonstrating scalability is highly desirable for any federated learning algorithm. However, due to our limited computational resources, we focused on experiments using datasets and models of the same scale as our fairly recent baselines. We hope to apply our approach to large-scale tasks in the future.\n\n**Sum of inclusion RV.** For both the unbiased and collective estimators, the sum of inclusion random variables $\\sum_{i=1}^N z_i$ is not random but deterministic and always equal to $n$. Note that in both  Theorem 1 and 2, we derive the probabilities for sampling without replacement with fixed sample size, see Eq. (19) and Eq. (25) respectively. Additionally, all the sampling methods we use guarantee that exactly $n = \\lceil N r \\rceil$ singular values are sampled for the client with keep ratio $r$. For that reason,  the variables $\\{z_i\\}$ are not mutually independent and do not lead to a Poisson distribution.\n\n\n\nReferences.\n\n[1] Mei et al. Resource-adaptive federated learning with all-in-one neural composition. In NeurIPS, 2022.", "claims_and_evidence": "", "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": "", "extra_scores": {}}, {"review_id": "yJqH2r41Cp", "reviewer": "Reviewer_qPbH", "summary": "This paper addresses the problem of heterogeneous clients in federated learning, where each client may have different constraints and capabilities. The authors propose two sampling strategies for spectral model sharding, which involves partitioning the model parameters into low-rank matrices based on singular value decomposition. The first strategy provides unbiased estimators of the original weights, while the second strategy aims to minimize the squared approximation error. The authors discuss how these estimators can be incorporated into the federated learning loop and provide practical considerations for local training. Empirical results show that both methods improve performance on various datasets. Overall, the paper presents novel sampling strategies for heterogeneous clients in federated learning and demonstrates their effectiveness in improving performance.", "questions": "1.\tConsidering that existing methods such as PriSM and FedHM have only been tested on models with parameters numbering up to tens of millions, there is a concern regarding the computational cost of singular value decomposition (SVD) and its applicability to larger-scale models. Therefore, it would be valuable to explore whether the proposed strategies in this paper can be extended to larger models, potentially utilizing approximations or other techniques to address the computational challenges associated with SVD. \n\n2.\tHow does the accuracy of the proposed methods compare to FedHM? Considering the absence of a direct comparison with FedHM in the baseline methods presented in the manuscript, evaluating the performance comparison would provide valuable insights into the effectiveness of the proposed strategies relative to an existing method.", "limitations": "N/A", "rating": 5, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "Strengths of this paper include the innovative approach to develop sampling strategies for spectral model sharding in federated learning. The authors provide two novel strategies derived from specific optimization problems, which aim to minimize the expected squared approximation error. These strategies are presented in closed form, making them easily implementable in practical settings. Additionally, the proposed methods do not require additional calculations on potentially weak client devices, enhancing their efficiency for on-device training. The incorporation of unbiased estimators and consideration of all client-specific sub-models distinguish this work from previous heuristic sampling designs, offering potential improvements in performance on various datasets.", "weaknesses": "The manuscript presents a natural and well-articulated approach, however, I have some concerns:\n\n1.\tThe reliance on the singular value decomposition (SVD) as a core component of the proposed strategies could raise apprehensions, especially considering the substantial computational cost associated with SVD. Furthermore, the need for repeated SVD decompositions after each communication round could limit the practical application of the method, particularly on large-scale models. It would be valuable for the authors to explore methods to mitigate computational costs, such as utilizing approximations or reducing the frequency of SVD decompositions, to enhance the feasibility of the approach, especially for larger and more complex models.\n\n2.\tThe absence of a comparison with the FedHM method based on random sampling in the experimental section is noteworthy. While the authors state their focus on analyzing the relative performance of the proposed spectral sharding strategies, the lack of comparison with the FedHM method could limit the comprehensive understanding of the strengths and weaknesses of the proposed approaches. Including such a comparison would provide deeper insights into the performance of the proposed strategies and their positioning in relation to existing methods.\n\n3.\tTypographical Errors: The manuscript contains minor typographical errors, such line 304 \"We leave this fo future work\" -> \"We leave this for future work\"", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the time and effort invested in our work. We are happy to take the feedback into account and improve our presentation accordingly. Below we address the questions raised in the review.\n\n**Comparison with FedHM.** According to the paper [1], FedHM is a method also based on spectral sharding, as it uses top-$n$ singular values and singular vectors of weight matrices of factorized layers. This corresponds to the Top-$n$ strategy in our paper. However, the FedHM method includes two additional details. First, the initial layers of the network are never decomposed, and the number of such layers is treated as a hyperparameter. Second, during the aggregation step on the server side, instead of averaging the updated singular vectors, the full weight matrices are re-materialized and then averaged.\n\nIn our work, we focused on exploring the sampling strategies and, for simplicity,  followed other works [2, 3] by not factorizing only the very first and the very last (i.e., classification head) layer. To test the strategies with this aggregation method, we conducted several experiments on CIFAR-100 with keep ratio of $0.1$. Their results are reported in Tab. R1 in the attached pdf. Interestingly, the strategies benefit from this kind of averaging, while Collective strategy still outperforms Top-$n$. Arguably, the performance of the Unbiased strategy is explained by the excessively large auxiliary multipliers used during the weight matrix re-materialization.\n\n**Larger models.** We admit that demonstrating scalability is highly desirable for any federated learning algorithm. However, due to our limited computational resources, we focused on experiments using datasets and models of the same scale as our fairly recent baselines. We hope to apply our approach to large-scale tasks in the future.\n\n**Computational cost of SVD.** Indeed, singular decomposition may appear to be a limiting factor for methods based on spectral sharding. However, it is important to note that this decomposition takes place solely on the server side and therefore does not impose any computational burden on low-capacity clients. In addition, in our method, the SVD decomposition for different layers can be performed in parallel, as they are independent of each other. For example, in case of transformer models which typically consist of homogeneous blocks, batched SVD implementation can already provide the speed up. Also, recent solvers for the SVD algorithm for GPU such as QDWH-SVD [4, 5] potentially allow for efficient implementation of our algorithm even for large models.\n\nAs a proof of concept, we simulated the training of a larger model with the same architectures on CIFAR-100 by decreasing the keep ratio to $0.05$ and applying fast randomized SVD algorithm [6] which estimates only half of the singular values and vectors. Also, in this experiment we used FedHM-like weight aggregation scheme on the server side, because otherwise PriSM-like scheme (used in the main text) would lead to the low-rank server model. As reported in Tab. R2, in this setting our strategies are still able to learn. In addition, we experimented with ODE-based SVD approximation [7] for the updated weight matrices but the results for large weight matrices were unsatisfactory.\n\nReferences.\n\n[1] Yao et al. FedHM: Efficient Federated Learning for Heterogeneous Models via Low-rank Factorization. 2021.\n\n[2] Khodak et al. Initialization and regularization of factorized neural layers. In ICLR, 2021.\n\n[3] Niu et al. Overcoming Resource Constraints in Federated Learning: Large Models Can Be Trained with only Weak Clients. TMLR, 2023\n\n[4] Nakatsukasa and Higham. Stable and Efficient Spectral Divide and Conquer Algorithms for the Symmetric Eigenvalue Decomposition and the SVD. SIAM Journal on Scientific Computing, 2013.\n\n[5] Sukkari et al. A High Performance QDWH-SVD Solver Using Hardware Accelerators. TOMS, 2016.\n\n[6] Haiko et al. Finding structure with randomness: Probabilistic algorithms for constructing approximate matrix decompositions. 2009.\n\n[7] Dieci and Eirola. On Smooth Decompositions of Matrices. 1999.", "claims_and_evidence": "", "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": "", "extra_scores": {}}]

Author rebuttal:
Dear reviewers, please find the results of requested experiments attached.
The comments on these results are provided in the individual replies.

Meta-review: Spectral model sharding is a technique used in federated learning to address system heterogeneity by breaking model parameters into low-rank matrices. This allows clients with varying computational capabilities to work with smaller, more manageable portions of the model. The authors build on this technique by introducing two new sampling strategies for spectral model sharding. The first strategy creates unbiased estimators of the original model weights. The second strategy focuses on minimizing the squared approximation error. 

The paper is well-articulated and provides meaningful insights, with the proposed strategies being more amenable to analysis than previous approaches. Empirical results demonstrate improvements over existing methods. Although there are some concerns about the focus on pre-client update errors and the computational cost of SVD on the server side, the reviewers are unanimously positive, leaning toward acceptance. I agree with this decision, as the paper makes a valuable contribution to addressing system heterogeneity in federated learning. As recommended by the reviewers, I encourage the authors to include post-client update experiments and a more detailed discussion on the collective estimator in the final version.
Decision: Accept (poster)