Stochastic Modified Equations and Dynamics of Dropout Algorithm

$**Point 1**$

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.

$**Reply**$

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}$"

$**Point 2**$

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 author's 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.

$**Reply**$

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

$**Point 3**$

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 ?

$**Reply**$

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

$**Point 4**$

It's not quite clear how the Hessian alignment with the covariance matrix is aiding to flatter minima. 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. Infact near convegence, $l_{i, 2}$ tends to 0, making the second term in the covariance term vanish.

$**Reply**$

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.

$**Point 5**$

How the effect of dropout is in aiding to flat minima correspond to that of SGD and parameter noise injection. In some sense, dropout has some correspondace to both. Parametrer 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.

$**Reply**$

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.

$**Point 6**$

Which SME is used in the experiments? The order-1 approximation or the order-2 approximation.

$**Reply**$

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.

Please, confirm that you have read the author's response and the other reviewers' comments and indicate if you are willing to revise your rating.

We appreciate the reviewer's suggestion to concisely address the issue. In our revised manuscript, we've added some brief remarks following Assumption 1, acknowledging the limitations and empirical findings: ``We remark that local existence of the solution to SDE and estimates of all $2l$-moments of the solution to SDE can be guaranteed for smooth coefficients and sufficiently small time $T^*>0$. Moreover, as the constants $C(T^*,\Theta_0)$ and $C(T^*,\theta_0,\eta_0)$ are exponential in time, the $2l$-moments of the solution might blow up for large enough $T^*$, which is unavoidable since we are unable to impose the uniform Lipschitz condition on $\nabla L_{S}$ and $\Sigma$. However, our empirical findings suggest that the SME still possess the desired approximation ability to dropout even for a large learning rate, as shown in Figure 1 (a).''

Furthermore, we are grateful for your willingness to consider raising the score for our work.

$**Point 1**$

Overall, there is a question of the impact of the contribution. Everything within the paper is well executed (up to some minor comments addressed in questions), but the main result seems to be writing down the modified loss and SDE. The results about the alignment of the Hessian and dropout noise seem somewhat incomplete; I have given suggestions for improving those analyses as well.

$**Reply**$

Dropout has been proven to be a useful and widely adopted technique for training neural networks. The SME framework serves as a rigorous tool for comprehending the noise structure associated with dropout and establishing connections between this structure, generalization, and parameter condensation phenomena (Zhang and Xu). In the revised manuscript, we have provided more elaborate experiments and analyses with the goal of expanding the scope of our SME framework and alignment structure. We hope these changes will meet your expectations. For specific discussions, please refer to the responses addressing the detailed questions below.

Zhang and Xu. Implicit regularization of dropout. arxiv 2207.05952

$**Point 2**$

A related question is whether or not the results hold in the SGD+dropout setting. Even additional numerical evidence would be sufficient here in my opinion.

$**Reply**$

We concur with the reviewer's point that theoretical analysis using GD may not comprehensively capture the behavior of dropout regularization when implemented with SGD. Nevertheless, the dropout algorithm implemented by SGD remains fundamentally a time-homogeneous Markov chain. Consequently, the first-order and second-order approximations of the SDE do not pose intrinsic challenges and the SME framework employed in this study is also well-suited for analyzing dropout algorithms implemented with SGD.

Given the nature of SGD as an unbiased estimator with respect to the full sample, we conjecture that an order-$1$ approximation utilizing SME for SGD combined with dropout shall be in the form:

wherein the drift term remains invariant regardless of GD or SGD, while the diffusion term $\widetilde{\Sigma} \left(\theta_t\right)$ combines noise from both dropout and SGD.

Based on the above SME, we numerically simulate the loss path of dropout implemented by SGD, and the simulation results further support the feasibility of our conjecture. Meanwhile, we added experiments to verify the Hessian-variance alignment combined with SGD noise. For detailed experimental settings and results, please refer to Appendix C in the revised manuscript.

$**Point 3**$

How do the results change for deeper networks? Is there evidence that the same structure should remain?

$**Reply**$

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.

$**Point 4**$

The theorems in the paper rely on small learning rate; however, typically neural networks are trained with large stepsizes (which are far from gradient flow for example). Which aspects of the theory will hold at larger learning rates?

$**Reply**$

We remark that regardless of any learning rate, dropout still updates the parameters recursively in the form of a time-homogeneous Markov Chain, i.e., $\theta_{N}=F(\theta_{N-1},\eta_N)$, therefore, our theory holds true for any learning rate. However, it's established in numerical analysis that algorithms with higher-order accuracies are preferred. This preference is based on the well-known principle that lower learning rates (smaller than 1) lead to a favorable relationship $\varepsilon^{k_1}\leq \varepsilon^{k_2}$ whenever $k_1\geq k_2$. Therefore, we prefer opting for a small learning rate ($\varepsilon<1$) to ensure that higher-order terms can be treated as small quantities compared to lower-order terms. Moreover, the commonly employed size of learning rates is usually less than 1.

It is noteworthy that as we set the learning rate to 1, the approximation ability of SME for dropout still remains robust, as is shown in both Figure 1(b) and Figure 2(b).

$**Point 5**$

Figure one and its related experiments could use improvement. For one, it would be better to plot e.g. panel 2 as a 2-D heatmap versus $\varepsilon$ and $p$, so one can understand the effects of varying each (and any interaction that occurs). Additionally, I think its important to probe some smaller values (1e-2, 1e-1 perhaps) as the effects of dropout get more interesting with sparser activations.

$**Reply**$

Drawing inspiration from the heatmap you mentioned, we generated point plots based on varying learning rates and dropout rates, providing a numerical verification of the accuracy of first-order and second-order SME approximation orders. The results are now included in the revised manuscript as Figure 1(c,d). In the figure, it is evident that an increase in learning rate and a decrease in $p$ both contribute to an escalation in error, attributed to heightened noise. We hope that this new presentation method aligns with your expectations.

Regarding the experimental study involving small values of $p$, such as $p=0.1$, we acknowledge its significance. However, we encountered challenges in training with such values. In our ongoing efforts and the preparation of a subsequent paper expanding on this conference paper, we are actively exploring avenues to secure additional computational resources to facilitate these experiments.

$**Point 6**$

For Figure 2, it is better to label each panel by the parameters directly rather than (setting 1, setting 2).

$**Reply**$

We have removed the statement for cases 1-4 and introduced the detailed setup for both $p$ and $\varepsilon$ in the caption and labels of Fig. 2.

$**Point 7**$

For the discussion of matrix alignment, it is worth pointing out that random PSD matrices have non-trivial correlation even in high dimensions... This will provide a better null model.

$**Reply**$

We thank the reviewer for pointing this out. We have corrected the construction method of the random matrix to compare the differences between the cosine similarity between two null models and the cosine similarity between the Hessian matrix and the covariance matrix fairly. Please refer to Section 5.3 in the revised manuscript for our change.

$**Point 8**$

I would suggest finding a different notation for the stepsize (currently $\varepsilon$). This will bring the notation more in line with what practitioners use, which I think will help the paper find a broader audience.

$**Reply**$

We thank the reviewer for pointing this out and we do prefer using the usual notation $\eta$ for the learning rate, however, it shall be noted that $\eta$ has also been widely employed in the context of dropout to denote the dropout random variable, and we have no choice but use $\varepsilon$ for the learning rate to avoid confusion. In our revised manuscript, we also add the following comments to clarify our settings: " We denote the learning rate of the dropout iteration by $\varepsilon$ to mitigate potential confusion with the scaling vector $\eta_N$."

$**Point 9**$

this line in the abstract: "These dual facets of our research, encompassing both theoretical derivations and empirical observations, collectively constitute a substantial contribution towards a deeper understanding of the inherent tendency of dropout to locate flatter minima.'' doesn't really add anything for the reader - the content of the paper is what will convince me if the contribution is substantial or not!

$**Reply**$

Thank you for your constructive feedback. We agree that the focus of the abstract should be on the content of the paper to demonstrate the substantial contribution. In response to your suggestion, we revised the abstract to ensure clarity and conciseness and removed the content mentioned in the point. The revised abstract can be seen in the revised manuscript.

$**Point 1**$

Theoretical analysis focused on two-layer neural networks, but it is indeed a meaningful step towards understanding dropout.

$**Reply**$

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.

$**Point 2**$

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?

$**Reply**$

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.

$**Point 3**$

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?

$**Reply**$

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.

Please, confirm that you have read the author's response and the other reviewers' comments and indicate if you are willing to revise your rating.

$**Point 1**$

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.

$**Reply**$

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:

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

$**Point 2**$

The findings and conclusions derived from the theoretical analysis using GD may not fully reflect the behavior and performance of dropout regularization when applied in practice with SGD.

$**Reply**$

Due to the space limitation, please refer to the response for Point 2 of Reviewer h4F1.

$**Point 3**$

Assumption 1 seems to be quite strong? Any concrete example for this?

$**Reply**$

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

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

$**Point 4**$

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

$**Reply**$

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

$**Point 5**$

The main theoretical result presented in this paper is an informal theorem, which may look a little bit weird.

$**Reply**$

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.

Please, confirm that you have read the author's response and the other reviewers' comments and indicate if you are willing to revise your rating.