Heterogeneous Stochastic Momentum ADMM for Distributed Nonconvex Composite Optimization

TL;DR

This paper introduces HSM-ADMM, a novel distributed optimization algorithm that adapts to heterogeneous network topologies, achieving optimal convergence rates with reduced communication and independence from global network parameters.

Contribution

The paper proposes a node-specific adaptive step-size strategy within HSM-ADMM, decoupling stability from global network properties and enabling robust, efficient convergence in heterogeneous networks.

Findings

Achieves optimal oracle complexity of $\\mathcal{O}(\epsilon^{-1.5})$.

Requires only a single primal variable transmission per iteration.

Demonstrates superior efficiency in numerical experiments.

Abstract

This paper investigates the distributed stochastic nonconvex and nonsmooth composite optimization problem. Existing stochastic typically rely on uniform step size strictly bounded by global network parameters, such as the maximum node degree or spectral radius. This dependency creates a severe performance bottleneck, particularly in heterogeneous network topologies where the step size must be conservatively reduced to ensure stability. To overcome this limitation, we propose a novel Heterogeneous Stochastic Momentum Alternating Direction Method of Multipliers (HSM-ADMM). By integrating a recursive momentum estimator (STORM), HSM-ADMM achieves the optimal oracle complexity of $\mathcal{O}(\epsilon^{-1.5})$ to reach an $\epsilon$-stationary point, utilizing a strictly single-loop structure and an $\mathcal{O}(1)$ mini-batch size. The core innovation lies in a node-specific adaptive step-size strategy, which scales the proximal term according to local degree information. We theoretically demonstrate this design completely decouples the algorithmic stability from global network properties, enabling robust and accelerated convergence across arbitrary connected topologies without requiring any global structural knowledge. Furthermore, HSM-ADMM requires transmitting only a single primal variable per iteration, significantly reducing communication bandwidth compared to state-of-the-art gradient tracking algorithms. Extensive numerical experiments on distributed nonconvex learning tasks validate the superior efficiency of the proposed HSM-ADMM algorithm.