Optimal Momentum Methods for Stochastic Multilevel Compositional Optimization
This paper studies stochastic multi-level compositional optimization, where the objective is a nested composition of $T$ smooth non-convex functions and only stochastic per-layer function and gradient values are available. The core idea is to use momentum-based estimators to track inner function values level-by-level, and then apply the chain rule to build a momentum gradient estimator for the overall objective. It matters because it establishes that the optimal $\mathcal{O}(\epsilon^{-4})$ sample complexity is achievable under standard smoothness assumptions, closing a gap with single-level lower bounds without relying on the stronger average-smoothness conditions common in prior multi-level literature.
The paper proves that finding an $\epsilon$-stationary point in stochastic multi-level non-convex optimization can be done with the optimal $\mathcal{O}(\epsilon^{-4})$ sample complexity under standard Lipschitz and bounded-variance assumptions, matching the single-level lower bound. The main MML algorithm theoretically requires large per-iteration mini-batches $B_1 = \mathcal{O}(K)$, while the MMLFA variant removes this requirement via a clipped first-order function-value approximation at the cost of introducing an additional boundedness condition on intermediate iterates.
The analysis is rigorous in dispensing with the average-smoothness assumption: the recursive lemmas (Lemmas 3.4–3.6) carefully bound the coupled function-value and gradient-estimation errors across levels. The batch-free MMLFA extension is a clever design, replacing the large-batch requirement with a first-order Taylor drift correction combined with a clipping operation to keep estimators bounded.
A key limitation is the mismatch between theory and practice for MML: the analysis requires per-iteration batches $B_1 = \mathcal{O}(K)$, yet the experiments use a fixed batch size of 128 for all methods without observing degradation, a discrepancy the authors note but do not explain.
Furthermore, MMLFA introduces an explicit boundedness requirement $\|f_i(\mathbf{u}_{i-1}^k)\| \le C$ that is absent from the base assumptions (Assumptions 3.1–3.3), and the analysis hides constants (e.g., $C_i$ and $D^i$) that scale exponentially with the composition depth $T$, which the paper does not discuss.
The theoretical claims are well-supported and the convergence rates match the established lower bound. Comparisons to Chen et al. (2021) and Zhang and Xiao (2021) correctly note their reliance on average smoothness. However, the strong claim that Balasubramanian et al. (2022) produced invalid results due to dependency issues is asserted in the main text without proof, deferred to an appendix. Empirically, the evaluation is limited to small-scale benchmarks with a single fixed batch size and no ablation on composition depth.
The algorithms and proof steps are written with enough detail to allow reimplementation, but no source code or data repository is provided. The extensive grid search for hyperparameters and the absence of reported random seeds make exact reproduction difficult. For MMLFA, the clipping threshold $C$ is argued to be robust when large, yet its theoretical guarantee depends on the a priori unknown bound of intermediate function norms.
This paper investigates stochastic multi-level optimization where the objective is a nested composition of several smooth non-convex functions. We assume that only stochastic estimates of the gradient and function values for each level are accessible. Consequently, obtaining an accurate estimate of the overall gradient is challenging due to the nested structure. To address this, we employ a momentum-based estimator with mini-batches to track the function values of each level, which are subsequently used to construct momentum gradient estimators. We establish an optimal sample complexity of $\mathcal{O}(\epsilon^{-4})$ for finding an $\epsilon$-stationary point, avoiding the stronger average smoothness assumption commonly relied upon in prior literature. Furthermore, by employing a normalization technique, we attain the same rate without requiring problem-dependent constants to set hyperparameters. To achieve the optimal rate without mini-batches, we further develop a batch-free method that incorporates a first-order approximation and a clipping technique for function value estimation. Finally, we validate the effectiveness of our proposed methods through experiments on risk-averse portfolio optimization and hierarchical tilted empirical risk minimization.
Pick a starting point or write your own. Challenges run in the background, so you can keep reading while the AI investigates.
No challenges yet. Disagree with the review? Ask the AI to revisit a specific claim.