CoBo: Collaborative Learning via Bilevel Optimization

Authors: diba  hashemi
martin  jaggi
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 20:10:31

Abstract: Collaborative learning is an important tool to train multiple clients more effectively by enabling communication among clients. Identifying helpful clients, however, presents challenging and often introduces significant overhead. In this paper, we model **client-selection** and **model-training** as two interconnected optimization problems, proposing a novel bilevel optimization problem for collaborative learning.
We introduce **CoBo**, a *scalable* and *elastic*, SGD-type alternating optimization algorithm  that efficiently addresses these problem with theoretical convergence guarantees. Empirically, **CoBo** achieves superior performance, surpassing popular personalization algorithms by 9.3% in accuracy on a task with high heterogeneity, involving datasets distributed among 80 clients.

Main paper: 
 introduction: 
 In a classic collaborative learning scenario, n clients, each with a distinct machine learning task, seek solutions that potentially outperform their individual solvers through a collective effort. Common collaborative learning frameworks generally alternate between training local models on individual datasets and synchronizing updates among collaborators. More concretely, during the computation step, client i ∈ [n] trains a d-dimensional model x i ∈ R d to minimize its loss function, f i : R d → R. In the subsequent communication step, client i exchanges updates with clients, potentially benefiting from collaboration.Despite the plethora of collaborative learning frameworks, the ideal approach to collaborate remains under-exploited. The FEDAVG [28, 18] algorithm learns a single global model over pooled datasets from all clients, i.e., min x∈R d 1 n n i=1 f i (x). However, due to heterogeneous data distributions among clients, a global model may significantly underperform compared to personal models trained on local datasets for certain clients, which can discourage their participation in collaborative training [29] . DITTO addresses this issue by training personal models with a regularization term that penalizes deviations from a global model [24] . Although DITTO enables personal models to leverage the global model, it offers only a coarse-grained level of collaboration. In instances where clients' data exhibit significant differences, the DITTO algorithm is constrained to facilitating collaboration at a global level, thereby neglecting the inherent heterogeneity among clients.Clustering-based federated learning algorithms have been developed to accommodate scenarios in which clients' data originate from multiple clusters [11, 39] . Nevertheless, these algorithms typically inherit the limitations associated with clustering techniques, including the need to predetermine the number of clusters, initialize cluster centers, and other such prerequisites, which can diminish their practical utility.In this paper, we propose a bilevel optimization framework to enhance collaborative learning by discovering better structural relationships among clients. The inner problem focuses on optimizing a binary collaborator selection variable w ij ∈ {0, 1}, determined based on a gradient alignment measure for each pair of clients. In the outer problem, we train personalized models X ∈ R n×d , incorporating a penalization term that accounts for the distances between clients, as dictated by the collaboration weights established in the inner problem.The contributions of this paper can be summarized as follows:• We model collaborative learning through a novel bilevel optimization formulation that yields more generalizable solutions by fully exploiting the inherent structure of collaboration.• We propose COBO, an SGD-type alternating optimization algorithm that efficiently solves the bilevel problem. COBO scales with the number of clients n and is elastic to the number of clients.• We prove that COBO enjoys theoretical convergence guarantees for collaborative learning with cluster structures.• Empirically, COBO surpasses popular personalized federated learning baselines in experiments involving highly heterogeneous federated learning settings and Large Language Models (LLMs). 
problem formulation: 
 In this paper, we model collaborative learning as a bilevel optimization problem, where personalized models X ∈ R d×n are trained in the outer problem, and collaborative weights W ∈ R n×n are determined by the inner problem. More concretely, EQUATIONw ij ∇f i x i + x j 2 , ∇f j x i + x j 2 ∀ i, j ∈ [n],(Client-Selection)where ρ > 0 is a hyperparameter for penalization. We break down the formulation as follows.Outer problem: training personalized models. In the outer problem (Model-Training), client i trains its model x i by minimizing its loss function f i , along with penalizing its distances to neighboring models, e.g. x j , as weighted by w ⋆ ij > 0. Our formulation is similar to DITTO [24] , but with two key differences: First, DITTO uses uniform and fixed collaboration weights and penalizes the distance between x i and a global model, whereas we penalize the distances between pairs of clients and adjust the collaboration weights during training. Consequently, when client tasks are heterogeneous-such as clients drawn from clusters-the performance of a global model deteriorates, and DITTO's local models cannot benefit from finegrained collaboration. In contrast, our method is able to exploit such structure and achieve better performance in diverse settings.Inner Problem: Finding Collaborators. In the inner problem, we decompose the task of optimizing W ∈ R n×n into independent sub-problems, one for each entry of W . We relax the binary collaborator selection variable w ij ∈ {0, 1} to a continuous weight w ij ∈ [0, 1]. As the objective function is linear with respect to w ij , and the domain is convex, optimization algorithms such as Frank-Wolfe [10, 17] or projected gradient descent can efficiently find the maximizers, which occur at 0 or 1.It is important to note that w ⋆ ij does not imply a permanent connection between clients i and j, but rather a temporary assessment based on the current states of x i and x j .A simple inner problem with two clients is illustrated in Figure 1 . The f 1 , f 2 are their loss functions, and suppose µ 1 , µ 2 , and µ are the minimizers of f 1 , f 2 , and1 x A 1 , x A 2 x B 1 , x B 2 2 1 x C 1 x C 2 x C Model weights x A , x B , x C Minimizers of f 1 and f 2 Minimizer of 1 2 (f 1 + f 2 ) Averaged model weight Grads for X-update f 1 (x 1 ), f 2 (x 2 )Grads for W-updatef 1( x1 + x2 2 ) , f 2( x1 + x2 2 )Figure 1 : Diagram of the inner problem (Client-Selection) represented through a contour of 1 2 (f 1 +f 2 ). The blue arrows → are gradients computed at middle point 1 2 (x 1 + x 2 ) to determine connectivity. The red arrows → represent gradients computed at local models to update model weights.• Point A: The model x A is far away from µ, i.e., ∥x A -µ∥ >> max i ∥µ i -µ∥. The descent directions of the clients have a positive inner product; therefore w 12 = 1. Collaboration at this stage speeds up training.• Point B: The model x B is closer to µ, i.e., ∥x B -µ∥ ∼ ∥µ i -µ∥. In this case, moving closer to the minimizer µ of 1 2 (f 1 + f 2 ) no longer helps both clients get closer to the minimizers of their own losses µ i . The inner problem yields w 12 = 0 and the clients disconnect.• Point C: The models x C 1 and x C 2 are already disconnected. The gradients computed at their midpoint suggest they should remain disconnected; thus w 12 = 0.Because collaboration weights in Client-Selection are determined in a pairwise fashion, our formulation, unlike clustering-based methods [11, 39] , does not require knowledge of cluster sizes, allowing clients to join and leave during collaborative training. This elasticity enables our method to be applicable in a wider range of scenarios. Remark 1 (Extensions). While Client-Selection is defined over a box constraint W ∈ [0, 1] n×n , it can be easily extended to other convex domains. For example, in all-for-one type collaborative training, the weights are optimized over a simplex. The experiment on language models is deferred to Section 4.3.We propose a novel stochastic gradient descent (SGD)-type alternating optimization algorithm, termed COBO, to solve the bilevel optimization problem defined by (Model-Training) and (Client-Selection). The algorithm alternates between updating the model variables X and the collaboration weights W .In each iteration t, we first fix the model variables {x t i } n i=1 and update the collaboration weights by applying projected gradient ascent with step size γ > 0 to the inner problem (Client-Selection):EQUATIONNext, with the updated collaboration weights {w t+1 ij } are fixed, we optimize the model variables {x i } n i=1 using the following update rule derived from the outer problem (Model-Training):EQUATIONwhere η > 0 is the step size for updating the model variables. This alternating process is repeated until convergence.The detailed implementation of the algorithm is provided in Algorithm 1. In this implementation, the full gradients {∇f i } i∈[n] in (1) and ( 2) are replaced by their stochastic estimates. Additionally, collaborative weights are updated with a probability of O( ∀ i ∈ [n] x 0 i = x 0 ∈ R d ; Penalization parameter ρ > 0; W 0 ∈ R n×n where w 0 ij = 1, ∀ i, j ∈ [n];Step size η, γ > 0. 1: for round t = 0, 1 . . . , T do2: Call W t+1 ← Client-Selection({x t i } i∈[n] , W t ) 3:for client i = 1, . . . n do 4:Draw sample ξ i ∼ D i and compute stochastic gradient g t i ∈ R d of f i (x t i ) and updateEQUATIONend for 6: end for 7: Output: Uniform randomly select s ∈ [T ] and return {x s 0 , . . . , x s n } and W s . for each pair of clients (i, j)where i ̸ = j ∈ [n] do 11:if with a probability 1/n, then 12:Compute the average model z ij = 1 2 (x i + x j ).Compute stochastic gradient g i←i and g i←j for f i (z ij ) and f j (z ij ) respectively,EQUATIONend if15:end for 16:return updated selection variables W 17: end procedure O(n) gradients. This results in an overhead comparable to that of standard decentralized learning methods [25, 19] , thereby enabling client selection with minimal additional cost.Compared to federated clustering algorithms, which require global synchronization before applying clustering oracles, the inner problem (Client-Selection) in COBO is solved in a pairwise fashion. This pairwise approach makes the algorithm non-blocking and robust to stragglers, providing greater flexibility and efficiency. Not all pairwise weights have to be computed in each iteration. In Table 1 we compare the performance of multiple edge-sampling strategies. 
theoretical results: 
 In this section, we define assumptions in collaborative learning settings and show that COBO converges to stationary points. The following assumptions regarding the local optimization objective f i are commonly adopted in the literature [1, 19] :(A1) L-smooth. For all x and y in R d and i ∈ [n], the loss function f i has L-Lipschitz gradients, i.e.∥∇f i (x) -∇f i (y)∥ ≤ L∥x -y∥ .(A2) Noise bound. For all x ∈ R d and i ∈ [n], there exists σ 2 > 0 such that the stochastic gradient has bounded noiseE ξ ∥∇f i (x; ξ) -E ξ [∇f i (x; ξ)]∥ 2 ≤ σ 2 .(A3) Global minimum. For all i ∈ [n], the loss function f i has a global lower bound f * i .The next assumption characterizes the possible relationships between clients. In the first case, when reaching the stationary point x of their joint objective f i + f j , then by (5) implies that ∇f i (x) = ∇f j (x) = 0 client i and j reach their own stationary points. In the second case, when client i reaches its stationary point, the gradient of j is lower bounded by a positive constant, meaning they don't share stationary points. This leads to eventual (A4) Collaborativeness. If clients i and j are collaborative, then there exists M ij > 0 such thatEQUATIONOtherwise, there exists ζ 2 ij > 0 such thatEQUATION)This assumption is similar to [39, Assumptions 4,5], but we define relations for pairs of clients instead of clusters. In the next example, we use quadratics to demonstrate (A4) Example 2. Assume that there are K clusters with[n] = ∪ k∈[K] C k and C k ∩ C k ′ = ∅ for all k ̸ = k ′ ∈ [K].Consider the k-th cluster with center µ k and client i ∈ C k , the loss function isf i (x) = ai 2 ∥x -µ k ∥ 2 2where a i > 0. Then for clients i, j in the same cluster, i.e. i, j ∈ C k∥∇f i (x) -∇f j (x)∥ 2 2 = (a i -a j ) 2 ∥x -µ k ∥ 2 2 = (a i -a j ) 2 (a i + a j ) 2 ∥∇f i (x) + ∇f j (x)∥ 2 2 .TheM ij = |ai-aj | ai+aj in this case. On the other hand, for i ∈ C k and j ∈ C k ′ and µ k ̸ = µ k ′ , ∥∇f i (x)∥ 2 2 + ∥∇f j (x)∥ 2 2 = a 2 i ∥x -µ k ∥ 2 2 + a 2 j ∥x -µ k ′ ∥ 2 2 = a 2 i a 2 j (a 2 i + a 2 j ) 2 ∥µ k -µ k ′ ∥ 2 2where the lower boundζ 2 ij = a 2 i a 2 j (a 2 i +a 2 j ) 2 ∥µ k -µ k ′ ∥ 2 2 > 0.Finally, we derive a convergence theorem with the assumption that clients are drawn from clusters, as e.g. in [33, Assumption 2] .(A5) Cluster. All clients are drawn from clusters where within each cluster clients share stationary points. Theorem I. Suppose Assumption 1,2,3,4,5 hold true. Suppose that COBO solves (4) with minibatch size b. Consider clients i and j in the same cluster C of size c. Suppose thatM 2 ij ∈ (0, 1 5 ), b ≥ 2 c 2 2Lη(c -2)σ 2 and ζ 2 ik ≥ ∥∇f i (x) + ∇f k (x)∥ 2 2 for all x and k. Let ρ ≥ √ 3L c and step size η ≤ min    2 σ √ LT 1 c 2 i,j∈C fij z 0 ij -f ⋆ ij , 1 2 √ 3L    .The consensus distance also converges to 0, i.e.1 c 2 T T -1 t=0 i,j∈C E x t+1 i -x t+1 j 2 2 ≤ 6M 2 ij ρ 2 c 2 Lσ 2 c 2 T i,j∈C fij z 0 ij -f ⋆ ij .Moreover, the gradient norm is upper bounded.1 c 2 T T -1 t=0 i,j∈C E ∇ fij z t ij 2 2 ≤3 Lσ 2 c 2 T i,j∈C fij z 0 ij -f ⋆ ij .This theorem suggests that clients inside the same cluster gradually reach consensus. This clusterlevel consensus model reaches stationary point of their losses by (5) . Note that a larger penalization parameter ρ and smaller values of M 2 ij lead to faster convergence, which aligns with our expectations. Note that M ij in (A4) measures how well i,j collaborate. A smaller M ij leads to better consensus distance in Theorem I, with M ij = 0 leading to identical data distribution. The following corollary states the convergence of norm of client gradient of model x i .Corollary II. Under same conditions as Theorem I, ∥∇f i (x t i )∥ 2 2 converges at a similar rate1 c 2 T T -1 t=0 i,j∈C ∇f i x t i 2 2 ≤4 Lσ 2 c 2 T i,j∈C fij z 0 ij -f ⋆ ij .Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?Answer: [Yes] Justification: The assumptions are given in Section 3. 
experiments: 
 In this section, we present three experiments to demonstrate the practical effectiveness of COBO. In the first two experiments, we benchmark COBO in both a cross-silo federated learning setup involving 8 clients and a cross-device setup with 80 clients, using the CIFAR-100 dataset for multi-task learning [21] . In the third experiment, we train language models on subsets of the Wiki-40B dataset while learning domain weights within a simplex [12] . Compared to state-of-the-art personalized federated learning baselines, COBO obtains personalized models of higher quality and correctly identifies cluster structures. Details of the experiments, including descriptions of the architectures and the system setup, are deferred to Appendix B.Throughout the experiments, we use popular federated learning baselines such as FEDAVG [28] , Federated Clustering (abbreviated as FC) [39] , DITTO [24] , IFCA [11] , and an oracle algorithm.The definition of the oracle baseline varies in each setup and will be discussed case by case. Note that we additionally provide clustering-based algorithms, i.e., FC and IFCA, with the actual number of clusters. Their experimental statistics reported in this section, such as accuracy and perplexity, include this advantage. In addition to previous baselines, we also compare COBO with a collaborative fine-tuning approach for large language models that leverages performance on validation data to determine the collaboration weights [38] (referred to as Validation Based in Table 2 ).In this experiment, we evaluate the performance of COBO by comparing the average accuracies of local models against those of established collaborative learning baselines. Our objective is to assess how effectively COBO discerns and leverages the structure of data clusters relative to other collaborative learning algorithms.We simulate a cross-silo multi-task environment where training a single model across all clients yields poor performance, thus highlighting the necessity for client selection. Our experimental configuration consists of 4 clusters, each containing 2 clients utilizing the ResNet-9 model [13] . To encourage collaboration within clusters, we randomly allocate half of the dataset to each client in a cluster. To differentiate between clusters, we introduce label diversity by flipping the image labels in each cluster using distinct random seeds. This process ensures that each class maintains unique labels across all clusters, effectively creating a scenario where a universally trained model would not be optimal, thereby necessitating personalized models that can cater to the specific label distribution of each cluster.In this context, collaboration among clients within the same cluster is advantageous, as their datasets are complementary. There are two primary reasons why collaboration between different clusters may not be beneficial: (1) the dataset available to clients within each cluster is identical, negating the incentive to collaborate with clients from other clusters; and (2) the label flipping across clusters could mean that inter-cluster collaboration might actually degrade local model performance.Given these considerations, we designate an oracle algorithm for our scenario: FEDAVG implemented separately within each cluster. This ensures that collaboration is confined to where it is most beneficial. Additionally, the oracle collaboration matrix is defined to be a block-diagonal matrix, with entries of 1 for pairs of clients within the same cluster (indicating collaboration) and entries of 0 for pairs from different clusters (indicating no collaboration). This matrix serves as a benchmark for the ideal collaboration structure in our simulated environment.To enable the practical application of COBO, we sample pairs of clients in each iteration to update their collaboration weights. We begin by examining the impact of various sampling strategies on the performance of COBO. The primary approach involves sampling with a constant probability of O(1/n). Additionally, we observe that COBO identifies an appropriate collaboration matrix early in the training process, motivating the use of a time-step-dependent sampling rate, O(1/t). We also implement a mixed strategy: employing the constant sampling rate, O(1/n), for the initial 0.2% of iterations, followed by a switch to the time-dependent sampling rate, O(1/t), for the remainder of the training. A comparison of these strategies with the non-sampling oracle, where all pairs are updated in every iteration, is presented in Table 1 . While COBO demonstrates consistent performance across all sampling strategies, achieving results close to those of the non-sampling oracle, the mixed strategy shows a slight performance advantage. To further assess performance, we trained COBO and other baseline algorithms for a total of 40,000 iterations. Figure 2b presents the accuracy diagram. We observe that COBO almost reaches the performance bound established by the Oracle. Moreover, COBO achieves a fixed accuracy of 60% in 4,500 iterations, which is 30% faster than DITTO. For better comparison, the values of accuracy and loss are reported in Table 2 . Additionally, the evolution of the collaboration matrix for clustering algorithms and COBO is illustrated in Figure 3 . COBO starts to identify clients with similar label permutations as early as 300 iterations and stabilizes in less than 5,000 iterations (12.5% of the training phase). IFCA always degenerates to one fully connected cluster, while FC periodically suffers from clustering mistakes even at the end of training.Figure 2a presents the results of the cross-silo experiment under various configurations to further assess the robustness of COBO. First, we modify the fraction of the dataset allocated to each client. Intuitively, the total amount of data available to a cluster directly impacts the performance of COBO.Next, we experiment with different numbers of clusters, each containing two clients, and observe that The collaboration matrix of COBO already starts to look similar to oracle matrix within as low as 300 iterations (0.75% of the total iterations), and converges to it within 5000 iterations (12.5% of the total iterations). On the other hand, IFCA yields a fully-connected matrix while FC occasionally diverges from the achieved cluster structures (e.g., iterations 300, 5000, and 40000), even at the end of training.the number of clusters does not significantly affect COBO 's accuracy. Additionally, we investigate the effect of varying the number of clients per cluster while maintaining a fixed total of four clusters. In this setup, the dataset is partitioned among clients within each cluster, resulting in less data per client as the cluster size increases. Despite this, COBO leverages collaboration to maintain robust performance even with larger cluster sizes.In this experiment, we demonstrate the performance of COBO in a challenging cross-device federated learning setting characterized by significant data heterogeneity. We create 10 clusters of varying sizes: 2 clusters consist of 6 clients each, another 2 comprise 7 clients each, and so on. Each cluster is allocated data from 10 distinct classes out of the total 100 classes available in the CIFAR-100 dataset, ensuring that the data across clusters are disjoint. Within each cluster, the data are distributed uniformly at random among the clients. We proceed to train individual ResNet-9 models [13] on each client's data for a total of 20,000 iterations. This setup allows us to observe the behavior of COBO and its ability to handle both the quantity and diversity of data across different client groups and cluster sizes.We define the oracle algorithm and the corresponding collaboration matrix in the same manner as in Section 4.1. Note that while we manually create the clusters, inter-cluster collaboration may still be helpful in practice, and it is impossible to know the actual ground truth in this case. Consequently, we recognize that the oracle may not correspond to the optimal performance. Nevertheless, this oracle still exhibits superior performance compared to other baselines that lack prior knowledge of the data distribution among clients, as evidenced by the results presented in Table 2 . The collaboration matrix and accuracy plots are deferred to Figure 5 and Figure 6 in Appendix B, respectively.In this challenging experiment, COBO surpasses all other baselines by at least 5.7% in accuracy. This supports the conclusion that COBO scales well with the size of collaborative learning and effectively exploits collaboration weights among clients at a fine-grained level. Recently, Large Language Models (LLMs) have gained significant popularity due to their ability to effectively solve challenging tasks. Their downstream performance can be further enhanced by finetuning; however, the scarcity of data often leads to inferior performance and necessitates collaboration. Therefore, we conduct an experiment with four clients, each having a pre-trained GPT-2 base model3 with 124 million parameters in total [32] , and a subset of articles from the Wiki-40B dataset [12] in one of the following four languages: Catalan, Spanish, German, or Dutch. We use LoRA for the Self-Attention and MLP layers for fine-tuning, which accounts for 0.47% of the full parameters [14] . For data-hungry tasks, such as those involving LLMs, contributions from all domains are valuable. Clustering methods fall short in this aspect due to their binary, discrete outputs, which do not capture the nuanced degrees of collaboration needed. COBO addresses this limitation by allowing for a continuous range of collaboration intensities, achieved by a simple yet effective modification to the projection domain in (1) . Specifically, we employ a probability simplex, denoted as ∆ i = {w ij ≥ 0, j w ij = 1}, as the domain of the inner problem.In Table 2 , we compare the perplexity of COBO with baselines after 500 iterations, when FE-DAVG converges. There are no oracle domain weights in this experiment due to the complex coherence among languages; therefore, we omit the oracle algorithm in the table. COBO achieves the best perplexity among all algorithms. In Figure 4 , we demonstrate the domain weights learned for the Catalan language. Overall, Catalan assigns the highest collaboration weight to Spanish, which is reasonable considering the similarity between the two languages.Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?Answer: [Yes] Justification: The experiments are clearly defined, and all the necessary hyperparameters for reproducing them are presented in the Appendix. 
related work: 
 Personalized federated learning. Personalized federated learning has received significant attention in recent years due to its potential to tailor models to individual user data while benefit from collaboration [33, 35, 36, 9, 22, 3] . There are various flavors of personalized federated learning. DITTO trains personalized models by incorporating a regularization term that penalizes the divergence from a global model [24] . Many personalization works assume that clients are drawn from clusters. For example, Marfoq et al. [27] use K-nearest neighbors (KNN) to determine collaborators. Mansour et al. [26] , Ghosh et al. [11] , Werner et al. [39] develop K personalized models and assign clients to clusters based on criteria such as minimum function values or gradient similarities. Additionally, Even et al. [7] provided theoretical insights by establishing lower bounds, which demonstrate that the optimal gradient filtering strategy involves clustering clients with identical optima.In federated learning, client selection is often performed by simultaneously minimizing task losses and collaborative weights in a single-level objective function. Zantedeschi et al. [40] minimize task losses augmented with a penalization term w ij ∥x ix j ∥ 2 2 , similar to our outer problem. However, optimizing w ij directly can lead to a degenerate solution (w ij = 0), which necessitates an additional penalization for small w ij values. Smith et al. [34] approach multi-task learning by minimizing task losses with a more sophisticated penalization term that accounts for the relationships between tasks. This formulation requires the client-selection function to be consistent with client selection, which can negatively impact performance. Apart from multi-task federated learning, a similar bilevel optimization formulation has been used by Le Bars et al. [23] to find a sparse mixing matrix while training a consensus model in the outer problem.Bilevel optimization and alternating optimization. Bilevel optimization is a powerful tool which models a broad range of problems, such as reinforcement learning [6, 30, 16, 15, 37, 31] , and linearlysolvable Markov decision process [5] , meta-learning [8, 20] , etc. A typical bilevel optimization problem, as the name indicates, consists of an outer and an inner optimization problem whose variables are inter-dependent. Typical bilevel optimization solvers requires hessian information which is usually expansive to acquire [8] . On the other hand, alternating optimization tools has been used be used to solve bilevel optimization problem [2, 4] . While in general there is no universal convergence guarantees for alternative optimizations, the special structure of our inner problem ensures the convergence of COBO to the stationary point.Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?Answer: [Yes] Justification: The code is available in supplementary materials as well as on an anonymous Github repository, and we only used publicly available datasets in our experiments. 
conclusions: 
 Existing collaborative learning algorithms only allow coarse-grained collaboration, which leads to inferior performance in practice. To address this issue, we model collaborative learning as a special bilevel optimization problem where client selection is based on the optimization of a linear function of gradient alignment measure for each pair of clients. In addition, we propose an efficient SGD-type alternating optimization algorithm COBO which is scalable, elastic, and enjoy theoretical guarantees. Besides, COBO empirically outperforms popular personalized federated learning algorithms in realistic collaborative learning problems.Limitations. In this work, we do not take privacy into consideration. The existing algorithm requires exchanging gradients between collaborators when updating weight w which may raise privacy concerns. We defer the discussion of privacy-preserving collaborative learning framework to future work.Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?Answer: [Yes] Justification: The full training setup is described in the main text. The hyper-parameters are also available both for COBO and other baselines either in the main text or appendix.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: Not applicable. 
 
 Appendix: Let z t ij := 1 2 (x t i + x t j ) be the average iterate of x t i and x t j and fij (x) := 1 2 (f i (x) + f j (x)) be an averaged objective function. Lemma 3. Suppose (A1) holds true. Let η ≤ 1 2L . Then for i, j in the same cluster C of size c1 c 2 i,j∈C ∇ fij z t ij 2 2 ≤ η 1 c 2 i,j∈C fij z t ij -E[ fij z t+1 ij ] + D 1 1 c 2 i,j∈C x t i -x t j 2 2 + 3n 2 ρ 2 1 cn i∈C n k=1 | E h [w t+1 ik ] -w ⋆ ik | 2 x t i -x t k 2 2 + Lησ 2 2 .where coefficient D 1 is defined asD 1 (L, c, η, b, σ 2 ) := 3L 2 4 + 3c 2 ρ 2 + Lηρ 2 (c -2)2σ 2 b .A. [•] = E h [E g [•]]. By the L-smoothness assumption (A1) and bounded noise assumption (A2)E h E g fij z t+1 ij ≤ fij z t ij + ∇ fij z t ij , E h E g z t+1 ij -z t ij + L 2 ∥E h E g [z t+1 ij -z t ij ]∥ 2 2 + L 2 E h E g ∥z t+1 ij -z t ij -E h E g [z t+1 ij -z t ij ]∥ 2 2 .The last quantity, i.e. variance of z t+1 ij -z t ij , can be expanded as followsE h E g ∥z t+1 ij -z t ij -E h E g [z t+1 ij -z t ij ]∥ 2 2 = E h E g ∥z t+1 ij -E h E g [z t+1 ij ]∥ 2 2 =E h E g ∥z t+1 ij ± E h [z t+1 ij ] -E h E g [z t+1 ij ]∥ 2 2 =E g ∥E h [z t+1 ij ] -E h E g [z t+1 ij ]∥ 2 2 + E h E g ∥z t+1 ij -E h [z t+1 ij ]∥ 2 2 =E g ∥E h [z t+1 ij ] -E h E g [z t+1 ij ]∥ 2 2 + E h E g ∥z t+1 ij -E h [z t+1 ij ]∥ 2 2 =η 2 E g 1 2 (g t i + g t j -∇f i (x t i ) -∇f j (x t j )) 2 2 + E h E g ∥z t+1 ij -E h [z t+1 ij ]∥ 2 2 ≤ η 2 σ 2 4 + E h E g ∥z t+1 ij -E h [z t+1 ij ]∥ 2 2 =:Tijwhere the linear term vanishes due to the expectation of h in the second equality and the definition of z t+1 ij and (3) for the third equality. ThenE fij z t+1 ij ≤ fij z t ij + ∇ fij z t ij , E z t+1 ij -z t ij + L 2 E[z t+1 ij -z t ij ] 2 2 + Lη 2 σ 2 4 + L 2 T ij .Expand the inner product with equality -⟨x, y⟩ = -EQUATIONEQUATIONE[z t+1 ij ] -z t ij η 2 2 - η 2 ∇ fij z t ij 2 2 + Lη 2 σ 2 4 + L 2 T ijwhere η ≤ 1 2L is applied in the last inequality. Then the upper bound ofEQUATIONThe T can be upper bounded by expanding z t+1 i with ( 2)EQUATION.Bound T 1 : Use L-smoothness of f i and f j . Take expectation with respect to g t i and g t j which are unbiased estimates of ∇f i (x t i ) and ∇f j (x t j )T 1 = ∇ fij z t ij - 1 2 ∇f i (x t i ) + ∇f j (x t j ) 2 2 ≤ L 2 2 z t ij -x t i 2 2 + L 2 2 z t ij -x t j 2 2 = L 2 4 x t i -x t j 2 2 .Bound T 2 : Use Cauchy-Schwarz inequality and C has a cluster size of cT 2 ≤ cρ 2 2 n k=1 w ⋆ ik x t i -x t k 2 2 + n k=1 w ⋆ jk x t j -x t k 2 2 .Bound T 3 : Use Cauchy-Schwarz inequalityT 3 ≤ nρ 2 2 n k=1 | E h [w t+1 ik ] -w ⋆ ik | 2 x t i -x t k 2 2 + n k=1 | E h [w t+1 jk ] -w ⋆ jk | 2 x t j -x t k 2 2 . Sum up ∇ fij z t ij 2 2for all of i, j in the same cluster C for (7) yields i,j∈CEQUATIONNow we expand T ij as followsT ij =E h E g ∥z t+1 ij -E h [z t+1 ij ]∥ 2 2 =E h E g η 2 (g t i + g t j ) + ηρ 2 n k=1 (w t+1 ik (x t i -x t k ) + w t+1 jk (x t j -x t k )) -E h η 2 (g t i + g t j ) + ηρ 2 n k=1 (w t+1 ik (x t i -x t k ) + w t+1 jk (x t j -x t k )) 2 2   =E h   ηρ 2 n k=1 ((w t+1 ik -E h [w t+1 ik ])(x t i -x t k ) + (w t+1 jk -E h [w t+1 jk ])(x t j -x t k )) 2 2   = η 2 ρ 2 4 k̸ =i,j E h w t+1 ik -E h [w t+1 ik ] 2 2 ∥x t i -x t k ∥ 2 2 + E h w t+1 jk -E h [w t+1 jk ] 2 2 ∥x t j -x t k ∥ 2 2where we use the independence of random variables in the last equality.Average the above equality over i, j∈ C yields1 c 2 ij T ij = η 2 ρ 2 (c -2) 2c 2 i,j E h w t+1 ij -E h [w t+1 ij ] 2 2 ∥x t i -x t j ∥ 2 2 ≤ η 2 ρ 2 (c -2) 2c 2 4σ 2 b i,j ∥x t i -x t j ∥ 2 2 .Therefore, (8) is upper bounded with i,j∈CEQUATIONx t+1 i -x t+1 j = x t i -x t j -η g t i -g t j -ηρ n k=1 w t+1 ik (x t i -x t k ) -w t+1 jk (x t j -x t k ) .As i and j belong to the same cluster (i.e., w ⋆ ij = 1), we add ±2ηρn k=1 w ⋆ ik (x t i -x t k ) x t+1 i -x t+1 j =(1 -2ηρc)(x t i -x t j ) -η g t i -g t j -ηρ n k=1 (w t+1 ik -w ⋆ ik )(x t i -x t k ) -(w t+1 jk -w ⋆ jk )(x t j -x t k ) .Compute the norm of x t+1 i -x t+1 j and choose ηρ ≤ 1 2c to use Jensen's inequalityEQUATION.Expand the right-hand side with Cauchy-Schwarz inequalityE g [x t+1 i -x t+1 j ] 2 2 ≤(1 -2ηρc) x t i -x t j 2 2 + 4ηρc 1 2c n k=1 (w t+1 ik -w ⋆ ik )(x t i -x t k ) -(w t+1 jk -w ⋆ jk )(x t j -x t k ) 2 2 + 4ηρc 1 2ρc ∇f i (x t i ) -∇f j (x t j ) 2 2 ≤(1 -2ηρc) x t i -x t j 2 2 + 8ηρc 1 2c n k=1 (w t+1 ik -w ⋆ ik )(x t i -x t k ) 2 2 + 8ηρc 1 2c n k=1 (w t+1 jk -w ⋆ jk )(x t j -x t k ) 2 2 + 4ηρc 1 2ρc ∇f i (x t i ) -∇f j (x t j ) 2 2 ≤(1 -2ηρc) x t i -x t j 2 2 + 2nηρ c n k=1 |w t+1 ik -w ⋆ ik | 2 ∥x t i -x t k ∥ 2 2 + 2nηρ c n k=1 |w t+1 jk -w ⋆ jk | 2 ∥x t j -x t k ∥ 2 2 + η ρc ∥∇f i (x t i ) -∇f j (x t j )∥ 2 2 =:T .The last term T can be upper bounded by adding ±∇f i z t ij ± ∇f j z t ij and use L-smoothness assumption (A1) of f i and that i, j belong to the same cluster (A4)EQUATION.Note thatE[∥x t+1 i -x t+1 j ∥ 2 2 ] = η 2 E[∥g t i -g t j -E[g t i -g t j ]∥ 2 2 ] + E h ∥E g [x t+1 i -x t+1 j ]∥ 2 2and independence between randomness on worker i and j where expectation is take to all of the randomness until time tBy averaging for all i, j ∈ CEQUATION.Use the previous Lemma 3 to bound i,j∈C ∇ fij (z t ij )2 2 1 c 2 i,j∈C E x t+1 i -x t+1 j 2 2 ≤ 1 -2ηρc + 3ηL 2 2ρc 1 c 2 i,j∈C x t i -x t j 2 2 + 4n 2 ηρ c 1 cn i∈C n k=1 E h |w t+1 ik -w ⋆ ik | 2 ∥x t i -x t k ∥ 2 2 + ησ 2 2ρc 1 c 2 + 3ηM 2 ij ρc   2 η 1 c 2 i,j∈C fij z t ij -E fij z t+1 ij + D 1 1 c 2 i,j∈C x t i -x t j 2 2   + 3ηM 2 ij ρc 3n 2 ρ 2 1 cn i∈C n k=1 E h |w t+1 ik -w ⋆ ik | 2 x t i -x t k 2 2 + Lησ 2 2 .Rearrange the terms1 c 2 i,j∈C E x t+1 i -x t+1 j 2 2 ≤ 1 -2ηρc + 3ηL 2 2ρc + 3ηM 2 ij ρc D 1 1 c 2 i,j∈C x t i -x t j 2 2 + 4n 2 ηρ c + 3ηM 2 ij ρc 3n 2 ρ 2 1 cn i∈C n k=1 E h |w t+1 ik -w ⋆ ik | 2 ∥x t i -x t k ∥ 2 2 + 3ηM 2 ij ρc 2 η 1 c 2 i,j∈C fij z t ij -fij z t+1 ij + ησ 2 2ρc 1 c 2 + 3ηM 2 ij ρc Lησ 2 2 .We would like to achieve 1-2ηρc + 3ηL 2 2ρc + 3ηM 2 ij ρc D 1 ≤ (1 -ηρc) by • letting ρ ≥ √ 3L c , s.t. 3ηL 2 2ρc ≤ ηρc 2 .• letting b ≥ 2 c 2 2Lη(c -2)σ 2 and ρ ≥ √ 3Lc and M ij ≤ 1 5 , the following inequality hold true3ηM 2 ij ρc 3L 2 4 + 3c 2 ρ 2 + Lηρ 2 (c -2)2σ 2 b ≤ 3ηM 2 ij ρc 15 4 ρ 2 c 2 ≤ 45 4 ρcηM 2 ij ≤ 1 2 ηρc.Using the same requirement that ρ ≥√ 3L c and M ij ≤ 1 5 4n 2 ηρ c + 3ηM 2 ij ρc 3n 2 ρ 2 ≤ 4n 2 ηρ 2 √ 3L + 3 √ 3ηM 2 ij n 2 ρ 2 L ≤ 5ηn 2 ρ 2 L .The upper bound of 1/c 2 i,j∈Cx t+1 i -x t+1 j 2 2 can be simplified to 1 c 2 i,j∈C E x t+1 i -x t+1 j 2 2 ≤ (1 -ηρc) 1 c 2 i,j∈C x t i -x t j 2 2 + 5nη 2 ρ 2 L 1 cn i∈C n k=1 E h |w t+1 ik -w ⋆ ik | 2 ∥x t i -x t k ∥ 2 2 + 6M 2 ij ρc 1 c 2 i,j∈C fij z t ij -E fij z t+1 ij + 1 c 2 + 3ηLM 2 ij ησ 2 2ρc .A.2 Proof of Theorem I Proof. Given Lemma 4 and average over time t = 0 over T -1 and take expectation to all randomness throughout training1 c 2 T T -1 t=0 i,j∈C E x t+1 i -x t+1 j 2 2 ≤ (1 -ηρc) 1 c 2 T T -1 t=0 i,j∈C E x t i -x t j 2 2 + 5nη 2 ρ 2 L 1 T cn T -1 t=0 i∈C n k=1 E |w t+1 ik -w ⋆ ik | 2 ∥x t i -x t k ∥ 2 2 + 6M 2 ij ρc 1 T c 2 T -1 t=0 i,j∈C E fij z t ij -E fij z t+1 ij + 1 c 2 + 3ηLM 2 ij ησ 2 2ρc . Rearrange 1 c 2 T T -1 t=0 i,j∈C E x t+1 i -x t+1 j 2 2 yields 1 c 2 T T -1 t=0 i,j∈C E x t+1 i -x t+1 j 2 2 ≤ 5nηρ Lc 1 T cn T -1 t=0 i∈C n k=1 E |w t ik -w ⋆ ik | 2 ∥x t i -x t k ∥ 2 2 + 6M 2 ij ηρ 2 c 2 1 c 2 T T -1 t=0 i,j∈C E fij z t ij -fij z t+1 ij + 1 c 2 + 3ηLM 2 ij σ 2 2ρ 2 c 2 . Consider bounding |w t ik -w ⋆ ik | 2 in two cases Case 1: w ⋆ ik = 1. Suppose M ik ∈ (0, 1), then ∥∇f i (z t ik ) -∇f k t )∥ 2 2 ≤ M 2 ij ∥∇f i (z t ik ) + ∇f j (z t ik )∥ 2 2 implies ⟨∇f i (z t ik ), ∇f k (z t ik )⟩ ≥ 1 -M 2 ik 2(1 + M 2 ik ) ∥∇f i (z t ik )∥ 2 2 + ∥∇f k (z t ik )∥ 2 2 ≥ 0. then w t+1 ik = w ⋆ ik = 1 and therefore |w t+1 ik -w ⋆ ik | 2 = 0. Case 2: w ⋆ ik = 0. Suppose ζ 2 ik ≥ ∥∇f i (x) + ∇f k (x)∥ 2 2 for all x then ∥∇f i (z t ik ) + ∇f k (z t ik )∥ 2 2 =∥∇f i (z t ik )∥ 2 2 + ∥∇f k (z t ik )∥ 2 2 + 2⟨∇f i (z t ik ), ∇f k (z t ik )⟩ ≥ζ 2 ik + 2⟨∇f i (z t ik ), ∇f k (z t ik )⟩ which means the inner product ⟨∇f i (z t ij ), ∇f j (z t ij )⟩ ≤ 0 is negative, i.e., w t+1 ij = 0 = w ⋆ ij .The above cases hold true for well initialized weights x 0 .Then with lower bound assumption of f i and f j (A3)EQUATIONBy the result of Lemma 3EQUATION2 on both sides of (5), and replace x with z ij we haveEQUATIONThen using the upper bound of M ij < 1/5 from Theorem I, and average over t and i, j yieldsEQUATIONBy Cauchy-Schwarz inequality and L-Lipschitz smoothness, we have that1 c 2 T T -1 t=0 i,j∈C ∇f i x t i 2 2 ≤ 1 c 2 T T -1 t=0 i,j∈C ∇f i z t ij 2 2 + 1 c 2 T T -1 t=0 i,j∈C ∇f i z t ij -∇f i x t i 2 2 ≤ 1 c 2 T T -1 t=0 i,j∈C ∇f i z t ij 2 2 + L 2 4 1 c 2 T T -1 t=0 i,j∈C x t i -x t j 2 2 .Applying ( 9) and ( 11) to the upper bound of the above inequality1 c 2 T T -1 t=0 i,j∈C ∇f i x t i 2 2 ≤ 1 + 1 25 3 Lσ 2 c 2 T i,j∈C fij z 0 ij -f ⋆ ij + L 2 4 6M 2 ij ρ 2 c 2 Lσ 2 c 2 T i,j∈C fij z 0 ij -f ⋆ ij .As ρc ≥ √ 3L and M ij < 1 5 as stated in Theorem I,1 c 2 T T -1 t=0 i,j∈C ∇f i x t i 2 2 ≤4 Lσ 2 c 2 T i,j∈C fij z 0 ij -f ⋆ ij .In Section 4, we present our results on two tasks with different properties. Here, we provide the full details of our experimental setup, alongside with additional experiments.We first describe the setup for cross-device and cross-silo experiments: we use the fix batch size of 128 for cross-device, and cross-silo experiments on CIFAR-100. We tune each method for the optimal learning rate individually: we use learning rate of 0.1 for ditto, 0.05 for Federated Clustering (FC), and 0.01 for all other methods. For Ditto, we use the hyper-parameter of λ = 1 as recommended in their paper. For Federated Clustering, we use the ground truth number of clusters and size of clusters as the hyper-parameter. We also use the ground truth number of clusters for IFCA, and sample all the clients in cross-silo experiment. We reduce the sampling rate to 10% for the cross-device experiment to ensure scalability and fairness for comparison to other methods. For cross-silo experiments we employed a single NVIDIA V-100 GPU with 32GB memory, and moved to four NVIDIA V-100 GPUs with 32 GB memory for cross-device experiment. With this setup, running COBO for cross-silo and cross-device experiment takes 9 hours and 28 hours respectively.For Language modeling experiment, we conducted the experiments with the learning rate of 0.002, batch size of 50, and 4 accumulation steps. Note that each agent only get a subset of the regarding language from Wiki-40B dataset, consisting of total of 800000 tokens. We also used the context length of 512, dropout rate of 0.1, and LoRA module with rank 4. Training is performed on a single NVIDIA A-100 GPU with 40GB memory. It takes 2.5 hours to run COBO for 500 iterations in this framework. We also use λ = 0.1 for Ditto which has higher performance for this experiment. There are 3 strategies proposed in [38] for collaborative fine-tuning, highlighting one as the most effective. However, this approach requires a portion of dataset to be shared among agents, which is incompatible with our experimental setup. Therefore, we use the second-best method in our experiments, which relies on validation performance. For the cross-device experiment with 80 agents in Section 4.3, we present the accuracy curve in Figure 6 . Our method outperform all other methods except the Oracle with a large margin. We can also observe the collaboration matrix of COBO in Figure 5 . The clusters are learned with COBO efficiently.NeurIPS Paper Checklist Justification: We list limitations at the end of the main text.(f 1 + f 2 ). Suppose µ 1 , µ 2 , and µ are the minimizers of f 1 , f 2 , and1 2 (f 1 + f 2 ) respectively. The model weights at points A, B, C demonstrate three scenarios for updating W .https://github.com/karpathy/nanoGPT 


Decision: This paper introduces a new approach to collaborative learning by framing it as a bilevel optimization problem, where the lower-level optimization problem aims to learn the communication topology, and the upper-level problem aims to train machine learning models. The authors provide a convergence analysis for the proposed method, and the experimental results confirm its efficacy.

Strengths:
1. Formulating collaborative learning as a bilevel optimization problem is innovative and allows for learning a problem-adaptive topology.
2. The paper includes a theoretical analysis of the proposed method, and the experimental results confirm its efficacy.

However, there are some minor concerns regarding the convergence analysis. In particular, the theoretical analysis in the original submission assumes the use of the full gradient in the lower-level optimization problem, which is inconsistent with Algorithm 1 and the experiments. During the rebuttal phase, the authors addressed this issue by providing a new analysis based on the stochastic gradient.

Overall, the idea is interesting, and the concerns regarding the convergence analysis have been addressed. Therefore, I suggest that the authors revise the paper as the reviewer suggested and recommend accepting it.

Reviews: review: paper_summary: paper_summary: The paper proposes bi-level training for heterogeneous federated learning, where heterogeneity is due to underlying clustered clients. The two levels of optimizations are model training and determining the client similarity. The authors prove a convergence result and prove empirical results on training vision and LLM models.
strength_weakness: strength_weakness: - For each experiment the setting is fixed, i.e. there are no ablations on how e.g. number of clusters, data samples per client, number of clients per cluster affects the performance. There is a sparsity in experiments.
- FedAvg+fine-tuning should be compared to as an additional method.
- The convergence result on f_ij instead of on f_i, I think this should be fixed since the optimization problem is on f_i's.
- The algorithm requires pairwise computations across clients which might be huge overhead for practical applications where there might be millions of clients. 
- More discussion is needed for the theorem, e.g. what does it shows us, how are different setting parameters like number of clusters, clients per clusters etc. affect the convergence.
questions: questions: - What does "each cluster clients share stationary points" mean? 
- Why is the convergence result is given on f_ij
- What is the performance of method when clustering assumption does not hold, for instance there is only one cluster?
limitations: limitations: Yes
paper_summary: paper_summary: In this paper, the authors model collaborative learning as a bilevel optimization problem, and propose CoBo, an SGD-type alternating optimization algorithm, to solve this problem. Theoretical convergence guarantees are provided, and experiments are conducted to evaluate the performance of CoBo.
strength_weakness: strength_weakness: 1. Although this paper provides theoretical convergence guarantees, the proof of Theorem I is given for a simplified scenario where the inner problem uses the full gradient, which is difficult to compute in practice.
2. This paper does not provide comparisons of the theoretical convergence performance between CoBo and state-of-the-art algorithms.
2. Repeat experiments are not conducted to ensure the robustness of the results.


Minor:
Line 120: "They leads to eventual" is not a full sentence.
questions: questions: 1. In Line 7 of Algorithm 1, why use this output rather than return $\\{ x_0^T,\cdots \cdots ,x_N^T \\}$ and $W^T$?
2. The collaborative learning scenario used in this paper doesn't involve a central server. Why do the authors only compare CoBo with personalized federated learning baselines? Why don't they also compare it with some fully decentralized algorithms?
3. The probability 1/n in Line 11 of Algorithm 1 is an important parameter. How does it influence the performance of CoBo?
limitations: limitations: The authors mention that the limitations (see Weaknesses 1 and 3) are due to time limitations. This suggests that the paper may not be fully ready for submission.
paper_summary: paper_summary: In this paper, the authors introduce a novel approach to collaborative learning by framing it as a bilevel optimization problem, enhancing the training efficacy of multiple clients. In conventional collaborative learning paradigms, clients engage in mutual model training through information exchange; however, pinpointing beneficial collaborators presents a challenge and can incur significant computational overhead. To address this, the authors model client selection and model training as two interrelated optimization issues, proposing the COBO algorithm—a scalable and elastic Stochastic Gradient Descent (SGD) type alternating optimization algorithm. Theoretically guaranteed for convergence, COBO demonstrates superior performance over popular personalization algorithms.
strength_weakness: strength_weakness: According to equation (1), the collaboration weights are updated by applying projected gradient descent with step size . However, an estimate of the step size  is not provided.
questions: questions: 1.Could you provide an estimate of the step size in equation (1)? (see weaknesses)

2.The paper mentions the application of the COBO algorithm in federated learning environments where privacy is a critical consideration. What privacy protection measures have been incorporated into the COBO algorithm to ensure the security of client data, and how is the balance between privacy protection and model performance achieved in the algorithm design?

3.The COBO algorithm and the DITTO algorithm are formulated within the framework of  bilevel optimization. Compared to the DITTO algorithm, what additional theoretical guarantees or advantages does COBO provide? Furthermore, does COBO demonstrate greater flexibility and robustness when dealing with diverse tasks?
limitations: limitations: 1.The proof of the COBO algorithm proposed in the paper is based on some simplified theoretical assumptions, such as the use of full gradient information in the inner problem and the assumption that the minimization problem is solved exactly. These may not fully reflect the complexities encountered in practical applications.

2.In the cross-device federated learning experiment involving 80 clients, each algorithm was executed only once. This implies that the results may not fully reflect the robustness of the algorithms under different random seeds or data partitions.

Comments: title: Response
Addressing weakness 1 (Part 2)
Addressing weakness 1
Response
comment: Thank the authors for the response. My concerns have been addressed, and I am willing to increase my score.
The chances are replaced Line 410. By taking $b\ge\frac{16}{3}\eta^2(c-2)\sigma^2$ and $\rho\ge \frac{\sqrt{3}L}{c}$ and $M_{ij}\le\frac{1}{5}$, the following inequality hold true
$$
\begin{align*}
    \frac{3\eta M^2_{ij}}{\rho c} \left(\frac{3L^2}{4} + 3c^2\rho^2 + \frac{4\eta^2\rho^2(c-2)\sigma^2}{b} \right) 
    \le \frac{1}{2} \eta\rho c.
\end{align*}
$$
Notice that the upper remain the same. Next is the full theorem where the only change is that we additional require the $b$ is not too small.

-----

Suppose Assumption 1,2,3,4,5 hold true. Suppose that Algorithm 1 solves (1) exactly.
Consider clients $i$ and $j$ in the same cluster $C$ of size $c$. Suppose  that $M _ {ij}^2\in(0,\frac{1}{5})$ and $\zeta^2 _ {ik}\ge \lVert \nabla f _ i(x) + \nabla f _ k(x) \rVert _ 2^2$ for all $x$ and $k$.
Let $\rho\ge \frac{\sqrt{3}L}{c}$ and  $\mathbf{b\ge\frac{16}{3}\eta^2(c-2) \sigma^2}$  and step size  is
$$
\begin{align*}
    \eta \le \min \left\\{ \frac{2}{\sigma\sqrt{LT}} \sqrt{ \frac{1}{c^2} \sum _ {i,j\in C}  \left( \tilde{f} _ {ij}\left(z _ {ij}^{0}\right) - \tilde{f} _ {ij}^\star  \right)}, \frac{1}{2\sqrt{3}L}  \right\\}.
\end{align*}
$$
The consensus distance also converges to 0, i.e.
$$
\begin{align*}
    \frac{1}{c^2T}\sum _ {t=0}^{T-1}\sum _ {i,j\in C}  \mathbb{E}\left[\lVert x _ i^{t+1} - x _ j^{t+1}\rVert _ 2^2 \right]
    \le& \frac{6M^2 _ {ij}}{\rho^2 c} 
        \sqrt{ \frac{L\sigma^2}{c^2T} \sum _ {i,j\in C} \left( \tilde{f} _ {ij}\left(z _ {ij}^{0}\right) - \tilde{f} _ {ij}^\star  \right)}.
\end{align*}
$$
Moreover, the gradient norm is upper bounded.
$$
\begin{align*}
    \frac{1}{c^2 T}\sum _ {t=0}^{T-1} \sum _ {i,j\in C} \mathbb{E}\left[\lVert\nabla \tilde{f} _ {ij}\left(z _ {ij}^{t}\right)  \rVert _ 2^2\right]
    \le& 3 \sqrt{ \frac{L\sigma^2}{c^2T} \sum _ {i,j\in C} \left( \tilde{f} _ {ij}\left(z _ {ij}^{0}\right) - \tilde{f} _ {ij}^\star  \right)}.
\end{align*}
$$
-----
Regarding the impact of stochastic gradient:  The theoretical results above suggest that as long as $b$ is not too small, that is $\Omega(\eta^2 c)$, the convergence results remain the same. The intuition is that these stochasticity gradients are only used to compute the inner product that approximates the direction $w_{ij}$ should head for. The perturbation may well not change the sign of the inner product. Besides, as $w_{ij}$ is bounded between [0,1], the inner produce of stochastic gradients can be seen as clipped, making it more robust to stochasticity. Throughout the experiments, we indeed use the stochastic gradients variant.

----
Regarding the comparison with other stochastic gradient based algorithm: There is no perfect baseline for comparison as our formulation is designed for client-selection and as far as we know different from other works.  But we can refer the rebuttal we posted to Reviewer xPeE that "We provide a comparison of the theoretical convergence performance between CoBo and state-of-the-art algorithms in our response. Specifically, we show that CoBo enjoys linear scalability with respect to cluster size, whereas federated clustering requires O(n^2) gradient computations per iteration. Ditto gives a coarse grained O(1/T) convergence rate for local models."
Thank the authors for providing more proof. Could the authors provide the final theorem and compare it with Theorem 1 to show the impact of the stochastic gradient and also compare it with stochastic gradient based state-of-the-art algorithms?
Average over i, j yields
$$
\begin{align*}
\frac{1}{c^2}\sum _ {ij}\mathbb{E} _ h\mathbb{E} _ g \left[\lVert z _ {ij}^{t+1} - z _ {ij}^{t} - \mathbb{E} _ h[z _ {ij}^{t+1} - z _ {ij}^{t}]  \rVert _ 2^2 \right] 
=\frac{\eta^2\rho^2(c-2)}{2c^2} \sum _ {i,j} \mathbb{E} _ h \left[ \left\lVert w _ {ij}^{t+1} - \mathbb{E} _ h[w _ {ij}^{t+1}] \right\rVert^2 _ 2 \right]  \lVert x _ i^t - x _ j^t\rVert _ 2^2  \\
\le&\frac{\eta^2\rho^2(c-2)}{2c^2} \frac{4\sigma^2}{b} \sum _ {i,j}  \lVert x _ i^t - x _ j^t\rVert _ 2^2.
\end{align*}
$$
Then
$$
    \begin{align*}
        \frac{1}{c^2}\sum _ {i,j\in c}\left\lVert\nabla \tilde{f} _ {ij}\left(z _ {ij}^{t}\right)\right\rVert _ 2^2
        \le \frac{2}{\eta c^2} \sum _ {i,j\in C}\left( \tilde{f} _ {ij}\left(z _ {ij}^{t}\right) - \tilde{f} _ {ij}\left(z _ {ij}^{t+1}\right)  \right)
        +  \left(\frac{3L^2}{4c^2} + 3\rho^2 + \frac{4\eta^2\rho^2(c-2)\sigma^2}{b}\right) \sum _ {i,j\in c} \lVert x _ i^t-x _ j^t\rVert _ 2^2 
         +  \frac{3n\rho^2}{c} \sum _ {i\in c} \sum _ {k=1}^n |w _ {ik}^{t+1} - w _ {ik}^\star|^2 \lVert x _ i^t - x _ k^t\rVert _ 2^2 + \frac{L\eta \sigma^2}{2}.
    \end{align*}
$$
Let $\eta \le \sqrt{\frac{b}{(c-2)\sigma^2}}$ we have that $\frac{4\eta^2\rho^2(c-2)\sigma^2}{b} \le 4 \rho^2$. The rest of the proof remain the same except for constants.

Due to character limit and the difficulty of editing long equations in markdown here, we show stop the proof here and note that the rest of the proof remain the same up to constants.
Our changes mainly applies to the proof of Lemma 3, starting from line 390 in the original submission. The current proof only consider the randomness of stochastic gradient $g_1^t, \ldots, g_n^t$ for updating model $x_i$. 

Now we additionally consider using stochastic gradients with batch size $b$ for the inner product. That is, let $h_i^t$ and $h_j^t$ be independent and unbiased estimates of 
$\langle\nabla f_i(z_{ij}^t), \nabla f_j(z_{ij}^t)\rangle$. The variance of ${h}_i^t$ has a variance of $\frac{\sigma^2}{b}$.

Let's denote $\mathbb{E} _ g:=\mathbb{E} _ {g _ 1,\ldots, g _ n}$ and $\mathbb{E} _ h:=\mathbb{E} _ {h _ 1,\ldots, h _ n}$ and replace the expecation $\mathbb{E}$ in the proof with $\mathbb{E} _ h\mathbb{E} _ g$. Notice that the only difference is that 

\begin{align*}
    \mathbb{E} _ h\mathbb{E} _ g\left[\tilde{f} _ {ij}\left(z _ {ij}^{t+1}\right)\right]
    \le& \tilde{f} _ {ij}\left(z _ {ij}^{t}\right) + 
    \left\langle \nabla \tilde{f} _ {ij}\left(z _ {ij}^{t}\right), \mathbb{E} _ h\mathbb{E} _ g\left[z _ {ij}^{t+1} - z _ {ij}^{t}\right] \right\rangle + \frac{L}{2} \lVert\mathbb{E} _ h\mathbb{E} _ g[z _ {ij}^{t+1} - z _ {ij}^{t}]\rVert _ 2^2 
     + \frac{L}{2} \mathbb{E} _ h\mathbb{E} _ g\left[\lVert z _ {ij}^{t+1} - z _ {ij}^{t} - \mathbb{E} _ h\mathbb{E} _ g[z _ {ij}^{t+1} - z _ {ij}^{t}] \rVert _ 2^2 \right].
\end{align*}

The last quantity can be expanded as follows:
\begin{align*}
 \mathbb{E} _ h\mathbb{E} _ g \left[\lVert z _ {ij}^{t+1} - z _ {ij}^{t} - \mathbb{E} _ h\mathbb{E} _ g[z _ {ij}^{t+1} - z _ {ij}^{t}] \rVert _ 2^2 \right] 
= \mathbb{E} _ h\mathbb{E} _ g \left[\lVert z _ {ij}^{t+1} - z _ {ij}^{t} \pm \mathbb{E} _ h[z _ {ij}^{t+1} - z _ {ij}^{t}] - \mathbb{E} _ h\mathbb{E} _ g[z _ {ij}^{t+1} - z _ {ij}^{t}] \rVert _ 2^2 \right] 
=  \mathbb{E} _ h\mathbb{E} _ g \left[\lVert \mathbb{E} _ h[z _ {ij}^{t+1} - z _ {ij}^{t}] - \mathbb{E} _ h\mathbb{E} _ g[z _ {ij}^{t+1} - z _ {ij}^{t}] \rVert _ 2^2 \right]+ \mathbb{E} _ h\mathbb{E} _ g \left[\lVert z _ {ij}^{t+1} - z _ {ij}^{t} - \mathbb{E} _ h[z _ {ij}^{t+1} - z _ {ij}^{t}]  \rVert _ 2^2 \right].
\end{align*}
Plug in the above equality to the above inequality
$$
\begin{align*}
    \mathbb{E} _ h\mathbb{E} _ g\left[\tilde{f} _ {ij}\left(z _ {ij}^{t+1}\right)\right]
    \le& \tilde{f} _ {ij}\left(z _ {ij}^{t}\right) + 
    \left\langle \nabla \tilde{f} _ {ij}\left(z _ {ij}^{t}\right), \mathbb{E} _ h\mathbb{E} _ g\left[z _ {ij}^{t+1} - z _ {ij}^{t}\right] \right\rangle + \frac{L}{2} \lVert\mathbb{E} _ h\mathbb{E} _ g[z _ {ij}^{t+1} - z _ {ij}^{t}]\rVert _ 2^2 
    + \mathbb{E} _ h\mathbb{E} _ g \left[\lVert \mathbb{E} _ h[z _ {ij}^{t+1} - z _ {ij}^{t}] - \mathbb{E} _ h\mathbb{E} _ g[z _ {ij}^{t+1} - z _ {ij}^{t}] \rVert _ 2^2 \right] 
    + \mathbb{E} _ h\mathbb{E} _ g \left[\lVert z _ {ij}^{t+1} - z _ {ij}^{t} - \mathbb{E} _ h[z _ {ij}^{t+1} - z _ {ij}^{t}]  \rVert _ 2^2 \right].
\end{align*}
$$
Among the 5 terms on the right hand side of the above inequality, the first 4 terms is bounded the same way as the original submission. The effect of using stochastic gradient for inner product is limited to 
\begin{align*}
\mathbb{E} _ h\mathbb{E} _ g \left[\lVert z _ {ij}^{t+1} - z _ {ij}^{t} - \mathbb{E} _ h[z _ {ij}^{t+1} - z _ {ij}^{t}]  \rVert _ 2^2 \right] \\
=\mathbb{E} _ h\mathbb{E} _ g \left[ \left\lVert \frac{\eta}{2} (g _ i^t + g _ j^t) + \frac{\eta\rho}{2}\sum _ {k=1}^n (w _ {ik}^{t+1} (x _ i^t - x _ k^t) + w _ {jk}^{t+1} (x _ j^t - x _ k^t)) \right.\right. 
 \qquad\left.\left. - \mathbb{E} _ h \left[\frac{\eta}{2} (g _ i^t + g _ j^t) + \frac{\eta\rho}{2}\sum _ {k=1}^n (w _ {ik}^{t+1} (x _ i^t - x _ k^t) + w _ {jk}^{t+1} (x _ j^t - x _ k^t)) \right]  \right\rVert _ 2^2 \right] 
=\mathbb{E} _ h\left[ \left\lVert \frac{\eta\rho}{2}\sum _ {k=1}^n ((w _ {ik}^{t+1} - \mathbb{E} _ h[w _ {ik}^{t+1}]) (x _ i^t - x _ k^t) + (w _ {jk}^{t+1}- \mathbb{E} _ h[w _ {jk}^{t+1}]) (x _ j^t - x _ k^t))  \right\rVert _ 2^2 \right] 
\end{align*}
Then
\begin{align*}
\mathbb{E} _ h\mathbb{E} _ g \left[\lVert z _ {ij}^{t+1} - z _ {ij}^{t} - \mathbb{E} _ h[z _ {ij}^{t+1} - z _ {ij}^{t}]  \rVert _ 2^2 \right] \\
=\frac{\eta^2\rho^2}{4} \sum _ {k\neq i,j} \left(\mathbb{E} _ h \left[ \left\lVert w _ {ik}^{t+1} - \mathbb{E} _ h[w _ {ik}^{t+1}] \right\rVert^2 _ 2 \right]  \lVert x _ i^t - x _ k^t\rVert _ 2^2 + \mathbb{E} _ h \left[ \left\lVert w _ {jk}^{t+1} - \mathbb{E} _ h[w _ {jk}^{t+1}] \right\rVert^2 _ 2 \right]  \lVert x _ j^t - x _ k^t\rVert _ 2^2 \right)
\end{align*}
where we use the independence of random variables in the last equality.
We are glad to have addressed your concern about FedAvg+Fine-tuning. Regarding the other points, we would like to clarify that:
- CoBo, as shown in Algorithm 1, **only computes O(n) gradients per iteration**, instead of pairwise O(n^2), making it very efficient in practice. In a setup with n clients, if we sample the pairs of clients with probability O(1/n), the expected number of selected pairs in each timestep is n, which is the same as the number of local computed gradients in all baselines. Hence, the complexity of our algorithm is the same as other collaborative learning algorithms. As is shown in the experiments sections of the submission pdf and rebuttal pdf, CoBo is able to reach **higher accuracies** compared to other methods **in all experiments** with **same order of complexity**. 

- As for point 3, we are not sure why the reviewer says that “we have not utilized markdown” as the equations render correctly to us. Could you try again or try other browsers?

- As for theoretical results, the technical challenges of the convergence proof are already addressed in the no-sampling theorem while the sampling variant is a simple corollary. We will add the simple corollary to the main text. 

We hope that our responses clear the reviewer’s main concern about complexity.
Thank the authors for the rebuttal. Since Weakness 1 has not yet been addressed, I would like to keep my current score.
Thanks for the rebuttal, comparison to FedAvg+ft is a good addition. Point 3. is very hard to follow as you have not utilized markdown. Despite sampling, I still think the pairwise computations are an important overhead in practice; also the theoretical results does not take into account such a sampling occurs. Hence, I would like to keep my score.