Title: Adaptive Variance Reduction for Stochastic Optimization under Weaker Assumptions
Venue: NeurIPS 2024
Accepted: yes
Keywords: Adaptive methods

variance reduction

finite-sum optimization

stochastic compositional optimization
Abstract: This paper explores adaptive variance reduction methods for stochastic optimization based on the STORM technique. Existing adaptive extensions of STORM rely on strong assumptions like bounded gradients and bounded function values, or suffer an additional $\mathcal{O}(\log T)$ term in the convergence rate. To address these limitations, we introduce a novel adaptive STORM method that achieves an optimal convergence rate of $\mathcal{O}(T^{-1/3})$ for non-convex functions with our newly designed learning rate strategy. Compared with existing approaches, our method requires weaker assumptions and attains the optimal convergence rate without the additional $\mathcal{O}(\log T)$ term. We also extend the proposed technique to stochastic compositional optimization, obtaining the same optimal rate of $\mathcal{O}(T^{-1/3})$. Furthermore, we investigate the non-convex finite-sum problem and develop another innovative adaptive variance reduction method that achieves an optimal convergence rate of $\mathcal{O}(n^{1/4} T^{-1/2} )$, where $n$ represents the number of component functions. Numerical experiments across various tasks validate the effectiveness of our method.

Adaptive Variance Reduction for Stochastic Optimization under Weaker Assumptions

Abstract

This paper explores adaptive variance reduction methods for stochastic optimization based on the STORM technique. Existing adaptive extensions of STORM rely on strong assumptions like bounded gradients and bounded function values, or suffer an additional O (log T ) term in the convergence rate. To address these limitations, we introduce a novel adaptive STORM method that achieves an optimal convergence rate of O ( T -1 / 3 ) for non-convex functions with our newly designed learning rate strategy. Compared with existing approaches, our method requires weaker assumptions and attains the optimal convergence rate without the additional O (log T ) term. We also extend the proposed technique to stochastic compositional optimization, obtaining the same optimal rate of O ( T -1 / 3 ) . Furthermore, we investigate the non-convex finite-sum problem and develop another innovative adaptive variance reduction method that achieves an optimal convergence rate of O ( n 1 / 4 T -1 / 2 ) , where n represents the number of component functions. Numerical experiments across various tasks validate the effectiveness of our method.

1 Introduction

This paper investigates the stochastic optimization problem

where f : R d ↦→ R is a smooth non-convex function. We assume that only noisy estimations of its gradient ∇ f ( x ) can be accessed, denoted as ∇ f ( x ; ξ ) , where ξ represents the random sample drawn from a stochastic oracle such that E [ ∇ f ( x ; ξ )] = ∇ f ( x ) .

Although aforementioned methods are equipped with convergence guarantees, their analyses rely on delicate configurations of hyper-parameters, such as the learning rate and the momentum parameter. To set them properly, the algorithm typically needs to know the value of the smoothness parameter L ,

Problem (1) has been comprehensively investigated in the literature [Duchi et al., 2011, Kingma and Ba, 2015, Loshchilov and Hutter, 2017], and it is well-known that the classical stochastic gradient descent (SGD) achieves a convergence rate of O ( T -1 / 4 ) , where T denotes the iteration number [Ghadimi and Lan, 2013]. To further improve the convergence rate, variance reduction methods have been developed, and attain an improved rate of O ( T -1 / 3 ) under a slightly stronger smoothness assumption [Fang et al., 2018, Wang et al., 2019]. However, these methods necessitate the use of a huge batch size in each iteration, which is often impractical to use. To eliminate the need for large batches, a momentum-based variance reduction method - STORM [Cutkosky and Orabona, 2019] is introduced, which achieves a convergence rate of O ( T -1 / 3 log T ) .

the gradient upper bound G , and the variance upper bound σ , which are often unknown in practice. Specifically, most algorithms require the learning rate η t smaller than O (1 /L ) , and for the STORM method, setting the momentum parameter to O ( L 2 η 2 t ) is crucial for ensuring convergence [Cutkosky and Orabona, 2019].

To overcome this limitation, many adaptive algorithms have been developed, aiming to obtain convergence guarantees without prior knowledge of problem-dependent parameters such as L , G and σ . Based on the STORM method, Levy et al. [2021] develop the STORM+ algorithm, which attains the optimal O ( T -1 / 3 ) convergence rate under the assumption of bounded function values and gradients. To remove the need for the bounded function values assumption, Liu et al. [2022] propose the META-STORM algorithm to attain an O ( T -1 / 3 log T ) convergence rate, but it still requires the bounded gradients assumption and includes an additional O (log T ) term. In summary, despite advancements in this field, existing adaptive STORM-based methods either depend on strong assumptions or suffer an extra O (log T ) term compared with the lower bound [Arjevani et al., 2023]. Hence, a fundamental question to be addressed is:

Is it possible to develop an adaptive STORM method that achieves the optimal convergence rate for non-convex functions under weaker assumptions?

We give an affirmative answer to the above question by devising a novel optimal Adaptive STORM method (Ada-STORM). The learning rate of our algorithm is set to be inversely proportional to a specific power of the iteration number T in the initial iterations, and then changes adaptively based on the cumulative sum of past gradient estimations. In this way, we are able to adjust the learning rate dynamically according to the property of stochastic gradients, and ensure a small learning rate in the beginning. Leveraging this strategy, Ada-STORM achieves an optimal convergence rate of O ( T -1 / 3 ) for non-convex functions. Notably, our analysis does not require the function to have bounded values and bounded gradients, which is a significant advancement over existing methods [Levy et al., 2021, Liu et al., 2022]. Additionally, our convergence rate does not contain the extra O (log T ) term, which is often present in STORM-based methods [Cutkosky and Orabona, 2019, Liu et al., 2022]. To highlight the versatility of our approach and its potential impact in the field of stochastic optimization, we further extend our technique to develop optimal adaptive methods for compositional optimization.

Finally, we investigate adaptive variance reduction for the non-convex finite-sum problems. Inspired by SAG algorithm [Roux et al., 2012], we incorporate an additional term in the STORM estimator, which measures the difference of past gradients between the selected component function and the overall objective. By changing the learning rate according to the sum of past gradient estimations, we are able to obtain an optimal convergence rate of O ( n -1 / 4 T -1 / 2 ) for finite-sum problems in an adaptive manner, where n is the number of component functions. Our result is better than the previous convergence rate of O ( n -1 / 4 T -1 / 2 log( nT )) obtained by adaptive SPIDER method [Kavis et al., 2022]. In summary, compared with existing methods, this paper enjoys the following advantages:

For stochastic non-convex optimization, our method achieves the optimal convergence rate of O ( T -1 / 3 ) under more relaxed assumptions. Specifically, it does not require the bounded function values or the bounded gradients, and does not include the additional O (log T ) term in the convergence rate.

Our learning rate design and the analysis exhibit broad applicability. We substantiate this claim by obtaining an optimal rate of O ( T -1 / 3 ) for stochastic compositional optimization, using the technique proposed in this paper.

For non-convex finite-sum optimization, we further improve our adaptive algorithm to attain an optimal convergence rate of O ( n 1 / 4 T -1 / 2 ) , which outperforms the previous result by eliminating the O (log( nT )) factor.

A comparison between our method and other STORM-based algorithms is shown in Table 1. Numerical experiments on different tasks also validate the effectiveness of the proposed method.

2 Related work

This section briefly reviews related work on stochastic variance reduction methods and adaptive stochastic algorithms.

2.1 Stochastic variance reduction methods

Variance reduction has been widely used in stochastic optimization to reduce the gradient estimation error and thus improve the convergence rates. The idea of variance reduction can be traced back to the SAG algorithm [Roux et al., 2012], which incorporates a memory of previous gradient values to ensure variance reduction and achieves a linear convergence rate for strongly convex finite-sum optimization. To avoid the storage of past gradients, SVRG [Zhang et al., 2013, Johnson and Zhang, 2013] proposes to calculate the full gradient periodically, obtaining the same convergence rate as the SAG algorithm. Subsequent advancement has been made by the SARAH method [Nguyen et al., 2017], which derives better convergence for smooth convex functions.

In the context of non-convex objectives, Fang et al. [2018] introduce the SPIDER estimator, which improves the convergence rate from O ( T -1 / 4 ) to O ( T -1 / 3 ) in stochastic settings, and to O ( n 1 / 4 T -1 / 2 ) in finite-sum scenarios, with n representing the number of components in the finite-sum. Following this, the SpiderBoost algorithm [Wang et al., 2019] refines the SPIDER approach by employing a larger constant step size and adapting it for composite optimization problems. However, a common limitation among these methods is their reliance on large batch sizes for each iteration, posing practical challenges due to high computational demands. To mitigate this issue, Cutkosky and Orabona [2019] introduce the STORM method, a momentum-based technique that achieves an O ( T -1 / 3 log T ) convergence rate without using large batches. Concurrently, Tran-Dinh et al. [2019] obtain the same result using a similar algorithm but through a different analysis.

2.2 Adaptive stochastic algorithms

For stochastic optimization, it is well-known that the SGD algorithm can obtain a convergence rate of O ( T -1 / 4 ) for non-convex objective functions with well-designed learning rates [Ghadimi and Lan, 2013]. Instead of using pre-defined iteration-based learning rates, many stochastic methods propose to adjust the learning rate based on past stochastic gradients. One of the foundational works is the AdaGrad algorithm [Duchi et al., 2011], which proves to be effective for sparse data. Further advancements include RMSprop [Tieleman and Hinton, 2012] and Adam [Kingma and Ba, 2015], demonstrating broad effectiveness across a wide range of machine learning problems. Later, the Super-Adam [Huang et al., 2021] algorithm further improves the Adam algorithm via the variance reduction technique STORM [Cutkosky and Orabona, 2019] and obtains a convergence rate of O ( T -1 / 3 log T ) . Nevertheless, to obtain the corresponding convergence rates, these methods still require knowledge of certain problem-dependent parameters to set hyper-parameters accurately, hence

Algorithm 1 STORM Algorithm

1: Input: time step T , initial point x 1

3: Set hyper-parameters β t and η t

2: for time step t = 1 to T do

4: Compute v t according to equation (2)

6: end for

5: Update the decision variable: x t +1 = x t -η t v t

7: Choose τ uniformly at random from { 1 , . . . , T }

8: Return x τ

not adaptive. 2 To solve this problem, many research aims to develop fully adaptive SGD methods that maintain the optimal convergence rate without knowing problem-specific parameters [Orabona, 2014, Chen et al., 2022, Carmon and Hinder, 2022, Ivgi et al., 2023, Yang et al., 2023].

Recently, adaptive adaptations of STORM have received considerable attention. A notable development is the introduction of STORM+ [Levy et al., 2021], which presents a fully adaptive version of STORM while attaining an optimal convergence rate. To circumvent the bounded function values assumption in STORM+, the META-STORM [Liu et al., 2022] approach is developed, equipped with a nearly optimal bound. However, META-STORM still requires the bounded gradients assumption, and it includes an additional O (log T ) in the convergence rate. Consequently, adaptive STORM with the optimal convergence rate and under mild assumptions still needs further explorations.

3 Adaptive variance reduction for non-convex optimization

In this section, we develop an adaptive STORM method for non-convex functions. We first outline the assumptions used, and then present our proposed method and analyze its convergence rate.

3.1 Assumptions

We introduce the following assumptions, which are standard and commonly adopted in the stochastic optimization [Nguyen et al., 2017, Fang et al., 2018, Cutkosky and Orabona, 2019, Li et al., 2021].

Assumption 1 (Average smoothness)

Assumption 2 (Bounded variance)

Assumption 3 f ∗ = inf x f ( x ) ≥ -∞ and f ( x 1 ) -f ∗ ≤ ∆ f for the initial solution x 1 .

Note that some additional assumptions are required in other STORM-based methods. Specifically, STORM [Cutkosky and Orabona, 2019], STORM+ [Levy et al., 2021], and META-STORM [Liu et al., 2022] assume the bounded gradients. Moreover, STORM+ makes an additional assumption of the bounded function values.

3.2 The proposed method

In this subsection, we aim to develop an adaptive STORM method that achieves an optimal convergence rate for non-convex functions under weaker assumptions. Our algorithm framework is the same as the original STORM [Cutkosky and Orabona, 2019], and the only difference is the setup of the momentum parameter β t and the learning rate η t . First, we present the STORM algorithm in Algorithm 1.

The core idea of STORM lies in a carefully devised variance reduced estimator v t , which effectively tracks the gradient ∇ f ( x t ) . For the first iteration ( t = 1 ), we set v 1 = ∑ B 0 i =1 1 B 0 ∇ f ( x 1 ; ξ i 1 ) , which is estimated within a batch B 0 = T 1 / 3 . Note that we use large batch only in the first iteration, and constant batch size in other iterations. In subsequent iterations ( t ≥ 2 ), estimator v t is defined as:

where the first two terms are similar to the momentum SGD, and the last term serves as the error correction, which ensures the variance reduction effect. By choosing the values of β t and η t carefully, STORM ensures that the estimation error E [ ∥ v t -∇ f ( x t ) ∥ 2 ] would decrease gradually. In the original STORM paper, these parameters are set up as:

where k = O ( G 2 / 3 L -1 ) , w = O ( G 2 ) and c = O ( L 2 ) . The settings of these hyper-parameters are crucial to the convergence analysis of STORM. However, it is worth noting that L is the smoothness parameter and G is the gradient upper bound, which are often difficult to determine in practice.

To address this problem, our approach defines the hyper-parameters as follows:

where 0 < α < 1 / 3 . Notably, our method does not rely on the parameters L and G , and also does not need the bounded gradients or bounded function values assumptions that are common in other methods. Although our formulation initially requires knowledge of the iteration number T , this can be effectively circumvented using the doubling trick, which will be explained later. The above learning rate η t can also be expressed in an alternative, more illustrative manner:

This formulation ensures that the learning rate starts sufficiently small in the initial stages and then changes dynamically based on the gradient estimator v t . This design makes our learning rate setup and convergence analysis distinctly different from previous methods. Next, we present the following theoretical guarantee for our algorithm.

Theorem 1 Under Assumptions 1, 2 and 3, Algorithm 1 with hyper-parameters in equation (3) guarantees that:

Remark: To ensure that E [ ∥∇ f ( x τ ) ∥ ] ≤ ϵ , the overall complexity is O ( ϵ -3 ) , which is known to be optimal up to constant factors [Arjevani et al., 2023]. Compared with existing STORM-based algorithms [Cutkosky and Orabona, 2019, Levy et al., 2021, Liu et al., 2022], our method does not have the extra O (log T ) term in the convergence rate, and our analysis does not require bounded gradients or bounded function values assumptions. Also note that the selection of α does not affect the order of T , and larger α leads to better dependence on parameter L and worse reliance on parameters ∆ and σ . Considering we require that 0 < α < 1 / 3 , we can simply set α = 0 . 3 in practice.

3.3 The doubling trick

While we have attained the optimal convergence rate using the proposed adaptive STORM method, it requires knowing the total number of iterations T in advance. Here, we show that we can avoid this requirement by using the doubling trick, which divides the algorithm into several stages and increases

the iteration number in each stage gradually. Specifically, we design a multi-stage algorithm over k = { 1 , 2 , · · · , K } stages. At the beginning of each new stage, we reset x t = x 0 . In each stage k , the STORM algorithm is executed for 2 k -1 iterations, effectively doubling the iteration numbers after each stage. In any step t , we first identify the current stage as 1 + ⌊ log t ⌋ and then calculate the iteration number for this stage as I t = 2 ⌊ log t ⌋ . Then, we can set the hyper-parameters as:

This approach eliminates the need to predetermine the iteration number T . By using the doubling trick, we can still obtain the same optimal convergence rate as stated in the following theorem.

Theorem 2 Under Assumptions 1, 2 and 3, Algorithm 1 with hyper-parameters in equation (4) guarantees that:

4 Extension to stochastic compositional optimization

To demonstrate the broad applicability of our proposed technique, we extend it to stochastic compositional optimization [Wang et al., 2017a,b, Yuan et al., 2019, Zhang and Xiao, 2019, 2021, Jiang et al., 2023, 2024a], formulated as:

where f and g are smooth functions. We assume that we can only access to unbiased estimations of ∇ f ( x ) , ∇ g ( x ) and g ( x ) , denoted as ∇ f ( x ; ξ ) , ∇ g ( x ; ζ ) and g ( x ; ζ ) . Here ξ and ζ symbolize the random sample drawn for a stochastic oracle such that E [ ∇ f ( x ; ξ )] = ∇ f ( x ) , E [ g ( x ; ζ )] = g ( x ) , and E [ ∇ g ( x ; ζ )] = ∇ g ( x ) .

Existing variance reduction methods [Hu et al., 2019, Zhang and Xiao, 2019, Qi et al., 2021] are able to obtain optimal O ( T -1 / 3 ) convergence rates for problem (5), but they require the knowledge of smoothness parameter and the gradient upper bound to set up hyper-parameters. In this section, we aim to achieve the same optimal convergence rate without prior knowledge of problem-dependent parameters. We develop our adaptive algorithm for this problem as follows. In each step t , the algorithm maintains an inner function estimator u t in the style of STORM, i.e.,

Then, we construct a gradient estimator v t based on u t also in the style of STORM:

After that, we apply gradient descent using the gradient estimator v t . The whole algorithm is presented in Algorithm 2, and hyper-parameters are set the same as in equation (3). For the first iteration, we simply set u 1 = ∑ B 0 i =1 1 B 0 g ( x 1 ; ζ i 1 ) and v 1 = ∑ B 0 i =1 1 B 0 ∇ f ( u 1 ; ξ i 1 ) ∇ g ( x 1 ; ζ 1 ) , where B 0 = T 1 / 3 . Next, we list common assumptions used in the literature of compositional optimization [Wang et al., 2017a,b, Yuan et al., 2019, Zhang and Xiao, 2019, 2021].

Assumption 4 (Average smoothness and Lipschitz continuity)

Assumption 5 (Bounded variance)

Algorithm 2 Compositional STORM

1: Input: time step T , initial point x 1

3: Compute u t according to equation (6)

2: for time step t = 1 to T do

4: Compute v t according to equation (7)

6: end for

5: Update the decision variable: x t +1 = x t -η t v t

7: Choose τ uniformly at random from { 1 , . . . , T }

8: Return x τ

Assumption 6 F ∗ = inf x F ( x ) ≥ -∞ and F ( x 1 ) -F ∗ ≤ ∆ F for the initial solution x 1 .

Remark: In Assumption 4, we further require standard Lipschitz continuity assumption, which is essential and widely required in the literature for stochastic compositional optimization [Wang et al., 2017b, Yuan et al., 2019, Jiang et al., 2022a,b]. This assumption is inherently introduced by the compositional optimization itself rather than by our adaptive techniques.

With the above assumptions, our algorithm enjoys the following guarantee.

Theorem 3 Under Assumptions 4, 5 and 6, our Algorithm 2 ensures that:

Remark: This rate matches the state-of-the-art (SOTA) results in stochastic compositional optimization [Hu et al., 2019, Zhang and Xiao, 2019, Qi et al., 2021], and our method achieve this in an adaptive manner. Note that our convergence rate aligns with the lower bound for single-level problems [Arjevani et al., 2023] and is thus unimprovable.

5 Adaptive variance reduction for finite-sum optimization

In this section, we further improve our adaptive variance reduction method to obtain an enhanced convergence rate for non-convex finite-sum optimization, which is in the form of

where each f i ( · ) is a smooth non-convex function. Existing adaptive method for this problem [Kavis et al., 2022] achieves a convergence rate of O ( n 1 / 4 T -1 / 2 log( nT )) based on the variance reduction technique SPIDER [Fang et al., 2018], suffering from an extra O (log( nT )) term compared with the corresponding lower bound [Fang et al., 2018, Li et al., 2021].

To obtain the optimal convergence rate for finite-sum optimization, we incorporate techniques from the SAG algorithm [Roux et al., 2012] into the STORM estimator. Specifically, in each step t , we start by randomly sample i t from the set { 1 , 2 , · · · , n } . Then, we construct a variance reduction gradient estimator as

where the first three terms align with the original STORM method, and the last term, inspired by the SAG algorithm, deals with the finite-sum structure. Here, g t tracks the gradient as

By such a design, we can ensure that the estimation error E [ ∥ v t -∇ F ( x t ) ∥ 2 ] reduces gradually. The whole algorithm is stated in Algorithm 3. In this case, we set the hyper-parameters as:

̸

where 0 < α < 1 / 3 . The learning rate η t is non-increasing and changes according to the gradient estimations, and the momentum parameter β remains unchanged throughout the learning process. Next, we show that our method enjoys the optimal convergence rate under the following assumptions, which are standard and widely adopted in existing literature [Fang et al., 2018, Wang et al., 2019, Li et al., 2021].

Assumption 7 (Smoothness) For each i ∈ { 1 , 2 , · · · , m } , function f i is L -smooth such that

Assumption 8 F ∗ = inf x F ( x ) ≥ -∞ and F ( x 1 ) -F ∗ ≤ ∆ F for the initial solution x 1 .

Theorem 4 Under Assumptions 7 and 8, our Algorithm 3 guarantees that:

Remark: Our result matches the lower bound for non-convex finite-sum problems [Fang et al., 2018, Li et al., 2021], and makes an improvement over the existing adaptive method, i.e., AdaSpider [Kavis et al., 2022]. Specifically, the convergence rate of the AdaSpider algorithm is O ( n 1 / 4 T -1 / 2 ( L 2 +∆ F ) · log (1 + nTL ) ) , and our result is better than theirs when 1 4 < α < 1 3 .

We can avoid storing past gradients by following the SVRG method [Zhang et al., 2013, Johnson and Zhang, 2013] to compute the full gradient periodically and incorporate it into STORM estimator. Instead of storing the past gradients as in SAG algorithm, we can avoid this storage cost by incorporating elements from the SVRG method. Specifically, we compute a full batch gradient at the first step and every I iteration (we set I = n ):

For other iterations, we randomly select an index i t from the set { 1 , 2 , · · · , n } and compute:

Note that the first three terms match the original STORM estimator, and the last term, inspired from SVRG, deals with the finite-sum structure. Compared with equation (8) in Algorithm 3, the difference is that we use ( ∇ f i t ( x τ ) -∇ f ( x τ )) instead of ( g i t t -1 n ∑ n i =1 g i t ) in the last term. The detailed procedure is outlined in Algorithm 4. This strategy maintains the same optimal rate, as stated below:

Theorem 5 Under Assumptions 7 and 8, our Algorithm 4 guarantees that:

Remark: The obtained convergence rate is in the same order as the results in Theorem 4, and Algorithm 4 does not require storing past gradients anymore.

0

50

100

150

200

Epoch

-

6

-

4

-

2

0

2

Training loss (log scale)

(a) Training loss

0

50

100

150

200

Epoch

90

92

94

96

98

100

Training accuracy

(b) Training accuracy

0

50

100

150

200

Epoch

-

1

.

0

-

0

.

5

0

.

0

0

.

5

Testing loss (log scale)

(c) Testing loss

0

50

100

150

200

Epoch

85

87

89

91

93

Testing accuracy

(d) Testing accuracy

Adam

AdaBelief

SGD

STORM

STORM+

META-STORM

Ada-STORM

6 Experiments

In this section, we evaluate the performance of the proposed Ada-STORM method via numerical experiments on image classification tasks and language modeling tasks. In the experiments, we compare our method with STORM [Cutkosky and Orabona, 2019], STORM+ [Levy et al., 2021] and METASTORM [Liu et al., 2022], as well as SGD, Adam [Kingma and Ba, 2015] and AdaBelief [Zhuang et al., 2020]. We use the default implementation of SGD and Adam from Pytorch [Paszke et al., 2019]. For STORM+, we follow its original implementation 3 , and build STORM, META-STORM and our Ada-STORM based on it. When it comes to hyper-parameter tuning, we simply set α = 0 . 3 for our algorithm. For other methods, we either set the hyper-parameters as recommended in the original papers or tune them by grid search. For example, we search the learning rate of SGD, Adam and AdaBelief from the set { 1 e -5 , 1 e -4 , 1 e -3 , 1 e -2 , 1 e -1 } and select the best one for each method. All the experiments are conducted on eight NVIDIA Tesla V100 GPUs.

Image classification task First, we conduct numerical experiments on multi-class image classification tasks to evaluate the performance of the proposed method. Specifically, we train ResNet18 and ResNet34 models [He et al., 2016] on the CIFAR-10 and CIFAR-100 datasets [Krizhevsky, 2009] respectively. For all optimizers, we set the batch size as 256 and train for 200 epochs. We plot the loss value and the accuracy against the epochs on the CIFAR-10 and CIFAR-100 datasets in Figure 1 and Figure 2. It is observed that, for training loss and training accuracy, our Ada-STORM algorithm achieves comparable performance with respect to other methods, and it outperforms the others in terms of testing loss and thus obtains a better testing accuracy.

Language modeling task Then, we perform experiments on language modeling tasks. Concretely, we train a 2-layer Transformer [Vaswani et al., 2017] over the WiKi-Text2 dataset [Merity, 2016]. We use 256 dimensional word embeddings, 512 hidden unites and 2 heads. The batch size is set as 20 and all methods are trained for 40 epochs with dropout rate 0 . 1 . We also clip the gradients by norm 0 . 25 in case of the exploding gradient. We report both the loss and perplexity versus the number of epochs in Figure 3. From the results, we observe that our method converges more quickly than other

95

0

50

100

150

200

Epoch

-

4

-

2

0

Training loss (log scale)

(a) Training loss

0

50

100

150

200

Epoch

90

92

94

96

98

100

Training accuracy

(b) Training accuracy

0

50

100

150

200

Epoch

0

.

25

0

.

50

0

.

75

1

.

00

1

.

25

Testing loss (log scale)

(c) Testing loss

0

50

100

150

200

Epoch

66

68

70

72

74

Testing accuracy

(d) Testing accuracy

Adam

AdaBelief

SGD

STORM

STORM+

META-STORM

Ada-STORM

0

10

20

30

40

Epoch

5

6

7

Training loss

(a) Training loss

0

10

20

30

40

Epoch

500

1000

Training perplexity

(b) Training perplexity

0

10

20

30

40

Epoch

5

.

0

5

.

5

6

.

0

Testing loss

(c) Testing loss

0

10

20

30

40

Epoch

200

300

400

500

Testing perplexity

(d) Testing perplexity

Adam

AdaBelief

SGD

STORM

STORM+

META-STORM

Ada-STORM

methods and obtains a slightly better perplexity compared with others, indicating the effectiveness of the proposed method.

7 Conclusion

In this paper, we propose an adaptive STORM method to achieve the optimal convergence rate for nonconvex functions. Compared with existing methods, our algorithm requires weaker assumptions and does not have the additional O (log T ) term in the convergence rate. The proposed technique can also be employed to develop optimal adaptive algorithms for compositional optimization. Furthermore, we investigate an adaptive method for non-convex finite-sum optimization, obtaining an improved convergence rate of O ( n 1 / 4 T -1 / 2 ) . Given that STORM algorithm has already been used in many areas such as bi-level optimization [Yang et al., 2021], federated learning [Das et al., 2022], min-max optimization [Xian et al., 2021], sign-based optimization [Jiang et al., 2024b], etc., the proposed methods may also inspire the development of adaptive algorithms in these fields.

References

Arjevani, Y. Carmon, J. C. Duchi, D. J. Foster, N. Srebro, and B. Woodworth. Lower bounds for non-convex stochastic optimization. Mathematical Programming , 199(1):165-214, 2023.

Carmon and O. Hinder. Making SGD parameter-free. In Proceedings of the 35rd Annual Conference on Learning Theory , pages 2360-2389, 2022.

Chen, J. Langford, and F. Orabona. Better parameter-free stochastic optimization with ODE updates for coin-betting. In Proceedings of the 36th AAAI Conference on Artificial Intelligence , pages 6239-6247, 2022.

Cutkosky and F. Orabona. Momentum-based variance reduction in non-convex SGD. In Advances in Neural Information Processing Systems 32 , pages 15210-15219, 2019.

Das, A. Acharya, A. Hashemi, S. Sanghavi, I. S. Dhillon, and U. Topcu. Faster non-convex federated learning via global and local momentum. In Proceedings of the 38th Conference on Uncertainty in Artificial Intelligence , pages 496-506, 2022.

Duchi, E. Hazan, and Y. Singer. Adaptive subgradient methods for online learning and stochastic optimization. Journal of Machine Learning Research , 12(61):2121-2159, 2011.

Fang, C. J. Li, Z. Lin, and T. Zhang. SPIDER: Near-optimal non-convex optimization via stochastic path-integrated differential estimator. In Advances in Neural Information Processing Systems 31 , pages 689-699, 2018.

Ghadimi and G. Lan. Stochastic first- and zeroth-order methods for nonconvex stochastic programming. SIAM Journal on Optimization , 23(4):2341-2368, 2013.

He, X. Zhang, S. Ren, and J. Sun. Deep residual learning for image recognition. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition , pages 770-778, 2016.

Hu, C. J. Li, X. Lian, J. Liu, and H. Yuan. Efficient smooth non-convex stochastic compositional optimization via stochastic recursive gradient descent. In Advances in Neural Information Processing Systems 32 , pages 6543-6553, 2019.

Huang, J. Li, and H. Huang. Super-adam: Faster and universal framework of adaptive gradients. In Advances in Neural Information Processing Systems 34 , pages 9074-9085, 2021.

Ivgi, O. Hinder, and Y. Carmon. Dog is SGD's best friend: A parameter-free dynamic step size schedule. In Proceedings of the 40th International Conference on Machine Learning , pages 14465-14499, 2023.

Jiang, G. Li, Y. Wang, L. Zhang, and T. Yang. Multi-block-single-probe variance reduced estimator for coupled compositional optimization. In Advances in Neural Information Processing Systems 35 , pages 32499-32511, 2022a.

Jiang, B. Wang, Y. Wang, L. Zhang, and T. Yang. Optimal algorithms for stochastic multi-level compositional optimization. In Proceedings of the 39th International Conference on Machine Learning , pages 10195-10216, 2022b.

Jiang, J. Qin, L. Wu, C. Chen, T. Yang, and L. Zhang. Learning unnormalized statistical models via compositional optimization. In Proceedings of the 40th International Conference on Machine Learning , pages 15105-15124, 2023.

Jiang, S. Yang, W. Yang, Y. Wang, Y. Wan, and L. Zhang. Projection-free variance reduction methods for stochastic constrained multi-level compositional optimization. In Proceedings of the 41st International Conference on Machine Learning , pages 21962-21987, 2024a.

Jiang, S. Yang, W. Yang, and L. Zhang. Efficient sign-based optimization: Accelerating convergence via variance reduction. In Advances in Neural Information Processing Systems 37 , 2024b.

Johnson and T. Zhang. Accelerating stochastic gradient descent using predictive variance reduction. In Advances in Neural Information Processing Systems 26 , 2013.

Kavis, S. Skoulakis, K. Antonakopoulos, L. T. Dadi, and V. Cevher. Adaptive stochastic variance reduction for non-convex finite-sum minimization. In Advances in Neural Information Processing Systems 35 , pages 23524-23538, 2022.

P. Kingma and J. L. Ba. Adam: A method for stochastic optimization. In International Conference on Learning Representations , 2015.

Krizhevsky. Learning multiple layers of features from tiny images. Masters Thesis, Deptartment of Computer Science, University of Toronto , 2009.

Y. Levy, A. Kavis, and V. Cevher. STORM+: Fully adaptive SGD with recursive momentum for nonconvex optimization. In Advances in Neural Information Processing Systems 34 , pages 20571-20582, 2021.

Li, H. Bao, X. Zhang, and P. Richtarik. Page: A simple and optimal probabilistic gradient estimator for nonconvex optimization. In Proceedings of the 38th International Conference on Machine Learning , pages 6286-6295, 2021.

Liu, T. D. Nguyen, T. H. Nguyen, A. Ene, and H. L. Nguyen. Meta-storm: Generalized fullyadaptive variance reduced SGD for unbounded functions. ArXiv e-prints , arXiv:2209.14853, 2022.

Loshchilov and F. Hutter. SGDR: Stochastic gradient descent with warm restarts. In International Conference on Learning Representations , 2017.

B. McMahan and M. Streeter. Adaptive bound optimization for online convex optimization. In Proceedings of the 23rd Annual Conference on Learning Theory , 2010.

Merity. The wikitext long term dependency language modeling dataset. Salesforce Metamind , 9, 2016.

M. Nguyen, J. Liu, K. Scheinberg, and M. Takac. SARAH: A novel method for machine learning problems using stochastic recursive gradient. In Proceedings of the 34th International Conference on Machine Learning , pages 2613-2621, 2017.

Orabona. Simultaneous model selection and optimization through parameter-free stochastic learning. In Advances in Neural Information Processing Systems 27 , pages 1116-1124, 2014.

Paszke, S. Gross, F. Massa, A. Lerer, J. Bradbury, G. Chanan, T. Killeen, Z. Lin, N. Gimelshein, L. Antiga, A. Desmaison, A. Kopf, E. Yang, Z. DeVito, M. Raison, A. Tejani, S. Chilamkurthy, B. Steiner, L. Fang, J. Bai, and S. Chintala. Pytorch: An imperative style, high-performance deep learning library. In Advances in Neural Information Processing Systems 32 , 2019.

Qi, Z. Guo, Y. Xu, R. Jin, and T. Yang. An online method for a class of distributionally robust optimization with non-convex objectives. In Advances in Neural Information Processing Systems 34 , pages 10067-10080, 2021.

L. Roux, M. Schmidt, and F. R. Bach. A stochastic gradient method with an exponential convergence rate for finite training sets. In Advances in Neural Information Processing Systems 25 , pages 2672-2680, 2012.

Tieleman and G. Hinton. Lecture 6.5-rmsprop: Divide the gradient by a running average of its recent magnitude. COURSERA: Neural Networks for Machine Learning, 4 , pages 26-31, 2012.

Tran-Dinh, N. H. Pham, D. T. Phan, and L. M. Nguyen. Hybrid stochastic gradient descent algorithms for stochastic nonconvex optimization. ArXiv e-prints , arXiv:1905.05920, 2019.

Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, L. u. Kaiser, and I. Polosukhin. Attention is all you need. In Advances in Neural Information Processing Systems 30 , pages 5998-6008, 2017.

Wang, E. X. Fang, and H. Liu. Stochastic compositional gradient descent: algorithms for minimizing compositions of expected-value functions. Mathematical Programming , 161(1-2): 419-449, 2017a.

Wang, J. Liu, and E. X. Fang. Accelerating stochastic composition optimization. Journal of Machine Learning Research , 18:105:1-105:23, 2017b.

Wang, K. Ji, Y. Zhou, Y. Liang, and V. Tarokh. SpiderBoost and momentum: Faster variance reduction algorithms. In Advances in Neural Information Processing Systems 32 , pages 2406-2416, 2019.

Xian, F. Huang, Y. Zhang, and H. Huang. A faster decentralized algorithm for nonconvex minimax problems. In Advances in Neural Information Processing Systems 34 , pages 25865-25877, 2021.

Yang, K. Ji, and Y. Liang. Provably faster algorithms for bilevel optimization. In Advances in Neural Information Processing Systems 34 , pages 13670-13682, 2021.

Yang, X. Li, I. Fatkhullin, and N. He. Two sides of one coin: the limits of untuned SGD and the power of adaptive methods. In Advances in Neural Information Processing Systems 36 , pages 1206-1216, 2023.

Yuan, X. Lian, C. J. Li, J. Liu, and W. Hu. Efficient smooth non-convex stochastic compositional optimization via stochastic recursive gradient descent. In Advances in Neural Information Processing Systems 32 , pages 14905-14916, 2019.

Zhang and L. Xiao. A stochastic composite gradient method with incremental variance reduction. In Advances in Neural Information Processing Systems 32 , pages 9075-9085, 2019.

Zhang and L. Xiao. Multilevel composite stochastic optimization via nested variance reduction. SIAM Journal on Optimization , 31(2):1131-1157, 2021.

Zhang, M. Mahdavi, and R. Jin. Linear convergence with condition number independent access of full gradients. In Advances in Neural Information Processing Systems 26 , pages 980-988, 2013.

Zhuang, T. Tang, Y. Ding, S. C. Tatikonda, N. Dvornek, X. Papademetris, and J. Duncan. Adabelief optimizer: Adapting stepsizes by the belief in observed gradients. In Advances in Neural Information Processing Systems 33 , pages 18795-18806, 2020.

A Proof of Theorem 1

First, we introduce the following lemma, which is frequently used in our proof.

Lemma 1 Suppose c i is positive for i = { 1 , 2 , · · · , n } , and let 0 < α < 1 . We can ensure that:

Proof 1 The proof mainly follows McMahan and Streeter [2010] and a similar analysis also appears in Levy et al. [2021]. First, we prove the right part by Induction.

(1) For n = 1 , we can easily show that the hypothesis holds:

(2) Next, assuming that the hypothesis holds for n = t -1 , then we show that it also holds for n = t . Define Z = ∑ t i =1 c i and X = c t . For n = t , we have

Taking the derivative concerning x , we know that

which indicates that h ( X ) decreases as X increasing. Since 0 ≤ X ≤ Z , which implies that the hypothesis is true for n = t .

Combining (1) and (2), we finish the proof for the right part. Then, we give the proof of the left part as follows:

Thus, we finish the proof of this lemma.

Next, we can obtain the following guarantee for our algorithm.

Lemma 2 Our method enjoys the following guarantee:

Proof 2 According to the definition of estimator v t , we can deduce that:

Note that η t is independent of random variable ξ t +1 . So we can guarantee that:

After rearranging, we have:

Letting H t be the history to time t , i.e., H t = { ξ 1 , · · · , ξ t } , we ensure that

By summing up and noting that η t is non-increasing such that η t +1 ≤ η t , we have:

Since we use a large batch size in the first iteration, that is, B 0 = T 1 / 3 , we can now ensure that E H 1 [ ∥ v 1 -∇ f ( x 1 ) ∥ 2 ] ≤ σ 2 B 0 = σ 2 T 1 / 3 . Due to the fact that η 1 ≤ T -1 / 3 and β = T -2 / 3 , the first term of the above inequality is less than σ 2 . So, we can finally have:

Also, due to the smoothness of f ( x ) , we know that:

By summing up and re-arranging, we have:

Then, by using equation (13) and the fact that f ( x 1 ) -f ∗ ≤ ∆ f , we have:

which finishes the proof of Lemma 2.

To effectively bound each term in the above lemma, we divide the algorithm into two stages. Suppose that starting from iteration t = s , the condition ∑ t i =1 ∥ v i ∥ 2 ≥ T 1 / 3 begins to hold. We refer to iterations t = { 1 , 2 , · · · , s -1 } as the first stage, and t = { s, · · · , T } as the second stage.

Bounding LHS: In the first stage, we know that

For the second stage, our analysis leads to:

where the first inequality stems from ∑ s -1 t =1 ∥ v t ∥ 2 ≤ T 1 / 3 , the second inequality results from ( x + y ) α ≤ x α + y α for positive x, y and 0 < α < 1 / 3 , and the forth inequality applies Lemma 1. Next, we bound (A) , (B) , (C) as follows.

Bounding (A) : In the first stage, with η t = T -1 / 3 , β = T -2 / 3 , and ∑ s -1 t =1 ∥ v t ∥ 2 ≤ T 1 / 3 , we can derive:

For the second stage, the analysis gives:

where the second inequality uses Lemma 1. Then we also have:

where the inequality employs Young's inequality, such that xy ≤ 3 αx 1 3 α +(1 -3 α ) y 1 1 -3 α for positive x, y . Very similarly, we have:

where the second inequality employs Young's inequality, such that xy ≤ 2 α 1 -α x 1 -α 2 α + 1 -3 α 1 -α y 1 -α 1 -3 α for positive x, y . Combining all above, we know that

Bounding (B) : In the first stage, with the learning rate set at η t = T -1 / 3 and ∑ s -1 t =1 ∥ v t ∥ 2 ≤ T 1 / 3 , we observe:

For the second stage, our analysis reveals:

where the second inequality leverages Lemma 1, and we have:

where the inequality is due to Young's inequality, such that xy ≤ 2 αx 1 2 α + (1 -2 α ) y 1 1 -2 α for positive x, y . Very similarly, we have:

where the second inequality employs Young's inequality, such that xy ≤ α 1 -α x 1 -α α + 1 -2 α 1 -α y 1 -α 1 -2 α for positive x, y . Combining all the above, we know that

Bounding (C) : Given that β = T -2 / 3 and η t ≤ T -1 / 3 , we can easily know that (C) ≤ 2 σ 2 . So far, we bound all terms in Lemma 2, and we can deduce

where

With the definition of Γ , we know that

These suggest that

which indicates the following by applying Jensen's inequality:

Also, because of Jensen's inequality, we have:

Summing up, we have proven that

Finally, we finish our proof by introducing the following lemma.

Lemma 3 Suppose 0 < β < 1 , our method ensures that

Proof 3 First note that we have already proven the following in equation (11).

After rearranging the items, we can get the following:

By summing up, we have:

Since we use a large batch size in the first iteration, that is, B 0 = T 1 / 3 , we can now ensure that E [ ∥ v 1 -∇ f ( x 1 ) ∥ 2 ] ≤ σ 2 B 0 = σ 2 T 1 / 3 . Note that β = 1 T 2 / 3 , so first term equals to σ 2 T 1 / 3 and the second term reduces to 2 σ 2 T 1 / 3 . To this end, we ensure

Thus we finish the proof for this lemma.

Here, we bound the term 2 L 2 β ∑ T t =1 η 2 t ∥ v t ∥ 2 as follows. In the first stage, with η t = 1 T 1 / 3 , we have:

For the second stage, the analysis gives:

where the second inequality uses Lemma 1, and the third one employs Young's inequality, such that xy ≤ 2 αx 1 2 α +(1 -2 α ) y 1 1 -2 α for positive x, y . We also know that

where the second inequality uses Lemma 1, and the third one employs Young's inequality, such that xy ≤ α 1 -α x 1 -α α + 1 -2 α 1 -α y 1 -α 1 -2 α for positive x, y .

As a result, we have:

By integrating these findings, we can finally have:

where

with C 0 defined in equation (15). We find that larger α leads to better dependence on L and worse reliance on parameters ∆ and σ . For α → 1 3 , we can obtain that

Since we require 0 < α < 1 3 , in practice, we can use α = 0 . 3 instead, which leads to a convergence rate of O ( ∆ 5 / 7 f + σ 10 / 7 + L 5 / 3 T 1 / 3 ) .

B Proof of Theorem 2

Since 2 0 + 2 1 + · · · + 2 K -1 < 2 K , running the algorithm for T iterations guarantees at least K = ⌊ log( T ) ⌋ complete stages. In the theoretical analysis, we can simply use the output of the last complete stage K = ⌊ log( T ) ⌋ , which has been at least run for 2 K -1 ≥ T/ 4 iterations. According to the analysis of Theorem 1, we have already known that running the Algorithm 1 for T/ 4 iterations leads to the following guarantee:

which is on the same order of the original convergence rate.

C Proof of Theorem 3

According to equation (14), we have already proven that

Then we bound the term ∑ T t =1 η t ∥∇ F ( x t ) -v t ∥ 2 as follows:

Define that G t = ∇ f ( u t ) ∇ g ( x t ) , then we have:

For the term ∑ T t =1 E [ η t ∥ v t -G t ∥ 2 ] , following the very similar analysis of equation (11), we have the following guarantee:

That is to say:

where the last inequality due to the fact that β = T -2 / 3 , η t ≤ T -1 / 3 , and we use a large batch size T 1 / 3 in the first iteration. Next, we further bound the term ∑ T t =1 E [ η t -1 ∥ u t -u t -1 ∥ 2 ] . First, we can ensure that:

So we know that

So far, we have

Next, we can bound ∑ T t =1 η t ∥ u t -g ( x t ) ∥ 2 following equation (11), as:

So we can have

where the last inequality due to the fact that β = T -2 / 3 , η t ≤ T -1 / 3 , and we use a large batch size T 1 / 3 in the first iteration. Combining all, we have:

Treating ∆ F , C, L, σ as constant, the above inequality is very similar to Lemma 2. Thus following the very similar analysis after Lemma 2, we can show that:

According to previous analysis, we also have that

which is similar to Lemma 3, and leads to ∑ T t =1 E [ ∥∇ f ( x t ) -v t ∥ /T ] ≤ O ( T -1 / 3 ) following the same analysis. Combing all these together, we can deduce that

which finishes the proof of Theorem 3.

D Proof of Theorem 4

Due to the smoothness of F ( x ) , we have proven the following in equation (14):

For the LHS, we have the following guarantee:

Then, we bound the terms in the RHS. First, we have the following lemma.

Lemma 4 Define that z t = ∇ f i t ( x t ) -g i t t + 1 n ∑ n i =1 g i t , we have:

Proof 4 According to the definition of z t , the estimator v t can be expressed as

Note that E i t +1 [ z t +1 ] = ∇ F ( x t +1 ) , and we have:

Rearrange the items and multiply the both sides by η t , we can get the following:

Note that η t is non-increasing. By summing up, we have:

Since we use a full batch in the first iteration, we can finish the proof of this lemma.

Next, we bound two terms in the above lemma.

Proof 5

where the last equation is due to the fact that i t +1 is randomly sample from { 1 , 2 , · · · , n } . Note that we also have that

That is to say, we can ensure that

By rearranging and summing up, we have

Since we use a full batch n in the first iteration, the second term equals zero, and thus we obtain:

which leads to the result of this lemma.

Lemma 6 We have the following guarantee:

Proof 6

where the last inequality employs Young's inequality, such that xy ≤ 2 α 1 -α x 1 -α 2 α + 1 -3 α 1 -α y 1 -α 1 -3 α for positive x, y .

Lemma 7 We can ensure the following guarantee:

Proof 7

where the last inequality employs Young's inequality, such that xy ≤ α 1 -α x 1 -α α + 1 -2 α 1 -α y 1 -α 1 -2 α for positive x, y . Combing all these, we have already proven that

where

So we can know that:

which indicate that as well as

To finish the proof, we also have to show the following lemma.

Lemma 8

Proof 8 According to the previous proof, we know that:

Also, we can deduce that:

where the last inequality employs Young's inequality, such that xy ≤ α 1 -α x 1 -α α + 1 -2 α 1 -α y 1 -α 1 -2 α for positive x, y .

As a result, we can ensure that

E Proof of Theorem 5

The analysis is very similar to that of Theorem 4, and the difference only appears in Lemma 5. For this new method, we can also prove the same lemma:

Lemma 9

Proof 9 This time, we have z t +1 = ∇ f i t +1 ( x t +1 ) -∇ f i t +1 ( x τ ) + ∇ F ( x τ ) . And we can know that:

By summing up, we have

The other analysis is exactly the same as that of Theorem 4.

NeurIPS Paper Checklist

1. Claims

Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?

Answer: [Yes]

Justification: The claims presented in the abstract and introduction accurately represent the contributions and scope of the paper.

Guidelines:

The answer NA means that the abstract and introduction do not include the claims made in the paper.

The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.

The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.

It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.

2. Limitations

Question: Does the paper discuss the limitations of the work performed by the authors?

Answer: [Yes]

Justification: The theoretical results demonstrated in the paper rely on specific assumptions, which have been clearly stated in the main text.

Guidelines:

The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.

The authors are encouraged to create a separate "Limitations" section in their paper.

The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.

The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.

The authors should reflect on the factors that influence the performance of the approach. For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.

The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.

If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.

While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.

3. Theory Assumptions and Proofs

Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?

Answer: [Yes]

Justification: The paper provides assumptions and proofs for each theoretical result.

Guidelines:

The answer NA means that the paper does not include theoretical results.

All the theorems, formulas, and proofs in the paper should be numbered and crossreferenced.

All assumptions should be clearly stated or referenced in the statement of any theorems.

The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.

Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.

Theorems and Lemmas that the proof relies upon should be properly referenced.

4. Experimental Result Reproducibility

Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?

Answer: [Yes]

Justification: The paper discloses the information necessary to reproduce the main experimental results.

Guidelines:

The answer NA means that the paper does not include experiments.

If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.

If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.

Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.

While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example

(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.

(b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.

(c) If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).

(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.

5. Open access to data and code

Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?

Answer: [No]

Justification: Due to privacy concerns and ongoing research, we do not include the code.

Guidelines:

The answer NA means that paper does not include experiments requiring code.

Please see the NeurIPS code and data submission guidelines ( https://nips.cc/ public/guides/CodeSubmissionPolicy ) for more details.

While we encourage the release of code and data, we understand that this might not be possible, so 'No' is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).

The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines ( https: //nips.cc/public/guides/CodeSubmissionPolicy ) for more details.

The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.

The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only a subset of experiments are reproducible, they should state which ones are omitted from the script and why.

At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).

Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.

6. Experimental Setting/Details

Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?

Answer: [Yes]

Justification: The paper describes the training and testing details.

Guidelines:

The answer NA means that the paper does not include experiments.

The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them.

The full details can be provided either with the code, in appendix, or as supplemental material.

7. Experiment Statistical Significance

Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?

Answer: [Yes]

Justification: The paper reports error bars.

Guidelines:

The answer NA means that the paper does not include experiments.

The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).

The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)

The assumptions made should be given (e.g., Normally distributed errors).

It should be clear whether the error bar is the standard deviation or the standard error of the mean.

It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified.

For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).

If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.

8. Experiments Compute Resources

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

Answer: [Yes]

Justification: We have provided the relevant information.

Guidelines:

The answer NA means that the paper does not include experiments.

The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.

The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.

The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).

9. Code Of Ethics

Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines ?

Answer: [Yes]

Justification: The research conducted in the paper conforms with the NeurIPS Code of Ethics.

Guidelines:

The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.

If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.

The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction).

10. Broader Impacts

Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?

Answer: [NA]

Justification: This is primarily a theoretical paper with no potential negative social impact.

Guidelines:

The answer NA means that there is no societal impact of the work performed.

If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.

Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.

The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.

The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.

If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).

11. Safeguards

Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?

Answer: [NA]

Justification: The paper poses no such risks.

Guidelines:

The answer NA means that the paper poses no such risks.

Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.

Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.

We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.

12. Licenses for existing assets

Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?

Answer: [Yes]

Justification: The creators or original owners of assets used in the paper are properly credited and the license and terms of use explicitly are properly respected.

Guidelines:

The answer NA means that the paper does not use existing assets.

The authors should cite the original paper that produced the code package or dataset.

The authors should state which version of the asset is used and, if possible, include a URL.

The name of the license (e.g., CC-BY 4.0) should be included for each asset.

For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.

For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.

If this information is not available online, the authors are encouraged to reach out to the asset's creators.

13. New Assets

Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?

Answer: [NA]

Justification: The paper does not release new assets.

Guidelines:

The answer NA means that the paper does not release new assets.

Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.

The paper should discuss whether and how consent was obtained from people whose asset is used.

At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.

14. Crowdsourcing and Research with Human Subjects

Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?

Answer: [NA]

Justification: The paper does not involve crowdsourcing nor research with human subjects.

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.

According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.

15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects

Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?

Answer: [NA]

Justification: the paper does not involve crowdsourcing nor research with human subjects.

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.

For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper introduces novel adaptive variance reduction methods for stochastic optimization, building on the STORM technique. The proposed Ada-STORM method closes the $O(\\log T)$ gap and achieves an optimal convergence rate of $O(T^{-1/3})$ for non-convex functions under assumptions weaker than previous approaches. Furthermore, the paper extends the method to stochastic compositional optimization, maintaining a similar $O(T^{-1/3})$ convergence rate. For non-convex finite-sum problems, the authors develop another innovative adaptive algorithm that attains the optimal $O(n^{1/4}T^{-1/2})$ rate. Numerical experiments effectively demonstrate the effectiveness of the proposed method.\n\nSTRENGTHS:\n1. The proposed Ada-STORM method overcomes major limitations in existing adaptive STORM methods. It does not require bounded gradient and function value assumptions, and achieves the optimal convergence rates without the additional $O(\\log T)$ term. The problem investigated in the paper is challenging and the results are substantial.\n2. The theoretical analysis is novel and easy to follow. Extending the method to stochastic compositional and finite-sum optimization problems illustrates its flexibility and potential impact across a wide range of optimization tasks.\n3. The numerical experiments on various tasks (e.g., image classification, language modeling) validate the theoretical results and highlight the method's superior performance compared to other algorithms.\n\nWEAKNESSES:\n1. In Theorem 4, while the authors clearly state the convergence rate of the proposed method in terms of $T$, $\\Delta_F$, and $L$, they only present the convergence of previous adaptive finite-sum methods in terms of $T$. A comparison of other important constants such as $\\Delta_F$ and $L$ with existing methods could enhance the discussion.\n2. The content in Lines 240 to 242, especially Algorithm 4, should be included in the main body of the paper rather than being delayed to the Appendix.\n3. Some typos in the paper:\n- Line 145: $(t>2)$ should be $(t \\geq 2)$.\n- Line 516, the first equality in Proof 7: $L^2$ should be $L$.\n\nQUESTIONS:\n1. Regarding Line 401, could the authors explain why the first inequality holds?\n2. In Theorem 4, is the dependency on $\\Delta_F$ and $L$ better or worse than that of previous methods?\n\nLIMITATIONS:\nThis paper does not present any negative societal impacts.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2024-07-10"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposes a novel adaptive STORM method that achieves an optimal convergence rate of $O(T^{-3})$ for nonconvex stochastic optimization, which requires weaker assumptions and attains the optimal convergence rate without the additional $O(\\log T)$ term.\n\nSTRENGTHS:\nFor stochastic non-convex optimization, this paper achieves the optimal convergence rate under more relaxed assumptions, which does not require the bounded function values or the bounded gradients and does not include the additional $O(\\log T)$ term in the convergence rate. The proposed technique has also been extended to stochastic compositional optimization with the same optimal rate. For non-convex finite-sum optimization, this papaer further improve the adaptive algorithm to attain an optimal convergence rate, which outperforms the previous result by $O(\\log(nT))$ factor.\n\nWEAKNESSES:\nAlthough this paper is theoretically sound in general, there are still some questions need to be discussed:\n1.\tThe experimental part is relatively limited. The proposed algorithm is parameter free, so it is best to provide some experimental results to demonstrate whether these compared algorithms are sensitive to parameter changes.\n2.\tIn theory, the output {x_ \\tau} of the proposed algorithms is randomly selected from {1,..,T}. Is the output in the experiment randomly selected?\n\nQUESTIONS:\n1.\tThe experimental part is relatively limited. The proposed algorithm is parameter free, so it is best to provide some experimental results to demonstrate whether these compared algorithms are sensitive to parameter changes.\n2.\tIn theory, the output {x_ \\tau} of the proposed algorithms is randomly selected from {1,..,T}. Is the output in the experiment randomly selected?\n\nLIMITATIONS:\nThe authors discussed the limitations and there is no negative societal impact about this work.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 2, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-10"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper studies non-convex stochastic optimization under the assumption of mean-square smoothness. It introduces, Ada-STORM, a variant of the STORM algorithm, which achieves the optimal rate $O(T^{-1/3})$. Unlike vanilla STORM, Ada-STORM eliminates the $O(\\log T)$ factor and does not require the Lipschitz assumption. The algorithm's success hinges on a carefully designed adaptive learning rate schedule. Instead of the standard AdaGrad learning rate $\\eta_t \\approx (\\sum_{i=1}^t \\\\|g_i\\\\|^2)^{-1/3}$ used in STORM, Ada-STORM employs $\\eta_t = \\min\\\\{T^{-1/3}, T^{-(1-\\alpha)/3}(\\sum_{i=1}^t \\\\|v_i\\\\|^2)^{-\\alpha}\\\\}$ where $\\alpha < 1/3$ and $v_t$ is the STORM update.\n\nSTRENGTHS:\nThe proposed algorithm, Ada-STORM, achieves the optimal rate $O(T^{-1/3})$ without the log factor. Unlike STORM+ which also eliminates the log factor, this algorithm does not assume $f$ to be Lipschitz and bounded above. This improvement represents a significant contribution. Moreover, the required modification is straightforward: it involves merely adjusting the learning rate scheduler $\\eta_t$ and the momentum factor $\\beta_t$. I particularly enjoy the constant momentum constant $\\beta_t\\equiv T^{-2/3}$ rather than time-varying and dependent on $\\eta_t$.\n\nIn addition, the proposed learning rate scheduler is of interest in its own right. The paper extends the application of this learning rate technique to related problems, achieving optimal rates in composite optimization and finite-sum optimization (ERM). Finally, empirical experiments also validate Ada-STORM's performance, which matches or outperforms existing STORM variants and other state-of-the-art optimizers.\n\nWEAKNESSES:\nOne major concern pertains to the technical correctness of the proof of Theorem 1, in particular from line 440-444. I don't immediately see how $\\mathbb{E}[\\sum_{t=1}^{s-1} \\\\|v_t\\\\|^2] \\le C_0T^{1/3}$ implies $\\mathbb{E}[\\frac{1}{T}\\sum_{t=1}^{s-1}\\\\|v_t\\\\|] \\le \\sqrt{C_0}T^{-1/3}$, and similarly how the equation in line 441 implies line 442. My concern arise from the fact that $s$ is a random variable dependent on the entire history $\\mathcal{H}_T$. As a result, applying Jensen's inequality like $\\mathbb{E}[\\frac{1}{T}\\sum\\_{t=1}^s\\\\|v_t\\\\|] \\le \\sqrt{\\mathbb{E}[\\frac{1}{T}\\sum\\_{t=1}^s\\\\|v_t\\\\|^2]}$ seems to be inappropriate when $s$ is a random variable dependent $v_t$'s. It's likely that I miss something, and I encourage the authors to elaborate this part of the proof to make it clear. \n\nAnother concern is about the doubling trick. In each stage $k$, the guarantee is \n$$\\frac{1}{2^{k-1}} \\sum_{t=2^{k-1}}^{2^k} \\mathbb{E}\\\\|\\nabla f(x_t)\\\\| \\lesssim \\frac{\\Delta_f^{3/4}+O(1)}{(2^{k-1})^{1/3}}.$$ \nSince the algorithm is not resetting $x_t=x_0$ in the new stage, the parameter $\\Delta_f$ is not constant. In particular, in stage $k$, $\\Delta_f = \\mathbb{E}[f(x_{2^{k-1}}) - \\inf f(x)]$. Consequently, the doubling trick implicitly assumes $f$ has bounded function value in order to bound $\\Delta_f$ in all stages. \n\nIf the authors can address my major concerns, I am inclined to revise my score upwards, given the otherwise strong results of the paper.\n\nQUESTIONS:\n- Is there a specific reason why the authors prefer $\\alpha$ as large as possible in the trade-off of $\\Delta^{1/2(1-\\alpha)} + \\sigma^{1/(1-\\alpha)} + L^{1/2\\alpha}$. Is it due to the assumption (or some empirical observation) that the smoothness constant $L$ is usually the dominating parameter?\n\n- In the experiments, did the authors apply any learning rate scheduler (e.g., linear decay or cosine decay) on top of the adaptive learning rate? Although somewhat tangential, it would be interesting to compare the new adaptive scheduler with other popular schedulers.\n\nLIMITATIONS:\nN/A", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-10"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper studies adaptive variants of STORM, a variance reduction technique proposed by Cutkosky and Orabona (2019), for nonconvex stochastic minimization problems. Through introducing a novel adaptive parameter and step size tuning method, the authors aim to remove the bounded gradients and bounded function values assumptions in existing literatures while achieving the optimal convergence rate without an additional penalty of $O(\\log T)$ term due to adaptiveness. Based on the same techniques, they further propose Compositional STORM for compositional functions and STORM for finite-sum problems. They support their theoretical claims with numerical experiments on image classification and small-scale language modelling tasks.\n\nSTRENGTHS:\n1. Their proposed method can achieve optimal convergence rate of $O(T^{-1/3})$ under more relaxed assumptions, namely by removing the need for bounded function values and bounded gradients. Their method does not incur a $\\log(T)$ term penalty like in other adaptive methods (Liu et al., 2022).\n2. They show that their learning rate design and analysis technique can be extended to stochastic compositional problems under further assumptions, namely Lipschitz continuity. \n3. Although rather unsurprising, they show that when the finite-sum structure is present, their method can be adapted to and showed improved convergence rates of a log factor over existing methods.\n\nWEAKNESSES:\n1. My understanding of the original non-adaptive STORM proposed by Cutkosky and Orabona (2019) does not have a log factor in their convergence rate when noise is present. Specifically their convergence rates in expectation is $O(\\log T / T^{1/2} + \\sigma^{1/3} / T^{1/3})$, thus the $\\log T$ term is dominated and should not appear in your big-O convergence rates in many places of the paper (e.g., Line 28, and Table 1). Can you please explain and check your comparisons?\n\n2. I am not sure about the significance of removing the bounded gradients and bounded function values assumption in practice. Cutkosky and Orabona (2019) proposed an adaptive $G_t$ approach to remove the bounded gradients assumption in their original paper, while bounded function values seems to be a very mild assumption under smoothness conditions when the final bounds all depend on the initial function value gap. Thus, I do not agree with the authors claiming them to be \"strong assumptions\". It would be useful if the authors can provide examples to justify why these assumptions can be problematic in practice.\n\n3. The authors should be clearer that their extension to stochastic compositional optimization is not exactly under weaker assumptions. They require standard Lipschitz continuity assumptions which have been standard in the literature, which in essence is same as bounded gradients.\n\nOverall I think this paper has potential but may benefit from addressing the above points.\n\nQUESTIONS:\nSee weaknesses.\n\nLIMITATIONS:\nYes.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-15"}