Efficient Adaptive Federated Optimization

We thank reviewer gLqi for their feedback.

[Missing Appendix?] In the original submission, we manually split the main paper and the appendix, and submitted the full paper along with the appendix and source code as a single supplementary zip file, per ICLR guidelines. We understand that this has led to the impression of a missing appendix, and would like to gently ask for a re-evaluation of the work given this context. In the revision, we directly uploaded the full paper as well as the appendix without splitting them for reviewers’ convenience.

[Convexity Assumptions] We use convex, online objectives in Section 3 only for motivation –but for the convergence of the proposed algorithm $FedAda^2$, we use the more general L-smoothness setting in Section 5.

[Table of $FedAda^2$ and Baselines + Ablations] The comparison table is given in Appendix G.5, with additional discussions. Regarding ablations, $FedAda^2$ and baselines can be viewed as natural extensions or ablations of one another, which incorporate local adaptive updates and memory-efficient strategies. That is, Direct Joint Adaptivity refers to a training paradigm where the server’s adaptive preconditioners are shared with clients during each communication round. Eliminating the transmission of server-side preconditioners and initializing client side preconditioners to zero gives ‘Joint Adaptivity without Preconditioner Communication’, which is more communication-efficient. Further compressing local preconditioners (e.g., SM3) to align with client memory constraints leads to the development of $FedAda^2$.

[Questions.]

[Q1. Empirical Setup] The empirical setup, including client count, synchronization rounds, and hyperparameter selection, is detailed in Appendix G, H (e.g., Table 1). Generally, we chose the device number to be large in relation to the compute we had available. For Stackoverflow (G.1), there are 400 devices, where each device is a site user. For CIFAR-100 (G.3), we used 1000 Devices, where highly non-IID client data distributions were induced via LDA partitioning (a topic modeling method [1]) with $\alpha = 0.001$. GLD-23K (G.2) consists of 233 real-world photographers, taking pictures of landmarks. The participation fractions were quite low, e.g., 0.01 for GLD-23K (only 2 clients per round) to model challenging, real-world non-IID settings.

[Q2. $FedAda^2$ and Baseline Comparisons] The communication/memory comparisons between $FedAda^2$ and other baselines are provided in Appendix G.5 in the original submission. For convenience, we also provide it here.

Table 1: Comparison of Baselines versus $FedAda^2$ with AdaGrad Instantiations.

For the ViT model for instance, we require just 0.48% memory to store the second moment EMA compared to the full gradient statistic during preconditioning, when SM3 is used. The variance between 1d and 2d in the $FedAda^2$ “Memory (client)” column depends on the instantiation of the client-side memory-efficient optimizer. Here, $d$ denotes the model dimensions.

[1] Blei et al., Latent dirichlet allocation. JMLR, 2003.

[Divergence of Client-Side Adaptivity in [1]] Firstly, the reviewer notes that as $FedAda^2$ uses local adaptivity, “computing a simple average at the global server may encounter divergence issues.” But in fact, FedAda^2 is performing globally adaptive updates, as opposed to “computing a simple average” at the server as the reviewer states, as described in Algorithm 1 (step 12). Hence, the issues with the client-only adaptive optimization the reviewer mentions do not necessarily hold in our setting. Our experiment results (Figures 1~3) in Section 6 validate the performance of $FedAda^2$, and prior works (e.g., [6]) have also shown the performance enhancement of client-side adaptive optimization.

[Experimental Improvements] Our goal is to show that joint adaptivity is practicable in real-world cross-device environments via efficient techniques. Therefore, we do not expect memory-efficient optimization to do even better than full-information optimization. The goal of this work is not to achieve the highest accuracy on all datasets, rather, it is to show that the communication- and memory-efficient $FedAda^2$ framework (a) converges theoretically, and (b) does not degrade the performance of expensive jointly adaptive optimization despite information compression. However, despite using approximation techniques, we found that $FedAda^2$ is (a) competitive compared with much more expensive baselines, and (b) sometimes does even better than the more expensive jointly adaptive baseline, "Direct Joint Adaptivity" (e.g., DP-Stackoverflow, Figure 1). When we report accuracies versus bits communicated, $FedAda^2$ widens its lead (Figure 2).

[6] Zhou et al., FedCAda: Adaptive Client-Side Optimization for Accelerated and Stable Federated Learning. Arxiv, 2024

We thank reviewer 3XLN for their review.

[Inaccurate Statements in Introduction] We respectfully disagree with the reviewer’s statement. The “explanatory context” the reviewer requests is given in Section 1, lines 033-055, Section 6, lines 374-387, etc. To summarize, there is a server-side optimizer and client-side optimizer in federated learning. Adaptivity can be employed in either the server-side or the client-side, where joint adaptivity (consisting of global and local adaptive updates) has been shown to play a pivotal role in accelerating convergence and enhancing accuracy in prior works. However, to our knowledge, there are no known convergence results in the jointly adaptive setup. The vast majority of work, including the work by Reddi et al that the reviewer cites, include proofs for server or client-only adaptive setups, e.g., FedAdam, FedAdaGrad, which we have implemented as baselines. In joint adaptivity settings, significant complications arise from the interaction between server- and client-side adaptive updates in the proof, which both leverages historical local gradient information. We would be happy to elaborate further upon request.

[On “Conflicting” Lipschitz Assumption] We respectfully disagree that the assumptions the reviewer refers to are conflicting, i.e., the Lipschitzness of the objectives (Section 5) do not conflict with the heavy-tailedness of the stochastic gradients (Section 3). This is because the Lipschitzness is imposed on the expected objective, which is defined as $F_i(x)=\mathbb{E}_{z \sim \mathcal{D}_i}\left[F_i(x, z)\right]$ (line 202). The heavy-tailedness of the gradients are modeled by $z$, which is integrated out. Therefore, an objective satisfying Lipschitzness can indeed “experience divergence during training” due to heavy tailed noise, as $z \sim \mathcal{D}$ is orthogonal to $F_i(x)$ after the expectation. Furthermore, Section 3 was included in order to theoretically study intriguing questions regarding the desirability of adaptive client-side optimization, something another reviewer has asked about. In contrast, the analysis in Section 5 aims to establish standard convergence results in the context of bounded gradients and L-smoothness as standard assumptions (e.g., [2-4]).

[Using Local Historical Preconditioner as in [1]] Using local historical preconditioners throughout training is not a viable approach for several real-world, cross-device settings, as noted in Section 4, lines 206-208. One determining property of cross-device federated settings is that the clients are not able to store or maintain ‘states’ across communication rounds [5]. For example, it is unclear that a client will participate in the training more than once, and it is unreasonable for resource-strapped clients (e.g., mobile phones) to maintain optimizer states in memory when not partaking in local updates. Additionally, server-side adaptivity is critical. Thus, we propose the $FedAda^2$ framework for communication and memory-efficient optimization, with rigorous convergence guarantees. We respectfully disagree with reviewer 3XLN that our work is “insignificant”, given the practical nature of this problem and the lack of theoretical convergence results in joint adaptive settings.

[2] Xie et al., Local adaalter: Communication efficient stochastic gradient descent with adaptive learning rates. Workshop on Optimization for Machine Learning, 2020.

[3] Wang et al., Tackling the objective inconsistency problem in heterogeneous federated optimization. NeurIPS, 2020.

[4] Li et al., On the convergence of FedAvg on non-iid data. ICLR, 2020.

[5] Kairouz et al., Advances and open problems in federated learning. Foundations and trends in machine learning, 2021.

Thanks for the response and I believe my question has not been effectively resolved for the following perspectives.

About the assumption: I greatly appreciate your analysis in sec.3, as it establishes the necessity of adaptive learning rates under label-inconsistent heterogeneity. If you could provide a convergence analysis under a long-tailed data distribution and highlight the advantages of using a dual-adaptive optimizer, I believe it would be a significant contribution. Currently, the theory can only be seen as an extension of FedAdam, and the convergence rate you’ve proven cannot surpass that of FedAdam because it already represents the optimal rate under general assumptions.

About the assumption conflicts: Authors use the assumption 2, that the gradient coordinates are bounded. It means the stochastic gradient $g$ is always bounded on each sample $z$, i.e. $\vert \nabla [F_i(x,z)(j)] \vert \leq G$. Therefore, $\Vert g \Vert^2 = \sum_{j=1}^d g_j^2 \leq dG^2$ is bounded. Then we consider $\mathbb{E}\Vert g - \nabla F(x) \Vert^\alpha$. We can take a sample as $\alpha=2$. $\mathbb{E}\Vert g - \nabla F(x) \Vert^2 = \mathbb{E}\Vert g\Vert^2 + \mathbb{E}\Vert \nabla F(x)\Vert^2 - 2\mathbb{E}\langle g, \nabla F(x)\rangle = \mathbb{E}\Vert g\Vert^2 - \Vert \nabla F(x)\Vert^2 \leq \mathbb{E}\Vert g\Vert^2 \leq dG^2$. In fact, it is bounded for all constant $\alpha$. This is the question I raised. These two assumptions seem to conflict. Can the author help me answer this question?

About the Zero Local Preconditioner Initialization: I emphasize that it is trivial because past work has already known relevant techniques. As mentioned above, [1] proposed that a local second-order momentum could be re-used. [2] proposed using a constant pre-conditioner as the adaptivity state at each initialization of local rounds (details seen in page 3 in their paper). Obviously, authors only set the constant in [2] to 0 and restart the local optimization state, and there is nothing new about it. Given that the adaptation is similar to the diagonalization estimation of the Fisher matrix, I do not think that the zero initialization in this paper can be regarded as an novel technique.

[1] Toward communication efficient adaptive gradient method

[2] Local Adaptivity in Federated Learning: Convergence and Consistency

Divergence of the naive global average: The problem pointed out by [1] is that if naive averaging is used, the merged direction $\Delta$ may be opposite to the gradient. which will cause training divergence. The author's answer is that adaptive correction is used globally, so this problem can be avoided. I would like to ask the author a question: can authors use the reverse gradient to train models? When the direction of the merged $\Delta$ is opposite to the gradient, is it still effective to use adaptive learning rate? My intuition for asking this question comes from the fact that I think the answer to both of the above questions is no. If the author believes that adaptation can correct the wrong direction, can provide me with more intuition or examples to explain why it can work. I am also very interested in this issue and happy to discuss it.

I understand the rate of this iclr are generally low, but I hope the author can read my questions carefully and then reply. I will seriously reply to your rebuttal during the whole discussion phase. At present, I think the concerns yet have not been well solved.

[Comparisons Between Our Work and [1]] The reviewer cites [1] closely, and we compare [1] with ours.

(a) We remove the need for preconditioner maintenance on the client-side, in contrast to [1]. For example, in our experiments for the FEMNIST dataset, only 2 writers out of the 500 clients participate in training each round, which is done for 131 global synchronization rounds until convergence. It is unlikely that any given client will participate more than a few times in such realistic, non-IID settings, thus it is not reasonable for clients to to maintain preconditioners when not directly participating in updates.

(b) We show convergence for general joint adaptive methods with local updates, and we remove the aggregation of preconditioners on the server-side for communication efficiency. The algorithm in [1] must transmit preconditioners for aggregation and incurs significant communication complexity, and additionally, [1] does not use joint adaptivity.

(c) We incorporate local memory-efficient compression, and study the more general cross-device setting in which only a fraction of the clients participate.

[Benefits of Adaptive Optimizers on the Client-Side] The reviewer has expressed concerns about naive local AMSGrad (an algorithm proposed in [1]), because the local preconditioners may lead to divergence. We agree this can be true in specifically constructed negative examples (e.g., using AMSGrad locally and setting $\beta_1=0$, $\beta_2 = 0.5$ with just 1 client update, and choosing an initial condition adversarial to a particularly designed 1-dimensional objective as in [1]). In light of such potential challenges of local adaptivity, our goal is to come up with a theoretical framework that avoids divergence while enjoying the benefits of joint adaptivity, in an efficient way. We establish convergence results of $FedAda^2$ in Section 5, which also covers scenarios in which participating clients use different local optimizer instantiations (hence the name “Blended Optimization” in Appendix D), which would help to mitigate adversarially constructed negative examples. Our work is different to the negative algorithm presented in [1], and one cannot directly attribute the divergence to $FedAda^2$. Our experiment results (Figures 1~3) in Section 6 validate the effectiveness of $FedAda^2$, and prior works (e.g., [6]) have also shown the performance enhancement of client-side adaptive optimization.

We welcome any additional questions or feedback.

[1] Chen et al., Toward communication efficient adaptive gradient method. Proceedings of the 2020 ACM-IMS on Foundations of Data Science Conference, 2020.

[2] Wang et al., Local Adaptivity in Federated Learning: Convergence and Consistency. Arxiv, 2021.

[3] Xie et al., Local adaalter: Communication efficient stochastic gradient descent with adaptive learning rates. Workshop on Optimization for Machine Learning, 2020.

[4] Wang et al., Tackling the objective inconsistency problem in heterogeneous federated optimization. NeurIPS, 2020.

[5] Kairouz et al., Advances and open problems in federated learning. Foundations and trends in machine learning, 2021.

[6] Zhou et al., FedCAda: Adaptive Client-Side Optimization for Accelerated and Stable Federated Learning. Arxiv, 2024

Thank you for your questions.

[Extending from FedAdam] We appreciate the reviewer’s enthusiasm on Section 3. We also respectfully argue that the theory cannot be seen as a simple extension of FedAdam, as FedAdam employs SGD on the client-side. Moving to the jointly adaptive client setting introduces significant theoretical challenges, as now expectations do not act linearly on the client gradients due to dependencies on the entire historical gradients of the client, rendering existing proof techniques inapplicable. Additionally, we incorporate memory-efficient local adaptive optimizers in the proof, such as SM3 (e.g., lines 1235-1347), which is significantly different from local SGD that FedAdam uses. The reviewer enquires about a proof of convergence of joint adaptivity under heavy-tailed noise; a step toward being able to achieve this is to first prove the convergence of joint adaptivity under the conventional L-smoothness (Assumption 1) and Bounded Gradients (Assumption 2) assumptions, which has not been done before and is by itself a significant complication in the convergence proof. Extending the proof to infinite-variance heavy-tailed distributions is beyond the scope of this paper, an entirely new paper altogether.

[About the Assumption Conflicts] In Section 3, we study infinite-variance heavy-tailed noise to model transformer training dynamics and elucidate the theoretical importance of client-side adaptivity. In Section 5, we assume bounded gradients (Assumption 2) to provide standard convergence results. It is straightforward to see that bounded gradient distributions cannot have infinite variance, and these two obviously “conflict”. But our understanding was that the reviewer’s original question was about the Lipschitzness of the gradients of the expected objective $F_i(x)$ (Assumption 1), which does not conflict with the heavy-tailed assumption. As noted in our response above, Section 3 was included in order to theoretically study intriguing questions regarding the desirability of adaptive client-side optimization, something another reviewer has asked about. In contrast, the analysis in Section 5 aims to establish standard convergence results in the context of bounded gradients and L-smoothness as standard assumptions (e.g., [3-4]), which in our view are essential to theoretically establishing the practicability of a proposed framework such as $FedAda^2$. Therefore, Sections 3 and 5 clearly serve different purposes, as stated in our paper. We hope the confusion has been clarified.

[About Zero Local Preconditioner Initialization] We respectfully disagree that our paper is "trivial" for the following reasons.

(a) The work [1] proposes that local preconditioners can be re-used and accumulated across communication rounds. As mentioned in our rebuttal above to reviewer 3XLN, maintaining local historical preconditioners throughout training is not a viable approach for several real-world, cross-device settings, as noted in Section 4, lines 206-208. One determining property of cross-device federated settings is that the clients are not able to store or maintain ‘states’ across communication rounds [5]. For example, it is unclear that a client will participate in the training more than once, and thus it is unreasonable for resource-strapped clients (e.g., mobile phones) to maintain optimizer states in memory when not partaking in local updates. Additionally, [1] studies only the cross-silo setting, where the maintenance of client preconditioners is not an issue. Therefore, we respectfully would like to argue that [1] does not render our work “trivial” or “insignificant”.

(b) We are aware that the paper [2], which we have cited, empirically uses zero preconditioner initialization. However, [2] does not have a proof for convergence of jointly adaptive setups under this setting. In Theorems 6, 20, and 25, we not only provide convergence bounds for server-adaptive and client SM3-adaptive setups under zero preconditioner initialization, but we also extend these results to provide a rigorous convergence guarantee for the $FedAda^2$ framework in the fullest generality, including a class of adaptive methods and using client-specific optimizers (Appendix D, Theorem 25).

(c) Additionally, zero preconditioner initialization is only one of the multiple techniques we utilize for $FedAda^2$. Although we have relegated details to the Appendix for interested readers due its technicality, the elements of SM3 (e.g., lines 1235-1347) and delayed preconditioner updates (e.g., lines 1269-1273, 1300-1361, and Appendix I.2) needed to be incorporated into the proof, which was non-trivial.

[Does Adaptive Optimization “Drop” Faster?] Adaptive optimizers, specifically for training transformer-based models, indeed demonstrate enhanced performance (e.g., [1]). This is why we have included Section 3 to analyze this theoretically in the distributed training setting. For differences between using “SGD... for more rounds” or adaptive optimizers, we refer to Section 3 (lines 113-191), which explains this extensively. For more details, please refer to Appendix A, lines 702-1230 where a self-contained exposition is provided. For experimental results confirming the benefits of using adaptive client-side optimizers, we point to Figures 1~3, c.f., comparison of jointly adaptive paradigms versus baselines.

[Novelty of Approach] We note that our work targets joint adaptivity on both the server side and client side. It is motivated by the fact that joint adaptivity may require additional communication cost and “simply using adaptive methods on the client side without preconditioner transmission” will induce significant memory limitations in resource-constrained clients (Section 1, lines 041-055, Appendix G). To mitigate these issues, we propose to simply avoid preconditioner communication, use various techniques for on-device compute/memory-compression, and successfully provide a novel convergence guarantee for such an efficient framework ($FedAda^2$) in the fullest generality. Empirically, we show better tradeoffs between systems efficiency (communication and memory costs) and model utilities.

[Quantitative Comparisons with SGD & Baselines] We refer to Figures 1~3 in the main text, where FedAvg (distributed local SGD) has indeed been included as a baseline. In all setups, FedAvg does far poorly in terms of convergence speed and final performance than its adaptive counterparts. For the theoretical analysis behind why this may be the case, it is for this reason that we have included Section 3 that the reviewer asked about. We also give an exposition in terms of communication costs/memory advantage of $FedAda^2$ versus baselines in various portions throughout the text, e.g., Appendix G.5, Section 6, lines 374-387, etc. We also refer to Figure 2 (main text), which plots accuracies versus bits communicated.

We welcome any additional questions or requests for clarifications.

[1] Zhou et al., FedCAda: Adaptive Client-Side Optimization for Accelerated and Stable Federated Learning. Arxiv, 2024

We thank reviewer EvCH for their review.

[Impact of Section 3] Section 3 is included to answer the other questions the reviewer has asked, such as: Why don’t we just use “a simple SGD on the client side for more rounds of training”? As discussed in Section 3 (lines 114-191), the discussion on heavy-tails justifies the importance of client-side adaptive optimization. In summary, we use Section 3 to motivate the studied problem (joint adaptivity), not the proposed solution. After analyzing the importance of client-side adaptivity, we propose efficient FL frameworks to mitigate the heightened resources induced by adaptive local optimizers, which is $FedAda^2$. We have clarified this further in the revision.

[On Core Contributions]

• Necessity of convergence analysis: We respectfully disagree with the reviewer, and argue that proving convergence via a rigorous convergence analysis (Section 5 and Appendix B~D) is both challenging and essential to theoretically establishing the practicability of a proposed framework.

• Reduction in communication and memory: We do analyze how communication and memory is reduced relative to other algorithms both theoretically and empirically. For example, in Appendix G.5, we provide a table summarizing communication, computation, and memory complexity across $FedAda^2$ and baselines, with detailed discussions. We provide the table below for convenience. We also refer to Figure 2 (c.f., performance of Direct Joint Adaptivity in Figure 1), which plots accuracies versus bits communicated. There, empirical results show $FedAda^2$ attains significant performance advantages in accuracy and convergence. For additional details, we also provide extensive discussions in Appendix B, lines 1232-1288, Appendix C, 1300-1345, and Section 6, 374-387, etc.

Table 1: Comparison of Baselines versus $FedAda^2$ with AdaGrad Instantiations.

For instance, for the ViT model, when using SM3, we require only 0.48% more memory to store the second moment estimates for preconditioning, as opposed to 1x more memory without SM3. The variance between 1d and 2d in the $FedAda^2$ “Memory (client)” column depends on the instantiation of the client-side memory-efficient optimizer and model architectures. Here, $d$ denotes the model dimensions.

[Difference Between $FedAda^2$ & Adaptive Methods on Client-Side] We have discussed these differences in Section 6, Appendix B, etc of the original submission. To summarize, ‘Direct Joint Adaptivity’ refers to a training paradigm where the server’s adaptive preconditioners are shared with clients for adaptive optimization during each communication round. By eliminating the transmission of server-side preconditioners and initializing client-side preconditioners to zero, we derive the ’Joint Adaptivity without Preconditioner Communication’ baseline. Further compressing local preconditioners to align with limited client memory constraints leads to the development of $FedAda^2$. Therefore, our framework is quite naturally motivated. However, the convergence analysis (Theorems 6, 20) is very challenging due to the generality of $FedAda^2$. We present convergence proofs for both the SM3-specific instantiation, as well as the fully general version, in Appendix C, D.

Dear reviewer EvCH,

Thank you for your time in writing your review. We are wondering if our responses have addressed the reviewer's concerns or answered clarification questions. We are happy to provide additional information and discussions.

(In case the Appendix was missed during the initial review, we would like to note that it was submitted as supplementary materials before, and now it has been re-uploaded together with the main submission.)

Thank you.

We thank reviewer 2y9A for their review, and for their constructive comments. We would like to mention two quick notes of clarification:

[Comparison with FedOPT, MIME] We have incorporated instantiations of FedOPT into our baselines (such as FedAdam, FedAdaGrad), and evaluated their performance against $FedAda^2$ in Section 6. Regarding MIME, their algorithm is not using joint adaptivity and can be significantly more expensive in terms of both communication and on-device memory. Thus we opted for the “Direct Joint Adaptivity” baseline instead. We refer to Table 2 in Appendix G.5 for more detailed explanation and comparisons.

[More Datasets from Reddi et al.,] Thanks for the suggestion. Our experiments do share some datasets with Adaptive Federated Optimization, Reddi et al., e.g., CIFAR-100, StackOverflow LR (specifically, these are in Figures 1, 3 in their work). We also use more challenging models or datasets such as ViT models and the GLD-23K dataset. GLD-23K is a real-world 203-class dataset where each client is a real-life photographer, and thus represents a more interesting task than CIFAR-10/100, or EMNIST classification (62-class). But it would indeed be interesting to see how $FedAda^2$ compares against datasets such as EMNIST as well. We have been able to implement these experiments on $FedAda^2$ and baselines, and we summarize their results in a table below. Convergence figures are added to the revision in Appendix section G.4 (highlighted in blue).

In summary, we observe that jointly adaptive approaches (Direct Joint Adap., Joint Adap. w/o Prec. Trans., and $FedAda^2$) outperform non-jointly adaptive methods by ~4% on the EMNIST dataset. Additionally, removing preconditioner transmission as well as compressing client-side gradient statistics for efficiency ($FedAda^2$) does not degrade the performance of direct joint adaptivity, while being more efficient. We refer to Appendix G.4 in Page 50 of the revision for more detailed discussions.

Thank you for your feedback.