Title: CoBo: Collaborative Learning via Bilevel Optimization
Venue: NeurIPS 2024
Accepted: yes
Keywords: collaborative learning

personalized federated learning

bilevel optimization

distributed learning
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.

COBO: Collaborative Learning via Bilevel Optimization

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

1 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).

2 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,

(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 , and 1 2 ( f 1 + f 2 ) . Suppose µ 1 , µ 2 , and µ are the minimizers of f 1 , f 2 , and 1 2 ( f 1 + f 2 ) respectively. The model weights at points A,B,C demonstrate three scenarios for updating W .

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

1

2

(

f

1+

f

2)

Averaged model weight

Grads for X-update

f

1(

x

1),

f

2(

x

2)

Grads for W-update

f

1

(

x

1+

x

2

2

)

,

f

2

(

x

1+

x

2

2

)

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.

2.1 Algorithm

Wepropose 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):

Next, 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):

where η > 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 ( 1 n ) , leading to an expected computation of

Algorithm 1 COBO: Co llaborative Learning via B ilevel O ptimization

Input: Model parameters ∀ 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 do

2: 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 update

̸

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.

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

(A2) Noise bound. For all x ∈ R d and i ∈ [ n ] , there exists σ 2 > 0 such that the stochastic gradient has bounded noise

(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 that

Otherwise, there exists ζ 2 ij > 0 such that

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 is f i ( x ) = a i 2 ∥ x -µ k ∥ 2 2 where a i > 0 . Then for clients i, j in the same cluster, i.e. i, j ∈ C k

̸

The M ij = | a i -a j | a i + a j in this case. On the other hand, for i ∈ C k and j ∈ C k ′ and µ k = µ k ′ ,

where 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 that M 2 ij ∈ (0 , 1 5 ) , b ≥ 2 c 2 2 Lη ( c -2) σ 2 and ζ 2 ik ≥ ∥∇ f i ( x ) + ∇ f k ( x ) ∥ 2 2 for all x and k . Let ρ ≥ √ 3 L c and step size

The consensus distance also converges to 0, i.e.

Moreover, the gradient norm is upper bounded.

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 rate

4 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).

4.1 Cross-silo federated learning experiment with 8 clients

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.

Acc.(%)

Loss

Constant (

O

(1

/n

)

)

73.05

1.104

Time-dependent (

O

(1

/t

)

)

73.18

1.226

Mixed

74.77

1.081

No Sampling (Oracle)

74.93

1.278

0.1

0.3

0.5

Fraction of dataset

60.0

67.5

75.0

3

4

5

Clusters

72.0

73.0

74.0

2

3

4

Clients per cluster

73.0

73.5

74.0

(a)

0

100

200

300

400

Iterations (×100)

0

10

20

30

40

50

60

70

CoBo

Finetuned FedAvg

Start of fine-tuning

Oracle

Ditto

FC

IFCA

Local

FedAvg

(b)

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

Iteration

300

Iteration

5000

Iteration

15000

Iteration

25000

Iteration

40000

FC

IFCA

CoBo

0.0

0.2

0.4

0.6

0.8

1.0

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.

4.2 Cross-device experiment experiment with 80 clients

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.

4.3 Collaborative fine-tuning on language models

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 model 3 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 FEDAVG 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.

0

25

50

75

100

Iterations

0.0

0.2

0.4

0.6

0.8

Domain Weights

Spanish

German

Dutch

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

Federated learning with client selection 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 i -x 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.

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

Acknowledgement. We acknowledge funding from Swiss National Science Foundation (SNSF) grant number 200020_200342, from Huawei Cloud Intelligent Cloud Technologies Initiative, and from Google Research Collaborations.

References

Yossi Arjevani, Yair Carmon, John C Duchi, Dylan J Foster, Nathan Srebro, and Blake Woodworth. Lower bounds for non-convex stochastic optimization. Mathematical Programming , 199 (1):165-214, 2023.

James C Bezdek and Richard J Hathaway. Convergence of alternating optimization. Neural, Parallel & Scientific Computations , 11(4):351-368, 2003.

Aymeric Capitaine, Etienne Boursier, Antoine Scheid, Eric Moulines, Michael I Jordan, ElMahdi El-Mhamdi, and Alain Durmus. Unravelling in collaborative learning. arXiv preprint arXiv:2407.14332 , 2024.

Tianyi Chen, Yuejiao Sun, and Wotao Yin. Closing the gap: Tighter analysis of alternating stochastic gradient methods for bilevel problems. Advances in Neural Information Processing Systems , 34:25294-25307, 2021.

Bo Dai, Niao He, Yunpeng Pan, Byron Boots, and Le Song. Learning from conditional distributions via dual embeddings. In Aarti Singh and Xiaojin (Jerry) Zhu, editors, Proceedings of the 20th International Conference on Artificial Intelligence and Statistics, AISTATS 2017, 20-22 April 2017, Fort Lauderdale, FL, USA , volume 54 of Proceedings of Machine Learning Research , pages 1458-1467. PMLR, 2017. URL http://proceedings.mlr.press/v54/ dai17a.html .

Bo Dai, Albert Shaw, Lihong Li, Lin Xiao, Niao He, Zhen Liu, Jianshu Chen, and Le Song. SBEED: convergent reinforcement learning with nonlinear function approximation. In Jennifer G. Dy and Andreas Krause, editors, Proceedings of the 35th International Conference on Machine Learning, ICML 2018, Stockholmsmässan, Stockholm, Sweden, July 10-15, 2018 , volume 80 of Proceedings of Machine Learning Research , pages 1133-1142. PMLR, 2018. URL http://proceedings.mlr.press/v80/dai18c.html .

Mathieu Even, Laurent Massoulié, and Kevin Scaman. On sample optimality in personalized collaborative and federated learning. Advances in Neural Information Processing Systems , 35: 212-225, 2022.

Chelsea Finn, Pieter Abbeel, and Sergey Levine. Model-agnostic meta-learning for fast adaptation of deep networks. In Doina Precup and Yee Whye Teh, editors, Proceedings of the 34th International Conference on Machine Learning, ICML 2017, Sydney, NSW, Australia, 6-11 August 2017 , volume 70 of Proceedings of Machine Learning Research , pages 1126-1135. PMLR, 2017. URL http://proceedings.mlr.press/v70/finn17a.html .

Gersende Fort. Federated expectation maximization with heterogeneity mitigation and variance reduction.

Marguerite Frank, Philip Wolfe, et al. An algorithm for quadratic programming. Naval research logistics quarterly , 3(1-2):95-110, 1956.

Avishek Ghosh, Jichan Chung, Dong Yin, and Kannan Ramchandran. An efficient framework for clustered federated learning. Advances in Neural Information Processing Systems , 33: 19586-19597, 2020.

Mandy Guo, Zihang Dai, Denny Vrandeˇ ci´ c, and Rami Al-Rfou. Wiki-40B: Multilingual language model dataset. In LREC - Proceedings of the Twelfth Language Resources and Evaluation Conference , pages 2440-2452. European Language Resources Association, May 2020. ISBN 979-10-95546-34-4. URL https://aclanthology.org/2020.lrec-1.297 .

Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition, 2015.

Edward J Hu, yelong shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang, Lu Wang, and Weizhu Chen. LoRA: Low-rank adaptation of large language models. In International Conference on Learning Representations , 2022. URL https://openreview. net/forum?id=nZeVKeeFYf9 .

Yifan Hu, Siqi Zhang, Xin Chen, and Niao He. Biased stochastic first-order methods for conditional stochastic optimization and applications in meta learning. In Hugo Larochelle, Marc'Aurelio Ranzato, Raia Hadsell, Maria-Florina Balcan, and Hsuan-Tien

Lin, editors, Advances in Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Systems 2020, NeurIPS 2020, December 612, 2020, virtual , 2020. URL https://proceedings.neurips.cc/paper/2020/hash/ 1cdf14d1e3699d61d237cf76ce1c2dca-Abstract.html .

Feihu Huang, Junyi Li, and Shangqian Gao. Biadam: Fast adaptive bilevel optimization methods. arXiv preprint arXiv:2106.11396 , 2021.

Martin Jaggi. Revisiting frank-wolfe: Projection-free sparse convex optimization. In International conference on machine learning , pages 427-435. PMLR, 2013.

Peter Kairouz, H Brendan McMahan, Brendan Avent, Aurélien Bellet, Mehdi Bennis, Arjun Nitin Bhagoji, Kallista Bonawitz, Zachary Charles, Graham Cormode, Rachel Cummings, et al. Advances and open problems in federated learning. Foundations and trends® in machine learning , 14(1-2):1-210, 2021.

Anastasia Koloskova, Nicolas Loizou, Sadra Boreiri, Martin Jaggi, and Sebastian U. Stich. A unified theory of decentralized SGD with changing topology and local updates. In Proceedings of the 37th International Conference on Machine Learning, ICML 2020, 13-18 July 2020, Virtual Event , volume 119 of Proceedings of Machine Learning Research , pages 5381-5393. PMLR, 2020. URL http://proceedings.mlr.press/v119/koloskova20a.html .

Guy Kornowski, Swati Padmanabhan, Kai Wang, Zhe Zhang, and Suvrit Sra. First-order methods for linearly constrained bilevel optimization. arXiv preprint arXiv:2406.12771 , 2024.

Alex Krizhevsky, Geoffrey Hinton, et al. Learning multiple layers of features from tiny images. 2009.

Viraj Kulkarni, Milind Kulkarni, and Aniruddha Pant. Survey of personalization techniques for federated learning. In 2020 Fourth World Conference on Smart Trends in Systems, Security and Sustainability (WorldS4) , pages 794-797. IEEE, 2020.

Batiste Le Bars, Aurélien Bellet, Marc Tommasi, Erick Lavoie, and Anne-Marie Kermarrec. Refined convergence and topology learning for decentralized sgd with heterogeneous data. In International Conference on Artificial Intelligence and Statistics , pages 1672-1702. PMLR, 2023.

T. Li, Shengyuan Hu, Ahmad Beirami, and Virginia Smith. Fair and robust federated learning through personalization. In 38th International Conference on Machine Learning , 2021.

Xiangru Lian, Ce Zhang, Huan Zhang, Cho-Jui Hsieh, Wei Zhang, and Ji Liu. Can decentralized algorithms outperform centralized algorithms? A case study for decentralized parallel stochastic gradient descent. In Isabelle Guyon, Ulrike von Luxburg, Samy Bengio, Hanna M. Wallach, Rob Fergus, S. V. N. Vishwanathan, and Roman Garnett, editors, Advances in Neural Information Processing Systems 30: Annual Conference on Neural Information Processing Systems 2017, December 4-9, 2017, Long Beach, CA, USA , pages 5330-5340, 2017. URL https://proceedings.neurips.cc/paper/2017/hash/ f75526659f31040afeb61cb7133e4e6d-Abstract.html .

Yishay Mansour, Mehryar Mohri, Jae Ro, and Ananda Theertha Suresh. Three approaches for personalization with applications to federated learning. 2021.

Othmane Marfoq, Giovanni Neglia, Laetitia Kameni, and Richard Vidal. Federated multi-task learning under a mixture of distributions. In Proceedings of the 35th International Conference on Machine Learning , volume 34, 2021.

Brendan McMahan, Eider Moore, Daniel Ramage, Seth Hampson, and Blaise Aguera y Arcas. Communication-efficient learning of deep networks from decentralized data. In Artificial intelligence and statistics , pages 1273-1282. PMLR, 2017.

Mehryar Mohri, Gary Sivek, and Ananda Theertha Suresh. Agnostic federated learning. In International Conference on Machine Learning , pages 4615-4625. PMLR, 2019.

Ofir Nachum and Bo Dai. Reinforcement learning via fenchel-rockafellar duality. ArXiv preprint , abs/2001.01866, 2020. URL https://arxiv.org/abs/2001.01866 .

Ieva Petrulionyte, Julien Mairal, and Michael Arbel. Functional bilevel optimization for machine learning. arXiv preprint arXiv:2403.20233 , 2024.

Alec Radford, Jeff Wu, Rewon Child, David Luan, Dario Amodei, and Ilya Sutskever. Language models are unsupervised multitask learners. 2019.

Felix Sattler, Klaus-Robert Muller, and Wojciech Samek. Clustered federated learning: Modelagnostic distributed multi-task optimization under privacy constraints. 2019.

Virginia Smith, Chao-Kai Chiang, Maziar Sanjabi, and Ameet S Talwalkar. Federated multi-task learning. Advances in neural information processing systems , 30, 2017.

Canh T Dinh, Nguyen Tran, and Josh Nguyen. Personalized federated learning with moreau envelopes. Advances in Neural Information Processing Systems , 33:21394-21405, 2020.

Alysa Ziying Tan, Han Yu, Lizhen Cui, and Qiang Yang. Towards personalized federated learning. IEEE Transactions on Neural Networks and Learning Systems , 2022.

Vinzenz Thoma, Barna Pásztor, Andreas Krause, Giorgia Ramponi, and Yifan Hu. Bilevel optimization with lower-level contextual mdps. In Agentic Markets Workshop at ICML 2024 .

Nicolas Wagner, Dongyang Fan, and Martin Jaggi. Personalized collaborative fine-tuning for on-device large language models. In First Conference on Language Modeling , 2024. URL https://openreview.net/forum?id=bwo3GVsgOv .

Mariel Werner, Lie He, Sai Praneeth Karimireddy, Michael Jordan, and Martin Jaggi. Provably personalized and robust federated learning. arXiv preprint arXiv:2306.08393 , 2023.

Valentina Zantedeschi, Aurélien Bellet, and Marc Tommasi. Fully decentralized joint learning of personalized models and collaboration graphs. In International Conference on Artificial Intelligence and Statistics , pages 864-874. PMLR, 2020.

A Theory

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 ˜ f ij ( x ) := 1 2 ( f i ( x ) + f j ( x )) be an averaged objective function.

Lemma 3. Suppose (A1) holds true. Let η ≤ 1 2 L . Then for i, j in the same cluster C of size c

where coefficient D 1 is defined as

A.1 Proof of Lemma 3

Proof. Let h t i and h t j be independent and unbiased estimates of ∇ f i ( z t ij ) , ∇ f j ( z t ij ) respectively. The variance of h t i has a variance of σ 2 b . Let's denote E g [ · ] := E g 1 ,...,g n [ ·| z t i ] and E h [ · ] := E h 1 ,...,h n [ · ] and let E [ · ] = E h [ E g [ · ]] . By the L-smoothness assumption (A1) and bounded noise assumption (A2)

The last quantity, i.e. variance of z t +1 ij -z t ij , can be expanded as follows

where 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. Then

where η ≤ 1 2 L is applied in the last inequality. Then the upper bound of ∥ ∥ ∥ ∇ ˜ f ij ( z t ij ) ∥ ∥ ∥ 2 2 is that

The T can be upper bounded by expanding z t +1 i with (2)

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 )

Bound T 2 : Use Cauchy-Schwarz inequality and C has a cluster size of c

Bound T 3 : Use Cauchy-Schwarz inequality

Sum up ∥ ∥ ∥ ∇ ˜ f ij ( z t ij ) ∥ ∥ ∥ 2 2 for all of i, j in the same cluster C for (7) yields

Now we expand T ij as follows

̸

where we use the independence of random variables in the last equality. Average the above equality over i, j ∈ C yields

Therefore, (8) is upper bounded with

Proof. Expand x t +1 i -x t +1 j with (2)

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 )

Compute the norm of x t +1 i -x t +1 j and choose ηρ ≤ 1 2 c to use Jensen's inequality

Expand the right-hand side with Cauchy-Schwarz inequality

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)

Note that E [ ∥ 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 2 and independence between randomness on worker i and j where expectation is take to all of the randomness until time t

By averaging for all i, j ∈ C

Use the previous Lemma 3 to bound ∑ i,j ∈C ∥ ∥ ∥ ∇ ˜ f ij ( z t ij ) ∥ ∥ ∥ 2 2

Rearrange the terms

We would like to achieve ( 1 -2 ηρc + 3 ηL 2 2 ρc + 3 ηM 2 ij ρc D 1 ) ≤ (1 -ηρc ) by

letting ρ ≥ √ 3 L c , s.t. 3 ηL 2 2 ρc ≤ ηρc 2 .

Using the same requirement that ρ ≥ √ 3 L c and M ij ≤ 1 5

The upper bound of

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 training

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 ( z t ik ) ∥ 2 2 ≤ M 2 ij ∥∇ f i ( z t ik ) + ∇ f j ( z t ik ) ∥ 2 2 implies

then w t +1 ik = w ⋆ ik = 1 and therefore | w t +1 ik -w ⋆ ik | 2 = 0 .

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)

Minimize the upper bound through choosing η

By the result of Lemma 3

A.3 Proof of Corollary II

Proof. By adding ∥∇ f i ( x ) + ∇ f j ( x ) ∥ 2 2 on both sides of (5), and replace x with z ij we have

Then using the upper bound of M ij < 1 / 5 from Theorem I, and average over t and i, j yields

By Cauchy-Schwarz inequality and L -Lipschitz smoothness, we have that

Applying (9) and (11) to the upper bound of the above inequality

√

B Experimental Details

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.

Iteration

300

Iteration

5000

Iteration

15000

Iteration

20000

CoBo

0.0

0.5

1.0

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

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 theoretical results are given in Section 3 and empirical evaluations are presented in Section 4.

2. Limitations

0

25

50

75

100

125

150

175

200

Iterations (×100)

10

20

30

40

50

60

70

80

CoBo

Oracle

Ditto

Finetuned

FedAvg

Local

FedAvg

IFCA

Start of

fine-tuning

Question: Does the paper discuss the limitations of the work performed by the authors?

Answer: [Yes]

Justification: We list limitations at the end of the main text.

3. Theory Assumptions and Proofs

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.

4. Experimental Result Reproducibility

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.

5. Open access to data and code

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.

6. Experimental Setting/Details

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.

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: We report proper confidence intervals for all experiments in our work.

8. Experiments Compute Resources

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

Answer: [Yes]

Justification: We explained the training framework thoroughly in appendix. All the mentioned information are discussed and we believe it is sufficient for reproducing the experiments.

9. Code Of Ethics

Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines ?

Answer: [Yes]

Justification: We preserve the NeurIPS Code of Ethics throughout the research project.

10. Broader Impacts

Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?

Answer: [NA]

Justification: Not applicable.

11. Safeguards

Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?

Answer: [NA]

Justification: Note applicable.

12. Licenses for existing assets

Question: 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?

Answer: [Yes]

Justification: We have cited the datasets.

13. New Assets

Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?

Answer: [NA]

Justification: Not applicable

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 to participants and screenshots, if applicable, as well as details about compensation (if any)?

Answer: [NA]

Justification: This paper does not involve crowdsourcing experiments.

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: Not applicable.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nIn 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.\n\nSTRENGTHS:\n1.The authors formulate collaborative learning with an innovative bilevel optimization framework that capitalizes on the intrinsic collaborative structure, yielding more universally applicable solutions.\n\n2.The authors introduce COBO, a scalable and elastic SGD-type alternating optimization algorithm that efficiently addresses the bilevel problem challenge, scaling with an increasing number of clients while maintaining elasticity in client count.\n\n3.The authors establish theoretical convergence guarantees for COBO in collaborative learning scenarios with clustered client structures.\n\n4.The authors demonstrate that COBO surpasses established personalized federated learning benchmarks, especially in heterogeneous federated learning environments and when fine-tuning Large Language Models (LLMs).\n\nWEAKNESSES:\nAccording 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.\n\nQUESTIONS:\n1.Could you provide an estimate of the step size in equation (1)? (see weaknesses)\n\n2.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?\n\n3.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?\n\nLIMITATIONS:\n1.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.\n\n2.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.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-12"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nIn 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.\n\nSTRENGTHS:\nIn this paper, the authors model client-selection and model-training as two interconnected optimization problems, proposing a bilevel optimization problem for collaborative learning. They introduce CoBo, an SGD-type alternating optimization algorithm designed to address this problem. CoBo is proven to have theoretical convergence guarantees for collaborative learning with cluster structures. In the experiments, CoBo outperforms popular personalized federated learning baselines.\n\nWEAKNESSES:\n1. 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.\n2. This paper does not provide comparisons of the theoretical convergence performance between CoBo and state-of-the-art algorithms.\n2. Repeat experiments are not conducted to ensure the robustness of the results.\n\n\nMinor:\nLine 120: \"They leads to eventual\" is not a full sentence.\n\nQUESTIONS:\n1. In Line 7 of Algorithm 1, why use this output rather than return $\\\\{ x_0^T,\\cdots \\cdots ,x_N^T \\\\}$ and $W^T$?\n2. 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?\n3. The probability 1/n in Line 11 of Algorithm 1 is an important parameter. How does it influence the performance of CoBo?\n\nLIMITATIONS:\nThe 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.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 2, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2024-07-13"}

comments: 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.

STRENGTHS:
- Paper is very easy to follow and understand due to the flow of language. 
- The presented results indicate the advantages compared to other works. 
- The high level idea is sound.

WEAKNESSES:
- 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:
- 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:
Yes
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 2
DATE: 2024-07-13