Stochastic Recursive Momentum Method for Non-Convex Compositional Optimization

Abstract: We propose a novel stochastic optimization algorithm called STOchastic Recursive Momentum for Compositional (STORM-Compositional) optimization that minimizes the composition of expectations of two stochastic functions, the latter being an optimization problem arising in various important machine learning applications. By introducing the momentum term in the compositional gradient updates, STORM-Compositional operates the stochastic recursive variance-reduced compositional gradients in an exponential-moving average way. This leads to an $O(\varepsilon^{-3})$ complexity upper bound for STORM-Compositional, that matches the best known complexity bounds in previously announced compositional optimization algorithms. At the same time, STORM-Compositional is a single loop algorithm that avoids the typical alternative tuning between large and small batch sizes, as well as recording of checkpoint gradients, that persist in variance-reduced stochastic gradient methods. This allows considerably simpler parameter tuning in numerical experiments, which demonstrates the superiority of STORM-Compositional over other stochastic compositional optimization algorithms.