Distributed Stochastic Consensus Optimization With Momentum for Nonconvex Nonsmooth Problems

Abstract

While many distributed optimization algorithms have been proposed for solving\nsmooth or convex problems over the networks, few of them can handle non-convex\nand non-smooth problems. Based on a proximal primal-dual approach, this paper\npresents a new (stochastic) distributed algorithm with Nesterov momentum for\naccelerated optimization of non-convex and non-smooth problems. Theoretically,\nwe show that the proposed algorithm can achieve an $\\epsilon$-stationary\nsolution under a constant step size with $\\mathcal{O}(1/\\epsilon^2)$\ncomputation complexity and $\\mathcal{O}(1/\\epsilon)$ communication complexity.\nWhen compared to the existing gradient tracking based methods, the proposed\nalgorithm has the same order of computation complexity but lower order of\ncommunication complexity. To the best of our knowledge, the presented result is\nthe first stochastic algorithm with the $\\mathcal{O}(1/\\epsilon)$ communication\ncomplexity for non-convex and non-smooth problems. Numerical experiments for a\ndistributed non-convex regression problem and a deep neural network based\nclassification problem are presented to illustrate the effectiveness of the\nproposed algorithms.\n