On Sampling Strategies for Spectral Model Sharding

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.

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.

In 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.

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.

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.

As 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.

References.

[1] Yao et al. FedHM: Efficient Federated Learning for Heterogeneous Models via Low-rank Factorization. 2021.

[2] Khodak et al. Initialization and regularization of factorized neural layers. In ICLR, 2021.

[3] Niu et al. Overcoming Resource Constraints in Federated Learning: Large Models Can Be Trained with only Weak Clients. TMLR, 2023

[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.

[5] Sukkari et al. A High Performance QDWH-SVD Solver Using Hardware Accelerators. TOMS, 2016.

[6] Haiko et al. Finding structure with randomness: Probabilistic algorithms for constructing approximate matrix decompositions. 2009.

[7] Dieci and Eirola. On Smooth Decompositions of Matrices. 1999.

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.

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.

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.

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.

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.

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.

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.

References.

[1] Mei et al. Resource-adaptive federated learning with all-in-one neural composition. In NeurIPS, 2022.

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.

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.

In 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.

As 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.

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.

The 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.

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.

References.

[1] Wang et al. Pufferfish: Communication-efficient Models At No Extra Cost. In MLSys, 2021.

[2] Wang et al. Cuttlefish: Low-Rank Model Training without All the Tuning. In MLSys, 2023.

[3] Yao et al. FedHM: Efficient Federated Learning for Heterogeneous Models via Low-rank Factorization. 2021.