Adaptive Federated Optimization Across Smooth and Nonsmooth Regimes

Zimeng Wang; Ali Dadras; Alp Yurtsever

math.OCOct 29, 2026

Adaptive Federated Optimization Across Smooth and Nonsmooth Regimes

Authors: Zimeng Wang, Ali Dadras, Alp Yurtsever

Organizations: Department of Mathematics and Mathematical Statistics, Ume˚a University, Sweden

Abstract

AdaGrad-style methods can simultaneously adapt to both the order and constants of smoothness in centralized optimization. However, existing extensions to federated optimization with heterogeneous client objectives and multiple local steps require the smoothness order to be specified in advance. We address this limitation with AdaFed, a federated method that adapts across several regimes under a single choice of parameters. In the deterministic setting, AdaFed achieves an $O(1/T)$ convergence rate for both smooth nonconvex and smooth convex objectives, and an $O(1/\sqrt{T})$ rate for Lipschitz nonsmooth convex objectives up to logarithmic factors. Under a light-tailed stochastic-gradient noise assumption, these guarantees extend to the stochastic setting with additional logarithmic factors and noise-dependent terms. Experiments across the three problem classes further support our theoretical findings.

Figures & tables

Problem class			Gradient type	Assumptions	Convergence
Smooth	Lipschitz	Convex
$\checkmark$	$\times$	$\times$	Exact	Assm. 4.1	$\raisebox{-0.43057pt}{$ O $}\big(\frac{1}{T}\big)$
$\checkmark$	$\times$	$\times$	Stochastic	Assms. 4.1 , 4.2	$\raisebox{-0.6458pt}{$ \widetilde{O} $}\big(\frac{1}{T}+\frac{\sigma}{\sqrt{mKT}}\big)$
$\checkmark$	$\times$	$\checkmark$	Exact	Assm. 4.1	$\raisebox{-0.43057pt}{$ O $}\big(\frac{1}{T}\big)$
$\checkmark$	$\times$	$\checkmark$	Stochastic	Assms. 4.1 , 4.2	$\raisebox{-0.6458pt}{$ \widetilde{O} $}\big(\frac{1}{T}+\frac{\sigma}{\sqrt{mKT}}\big)$
$\times$	$\checkmark$	$\checkmark$	Exact	Assm. 4.1	$\raisebox{-0.6458pt}{$ \widetilde{O} $}\big(\frac{1}{\sqrt{T}}\big)$

Table 1: Summary of convergence results for Algorithm 1 .

\raisebox{-0.58122pt}{

\widetilde{O}

}(\cdot)

hides logarithmic factors.

\displaystyle x_{t+1}:=x_{t}+\frac{1}{mK\tau_{t}}\sum_{i=1}^{m}\Delta_{i}^{t},\qquad\frac{1}{\eta_{t+1}^{2}}:=\frac{1}{\eta_{t}^{2}}\left(1+\frac{\|x_{t+1}-x_{t}\|^{2}}{\eta^{2}}\right).

Algorithm 1 Adaptive Federated Learning Method (AdaFed)

Figure 1: Performance of AdaFed on IJCNN1 across the three problem classes. Solid and dashed curves show the uniform-random and last-iterate metrics, respectively. Stochastic results show the mean over twenty seeds, with shaded regions indicating the upward and downward RMS deviations.

Figure 2: Performance of AdaFed compared with representative adaptive FL methods on IJCNN1 with

K=10

local steps. FedAdagrad, FedLADA, and FAFED represent server-, client-, and both-level adaptation, respectively. Solid and dashed curves show the uniform-random and last-iterate metrics, respectively. Stochastic results show the mean over twenty seeds, with shaded regions indicating the upward and downward RMS deviations.

Appendix figures & tables18 assets

Supplementary material from the paper’s appendix.

Appendix

\displaystyle x_{t+1}:=x_{t}+\frac{1}{mK\tau_{t}}\sum_{i\in\mathcal{S}_{t}}\frac{\Delta_{i}^{t}}{p_{i}},\qquad\frac{1}{\eta_{t+1}^{2}}:=\frac{1}{\eta_{t}^{2}}\left(1+\frac{\|x_{t+1}-x_{t}\|^{2}}{\eta^{2}}\right).

Appendix

Algorithm 2 Adaptive Federated Learning Method (AdaFed) with Partial Participation

Figure 3: Performance of AdaFed on synthetic logistic problems with Gram condition number

\kappa=10

across the three problem classes and different numbers of local steps

K

. Rows show exact gradients and stochastic gradients with minibatch sizes

b=1

and

b=8

, respectively. Solid and dashed curves show the uniform-random and last-iterate metrics, respectively. Stochastic results show the mean over ten seeds, with shaded regions indicating the upward and downward RMS deviations.

Figure 4: Performance of AdaFed on IJCNN1 across the three problem classes and different numbers of local steps

K

. Rows show exact gradients and stochastic gradients with minibatch sizes

b=1

and

b=32

Figure 5: Sensitivity of AdaFed to

\eta

on synthetic logistic problems with Gram condition number

\kappa=10

, with

\gamma=1

and

K=10

. Solid and dashed curves show the uniform-random and last-iterate metrics, respectively. Stochastic results show the mean over ten seeds, with shaded regions indicating the upward and downward RMS deviations.

Figure 6: Sensitivity of AdaFed to

\eta

on IJCNN1, with

\gamma=1

and

K=10

Figure 7: Sensitivity of AdaFed to

\gamma

on synthetic logistic problems with Gram condition number

\kappa=10

, with

\eta=1

and

K=10

Figure 8: Sensitivity of AdaFed to

\gamma

on IJCNN1, with

\eta=1

and

K=10

Figure 9: Performance of AdaFed under partial participation on synthetic logistic problems with Gram condition number

\kappa=10

and different participation rates

p

Figure 10: Performance of AdaFed under partial participation on IJCNN1 with different participation rates

p

Method	Tuned quantities	Finite search bounds
AdaFed	$\eta,\gamma$	$\eta\in[10^{-3},10^{3}]$ ; $\gamma\in[10^{-3},10^{3}]$
FedAdagrad	$c,\eta_{g}$	$c\in[1/128000,20]$ , $\eta_{g}\in[10^{-3},100]$
FedAdam	$c,\eta_{g}$	$c\in[1/128000,20]$ , $\eta_{g}\in[10^{-3},100]$
FedUR	$c_{0},\alpha$	$c_{0}\in[10^{-4},10]$ , $\alpha\in[0.1,1]$
Local AdaAlter	$c,\epsilon$	$c\in[10^{-5},100]$ , $\epsilon\in[10^{-4},10]$
FedSPS	$c,g$	$c\in[10^{-3},1000]$ , $g\in[10^{-4},100]$

Appendix

Table 2: Tuned quantities and finite search bounds for the adaptive-method comparisons.

\eta_{g}

is a server rate,

c

c_{0}

is a local scale,

g

is a FedSPS cap scale, and

\epsilon

is Local AdaAlter’s placeholder.

Figure 11: Performance of AdaFed compared with representative adaptive FL methods on synthetic data with

\kappa=1

and

K=10

Figure 12: Performance of AdaFed compared with server-adaptive FL methods on synthetic data with

\kappa=1

and

K=10

local steps. FedAdagrad, FedAdam, and FedUR are server-level adaptive methods. Solid and dashed curves show the uniform-random and last-iterate metrics, respectively. Stochastic results are shown as the mean over three seeds.

Figure 13: Performance of AdaFed compared with server-adaptive FL methods on IJCNN1 with

K=10

local steps. FedAdagrad, FedAdam, and FedUR represent server-level adaptation. Solid and dashed curves show the uniform-random and last-iterate metrics, respectively. Stochastic results are shown as the mean over three seeds.

Figure 14: Performance of AdaFed compared with client-adaptive FL methods on synthetic data with

\kappa=1

and

K=10

local steps. Local AdaAlter, FedSPS, and FedLADA represent client-level adaptation. Solid and dashed curves show the uniform-random and last-iterate metrics, respectively. Stochastic results are shown as the mean over three seeds.

Figure 15: Performance of AdaFed compared with client-adaptive FL methods on IJCNN1 with

K=10

Figure 16: Performance of AdaFed compared with FL methods using both server- and client-level adaptation on synthetic data with

\kappa=1

and

K=10

local steps. FAFED and FedAda 2 represent both-level adaptation. Solid and dashed curves show the uniform-random and last-iterate metrics, respectively. Stochastic results are shown as the mean over three seeds.

Figure 17: Performance of AdaFed compared with FL methods using both server- and client-level adaptation on IJCNN1 with

K=10

Paper / result	Objective assumptions	Stochastic / heterogeneity	Method setting	Main result / notes
[ Wang et al., 2025 ] Thm. 1	Class: nonconvex Smoothness: $L$ -smooth in $\\|\cdot\\|_{\alpha}$ Strong conv.: ${\times}$ Grad. bound: $\\|\nabla f_{n}(w)\\|_{\infty}\leq C$	Variance/noise: $\mathbb{E}\\|\xi_{t}\\|_{\alpha}^{\alpha}\leq G$ Heterogeneity: ${\times}$	Participation: full participation Adaptation: In server Local updates: one local step	Thm. 1: AdaGrad-OTA; rate $O(\log(T)/T^{1-1/\alpha})$
[ Wang et al., 2025 ] Thm. 2	Class: nonconvex Smoothness: $L$ -smooth in $\\|\cdot\\|_{\alpha}$ Strong conv.: ${\times}$ Grad. bound: $\\|\nabla f_{n}(w)\\|_{\infty}\leq C$	Variance/noise: $\mathbb{E}\\|\xi_{t}\\|_{\alpha}^{\alpha}\leq G$ Heterogeneity: ${\times}$	Participation: full participation Adaptation: In server Local updates: one local step	Thm. 2: Adam-OTA; rate $O(1/T)$
[ Chen et al., 2025 ] Thm. 5.2	Class: nonconvex Smoothness: each $f_{i}$ is $L$ -smooth Strong conv.: ${\times}$ Grad. bound: ${\times}$	Variance/noise: $\mathbb{E}\\|n_{i}^{(t,k)}\\|^{2}\leq\sigma^{2}$ Heterogeneity: ${\times}$	Participation: client sampling $S^{t}$ ; minibatch Adaptation: In client Local updates: $K$ local steps	Thm. 5.2: FAdamGC, $\beta_{2}=\epsilon=0$ ; rate $O\!\left(\sqrt{\frac{LF\sigma^{2}}{nKT}}+\frac{LF}{T}\allowbreak+\frac{LK}{T}+\frac{K\sigma}{T}\right)$
[ Chen et al., 2025 ] Thm. 5.4	Class: nonconvex Smoothness: each $f_{i}$ is $L$ -smooth Strong conv.: ${\times}$ Grad. bound: $\\|g_{i}^{(t)}\\|\leq G$	Variance/noise: $\mathbb{E}\\|n_{i}^{(t,k)}\\|^{2}\leq\sigma^{2}$ Heterogeneity: ${\times}$	Participation: client sampling $S^{t}$ ; minibatch Adaptation: In client Local updates: $K$ local steps	Thm. 5.4: FAdamGC, general Adam case; rate $O\!\Big(\sqrt{\frac{LF\sigma^{2}}{nKT}}+\frac{LF}{T}+\frac{KG^{6}}{\epsilon^{2}T}\allowbreak+\frac{K(\sigma^{2}+(1+\epsilon^{2})G^{2})}{\epsilon^{2}T}\Big)$
[ Chen et al., 2025 ] Thm. C.2	Class: nonconvex Smoothness: each $f_{i}$ is $L$ -smooth Strong conv.: ${\times}$ Grad. bound: $\\|g_{i}^{(t)}\\|\leq G$	Variance/noise: $\mathbb{E}\\|n_{i}^{(t,k)}\\|^{2}\leq\sigma^{2}$ Heterogeneity: ${\times}$	Participation: client sampling $S^{t}$ ; minibatch Adaptation: In client Local updates: $K$ local steps	Thm. C.2: FA-NT, general Adam case; rate has worse $K^{2}$ final term
[ Chen et al., 2025 ] Thm. C.3	Class: nonconvex Smoothness: each $f_{i}$ is $L$ -smooth Strong conv.: ${\times}$ Grad. bound: ${\times}$	Variance/noise: $\mathbb{E}\\|n_{i}^{(t,k)}\\|^{2}\leq\sigma^{2}$ Heterogeneity: $\\|\nabla f_{i}(x)-\nabla f(x)\\|\leq B$	Participation: client sampling $S^{t}$ ; minibatch Adaptation: In client Local updates: $K$ local steps	Thm. C.3: FA-NT, $\beta_{2}=\epsilon=0$ ; requires bounded data heterogeneity

Appendix

Table 3: Detailed comparison of assumptions, adaptive mechanisms, and theoretical guarantees in related federated optimization methods.

Federated Learning

Federated learning trains shared models across data holders without requiring them to pool their raw data.

See all 961 papers

Mar 3, 2026

Convex and Non-convex Federated Learning with Stale Stochastic Gradients: Diminishing Step Size is All You Need

We propose a general framework for distributed stochastic optimization under delayed gradient models. In this setting,

n

local agents leverage their own data and computation to assist a central server in minimizing a global objective composed of agents' local cost functions. Each agent is allowed to transmit stochastic-potentially biased and delayed-estimates of its local gradient. While a prior work has advocated delay-adaptive step sizes for stochastic gradient descent (SGD) in the presence of delays, we demonstrate that a pre-chosen diminishing step size is sufficient and matches the performance of the adaptive scheme. Moreover, our analysis establishes that diminishing step sizes recover the optimal SGD rates for nonconvex and strongly convex objectives.

Xinran Zheng, Tara Javidi, Behrouz Touri

University of Illinois Urbana-Champaign · University of California San Diego

Stochastic Optimization Convergence Stochastic Gradient Descent Nonconvex Stochastic Optimization Federated Learning Asynchronous Federated Learning

Ask PDF

Feb 4, 2025

Correcting Before Adapting: Gradient Correction for Federated Adaptive Optimization

Client-drift compensation methods mitigate the impact of data heterogeneity in federated learning (FL). While drift correction is well studied for stochastic gradient descent (SGD)-based FL, its interaction with adaptive client optimizers remains less understood. Adaptive optimizers such as Adam build each update from running averages of past gradients, so correcting only the final update direction leaves those averages, and the drift they contain, unchanged. We propose Federated Adaptive Gradient Correction ({\tt FA-GC}), which corrects each gradient before it enters these running averages, and compare it with Federated Adaptive Naive Tracking ({\tt FA-NT}), which adds the correction to the final direction. We show that, in an idealized setting without gradient noise, the direction FA-GC follows vanishes only at stationary points of the global objective, whereas the direction FA-NT follows can vanish elsewhere when the optimizer is nonlinear. We also prove convergence of FA-GC for Adam-family methods, gradient normalization, and Muon, which orthogonalizes matrix updates. Unlike our bound for FA-NT, these bounds contain no term for client heterogeneity. On seven image and language tasks, FA-GC with Adam needs the fewest rounds on average to reach the target accuracy, and on the three image tasks it also uses less total communication than uncorrected local Adam and FA-NT. On the QNLI question-entailment task with SOAP, which applies Adam in a rotated basis, and Muon under strong heterogeneity, FA-GC gives the highest accuracy, while FA-NT falls below the uncorrected optimizers.

Evan Chen, Shiqiang Wang, Dong-Jun Han +2

Elmore Family School of Electrical and Computer Engineering, Purdue University, West Lafayette, IN · IBM T. J. Watson Research Center, Yorktown Heights, NY, USA · Department of Computer Science and Engineering, Yonsei University, Seoul, Korea

Federated Learning Heterogeneous Federated Learning Adaptive Gradient Methods

Ask PDF

May 28, 2026

Convergence of Steepest Descent and Adam under Non-Uniform Smoothness

Recent work has analyzed the convergence of first-order methods under non-uniform smoothness assumptions that better model the loss landscape in machine learning tasks. We generalize this assumption to objectives whose curvature is an affine function of the objective value. This property is satisfied by a broad class of problems, including logistic regression, generalized linear models with a logistic link function, softmax policy gradient in reinforcement learning, and a class of neural networks. Under this assumption and gradient domination conditions, we establish a general convergence rate for the steepest descent method, and deterministic, diagonal variants of RMSProp and Adam. Our results imply that for logistic regression on separable data and the softmax policy gradient objective, sign GD converges linearly and is provably faster than GD. Furthermore, we show that for a class of two-layer neural networks on separable data, RMSProp and Adam can converge at a linear rate with a constant step-size and momentum parameter. Finally, we present a lower bound demonstrating that, under our assumption, RMSProp and Adam are provably faster than AdaGrad, AMSGrad, gradient descent, and heavy-ball momentum.

Sharan Vaswani, Yifan Sun, Reza Babanezhad

1Simon Fraser University · 2Stony Brook University · 3Samsung AI, Montreal.

Optimization Convergence Analysis Gradient Descent Nonconvex Optimization Adaptive Gradient Methods

Ask PDF

Abstract

Federated Learning

Explore similar work

Convex and Non-convex Federated Learning with Stale Stochastic Gradients: Diminishing Step Size is All You Need

Correcting Before Adapting: Gradient Correction for Federated Adaptive Optimization

Convergence of Steepest Descent and Adam under Non-Uniform Smoothness