Last-Iterate Convergence Analysis of Federated Adaptive Optimization using CIFAR-10 and FEMNIST Benchmarks

Abstract

Deep models over a heterogeneous set of resource-constrained clients have gained widespread popularity for training with federated adaptive optimization methods like FedAdam and FedYogi, but most guarantees are on convergence of an averaged iterate and not the actual last iterate that is returned to practitioners. This gap is significant in practice as it's not an average but the final global model that is deployed. We propose a last-iterate convergence analysis for a large class of federated adaptive optimizers with general non-convex smoothness conditions, partial participation and bounded client drift. We motivate a drift-corrected adaptive aggregation rule, LI-FedAdapt, which incorporates serverside momentum correction and a decaying schedule of local-steps, and establish a last-iterate stationarity bound of 1/√T up to constants determined by client heterogeneity that agrees with the optimal rate of averaged-iterate federated adaptive methods. We test the theory on two common federated benchmarks, the partitioned CIFAR-10 dataset with label skew Dirichlet(0.3), and the naturally-partitioned FEMNIST dataset, with 100 simulated clients and 10% participation per round, for 200 communication rounds. In practice, LI-FedAdapt closely follows the rate in theory, and is significantly more efficient in terms of final test accuracy and communication efficiency as compared to FedAvg, FedAdam and FedYogi, and a per-iteration ablation in participation rate and local-epoch count shows a visually smaller gap at the last iteration. The above results show that it is indeed possible to achieve these last iterate guarantees and in fact even achievable for adaptive federated optimization with no extra communication cost.