Optimal Momentum Methods for Stochastic Multilevel Compositional Optimization

Wei Jiang; Rui Yan; Sifan Yang; Yuanyu Wan; Lijun Zhang; Zechao Li

Source-linked AI summary

Optimal Momentum Methods for Stochastic Multilevel Compositional Optimization

Wei Jiang, Rui Yan, Sifan Yang, Yuanyu Wan, Lijun Zhang, Zechao Li

arXiv:2610.01572v1Oct 1, 2026math.OCcs.LG

TL;DR

Nested stochastic compositions make accurate overall gradient estimation difficult, while prior optimal-rate analyses often require stronger average smoothness assumptions. The paper develops momentum-based estimators and variants, establishing O(ϵ^-4) sample complexity under standard smoothness assumptions. It also reports experiments on portfolio optimization and hierarchical tilted empirical risk minimization.

Problem
In nested stochastic optimization, only unbiased stochastic function-value and gradient estimates at each level are available, making an unbiased overall gradient estimate generally infeasible; prior work often assumes average smoothness.
Source: §1 · p. 1
Method
Momentum-based estimators track function values and gradients, with normalization for hyperparameter setting and first-order approximation plus clipping for a constant-batch variant.
Source: §1 · p. 2
Results
O(ϵ^-4) sample complexity is achieved under standard smoothness assumptions without the stronger average smoothness conditions used in prior work.
Source: §3.2 · p. 5, §6 · p. 10
Takeaways & Limitations
Experiments demonstrate the methods’ effectiveness on risk-averse portfolio optimization and hierarchical tilted empirical risk minimization.
Source: §1 · p. 2, §6 · p. 10
Theorem 3.7 requires knowledge of problem-dependent constants to set hyperparameters.
Source: §3.2 · p. 5

Abstract

from arXiv · show

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.

1 INTRODUCTION

Nested stochastic compositions make accurate overall gradient estimation difficult, and prior optimal-rate analyses often use stronger average smoothness assumptions. The paper presents momentum-based methods that achieve O(ϵ^-4) sample complexity without those assumptions, with normalized and constant-batch variants.

1 INTRODUCTION

Multi-level nesting makes unbiased estimates of the overall gradient and function value generally infeasible, and estimation errors can accumulate with depth.Prior work often imposes average smoothness and average Lipschitz conditions on stochastic layer estimates.Source: §1 · p. 1
Momentum estimators track function values and gradients to achieve O(ϵ^-4) sample complexity without average smoothness or Lipschitz assumptions.The method uses mini-batches.Source: §1 · p. 2
MML-v2 preserves the optimal sample complexity while removing the need for problem-dependent constants to set hyperparameters.Source: §1 · p. 2
MMLFA combines first-order function-value tracking with clipping to preserve the optimal rate using a constant batch size.Source: §1 · p. 2
Experiments validate the proposed methods on risk-averse portfolio optimization and hierarchical tilted empirical risk minimization.Source: §1 · p. 2

2 RELATED WORK

The related-work section reviews stochastic multi-level optimization and its decentralized and federated variants. It contrasts convergence rates and assumptions across prior methods and summarizes extensions to other settings.

2 RELATED WORK

The section reviews stochastic multi-level optimization alongside decentralized and federated variants.Source: §2 · p. 2

2.1 MULTI-LEVEL OPTIMIZATION

Yang et al. (2019) initiated multi-level compositional optimization; later methods report O(ϵ^-4) under average smoothness or O(ϵ^-3) with nested variance reduction and large batches.Other work studies non-smooth, convex, constrained, stability, and minimax settings (Ruszczy´nski, 2021; Zhang & Lan, 2025; Xiao et al., 2022; Jiang et al., 2024; Pan et al., 2024; Zhang & Gao, 2025).Source: §2.1 · p. 2

2.2 DECENTRALIZED/FEDERATED MULTI-LEVEL OPTIMIZATION

Decentralized and federated work covers level-independent convergence, communication and sample complexity, imbalanced-data AUC maximization, and linear speedup under heterogeneity.The approaches include a double-loop federated method and a Jacobian-vector product estimator.Source: §2.2 · p. 3

3 MOMENTUM-BASED MULTI-LEVEL ALGORITHM

Momentum-based estimators track layer-wise function values and gradients to address the difficulty of estimating gradients through nested expectations. The resulting methods achieve O(ϵ^-4) sample complexity under standard smoothness, with a normalized variant that avoids problem-dependent hyperparameter constants.

3.1 ASSUMPTIONS

The analysis assumes Lipschitz-continuous functions and gradients, plus unbiased function and gradient estimates with bounded variance, but does not require average smoothness and Lipschitz conditions.The paper identifies those stronger conditions as assumptions used in some prior multi-level work.Source: §3.1 · p. 3

3.2 THE PROPOSED METHOD

MML uses momentum estimators to track each layer’s function values and construct a gradient estimator, addressing nested expectations that prevent direct access to an unbiased overall gradient estimate.The method also uses level-wise estimators because estimation errors may accumulate with composition depth.Source: §3.2 · pp. 3–4
O(ϵ^-4) sample complexity is achieved with MML-v1 under standard smoothness assumptions, matching the single-level lower bound.The iteration complexity is O(ϵ^-2), with overall sample complexity O(ϵ^-4).Source: §3.2 · p. 5
MML-v2 preserves the same rate while normalization removes the need for problem-dependent constants when setting hyperparameters.The stated parameter choices include η = K^-1/2 and α = 1/2.Source: §3.2 · p. 5

4 MOMENTUM-BASED MULTI-LEVEL OPTIMIZATION WITH FIRST-ORDER APPROXIMATION

MMLFA addresses the large-batch requirement in stochastic multi-level optimization by tracking function values with first-order approximations and clipping, while retaining optimal sample complexity.

4 MOMENTUM-BASED MULTI-LEVEL OPTIMIZATION WITH FIRST-ORDER APPROXIMATION

Large batches can be impractical when hardware or memory constraints prevent using B1 = O(K).The section introduces MMLFA as an alternative to methods relying on large batches.Source: §4 · p. 5
First-order Taylor tracking plus clipping keeps function estimates controlled and permits small gradient-estimation error with a constant batch size.The function-value tracker uses a first-order approximation; the estimator is clipped to a controlled range.Source: §4 · pp. 5–7
The function-estimation error changes from O(σ²/B1) to O(ασ²), while the recurrence factor improves from 1/α² to 1/α.The smaller α-dependent terms allow control without relying on a large batch size.Source: §4 · p. 6
O(ϵ^-4) sample complexity is retained by MMLFA while eliminating the need for large batches.The guarantee uses η = O(K^-1/2) and α = O(K^-1/2).Source: §4 · p. 7

5 EXPERIMENTS

The experiments compare the proposed methods with existing stochastic multilevel algorithms on portfolio optimization and hierarchical TERM. Both applications report faster decreases in objective-related measures for MML and MMLFA than for baselines.

5 EXPERIMENTS

The evaluation compares against A-TSCGD, NLASG, Nested-SPIDER, and SCSC, selecting hyperparameters by grid search and averaging results over 20 runs.Each method uses batch size 128; a sufficiently large clipping threshold works well without tuning.Source: §5 · p. 7

5.1 MEAN-DEVIATION RISK-AVERSE PORTFOLIO MANAGEMENT

Across Industry-10, Industry-12, and Industry-17, MML and MMLFA reduce loss and gradient norms faster and stabilize at lower levels than other methods.Their overall performance is similar: MMLFA is slightly better on Industry-10/12, while MML is slightly better on Industry-17.Source: Figure 1 · p. 8, §5.1 · p. 8
The portfolio experiments use three real-world datasets with 10, 12, and 17 industrial assets, respectively, and set λ = 0.2.Figure 1 plots loss values and gradient norms against iterations; with batch size 128, sample complexity is proportional to iteration count.Source: §5.1 · pp. 7–8

5.2 HIERARCHICAL TILTED EMPIRICAL RISK MINIMIZATION

MML and MMLFA achieve the best rare-class and overall classification accuracy, while decreasing the objective and gradient norm faster than baselines.Table 1 reports classification accuracy as mean±std; experiments use three LIBSVM datasets with about 30% label reshuffling and a 1:20 rare-to-common class ratio.Source: Table 1 · p. 9, §5.2 · p. 9

6 CONCLUSION

The paper concludes that momentum-based estimators track function values and gradients in stochastic multilevel optimization. It reports optimal sample complexity under standard smoothness, alongside normalized and batch-free variants.

6 CONCLUSION

The mini-batch method achieves optimal sample complexity O(ϵ^-4) under standard smoothness without the stronger average-smoothness conditions used in prior work.A normalized variant preserves the rate without problem-specific constants for hyperparameter setting; a first-order-approximation and clipping approach retains the rate with constant batch size.Source: §6 · p. 10
The proposed methods are reported effective in experiments on risk-averse portfolio optimization and hierarchical tilted empirical risk minimization.The methods use momentum-based estimators to track function values and gradients.Source: §6 · pp. 9–10

A PROOF OF LEMMA 3.4

The proof bounds the difference between function estimators at adjacent steps and then sums the resulting bounds.

A PROOF OF LEMMA 3.4

The proof considers the difference between function estimators at adjacent steps.The passage introduces this bound as the next step in the proof.Source: A PROOF OF LEMMA 3.4 · p. 12
The proof sequence moves from adjacent-step estimator differences to a summed bound.This connects the two stated proof steps.Source: A PROOF OF LEMMA 3.4 · p. 12
The proof then sums the bounds to obtain its stated result.The supplied passage does not specify the result further.Source: A PROOF OF LEMMA 3.4 · p. 12

B PROOF OF LEMMA 3.5

The proof uses induction, states a guarantee, and specifies a recursive choice for constants C_i.

B PROOF OF LEMMA 3.5

The lemma is proved by induction, with a guarantee stated in the proof.The supplied passages do not give the guarantee's content.Source: B PROOF OF LEMMA 3.5 · p. 13
For level i = n + 1, the proof establishes the next induction step.The passage introduces this level-specific argument before stating a further fact.Source: B PROOF OF LEMMA 3.5 · pp. 13–14
Setting C_0 = 1 gives C_(i+1) = 5 + 6L^2_f C_i for i ≥ 1.The passage also states C_1 = 5 + 6L^2_f.Source: B PROOF OF LEMMA 3.5 · pp. 13–14

C PROOF OF LEMMA 3.6

The proof defines the level-i composition and its gradient, then uses intermediate steps to establish its conclusion.

C PROOF OF LEMMA 3.6

F_i is defined as the composition f_i ◦ f_{i−1} ◦ … ◦ f_1 evaluated at x_k, with gradient ∇F_i(x_k).Source: C PROOF OF LEMMA 3.6 · p. 14
The proof next considers the definition of z_{k+1} before deriving a further intermediate result.Source: C PROOF OF LEMMA 3.6 · p. 15
The proof sums and rearranges intermediate bounds, using B_1 ≤ α^2 B_0 K to establish its conclusion.Source: C PROOF OF LEMMA 3.6 · p. 15

D PROOF OF THEOREM 3.7

These excerpts give proof steps for several theorems and lemmas, then identify a dependence issue in a prior analysis.

D PROOF OF THEOREM 3.7

Theorem 3.7’s proof invokes smoothness of F and sets B_0=1 and B_1=α^2K=O(K).Source: D PROOF OF THEOREM 3.7 · p. 16

E PROOF OF THEOREM 3.8

Theorem 3.8’s proof uses the update x_{k+1}=x_k−η z_k/∥z_k∥ and sets B_0=1, B_1=⌊K/4⌋, and α=1/2.Source: E PROOF OF THEOREM 3.8 · p. 17

F PROOF OF LEMMA 4.1

Lemma 4.1’s proof treats an operation as a projection and invokes non-expansiveness while bounding adjacent-step function-estimator differences.Source: F PROOF OF LEMMA 4.1 · pp. 18–19

G PROOF OF LEMMA 4.2

The Lemma 4.2 proof excerpt refers to a previously proved bound and then derives a further result.Source: G PROOF OF LEMMA 4.2 · p. 19

H PROOF OF LEMMA 4.3

The Lemma 4.3 proof excerpt refers to prior analysis and sums and rearranges its terms.Source: H PROOF OF LEMMA 4.3 · p. 20

I PROOF OF THEOREM 4.4

Theorem 4.4’s proof sets α=O(K^−1/2) and a learning rate of O(K^−1/2), using smoothness of F.Source: I PROOF OF THEOREM 4.4 · p. 20

J DEPENDENCE ISSUES IN PREVIOUS LITERATURE

The prior analysis is described as relying on an invalid inequality that treats dependent terms as if their expectations factorized.Source: J DEPENDENCE ISSUES IN PREVIOUS… · p. 21

K ADDITIONAL EXPERIMENTAL RESULTS

The additional experimental figures cover batch-size results for portfolio optimization and clipping-threshold sensitivity on two tasks.

K ADDITIONAL EXPERIMENTAL RESULTS

Figure 3 reports risk-averse portfolio optimization results across different batch sizes.Source: Figure 3 · p. 21
Figure 4 reports sensitivity to clipping threshold C on the portfolio task.Source: Figure 4 · p. 22
Figure 5 reports sensitivity to clipping threshold C on hierarchical TERM.Source: Figure 5 · p. 22