CoBo: Collaborative Learning via Bilevel Optimization
Authors: Diba Hashemi
Lie He
Martin Jaggi
Venue: NeurIPS 2024 poster (2024)
Keywords: collaborative learning
personalized federated learning
bilevel optimization
distributed learning
Primary area: optimization
TL;DR: Collaborative learning through solving a bilevel optimization problem.

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
Diba Hashemi
Lie He
Martin Jaggi EPFL martin. jargin.jagg@@ppll... ch
EPFL diba. hwwhahsshhmh.ccc.mmmmmmaa....hh @epfl.cl
Tencent Inc. 1i hew5@gmail...mmmmmmmmmmmmmmmammm com
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 interconnected optimization problems, proposing a novel bilevel optimization problem for collaborative learning. We introduce CoBo, 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 task with high heterogeneity, involving datasets distributed among 80 clients.?
1 Introduction
In classic collaborative learning scenario, 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 E [n) trains a dddiimensional model æ; Rd to minimize its loss function, fi IRd _, IR. In the subsequent communication step, client exchanges updates with clients, potentially benefiting from collaboration.
Despite the plethora of collaborative learning frameworks, the ideal approach collaborate remains under-exploited. The FEDAVG [28, 18] algorithm learns single global model over pooled datasets from all clients, i.e., min Σx= f((x). However, due to heterogeneous data distributions among clients, 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 global model [24]. Although DITTO enables personal models leverage the global model, it offers only coarse-graines level of collaboration. In instances where clients" data exhibit significant differences, the DITTO algorithm constrained to facilitating collaboration at global level, thereby neglecting the inherent heterogeneity among clients.
Clustering-bant based federated learning algorithms have been developed 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 predetermine the number of clusters, initialize cluster centers, and other such prerequisites, which can diminish their practical utility.
*Corresponding author.
"The code is available at: https: h/tittuu//ttubb.... cttpaa///com//me//eeeeeee
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

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 wij (0, 1), determined based on a gradient alignment measure for each pair of clients. In the outer problem, we train personalized models X IR" incorporating 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 follows:
We model collaborative learning through 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 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 (LL.Ms..
2 Problem formulation
In this paper, we model collaborative learning as a bilevel optimization problem, where personalized models X [ trained in the outer problem, and collaborative weights W gax= are determined by the inner problem. More concretely,
min Cn e
f((2s)
Σ Isicisn
wiy
xا112 ))
(Model-Training)
2; +
c; + Fÿ
where wi E arg max Waj [.
Vf
Vfx
Vi,j E [n].
2
2
(Client-Selection)
where p > a hyperparameter for penalization. We break down the formulation follows.
Outer problem: training personalized models. In the outer problem (Model-Trainingg client trains its model æ, by minimizing its loss function S.s along with penalizing its distances to neighboring models, c.g. 2y, weighted by wiy 0.
Our formulation similar DITTO [24], but with two key differences: First, DITTO uses uniform and fixed collaboration weights and penalizes the distance between æ; and global model, whereas we penalize the distances between pairs of clients and adjust the collaboration weights during training. Consequently, when client tasks are heterogeneous sucouuss--chhhhhe performance global model deteriorates, and DITTO's 's local models cannot benefit from fine- grained collaboration. In contrast, our method is able to exploit such structure and achieve better performance in diverse settings.
clients drawn from clusters the
Finding Collaborators. In the inner problem, decompose the task of opti- mizing W E RR into independent sub-problems, one for each entry of W. We relax the binary collaborator selection variable Wi E {0, 1} to continuous weight Wiy [0, 1]. As the objective function is linear with respect to Wajs 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 or 1.
Inner Problem:
It is important to note that wi, does not imply permanent connection between clients i and j. but rather a temporary assessment based on the current states of æ; and Øz..
A simple inner problem with two clients is illustrated in Figure 1. The fi, f2 are their loss functions, and suppose μr, M2, and μ are the minimizers of f1, fz, and ] (f1 fz). Suppose μr, p2, and μ are the minimizers of fi, f2, and (fi + f2) respectively. The model weights at points A, B,C C demonstrate three scenarios for updating W.
2

---

x x3
Model weights x xx,x9xx
Minimizers of and
Minimizer of 21h + f2)
Averaged model weight
Grads for X-update Vf((x),, V2((2)
x
Grads for W-update
vS (-
10
Figure I: Diagram of the inner problem (Client-Selection) represented through contour of 1 (fi The blue arrows gradients computed at middle point (at1 2t2) to determine connectivity. The red arrows + represent gradients computed at local models to update model weights.
Point A: The model A is far away from μ., i.e., | p >>
pl|- The descent directions of the clients have positive inner product; therefore W12 = 1. Collaboration at this stage speeds up training.
max; 14;
Point B: The model BB is closer μs, i.e., ||æ:B pal ~ || pul- In this case, moving closer to the minimizer μa of 1 (f1 + f2) no longer helps both clients get closer to the minimizers of their losses μ4, The inner problem yields W12 O and the clients disconnect.
Point C: The models of and æ. are already disconnected. The gradients computed at their midpoint suggest they should remain disconnected; thus w12 =
Because collaboration weights in Client-Selection are determined in pairwise fashion, our for- methods [11, 39], does not require knowledge cluster sizes, allowing clients to join and leave during collaborative training. This elasticity enables method to be applicable a wider range of scenarios.
mulation, unlike clustering-
Remark (Extensions). While Slient-Selection is defined over a box constraint W E [0, 1]mxn, 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
We propose novel stochastic gradient descent (SGD)-type alternating optimization algorithm, termed CoBo, to solve the bilevel optimization problem defined by (Model- Model-Trainingg) and (Client-Selection The algorithm alternates between updating the model variables X and the collaboration weights W.
In each itcration f, we first fix the model variables { and update the collaboration weights by applying projected gradient ascent with step size y 0 the inner problem (Cliint-Selectionn):
++ w
Projjo.1]
J +
V
x
+ x
V
x + x
Vi,j E [n].
(1)
2
Next, with the updated collaboration weights (u(+1) are fixed, we optimize the model variables (æ;) J using the following update rule derived from the outer problem (Model-Training):
# ++1 x. Vi((aa)) +
Σ
e+1 (x)) w
ax))
Vi E [n].
A=1
where 7) > 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. this implementation. the full gradients [Vfthefe in (1) and (2) are replaced by their stochastic estimates. Additionally, collaborative weights updated with probability of O 1 leading to expected computation of

---

Algorithm CoBo:
Collaborative Learning via Bilevel Optimization
Input: Model parameters Vi E [n]
x = x E Rd Penalization parameter > 0; W E R
where w/t = 1, v i.j [n]; Step size 7,, 0. for round
..... do Call W + Client-Selectee (f= W
2:
3:
for client 1.... do
4:
Draw sample & ~ D; and compute stochastic gradient of E Rd of fi(2)) and update
i m + ar 17 g +
Wik
++1
(ar ark)
(3)
R=1
5: end for
6: end for
7:
Output: Uniform randomly select [T' and return ( ........,. an) and W
8:
9: procedure CLIEN-SSELECTION
W) for each pair clients (i.j) where * E [n] do
10:
11:
if with probability 1/n, then
12:
Compute the average model žij j/ (x, + xz))
13:
Compute stochastic gradient gi-i and Ji-j for fi(zz)) and (ziy) respectively,
Wij + Proj (wiy + (gi-s,
(4)
14:
end if
15:
end for
16:
return updated selection variables W
17: end procedure
(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-Sellccion) in CoBo is solved in pairwise fashion. This pairwise approach makes the algorithm non-blocking and robust to stragglers, providing greater flexibility and efficiency. Not all pairwise weights have be computed in each iteration. In Table we compare the performance of multiple edge-sampling strategies.
Theoretical results
In this section, define assumptions in collaborative learning settings and show that CoBo con- verges to stationary points. The following assumptions regarding the local optimization objective fi commonly adopted in the literature [1, 19]:
(A1) L-smooth. For all x and y in Rd and i [n), the loss function fo has L-Lipschitz gradients, i.e.
(v) ff((i)l L| yll-
(A2) Noise bound. For all a E Rd and E [n)]. there exists o² O such that the stochastic gradient has bounded noise
Eg [IV Ee [Vl E)]]1 s
(A3) Global minimum. For all i [ri], the loss function f. has global lower bound f²-
The next assumption characterizes the possible relationships between clients. In the first case, when reaching the stationary point æ their joint objective ʃ; + Jj+ then by (5) implies that [((()) (x) 0 client i and reach their own stationary points. In the second case, when client reaches its stationary point, the gradient of j lower bounded by positive constant, meaning they don't it share stationary points. This leads to eventual

---

(A4) Collaborativeness. If clients and j are collaborative, then there exists Mi > 0 such that
[VVf(()= (f,())] $ (() + (1)(2))
E
(5)
Otherwise, there exists 23 > 0 such that
V [f
C3 E
(6)
This assumption is similar to [39, Assumptions 4.5], but we define relations for pains of clients instead of clusters. In the next example, we use quadratics to demonstrate (A4)
K clusters with [ni] UnE(k] Ck and Cj n Cx. 0 for all k + k' E [K] Consider the k-th cluster with center Pa and client E Ck. the loss function is fx( b x psvl12 where o₁ >> O. Then for clients j, in the same cluster, i.e. i.j E Cs
Example 2. Assume that there
(a: a;)² [V + V15(2))1
[[: - ( (a: - -a,) - 1
(a: + az)²
The My la:=a: in this case. On the other hand, ford Ck and j Cw and μı: + pkx
a;toj
a?
[VL(0) IV = at ||:0 pa 112 + azl1E: 1ll1
(af + agy3 ppx
pax l12
o?a (a Fagy
where the lower bound Si
lIme
Mer 11 2 > O.
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,45 hold true. Suppose that CoBO solves (4) with mini- batch size Consider clients and j in the same cluster C of size c. Suppose that Mf E (0, 2).
b 2 2Taa(c- 2) and C2 2 [V[((a)) + Vx((x))] (B for all a and Let p 2 s/L and step size
n S min
2
Jo (2%)) 115)
1
oVLT
2/3L
The consensus distance also converges to 0, i.e.
B 111
T-1
6.M2
1j p2o²
Loz Σ
Jis (2 131
TT
c2T t.jEC
L= IJEC
Moreover, the gradient norm is upper bounded.
ΣΣ I=0 LeE
T E
Vi (z))
<3
Loz PT
Σ 1jEC
(Jo f1))
T.T
This theorem suggests that clients inside the same cluster gradually reach consensus. This cluster- level consensus model reaches stationary point of their losses by (5). Note that a larger penalization parameter and smaller values of M3 lead to faster convergence, which aligns with our expectations. Note that Mij in (A4) measures how well i.j collaborate. A smaller Mi, leads better consensus distance in Theorem I, with M,j- = O leading to identical data distribution. The following corollary states the convergence of norm of client gradient of model x;
Corollary II. Under same conditions as Theorem I, V fa (a!))2 converges at a similar rate
Σ 1= eeeC
Σ
V
(x))
Lo² T2T ijec
Σ
(fg
-ft)
TT
5

---

4 Experiments
In this section, we present three experiments demonstrate the practical effectiveness of CoBo. In the first two experiments, we benchmark CoBo in both cross-silo federated learning setup involving 8 clients and 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 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, 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 accuracy and perplexity, include this advantage. In addition to previous baselines, we also compare CoBo with 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, 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 other collaborative learning algorithms.
We simulate cross-silo multi-task environment where training 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 scenario where 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 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 diigonal 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 benchmark for the ideal collaboration structure 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 constant probability of (L/n)). Additionally, we observe that CoBo identifies an appropriate collaboration matrix early in the training process, motivating the use of time-step -te-tteeedeeeeentt sampling rate, ((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 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.

---

Table 1: Comparison of the average performance of CoBo across different sampling strategies for updating the weights of client pairs in the collaboration matrix. All strategies demonstrate performance close to that of the non-sampling oracle. However, mixed strategy, which combines constant sampling rate at the start with time-dependent rate during later training phases, shows superior performance.
Acc.(%)
Loss
Constant (0(1
73.05
1.104
Time-dependent (O(1/L)
73.18
1.226
Mixed
74.77
1.081
No Sampling (Oracle)
74.93
1.278
75.0
Accuracy
70
67.5
60
60.0
0.1
0.3
0.5
50
Fraction of dataset
Accuracy
74.0
40
Accuracy
73.0
30
72.0
20
4
5
Clusters
10
74.0
Accuracy
100
200
300
400
73.5
Iterations (x100)
73.0
CoBo
Oracle
FC
Local
3
4
Finetunod PedAvg
IFCA
FodAvg
Start of fine-tuning
Clients per cluster
(a)
(b)
Figure 2: (2a) Average accuracy in cross-silo experiments with varying factors, including the fraction of the dataset available clients, the number of clusters, and the number of clients per cluster. (2b) Average accuracy of personalized models for cross-silo federated learning with clients. The "Oracle" denotes applying FedAvg to the clients with the same label permutation.
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 fixed accuracy of 60% in 4,500 iterations, which 30% faster than DITTO. For better comparison, the values accuracy and loss reported in Table 2. Additionally, the evolution of the collaboration matrix for clustering algorithms and CoBo illustrated in Figure CoBo starts to identify clients with similar label permutations as early 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
the end of training.
Figure presents the results 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 cluster directly impacts the performance 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
1.0
FC
0.8
0.6
IFCA
0.4
CoBo
0.2
0.0
Figure 3: Collaboration matrices learned by Federated Clustering (FC), IFCA, and CoBo different stages of training for cross-silo experiment with 8 clients. The diagonals are masked out. The oracle matrix is block diagonal matrix with blocks size 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 fully-connected matrix while FC occasionally diverges from the achieved cluster structures (e.g., iterations 300, 5000, and 40000), even at the end training.
the number of clusters does not significantly affect CoBo accuracy. Additionally, we investigate the effect of varying the number of clients per cluster while maintaining 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, demonstrate the performance of CoBo in challenging cross-device federated learning setting characterized by significant data heterogeneity. We create 10 clusters of varying sizes: clusters consist of 6 clients each, another 2 comprise 7 clients each, and allocated data from 10 distinct classes out of the total 100 classes available in the CIFAR- 100 dataset, ensuring that the data across clusters disjoint. Within each cluster, the data distributed uniformly at random among the clients. We proceed to train individual ResNet-9 models [13] on each client's s data for total of 20,000 iterations. This setup allows 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.
Each cluster
We define the oracle algorithm and the corresponding collaboration matrix in the same manner as in Section 4. Note that while manually create the clusters, inter-cluster collaboration may still be helpful in practice, and it is impossible to know the actual ground truth 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, evidenced by the results presented in Table 2. The collaboration matrix and accuracy plots are deferred to Figure and Figure 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 fine-gr fine-grained level.

---

Table Comparisons of model quality and fairness measure of personalized models for cross-silo experiment with 8 clients, and cross cross-device experiment with 80 clients, and the language modelling experiment with 4 clients having different languages. Federated clustering (FC) is not scalable with number of clients due to its (nn²) complexity, and therefore ignored in the cross-device fl experiment. The clustering algorithms IFCA and FC are not applicable to L.LMs and there ignored. Note that Oracle is not defined in the LLMs experiment. The column "Imp.(%)" demonstrates the percentage of clients with improved performance compared local training.
Cross-silo
Cross-device
Fine-tuning LLMs
Acc.(%)
Loss
Imp.(%)
Acc. (b)
Loss
Imp.(%)
Perplexity
Imp. (%)
Local
64.9
1.67
54.9 0.1
1.40
41.26 0.38
FedAvg
18.8
2.66
53.9 0.1
1.79
64.86 84 0.00
Fine-tuning FedAvg
0.2
58.9 + 0.1
1.88
46.70 0.07
Ditto
0.3
1.55
100
70.3 0.1
1.21
40.05 4 0.01
IFCA
18.6
2.75
0.8
2.15
FC
5..1
1.79
Validation Based
42.90 + 1.68
CoBo
74.6 0.2
1.08
100
79.6 0.4
0.97
39.28 + 0.01
Oracle
7.24
1.07
100
83.6
0.70
100
4.3 Collaborative fine-tuning 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 fine- tuning; however, the scarcity of data often leads to inferior performance and necessitates collaboration. Therefore, we conduct an experiment with four clients, each having pre-trained GPT-2 base model? with 124 million parameters in total [32], and subset of articles from the Wiki-40B dataset [12] in one 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 valu- able. Clustering methods fall short in this aspect due to their binary, discrete outputs, which do not capture the nuanced degrees of collabora- tion needed. CoBo addresses this limitation by allowing for continuous range of collabora- tion intensities, achieved by simple yet effec- tive modification the projection domain in (1).
0.8
Domain Weights
0.6
0.4
Specifically, we employ probability simplex, denoted as A (wi 2 0, E, wiy = 1). the domain of the inner problem.
0.2
0.0
In Table 2, we compare the perplexity of CoBo with baselines after 500 iterations, when FE- DAYG converges. There are no oracle domain weights in this experiment due to the com- plex coherence among languages; therefore, we omit the oracle algorithm in the table. CoBo achieves the best perplexity among all algo- rithms. In Figure 4, we demonstrate the domain weights learned for the Catalan language. Over- all, Catalan assigns the highest collaboration weight to Spanish, which is reasonable consid- ering the similarity between the two languages.
0
25
50
75
100
Iterations
Spanish
German
Dutch
Figure 4: Domain weights found by CoBo for Catalan language. There are 4 domains in total: Catalan, Spanish, German, and Dutch. The curves are smoothed by exponential moving average.
http://ghhhbb.ommmmaaaaaaaaaaaaa

---

collaborative learning framework to 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 regularization term that penalizes the divergence from 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 al. [26], Ghosh 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 single-level objective function. Zantedeschi et [40] minimize task losses augmented with a penalization term Wiylaa | 21112. similar our outer problem. However, optimizing wi directly can lead to degenerate solution (we = 0), which necessitates an additional penalization for small Wij 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- selectinn function to be consistent with client selection, which can negatively impact performance. Apart from multi-task federated learning, similar bilevel optimization formulation has been used by Le Bars et al. [23] to find sparse mixing matrix while training consensus model in the outer problem.
Bilevel optimization and alternating optimization. Bilevel optimization is powerful tool which models a broad range of problems, such as reinforcement learning [6, 30, 16, 15, 37, 31], and linearly- solvable Markov decision process [5], meta-learning [8, 20], etc. A typical bilevel optimization problem, as the name indicates, consists an outer and an inner optimization problem whose variables are inter-dependent. Typical bilevel optimization solvers requires hessian information which 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-graid searse-grained collaboration, which leads to inferior performance in practice. To address this issue, 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 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 hetween collaborators when updating weight w which may raise privacy concerns. We defer the discussion of future work.
Acknowledgement. We acknowledge funding from Swiss National Science Foundation (SNSF) grant number 200020 200020_200342 from Huawei Cloud Intelligent Cloud Technologies Initiative, and from Google Research Collaborations.
References
[1] Yossi Arjevani, Yair Carmon, John C Duchi, Dylan J Foster, Nathan Srebro, and Blake Wood- worth. Lower bounds for non-convex stochastic optimization. Mathematical Programming, 199 (1):165-214, 2023.
10

---

[2] James C Bezdek and Richard Hathaway. Convergence of alternating optimization. Neural, Parallel Scientific Computations, 114))335-366, 2003.
[3] Aymeric Capitaine, Etienne Boursier, Antoine Scheid, Eric Moulines, Michael I Jordan, El- Mahdi El-Mhamdi, and Alain Durmus. Unravelling in collaborative learning. arXiv preprint arXiv:240.114332, 2024.
[4] 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.
[5] 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:// http:://proceednng. mlr www..ppre//vee.e. press/v54/ dai17a 17a. html.
[6] 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 Jen- nifer G. Dy and Andreas Krause, editors, Proceedings of the 35th International Conference on Machine Learning, ICMI. 2018, Stockhollssüussan, Stockholm, Sweden, July 10-15, 2018, volume 80 of Proceedings of Machine Learning Research, pages 1133-1142. PMLR, 2018. URL http: /ttp://///.occceng....mm mlr. pwwss/eessseeeee..p /voss///v///e /dai8a111e html.
[7] Mathieu Even, Laurent Massoulié, and Kevin Scaman. On sample optimality in personalized collaborative and federated learning. Advances in Neural Information Processing Systems, 35: 2 212-225, 2022.
[8] Chelsea Finn, Pieter Abbeel, and Sergey Levine. Model-agnostic meta-learning for fast adapta- tion 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: http///preeennn...../ wttpesmmr..pprrree.e walr.press/vvi /voss//vvO//aa iinnina..tha
[9] Gersende Fort. Federated expectation maximization with heterogeneity mitigation and variance reduction.
[10] Marguerite Frank, Philip Wolfe, et al. An algorithm for quadratic programming. Naval research logistics quarterly, 1100,, 1956.
[11] 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.
[12] Mandy Guo, Zihang Dai, Denny Vrandečié, and Rami Al-Rfou. Wiki-40B: Multilingual language model dataset. In LREC Proceedings of the Twelfth Language Resources and Evaluation Conference, pages 2440-2440 2452. European Language Resources Association, May 2020. ISBN 779-110955446344- URL https: http:///aaaaallloao//a org/2020 1rec- 102977.
[13] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition, 2015.
[14] Edward 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: httpss///ppnrreieet sst/foormmiddnneeWeeee
[15] Yifan Hu, Siqi Zhang, Xin Chen, and Niao He. Biased stochastic first-order meth- ods for conditional stochastic optimization and applications in meta learning. In Hugo Larochelle, Mare' Aurelio Ranzato, Raia Hadsell,
Balcan, and Eisuan-Tien
11

---

Lin. editors. Advances in Neural Information Processing Systems 33: Annual Con- ference on Neural Information Processing Systems 2020. NeurIPS 2020, December 6- 12, 2020, virtmal, 2020. URL https::/pproceedinggs www.ssnnrrppp... wttpripp.cppppppee0eee http:ppppppppp//pppeee 1cdf 1444444999111117777767661ee
Abstract. html.
[16] Feihu Huang, Junyi Li, and Shangqian Gao. Biadam: Fast adaptive bilevel optimization methods.
arXiv preprint arXiv:2 106. 1396, 2021.
[17] Martin Jaggi. Revisiting frank-wolfe: Projection-free free sparse convex optimization. In Interna-
tional conference on machine learning, pages 427-435. PMLR, 2013.
[18] Peter Kairouz, H Brendan McMahan, Brendan Avent, Aurélien Bellet, Mehdi Bennis, Ar- jun 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.
[19) Anastasia Koloskova, Nicolas Loizou, Sadra Boreiri, Martin Jaggi, and Sebastian U. Stich. A unified theory of decentralized SGD with changing topology and local updates. 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 pttt:://proceedings. wwa.eeediinnmmm........m.mm.mm.. mlr. wwwpsseeeiiiieeeee htt11199/koloskoaaaaaaaa html.
[20] Guy Kornowski, Swati Padmanabhan, Kai Wang, Zhe Zhang, and Suvrit Sra. First-order methods for linearly constrained bilevel optimization. arXiv preprint M4Xv:24006.1271,, 2024.
[21] Alex Krizhevsky, Geoffrey Hinton, et al. Learning multiple layers of features from tiny images. 2009.
[22] 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.
[23] Batiste Le Bars, Aurélien Bellet, Mare 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.
[24] T. Li, Shengyuan Hu, Ahmad Beirami, and Virginia Smith. Fair and robust federated learning through personalization. In 38th International Conference Machine Learning, 2021.
[25] Xiangru Lian, Ce Zhang, Huan Zhang, Cho-Jui Hsich, Wei Zhang, and Ji Liu. Can de- centralized algorithms outperform centralized algorithms? A case study for decentralized parallel stochastic gradient descent. In Isabelle Guyon, Ulrike von Luxburg, Samy Ben- gio, Hanna M. Wallach, Rob Fergus, S. V. N. Vishwanathan, and Roman Gamett, edi- tors, Advances in Neural Information Processing Systems 30: Annual Conference on Neu- ral Information Processing Systems 2017, December 4-9, 2017. Long Beach, CA, USA, pages 5330-5340, 2017. URL https: http://prooeeennaeee/eeeeeee neurips hpppp//ppe////////// www.ee44eed----s-taeeaeeee Abstract httml.
[26] Yishay Mansour, Mchryar Mohri, Jae Ro, and Ananda Theertha Suresh. Three approaches for personalization with applications to federated learning. 2021.
[27] Othmane Marfoq, Giovanni Neglia, Lactitia Kameni, and Richard Vidal. Federated multi-task learning under mixture of distributions. In Proceedings of the 35th International Conference on Machine Learning, volume 34, 2021.
[28] Brendan McMahan, Eider Moore, Daniel Ramage, Seth Hampson, and Blaise Aguera y Arcas. Communication-effici lcarning of deep networks from decentralized data. In Artificial intelligence and statistics, pages 1273-1282. PMLR, 2017.
[29] Mehryar Mohri, Gary Sivek, and Ananda Theertha Suresh. Agnostic federated learning. In International Conference Machine Learning, pages 4615-4625. PMLR, 2019.
12

---

[30] Ofir Nachum and Bo Dai. Reinforcement learning via fenchel-rockafellaa duality. ArXiv preprint, 0ts/2001.01866 2020. URL
01866.
[31] leva Petrulionyte, Julien Mairal, and Michael Arbel. Functional bilevel optimization for machine learning, arXiv preprint arXiv:2403.2022 2024.
[32] Alec Radford, Jeff Wu, Rewon Child, David Luan, Dario Amodei, and Iyya Sutskever. Language models are unsupervised multitask learners. 2019.
[33] Felix Sattler, Klbbrr Muller, and Wojciech Samek. Clustered federated learning: Model- agnostic distributed multi-task optimization under privacy constraints. 2019.
[34] Virginia Smith, Chao-Kai Chiang. Maziar Sanjabi, and Ameet S Talwalkar, Federated multi-tase learning. Advances in neural information processing systems, 30, 2017.
[35] Canh T Dinh, Nguyen Tran, and Josh Nguyen. Personalized federated learning with moreau envelopes. Advances Neural Information Processing Systems, 33:21394-21405,, 2020.
[36] Alysa Ziying Tan, Han Yu, Lizhen Cui, and Qiang Yang. Towards personalized federated learning. IEEE Transactions Neural Networks and Learning Systems, 2022.
[37] 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.
[38] 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://openreviee net.fee//feeeeee forum? id=bwo3GVsgOv..
[39] Mariel Werner, Lie He, Sai Praneeth Karimireddy, Michael Jordan, and Martin Jaggi. Provably personalized and robust federated learning. arXiv preprint arXiv:2306.08300 2023.
[40] Valentina Zantedeschi, Aurélien Bellet, and Mare Tommasi. Fully decentralized joint learning of personalized models and collaboration graphs. In International Conference on Artificial Intelligence and Statistics, pages 864-874. PMLR, 2020.
13

---

A Theory
Let Z'z := (ar) + =3)) be the average iterate of æf and z and fly(2) == 2 + fz()) be an averaged objective function.
Lemma 3. Suppose (AI) holds true. Let 1j S 21. Then for i,j in the same cluster C. of size c.
1
Σ
(z
2 1
Σ
fis (21)) Lfff (z) D1
Σ
| x3) 112
2
c2
JEC
WEC
i.ECC
12
Lno²
3n
p
En (wit wtal 1
en
IEC
where coefficient Dj is defined
31,
+ +322p² + Imp (0le-2)0 -2)20²
DI(L,,,,,,o0)
4
b
A.1 Proof of Lemma
Proof. Let hy and ht be independent and unbiased estimates of fi((z()) (=f,) respectively. The
variance of h has variance of o²
6 Let's denote E-1-] := Eg ......... [|2]]] and Ex == BA A......... [-]
and let [E[1] = E, [-I]- By the L-smoothness assumption (Al) and bounded noise assumption (A2)
E,E, fo ( Sig (=() +
faz (z/)) Eg |
1)
ER Eggg - =1,1113
+
L
En || - و\ح - E. F.g 1 1|,1
2
The last quantity, i.e. variance of = = can be expanded as follows
E, [||=|]] 21- - E, Ex ++1 -111 EE [1 E, 111
E, [||=(]] + RE -,, [|[]]
-E, [1|En |=ttt] E,E,,22/-1111] E, [1|= 1hlat-11113
-Eg [|IE 11 - -1,,,,, + E, [1l2 - -Enlatt
1 (g) + 9 vV(((())) v)(()))) + E, |||2511 - E, E2
-y
2
11002
4 + E,
||| (+
E. 12
T,
where the linear term vanishes due to the expectation of h in the second equality and the definition of 1+1 and (3) for the third equality. Then
LIRo²
L
E
Fi
(z())
E
[2
c+1 - 21,1> +
+
Tsj.
2
4
2
Expand the inner product with equality - (x. y)
= 11=:111 1leel
+ -l1l==
=11 yields
E
(+1
=2/
B
E
Jis
(2
<in (z1g)
+
?
+ V (2())
"
2
- 2
V Jes
Tay
E B
z
Inda²
L
Tes
"
4
2
|3|
1+11
21
<ky (z() + ?
3j
+ V (z())
2
"
17 4
B(2
2//
2
FFu (2/))
1|2
LIRO²
+
L
T
2
2
14

---

where 1) S 2T. is applied in the last inequality. Then the upper bound of (2()) ||
L
|
is that
fiy (2t))1 11 =
2 (frs (2/)) Elfo (-)))))
B
1+11
V
fig
(z/z)
n
(7)
Inj² 2 L +
Tíj
2
17 2
The T can be upper bounded by expanding 2 with (2)
E E zi
V fis (z/})
2
V hig (2/))
1/2 (Vfi(at)) (a)(a))
0
Σ
EA
++
[ - xx)) +
Balw jk
1 1[00] - stih
2
k=1
<3
V fo (2))
(v (zz())))
+3
w (2f - xx)1 +
Σ
Win (af - xx/2
2
2
km
Ti
To
2
+3
3
P
Σ
En [w (x)
x1) +
EE [w 1u44
xx/²
2
Bound Ti: Use L-smoothness of fa and fy. Take expectation with respect to g and 9 which are unbiased estimates of f((()) and (z))
1 (Vi((z))) + v.s(x²)))
T
Vis (2/1)
2
2
|| +1112++
|2// 2311
2
2
1.2
|
2 11
4
Bound Tz: Use Stuchy-Schwarc inequality and C has a cluster size of c
Ta < cp² Σ wie lot
2ll12 + Σ
2
1=1
Bound Ts: Use Cauchy-Schwarc inequality
Tg <
np²
Σ
EA (w
+
wijx
xil12 + Σ
En
(wo 11 120 - ardl(]
2
km
15

---

Viy (z1g)
Sum up
for all of j in the same cluster C for (7) yields
Σ
Vis (
(1
=2
2 Σ (Z))) ELio (z)
3
| 20/11
AJECC
iJEC
1.EC
+ 3622p² 2
Σ
||0
2 + 3nep 2
Ex[w (u [+] - wottl 2 a2 - x 11
iJEC
iEC k=
e² Lmjo²
2
2
Tis
2
17 2
2
Trs (2/s)
Jig (z))))) +
31.2
+ +e2pp
Σ AEEC
|| - 2
"1 WEC
4
+ 3ncppp
En [wttt
woee 12 1|
x 1/21
c² Lnd?
2 L
Σ Tt.
17 2
IEC Am
1,JC
(8)
Now we expand Tij as follows
Ti - E,,,E B,
(+1 - En 1 ij
= En
(ot + .b)) + 70
Σ
Wikk
++1
(xf - =//) + w ++1 (at) xx(1)
2
2
k=1
- EA,
(of + 93) +
?p8
(w ++
(af x + uitt (an) =xt1
2
2
dk
k=1
27
-Ep
7P
(w - (u (005
at) + (wft
En En M22 - z1.))
2
k=1
2
n²p p2P²
En
| - En
ww |112]
|| - a + EE
||u - Ex (u (+1
1|
27
1105
- 1113
4
where we use the independence of random variables in the last equality. Average the above equality over i, j C yields
pp 2)) 2)
Σ
Ty
Bn
[|ut - Es []]]]] 27
|| - =1,113
20t
i.j
112 p2 2)
402
Σ
|aa 13111.
2c²
16

---

Therefore, (8) is upper bounded with
(z 112 2 Σ
Jis (2/) - Blfas (=(")) + (3L2
4 *34 + +2222²
Σ
|oc -0911
1) 1,EEC
5.3EC
a
c² Lnjo²
+ 3ncp²
En [w
1
w 12 क z +
2
iEC k=1
L 12 p²
2) 40²
Σ | 2112
17 2
2
b
2
(Ts (2/)
- Bisis (z 11
11 J,JEC
3L2
+ Lmp (c- - 2)20² b
+
32² p²
Σ
| 23
4
WE
Lnoo
+ 3ncpp
Σ
Bn [uitt 11- - wirl 1 a0
xiil2 +
2
/EC
km
Lemma 4. Suppose Mij S 5.2 Let p 2 L31, b 22Lm(C- - 2)0²² and 1) S
1 S
then
2pc
2/3L
[[]] S - (ppc)) 1
Σ E
1E -
Σ
| - 2
AJEC
5mm/ /²
En [w - w.al
|ac| xf112
L
on
iEC R=1
6M 1
fis (z//) EE [is (1)
pc 2²
ijEC
1
33LLM3
n0²
2pc
Proof. Expand xf atft1 with (2)
x 1+1 20/t [+1 = x - 21 17 (g)- - gÿ) - -10 Σ
1+1
(z) - -x/))= - -wst 1 (2.) - xx))
sk
k=
As i and j belong to the same cluster (i.e., wij = 1). we add +2mp Σi_ (2s) x1))
x/- 2/-1 2 1+1 =(1 (xpc) =x/)== 17 (g) - 9¹)
(wttt ++1 (at - x - (wf
++1
(ay - axk))
mp
Compute the norm of ar!11
t+1 and choose mp S 1 26 to use Jensen's inequality
B. [2] - 25511 "112 5(1 - 2mpc) |لم
20/112
+ 2:ppcc
(u - (a)) a/)) (w wjy(as a1))
2r
k=1
(Vi((z))) - v.(((a))) 12 2
1
+ 2pc
17

---

Expand the right- hand side with Cauchy-Schwarz inequality
| [ - -x² 213"1112 5(1 - -2mpc)
x xأل12
+ anpc
Σ
(10tx - - ax)) - (wf - -3je) -
- x
2c
k=1
|
+ 4npc
(Vf((x)))
vfJ(zz}))
2pc
1 n
5(1 --mpc))|
| xiil2 + 8npc
Σ
2c
1
Snpe
(wott - (x0) 2xk/
20
1
1
4mpc
(Vfi(x)))
(ff(a))))
2pc
2
<11 1-mpc))plla
it x3i12² +
2nmp
w - xx11 -
2nmp
1 -
x
"
- ( [fi((
pc
The last term T can be upper bounded by adding + (z/z) + VJ (2/y) and use L-smoothness assumption (A1) of fi and that i, j belong to the same cluster (A4)
7 -= [=) ( -
/3/lV/
V (++) +3 (z1y) vis (2
+ (z) v |
3L
|00| 2//18 + 3M31 V / (=1)) V/ (=
3I22
|| 10/12 + +M3 V Vl (2 ||
2
Note that E|[ 1+1
9j -Egi + + +E,1|E,(= -
1-p-l==(+11 11 and inde-
pendence between randomness on worker and j where expectation take to all of the randomness until time t By averaging for i E C
3,1.2
2
< 1 1-22pp
2mpc +
Σ | 23)3 WEEC
2pc
IJEC
4n ap
+
En [lutt 1 - - 22] ||00
2
c cn
iEC k=1
11 20
3n 1j
1
+
+
pc
pc
iJEC
18

---

Use the previous Lemma 3 to bound Eweec
(2/))
1
Σ
EE
++
++1
21 2]
2²
SEEC
S
2npe +
31112
1
Σ
||a|
4n2np mp
n
Bn,
[lust
-wi 2]] ||,55
sef. 11
2pc
2²
c
cn
IE
IEC 1=1
+
1102² 1
+
Mr
12
2
Σ
(2))) - E
Pos ( + DI
1
Σ
|
22,12
2pc c²
pC
17
1.jEC
K.EE
3
Layo²
+
1j
3n
2
EBs
| e+1 - wirl
|2 - 2i1122 +
pc
2
REC
km
Rearrange the terms
E
11
112]
IJE
1
2npc+
+
31L2
M D₁
1
|ag| 23112
2pc
pc
L.EEC
An?up
M 3n 2 p²
En || - -wotel ||sc; x111
+
pe
iEC 1=
3n,1 2
(fss (z/z) - fis (z1) + 11²2 1
2pc c² + Mr Lndd
Σ
pc
c²2
pc 2
1,ECC
We would like to achieve
(1-2npe 2npc +
(1-
3n12 200
+ pc
anM Di) S - (p-) by
letting p 2 33L s.t.
3q1,² 2pc
5
ape 2
letting b 62 2LL - 2)0² and p 2 /3L
and Mi S _. the following inequality hold true
M 1j
+ 3e2p² + Lnp²(c 2]]20²
Mr 5j
15
43
1
Mo S npc.
pC
4
b
pc 4
4
2
Using the same requirement that p 2
33L
and
Mij S
An²mp
Mr 3n²p²
4n²np²
3
3/3n p²
5mm2²p²
c
pc
V3L
L
L
The upper bound of 1/62 Eisecll
1//// - xz// x2/+111 11/
can be simplified to
Σ
|
+1 - 2551113 [[2] S(1 - <(1-npc)
Σ
| - <²) 2²
c² SEEC
c²
WJEC
5 p² 1
En [1ut wisa?
|
L
iEC k=1
6M2 1
Σ
(Tss (2(y) -E [As (2)
+
pc
ijEC
1
33LM33
no?
+
x
2pc.
19

---

A.2
Proof of Theorem I
Proof. Given Lemma 4 and average over time Over T and take expectation to all randomness throughout training
T 1
||
$ npe)
|| 25/12]
e²T
2T
EJEC
t=
iJEC
5²
E [ [ - wnje 1| - or 11]
-
L
Ten
1=0 iEC k=1
6M²
T-1
E Jey (=t)) - E [Tt (2(+)))
Te
pc
1= SJEC
1102
+ c +
2pc
Rearrange T²T Ers Eisec
E
||
2
2 [[]]
yields
+ [
2,,1112
5mmp 1
EE [
woll?110 - x [12)
2T 1=0 iJECC
Le Ten
1=0 sEC k=
6M
T-1
E
( (is
Fo Fu
n2 eT
1= i.jEC
[ +3LLM3
0²
20²
Consider bounding |w/r-- wiyl 1 in two cases
Case 1: wit = 1. Suppose Mis E (0, then [ VALEEE S MBIT)CH) -
(2( implies
1 1- M
(V
2(1 + M (1 112 +
20.
then w = w = and therefore | - 1 = O.
Case 2: w O. Suppose C. 2 f((x)) (=) for all then
[+ + V.
013 + 2(V V
> c4 + (V)) Vf+(())))
which means the inner product (V (+)(())))) < 0 is negative, i.e., wij 1+1 = 0 = w/y- The above cases hold true for well initialized weights x..
Then with lower bound assumption of St and /j (A3)
6M 1
[|a
25 1|
Σ
Σ
E
(Te (21) - Bo (xff'))
TT
c CT
KJECC
1=0 t.jEC
( x + LLL
o²
2p2c2
6M 1
Σ R.
Jo (=+) (1b
(21))
np²c cT
WJEC
(12 3nLM3
o²
2p2c2
20

---

Minimize the upper bound through choosing
2
1
Σ
(29) - 13)
such that
E
21
Lo²
Σ
f (2 ft) )
(9)
T 1=0 I.eCC
2
p²
o2T
WEE
By the result of Lemma 3
1
Σ
Σ
Vi (+
[[]]
-
21
Σ LEEC
(is (2/)) - 1ts) + 462 /² T
||ac a?][?]]
+
6² Lajo
T
2
t= EJEC
j.jEC
$26²
Lo²
Σ
fi (2.) - fit))
+ 2e² p2
E
||a|
- x [1
C2T
T
iJEC
1=0 iJeC
Lo²
$ <220² + +2Ac² Mij)
Σ
(fss (22) - fis)
c2T
ijEC
Lo²
S3c²
Σ
fi (2ß)) - -í3)
c²T
iJEC
A.3 Proof of Corollary II
Proof. By adding [[V[(() + V on both sides of (5), and replace & with Zij we have
2 (IV + (1))))))) < 4(1 + MB) V (in(
)
E
(10)
Then using the upper bound of My < 1 1/5 from Theorem 1, and average over and i,j yields
[IVA & (z 2 ( 25
Loz TT S.EC
Σ
(=i)
1t)-
(11)
TT
1=
+JEC
By Mouchy-Schwarz inequality and L-Lipschitz smoothness, we have that
T-
T-1
T-1
V
(x))
| (1
Σ
Σ
( - V.. (2:)
2T
TT
1=
SCEEC
1=
IJEC
m0 IJE
L2
T
M
Σ
| - 20y
+
2T
IJEC
1=0 KJEC
Applying (9) and (11) to the upper bound of the above inequality
T
Lo²
(x:))l2 <
1 + 25
(26) - ft)
2T
eT iJEC
Cee
L 6M2
Lo² Σ
(iy (zg) - ft))
p²c²
T2T isEe
As pc 2 v3L and Mij- < as stated in Theorem I.
VI (x:) 54
Lo
Σ
(Jo - -tt)
T
c²T
LJEC
IJEC
21

---

B Experimental Details
In Section 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: use the fix batch size of 128 for cross-device, and cross-silo experiments on CIFAR-100 100. We tune each method for the optimal learning rate individually: and 0.01 for all other methods. For Ditto, We use the their paper. For Federated Clustering, 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 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 four NVIDIA V-100 GPUs with 32 GB memory for cross-devies experiment. With this setup, running CoBo for cross and cross device experiment takes 9 hours and 28 hours respectively.
use learning rate of 0.1 for ditto, 0.05 for Federated Clustering (FC), of À = as recommended in the ground truth number of clusters and size of clusters
For Language modeling experiment, we conducted the experiments with the learning rate 0.002, batch size of 50, and 4 accumulation steps. Note that each agent only get subset of the regarding language from Wiki-40B dataset, consisting of total 800000 tokens. We also used the context length of 512, dropout rate of 0.1, and LoRA module with rank 4. Training performed on 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 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
1.0
CoBo
0.0
Figure 5: Collaboration matrices learned by CoBo at different stages of training for cross-device experiment with 80 clients. The diagonals are masked out. The oracle matrix a block diagonal matrix, consisting of 10 blocks: two blocks of size 10, two blocks of size 9, and on. The collaboration matrix CoBo already starts to look similar to oracle matrix within as low as 300 iterations. (1.5% of the total iterations)
For the cross-devics 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]
The theoretical results are given in Section 3 and empirical evaluations are presented in Section 4.
Justification:
2. Limitations
22

---

80
70
CoBo
Oracle
60
Ditto
Accuracy
Finetuned
50
FedAvg
Local
40
FedAvg
IFCA
30
Start
fine-tuning
20
10
0
25
50
75
100
125
150
175
200
Iterations (x100)
Figure 6: Averaged accuracies personalized models for cross-device federated learning with 80 clients. The "Oracle" denotes applying FedAvg to the clients having the data from the same classes of CIFAR-100 dataset.
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 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 ex- perimental results of the paper to the extent that affects the main claims and/or conclusions of the paper (regardless 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 instruc- tions faithfully reproduce the main experimental results, 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, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary understand the results?
Answer: [Yes]
Justification: The full training setup described the main text. The hyper-parameters are also available both for CoBo and other baselines either in the main text appendix.
23

---

7. Experiment Statistical Significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
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 com- puter 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 men- tioned information are discussed and we believe it is sufficient for reproducing the experi- ments.
9. Code Of Ethics
Question: Does the research conducted the paper conform, in every respect, with the NeurlPS Code of Ethics https: http:::////uiippp/e httpol/i/eehhiscueeeelllele.
Answer: [Yes]
Justification: We preserve the NeurIPS Code of Ethics throughout the research project.
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
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)?
Question:
Answer: [NA]
Justification: Note applicable.
12. Licenses for existing assets
Question: Are the creators 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 details about compensation (if any)?
Answer: [NA]
Justification: This paper does not involve crowdsoureing experiments.
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.
25


Reviews:
[{"review_id": "k5nUNeqL9r", "reviewer": "Reviewer_xPeE", "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.", "questions": "- What does \"each cluster clients share stationary points\" mean? \n- Why is the convergence result is given on f_ij\n- What is the performance of method when clustering assumption does not hold, for instance there is only one cluster?", "limitations": "Yes", "rating": 4, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "- Paper is very easy to follow and understand due to the flow of language. \n- The presented results indicate the advantages compared to other works. \n- 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.\n- FedAvg+fine-tuning should be compared to as an additional method.\n- 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.\n- The algorithm requires pairwise computations across clients which might be huge overhead for practical applications where there might be millions of clients. \n- 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their thoughtful feedback on our paper. We address each of the comments below:\n1. Sparsity in experiments and ablations:\nWe have added more experiments to investigate the effect of different settings, such as the number of clusters, data samples per client, and number of clients per cluster, on the performance of our method. This provides a more comprehensive understanding of the robustness and generalizability of CoBo. The full experiment is available in the global PDF file. As expected, the total amount of data available to each cluster has a direct influence on the performance of the method. Additionally in the experiments, CoBo is almost invariant to the number of clusters. In the experiment for different number of clients per cluster, the dataset is partitioned among clients in each cluster, meaning that larger cluster sizes result in less data available to each individual client. Nevertheless, CoBo can utilize collaboration to preserve the performance for larger cluster sizes, and remain invariant to cluster size as well.\n2. Comparison to FedAvg+fine-tuning:\nWe have included a comparison of our method with FedAvg+fine-tuning in the revised manuscript, allowing for a more complete evaluation of the performance of CoBo relative to existing methods. The performance of this baseline could be found in the table available in the global PDF file. We confirm that CoBo can outperform this method in all experiments.\n3. Convergence result on f_i:\nWe agree with the reviewer that the results should be given for $f_i$. We note that the convergence results for $f_i$ can be easily derived by combining the last equation of Theorem I with Equation (5). That is, by adding \n$\\lVert \\nabla f_i(x) + \\nabla f_j(x) \\rVert_2^2$ to both sides of the Equation (5), we have that \n$$2(\\lVert \\nabla f_i(z_{ij}) \\rVert_2^2 + \\lVert \\nabla f_j(z_{ij}) \\rVert_2^2) \\le 4 (1+ M_{ij}^2)  \\lVert \\nabla  f_i (z_{ij})  \\rVert_2^2 $$\nThen using the upper bound of $M_{ij} < ⅕$ in Theorem 1 and average over t and i,j yields\n$$\\frac{1}{c^2T} \\sum_{t=0}^{T-1} \\sum_{ij\\in c} \\lVert \\nabla f_i(z_{ij}) \\rVert_2^2 \\le (1+1/25) RHS$$\nwhere RHS refers to the right hand side of the last equation of Theorem 1. The convergence of $\\lVert \\nabla f_i(x_i) \\rVert_2^2$ can also be derived using Cauchy-Schwarz inequality and apply the consensus distance inequality in Theorem 1. \n4. Pairwise computation overhead:\nWe agree with the reviewer that the pairwise gradient computation in the inner problem can be expansive. In fact, we had already addressed this issue in Section 2.1 of the manuscript by uniformly sampling only O(1/n) percent of edges in each iteration, leaving the computation complexity of the inner problem the same order as the outer problem. We also included an experiment to empirically compare different sampling strategies, available in the global PDF file. The experiment demonstrates that CoBo is robust to different sampling strategies, while proposing a sampling method that slightly increases the performance. \n\n5. Discussion of the theorem:\nWe thank the reviewer for the insightful suggestion. We will incorporate the discussions into the paper. As M_ij measures how well i,j collaborate, in fact, smaller M_{ij} leads to better consensus distance, with M_ij=0 leading to always identical as expected. The gradient norms does not scale linearly with the number of clients because we are considering the norm of individual gradients. (This scaling will appear when we consider the averaged gradient among all clients in this cluster. We will add the proof in the future version.) \n\n\n\nRegarding questions:\n1. Clients sharing stationary points:\nThe statement \"each cluster clients share stationary points\" means that, despite having different data distributions, the clients within a cluster can simultaneously reach a stationary point. This assumption is indeed a relaxation of the i.i.d. assumption, allowing for more realistic and heterogeneous data distributions.\n3. Performance without clustering assumption:\nOur algorithm still applies when there is only one cluster as it does not require knowledge of the exact number of clusters or balanced clusters. \n\nWe hope that our revisions have  addressed the concerns raised by the reviewer, and increase the score if possible.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "C3m8KoSIFs", "reviewer": "Reviewer_H7Ja", "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.", "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$?\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?", "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.", "rating": 6, "confidence": 2, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "In 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.", "weaknesses": "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.\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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We would like to thank you for your thoughtful reviews and suggestions on our paper. We appreciate your feedback and are happy to address the concerns raised.\n\nRegarding Weaknesses:\n\n**Theorem I proof:** We agree with the reviewer that using full gradients for the inner problem simplifies the proof. In practice, we still use mini-batch SGD as the unbiased estimate of the full gradient. We plan to update the proof in the future to reflect this practical consideration.\n\n**Theoretical convergence performance comparison:** 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.\n\n**Robustness of results:** We acknowledge the importance of repeat experiments for ensuring the robustness of our results. We have repeated the experiments, and the results are available in the global PDF file.\n\n**Regarding Minor Issues:**\nLine 120: We apologize for the mistake and will ensure that the sentence is complete in the final version of the paper.\n\nQuestions:\n\n\n**Output of Algorithm 1:** Non-convex optimization usually tries to prove that the time-averaged gradient norm, such as $\\frac{1}{T} \\sum_{t=0}^{T-1} \\lVert \\nabla f(x^t) \\rVert_2^2$, is upper bounded by a very small value. However, the theoretical results do not guarantee that the last iterate $x^T$ is also  upper bounded. Instead, the $\\frac{1}{T} \\sum_{t=0}^{T-1} \\lVert \\nabla f(x^t) \\rVert_2^2$ can be seen as the expected gradient norm of a uniform randomly drawn $x^t$, therefore this drawn iterate has the upper bound in expectation. This is a well-established technique in the literature (see [Fang et al. 2018] for an example).\nComparison with fully decentralized algorithms: We appreciate your suggestion to compare CoBo with fully decentralized algorithms. However, our baseline federated clustering algorithm is already fully decentralized and does not rely on a central server. In contrast, personalized federated learning with a central server can be seen as a fully-connected decentralized graph, which is a more challenging baseline to beat.\n\n**Influence of sampling rate:** You are correct that the sampling rate of 1/n is an important parameter. We have carefully chosen this value to minimize the pairwise computation overhead while preserving the quality of the solution.\n\na. Complexity: Minimizing outer n objectives for $T$ iterations require computing n * T gradients. Naively computing pairwise gradients for the inner objective requires n^2 T gradients computation which dominate the complexity. With a sampling rate of O(1/n), the overall complexity of the bilevel formulation remain the same as single level optimization. \n\nb. Preserved quality: uniform randomly sampling O(1/n) of the edges do not harm the performance of CoBo. As the underlying connectivity matrix is fixed throughout training,each edge will be sampled $O(T/n)$ times to determine the connectivity. As the number of iterations $T$ is usually much larger than the number of clients, it is enough to determine the connectivity.\n\nFang C, Li C J, Lin Z, et al. Spider: Near-optimal non-convex optimization via stochastic path-integrated differential estimator[J]. Advances in neural information processing systems, 2018, 31.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Ui7lPej37P", "reviewer": "Reviewer_4QtE", "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.", "questions": "1.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?", "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.\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.", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1.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).", "weaknesses": "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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their insightful questions and comments.\n\n**Regarding Question 1**\n\nWe appreciate the reviewer's interest in the step size selection for the projected gradient descent in CoBo. We choose $O(1/LT)$ as the step size, one intuitive explanation is motivated by the similarity between our formulation and the Frank-Wolfe algorithm. As noted in the Frank-Wolfe literature, step sizes of the order $O(1/(t+2))$ are commonly used. \n\n**Regarding Question 2**\n\nWe acknowledge the importance of privacy in federated learning and appreciate the reviewer's concerns. While we do not incorporate specific privacy measures into CoBo in this paper, we note that differential privacy can be combined with CoBo directly. However, achieving output privacy may require additional cryptographic tools. We consider this a promising direction for future work. We hope the reviewer will consider this a natural next step rather than a drawback of our paper.\n\n**Regarding Question 3**\n\nWe are glad to highlight the key differences between Ditto and CoBo, as described in Section 2 (between lines 57 and 62). In scenarios where clients are drawn from diverse distributions, possibly with conflicting labels, CoBo's flexible formulation allows clients to identify fine-grained cluster structures, leading to better local models and improved performance. In contrast, Ditto's formulation may penalize the distance between local models and a meaningless global model, potentially harming local model performance. We believe this example illustrates the advantages of CoBo in handling diverse and complex client distributions.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
We thank the reviewers for their insightful reviews and feedback. **The attached PDF**, contains additional experiments allowing us to gladly address the following concerns:

**1. FedAvg+fine-tuning:** A new baseline is included for all experiments. This baseline is similar to FedAvg for the first 80\% of the training, and then it fine-tunes each client on their local dataset for the remainder of the iterations. Although this method performs better than FedAvg, CoBo's performance is still superior in all experiments.

**2. Ablation study on the Cross-silo experiment:** We have added new experiments to see how CoBo behaves with different experimental setups. As expected, CoBo is sensitive to the total data available to each cluster. We also found out that increasing the number of clusters does not affect the performance of CoBo significantly. We also experimented with the different number of clients per clusters, maintaining the fixed number of 4 clusters. Since the data is partitioned among clients of each cluster, larger cluster size means lower data available to each individual client. However, we are glad that CoBo can also preserve the accuracy for larger clusters, by enabling collaboration within clusters.

**3. Sampling strategies for updating collaboration matrix:** We agree that the sampling method is a very important part of our algorithm. We therefore conduct experiments on different sampling strategies in the Cross-silo experiment. Our proposed methods almost perform similarly, suggesting the robustness of CoBo on its sampling method. We observed that CoBo can find collaborators in the early stages of the training, leading us to propose a new sampling method based on sampling more frequently in the early stages and reduce the frequency in the later stages. We show that this mixed strategy can slightly perform better than other strategies in the Cross-silo experiment.

**4. Repeated Experiments:** We acknowledge the importance of the repeated experiments to examine the robustness of methods, and we thank the reviewers for raising this point. We have added the confidence intervals for the Cross-device experiment as well in order to address this matter.

Meta-review: 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.
Decision: Accept (poster)