Decentralized Nonconvex Composite Federated Learning with Gradient Tracking and Momentum

Abstract

Decentralized Federated Learning (DFL) enables collaborative model training without relying on a central server. When local objectives are nonconvex and coupled with nonsmooth weakly convex regularization, DFL gives rise to a challenging decentralized nonconvex composite optimization problem involving data heterogeneity, stochastic-gradient noise, and consensus error. We propose DEPOSITUM, a decentralized composite optimization algorithm for this problem. DEPOSITUM maintains momentum-filtered stochastic gradient estimates via a tracking mechanism, which accommodates both Polyak and Nesterov momentum. It further allows multiple local updates between communication rounds to improve communication efficiency. Theoretical analysis demonstrates that it achieves an expected $\epsilon$-stationary point with an iteration complexity of $\mathcal{O}(1/{\epsilon }^2)$ without imposing bounded gradient heterogeneity or mean-squared smoothness assumptions. With an appropriate stepsize and momentum schedule, the averaged stationarity measure further achieves a rate of \(\mathcal{O}(1/\sqrt{nT})\) after a network-dependent transient, using a mini-batch size independent of $T$. Experiments on different benchmark datasets validate the effectiveness of DEPOSITUM and demonstrate competitive performance against representative federated composite optimization methods.