Stochastic Anderson Mixing for Nonconvex Stochastic Optimization

Abstract: Anderson mixing (AM) is an acceleration method for fixed-point iterations. Despite its success and wide usage in scientific computing, the convergence theory of AM remains unclear, and its applications to machine learning problems are not well explored. In this paper, by introducing damped projection and adaptive regularization to the classical AM, we propose a Stochastic Anderson Mixing (SAM) scheme to solve nonconvex stochastic optimization problems. Under mild assumptions, we establish the convergence theory of SAM, including the almost sure convergence to stationary points and the worst-case iteration complexity. Moreover, the complexity bound can be improved when randomly choosing an iterate as the output. To further accelerate the convergence, we incorporate a variance reduction technique into the proposed SAM. We also propose a preconditioned mixing strategy for SAM which can empirically achieve faster convergence or better generalization ability. Finally, we apply the SAM method to train various neural networks including the vanilla CNN, ResNets, WideResNet, ResNeXt, DenseNet and LSTM. Experimental results on image classification and language model demonstrate the advantages of our method.


Reviews and author responses:

This paper proposes a stochastic Anderson mixing method to solve non-convex stochastic optimization problems. The main contribution is a  convergence theory and the applications of this method for deep learning problems. Adaptive, variance reduced, and preconditioned versions of this methods are also studied. 
  
This paper proposes a new second-order method, SAM, to improve the convergence rate and the generation performance over existing methods. Similar to stochastic quasi-Newton methods, this method is originated from multi-secant quasi-Newton methods. Currently it remains open to show convergence of Anderson mixing methods for non-convex problems. The theoretical contribution of this paper provides an answer to this problem. Numerical results show that this idea can achieve good performance, in terms of convergence rate, or generalization accuracy. 
The paper is well written, but I think it has a lot, maybe too many results, for a conference paper. 
I have some concerns about the technical contributions: 

* What is the technical contribution compared to the work of stochastic quasi-Newton method in the paper [60]. As Theorem 1-5 have a strong resemblance to the results in this paper, it is better to add more words on line 173 to highlight the main technical challenges. 
* Given this strong similarity, a numerical comparison of SAM with the method SdLBFGS in [60] is indeed performed, but I am not sure if it is a fair comparison. In Figure 1, it seems that only AdaSAM is used to compare with SdLBFGS which is not adaptive? 
* Is there any guarantee that \lambda_k in (15) is always positive (i.e. \lambda_k >0)? Otherwise, it is not always possible to find a positive alpha_k. 
* What is the impact of the hyper-parameters such as m in Algorithm 1 on the main results, e.g. is there any performance analysis with respect to the choice of m? Also is N = max_iter? 
* Are the moving average, alternating iteration used in the experiments, as mentioned in remark 2? If yes, why?
* Why only a subset of MNIST training data is used, rather than the full dataset, in Figure 1? What about the generalization performance? 

Minor issue:
* What is SGDM on line 237? It is more clear to add in words if it is SGD with momentum. 
* What is SFO-calls in Theorem 4? 
* What is RAM on line 237? Isn’t it SAM?  
  yes

 Dear reviewers,

Thank you very much for your efforts in reviewing this paper and providing constructive suggestions.
We tried our best to address all the proposed concerns and problems. Are there any unclear explanations that we could further clarify or expand upon? We are happy to address them as well.

Yours faithfully,

The Authors

 Thank you for your comments and advice. We hope our following answers can clarify our contributions and address your concerns.

Q1:  The technical contribution compared with the work of stochastic quasi-Newton method in the paper [60].  
A1:  Although our theoretical results are similar to those in [60], there are some fundamental differences that make our analysis more difficult.
1. The basic assumptions are different.  
(1) Recalling that the update is $x_{k+1}=x_k-H_kg_k$, where $g_k = \nabla f_{S_k}(x_k)$ is the stochastic gradient evaluated in the mini-batch $S_k$. Let $S_{[k-1]}:=(S_0,\dots,S_{k-1})$ denote the random samplings in the first $k$ iterations. One of the key assumption in [60] is that $H_k$ depends only on $S_{[k-1]}$, which guarantees $\mathbb{E}[H_kg_k|S_{[k-1]}]=H_k\nabla f(x_k)$ that can greatly simplify their derivations. This assumption is formulated as AS.4 in [60].  In contrast, in our methods SAM and SAM-VR, $H_k$ depends on $S_k$ since $r_k = -g_k$ participates in constructing $H_k$ (Lines 130-131), thus violating this assumption. The absence of this assumption makes our proofs more complicated and technical since some relations like $\mathbb{E}[H_kg_k|S_{[k-1]}]=H_k\nabla f(x_k)$ fail to hold.  
(2) The proposed algorithm SdLBFGS in [60] needs twice continuously differentiable assumption of the objective function and the boundedness of the Hessian (AS.5 in [60]), which is not required for our algorithm.  

2. The mechanisms of constructing the secant equations are different.  
In the deterministic case, a secant equation is $H_k s_k = y_k$, where $s_k=x_k-x_{k-1}$ and $y_k=g_k-g_{k-1}$. In the stochastic case, various choices exist in forming $y_k$. For SAM, the choice of $y_k$ is the same as that of the deterministic case. That is why $H_k$ of SAM depends on $S_k$. For the stochastic quasi-Newton (SQN) methods in [60], $y_k = \nabla f_{S_{k-1}}(x_k) - \nabla f_{S_{k-1}}(x_{k-1})$ (Eq. (3.1) in [60]), which ensures $H_k$ is independent of $S_k$. Such choice can reduce variance of noise on the secant equation but demands evaluating gradients on the same mini-batch $S_{k-1}$ twice, thus incurring 2x cost in gradient evaluations. Some other SQN methods [8] use subsampled Hessian-vector products to form $y_k =\nabla^2 f_{S_k'}(x_k-x_{k-1})$, where $S_k' $ is the data samples independent of $S_k$, which also ensures $H_k$ is independent of $S_k$. Therefore, the simpleness of SAM in forming $y_k$ can incurs difficulties in convergence analysis.  

3. The construction of $H_k$ is different.  
For LBFGS, preserving the positive definiteness of $H_k$ is easier as long as ensuring $H_{k-1}\succ 0$ and $s_k^\mathrm{T}y_k>0$. For SAM, $H_k$ is generally not symmetric and no simple relation exists to guarantee the positive definiteness of $H_k$.

We will highlight the distinction and our technical contribution in the revision.

Q2. The numerical comparison between SAM and SdLBFGS.  
A2: AdaSAM is a specific instance of SAM with $\delta_k$ chosen as Eq. (13). Similarly, SdLBFGS is a specific instance of the framework of stochastic quasi-Newton (SQN) methods in [60].To the best of our knowledge, there is no adaptive version of SdLBFGS in the literature and SdLBFGS is the suggested version in SQN. Thus, we compare AdaSAM with SdLBFGS and a different choice of $\delta_k$ for SAM is present in Sec. D.3 in supplemental material.

Q3. The concern about the positiveness of $\lambda_k$ and $\alpha_k$.  
A3: Recalling that the condition (15) in the paper is
$$
\alpha_k\lambda_k \leq 2\beta_k(1-\mu).
$$
Since the right hand side of is always nonnegative if $\mu<1$ and $\beta_k>0$, $\alpha_k$ can be any positive value to satisfy this condition if $ \lambda_k \leq 0 $.

Q4. The role of $m$ and the number of iterations.  
A4: 
1. $m$ denotes the historical length of SAM. Although $m$ does not explicitly participate in our theorems, a too small $m$ may lead to deterioration of the practical performance of SAM. In Sec. D.6 in the supplemental material, we give some numerical tests comparing different $m$.  
2. In practice, we can set $max\\_iter$ as a large number and terminate the training algorithm by considering its first order optimality condition and the consecutive loss of two steps.

Q5. Moving average and alternating iteration in the experiments.  
A5: Moving average and alternating iterations are incorporated in our implementation as two options. Both are irrelevant with our theoretical analysis. 
(1) Moving average was used and we gave the ablation study in Sec. D.4 in the supplemental material. As a common low-cost procedure, moving average can smooth the stochasticity in gradient estimations. Thus, the moving average technique can be helpful when the batch size is very small (batch size = 20 in tests on Penn TreeBank). In the tests on CIFAR, as the batchsize = 128, the effect of moving average is minor (usually less than 0.1% difference).
 (2) Alternating iteration serves as an option to reduce the computational cost. We tested it on CIFAR and found it may achieve comparable test accuracy even without applying AdaSAM in each step. The results are shown in the last two rows in Table 1 and in Figure 9. 

Q6. The tests on MNIST.  
A6: Thanks for pointing it out. The goal of the experiment on MNIST is to demonstrate the convergence performance of SAM and its enhanced version including variance reduction and preconditioned version. Thus, we consider training loss with large batch size: full gradient, 6K, 3K, 2K. The full gradient result is shown in Figure 14 in the supplemental material. Using large mini-batch training with the full datasets can run out of CUDA memory in our device, so we randomly selected a subset from MNIST. For the more practical and complicated deep learning tasks on CIFAR10/CIFAR100 and Penn TreeBank, we used full datasets and followed the standard setting in all tasks. 

Minor issues

(1) SGDM denotes SGD with momentum. Thanks for pointing out this negligence.

(2) $\mathcal{SFO}$ is first introduced in Line 23 and one $\mathcal{SFO}$-call means one gradient evaluation for a single data sample. For example, to form a full gradient of Eq. (2), we need $T$ $\mathcal{SFO}$-calls. We will clarify it in the revision.

(3) RAM is an ablation study introduced in the main paper (See Lines 112-116). RAM uses a constant regularization strategy, which is originally proposed in [50, 53]. From our experiments, it is shown that our proposed adaptive regularization has better performance over RAM.

 Thank you for your review. We hope our following answers can address your concerns.

Q1. The relationship between the proposed algorithm and Nesterov acceleration.  
A1: Our proposed method indeed is reminiscent of Nesterov's accelerated gradient (NAG) and Polyak's heavy-ball (HB) method. Recall that for HB method, the iterative update is
$$
y_k = x_k+\alpha_k(x_k-x_{k-1}) \quad \text{and} \quad x_{k+1}=y_k-\beta_k \nabla f(x_k), \quad \text{(HB)}
$$
and for NAG, it is
$$
y_k = x_k+\alpha_k(x_k-x_{k-1}) \quad \text{and} \quad x_{k+1}=y_k-\beta_k \nabla f(y_k), \quad \text{(NAG)}
$$
where $\alpha_k$ and $\beta_k$ are hyper-parameters. These methods can be viewed as multi-step methods which use extrapolation to accelerate the simple gradient descent method. In our paper, we introduce Anderson mixing by decoupling it into two steps, as described in Eq. (4a) and Eq. (4b). 
This two-step procedure provides a general framework for extrapolation based methods. By choosing different $m$, $\Gamma_k$ and $\bar{r}_k$, we can obtain a different method. More concretely, HB and NAG set the historical length $m=1$ and $\Gamma_k$ is a varying parameter. As for $\bar{r}_k$, HB sets $\bar{r}_k = r_k = -\nabla f(x_k)$, while NAG sets $ \bar{r}_k=-\nabla f(\bar{x}_k) $. For our method, $m\geq 1$ is the finite historical length; $\Gamma_k$ solves the regularized least squares problem (see (11) and Eq. (12));  $\bar{r}_k$ is a linear combination of $r_k$ and the past gradients (see Line 83). 

Q2: The concern about the constant $\mu$.  
A2:  The constant $\mu>0$ serves as a threshold for the positive definiteness of our Hessian approximation $ H_k $ (See Eq. (14)). In stochastic optimization, the positive definiteness of $H_k$ can guarantee that $-H_kg_k$ is a descent direction. In practice, $\mu$ is preset, and $\alpha_k$ needs to be dynamically adjusted to meet the condition (15). 

 Thanks for your comments and suggestions.
1. The length of the paper.
We apologize for the long length of our paper that may spend you too much time to evaluate it. Since this is the first attempt for developing the stochastic version of Anderson mixing, for the completeness of this paper, we show the asymptotic and non-asymptotic convergence properties of SAM, present the full proof and its detailed application for training networks. In revision, we will try our best to shorten the paper.
2. Recent related work on non-convex optimization.
Thanks a lot for your suggestion. In fact, there are many works in the non-convex optimization field. Due to the space limit, we only cite the most related works in our manuscript. In revision, we will include more recent works as suggested, especially the algorithms inspired from the dynamic perspective. 

 Thank you for your comments.  We hope our following answers can address your concerns.

Q1. The effectiveness of the proposed method in theory and in practice.  
A1: Our method is effective both in theory and in practice.  
1. Theoretical perspective.
Due to the hardness of nonconvex stochastic optimization, it is known that the complexity $O(1/\epsilon^2)$ of SGD cannot be further improved even when the function is (non-strongly) convex ([Nemirovski, Arkadi and Yudin, D. Problem Complexity and Method Efficiency in Optimization. John Wiley and Sons, 1983.]). The difficulty is from the variance of the estimated gradient [44]. Hence, our proposed stochastic Anderson mixing (SAM) is optimal since it also achieves the same complexity in nonconvex stochastic optimization. 
In several special cases, SAM has better convergence rate than the gradient descent methods. In deterministic quadratic optimization, SAM and its preconditioned version have strong relationship to the (preconditioned) GMRES method, a powerful Krylov subspace solver that has superlinear convergence rate in solving certain linear systems. In nonlinear optimization, (preconditioned) SAM is also closely related to multi-secant quasi-Newton methods that can have superlinear convergence rate in some cases. 

2. Practical perspective.
Our experiments on image classification (MNIST, CIFAR) and language modeling (Penn TreeBank) demonstrate the advantages of SAM in terms of epochs (Figure 1, 3 in the main article and Figure 4, 8 in the supplemental material) and generalization error (Table 1 and Table 2), comparing with SGD, ADAM, and other optimizers. 

Q2. The numerical comparison with AdaHessian.  
A2: AdaHessian is an interesting and very inspiring work. We will cite this work in revision. It is noted that the convergence property of AdaHessian is limited to the strongly convex case and no asymptotic convergence rate is given there. For the numerical comparison, we use the official implementation of AdaHessian and run it on CIFAR10/CIFAR100 following the instructions in their paper. The numerical results are as follows:

method | ResNet18 | ResNet20 | ResNet32 | ResNet44 | ResNet56 |WideResNet16-4 | CIFAR100/ResNet18
:-: | :-: | :-: | :-: | :-: | :-: | :-: | :-: 
AdaHessian |  94.36±0.09 |  91.92±0.32 |  92.18±0.18 |   92.74±0.11 | 92.40±0.06 |  94.04±0.12 |   76.59±0.42
AdaSAM | 95.17±0.10 |  92.43±0.19 |  93.22±0.32 |   93.57±0.14 | 93.77±0.12 |  95.23±0.07 |   78.13±0.14

AdaHessian runs out of CUDA memory in our device (GeForce RTX 2080 Ti GPU) when testing ResNeXt50, DenseNet121 on CIFAR100. The training time of AdaHessian for ResNet20, ResNet32, ResNet56 on CIFAR10 is 2.7x, 3.2x, 4x as that of AdaSAM.

The paper applies Anderson mixing to stochastic optimization problems and shows its effectiveness in training neural networks. To make Anderson mixing work in stochastic settings, the paper proposes a series of strategies to handle noise and uncertainty, including damped projection, adaptive regularization, variance reduction, and preconditioned mixing. Proof of convergence is provided, together with numerical evaluations.    The main strength of the paper is the proposed strategies to make Anderson mixing work in stochastic settings, and the theoretical analysis provided. The paper provides a set of reasonable strategies to overcome the instability of vanilla Anderson mixing, and the numerical evaluation appears to verify its effectiveness on neural network training. This can potentially be useful in speeding up neural network training.

Looking at the convergence graphs, it seems that the asymptotic convergence rate of the proposed method is the same as stochastic gradient descent. Can the author(s) clarify whether the proposed method improves the convergence rate? Is there any theoretical results to quantify the acceleration from the proposed method?

Recently, there are other methods that attempt to improve the convergence rate compared to SGD. One notable example is:
- Yao, Z., Gholami, A., Shen, S., Mustafa, M., Keutzer, K. and Mahoney, M., 2021, May. ADAHESSIAN: An Adaptive Second Order Optimizer for Machine Learning. In Proceedings of the AAAI Conference on Artificial Intelligence (Vol. 35, No. 12, pp. 10665-10673).
The current paper should cite this work and provide a comparison in the numerical evaluation. There is no discussion.

The paper under review consider the stochastic Anderson Mixing algorithm and  shows its convergence analysis. An enhanced variance reduction of the algorithms is also derived. The algorithms are tested on extensive data sets with other existing method. This paper seems to be the first one to address the theoretical convergence analysis in this direction.  Positive.

1. The proposed SAM algorithm seems to be new and the variance reduction seems to be a plus.

 2. The paper provides interesting convergence analysis for SAM to its stationary point. 

3. Experiments are impressive.  Negative.

1. The paper is really long and proof is not easy to check for all the details. 

2. The paper claims to compare with the state-of-art algorithm in the field for solving non-convex optimization problems. However, the cited paper related to SGD and non-convex optimization are not the latest ones. There are many recent works solving non-convex optimization problems involving mini-batch and variance reduction, e.g. replica-exchange Langevin dynamics based algorithm, underdamped Langevin dynamics based algorithm, etc. It might be interesting to see the performance differences. It might be a good idea that the authors include more recent related works. 

In this paper, the author propose a stochastic version of Anderson acceleration (widely used in fixed point iterations)  to solve non convex optimization problem. Some important modifications of the deterministic  version are introduced in order to stabilize and improve the algorithm. The paper theoretically shows that the new algorithm can achieve O(epsilon^{-2}) sample complexity. Numerical results are provided to verify the performance of SAM.  Originality: I am not quite familiar with Anderson mixing, but the intuition and final theoretical results look interesting to me. One thing I would like the author to comment is what is the relationship between the proposed algorithm and Nestrov Acceleration? From equation (6) and (9), the update of SAM seems to be quite similar to Nestrov. 

Clarity: The paper overall is well written and all the results are presented clearly. The remarks are helpful for the reader to understand the intuition behind the algorithm 

Significance: The stochastic version of AM is quite essential for an efficiency application. The theoretical result guarantees the potential benefit of AM in the training of deep neural networks. The paper also provide some important modifications. Overall, the contribution is significant. 

Other comments:
1. The constant $mu$ appears in all the rates and it looks like the constant can only be verified by the numerical computation (if I don't miss anything). Is there any theoretical bound on this constant?
 The paper will not have any potential negative social impact. 

Meta-review: This paper studies Anderson mixing for stochastic optimization. While the paper is on the borderline, the reviewers mostly found the paper to be interesting and believe that the approach can potentially have some practical use in training neural networks. The presentation of the paper is one of the primary criticisms of the paper. After going through the paper, I tend to agree with the reviewers. In particular:

(1) The paper needs to be revised in order to make the theoretical details easily digestible to the readers. 

(2) The authors also need to add more discussion on comparison with quasi Newton methods. 

(3) Also, I would like to see some empirical comparison of optimizers in the main paper with respect to the *wall-time* rather than number of epochs.

I expect the authors to make these changes in the final version.

Decision: accept