A Mini-batch Proximal Adam Algorithm for Nonsmooth Nonconvex Stochastic Composite Optimization via MCMC Sampling

We propose MCMC-ProxAdam, a novel stochastic proximal gradient algorithm designed to solve constrained nonconvex stochastic composite optimization problems with a nonsmooth convex regularizer, in settings where the gradient of smooth part is intractable and must be estimated via MCMC sampling. We depart from the traditional assumption of bounded stochastic gradient by introducing bounded sub-exponential norm for the stochastic gradient, then we derive a Bernstein-type inequality for the mini-batch stochastic estimator under non-stationary Markov chains, allowing our framework to accommodate biased and unbounded stochastic gradient common in practical scenarios. Leveraging this Bernstein-type inequality, we derive error bounds for the mini-batch gradient estimator and establish the convergence result for MCMC-ProxAdam in the nonsmooth nonconvex setting. Numerical experiments on energy-based models (EBMs) for image generation tasks across the MNIST, CIFAR-10, and CelebA datasets demonstrate that the proposed algorithm is capable of solving nonconvex composite optimization problems with a convex nonsmooth term, with numerical performance empirically validating our theoretical convergence results.