Stochastic Momentum ADMM for nonconvex and nonsmooth optimization with application to PnP algorithm

This paper introduces a single-loop Stochastic Momentum Alternating Direction Method of Multipliers (SMADMM) for tackling a class of nonconvex and nonsmooth optimization problems. We establish that SMADMM achieves an optimal oracle complexity of $\mathcal{O}(\epsilon^{-\frac{3}{2}})$ in the online setting, where only stochastic first-order oracle, is available. In particular, SMADMM requires only $\mathcal{O}(1)$ stochastic gradient evaluations per iteration and avoids the need for restarting with large batch gradient estimates. This is the first stochastic ADMM method achieving optimal oracle complexity for nonconvex and nonsmooth problems, requiring $\mathcal{O}(1)$ batch size. Furthermore, we extend our method by integrating it with plug-and-play (PnP) priors, resulting in the PnP-SMADMM algorithm. Numerical experiments on classification, CT image reconstruction and phase retrieve demonstrate the practical effectiveness of our approach and validate the theoretical findings.