Optimal Primal-Dual Algorithm with Last iterate Convergence Guarantees for Stochastic Convex Optimization Problems

TL;DR

This paper introduces Aug-ConEx, a novel primal-dual algorithm that guarantees optimal convergence rates on the last iterate for stochastic convex optimization with nonsmooth functions and constraints, outperforming existing methods.

Contribution

The paper presents the first method achieving optimal last-iterate convergence rates for stochastic composite nonsmooth convex problems with constraints, using a new augmented Lagrangian approach.

Findings

Achieves $ ext{O}(1/\sqrt{K})$ convergence rate without strong convexity.

Attains $ ext{O}(1/K)$ convergence rate with strongly convex objectives.

Demonstrates superior performance over state-of-the-art algorithms in numerical experiments.

Abstract

This paper proposes a novel first-order algorithm that solves composite nonsmooth and stochastic convex optimization problem with function constraints. Most of the works in the literature provide convergence rate guarantees on the average-iterate solution. There is growing interest in the convergence guarantees of the last iterate solution due to its favorable structural properties, such as sparsity or privacy guarantees and good performance in practice. We provide the first method that obtains the best-known convergence rate guarantees on the last iterate for stochastic composite nonsmooth convex function-constrained optimization problems. Our novel and easy-to-implement algorithm is based on the augmented Lagrangian technique and uses a new linearized approximation of constraint functions, leading to its name, the Augmented Constraint Extrapolation (Aug-ConEx) method. We show that Aug-ConEx achieves $\mathcal{O}(1/\sqrt{K})$ convergence rate in the nonsmooth stochastic setting without any strong convexity assumption and $\mathcal{O}(1/K)$ for the same problem with strongly convex objective function. While optimal for nonsmooth and stochastic problems, the Aug-ConEx method also accelerates convergence in terms of Lipschitz smoothness constants to $\mathcal{O}(1/K)$ and $\mathcal{O}(1/K^2)$ in the aforementioned cases, respectively. To our best knowledge, this is the first method to obtain such differentiated convergence rate guarantees on the last iterate for a composite nonsmooth stochastic setting without additional $\log{K}$ factors. We validate the efficiency of our algorithm by comparing it with a state-of-the-art algorithm through numerical experiments.