Source-linked AI summary
Accelerating Stochastic Composition Optimization
Mengdi Wang, Ji Liu, Ethan X. Fang
TL;DR
The paper studies stochastic composition optimization with nonsmooth regularization, where standard stochastic gradients are unavailable. It introduces ASC-PG, a proximal method using stochastic first-order oracle queries and extrapolation-smoothing, and reports faster convergence with optimal sample-error complexity in important special cases, including reinforcement learning.
Problem
Standard stochastic gradient methods do not apply directly to nonlinear compositions of expected-value functions, including problems with nonsmooth penalties and applications such as reinforcement learning.
Method
ASC-PG is a stochastic compositional proximal-gradient method that uses black-box first-order oracle queries, two-timescale updates, and extrapolation-smoothing to estimate the inner expected function.
Results
The method improves known convergence rates, achieves optimal sample-error complexity in important linear-function cases, and yields a faster-converging on-policy reinforcement-learning algorithm.
Takeaways & Limitations
ASC-PG extends stochastic composition optimization to nonsmooth regularization and supplies rate-guaranteed tools for reinforcement learning and related applications.
Abstract
from arXiv · showhide
Consider the stochastic composition optimization problem where the objective is a composition of two expected-value functions. We propose a new stochastic first-order method, namely the accelerated stochastic compositional proximal gradient (ASC-PG) method, which updates based on queries to the sampling oracle using two different timescales. The ASC-PG is the first proximal gradient method for the stochastic composition problem that can deal with nonsmooth regularization penalty. We show that the ASC-PG exhibits faster convergence than the best known algorithms, and that it achieves the optimal sample-error complexity in several important special cases. We further demonstrate the application of ASC-PG to reinforcement learning and conduct numerical experiments.
1 Introduction
The paper addresses stochastic optimization problems involving compositions of expected-value functions, where standard stochastic gradients are unavailable. It develops ASC-PG to obtain efficient convergence guarantees for broader penalized problems and applications including reinforcement learning and risk-averse learning.
- Standard stochastic gradient methods do not apply directly to nonlinear objectives involving compositions between expected-value functions.
- The problem allows smooth inner and outer functions, a nonsmooth closed convex penalty such as the ℓ1-norm, nonconvexity, nonmonotonicity, and dependent random variables.
- The stochastic composition framework covers reinforcement learning, where Bellman equations can be solved from sampled transition information.
- Risk-averse learning is another application because its variance term composes a mean-square function with an expected loss function.
- ASC-PG is proposed as a proximal-gradient method with significantly improved sample-error complexity, including optimal complexity when the outer or inner function is linear.
- In the inner-linear case, ASC-PG achieves O(1/K), and the resulting Bellman-equation method reaches O(1/K) based on K observed state transitions.
2 Algorithm
ASC-PG operates in a black-box sampling setting with stochastic first-order information and combines proximal updates with extrapolation-smoothing auxiliary iterates. The auxiliary scheme produces nearly unbiased estimates of the inner expected function, enabling faster convergence.
- The sampling oracle returns random inner-function evaluations, noisy inner subgradients, and noisy gradients of the outer function.
- ASC-PG generalizes SCGD by replacing its projection step with a proximal step, allowing treatment of penalized stochastic composition problems.
- The algorithm updates the main iterate and auxiliary iterates through a proximal-gradient step and an extrapolation-smoothing scheme.
- The auxiliary y- and z-updates use weighted smoothing and extrapolation so that y_k tracks the unknown quantity g(x_k) efficiently.
- In the linear-inner-function case, the scaled error becomes a zero-mean, zero-drift martingale under additional technical assumptions.
- The zero-drift property follows from the near-unbiasedness created by extrapolation-smoothing and is identified as the key source of faster convergence.
3 Main Results
The paper establishes convergence guarantees for ASC-PG under general and optimally strongly convex stochastic composition settings, including nonsmooth penalties and dependent sampling variables. It improves prior rates in the general case and reaches optimal sample-error complexity when either the inner or outer function is linear.
- Assumptions: The analysis assumes unbiased sampling, bounded sample gradients and penalty gradients, and Lipschitz conditions on the inner and outer functions.The stochastic-oracle assumptions permit dependence between the inner and outer random variables.
- General optimization: Theorem 1 covers general optimization problems that need not be convex and establishes convergence for the smooth-optimization case with R(x)=0.The stepsizes use α_k=k^-a and β_k=2k^-b under the stated assumptions.
- Special cases: When Lg=0 or Lf=0, Theorem 1 selects a=b=1/2 and obtains the corresponding improved special-case rate.These cases respectively include a linear inner mapping or a linear outer function.
- General optimization: O(k^-4/9) replaces O(k^-2/7) in the general case, strictly improving the previous convergence rate of Wang et al. [2016].The paper presents this as a new benchmark for stochastic composition optimization.
- Strongly convex optimization: For optimally strongly convex objectives, ASC-PG has a convergence guarantee under stepsizes with Ca>4λ and Cb>2.Theorem 2 measures convergence through the distance from xk to the projection of xk onto the optimal solution set.
- Special cases: With a linear inner mapping, ASC-PG achieves O(1/K) for Bellman error minimization with ℓ1 regularization, identified as optimal sample-error complexity for on-policy reinforcement learning.The paper argues more broadly that inner linearity makes the composition problem no harder in convergence rate than expectation minimization.
4 Application to Reinforcement Learning
The paper applies ASC-PG to on-policy reinforcement-learning value evaluation by solving a Bellman residual minimization problem. Experiments compare convergence across unregularized and ℓ1-regularized settings, with faster empirical convergence and an O(1/k) rate in the linear-inner-function case.
- Application: ASC-PG is applied to policy-value evaluation in reinforcement learning through Bellman equations.The method is used to solve a Bellman residual minimization formulation.
- Bellman formulation: The Bellman equation defines each state value as an expected discounted reward plus the next-state value under policy π.The value function is approximated using features and a parameter vector w∗.
- Experimental setup: The first two experiments compare ASC-PG with SCGD and GTD2-MP without regularization, while Experiment 3 adds an ℓ1 penalty.Experiment 1 uses Baird’s example; Experiment 2 uses randomly generated 100-state MDPs; Experiment 3 uses a 100-dimensional feature space with sparse w∗.
- Results: ASC-PG achieves the fastest empirical convergence in both unregularized experiments across the reported parameter and value-function errors.The errors are averaged over 100 runs.
- Results: O(1/k) convergence is empirically supported for ASC-PG when the inner mapping is linear, while the ℓ1-regularized experiment also shows fast empirical convergence.The regularized experiment uses an approximate optimum obtained after one million iterations.
5 Conclusion
The paper develops a stochastic proximal-gradient method for penalized stochastic composition problems and applies it to reinforcement learning. It reports faster convergence than the best known algorithms while leaving further improvement of the general O(K−4/9) rate open.
- Method: ASC-PG is a proximal-gradient method for penalized stochastic composition problems that interacts with a stochastic first-order oracle.The method is designed to handle nonsmooth regularization in the composition setting.
- Results: The method establishes convergence rates under a variety of assumptions and yields a reinforcement-learning algorithm with faster convergence than the best known algorithms.The reinforcement-learning application is described as an on-policy learning algorithm.
- Future directions: Further improvement of the general O(K−4/9) rate remains open, and adapting the method to more specific reinforcement-learning and risk-averse problems is identified as future work.The proposed adaptations are intended to exploit more problem-specific structure.
Supplemental Materials
The supplemental material proves convergence bounds for the algorithm under stated assumptions using iterate inequalities, proximal optimality conditions, and bounds on auxiliary sequences. The selected results identify step-size choices yielding an O(K−4/9) convergence rate.
- Convergence analysis: The supplemental analysis bounds consecutive iterates and auxiliary sequences under the paper’s assumptions.The proofs use induction, convexity inequalities, proximal-operation definitions, and previously established lemmas.
- Proximal analysis: The proof handles the nonsmooth regularizer through subdifferential vectors and proximal optimality conditions.The analysis also bounds distances to the optimal solution set using decompositions of successive iterates.
- Error bounds: The analysis controls the running-estimation error and related auxiliary terms through decay bounds derived from the supporting lemmas.The estimates are combined with expectation bounds and concavity or Lipschitz arguments.