Paper: Stochastic Modified Equations and Dynamics of Dropout Algorithm (ICLR-2024)

Reviewer comment (Writing):
In section-4.1, how the modified loss was derived is not very clear. It seems clear from the appendix. Better to refer it to the appendix or provide a statement on how it was derived.

Author response:
We thank the reviewer for pointing this out, and we have added some brief derivation and discussion about $L_{S}$ in Section 4.1 of the revised manuscript to help readers understand $L_{S}$. Meanwhile, we have added the following content at the end of 4.1 to help readers refer to the detailed derivation part more conveniently: "Please refer to Appendix E.1 for the detailed derivation of $L_{S}$"

Reviewer comment (Theory):
It's not evident how the authors get equation 8 from section 4.1. I missed the part where the stochasticity V comes into the picture from section 4.1. From Theorem-1, it is not clear what effect dropout is having as the authors do not clearly mention the source of stochasticity. It is well known a stochastic discrete algorithm can be well approximated by Euler-Maruyama discretization of an SDE, but in the manuscript the connection/derivation is not clear.

Author response:
$V$ stands for the fluctuation term of dropout,   and we point out in the very front of Section 4.1 that the stochasticity originates from the random vector $\eta=(\eta_1, \eta_2, \cdots, \eta_m)^T$ that randomly drops the parameters $\theta$ from time to time at each step. To avoid such confusion,  we add some comments at the very start of Section 4.2 "In  pursuit of a more comprehensive understanding of the dynamics of dropout, we  integrate the fluctuation term of dropout into our analysis."  We remark that the modified equation approach relies heavily on the Euler-Maruyama discretization of   SDE. As is shown already in   Equation (9) and (10), in order to align the dropout iteration with the   Euler-Maruyama discretization of the SDE, we thereby choose the coefficients $b$ and $\sigma$ to match the coefficients in the discrete iteration process of dropout. This meticulous alignment ensures that the trajectories of the dropout iteration and the Euler-Maruyama discretization remain proximate up to the first order.  It is crucial to highlight that Theorem 1 serves as an approximation result, with the emphasis placed on providing an effective approximation rather than expounding on the origin of stochasticity.

Reviewer comment (Theory):
The role of the modified loss $L_{S}$ from the SDE is not clear in terms of generalization. Authors only use the flatness argument from the Hessian-covariance alignment, but the deterministic part driving the SDE seems to not play any role.

Author response:
Zhang and Xu study the influence of the drift term, i.e., $L_S$, on training, and verify experimentally and theoretically that the drift term can significantly improve the flatness of the model loss landscape compared with GD without dropout, and promote the condensation. Meanwhile, they experimentally verify that $L_S$ has better performance than GD without dropout at a small learning rate. Meanwhile, in this work, we demonstrate the promotion effect of dropout noise on generalization ability, which is supported by Fig. 2(b) in the main text, i.e., the comparison between the model trained with $L_S$ and the model trained with dropout. Based on the above two points, we believe that both the drift term and the dropout noise contribute to generalization. 

Zhang and Xu. Implicit regularization of dropout. arxiv 2207.05952

Reviewer comment (Experiments):
Which SME is used in the experiments? The order-1 approximation or the order-2 approximation.

Author response:
We use the order-1 approximation in our numerical validation. We add the clarification of the simulation setup in the first paragraph of Section 4.3: "For the numerical simulation of the SME, unless otherwise specified, we employ the Euler-Maruyama method to approximate its dynamic evolution by the order-1 approximation." In addition, we add validation experiments for second-order SMEs.

Reviewer comment (Theory):
It's not quite clear how the Hessian alignment with the covariance matrix is aiding to flatter minimas. Does the eigenvector directions of the covariance matrix also correspond exactly as that of the Hessian? Given the equations in 5.2, it is not clear how the structure of the Hessian and the Cov matrix are the same. The coefficients (involving $e_{i}$) seem to suggest that the alignment may not be same. In fact, near convergence, $l_{i,2}$ tends to 0, making the second term in the covariance term vanish.

Author response:
The Inverse Variance-Flatness relation, extensively discussed in Appendix B, delves into a more intricate exploration of the interplay between the Hessian matrix and the covariance matrix eigenspace. Additionally, we have conducted numerical experiments in Appendix B.3 to investigate the alignment of eigenvectors with a significant energy proportion between the Hessian matrix and the covariance matrix. The experimental findings underscore a pronounced alignment property in the eigenspace of both matrices.

Regarding the equation in Section 5.2, we acknowledge your point. This equation does not offer a detailed portrayal of the similarity between the Hessian matrix and the covariance matrix; rather, it serves as a heuristic tool. Our primary objective is to empirically examine the alignment and structural similarity of their eigenspaces. In this equation, $\nabla_{\theta} f_{\theta}\left(x_{i}\right)$ and $\nabla_{q_r}\left(a_r\sigma(w_r^{T}x_i)\right)$ can be analogized to eigenvectors, inspiring us to explore the similarity in their eigenspace structures.

Moreover, based on the experimental results in Appendix B.3, we observe that towards the end of training, both the Hessian matrix and the covariance matrix typically possess only a few significant eigenvalues (nine in our experiment). During this phase, the eigenvectors corresponding to these significant eigenvalues in both matrices dictate the level of similarity between them. Consequently, the impact of $\nabla_{q_r}\left(a_r\sigma(w_r^{T}x_i)\right)$ on the model structure may become negligible at the end of training. This observation elucidates why, even in cases where $l_{i,2} \approx 0$, the two matrices still exhibit high similarity.

Reviewer comment (Experiments):
How the effect of dropout in aiding to flat minima corresponds to that of SGD and parameter noise injection. In some sense, dropout has some correspondence to both. Parameter noise injection is well known as the explicit Hessian regularizer (more explicitly on the trace of the Hessian https://arxiv.org/abs/2206.04613). In my opinion, a comparison between SGD, parameter noise-injection, and dropout are critical.

Author response:
The above three methods are all common training techniques. They improve the generalization ability of the model by adding different forms of noise during the training process. Among them, the noise added by SGD is an unbiased noise, while the noise added by parameter noise injection and dropout is biased noise. The bias term and the variance term in the noise jointly affect the training process. Meanwhile, the variance terms of the three, even if they have different structures, appear to be able to improve the flatness of the model, thus affecting the generalization ability of the model. 

We experimentally compared the differences in output, flatness, and generalization capabilities of the models obtained by the three training strategies. The parameter noise-injection method follows the mentioned article (https://arxiv.org/abs/2206.04613) and the corresponding code in github. Under the settings we are concerned about, i.e., the one-dimensional fitting task and the MNIST classification task, dropout has achieved better performance in multiple parallel experiments. For specific settings and results, please refer to Appendix D in the revised manuscript.

Reviewer comment (Theory):
Theoretical analysis focused on two-layer neural networks, but it is indeed a meaningful step towards understanding dropout.

Author response:
Thank the reviewer for the positive comment. We have a little more discussion about the difficulty of extending our analysis to the deeper or more complex NN architecture in the section "Conclusion and Discussion". Due to space limitations, please refer to Point 1 of reviewer 25p8's reply and Appendix C in the revised manuscript for detailed discussion and numerical validation of this issue.

Reviewer comment (Theory):
In Wu et al (2022), they provide an upper bound for the Hessian of the solution found by SGD based on the analysis of SGD noise, using dynamic stability analysis. Is it possible to develop a similar analysis that can offer an estimate of the Hessian for the solution found by dropout?

Author response:
The dynamic stability analysis, also known as linear stability analysis, as presented in Wu et al. (2018) and Wu et al. (2022), serves as a valuable technique for approximating the dynamical behavior of stochastic gradient descent (SGD). The foundation for the Hessian estimate of the solution found by SGD   arises from Lemma 3.2 in Wu et al. (2022), wherein for a general SGD  $\theta_{t+1}=\theta_t-\eta(\nabla L(\theta_t)+\xi_t)$ where $ \xi_t$ are any noises satisfying $\mathbb{E}[\xi_t]=0,~~\mathbb{E}[\xi_t \xi_t^T]=S(\theta_t).$ Then we have for the linearized model equipped with learning rate $\eta$,  the following holds:
$
\mathbb{E}[L(\theta_{t+1})]=\mathbb{E}[r(\theta_t) L(\theta_t)+\eta^2 \nu(\theta_t)],
$
where $\nu(\theta)=\operatorname{Tr}(H S(\theta)) / 2$ and $r(\theta) \geq 0$.

In our work, we demonstrate that    the dynamics of dropout follow close to the gradient descent   trajectory of the modified loss
$
L_S(\theta):=L(\theta)+L_1(\theta):=\frac{1}{2n}\sum_{i=1}^ne_i^2 +\frac{1-p}{2np}\sum_{i=1}^n \sum_{r=1}^m a_{r}^2\sigma(w_{r}^{T}x_i)^2,
$
where $L(\theta)$ is the same loss in SGD, and the extra term $L_1(\theta)$ accounts for the explicit regularization effect introduced by dropout. Due to the presence of term $L_1(\theta)$ in the dropout dynamics,  direct application of the linear stability analysis technique is not straightforward. However, we propose a reasonable approach by considering that as the dropout dynamics follow the gradient descent trajectory of the modified loss $L_S(\theta)$,    linearizing the dropout dynamics leads to an expression of the form: $\theta_{t+1}\approx\theta_t-\eta \nabla^2 L_S(\theta^*)(\theta_t-\theta^*),$ where $\theta^*$ is now one of the global minimum of $L_S(\theta)$. Following the high-level idea from the proof of Lemma 3.2, and considering the dropout dynamics as $\theta_{t+1}=\theta_t-\eta\left(\nabla L_S\left(\theta_t\right)+\xi_t\right),$
for the above linearized model, where $\xi_t$ are any noises satisfying $\mathbb{E}[\xi_t]=0$ and  $\mathbb{E}[\xi_t \xi_t^T]=S(\theta_t)$.
Then we have
$\mathbb{E}[L_S(\theta_{t+1})]=\mathbb{E}[r_S(\theta_t) L_S(\theta_t)+\eta^2 v(\theta_t)],$
where $\nu(\theta)=\operatorname{Tr}(\nabla^2 L_S\left(\theta^*\right) S(\theta)) / 2$ and $r_S(\theta) \geq 0$. Similarly, the two terms $r_S(\theta_t) L_S(\theta_t)$ and $\eta^2 v(\theta_t)$ denote the contributions from the modified gradient $\nabla L_S(\theta_t)$ and the noise $\xi_t$, and linear stability is affected by both terms simultaneously.  

To sum up, the above paragraph outlines a rough idea of the extension from linear stability analysis on SGD to dropout. In SGD, as the noise arises from the stochasticity involved in the selection of training samples, the noisy gradient is an unbiased estimator for the loss $L(\theta)$. However, as the dropout algorithm introduces noise through the stochastic removal of parameters,  the parameter-removal gradient is a   biased estimator for the loss $L(\theta)$, which makes the direct extension of the linear stability analysis technique hard. But as we observe that the parameter-removal gradient is an unbiased estimator for the modified loss $L_S(\theta)$, hence with slight modifications, the extension of the linear stability analysis technique to dropout becomes plausible.

Reviewer comment (Evaluation):
As shown in Figure 3, the alignment measure $\alpha(\theta_t)$ appears to strengthen after the initial training phase. What could be the potential factors contributing to this observed phenomenon?

Author response:
The alignment between the covariance matrix and the Hessian matrix is inspired by the equations in Section 5.2. The establishment of the similarity between these two equations requires the training process at the final stage, i.e. $L_S \approx 0$. Therefore, as training proceeds, the loss value gradually decreases, and the degree of alignment gradually increases. Meanwhile, the increase in alignment helps the model escape the sharp minimum point faster, thereby improving the flatness of the model. We add a remark at the end of Section 5.3.

Reviewer comment (Theory):
The results presented in this paper are specifically applicable to shallow neural networks. The analysis and findings may not directly extend to deeper or more complex neural network architectures.

Author response:
We acknowledge the importance of broadening the analysis to encompass other neural network architectures. Meanwhile, we want to inform that the analysis of complex network architectures is a notoriously difficult problem for deep learning theory. Below, we demonstrate the obstacles of complex network structures.

In the case of multi-layer neural networks, if dropout is applied only to the outermost layer, we can still calculate the explicit expression for the modified loss due to the linear structure. For deep neural network $f_{\theta}(x)=\sum_{r=1}^m a_r(\eta)_{r}\sigma(w_r^{T}x^{[L]})$, where  $x^{[L]}$ is the output function of a $L$-layer neural network. The modified loss $L_S(\cdot)$ reads:

$$ 
L_S(\theta):=\frac{1}{2n}\sum_{i=1}^ne_i^2 +\frac{1-p}{2np}\sum_{i=1}^n \sum_{r=1}^m a_{r}^2\sigma(w_{r}^{T}x_i^{[L]})^2.
$$

However, the situation becomes significantly more challenging when dropout is applied to the inner layers of the multi-layer neural network. In this scenario, obtaining a closed-form expression for the expectation $\mathbb{E}[h(\eta)]$, where $h$ is a highly nonlinear function with respect to $\eta$, becomes nearly impossible, which is a well-known difficulty in deep learning theory. However, the SME framework in this work is still valid if the drift item can be obtained explicitly. In this case, we design numerical experiments to verify the approximation ability of deep network SMEs in Appendix C.2.1 in the revised manuscript by numerically approximating the drift term. Meanwhile, we study the Hessian-variance alignment property under the deep network setting. Meanwhile, we numerically verify the Hessian-variance alignment properties in deep network settings (Fig. 3 in Section 5.3 and Appendix C.2.2), and conduct a detailed analysis of the feasibility of SME application in complex network structures (Appendix C.1).

Reviewer comment (Theory):
Assumption 1 seems to be quite strong? Any concrete example for this?

Author response:
Assumption 1 plays a crucial role in ensuring the local existence of the solution to the SME and estimates for its second, fourth, and sixth moments.   As shown in the proof of Theorem 5.2.1 in [1], a Picard iteration sequence is constructed to demonstrate the existence of the solution to the SDEs. We remark that in  Theorem 5.2.1  in [1], they enforce a linear growth condition, i.e.,
$$
||b \left(\theta_t\right)||\leq C(1+||\theta_t||),~~~~||\sigma \left(\theta_t\right)||\leq C(1+||\theta_t||),
$$
to ensure the global existence of the solution for any time $t>0$. In contrast, our requirement is more modest, necessitating the existence of a solution only within a finite time interval $[0, T^*]$. Consequently, the validity of the Picard iteration sequence persists for sufficiently small  $T^*$, thus indicating that  Assumption 1 automatically holds for small enough $T^*$. 

We also would like to point out that in the work by Li et al. [2], which is the major motivation for our paper, their imposition of a uniform Lipschitz condition on  $b \left(\theta_t\right)$ and $\sigma \left(\theta_t\right)$ is noteworthy. In Remark 11, Li et al. acknowledge that "In the above results, the most restrictive condition is probably the Lipschitz condition." In comparison, we circumvent the uniform Lipschitz condition by shrinking the time interval from $[0,\infty)$ to  $[0, T^*]$. Finally, we add a line of remark right below Assumption 1 in our revised manuscript to avoid potential confusion.


[1] Oksendal, B. Stochastic differential equations: an introduction with applications. 

[2] Li Q, Tai C, Weinan E. Stochastic modified equations and dynamics of stochastic gradient algorithms i: Mathematical foundations[J].

Reviewer comment (Experiments):
Do the results hold for any activation functions? Considering the impact of different activation functions on the behavior of dropout regularization would contribute to a more comprehensive understanding.

Author response:
The validity of the presented results is not contingent on the specific choice of activation function,   as long as the activation function  $\sigma$  is continuously differentiable up to order $6$. This condition ensures the applicability of Taylor's theorem with the Lagrange form of the remainder. However, in numerical experiments, the requirement of smoothness of the activation does not present an obstacle. Experiments using ReLU as the activation function also achieved good approximation results shown in Figs. 1, 2. The effect of different activation functions on the noise structure may not be obvious - they all work towards a flat solution. The drift term may have different performances under different activation functions, and we leave it to subsequent work.

To preempt any potential confusion, we have incorporated a clarifying remark at the bottom of Page 2 in the revised manuscript, explicitly stating that  "we impose hereafter that the activation function..."

Reviewer comment (Theory):
The main theoretical result presented in this paper is an informal theorem, which may look a little bit weird.

Author response:
We thank the reviewer for pointing this out and we have replaced the informal statement of Theorem 1* with the formal Theorem 1 in our revised manuscript.