Paper: Communication-Efficient Federated Non-Linear Bandit Optimization (ICLR-2024)

Reviewer comment (Experiments):
Comparison with State-of-the-Art: While the paper demonstrates the superiority of Fed-GO-UCB over existing federated bandit algorithms, a more comprehensive comparison with state-of-the-art methods in decentralized machine learning would strengthen the paper's contributions.

Author response:
**[Q2] Comparison with decentralized machine learning methods**

Thanks for the suggestion. We have added a section in Appendix B to discuss existing works in offline federated learning / decentralized machine learning to help further highlight the contribution of our paper.

However, we want to clarify that, due to the fundamental differences in the problem formulation and learning objectives between federated bandit learning and decentralized machine learning, direct comparison in theoretical analysis and experiments is not possible. 
Specifically, the focus of decentralized machine learning is to collaboratively learn a good *point estimate* over a fixed dataset, i.e., convergence to the minimizer with fewer communications/iterations, while federated bandit learning requires collaborative *confidence set estimation* for efficient regret reduction over a finite time horizon $T$.

One perspective to understand their relation is that federated bandit learning studies decision making over the whole time horizon $T$, while decentralized machine learning aims to solve the optimization problem on the dataset collected up to a fixed time step $t \in [T]$, i.e., the latter can be viewed as a component of the former. For example, in our Fed-GO-UCB algorithm, we can adopt any decentralized machine learning method as the *Oracle* to optimize Equation (2), but we still need the other components, e.g., confidence set construction (line 4) and OFU principle (line 3) for effective decision making. Therefore, these two are not directly comparable.

Reviewer comment (Evaluation):
Over-reliance on Communication: The need for occasional communications to aggregate local learning parameters may lead to potential inefficiencies or delays.

Author response:
**[Q3] The need for occasional communications may lead to inefficiencies or delays**

We definitely agree with the reviewer that such communication “may lead to potential inefficiencies or delays”, but as discussed in our response to CQ2, in order to avoid the suboptimal regret $\tilde O(N \sqrt{T})$, communication among clients is necessary.

This is why we need a well-designed federated bandit optimization method like Fed-GO-UCB to attain a good tradeoff between the conflicting objectives of regret minimization and communication efficiency, i.e., we wish to have the most regret reduction per communication. We have shown both theoretically and empirically, that our Fed-GO-UCB algorithm achieved this goal for the federated bandit optimization of nonlinear functions.

Reviewer comment (Theory):
Can Fed-GO-UCB be extended to handle heterogeneous clients with different reward functions, and if so, what modifications would be necessary?

Author response:
**[Q5] Necessary modification to handle heterogeneous clients**


Yes, Fed-GO-UCB can be extended to handle heterogeneous clients.

Thanks to our confidence set construction, methods developed for federated linear bandits can now be readily applied on top of our framework to deal with nonlinear functions. Therefore, to handle heterogeneous clients, one can adopt the clustering methods proposed by Gentile et. al. (2014) and Li et. al. (2021) to cluster the clients into different groups (with clients sharing similar reward functions being in the same group), and then aggregate the statistics of clients in the same group for confidence set construction. Moreover, the communication protocol in our Algorithm 1 can still be used within each group, to reduce communication cost.


A recent work by Liu et. al. (2022) does exactly this for federated linear bandits, via a combination of the clustering method of Gentile et al. (2014) and the asynchronous communication protocol of Li and Wang (2022a). With the confidence sets constructed in our paper, a similar algorithm design as Liu et. al. (2022) can be used for heterogeneous clients with different nonlinear reward functions.


**Additional References**

- Xutong Liu, Haoru Zhao, Tong Yu, Shuai Li, and John CS Lui. "Federated online clustering of bandits." In Uncertainty in Artificial Intelligence, pp. 1221-1231. PMLR, 2022.
- Claudio Gentile, Shuai Li, and Giovanni Zappella. "Online clustering of bandits." In International conference on machine learning, pp. 757-765. PMLR, 2014.
- Chuanhao Li, Qingyun Wu, and Hongning Wang. "Unifying clustered and non-stationary bandits." In International Conference on Artificial Intelligence and Statistics, pp. 1063-1071. PMLR, 2021.

Reviewer comment (Theory):
What are the practical implications of the “mild conditions” under which Fed-GO-UCB operates, and how do these conditions influence the algorithm's applicability in real-world scenarios?

Author response:
**[Q1] Discussion on assumptions and their practical implications**

Please find our answer to CQ1 in the common response for discussions about the assumptions made in our paper. 

Here we want to emphasize that the function class defined by Assumptions 1-3 generalizes the parametric function classes studied in existing federated bandit papers (Wang et. al. 2020; Li and Wang, 2022a; He et. al., 2022; Li and Wang, 2022b), and also covers additional nonlinear, nonconvex functions that have not been considered in existing works in federated bandits. We want to emphasize that this is already a nontrivial improvement.

In terms of applicability in real-world scenarios, it is worth noting that linear and generalized linear models (both are special cases of ours) have already been shown to perform well in applications like news recommendation (Li et. al., 2010), and our generalization offers even more flexibility in the modeling choices during federated optimization, i.e., any parametric functions, including neural networks, that satisfy our local strongly convexity assumption.

Reviewer comment (Novelty):
Novelty: Though the setting is new and the two-phase design is interesting, after the uniform exploration (phase 1), it seems to me that one can combine the techniques from centralized non-linear bandit optimization problem with the communication protocol from the federated linear bandits (Li & Wang, 2022a) and federated generalized linear bandits (Li & Wang, 2022b).

Author response:
**[Q1] Novelty**

As mentioned in the Technical novelties paragraph in Section 1, the approximation method from centralized non-linear bandit optimization (Liu and Wang, 2023) cannot be applied to the federated setting. Specifically, the confidence sets for the optimal parameter $\omega^\star$ constructed in (Liu and Wang, 2023) is $\{\omega: \lVert \omega - \hat \omega_t \rVert_{\Sigma_{t}} \leq \beta_{t}\}$, where $\Sigma_t =\lambda \mathbf{I} + \sum_{s=1}^{t} \nabla f_{x_s}(\hat \omega_{s}) \nabla f_{ x_s}(\hat \omega_{s})^\top$ and $b_t =\sum_{s=1}^{t} \nabla f_{x_s}(\hat \omega_{s}) \bigl[\nabla f_{ x_s}(\hat \omega_{s})^\top \hat \omega_{s} + y_s - f_{x_s}(\hat \omega_{s}) \bigr]$, where the gradients $\nabla f_{x_s}(\hat \omega_{s})$ have to be calculated w.r.t. the model $\hat \omega_{s}$ at step $s$.

Due to local model updates in federated settings, now each client $i$ has a different sequence of local models $\hat \omega_{s,i}$ for $s=1,2,\dots,t$. By directly aggregating the statistics $\Sigma_{t,i}$ and $b_{t,i}$ computed based on such $N$ different sequences of local models, one cannot obtain a valid confidence set. To address this issue, we improve the analysis of Liu and Wang (2023), and show that we can construct a valid confidence set using the gradients w.r.t. a common initial model $\hat \omega_0$, instead of clients’ locally updated ones. 

Therefore, after obtaining $\hat \omega_0$ at the end of Phase I, the clients can collaborate on confidence set estimation via a common embedding function defined by the gradient $\nabla f_{x_s}(\hat \omega_{0})$. This is why we are able to utilize the communication protocol designed for federated linear bandits (Wang et. al. 2020). We want to emphasize that being able to utilize communication protocols designed for federated linear bandits for nonlinear functions via our novel confidence set construction is already non-trivial and may be of independent interest for future research. For example, one can also extend our method to asynchronous protocols (Li and Wang, 2022a; He et. al, 2022), as well as heterogeneous clients setting (Liu et. al. 2022). Moreover, as discussed in our response to CQ2, this new algorithm design helps us improve upon FedGLB-UCB (Li and Wang, 2022b) by a factor of $d_x \sqrt{N}$ for federated optimization of nonlinear functions, because in Phase II we are able to use the efficient closed-form update as federated linear bandits, instead of the expensive iterative optimization as FedGLB-UCB.

Reviewer comment (Theory):
Comparison with federated generalized linear bandits (Li & Wang, 2022b): From the algorithmic perspective, Li & Wang 2022b also uses some global updates (similar to uniform exploration) after the communication condition (line 6 of Algorithm 1) is satisfied, I am wondering if the current paper also uses this way to update $\hat{w}_0$, will there be any difference or improvement in the regret analysis? Moreover, I find that Li & Wang, 2022b can only achieve $O(dN^2\sqrt{T})$ communication. Since I suppose the current paper is more general and can cover Li & Wang 2022b, does this mean that the current paper can achieve a $O(\sqrt{N})$ improvement because of the difference in the algorithm design?

Author response:
**[Q2] Improvement compared with federated generalized linear bandits (Li & Wang, 2022b)**

Yes, our paper has an $\tilde O(d_x \sqrt{N})$ improvement on communication cost thanks to the new algorithm design. Please see our answer to CQ2 in the general response for more details about how this is achieved.

Reviewer comment (Theory):
Lower bound: I do not find any discussion about the lower bound result or any discussion about it. Without it, one cannot see how tight the results are in terms of $$N, T$$ and all other parameters. I hope the authors can discuss the tightness of their results during the rebuttal.

Author response:
**[Q3] Tightness of communication cost upper bound**

Please see our answer to CQ2 in the common response for discussions on the tightness of our communication cost upper bound in terms of $N$ and $T$.

Reviewer comment (Novelty):
Though this work features considering non-linear function optimization, it seems that its main contribution is on addressing the challenges of federated setting. Could the authors discuss more about the special challenges of non-linear function optimization?

Author response:
**[Q2] Challenges of nonlinear function optimization**

We want to clarify that the unique challenge addressed in this paper is bandit optimization of **nonlinear functions under federated settings**. We cannot view them separately as suggested by the reviewer. Specifically, federated setting requires communication efficiency, and bandit optimization of nonlinear functions requires construction of confidence sets. Therefore, the unique question answered in our paper is, how to construct confidence sets for the nonlinear functions collaboratively by all clients in a communication efficient manner, which cannot be addressed by prior works. 


As mentioned in the Technical novelties paragraph in Section 1 of our paper, the approximation method we propose is a non-trivial extension of (Liu and Wang, 2023), i.e., their method cannot be applied to federated settings as their approximation relies on a sequence of continuously updated models $\hat \omega_{s}$ for $s=1,2,\dots,t$. But in federated settings, each client has a different sequence of locally updated models, which leads to difficulty in aggregating local statistics (see also our answer to Q1 of Reviewer yBxF). This problem does not exist for federated linear bandits, since the confidence set for linear models can be directly constructed using feature vector $x$, instead of gradient $\nabla f_{ x_s}(\hat{\omega}_{s})$.


To address this issue, we propose a new approximation procedure with improved analysis over Liu and Wang (2023), such that the statistics computed by all clients are now based on a common embedding function $\nabla f_{x_s}(\hat{\omega}_{0})$, i.e., the gradient w.r.t. the shared model $\hat \omega_0$. Note that $\hat \omega_0$ is obtained at the end of Phase I and then always remains fixed.


Compared with prior works in federated linear bandits, which conveniently enjoys closed-form solutions, nonlinear function faces another challenge, i.e., it requires iterative optimization for each model update, which is expensive in federated settings. As mentioned in our answer to CQ2, prior work that studies generalized linear models [Li and Wang 2022b] needs to call the distributed regression oracle during each global synchronization to execute such an expensive iterative optimization procedure. We alleviate this issue by only executing the iterative optimization at the end of Phase I, and then resorting to approximated function value during Phase II, so that more efficient communication is enabled thanks to its closed-form update.

Reviewer comment (Theory):
Could the authors comment on the tightness of the communication cost bound?

Author response:
**[Q4] Tightness of the communication cost upper bound**

Please see CQ2 in the common response for discussions on tightness of the communication cost bound.

Reviewer comment (Theory):
The theoretical guarantee of Fed-GO-UCB relies on several critical assumptions (Assumptions 1 to 3), and it may be challenging to establish whether these assumptions hold in common situations, such as neural networks with ReLU activation functions. Even if the assumptions hold, determining the values of the parameters in Fed-GO-UCB that depend on these assumptions can be a non-trivial task. It remains an open question how to calculate these constants and set the algorithm's parameters effectively in practical applications.

Author response:
**[Q1] Whether the assumptions hold in common situations and how to determine the values of the parameters**
 
Please find our answer to CQ1 in the common response for discussions about the assumptions made in our paper. 

Here we want to emphasize that, though still with its limitations, the function class defined by Assumptions 1-3 generalizes the parametric function classes studied in existing federated bandit papers (Wang et. al. 2020; Li and Wang, 2022a; He et. al., 2022; Li and Wang, 2022b), and also covers additional nonlinear, nonconvex functions that have not been considered in prior works. It is worth noting that linear and generalized linear models (both are special cases of ours) have already been shown to perform well in real-world applications like news recommendation (Li et. al., 2010), and our generalization offers even more flexibility in the modeling choices during federated optimization.

Moreover, we want to clarify that, Fed-GO-UCB algorithm does not need to know the values of the constants appearing in Assumptions 1-3, except for the bound of function value $F$ and the strong convexity parameter $\mu$ in order to set the width of the confidence set $\beta$ (as given in Lemma 8). Both constants are assumed to be known in bandit learning literature, e.g. linear bandits (Abbasi-yadkori, 2011; Wang et. al., 2020), logistic bandits (Faury et. al., 2020), and generalized linear bandits (Filippi et. al., 2010; Li and Wang, 2022b), so we are not introducing any additional assumption on the knowledge of these constants compared with prior works. 

In practical applications, since these two constants, and the standard deviation $\sigma$ of reward noise that also appears in $\beta$, are unknown before observing the data, people typically directly tune the value of $\beta$ when applying bandit algorithms under different scenarios. But we agree that it would be an interesting future direction to investigate whether one can design no-regret bandit algorithms that do not require such knowledge as input, i.e., automatically adapt the strength of exploration to these unknown constants in an online manner.

Reviewer comment (Theory):
The communication complexity, as indicated by Theorem 5, is stated as $O(\sqrt{T})$, which is significantly higher than the communication complexity of $O(\log T)$ achieved by previous works. Presenting this level of communication cost as efficient without providing lower bound results for communication complexity can be misleading. Furthermore, the experimental results in Figure 2 show a communication complexity of $10^7$ for round $T=500$, which seems highly inefficient in practice.

Author response:
**[Q2] Tightness of communication cost upper bound and comparison with prior works**

Please find our answer to CQ2 in the common response about the communication lower bound and detailed comparison with prior works. 

Specifically, we want to clarify that, the mentioned $O(\log(T))$ rate is attained by prior works on federated linear bandits (Wang et. al., 2020; Li and Wang, 2022a; He et al., 2022; Min et al. 2023), and federated kernelized bandits (Li et al., 2022; Li et al., 2023), where closed-form solutions are available, so that only one round of communication is needed to aggregate the local sufficient statistics for model update. 

In comparison, our work is more similar to (Li and Wang, 2022b) that studies federated generalized linear bandits, in the sense that, both consider nonlinear functions with no closed-form solution, and thus iterative optimization is needed to compute model update. This is intrinsically more expensive, because at least $\Omega(\sqrt{NT})$ rounds of communication is needed to collect and aggregate local gradients/models when calling the distributed regression oracle (see the discussions in Section 4.1 and 4.2 of Li and Wang (2022b), as well as the lower bound result of Arjevani and Shamir (2015)). Moreover, as discussed in our answer to CQ2, we strictly improve upon the result of Li and Wang (2022b) by a factor of $d_x \sqrt{N}$, since we only need to call the regression oracle once at the end of Phase I, instead of at each global synchronization as Li and Wang (2022b). As shown in the table in CQ2, Fed-GO-UCB is currently the most communication efficient solution for federated bandit optimization of nonlinear function, which is already a nontrivial contribution to the literature.

Reviewer comment (Experiments):
In the context of federated learning, there is often a tradeoff between communication complexity and performance, specifically in terms of regret guarantee. It's a common observation that algorithms with higher communication complexity can potentially achieve better performance. To gain a clearer understanding of the algorithms' relative performances, it would be beneficial for the author to conduct experiments where the selected parameters makes different algorithms have similar communication complexity. This approach would enable a direct and fair comparison that isolates the impact of communication complexity from the approximation capabilities, providing more conclusive insights into the algorithm's efficiency and effectiveness in a practical federated learning setting.

Author response:
**[Q3] Empirical comparison of communication efficiency**

Here we want to clarify a potential misunderstanding in the `fair comparison’ suggested by Reviewer rFun. 
We definitely agree that higher communication cost leads to better (at least no worse) performance, but this monotonicity is only true when we fix the function class unchanged. Comparison across different function classes introduce additional intricacies. 

Intuitively, a function that involves more parameters requires more samples to estimate, e.g., an overparameterized neural network vs a linear regression model. Therefore, a federated bandit algorithm using the overparameterized neural network may perform worse than that with a simple linear model, when there is not enough communication, or when the dataset is `easy’, i.e., linear function already fits it pretty well. However, this does not mean the former is a worse algorithm than the latter.

The proper choice of function class depends on the dataset (i.e., the underlying reward mapping), as well as the communication bandwidth of the application at hand. The advantage of our Fed-GO-UCB over prior works, e.g., DisLinUCB and FedGLB-UCB, is that **it offers a larger variety of modeling choices**, so that it would be easier to find a suitable function class for different application scenarios. In general, for datasets that can be well-fitted by a linear model, DisLinUCB enjoys the advantage of its closed-form model updates, which is much more communication-efficient than gradient-based updates. However, in the more common situations where linear models give poor performance and one has to employ more complicated models at the cost of communication efficiency (because closed-form solution no longer exists), DisLinUCB no longer applies, while Fed-GO-UCB addresses this in a provably efficient manner.

To justify this argument, we have **conducted additional experiments that add One-LinUCB** (DisLinUCB algorithm with communication threshold set to $-\infty$, so that communication happens in each time step) into comparison. Please see the updated Figure 2 and Figure 3 in our paper for more details. From the experiment results, we can see that, though with maximum communication possible, One-LinUCB is still outperformed by Fed-GO-UCB by a large margin in terms of cumulative regret, which clearly shows the limited capability of linear models in fitting the dataset. In this case, more complex models are preferred, and our Fed-GO-UCB algorithm offers a principled way to balance regret and communication.

Reviewer comment (Reproducibility):
In the Fed-GO-UCB algorithm, it's important to clarify that agents should reset $\Delta\Sigma$ and $\Delta b$ to zero after uploading their respective datasets.

Author response:
**[Q4] Reset $\Delta \Sigma$ and $\Delta b$ to zero after uploading**

Thanks for the suggestion. We have updated the description of Algorithm 1 to clarify this.