Title: Convergence of Adafactor under Non-Convex Smooth Stochastic Optimization
Venue: NeurIPS 2024
Accepted: no
Keywords: Adafactor

stochastic optimization

convergence

non-convex smoothness
Abstract: As model sizes in deep learning continue to expand, memory-efficient optimizers are increasingly critical to manage the substantial memory demands of popular algorithms like Adam and AdamW. Among these, Adafactor has emerged as one of the widely adopted choices for training deep learning tasks, particularly large language models. However, despite its practical success, there is limited theoretical analysis on Adafactor's convergence. This paper presents a comprehensive analysis on Adafactor in a non-convex smooth setting, demonstrating its convergence to find a stationary point at a rate of $\tilde{\mathcal{O}}(1/\sqrt{T})$. We find that the default hyper-parameter setting results in a sub-optimal rate in our framework, and propose an alternative setting that could theoretically achieve optimal convergence rate. This finding is further supported by some experimental results.
We also prove that Adafactor with a suitable time-varying clipping threshold could also converge, achieving performance in experiments comparable to that of the standard constant setting.

Abstract

As model sizes in deep learning continue to expand, memory-efficient optimizers are increasingly critical to manage the substantial memory demands of popular algorithms like Adam and AdamW. Among these, Adafactor has emerged as one of the widely adopted choices for training deep learning tasks, particularly large language models. However, despite its practical success, there is limited theoretical analysis on Adafactor's convergence. This paper presents a comprehensive analysis on Adafactor in a non-convex smooth setting, demonstrating its convergence to find a stationary point at a rate of ˜ O (1 / √ T ) . We find that the default hyper-parameter setting results in a sub-optimal rate in our framework, and propose an alternative setting that could theoretically achieve optimal convergence rate. This finding is further supported by some experimental results. We also prove that Adafactor with a suitable time-varying clipping threshold could also converge, achieving performance in experiments comparable to that of the standard constant setting.

1 Introduction

The adaptive gradient-based methods, such as the well-known AdaGrad [9, 29], RMSProp [30], Adadelta [35], Adam [15] and AdamW [22], have become the preferred approaches in solving the following unconstrained stochastic optimization problem in deep learning fields:

where the object function f is non-convex and P denotes a probability distribution. During the training process, these adaptive methods require to store the historical gradients' information so as to adaptively tune their step-sizes. For example, both Adam and AdamW maintain the exponential average of gradients and squared gradients, and AdaGrad stores the cumulative of squared gradients. Despite their effectiveness, these algorithms pose substantial memory challenges for GPUs to save these additional gradients' information, especially when training large language models (LLMs), such as GPT-3 [5], which contains over 175 billion parameters.

To address memory constraints, several memory-efficient optimization algorithms have been developed, e.g., [26, 1, 23, 17]. One of the most popular optimizers is Adafactor [26] which employs a rank-1 matrix factorization to approximate the second moment matrix in Adam. For an n × m weight matrices, this technique reduces memory usage from O ( mn ) to O ( m + n ) by only tracking the moving averages of the row and column sums of the squared gradients matrix. Additionally, Adafactor eliminates the first-order momentum used in Adam and incorporates update clipping to enhance training stability.

The empirical results reveal that Adafactor achieves comparable performance to Adam in training 32 Transformer models [26] . In real applications, several LLMs including PaLM [6] and T5 [24] have 33

34

35

36

37

38

39

40

41

42

43

44

45

46

47

48

49

50

51

52

53

54

55

56

57

58

59

60

61

62

63

64

65

66

67

68

69

70

71

72

73

74

75

76

77

78

79

80

applied Adafactor as their main optimizers [38]. In spite of Adafactor's widely usage, there is still limited understanding on its convergence in theory, especially the effect of the matrix approximation and update clipping, and the explanation for its hyper-parameter setting in experiments.

In this paper, we take a closer look on Adafactor's convergence under non-convex smooth optimization problems, considering the typical bounded gradient setting as those for AdaGrad [19, 32] and Adam [34]. We aim to provide a convergence rate for Adafactor and explain the influence of the hyperparameters for the convergence speed. We also prove in theory why the default parameter setting is effective in practical scenarios. The analysis to Adafactor is non-trivial compared to other adaptive methods such as AdaGrad and Adam due to the unique matrix factorization and update clipping mechanisms. Based on a new proxy step-size construction and some new compositions as well as estimations, we analyze the additional error terms in the Descent Lemma introduced by the matrix approximation and update clipping. Our main contributions are summarized as follows.

Contributions

We provide a convergence analysis for the full-batch Adafactor considering bounded gradients and a broader range of parameter setting which covers the default one in [26]. The result shows that Adafactor could converge to find a stationary point with a rate of ˜ O (1 / √ T ) where T denotes the total iteration number.

We further investigate the more realistic stochastic Adafactor. It's found that a simple variant of Adafactor, which drops the update clipping, could attain the best convergence rate of ˜ O (1 / √ T ) when the second moment decay rate is 1 -1 /k . We also verify that the default decay rate 1 -1 /k 0 . 8 could lead to a sub-optimal convergence rate in our framework. To illustrate this finding, we provide some empirical results, showing that the potential best hyper-parameter setting in theory could perform better than the default one used in experiments.

We extend our study to include a time-varying clipping threshold. Our analysis implies that with proper selections of clipping threshold and hyper-parameters, Adafactor could also achieve the best convergence rate of ˜ O (1 / √ T ) . We also do some experiments to show that the new clipping threshold scheme achieves comparable performance and training stability to the original constant threshold setting.

The rest of the paper is organized as follows. The next section provides some most relevant works. Section 3 presents some necessary notations definitions and problem setup. Section 4 reviews Adafactor and introduces its essential mechanism. In Section 5 and Section 6, we separately provide convergence bounds for full-batch Adafactor and stochastic Adafactor without update clipping. We further discuss the hyper-parameters' dependency. In Section 7, we investigate Adafactor using a time-increasing update clipping threshold. Section 8 provides experimental results to support our theory. All the detailed proof could be found in the appendix.

2 Related work

In this paper, we mainly investigate the theoretical convergence of Adafactor. Although there is limited works on Adafactor in theory, it's necessary to briefly discuss related works on the convergence of other adaptive methods, particularly on non-convex smooth optimization. Here, we briefly list some of the most related works.

Convergence of adaptive methods Several studies address the convergence of AdaGrad in nonconvex settings. For example, [19] considered a simple variant with delayed step-size, while [32] and [39] assumed bounded stochastic gradients. Other works [14, 10, 21, 3, 31, 27, 33] derived convergence bounds under more relaxed assumptions. Another line of research has investigated the convergence of Adam. For instance, [34, 7, 39, 11, 8] assumed bounded gradients. [28, 36, 31] considered more relaxed noise assumptions without relying on bounded gradients. Additionally, [18] derived convergence bounds for Adam under generalized smooth conditions.

Overall, the convergence analysis of optimizers typically starts with standard assumptions, such as 81 bounded gradients and smooth objective functions. In subsequent studies, these assumptions are 82 gradually relaxed to investigate the convergence properties of the optimizers under less stringent 83 conditions. 84

Memory efficient algorithms As large models are increasingly used in deep learning, memory 85 constraints have become a central issue during training. Consequently, several memory-efficient 86 optimizers have been developed to address this challenge. 87

88

89

90

91

92

93

94

One approach to save memory involves applying matrix factorization to oeptimization algorithms. For instance, [25] used matrix factorization in the second moment estimator of gradients in Adam, similar to the concept behind Adafactor. [23] introduced CAME, a variant of Adafactor, which incorporates a confidence-guided strategy to mitigate instability caused by erroneous updates. [37] proposed Adapprox, leveraging randomized low-rank matrix approximation for Adam's second moment estimator, demonstrating superior performance and reduced memory usage compared to AdamW.

95

96

97

98

99

100

101

102

103

104

105

106

107

108

109

110

111

112

113

114

115

116

117

118

119

120

121

122

123

124

125

126

127

128

129

There are some other techniques to save the memory. For example, [12] relied on a 'Shampoo" technique to reduce the storage requirement of full-matrix preconditioning methods. Notably, their method could be further extended to the more realistic tensor case. [1] presented a memory-saved version of AdaGrad, called SM3, by maintaining k sets gradient accumulator. They proved the convergence guarantee of SM3 on online convex optimization and the effectiveness in experiments. Recently, [17] built a 4-bit Adam using quantization techniques to compress the first and second moment estimators in Adam, also reducing memory usage.

In summary, many existing optimizers, particularly adaptive methods like AdaGrad and Adam, face memory overhead. In response, the discussed works have designed memory-efficient optimizers that aim to achieve comparable performance to these existing methods while achieving memory benefits.

3 Problem setup

To start with, we introduce some necessary notations.

Notations For any two matrices X = ( x ij ) ij , Y = ( y ij ) ij ∈ R n × m , we define ⟨ X , Y ⟩ = ∑ n i =1 ∑ m j =1 x ij y ij . X ⊙ Y , X / Y and √ X denote the coordinate-wise product, quotient and squared root respectively. 0 n and 1 n denote the zero and one n -dimensional vector respectively, and 1 n × m denotes the one n × m -dimensional matrix. The index set [ n ] denotes { 1 , 2 , · · · , n } . ∥ · ∥ F denotes the Frobenius norm. For a positive sequence { α i } i ≥ 1 , we define ∑ b i = a α i = 0 and ∏ b i = a α i = 1 if a > b . The operator RMS ( · ) denotes

We consider unconstrained stochastic optimization (1) over R n × m with the Frobenius norm. The objective function f : R n × m → R is differentiable. Given an n × m matrix X , we assume a gradient oracle that returns a random matrix g ( X , Z ) ∈ R n × m dependent by the random sample Z . The deterministic gradient of f at X is denoted by ∇ f ( X ) ∈ R n × m .

Assumptions We make the following standard assumptions throughout the paper.

( A1 ) L -smoothness: For any X , Y ∈ R n × m , ∥∇ f ( Y ) -∇ f ( X ) ∥ F ≤ L ∥ Y -X ∥ F ;

( A2 ) Bounded below: There exists f ∗ > -∞ such that f ( X ) ≥ f ∗ , ∀ X ∈ R n × m ;

( A3 ) Unbiased estimator: The gradient oracle provides an unbiased estimator of ∇ f ( X ) , i.e., E Z [ g ( X , Z )] = ∇ f ( X ) , ∀ X ∈ R n × m ;

( A4 ) Almost surely bounded stochastic gradient: for any X ∈ R n × m , ∥ g ( X , Z ) ∥ F ≤ G , a.s..

Combining (A3) and (A4), it's easy to verify that ∥∇ f ( X ) ∥ ≤ G, ∀ X ∈ R n × m . Assumptions (A1)-(A3) are standard in the non-convex smooth convergence analysis. Although Assumption (A4) is a bit strong since it requires an almost surely bounded stochastic gradients instead of an expected one, it's still commonly used to derive the high probability convergence bound, see e.g., [32, 14], which is a stronger result than an expected convergence. In coordinate-wise algorithm, another standard assumption is l ∞ -bounded gradient where ∥ g ( X , Z ) ∥ ∞ ≤ G ∞ , see e.g., [8]. These two types of assumption are equivalent up to dimension factors.

4 Review of Adafactor 130

In this section, we briefly discuss Adafactor based on the reference [26]. The pseudocode for 131 Adafactor is presented in Algorithm 1. 132

Algorithm 1 Adafactor

Input: Initialization point X 1 ∈ R n × m , R 0 = 0 m , C 0 = 0 ⊤ n , relative step-sizes { ρ k } k ≥ 1 , decay rate { β 2 ,k } k ≥ 1 ∈ [0 , 1) , regularization constants ϵ 1 , ϵ 2 > 0 , clipping threshold d . for k = 1 , · · · , T do

G k = g ( X k , Z k ) ;

R

k

2

,k

k

1

-

R

+(1

=

β

C

2

,k

k

1

k

-

C

+(1

=

β

W

k k k U k = G k / √ W k ;

⊤

n

-

-

k

;

R

C

1

R

= (

)

/

η k = max { ϵ 2 , RMS ( X k ) } ρ k / max { 1 , RMS ( U k ) /d } ; √

k

+1

X

X

=

η

-

end for

Matrix factorization Adafactor could be severed as a saved-memory version of Adam. Throughout the training process, Adam maintain two n × m matrices M k and V k using exponential moving average update,

where β 1 ,k , β 2 ,k ∈ (0 , 1) , thereby tripling the memory usage. The innovation in Adafactor lies in its method of approximating V k by factoring it into two rank-1 matrices, specifically the row sums and column sums of V k . This approximation is guided by maintaining a minimal general Kullback-Leibler (KL) divergence as follows,

where d ( p, q ) = p log( p/q ) -p + q . The choice of KL-divergence over the more typical Frobenius norm allows for an analytical solution to be derived, specifically given by

Therefore, Adafactor only requires to maintain two vectors R k = V k 1 m , C k = 1 ⊤ n V k , sufficiently reducing the memory from 2 mn to m + n . Although this factorization sacrifices some information of the squared gradients, Adafactor still delivers performance comparable to Adam in many real application tasks, making it a practical choice where memory is a constraint.

Increasing decay rate In Adam, corrective terms are introduced into M k and V k , resulting in two increasing-to-one decay rates. Theoretically, it has been demonstrated that a value closed to one for β 2 ,k would ensure the convergence, e.g., [8, 39, 36]. Inspired by this observation, Adafactor used an increasing second moment decay rate β 2 ,k = 1 -1 /k c , c ≥ 0 , and the empirical default setting is c = 0 . 8 . As pointed out by [26], this setting allows for enjoying the stability of a low β 2 ,k at the early stage of training and the insurance of convergence from a high β 2 ,k as the run progresses. Moreover, it also leverages the bias correction.

Update clipping Adafactor modifies the update process by discarding the first-order moment M k and instead applies an update clipping technique inside the step-size η k . This involves dividing the root-mean-square of the update U k , denoted as RMS ( U k ) , when it exceeds a threshold d . This mechanism helps to calibrate the second moment estimator W k when it's larger-than-desired G k ⊙ G k . Empirical findings in [26] indicated that implementing update clipping leads to significant performance improvements when the warm-up technique is not used.

Relative step-sizes Adafactor incorporates a step-size proportional to scale of X k , denoted by RMS ( X k ) , which is shown in experiments more resilient to the more naive parameter initialization and scaling schemes [26].

k

k

k

k

G

/

·

133

134

135

136

137

138

139

140

141

142

143

144

145

146

147

148

149

150

151

152

153

154

155

156

157

158

159

160

161

β

2

β

,k

)(

2

,k

)

G

1

W

k

⊤

n

(

;

⊙

G

k

G

k

⊙

+

G

k

ϵ

1

1

+

n

ϵ

1

1

⊤

m

1

)

n

1

m

1

⊤

m

;

)

;

162

163

164

165

166

167

168

169

170

171

172

173

174

175

176

177

178

179

180

181

182

183

184

185

186

187

188

189

190

191

5 Convergence result for full-batch Adafactor

We first provide the convergence bound for full-batch Adafactor. At each iteration, full-batch Adafactor obtains the deterministic gradient ∇ f ( X k ) and then updates R k , C k using ∇ f ( X k ) instead of G k in Algorithm 1.

Theorem 5.1. Let { X k } k ≥ 1 be generated by Algorithm 1 with g ( X k , Z k ) = ∇ f ( X k ) , ∀ k ≥ 1 . If Assumptions (A1) and (A2) hold, ∥∇ f ( X k ) ∥ F ≤ G, ∀ k ≥ 1 , β 2 , 1 = 1 2 and

for some positive constant ρ 0 , then for any T ≥ 1 ,

The result indicates that full-batch Adafactor could find a stationary point at a rate of O (log T/ √ T ) under the non-convex smooth case, which is similar to gradient descent but with a sub-optimal rate compared to O (1 /T ) [4]. The hyper-parameter setting in (3) only requires β 2 ,k ∈ (0 , 1) , denoting a much wider range including the default one which requires β 2 ,k to increase to one. The detailed version for the above result can be found in Theorem A.1 from the appendix.

6 Stochastic Adafactor without update clipping

In the stochastic case, we start from the simple scenario of

without considering the update clipping 1 / max { 1 , RMS ( U k ) /d ) } in Algorithm 1, where the main reasons are as follows.

As pointed out in the experiments from [26], Adafactor's performance shows little difference with and without update clipping when implementing learning rate warm-up. Since the warm-up technique is a popular method in deep learning [38], it's reasonable to drop the update clipping.

In stochastic Adafactor, the correlation between G k and η k would be more complex if the update clipping is involved. The proof would be simpler when dropping the update clipping, which could help to better understand the analysis for Adafactor.

We now present the probabilistic convergence bound for Adafactor without update clipping as follows, where we summarize different convergence rate with respect to the factor c from β 2 ,k = 1 -1 /k c , c ∈ [1 / 2 , 1] .

Theorem 6.1. Let { X k } k ≥ 1 be generated by Algorithm 1 without update clipping where η k is given by (5) for each k ≥ 1 . If Assumptions (A1)-(A4) hold, and

for some constants 1 / 2 ≤ c ≤ 1 , ρ 0 > 0 , then for any T ≥ 1 , δ ∈ (0 , 1) , with probability at least 1 -δ ,

The above result indicates that with appropriate hyper-parameters, Adafactor without update clipping 192 could approximately find a stationary point. When the decay rate β 2 ,k is 1 -1 /k , the convergence 193 rate could attain to O (log T/ √ T ) , matching the rate of stochastic gradient descent [4] and the lower 194 rate in [2] up to only a logarithm factor. The hyper-parameter setting in (6) covers the experimental 195 default setting where c = 0 . 8 . The result shows a sub-optimal rate of O (log T/T 0 . 3 ) under the 196 default setting. This finding is further complemented by the coming numerical experiments in Section 197 8. The detailed version of the above results can be found in Theorem B.1 from the appendix. 198

199

200

201

202

203

204

205

206

207

208

209

210

211

212

213

214

215

216

217

218

219

220

221

222

223

224

225

226

227

228

229

230

231

232

233

234

235

236

237

238

6.1 Discussion of the hyper-parameter dependency

In this section, we discuss the dependency of several important hyper-parameters in Theorem 6.1 and the detailed version in Theorem B.1 in the appendix. It's worthy to mention that the dominated order in our convergence bound is determined by the total iteration number T , whereas other hyperparameters could be regarded as constants. However, we hope to improve the dependency of these hyper-parameters as much as possible to make the convergence bound tight.

Discussion of c and the optimal rate The convergence bound in Theorem 6.1 reveals that when c = 1 , β 2 ,k = 1 -1 /k and ρ k = ρ 0 / √ k , the convergence rate attains the optimal rate matching the lower bound. In addition, when c is closed to 1 / 2 , the convergence rate deteriorates. This phenomenon somehow explains that a small decay rate β 2 ,k ( c is low) may harm the convergence speed, as β 2 ,k should be closed enough to 1 to ensure the convergence, which has been pointed out similarly for Adam in [8, 39, 36].

The theoretical best parameter setting remains a small gap to the default one of c = 0 . 8 . To verify our theoretical finding, we provide some empirical evidence in Section 8, showing that β 2 ,k = 1 -1 /k performs even better than the default one and the performance would be better when c increases from 1 / 2 to 1 .

Dependency to mn It's clear to see that the convergence bounds in Theorem A.1 and Theorem B.1 are free of the curse of the dimension factor mn as mn only appears on the denominator in each coefficient. We think that solving the curse of dimension is vital since the applied range for Adafactor includes many deep learning tasks where mn are comparable large to T .

Dependency to ϵ 1 , ϵ 2 The convergence bounds in (37) and (39) from Theorem B.1 has a dependency of O ( ϵ -1 1 log(1 /ϵ 1 )) on ϵ 1 . 1 Although the polynomial dependency to ϵ 1 is a bit worse since ϵ 1 ususally takes a small value in experiments, e.g., the default setting is 10 -30 , it's still common in some theoretical convergence results, e.g., [34, 18]. We also perform some experiments to show that a relatively large ϵ 1 , roughly 10 -3 , makes no observable effect on the performance. Thereby, ϵ 1 could be regarded as a constant in comparison to T and the influence brought by 1 /ϵ 1 could be somehow acceptable.

Since the default value of ϵ 2 is 10 -3 in experiments, it could also be regarded as a constant compared to T . Therefore, the dependency O (1 /ϵ 2 ) on ϵ 2 shows little effect on convergence bounds given the sufficiently large T .

Dependency to the scale of parameters. The convergence bounds in Theorem B.1 contain a O (Θ max ) factor where Θ max denotes the maximum values of ∥ X k ∥ ∞ along the training process. It's reasonable to assume that Θ max ≤ G 0 for a comparable large constant G 0 in practice.

7 Convergence of Adafactor with update clipping

In this section, we take a closer look on the comprehensive Adafactor with both matrix factorization and update clipping. We slightly change the update clipping threshold d in Algorithm 1 to a timevarying threshold d k . The step-size η k then becomes

Then, we present the following convergence bound.

Theorem 7.1. Let { X k } k ≥ 1 be generated by Algorithm 1 with η k given by (7) for each k ≥ 1 . If Assumptions (A1)-(A4) hold, and

for some constants α > 1 , 1 / 2 ≤ c ≤ 1 , ρ 0 > 0 , then for any T ≥ 1 , δ ∈ (0 , 1) , with probability at 239 least 1 -δ , 240 241

Discussion of Theorem 7.1 The convergence result indicates that with a proper selection of the clipping threshold, along with the commonly used step-size ρ k and decay rate β 2 ,k , Adafactor can find a stationary point when T is sufficiently large. The dependency of convergence bound on c remains consistent with Theorem 6.1, achieving the optimal order when c = 1 . In addition, the convergence bound can still avoid the curse of dimension, which is shown in the detailed version Theorem C.1 from the appendix.

The additional hyper-parameter α primarily influences the dependency on ϵ 1 , specifically as O ( ϵ -α 1 log(1 /ϵ 1 ) ) . Thus, our convergence bound may deteriorate as α increases, possibly due to the limitation of our proof framework. This dependency could be potentially improved to O ( ϵ -1 1 log(1 /ϵ 1 ) ) when mn is comparable to 1 /ϵ 1 , which is practical when implementing a largesize model. 2 In our experiments, we tested different values of α and found that suitably small values, such as α = 4 , 6 , 7 , 8 can lead to performance and training stability comparable to the default setting, even without implementing the warm-up technique. This finding suggests that our new threshold setting plays a similar role in enhancing training stability as the default one, which is also the main motivation of update clipping. Since ϵ 1 can be set to a relatively large value, e.g., 10 -3 , a dependency like O ( ϵ -4 1 log(1 /ϵ 1 )) is somewhat acceptable for sufficiently large T .

The time-increasing d k provides the following intuition: As shown in [26, Figure 1], during the early stages of training, a high decay rate β 2 ,k can cause larger-than-desired updates and training instability. Therefore, we set a low threshold d k to ensure that the update clipping mechanism effectively calibrates these larger-than-desired updates. As training progresses, the sequences and updates become more stable, and the second moment estimator W k becomes more accurate in estimating the squared gradients, which is also shown in [26, Figure 1]. Consequently, there is less need for update clipping, corresponding to a relatively large d k . We have also verified through experiments that our setting can achieve performance comparable to the default setting of d = 1 .

8 Experiments

In this section, we will report our experimental results based on the insights obtained in our theory. We will mainly provide the following three experiments:

We test Adafactor without update clipping under different decay rate parameters c , aiming to demonstrate performance improvement as c increases from 0 . 5 to 1 with optimal performance at c = 1 , as indicated in Theorem 6.1 and Theorem 7.1.

We evaluate the sensitivity of Adafactor to different values of ϵ 1 , particularly showing that a relatively large ϵ 1 does not significantly impact performance.

We assess the performance of Adafactor with a time-increasing d k setting, as described in Theorem 7.1, and compare it to the default constant setting.

8.1 Experiment setup

In all experiments, the initialization is R 0 = 0 m and C 0 = 0 ⊤ n . We use a learning rate with the warm-up technique as described in [26], specifically ρ k = min { 10 -6 · k, 1 / √ k } for all experiments unless otherwise specified. The batch size is set to 256, and the total number of epochs is 400 by default. Our models are ResNet-20 and ResNet-110 [13], and we use the CIFAR-10 and CIFAR-100 datasets [16] without any data augmentation. The experiments are conducted using the PyTorch implementation of Adafactor on a single NVIDIA GeForce RTX 4090 GPU.

242

243

244

245

246

247

248

249

250

251

252

253

254

255

256

257

258

259

260

261

262

263

264

265

266

267

268

269

270

271

272

273

274

275

276

277

278

279

280

281

282

283

284

285

286

287

288

289

290

291

292

293

294

295

296

297

298

299

300

301

302

303

304

305

306

0.5

0.6

0.7

0.8

0.9

1.0

decay rate parameter

c

0.5

0.6

0.7

Test Accuracy

0.5

0.6

0.7

0.8

0.9

1.0

decay rate parameter

c

0.15

0.20

0.25

0.30

0.35

0.40

Test Accuracy

0.5

0.6

0.7

0.8

0.9

1.0

decay rate parameter

c

0.1

0.2

0.3

0.4

0.5

Test Accuracy

(a) ResNet-20 on CIFAR-10

(b) ResNet-20 on CIFAR-100

(c) ResNet-110 on CIFAR-100

0

20000

40000

60000

80000

Step

t

0

1

2

3

4

Training Loss

1 =10 30

1 =10 15

1 =10 8

1 =10 5

1 =10 3

1 =10 1

(a) ResNet-20 on CIFAR-10

0

20000

40000

60000

80000

Step

t

0

2

4

6

Training Loss

1 =10 30

1 =10 15

1 =10 8

1 =10 5

1 =10 3

1 =10 1

(b) ResNet-20 on CIFAR-100

0

20000

40000

60000

80000

Step

t

0

2

4

6

8

Training Loss

1 =10 30

1 =10 15

1 =10 8

1 =10 5

1 =10 3

1 =10 1

(c) ResNet-110 on CIFAR-100

8.2 Report on Experiment 1

We test Adafactor without update clipping using decay rate parameter c ranging from 0 . 5 to 1 . 0 in increments of 0 . 05 , while keeping other hyper-parameters at their default values. Each experiment is run 10 times with 100 epochs, and we plot the average test accuracy and standard deviation (shallow blue region) in Figure 1. The results indicate that c = 1 . 0 yields better performance and stability compared to c < 1 . 0 on different models and datasets, corresponding to the highest test accuracy and thinner shallow blue band. These performances show a noticeable improving trend as c increases from 0.5 to 1.0, aligning roughly with the results in Theorem 6.1.

8.3 Report on Experiment 2

In the second experiment, we test Adafactor without update clipping under different ϵ 1 values. We plot the training loss against the step t on different models and datasets in Figure 2. The performance for ϵ 1 = 10 -8 and ϵ 1 = 10 -5 is nearly identical to that for ϵ 1 = 10 -30 . Moreover, even a larger value of 10 -3 achieves comparable training performance, though with a slower decrease in loss. Notably, ϵ 1 = 10 -3 requires approximately the same number of steps ( t ≈ 20000 ) as ϵ 1 = 10 -30 to achieve near-zero training loss. We conclude that Adafactor is not sensitive to the choice of ϵ 1 , and a relatively large ϵ 1 can still lead to convergence, making the polynomial dependency O (1 /ϵ 1 ) in our convergence bounds acceptable.

8.4 Report on Experiment 3

In this experiment, we explore the appropriate values of α in Theorem 7.1 to achieve performance comparable to the default setting of d = 1 . As indicated by Theorem 7.1, a relatively small α is desirable for better dependency on ϵ 1 . We train models with α set to 4, 6, 7, 8, and 9, keeping other hyper-parameters at their default values. We also train models with the default d = 1 setting as the baseline. We plot the training loss against the steps in Figures 3 without step-size warm-up and 4 with step-size warm-up.

0K

2K

4K

6K

8K

10K

Step

t

0.0

0.5

1.0

1.5

2.0

Training Loss

=9.0

=8.0

=7.0

=6.0

=4.0

Baseline

(a) ResNet-20 on CIFAR-10

0K

5K

10K

15K

20K

Step

t

0

1

2

3

4

5

Training Loss

=9.0

=8.0

=7.0

=6.0

=4.0

Baseline

(b) ResNet-20 on CIFAR-100

0K

2K

4K

6K

8K

10K

Step

t

0

2

4

6

Training Loss

=9.0

=8.0

=7.0

=6.0

=4.0

Baseline

(c) ResNet-110 on CIFAR-100

0K

5K

10K

15K

20K

25K

30K

Step

t

0

1

2

3

Training Loss

=9.0

=8.0

=7.0

=6.0

=4.0

Baseline

(a) ResNet-20 on CIFAR-10

0K

10K

20K

30K

40K

50K

Step

t

0

2

4

6

Training Loss

=9.0

=8.0

=7.0

=6.0

=4.0

Baseline

(b) ResNet-20 on CIFAR-100

0K

5K

10K

15K

20K

25K

30K

Step

t

0

2

4

6

Training Loss

=9.0

=8.0

=7.0

=6.0

=4.0

Baseline

(c) ResNet-110 on CIFAR-100

The results indicate that, for these values of α , Adafactor achieves comparable or even better 307 convergence speed compared to the default threshold (represented by "Baseline"). The comparable 308 results to the "Baseline" in Figure 3 further suggest that the time-increasing d k in Theorem 7.1 plays a 309 role similar to that of the default setting, enhancing training stability even when the step-size warm-up 310 is turned off. 311

312

313

314

315

316

317

318

319

320

321

322

323

324

9 Conclusions

In this paper, we investigate the convergence behavior of Adafactor on non-convex smooth landscapes, considering bounded stochastic gradients. We introduce a new proxy step-size to decouple the stochastic gradients from the unique adaptive step-size. Additionally, we use new estimations to control the errors introduced by matrix factorization and update clipping in Adafactor.

Our findings reveal that full-batch Adafactor is capable of finding a stationary point, requiring only a step-size η k ∼ O (1 / √ k ) and a second moment decay rate β 2 ,k ∈ (0 , 1) , denoting a wide range including the default setup. In the case of stochastic Adafactor without update clipping, the convergence rate can achieve the optimal order ˜ O (1 / √ T ) when β 2 ,k = 1 -1 /k c , c = 1 . However, performance deteriorates as c decreases. This finding is supported by experimental results. We also explore Adafactor with a time-increasing clipping threshold and derive similar convergence results. The empirical results demonstrate that the new clipping threshold provides performance comparable to the default constant setting.

Limitations There are several limitations in our work that warrant further investigation. First, 325 the polynomial dependency on ϵ 1 in our convergence bounds may be further improved to a better 326 dependency, such as log(1 /ϵ 1 ) . Second, although we provide convergence results for several variants 327 of Adafactor and demonstrate comparable performance to the original one in experiments, the 328 convergence bound for stochastic vanilla Adafactor remains unknown. Finally, our experimental 329 results primarily focus on traditional deep learning tasks due to our GPU limitations. It would be 330 beneficial to test the scalability of our theoretical results, e.g., on large language models. 331

332

333

334

335

336

337

338

339

340

341

342

343

344

345

346

347

348

349

350

351

352

353

354

355

356

357

358

359

360

361

362

363

364

365

366

367

368

369

370

371

372

373

374

375

References

Rohan Anil, Vineet Gupta, Tomer Koren, and Yoram Singer. Memory efficient adaptive optimization. In Advances in Neural Information Processing Systems , 2019.

Yossi Arjevani, Yair Carmon, John C Duchi, Dylan J Foster, Nathan Srebro, and Blake Woodworth. Lower bounds for non-convex stochastic optimization. Mathematical Programming , 199(1-2):165-214, 2023.

Amit Attia and Tomer Koren. SGD with AdaGrad stepsizes: full adaptivity with high probability to unknown parameters, unbounded gradients and affine variance. In International Conference on Machine Learning , 2023.

Léon Bottou, Frank E Curtis, and Jorge Nocedal. Optimization methods for large-scale machine learning. SIAM Review , 60(2):223-311, 2018.

Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in Neural Information Processing Systems , 33:1877-1901, 2020.

Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, Maarten Bosma, Gaurav Mishra, Adam Roberts, Paul Barham, Hyung Won Chung, Charles Sutton, Sebastian Gehrmann, et al. Palm: Scaling language modeling with pathways. Journal of Machine Learning Research , 24(240):1113, 2023.

Soham De, Anirbit Mukherjee, and Enayat Ullah. Convergence guarantees for RMSProp and Adam in non-convex optimization and an empirical comparison to Nesterov acceleration. arXiv preprint arXiv:1807.06766 , 2018.

Alexandre Défossez, Leon Bottou, Francis Bach, and Nicolas Usunier. A simple convergence proof of Adam and Adagrad. Transactions on Machine Learning Research , 2022.

John Duchi, Elad Hazan, and Yoram Singer. Adaptive subgradient methods for online learning and stochastic optimization. Journal of Machine Learning Research , 12(7):2121-2159, 2011.

Matthew Faw, Isidoros Tziotis, Constantine Caramanis, Aryan Mokhtari, Sanjay Shakkottai, and Rachel Ward. The power of adaptivity in SGD: self-tuning step sizes with unbounded gradients and affine variance. In Conference on Learning Theory , 2022.

Zhishuai Guo, Yi Xu, Wotao Yin, Rong Jin, and Tianbao Yang. A novel convergence analysis for algorithms of the Adam family. In Annual Workshop on Optimization for Machine Learning , 2021.

Vineet Gupta, Tomer Koren, and Yoram Singer. Shampoo: Preconditioned stochastic tensor optimization. In International Conference on Machine Learning , pages 1842-1850. PMLR, 2018.

Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition , 2016.

Ali Kavis, Kfir Yehuda Levy, and Volkan Cevher. High probability bounds for a class of nonconvex algorithms with AdaGrad stepsize. In International Conference on Learning Representations , 2022.

Diederik P Kingma and Jimmy Ba. Adam: a method for stochastic optimization. In International Conference on Learning Representations , 2015.

Alex Krizhevsky, Geoffrey Hinton, et al. Learning multiple layers of features from tiny images. 2009.

Bingrui Li, Jianfei Chen, and Jun Zhu. Memory efficient optimizers with 4-bit states. Advances 376 in Neural Information Processing Systems , 36, 2024. 377

Haochuan Li, Ali Jadbabaie, and Alexander Rakhlin. Convergence of Adam under relaxed 378 assumptions. In Advances in Neural Information Processing Systems , 2023. 379

380

381

Xiaoyu Li and Francesco Orabona. On the convergence of stochastic gradient descent with adaptive stepsizes. In International Conference on Artificial Intelligence and Statistics , 2019.

382

383

384

385

386

387

388

389

390

391

392

393

394

395

396

397

398

399

400

401

402

403

404

405

406

407

408

409

410

411

412

413

414

415

416

417

418

419

420

421

422

423

424

Xiaoyu Li and Francesco Orabona. A high probability analysis of adaptive SGD with momentum. In Workshop on International Conference on Machine Learning , 2020.

Zijian Liu, Ta Duy Nguyen, Thien Hang Nguyen, Alina Ene, and Huy Nguyen. High probability convergence of stochastic gradient methods. In International Conference on Machine Learning , 2023.

Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. In International Conference on Learning Representations , 2019.

Yang Luo, Xiaozhe Ren, Zangwei Zheng, Zhuo Jiang, Xin Jiang, and Yang You. CAME: Confidence-guided adaptive memory efficient optimization. In Proceedings of the 61st Annual Meeting of the Association for Computational Linguistics , 2023.

Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J Liu. Exploring the limits of transfer learning with a unified text-to-text transformer. Journal of Machine Learning Research , 21(140):1-67, 2020.

Noam Shazeer, Azalia Mirhoseini, Krzysztof Maziarz, Andy Davis, Quoc Le, Geoffrey Hinton, and Jeff Dean. Outrageously large neural networks: The sparsely-gated mixture-of-experts layer. In International Conference on Learning Representations , 2017.

Noam Shazeer and Mitchell Stern. Adafactor: Adaptive learning rates with sublinear memory cost. In International Conference on Machine Learning , 2018.

Li Shen, Congliang Chen, Fangyu Zou, Zequn Jie, Ju Sun, and Wei Liu. A unified analysis of AdaGrad with weighted aggregation and momentum acceleration. IEEE Transactions on Neural Networks and Learning Systems , 2023.

Naichen Shi, Dawei Li, Mingyi Hong, and Ruoyu Sun. RMSProp converges with proper hyper-parameter. In International Conference on Learning Representations , 2020.

Matthew Streeter and H Brendan McMahan. Less regret via online conditioning. arXiv preprint arXiv:1002.4862 , 2010.

Tijmen Tieleman and Geoffrey Hinton. Lecture 6.5-RMSProp: Divide the gradient by a running average of its recent magnitude. COURSERA: Neural Networks for Machine Learning , 2012.

Bohan Wang, Jingwen Fu, Huishuai Zhang, Nanning Zheng, and Wei Chen. Closing the gap between the upper bound and lower bound of Adam's iteration complexity. In Advances in Neural Information Processing Systems , 2023.

Rachel Ward, Xiaoxia Wu, and Leon Bottou. Adagrad stepsizes: sharp convergence over nonconvex landscapes. Journal of Machine Learning Research , 21(1):9047-9076, 2020.

Junchi Yang, Xiang Li, and Niao He. Nest your adaptive algorithm for parameter-agnostic nonconvex minimax optimization. In Advances in Neural Information Processing Systems , 2022.

Manzil Zaheer, Sashank Reddi, Devendra Sachan, Satyen Kale, and Sanjiv Kumar. Adaptive methods for nonconvex optimization. In Advances in Neural Information Processing Systems , 2018.

Matthew D Zeiler. Adadelta: an adaptive learning rate method. arXiv preprint arXiv:1212.5701 , 2012.

Yushun Zhang, Congliang Chen, Naichen Shi, Ruoyu Sun, and Zhi-Quan Luo. Adam can converge without any modification on update rules. In Advances in Neural Information Processing Systems , 2022.

Pengxiang Zhao, Ping Li, Yingjie Gu, Yi Zheng, Stephan Ludger Kölker, Zhefeng Wang, and 425 Xiaoming Yuan. Adapprox: Adaptive approximation in adam optimization via randomized 426 low-rank matrices. arXiv preprint arXiv:2403.14958 , 2024. 427

428

429

430

Wayne Xin Zhao, Kun Zhou, Junyi Li, Tianyi Tang, Xiaolei Wang, Yupeng Hou, Yingqian Min, Beichen Zhang, Junjie Zhang, Zican Dong, et al. A survey of large language models. arXiv preprint arXiv:2303.18223 , 2023.

431

432

433

Fangyu Zou, Li Shen, Zequn Jie, Weizhong Zhang, and Wei Liu. A sufficient condition for convergences of Adam and RMSProp. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition , 2019.

A Proof detail for full-batch case 434

We first provide the full-batch Adafactor as follows. The only difference to Algorithm (1) is the 435 replacement of stochastic gradient by deterministic gradient ∇ f ( X k ) at each iteration. 436

Algorithm 2 Full-batch Adafactor

Input: Initialization point X 1 ∈ R n × m , R 0 = 0 n , C 0 = 0 ⊤ m , relative step-sizes { ρ k } k ≥ 1 , decay rate { β 2 ,k } k ≥ 1 ∈ [0 , 1) , regularization constants ϵ 1 , ϵ 2 > 0 , clipping threshold d . for k = 1 , · · · , T do

ˆ η k = max { ϵ 2 , RMS ( X k ) } ρ k / max { 1 , RMS ( ¯ U k ) /d } ;

Then, we provide the detailed version of Theorem 5.1 as follows. 437

Theorem A.1. Let { X k } k ≥ 1 be generated by Algorithm 2. If Assumptions (A1), (A2) hold, 438 ∥∇ f ( X k ) ∥ F ≤ G, ∀ k ≥ 1 and 439

for some positive constant ρ 0 , then for any T ≥ 1 , 440

where we define 441

and the other constant parameters are given by 442

A.1 Preliminary 443

We first denote the auxiliary matrix ¯ G 2 k,ϵ 1 = ¯ G k ⊙ ¯ G k + ϵ 1 1 n 1 ⊤ m . In addition, we define ¯ V k = 444 ( ¯ v ( k ) ij ) ij as follows, 445

To simplify the notation, we let ¯ G k = ( ¯ g ( k ) ij ) ij , R ( i ) ¯ V k , C ( j ) ¯ V k and S ¯ V k be the i -th row sum, j -th column sum and the coordinate sum of ¯ V k respectively. The same definition principal is applied to the notation R ( i ) ¯ G 2 k,ϵ 1 and C ( j ) ¯ G 2 k,ϵ 1 . We also use ¯ w ( k ) ij , ¯ v ( k ) ij , ¯ u ( k ) ij to denote the coordinates of ¯ W k , ¯ V k , ¯ U k in Algorithm 2 respectively. We also define values G 1 , G 2 , G as follows:

446

447

448

449

A.2 Technical lemmas 450

Following the descent lemma for a L -smooth objective function f , we derive that 451

452

In the following, we will provide some necessary technical lemmas.

Lemma A.1. Let β 2 ,k ∈ (0 , 1) and Γ k be defined by 453

Proof. We could prove the result by induction. Since Γ 0 = 0 , it's easy to derive that (1 -β 2 , 1 ) = Γ 1 ≤ 1 . Suppose that for any j ∈ [ k -1] , (1 -β 2 , 1 ) ≤ Γ j ≤ 1 . Then

454

455

456

The induction is then complete. 457

Lemma A.2. Let ¯ V k be defined in (12) . For any k ≥ 0 , it holds that 458

As a consequence, 459

Proof. Note that ¯ R 0 = ¯ V 0 1 m = 0 n and ¯ C 0 = 1 ⊤ n ¯ V 0 = 0 ⊤ m . Suppose that for any j ≤ k -1 , 460 ¯ R j = ¯ V j 1 m , ¯ C j = 1 ⊤ n ¯ V j . Then using the updated rule in Algorithm 2 and (12), 461

Since S ¯ V k represents the coordinate sum of ¯ V k , we could derive that 462

Since R ( i ) ¯ V k denotes the i -th row sum of ¯ V k , it's the i -th coordinate of ¯ R k . Hence, for each coordinate 463 of ¯ R k , using (16), 464

Similarly, we could derive the results related to C ( j ) ¯ V k . 465

Lemma A.3. Following the parameter setting in (3) , for any i ∈ [ n ] , j ∈ [ m ] , k ≥ 1 , it holds that 466

Proof. Recalling the definition of ¯ V k in (12) and ∥∇ f ( X k ) ∥ F ≤ G, ∀ k ≥ 1 , we derive that 467

where the last inequality comes from Lemma A.1. Following (18) and Lemma A.1, we also derive 468 that 469

We also derive the upper bounds for R ( i ) ¯ V k and C ( j ) ¯ V k as follows, 470

Similarly, the lower bound could be derived by 471

472

A.3 Proof of Theorem A.1 473

Now we move to prove the main result. Using (14) and the updated rule in Algorithm 2, 474

We then re-arrange the order, sum up both sides over k ∈ [ t ] and apply f ( X t +1 ) ≥ f ∗ from 475 Assumption (A2) to get, 476

Since Θ min ≤ ∥ X k ∥ ∞ ≤ Θ max , we have Θ min ≤ RMS ( X k ) ≤ Θ max for any k ≥ 1 . Hence, using 477 ˆ η k defined in Algorithm 2, 478

Using (21), ¯ U k = ¯ G k / √ ¯ W k , ∆ 0 in (11) and ρ k = ρ 0 / k , we thus derive that 479

To lower bound (a) , we first discuss the maximum operator inside ˆ η k . Let 480

When k ∈ E 1 , since ∥ X k ∥ ∞ ≥ Θ min , it derives that 481

Using Lemma A.2, we first derive that ¯ w ( k ) ij = ( R ( i ) ¯ V k C ( j ) ¯ V k ) /S ¯ V k . Then, applying Lemma A.3 and 482 ∥∇ f ( X k ) ∥ F ≤ G , we could upper bound ∥ ¯ U k ∥ 2 F as follows, 483

Hence, combining with (23) and (24), we have 484

When k ∈ E 2 , we obtain that ˆ η k = max { ϵ 2 , RMS ( X k ) } ρ k ≥ max { ϵ 2 , Θ min } ρ k and thus 485

Combining with (25) and (26), we derive that 486

We also derive from Lemma A.2 and Lemma A.3 that for any i ∈ [ n ] , j ∈ [ m ] , 487

Using (28), we have 488

491

where A 1 has been defined in (11). Plugging (29) into (27), we derive that 489

Plugging (22) and (30) into (20), and using ρ k = ρ 0 / k , we thus derive that 490

where A 0 is given in (11). Moreover, we have the following results,

We thus derive the first desired result in (9) as follows, 492

Avoiding the curse of dimension To derive a free-dimension numerator bound, we first derive 493 from (21) and (24) with ρ k = ρ 0 / √ k that 494

where ˜ ∆ 0 has been defined in (11). In addition, we derive from Lemma A.2, Lemma A.3 and (13) 495 that 496

where we use m + n ≤ mn and A ′ 1 in (11). Thereby, we have 497

Combining with (27), we thus derive that 498

Plugging (33) and (35) into (20), and using ρ k = ρ 0 / √ k , we derive that 499

where A 0 has been defined in (11). Using (31), we derive the second desired result in (9). 500

B Proof detail for stochastic Adafactor without update clipping

We first provide the detailed version of Theorem 6.1.

Theorem B.1 ( Formal statement of Theorem 6.1 ) . Let { X k } k ≥ 1 be generated by Algorithm 1 without update clipping where η k is given by (5) for each k ≥ 1 . If Assumptions (A1)-(A4) hold, and

for some constants 1 / 2 ≤ c ≤ 1 , ρ 0 > 0 , then for any T ≥ 1 , δ ∈ (0 , 1) , we have the following 505 results. 506

When c = 1 , with probability at least 1 -δ , 507

When 1 / 2 ≤ c < 1 , with probability at least 1 -δ , 508

Here, Θ min , Θ max and G are as in (10) , and 509

The C 0 , C 2 , C 3 are constants defined as 510

501

502

503

504

The C ′ 0 , C ′ 2 , C ′ 3 are positive constants (that could be further upper bounded by constants independent 511 from m,n ), defined by 512

and G 1 , G 2 , G 3 are given by 513

Calculation of hyper-parameter dependency To derive a free dimension bound, we shall use the convergence bounds in (38) and (40). From (43), it's easy to show that m,n could only exist in the denominator of C ′ 0 , C ′ 2 , C ′ 3 , which could avoid the curse of dimension.

To calculate the dependency of ϵ 1 , we first show that its dependency in coefficients C 0 , C 1 , C 2 , C 3 as follows, based on the assumption that 0 < ϵ 1 < 1 ,

Thereby, with the convergence bounds in (37) and (39), it's easy to show that

Proposition B.1. Following the same assumptions and settings in Theorem 6.1, then with probability at least 1 -δ ,

and with probability at least 1 -δ ,

514

515

516

517

518

519

520

521

522

where all constants are given as in Theorem B.1. 523

B.1 Preliminary 524

525

526

527

528

529

We first follow the notations of ¯ G k = ( ¯ g ( k ) ij ) ij and G , G 1 , G 2 in (13). Let G k = ( g ( k ) ij ) ij and ξ k = G k -¯ G k . We also define G 2 k,ϵ 1 = G k ⊙ G k + ϵ 1 1 n 1 ⊤ m and V k = ( v ( k ) ij ) ij as follows,

We also define R ( i ) V k , C ( j ) V k and S V k as the i -th row sum, j -th column sum and coordinate sum of V k respectively. R ( i ) G 2 k,ϵ 1 and C ( j ) G 2 k,ϵ 1 represent the same definitions with respect to G 2 k,ϵ 1 . Then, using a similar deduction in Lemma A.2, we also obtain that for all k ≥ 1 ,

As a consequence of (48), each coordinate of W k satisfies that 530

Next, we introduce a proxy step-size matrix A k = ( a ( k ) ij ) ij such that 531

532

533

534

535

536

537

538

539

540

541

The proxy step-size technique is a standard way in the convergence analysis of adaptive methods, e.g., [32, 8]. We provide a new proxy step-size in (50) to handle the matrix factorization in Adafactor. This construction satisfies two properties. First, it's independent from Z k in order to disrupt the correlation of stochastic gradients and adaptive step-sizes. Second, it needs to remain sufficiently close to the original adaptive step-size w ( k ) ij to avoid generating divergent terms.

B.2 Technical lemmas

In the following, we first provide some more necessary technical lemmas. We introduce a concentration inequality for the martingale difference sequence, see [20] for a proof.

Lemma B.1. Suppose that { Z s } s ∈ [ T ] is a martingale difference sequence with respect to ζ 1 , · · · , ζ T . Assume that for each s ∈ [ T ] , σ s is a random variable dependent on ζ 1 , · · · , ζ s -1 and satisfies that

Then for any λ > 0 , and for any δ ∈ (0 , 1) , it holds that 542

Lemma B.2. Following the parameter setting in (6) , for any i ∈ [ n ] , j ∈ [ m ] , k ≥ 1 , it holds that 543

Proof. First, using Assumption (A4), we derive that 544

545

546

547

548

549

Using the similar deduction for Lemma A.3, we could show that mϵ 1 (1 -β 2 , 1 ) ≤ R ( i ) V k ≤ G 1 . Since β 2 , 1 = 1 / 2 from (6), we then obtain the desired result. The bounds for C ( j ) V k , S V k could be also derived by using similar arguments.

We have the following lemma to upper bound each coordinate of the proxy step-size matrix A k defined in (50) .

Lemma B.3. For any k ≥ 1 , it holds that 550

Proof. We first have 551

Then, recalling the definition of a ( k ) ij in (50) and Lemma B.2, it derives that C ( j ) V k -1 ≤ G 2 and thereby 552 β 2 ,k C ( j ) V k -1 +(1 -β 2 ,k ) G 2 ≤ G 2 ≤ G . Then combining with (51), we derive a ( k ) ij ≤ 2 G . We also 553 derive a free dimension bound from Lemma B.2 for a ( k ) ij as follows, 554

where we use m + n ≤ mn when m,n ≥ 2 and β 2 ,k S V k -1 +(1 -β 2 ,k ) G ≥ mnϵ 1 / 2 . To lower 555 bound a ( k ) ij , we derive from Lemma B.2 that β 2 ,k S V k -1 +(1 -β 2 ,k ) G ≤ G . Thereby, 556

557

Lemma B.4. Let W k and V k be defined in Algorithm 1 without update clipping where η k is given by 558 (5) and (47) respectively. For any k ≥ 1 , it holds that 559

Proof. Recalling (49), v ( k ) ij ≤ R ( i ) V k , v ( k ) ij ≤ C ( j ) V k and Lemma B.2, one could verify that 560

which leads to the desired result that 561

562

The following lemma is inspired by [8, Lemma 5.2] where they considered a constant β 2 ,k . Here, we 563 generalize the result to the case of time-varying β 2 ,k and provide the proof detail. 564

Lemma B.5. For any t ≥ 1 , if β 2 ,k are as in (6) , then it holds that 565

Proof. Recalling the definition of V k and since V 0 = 0 n × m , we have that for any k ≥ 1 , 566

Then, we have 567

where we set y 0 = 0 , θ 0 = 0 and 568

Then we have y k -x k = β 2 ,k y k -1 , ∀ k ≥ 1 . Moreover, since y k ≥ x k , we could use log x ≥ 569 1 -1 /x, ∀ x ≥ 1 to derive that 570

Noting that θ k = β 2 ,k θ k -1 +(1 -β 2 ,k ) ϵ 1 , which leads to β 2 ,k θ k -1 ≤ θ k . Hence, we further have 571

Hence, summing up on both sides of (52) and (53) over k ∈ [ t ] , and noting that x 1 = y 1 , we obtain 572 that 573

Note that y 1 + θ 1 ≥ (1 -β 2 , 1 ) ϵ 1 = ϵ 1 / 2 . Moreover, using Lemma A.1 and Assumption (A4), we 574 have θ t = Γ t ϵ 1 ≤ ϵ 1 and y t ≤ Γ t G 2 ≤ G 2 . We then derive that 575

Noting that for k ≥ 2 , c ∈ [1 / 2 , 1] , β 2 ,k ≥ β 2 , 2 = 1 -1 / 2 c ≥ 1 -1 / √ 2 , we then derive that 576

Finally, plugging (55), (56) into (54), and then summing (54) up over i ∈ [ n ] , j ∈ [ m ] , we obtain the desired result.

Next, we have the following probabilistic result relying on the property of the martingale difference sequence which is commonly used in the analysis of adaptive methods.

Lemma B.6. Following the parameter setting in (6) , for any T ≥ 1 and λ > 0 , with probability at least 1 -δ , ∀ t ∈ [ T ] ,

577

578

579

580

581

582

Proof. Let ζ k = -η k 〈 ¯ G k , ξ k √ A k 〉 and the filtration F k = σ ( Z 1 , · · · , Z k ) where σ ( · ) denotes the 583 σ -algebra. Note that η k , ¯ G k and A k are dependent by { X 1 , · · · , X k -1 } and thereby F k -1 . Since 584 ξ k is dependent by F k , we could prove that { ζ k } k ≥ 1 is a martingale difference sequence since 585

where we apply that E [ ξ k | F k -1 ] = E Z k [ ξ k ] = 0 from Assumption (A3). Then, using Assumption 586 (A3) and Assumption (A4), we have 587

588

Let ω k = 2 Gη k ∥ ∥ ∥ ¯ G k √ A k ∥ ∥ ∥ F . We thus derive from the Cauchy-Schwarz inequality that

Then, using Lemma B.1, it leads to that for any λ > 0 , with probability at least 1 -δ , 589

√

Meanwhile, when Θ min ≤ ∥ X k ∥ ∞ ≤ Θ max , ρ k = ρ 0 / k , we have 590

Combining with Lemma B.3, we derive that 591

√

where we use β 2 , 1 = 1 / 2 , β 2 , 2 = 1 -1 / 2 c ≥ 1 -1 / 2 , c ∈ [1 / 2 , 1] from (6) in the last inequality. 592 Hence, plugging (60) into (57) and then re-scaling the δ , we found that with probability at least 1 -δ , 593 for all t ∈ [ T ] , 594

595

Setting λ = ϵ 1 / (24 G 2 ( ϵ 2 +Θ max ) ρ 0 ) , we derive the desired result.

The following key lemma provides an upper bound for the error brought by the proxy step-size a ( k ) ij , 596 illustrating the error is controllable. 597

Lemma B.7. For any k ≥ 1 , i ∈ [ n ] , j ∈ [ m ] , it holds that 598

where G is as in (13) and G 1 , G 2 are as in (44) . 599

Proof. To simplify the notation, we let 600

Then we have 601

Applying Lemma B.2, we could verify that X,Y,Z ≥ 0 and 602

0 ≤ ∆ X ≤ (1 -β 2 ,k ) G 1 , 0 ≤ ∆ Y ≤ (1 -β 2 ,k ) G 2 , 0 ≤ ∆ Z ≤ (1 -β 2 ,k ) G . (63)

Hence, we derive that 603

Since XY ≥ 0 from (62), Term (I) could be bounded as 604

Recalling the definition, we have R ( i ) V k -1 ≤ S V k -1 , C ( j ) V k -1 ≤ S V k -1 for any i ∈ [ n ] , j ∈ [ m ] . Further, 605 applying Lemma B.2 and (63), we derive that 606

We then derive from (65), G 1 ≤ G and G ≤ G that 607

To derive a free dimension bound, we could obtain from Lemma B.2, (63) and G ≥ mnϵ 1 / 2 that 608 Z +∆ Z ≥ mnϵ 1 / 2 . Hence, 609

We then derive that 610

where we used m + n ≤ mn , and G 1 is defined in (44). Then, combining with (66) and (67), we 611 have 612

where we applied that m + n ≤ mn when m,n ≥ 2 . Then we move to bound (II) . Recalling the 613 definitions in (62), we have X ≤ Z, Y ≤ Z . Applying (63), we have 614

Similarly, we derive from Lemma B.2 that Z ≥ mnϵ 1 / 2 , X ≤ G 1 , Y ≤ G 2 . Hence, 615

G 616

Combining (68) with (69), we then derive the desired result. 617

B.3 Proof of Proposition B.1 618

Using the inequality in (14), we have 619

Introducing the proxy step-size matrix A k in (50) and then summing up both sides over k ∈ [ t ] , we 620 derive that 621

Estimation for A We first introduce ξ k into A , 622

Then, using Lemma B.6, with probability at least 1 -δ , for all t ∈ [ T ] , 623

Estimation for B Term B is essentially the error brought by the proxy step-size A k . We will first 624 calculate the gap of 1 / √ w ( k ) ij and 1 / √ a ( k ) ij as follows, 625

We then apply (73) and Young's inequality, 626

Thus, plugging (61) in Lemma B.7 into (74), we derive that 627

√

where we used (58) in the second inequality and 1 / k ≤ 1 /k c/ 2 , c ∈ [1 / 2 , 1] . Furthermore, using 628 Lemma B.4 and Lemma B.5, we derive that 629

Estimating C Using the similar deduction in (75) and (76), we derive that 630

Putting together We first re-arrange the order in (70) and use f ( X t +1 ) ≥ f ∗ in Assumption (A2) 631 to derive that 632

We then plug (72), (76), (77) into (78) and set t = T , which leads to that with probability at least 633 1 -δ , 634

where C 1 , C 2 , C 3 are as in Theorem B.1. Moreover, using Lemma B.3 and (58), we have 635

Combining with (80) and (79), and using ∑ T k =1 1 / √ k ≥ √ T , we derive that 636

where C 0 has already been defined in (42). We then derive the first desired result that 637

Free dimension bound We follow the similar deduction in (75) and use Lemma B.7 to derive that 638

Recalling the definition of w ( k ) ij in (49) and Lemma B.2, we derive that 639

where G 3 is as in (44). We thus derive from (82) and (83) that 640

Using (58) and (83), we derive that 641

Plugging the unchanged estimation for A in (72), (84) and (85) into (70), we have that with probability 642 at least 1 -δ , for all t ∈ [ T ] , 643

where C ′ 2 , C ′ 3 are given as in (43) and C 1 is as in (41). Further, using Lemma B.3 and the similar 644 deduction for (80), 645

where C ′ 0 is as in (43). Combining with (86) and (87), and setting t = T , we derive the second 646 desired result in Proposition B.1 that 647

B.4 Proof of Theorem B.1 648

Now based on the result in Proposition B.1, we could further derive the final convergence rate. Noting 649 that when c = 1 , we could bound that 650

Then, we obtain that 651

When 1 / 2 ≤ c < 1 , we have 652

Then, we obtain that 653

654

655

656

657

663

664

665

666

667

668

C Proof detail for stochastic Adafactor with update clipping

We first provide the detailed version of Theorem 7.1 as follows.

Theorem C.1. Let { X k } k ≥ 1 be the sequence generated by Algorithm 1 with (7) . If Assumptions (A1) -(A4) hold, and

When c = 1 , with probability at least 1 -δ , 658

When 1 / 2 ≤ c < 1 , with probability at least 1 -δ , 659

where C 1 , C 2 , C 3 , C ′ 2 , C ′ 3 are as in Theorem B.1 and 660

Calculation of hyper-parameters' dependency We first calculate the dependency on m,n,ϵ 1 , α 661 in the additional coefficient D 1 ( α ) as follows, 662

which is free of the curse of dimension since mn exists in the denominator. Recalling the definitions of C ′ 0 , C 1 , C ′ 2 , C ′ 3 in (41) and (43), it's easy to verify that these coefficients are also free of the curse of dimension factor m,n since m,n exist in the denominator. Thereby, we also derive a free dimension bound selecting (91) and (93).

To calculate the dependency on ϵ 1 , we could combine with (45) and (95) to derive that

Thereby, selecting the bounds in (90) and (92) and noting that α > 1 , we derive that the order on ϵ 1 is

Moreover, it's clear to reveal that there exist mn in denominator, which could improve the dependency 669 on ϵ 1 . If we suppose that mn is comparable to ϵ 1 , then we derive that C 0 D 1 ( α ) ∼ O ( ϵ -1 / 2 1 ) and the 670 order on ϵ 1 is 671

C.1 Proof of Theorem C.1 672

We define 673

Since RMS ( U k ) = ∥ U k ∥ F / √ mn , Θ min ≤ RMS ( X k ) ≤ Θ max , we derive that 674

where we applied that 1 / √ k ≤ 1 /k c/ 2 , c ∈ [1 / 2 , 1] and β 2 ,k = 1 -1 /k c in the last inequality. Using 675 the inequalities in (14) and (99), we have 676

Summing up both sides over k ∈ [ t ] and using f ( X t +1 ) ≥ f ∗ from Assumption (A2), we derive that 677

Introducing A k in (50), we further have the following decomposition, 678

Estimating E Hence, using (98), (99), Lemma B.4 and Lemma B.5, we derive that 679

To avoid the curse of dimension, we drive from (98) and (83) that 680

Then, using (99) and (103), we derive that 681

Estimating D.1 We could follow the similar deduction in (73) and (74) to derive that 682

Using Lemma B.7 and (105), we further derive that 683

Using (99), Lemma B.4 and Lemma B.5, we further have 684

To avoid the curse of dimension, we apply Lemma B.7, (99) and (83) to derive that 685

Estimating D.2 Since A k is independent from Z k , it further leads to 686

Then, the deduction for estimating D.2 follows the similar idea as in Lemma B.6, relying on a 687 martingale difference sequence. 688

Let us set φ k = -ˆ ρ k 〈 ¯ G k √ A k , ˜ G k -E Z k [ ˜ G k ]〉 and the filtration F k = σ ( Z 1 , · · · , Z k ) . Noting that 689 ˆ ρ k , ¯ G k and A k are dependent by F k -1 . Since ξ k is dependent by F k , we could prove that { φ k } k ≥ 1 690 is a martingale difference sequence by showing that 691

In addition, using Assumptions (A3), (A4) and Jensen's inequality, we have 692

Therefore, we derive that 693

Let ω ′ k = 2 G ˆ ρ k ∥ ∥ ∥ ¯ G k √ A k ∥ ∥ ∥ F . We thus derive from the Cauchy-Schwarz inequality and (108) that

694

Then, using Lemma B.1, it leads to that for any λ > 0 , with probability at least 1 -δ , 695

Since { β 2 ,k } k ≥ 2 is non-decreasing, we could apply Lemma B.3 to derive that 696

Then, we apply (99), and re-scale δ to obtain that for any λ > 0 , with probability at least 1 -δ , for 697 all t ∈ [ T ] , 698

Setting λ = √ ϵ 1 / (24 G 2 ρ 0 ( ϵ 2 +Θ max )) , we derive that 699

Estimating D.3 First, since A k is independent from Z k and E Z k [ G k ] = ¯ G k , we have 700

We define the random variable S (1) k , S (2) k and ˜ S (1) k using the indicator function χ and G 4 in (94) as 701 follows, 702

S (1) k = χ {∥ U k ∥ F >d k √ mn } , S (2) k = χ {∥ U k ∥ F ≤ d k √ mn } , ˜ S (1) k = χ {∥ G k ∥ F ≥ d k G 4 } .

From (83), we derive that 703

Hence, S (1) k ≤ ˜ S (1) k , ∀ k ≥ 1 . Note that when S (2) k = 1 , it's equivalent to Ω k = 0 . Then, we derive 704 that 705

Furthermore, we use Assumption (A4) and Lemma B.2 to derive a lower bound for a ( k ) ij where 706

Combining with (99), (110), (111) and (112), we thus derive that 707

Putting together Both E and D.1 are bounded with two estimations, one of which owns a better 708 dependency to 1 /ϵ 1 and the other avoids the curse of the dimension. We thereby derive two results. 709 Plugging (106), (109) and (113) into (101) and then combining with (102) and (100), we then derive 710 that with probability at least 1 -δ , for all t ∈ [ T ] , 711

where C 1 , C 2 , C 3 are as in Theorem B.1 and D 1 ( α ) is as in (94). Plugging (107), (109) and (113) 712 into (101), then combining with (104) and (100), we then derive that with probability at least 1 -δ , 713 for all t ∈ [ T ] , 714

where C ′ 2 , C ′ 3 are as in Theorem B.1. Moreover, using (99), we reveal that the lower bound for ˆ ρ k is 715 the same the one for η k in (58). Thereby, following the same deduction in (80) and (86), we derive 716 that 717

where D 0 = min { C 0 , C ′ 0 } that has been defined in (94). Setting t = T on (114) and (115), and then 718 using (116), we then derive that 719

Then, using the results in (88) and (89), we could derive the desired result in Theorem C.1. 720

721

722

723

724

725

726

727

728

729

730

731

732

733

734

735

736

737

738

739

740

741

742

743

744

745

746

747

748

749

750

751

752

753

754

755

756

757

758

759

760

761

762

763

764

765

766

767

768

NeurIPS Paper Checklist

The checklist is designed to encourage best practices for responsible machine learning research, addressing issues of reproducibility, transparency, research ethics, and societal impact. Do not remove the checklist: The papers not including the checklist will be desk rejected. The checklist should follow the references and follow the (optional) supplemental material. The checklist does NOT count towards the page limit.

Please read the checklist guidelines carefully for information on how to answer these questions. For each question in the checklist:

You should answer [Yes] , [No] , or [NA] .

[NA] means either that the question is Not Applicable for that particular paper or the relevant information is Not Available.

Please provide a short (1-2 sentence) justification right after your answer (even for NA).

The checklist answers are an integral part of your paper submission. They are visible to the reviewers, area chairs, senior area chairs, and ethics reviewers. You will be asked to also include it (after eventual revisions) with the final version of your paper, and its final version will be published with the paper.

The reviewers of your paper will be asked to use the checklist as one of the factors in their evaluation. While "[Yes] " is generally preferable to "[No] ", it is perfectly acceptable to answer "[No] " provided a proper justification is given (e.g., "error bars are not reported because it would be too computationally expensive" or "we were unable to find the license for the dataset we used"). In general, answering "[No] " or "[NA] " is not grounds for rejection. While the questions are phrased in a binary way, we acknowledge that the true answer is often more nuanced, so please just use your best judgment and write a justification to elaborate. All supporting evidence can appear either in the main paper or the supplemental material, provided in appendix. If you answer [Yes] to a question, in the justification please point to the section(s) where related material for the question can be found.

IMPORTANT, please:

Delete this instruction block, but keep the section heading 'NeurIPS paper checklist" ,

Keep the checklist subsection headings, questions/answers and guidelines below.

Do not modify the questions and only use the provided macros for your answers .

1. Claims

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

Answer: [Yes]

Justification: [NA]

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: [NA]

769

770

771

772

773

774

775

776

777

778

779

780

781

782

783

784

785

786

787

788

789

790

791

792

793

794

795

796

797

798

799

800

801

802

803

804

805

806

807

808

809

810

811

812

813

814

815

816

817

818

819

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: [NA]

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: [NA]

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: Our code is based on Pytorch package which is standard. In addition, we have clarified the detailed experimental setup in our paper and the experiments are easy to reproduce.

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).

876

877

878

879

880

881

882

883

884

885

886

887

888

889

890

891

892

893

894

895

896

897

898

899

900

901

902

903

904

905

906

907

908

909

910

911

912

913

914

915

916

917

918

919

920

921

922

923

924

925

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: [NA]

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: [NA]

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: [NA]

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.

926

927

928

929

930

931

932

933

934

935

936

937

938

939

940

941

942

943

944

945

946

947

948

949

950

951

952

953

954

955

956

957

958

959

960

961

962

963

964

965

966

967

968

969

970

971

972

973

974

975

976

977

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: [NA]

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: [NA]

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: [NA]

Guidelines:

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

978

979

980

981

982

983

984

985

986

987

988

989

990

991

992

993

994

995

996

997

998

999

1000

1001

1002

1003

1004

1005

1006

1007

1008

1009

1010

1011

1012

1013

1014

1015

1016

1017

1018

1019

1020

1021

1022

1023

1024

1025

1026

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: [NA]

Justification: [NA]

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: [NA]

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]

1027

Justification: [NA]

1028

1029

1030

1031

1032

1033

1034

1035

1036

1037

1038

1039

1040

1041

1042

1043

1044

1045

1046

1047

1048

1049

1050

1051

1052

1053

1054

1055

1056

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: [NA]

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:\nThe paper examines the convergence properties of Adafactor, an adaptive learning rate optimizer designed for deep learning tasks, particularly in memory-constrained environments. The study focuses on Adafactor’s performance in non-convex optimization scenarios and provides theoretical convergence proofs under smooth conditions. Despite its widespread practical use, especially for training large language models, Adafactor’s theoretical understanding has been limited. This research fills that gap by proving that Adafactor can reach a stationary point with a specific convergence rate, highlighting both its efficiency and the impact of different parameter settings on its performance. The paper also introduces modifications to the default hyper-parameter settings based on theoretical insights, which are validated through empirical tests, showing potential improvements over traditional setups.\n\nSTRENGTHS:\nThe analysis of Adafactor’s convergence is crucial for its application in training extremely large models, such as large language models (LLMs). This study significantly contributes to the theoretical foundations of Adafactor, supporting its practical use with a solid mathematical framework.\n\nWEAKNESSES:\nThe manuscript could benefit from a deeper discussion on the proving techniques and the tightness of the provided convergence bounds.\n\nQUESTIONS:\n1. Could you clarify the primary differences or challenges encountered in proving the convergence of Adafactor compared to other adaptive gradient-based methods?\n\n2. The convergence bound suggests that Adafactor is not affected by the curse of dimensionality related to the factor $mn$. Have you observed this phenomenon in practical experiments or applications?\n\n3. Is there potential to improve the current convergence rate of $\\log T/\\sqrt{T}$ stated in Theorem 5.1? If so, what changes or assumptions might lead to such an improvement?\n\nLIMITATIONS:\nYes.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-13"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper studies the convergence of Adafactor for non-convex smooth objectives. The paper looks at both full batch and stochastic cases and analyze the convergence rate. Experiments are provide to validate some of the findings about the hyperparameters. \n\nThe main contributions of this paper are: (1) convergence rates for both full-batch and stochastic cases for Adafactor (2) Provide empirical evidence that the hyperparameters leading to optimal convergence rates yields better empirical performance. My main concern regarding this paper is about novelty and lack of comprehensive experimental evidence. First regarding novelty, convergence of Adaptive methods has been studied in several earlier papers for full-batch settings (e.g. De et.al., 2018). Similar, the issue of second moment decay parameter increasing at the rate of 1 - 1/k is well known (e.g. Reddi et al., 2018) and under this particular schedule, the algorithms roughly boils down to Adagrad like schedule (instead of exponential moving average). Both these contributions are quite well-known for adaptive methods. Thus, it is not entirely clear to me if these contributions for Adafactor are novel enough to warrant acceptance.\n\nFurthermore, Shazeer et.al., in Section 7.2 of their paper already discuss about the aspect of second moments decay. The experiments in the current paper are neither comprehensive or convincing that this leads to he optimal convergence in practice. It is important to do a very thorough investigation if the authors have to demonstrate this phenomenon (e.g. try it in different NLP settings where adaptive methods are very effective). Overall, in my opinion, the main weaknesses of this paper are novelty and poor empirical study.\n\nSTRENGTHS:\nSee summary\n\nWEAKNESSES:\nSee summary\n\nQUESTIONS:\nSee summary\n\nLIMITATIONS:\nSee summary", "TITLE": "Official Review", "RECOMMENDATION": 3, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-29"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper studies the convergence of a memory-efficient, adaptive algorithm, Adafactor, under non-convex smooth settings. First, the authors show that in the full-batch setting (with appropriate hyperparameters), Adafactor converges to a stationary point at an $\\tilde{O}(1/\\sqrt{T})$ rate. For the stochastic setting, they study two regimes: with and without clipping $\\eta_k$, and show that under appropriate selection of hyperparameters, Adafactor attains an $\\tilde{O}(1/\\sqrt{T})$ rate of convergence to the stationary point, matching SGD up to logarithmic factors. The observations are complemented with empirical findings.\n\nSTRENGTHS:\nConsidering large-scale language model training, the work focuses on an important memory-constrained practical optimizer for which we have limited theoretical understanding. The work is clearly motivated, the introduction to Adafactor and its connection to Adam is concisely discussed, and the paper is easy to follow in general. Even though it’s not exactly vanilla Adafactor, it’s exciting to see the authors have established bounds matching SGD in the stochastic case. I also appreciate the authors not shying away from discussing the impact of various hyperparameters (and potential negative points).\n\nWEAKNESSES:\nEven though there are space constraints, I would like to see some proof sketch (at least for the full-batch case) in the main body to provide a general sense for the reader. For example, it could be as simple as starting from smoothness in Taylor series (Eq. 14), telescoping over $k$, lower bounding (a), upper bounding (b) in Eq 20. The more critical step appears to be the lower bound, and a general flavour of how Lemmas $A.2, A.3$ are used to achieve that would be nice to see.\n\nThis is minor, but it would be helpful for the reader if figures are referenced whenever a discussion about some experiment is invoked. Two instances I noticed are line 222, the effect of $\\epsilon_1$, and line 265, the effect of time-increasing $d_k$. Please look for other instances, if any.\n\nFigure 1, experiment on the effect of different decay rate parameter $c$: Showing the test performance is nice, but as the discussion is about train-loss convergence, their time-evolution should be included. It’s fine if the convergence doesn’t speed up with increasing ($c$) and doesn’t match theory, but its important to have them as the entire work deals only with train-loss convergence.\n\nQUESTIONS:\nOn Assumption A4: Did you try considering the less restrictive expectation bound and establishing expectation convergence results instead of high-probability results for Thm. $6.1$, $7.1$? I am curious to hear your thoughts on this.\n\nOn Thm $5.1$: Why you think this bound is loose? What part of the analysis is preventing the establishment of at least the $O(1/T)$ rate? \n\nLine 178: Under learning rate warmup, update clipping shows little effect on performance, and this is used as one of the reasons to drop update clipping. However, in Thm $6.1$, I see no warm-up on $\\rho_k$, it just decreases with $k$. Please correct me if I am misunderstanding something. If not, I don’t mind analysing the a slightly modified algorithm for theory, but the reasons should justify the choices made.\n\nLIMITATIONS:\nYes, the authors have addressed the limitations.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-30"}