Cohort Squeeze: Beyond a Single Communication Round per Cohort in Cross-Device Federated Learning

The core novel contributions of this paper are still unclear to me. I appreciate the detailed explanations of the theoretical results and the examples with various concrete sampling strategies. However, the novel algorithm SPPM-AS seems to heavily rely on SPPM, which can already be applied directly to the federated learning setting (Equation 1). I want to understand the technical differences and contributions of SPPM-AS compared to the SPPM algorithm. Please provide a more explicit comparison between SPPM-AS and SPPM, highlighting the key technical differences and innovations.

Moreover, It appears that this paper eliminates the need for the second-order similarity condition in SPPM. How eliminating the second-order similarity condition can be achieved in your proof is of great interest to me.

Last but not least, explain in more detail how the multiple communication rounds within cohorts contribute to the novelty of the approach.

• In section 2.2, the author(s) discussed some properties of the SPPM-AS, but I cannot find the communication cost analysis of SSPM-AS, which is the most important factor of FL algorithms. Theoretically, how does the total communication cost of SPPM-AS compared with existing FL algorithms such as FedAvg, FedProx, SCAFFOLD, etc.
• Similar to the question above, how is $prox_{\gamma f_{S_t}} ( x_t )$ being solved? There must be some communication between $S_t$ during the optimization, how expensive is the communication?
• Table 1 is not very easy to read. I did not fully get the meaning between 313-323 when I read it for the first time.
• In line 340, how is $\tilde{f}_i$ defined?
• In experiments, how to solve the proximal point problem is kind of vague, what is the local communication cost and how is the local communication cost being controlled in each experiment?
• Federated leaning has been studied many years. The baseline methods in the experiments is limited (FedAvg), the author(s) should include some more recent FL algorithms.

1. The presentation of the paper needs to be improved. For example, it is not easy to follow the paper since a lot of important discussions and results are presented in the appendix.
2. A detailed explanation of Algorithm 1 or SPPM should be provided to improve better reader understanding.
3. The theoretical analysis is based on the strongly convex assumption. Extending the analysis to a more general non-convex setting would strengthen the paper.
4. The comparisons between different sampling methods are based on simplified settings, e.g., $b$ clusters of uniform size $b$, with blocking size and the number of blocks set as 2.
5. The authors only provide experiments on logistic regression using datasets from the LibSVM repository and on CNN with FEMNIST dataset, which are relatively simple. To better demonstrate the performance of SPPM-AS, experiments on more complex datasets (e.g., CIFAR-100, Shakespeare) and tasks (e.g., NLP) are recommended.

Minor:

1. Notations should be explained when they first appear in the paper, e.g., $n$.
2. In line 93, "dashed line" should be corrected to "dashed red line".
3. Abbreviations should be defined upon their first appearance in the paper, e.g., $HP$.

1. I don't see any significant theoretical improvement of the proposed algorithms.

   1a. The iteration complexity is $1/\epsilon$ (Line 212), which does not seem to improve upon FedAvg expect possibly in terms of constant factors. This is not too surprising considering that FedProx does not improve upon FedAvg, but there are already works in FL showing that the order of client sampling can improve convergence rate in terms of the $\epsilon$ dependence [1]. Therefore the iteration complexity shown in this paper is not a significant improvement.

   1b. The algorithm which enjoys theoretical guarantees cannot actually be implemented, because the choice of hyperparameters requires knowledge of the global minimum. The iteration complexity (Line 212) requires to choose the stepsize based on $\sigma_{\*,\text{AS}}^2$, which can only be computed with knowledge of $x_*$.

   1c. The only possibility for theoretical improvement is an improvement of constant factors from optimal stratified sampling (Lemma 2), but optimal stratified sampling cannot be computed without knowledge of the global minimum.

2. I don't see any significant experimental improvement due to issues with the experimental methodology.

   2a. The experimental evaluation only compares against naive baselines of Local GD and Minibatch GD. There are a huge number of works in FL that try to improve optimization with different client selection strategies, and these works are essentially ignored by this experimental evaluation (see [2] and its references). The authors' claim of $74$% improvement compares against the naive baseline, not against state-of-the-art (or any of the relevant existing work).

   2b. SPPM-SS cannot be implemented with real data. As I pointed out in Weakness #1c, optimal stratified sampling cannot be computed without knowledge of the global minimum. To run experiments, the authors instead use a clustering heuristic that stratifies clients according to features, clustered using K-means. However, it is unclear whether such a clustering procedure can be executed in a real federated learning scenario when client data must remain on-device. Without this, a significant portion of the experimental results (Figures 1, 2, part of 3) only describe an algorithm which cannot be implemented in practice.

   2c. The neural network experiments (Figure 4) may not be a fair comparison between LocalGD and SPPM-NICE. SPPM-NICE uses Adam as a local prox solver, which may not be a reasonable comparison against LocalGD, since LocalGD does not include any update preconditioning (known to be important for neural network training). It would be more appropriate to compare SPPM-NICE against LocalGD when SPPM-NICE uses GD as a local prox solver. An alternative is to compare SPPM-NICE w/ Adam against a local version of Adam, for example FedAdam. Appendix E.4 contains NN experiments with different local solvers (Figure 16), but I don't see exactly how these results related to those in Figure 4. It looks like the choice of local solver can create a gap of about 6% in train accuracy, and this is described as "all methods perform similarly" (Line 1526), whereas a similar gap between LocalGD vs. SPPM-NICE in Figure 4 is described as "enhanced performance" (Line 513).

3. The paper exaggerates its own contribution and ignores relevant previous work. There are a huge number of works that improve federated optimization with different client selection strategies, which are ignored by this paper in terms of theory, experiments, and general framing (see [2] and its references). Some examples of exaggerated language that I find inappropriate:

• Abstract: "Virtually all FL methods operate in the following manner..." This claim is not accurate; there are many works in FL that use peer-to-peer communication [3], asynchronous communication [4], etc. Further, I fail to see how the proposed algorithms of this paper do not also fall into the category described in the abstract.
• Line 524: "This foundational work showcases a pivotal shift in federated learning strategies". I don't believe that this work departs very far at all from previous work in FL (e.g. [5] and related works). In my opinion, this kind of self-aggrandizing is not appropriate for a scientific publication.

4. The message of the paper is not totally coherent. The abstract talks about "cohort squeeze" and novel communication principles, but most of the paper actually deals with client selection strategies within the standard intermittent communication structure. The experiments discuss local vs. global communication (Section 3.6), which seems to be the connection to the "cohort squeeze" of the title and abstract, but this section makes up a very small part of the paper's technical content. Perhaps I have missed a connection between the content of the abstract and the content of the main text.

[1] Cho, Yae Jee, et al. "On the convergence of federated averaging with cyclic client participation." International Conference on Machine Learning. PMLR, 2023.

[2] Fu, Lei, et al. "Client selection in federated learning: Principles, challenges, and opportunities." IEEE Internet of Things Journal (2023).

[3] Beltrán, Enrique Tomás Martínez, et al. "Decentralized federated learning: Fundamentals, state of the art, frameworks, trends, and challenges." IEEE Communications Surveys & Tutorials (2023).

[4] Xu, Chenhao, et al. "Asynchronous federated learning on heterogeneous devices: A survey." Computer Science Review 50 (2023): 100595.

[5] Grudzień, Michał, Grigory Malinovsky, and Peter Richtárik. "Improving accelerated federated learning with compression and importance sampling." arXiv preprint arXiv:2306.03240 (2023).

1. The federated learning setting considered in this paper are limited. Specifically, this paper assume that the objective function is strongly convex (Assumption 2.2). This not only significantly simplifies the theoretical analysis to achieve stronger convergence results, but also is not of practical interests since most ML models are non-convex.

2. The algorithm design is only a minor variation of the existing algorithmic framework in federated learning, which has been well-explored. Algorithmic ideas such as proximal point optimization, various types of client sampling, and multiple local updates for each cohort are not new and have been considered in the literature. It's unclear what are the major novelty in this paper.

3. As a consequence of the above factors, the theoretical contributions of this paper are marginal, since most of the proof details in Appendix F are quite standard.