The Implicit and Explicit Regularization Effects of Dropout

Introduction

At training time, dropout sets a random subset of activations to zero, perturbing the network output with a remarkable amount of noise. Testing is performed on the full model, and it is somewhat mysterious that dropout works so well despite this difference between train and test. The esoteric nature of dropout has inspired a large body of work studying its regularization effects: Wager et al. (2013); Helmbold & Long (2015); Cavazza et al. (2017); Mianjy et al. (2018); Mianjy & Arora (2019) study dropout for linear models, matrix factorization, and linearized networks; Arora et al. (2020) study deep networks with dropout only at the last layer. These works primarily study simpler settings than those used in practice, and, as we demonstrate, there is an implicit regularization effect of dropout that is not adressed by prior work.

A large body of recent work has studied implicit, or algorithmic regularization in deep learning, defined to be a regularization effect imposed by the training algorithm, not by the objective (see for example (Gunasekar et al., 2017; Li et al., 2017; Gunasekar et al., 2018b; Arora et al., 2019) and references therein). One notable example of this is in comparing the generalization performance of SGD vs GD: the implicit regularization effect of stochasticity in SGD has been empirically studied in the context of small v.s. large batch training Keskar et al. (2016), where it is observed that noisier small-batch SGD converges to “flatter” local minima which generalize better, whereas large-batch SGD converges “sharper” local minima which generalize more poorly. The starting point of this work is observing that in practice, dropout also introduces an implicit source of regularization because it adds noise to the gradient updates (somewhat analogous to the small v.s. large batch training). Prior studies of dropout only analyze its explicit regularization effect, focusing on how it modifies the expected loss.Prior work (Mianjy et al., 2018) refers to this as the “implicit bias” of dropout. We refer to this as explicit regularization and reserve the term “implicit” to mean algorithmic regularization effect which does not change the objective. Understanding dropout in practical settings requires studying both regularization effects.

This paper focuses on a sharp characterization of the regularization effects in dropout, where we:

disentangle and analytically characterize the explicit and implicit regularization effects of dropout.

derive simplified, analytical, and interpretable regularizers which completely replace dropout for language modeling tasks.

More concretely, this work makes the following contributions:

1. This work empirically shows that dropout provides both explicit and implicit regularization effects. Dropout modifies the expected training objective, and it is natural to define the explicit regularizer as the difference between the expected training objective and the standard objective, as follows:

Here $F_{\textup{drop}}$ denotes the dropout model and drop denotes the randomness from dropout. Moreover, the optimization uses a stochastic approximation of the expected training loss by sampling the dropout noise, which gives rise to an implicit regularization effect.

In practice, the two regularization effects are entangled and easy to conflate. Section 3 provides results of experiments which disentangle these effects.

2. We then distill these two regularization effects, providing simpler and more interpretable regularizers that depend on the derivatives of the model and loss (Section 4). Intuitively, dropout regularizes the stability of the model and loss output evaluated on each training datapoint. Theoretically (in Section 4.3), we provide a generalization bound which helps justify the dependencies of these regularizers on the loss derivatives.

3. Empirically, detailed experiments are provided in Section 5 showing that these simplified, analytical regularizers can faithfully match and replace dropout for both LSTM and Transformer architectures, on the Penn Treebank, Wikitext-2, and Wikitext-103 datasets. To our knowledge, these are the most accurate empirical demonstrations of theory matching practice with regards to the analysis of dropout.Our code is available at https://github.com/cwein3/dropout-analytical.

4. Finally, the form of the derived explicit regularizer provides detailed intuition on how to regularize the stability of a deep model. When the number of output classes (i.e. vocabulary in language modeling) is large, dropout regularizes most heavily the stability of predictions corresponding to classes to which the model assigns a prediction probability that is not too certain (i.e., not close to either 0 or 1). Our ablation experiments in Section 5.2 reveal this is critical for the effectiveness of dropout, and our theory in Section 4.3 offers additional justification for this perspective.

More generally, we hope that the precise methodological derivations that we provide can inform the future study and derivation of data-dependent regularizers in deep learning.

Preliminaries

Here $k$ indexes a coordinate of $h$. Note that $\eta$ is a zero mean random variable. We then apply dropout by computing

and using ${h}_{\textup{drop}}$ instead of $h$. With slight abuse of notation, we let $\eta$ denote the collection of such vectors over all layers. $F(x,\eta)$ denotes the output of model $F$ on input $x$ using dropout noise $\eta$.

Disentangling Explicit and Implicit Regularization in Dropout

We now present an experimental study designed to disentangle the two regularization effects, which confirms the existence of implicit regularization in dropout. Furthermore, this approach allows us to study each effect in isolation.

Consequently, the expected training objective also differs from $L(F)$:

It is natural to define the explicit regularizer as the difference between the expected training objective (averaged over both $x$ and $\eta$) and the standard objective, i.e.

Due to the fact that in practice, we only have access to a finite training sample (and not the population), it is helpful to define explicit regularizer on a single example as follows:

where $\gamma$ is the stepsize, $x$ is a randomly sampled datapoint, and $\eta$ is a randomly sampled dropout noise variable.

If there were no implicit regularization from the stochasticity of dropout, then we would expect $\textsc{Dropout}_{k}$ to have similar test performance to $\textsc{Dropout}_{1}$, which is equivalent to standard dropout. In Figure 1, we plot the validation accuracy vs. training steps for models trained using $\textsc{Dropout}_{k}$ for various values of $k$. Figure 1 shows that, perhaps surprisingly, performance degrades quite sharply for larger choices of $k$. However, the explicit regularizer is still helpful, as $\textsc{Dropout}_{32}$ does not overfit as severely as the model trained without dropout (for Penn Treebank, the best perplexity without dropout is around 120, which is outside the bounds of the graph). $\textsc{Dropout}_{k}$ and $\textsc{Dropout}_{1}$ optimize the same expected objective, so the change in algorithm must be the cause of these performance discrepancies.

Our proposed explanation for Figure 1 is that the gradient noise induced by dropout provides an implicit regularization effect. We verify this constructively by adding noise to the $\textsc{Dropout}_{k}$ updates in order to recover the performance of standard dropout. Let $\xi_{\textup{drop}}$ denote the fluctuation of the stochastic dropout gradient around its mean:

Note that $\textup{Cov}(\widetilde{\xi}_{\textup{drop}})=2\textup{Cov}(\xi_{\textup{drop}})$. Thus, by adding the term $\sqrt{\frac{1}{2}(1-\frac{1}{k})}\widetilde{\xi}_{\textup{drop}}$ to the $\textsc{Dropout}_{k}$ gradient, we obtain a gradient estimate with the same covariance as $\xi_{\textup{drop}}$.

In Figure 2, we verify that this correction procedure recovers the test performance of $\textsc{Dropout}_{1}$. Thus, we have constructed a (complicated) implicit regularizer which explains the discrepancy between $\textsc{Dropout}_{k}$ and $\textsc{Dropout}_{1}$. In Section 4.2, we will explore its simplifications.

Characterizing the Dropout Regularizers

In Section 4.1, we present and derive our explicit regularizer. In Section 4.2, we derive an update noise distribution which captures the implicit regularization effect in dropout. In Section 4.3, we prove a generalization bound for the cross-entropy loss which further justifies our stability-based regularizers. In Section 5, we empirically demonstrate that our derivations accurately capture the regularization effects in dropout – we can match the performance of dropout for language modeling tasks by using only our regularizers.

For simplicity, we focus on node dropout (Hinton et al., 2012; Srivastava et al., 2014), though our analysis applies to variants such as DropConnect as well (Wan et al., 2013).

Single-layer Dropout. For simplicity, we start by considering node dropout applied to a single layer $i$ of the network. For the rest of the paper, we use $h_{i}$ to denote the $i$-th hidden layer of the network and let $F_{i}$ denote the composition of the layers after $h_{i}$, that is, the function that takes in $h_{i}$ as input, and outputs the model prediction. (Thus, $F_{i}(h_{i})=F(x)$).

This provides an approximate version of the dropout explicit regularizer $R_{\textup{drop}}$:

Here the expectation over the linear term in (4.1) vanished because $\delta=\eta_{i}\odot h_{i}(x)$ is a mean-zero vector. Next we take expectation over $\delta$:

We obtain an analytical approximation for the explicit regularizer $R_{\textup{drop}}(F,x)$ by combining the equations above. Next we will rewrite the RHS of (4.2) in a more interpretable form by further dropping some terms.

For notational simplicity, let $J_{F,{i}}(x)$ be the Jacobians of the network output with respect to the hidden layers, and $H_{\textup{out}}(x)$ be the Hessian of the loss with respect to the network outputs:

We claim that $R_{\textup{drop}}(F,x)$ (or the RHS of (4.2)) can be replaced by the following analytical form

Multi-layer Dropout. To deal with dropout on all layers, we simply take Taylor expansion with all the $\delta$’s at every layer. Cross terms cancel because the masks of different layers are independent, and the resulting regularizer is a sum of equation (4.3) over $i$, giving our analytical explicit regularizer:

Interpretation. Our regularizer ensures that the Jacobians and hidden layers of the model output are small when measured in the norm of $H_{\textup{out}}$. We note that for cross entropy loss, $H_{\textup{out}}(x)=\textup{diag}(p)-pp^{\top}\succcurlyeq 0$, where $p$ is the probability vector predicted by the model encoding the distribution over output class labels. As the diagonal entries take the form $p_{k}(1-p_{k})$, this Hessian places stronger emphasis on output classes which the model believes are plausible but not certain. Our experiments in Section 5.2 demonstrate that this particular weighting is an important factor for the success of dropout – alternative ways to weight the stability of each output class in the regularizer do not perform as well.

Keskar et al. (2016); Yao et al. (2018); Jastrzebski et al. (2018) study the relationship between SGD batch size and notions of “flatness” of local minima via metrics related to the magnitudes of the eigenvalues of the second derivative of the loss with respect to the model parameters. They observe that flatter local minima tend to correlate with better generalization. Our regularizer encourages a notion of flatness that depends on the second derivative of the loss with respect to the hidden layers (see (4.2) in our derivation). These quantities are closely related. For example, consider weight matrix $Z$ parametrizing some linear transformation layer, such that $F(x)=F^{\prime}(Zh(x))$, where $h$, $F^{\prime}$ denote the compositions of the layers before and after the application of $Z$. Then defining $h^{\prime}(x)=Zh(x)$, we have

Thus, the loss derivatives with respect to model parameters can be expressed in terms of those with respect to the hidden layers.

We emphasize that one benefit of $R_{\textup{approx}}(F,x)$ is that it provides an interpretable and detailed characterization of the explicit regularization effect of dropout. We hope this can help provide theoreticians and practictioners alike with precise intuitions on why dropout works, and, more broadly, how to design effective stability regularizers in practice.

2 Characterizing the Implicit Regularization Effect

In this section, we derive a gradient noise distribution which can replace the mean-zero gradient noise in dropout, $\xi_{\textup{drop}}$.

Multi-layer Dropout. To handle multi-layer dropout, we Taylor expand over all the layers, obtaining a sum of (4.5) over the layers:To make tuning slightly simpler, we compute the noise by sampling the coordinates of $\eta_{i}$ uniformly from $\{-1,+1\}$ and scaling by $\sqrt{\frac{q}{q-1}}$, as this preserves the covariance.

To replace the implicit effect of dropout, we add the mean-zero noise $\xi_{\textup{approx}}$ to the gradients of the objective.

Interpretation: It is a major open question in deep learning theory to understand the regularization effects of noise (Li et al., 2019). For example, it is even unclear why mini-batch noise in SGD empirically helps in general. Prior works (Yaida, 2018; Wei & Schwab, 2019) have (heuristically) suggested that the noise encourages the algorithm to find a solution that minimizes the trace of the covariance of the noise. As the covariance of $\xi_{\textup{approx}}$ is some function of $\{J_{\textup{loss},{i}}\},\{h_{i}\}$, and their gradients with respect to $W$, the induced regularizer controls some data-dependent stability of the model. Note the conceptual difference with the explicit regularizer, which multiplies the model Jacobian with the loss Hessian, whereas $\xi_{\textup{approx}}$ multiplies the model Jacobian with the loss Jacobian. More precise interpretations are left for future work.

In Section 5, we demonstrate that a combination of our explicit and implicit regularizer can successfully replace dropout. The general update rule which applies these regularizers in lieu of dropout is described in Algorithm 2.

3 Theoretical Support for Stability-based Regularization

Recent works (Arora et al., 2018; Nagarajan & Kolter, 2019; Wei & Ma, 2019a, b) support our stability-based regularization by bounding generalization of the model in terms of its Jacobian norms on the training data. These bounds align with the Jacobian terms in the regularization (4.4). However, they miss a crucial aspect of the regularizers derived in Section 4.1 as they only consider derivatives of the model output, ignoring the loss derivatives (the $H_{\textup{out}}(x)$ term in equation (4.4)). Though this is a subtle distinction, in Section 5.2 we demonstrate that the loss derivatives are necessary on language modeling tasks.

In this section, we prove a new generalization bound for cross entropy loss on linear models. Our bound helps further justify the forms of our regularizers in (4.4) and (4.6), as every term in our bound is scaled by a derivative of the loss.

With probability $1-\delta$ over the training examples, for all weight matrices $W$ satisfying the norm bound $\|W^{\top}\|_{2,1}\leq A$, the following holds:

Here $\mu(W),\nu(W)$ measure the Jacobians and Hessians of the loss and are defined by

Additionally, we define $\theta\triangleq\log^{3}(nc)\max_{i}\|x_{i}\|_{2}^{2}$ and $\zeta\triangleq\frac{B(\log(1/\delta)+\log\log n)}{n}$ is a low order term.

Experiments

In this section, we empirically confirm that our derivations in Section 4 provide accurate characterizations of dropout. Our focus is on language modeling tasks using the LSTM and Transformer architectures.

In this section, we show that the regularizers derived in Section 4 can replace dropout for LSTMs on language modeling tasks. We work with Penn Treebank (Marcus et al., 1994), a corpus of 887,521 tokens and Wikitext-2 (Merity et al., 2016), a corpus of 2,088,628 tokens. In Section 5.3, we study whether our findings can also scale to larger datasets and architectures such as Transformer-XL (Dai et al., 2019).

For the experiments in this section, we base our model and code on Merity et al. (2017a, 2018). For the dropout-trained models, we use node dropout on the output, hidden, and embedding layers as well as DropConnect (Wan et al., 2013) on the weight matricess. We fix the dropout probability to $q=0.4$ for these experiments. To compute the update gradients for our regularizers, we follow the general rule described in Algorithm 2. We specify additional hyperparameters in Section D.

We study three settings described in detail below: our explicit regularizer $R_{\textup{approx}}$ only, adding our noise $\xi_{\textup{approx}}$ to $\textsc{Dropout}_{k}$ updates, and combining our explicit and implicit regularizers. Tables 3 and Tables 4 in Section D summarize the experimental results on our regularizers for the Penn Treebank and Wikitext-2 datasets. We obtain our results without tuning, as we use the regularization coefficient suggested in Section 4 to match the dropout strength. The Jacobian optimization required for the analytical regularizers results in around 3x runtime slowdown compared to dropout, though we note that the analytical regularizers appear to optimize in fewer iterations (see Figure 5).

Completely Replacing Dropout. We demonstrate that the combination of our regularizers can completely replace dropout. We apply algorithm 2, setting $R=R_{\textup{approx}}$ and $\xi=\xi_{\textup{approx}}$. In Figure 5, we plot the validation perplexity vs. time of a model trained with our regularization vs. $\textsc{Dropout}_{1}$. Figure 5 demonstrates that our regularization is enough to replace dropout, confirming the validity of our derivations. We note that our regularizer appears to require fewer iterations to decrease the validation perplexity. This raises the exciting possibility of designing more efficient regularizers than dropout, which we leave for future work.

2 Regularizing the Loss Hessian is Necessary

We argue that simply regularizing the stability of the model outputs is not sufficient. As argued in Section 4.1, our derivations show that dropout enforces stronger stability for output coordinates where the model assigns non-trivial probability mass but is not extremely confident. To demonstrate this is helpful, we experiment with replacing $H_{\textup{out}}$ in our explicit regularizer (see (4.4)) with two alternative quantities.

3 Additional Settings

We test how well our findings translate to larger datasets and different architectures. We use the Wikitext-103 dataset, which contains 103,227,021 tokens, and the Transformer-XL (Dai et al., 2019) and QRNN (Bradbury et al., 2016) architectures. First, we explore whether the implicit regularization effect of dropout is as important on larger datasets. We train the Transformer-XL and QRNN architectures on the Wikitext-103 corpus using $\textsc{Dropout}_{k}$ for $k=1,2,4$. Table 2 shows that for Transformer-XL trained on the full dataset, the implicit regularization effect disappears. We observe the same for QRNN (see Section D).

In Table 2, we also demonstrate that there is an implicit regularization effect when we downsample Wikitext-103 by a factor of 5, though it is not as crucial. Thus, the importance of the implicit regularization depends on the dataset size.

Finally, we confirm that our explicit regularizer is effective on a larger dataset. For Wikitext-103 and Transformer-XL, Table 2 shows that our explicit regularizer achieves validation perplexity of 24.12, within $0.7$ of dropout.

Related Work

Dropout has been the focus of several theoretical works studying its properties for both optimization and generalization (Wager et al., 2013, 2014; Baldi & Sadowski, 2013; Helmbold & Long, 2015; Gal & Ghahramani, 2016a; Helmbold & Long, 2017; Cavazza et al., 2017; Mianjy et al., 2018; Mianjy & Arora, 2019; Arora et al., 2020). Wang & Manning (2013); Maeda (2014); Gal & Ghahramani (2016a); Ma et al. (2016) study dropout from a Bayesian perspective. Gao et al. (2019) empirically study the effect of applying dropout masks in one only direction of the network (either the forward or backward pass).

Wager et al. (2013); Helmbold & Long (2015) use a Taylor expansion to analyze dropout in linear models, and our work extends their analysis to neural networks. Cavazza et al. (2017); Mianjy et al. (2018); Mianjy & Arora (2019) study the expected dropout objective for matrix factorization and linearized neural nets, respectively. Recent work (Arora et al., 2020) studies dropout applied to only the last layer of a deep neural net, computing an exact expression for the explicit regularizer which depends on the magnitudes of coordinates in the last hidden layer. Our analysis of the explicit regularization results in a more general expression which contains similar quantities, but considers dropout at all layers of the network. These prior works focus on explicit regularization and do not study the implicit regularization effects of dropout.

There has been a large body of prior work studying the relationship between gradient noise and generalization (Keskar et al., 2016; Keskar & Socher, 2017; Smith & Le, 2017; Jastrzebski et al., 2018; Xing et al., 2018; Li et al., 2019; Chaudhari & Soatto, 2018). Jastrzębski et al. (2017); Zhu et al. (2018); Wen et al. (2019) study how the noise distribution in SGD affects generalization. Wen et al. (2019) inject noise with an appropriate covariance structure into the updates of large-batch SGD, making it match the behavior of small-batch SGD. We inject appropriate noise to make large-sample dropout updates match standard dropout.

Prior works have also studied data-dependent regularizers of the model and loss stability. Sokolić et al. (2017); Hoffman et al. (2019) apply Jacobian-based regularization to train robust classifiers. Krueger & Memisevic (2015) propose a data-dependent regularizer for the stability of RNN activations. Novak et al. (2018); Arora et al. (2018); Nagarajan & Kolter (2019); Wei & Ma (2019a, b) study the relationship between model stability and generalization.

Finally, regularization for deep models is an important issue in NLP. Zaremba et al. (2014) demonstrated that dropout can be very helpful for NLP tasks. Semeniuta et al. (2016); Gal & Ghahramani (2016b) propose variants of dropout designed for recurrent neural networks. Krueger & Memisevic (2015); Merity et al. (2017b) study temporal activation stability regularization. Merity et al. (2017b); Melis et al. (2017); Merity et al. (2018) demonstrate that the proper tuning of regularizers can greatly impact performance.

On the broader topic of generalization theory of neural networks, Zhang et al. (2016); Neyshabur et al. (2018) observe that deep learning defies a lot of conventional statistical wisdom. Several works have studied generalization bounds for deep networks (see (Bartlett et al., 2017; Neyshabur et al., 2017; Golowich et al., 2017; Dziugaite & Roy, 2017; Arora et al., 2018; Wei & Ma, 2019b) and references therein). Another line of work studies implicit regularization in deep learning (see (Gunasekar et al., 2017, 2018a, 2018b; Soudry et al., 2018; Woodworth et al., 2019; Arora et al., 2019) and references therein).

Conclusion

In this work, we show that dropout actually introduces two entangled sources of regularization: an explicit one which modifies the expected objective, and an implicit one due to stochasticity in the updates. We empirically disentangle these regularizers and derive analytic simplifications which faithfully distill each regularization effect. We demonstrate that our simplified regularizers can replace dropout in practice. Our derivations show that dropout regularizes the stability of the model and loss around the training data.

More broadly, our analytic characterizations of dropout can provide intuition on what works and what doesn’t for stability-based regularizers in deep learning. We hope that these intuitions can help inform and motivate the design of more principled regularizers for deep networks.

Acknowledgements

We would like to thank Michael Xie for suggesting the experiment which disentangled the implicit and explicit regularization effects. Sham Kakade would also like to thank Xinyi Chen, Cyril Zhang, and Yi Zhang for numerous helpful discussions and help with earlier experimentation. Colin Wei acknowledges support from an NSF Graduate Research Fellowship. The work is also partially supported by SDSI and SAIL at Stanford. Sham Kakade acknowledges funding from the Washington Research Foundation for Innovation in Data-intensive Discovery, and the NSF Awards CCF-1703574, and CCF-1740551.

References

Appendix A Full Derivations in Section 4

To ensure that our regularizer is nonnegative, we ignore the second term containing the non-PSD matrix $M$. Ignoring the non-PSD term in this kind of decomposition was also suggested in Sagun et al. (2017). We also omit the factor of $\frac{q}{2(q-1)}$ in (4.3) for simplicity.

A.2 Justification of Taylor Expansion

As dropout introduces a change that has magnitude which is multiplicative in the size of the coordinates of the hidden layers, the perturbation due to dropout might not be small. Since Taylor expansions typically require a small level of perturbation, in this section we argue that when the application of dropout is followed by a linear transformation layer, the perturbation to the linear layer could be small. Furthermore, we demonstrate that performing Taylor expansion with respect to this layer will ultimately give the same regularizer.

We work in the same setting of the derivation in Section 4.1. We add the additional assumption that $h_{i}$ is followed by a linear transformation parameterized by weight matrix $Z$. Thus, we can express $F(x)=F_{i}(h_{i}(x))=F_{i+1}^{\prime}(Zh_{i}(x))$ where $F_{i+1}^{\prime}$ denotes all the computation after the matrix multiplication $Z$. ($F_{i+1}^{\prime}$ differs from $F_{i+1}$ just by an additional activation layer that follows the matrix multiplication by $Z$.) Now we can compute the loss after applying dropout on $h_{i}$ by

Our key observation, as detailed below, is that although the perturbation $\delta$ to $h_{i}(x)$ could be large relative to the magnitudes of the coordinates of $h_{i}(x)$, the perturbation $Z\delta$ may be much smaller relative to the magnitudes of the coordinates of $Zh_{i}(x)$. Thus, the effect of the dropout noise can be mitigated as it passes through linear layers of the network, making the Taylor expansion more realistic.

Concretely, consider the standard deviation of the $j$-th coordinate of $Z\delta$:

In the case where $\{Z_{jk}\}_{k}$ and $h_{i}$ share the same sign on each coordinate, the signal $(Zh_{i}(x))_{j}$ (A.6) can be larger than the size of the perturbation, that is, $\textup{std}[(Z\delta)_{j}]$ (A.5), by a factor of $\sqrt{d}$. For example, consider the case where all entries of $Z$ and $h_{i}$ are 1. In other words, even though the vector $\delta$ seems to be comparable to $h_{i}$ in the norm, after passing through the linear transformation, due to the cancellation arising from the randomness in $\delta$, $Z\delta$ can be much smaller compared to $Zh_{i}$.

Thus, when the weight matrix and hidden layer are well-aligned, the level of perturbation caused by dropout to the subsequent linear layer might not be too large. This supports our use of Taylor expansion. Now we can also check that Taylor expanding around $Z\delta=0$ gives the same regularizers. (This is unsurprising because the form of Taylor expansion is invariant to linear transformation.) Using $h_{i+1}^{\prime}$ to refer to $Zh_{i}$, we have

Substituting these back into (A.7) brings us back to (4.1), which served as the starting point for the derivations of our explicit and implicit regularizers. Thus, we obtain the same analytic expressions by performing Taylor expansion around $Z\delta$.

Appendix B Proof of Theorem 4.1

Our proof of Theorem 4.1 will rely on the following slightly more general statement for loss functions with an exponential tail.

where $\theta\triangleq\max_{i}\|x_{i}\|^{2}\log^{3}(nc)$ and $\mu$, $\nu$ measure the Jacobians and Hessians of the loss and are defined by

We provide the full proof of Lemma B.1 in Section B.3.

Thus, it suffices to prove Theorem B.1. To do so, we will rely on the following lemmas.

where $G_{W,\alpha}$ is the data-dependent function defined by

where $\mu(W),\nu(W)$ are defined as in (B.4) and (B.5) which (implicitly) depend on the training data, and $\rho^{\prime}>0$ is some universal constant.

In the setting of Lemma B.2, let $G_{W,\alpha}(\beta)$ be defined as in (B.7). Define $\theta\triangleq\log^{3}(nc)\kappa^{2}$. Then

where $\mu,\nu$ are defined in (B.4) and (B.5).

We prove this Lemma in Section B.2. We now can prove Theorem B.1 by combining the lemmas above.

Combining Lemmas B.2 and Lemma B.3 and choosing $\alpha=100$, we get the desired result. ∎

We apply Lemma B.4 for $\sigma=\beta_{0},\ldots,\beta_{J}$ using probability $\delta/|\mathcal{S}|$ and union bound over the failure probability. This allows us to conclude that with probability $1-\delta$, for all $\|W^{\top}\|_{2,1}\leq A$ and $\beta\in\mathcal{S}$,

where $\rho_{1},\rho_{2}$ are universal constants and $G_{W,\alpha}(\beta)$ is defined in (B.7). (Note that $G_{W,\alpha}(\beta)$ depends on $W$ and the training data.) For a fixed choice of $W$, let $\beta^{\star}\triangleq\arg\min_{\beta>0}G_{W,\alpha}(\beta)$. First, if $\beta^{\star}\in[\beta_{0},\beta_{J}]$, then by construction $\exists\widebar{\beta}\in\mathcal{S}$ such that $\widebar{\beta}\in[\beta^{\star}/e,\beta^{\star}]$ where $e$ denotes the mathematical constant. For this choice of $j$, we have $G_{W,\alpha}(\widebar{\beta})\leq e^{2}G_{W,\alpha}(\beta^{\star})$ (as the third term in $G_{W,\alpha}(\beta)$ is the only decreasing term in $\beta$, and it differs by a factor of at most $e^{2}$ from $\widebar{\beta}$ to $\beta^{\star}$.)

Finally, in the case when $\beta^{\star}>\beta_{J}$, we note that $G_{W,\alpha}(\beta_{J})-G_{W,\alpha}(\beta^{\star})\lesssim\frac{B}{n}$. This is again because only the third term in $G_{W,\alpha}(\beta)$ is decreasing in $\beta$, and for $\beta>\beta_{J}$, this term is at most $(1+\alpha)B/n$.

Thus, for all choices of $\beta^{\star}$, we can conclude that

for universal constants $\rho_{3},\rho_{4}>0$. This gives the desired result. ∎

satisfying the following: for all $W\in\mathcal{W}(r)$ , there exists $\widebar{W}\in\widebar{\mathcal{W}}_{\epsilon}$ such that

Applying Claims B.1 and B.2 lets us bound the covering number of $\widehat{\mathcal{L}}_{\sigma}(r)$.

In the above setting, we have the covering number bound

Let $\widebar{\mathcal{W}}_{\epsilon}$ be the set of matrices inducing the $\epsilon$ cover of $\{Wx_{i}/\|x_{i}\|_{2}\}_{i=1}^{n}$ whose cardinality is bounded in Claim B.2. For any $\widebar{W}\in\widebar{\mathcal{W}}_{\epsilon}$, we can compute

This translates into the following Rademacher complexity bound for $\widehat{\mathcal{L}}_{\sigma}(r)$:

We apply Dudley’s entropy integral using the covering number bound in Claim B.3. This mirrors the calculation used to prove Lemma 2.2 in (Srebro et al., 2010). From Lemma A.1 of (Srebro et al., 2010), we have

Now we perform a change of variables $\epsilon=\frac{\sqrt{8Br}}{\sigma}\epsilon^{\prime}$, after which (B.25) becomes

To change the upper limit of the integral from $\sigma/\sqrt{8}$ to $\min\{2A,\sigma/\sqrt{8}\}$, we used the fact that we only need to integrate $\epsilon^{\prime}$ to $2A$, because for $\epsilon^{\prime}>2A$ the log covering number is 0 by Claim B.3. Now we plug in $\alpha=\frac{A\sqrt{8Br}}{\sigma\sqrt{n}}$ into (B.27) and use the fact that we only integrate over $\epsilon^{\prime}>A/\sqrt{n}$ to obtain (after simplification):

Using Claim B.4, we can complete the proof of Lemma B.4 using the technique of (Srebro et al., 2010), which is essentially local Rademacher complexity (Bousquet, 2002).

Define $\psi(r)\triangleq\frac{3500A\sqrt{Br\log^{3}(35nc)}}{\sigma\sqrt{n}}$. By Claim B.4, we have $\widehat{\mathfrak{R}}(\widehat{\mathcal{L}}_{\sigma}(r))\leq\psi(r)$.

By the AM-GM inequality, for all $\alpha>0$, we have

where $r^{\star}$ satisfies $\psi(r^{\star})=r^{\star}$. Now using the fact that $r^{\star}\lesssim\frac{BA^{2}\log^{3}(nc)}{\sigma^{2}n}$, we obtain (B.12) for some $\rho>0$. ∎

Let $\delta^{\star}$ be the optimal perturbation for $W$ and $x$, i.e.

Using the same reasoning, we can also obtain

Squaring both sides gives the desired result. ∎

Furthermore, setting $\widetilde{\mathcal{W}}_{\epsilon}\triangleq\cup_{v\in\widebar{\mathcal{V}}}\widetilde{\mathcal{W}}_{\epsilon}(v)$, we thus have

To obtain the last line, we use the fact that $\widetilde{W}_{\epsilon}$ has cardinality 0 when $\epsilon\geq A$. It remains to show that $\widetilde{W}$ satisfies the desired error properties. For any $W$ satisfying $\|W^{\top}\|_{2,1}\leq A$, there exists $\widebar{v}\in\widebar{\mathcal{V}}$ satisfying

Furthermore, by construction there exists $\widebar{W}\in\widetilde{\mathcal{W}}_{\epsilon}(\widebar{v})$ satisfying

It follows that for all $i=1,\ldots,n$, we have

To conclude the statement of the lemma, we note that for each element $\widebar{W}\in\widetilde{\mathcal{W}}_{\epsilon/2}$, we can add to $\widebar{\mathcal{W}}_{\epsilon}$ a single $\widebar{W}^{\prime}\in\mathcal{W}(r)$ satisfying

Then $\widebar{\mathcal{W}}_{\epsilon}$ will be the desired cover with cardinality bounded by

B.1.2 Proof of Lemma B.5

We will rely on the following bound on the change of a function satisfying (B.1).

The proof of this claim mirrors the proof of Proposition 1 in (Bach et al., 2010). ∎

B.2 Proof of Lemma B.3

We will rely on the following statement regarding the optimum of a function which shows up in our proof for Lemma B.3.

We drop the $W,\alpha$ dependency in the notation for simplicity. Define

To see this, let $\beta^{\star},\beta_{1},\beta_{2}$ be the minimizers for $G,G_{1}+G_{3},G_{2}+G_{3}$, respectively, and assume without loss of generality that $\beta_{2}\geq\beta_{1}$. Then we have

which gives the desired statement from the definitions of $\beta_{1},\beta_{2}$. Thus, it suffices to minimize $G_{1}+G_{3}$, $G_{2}+G_{3}$ separately.

To bound the minimum of $G_{1}+G_{3}$, we observe that to minimize any function of the form $a_{1}\beta+a_{2}/\beta^{2}$, we can set $\beta=(a_{1})^{-1/3}a_{2}^{1/3}$. Applying this to $G_{1}+G_{3}$ gives

Next, we bound the minimum of $G_{2}+G_{3}$. As the function $\exp(a)-a-1$ is increasing in $a$, we have

where $\nu$ is defined in (B.5). Let $\widebar{G}_{2}(\beta)$ denote the right-hand side of the above equation. Now we can invoke Claim B.6 on the variable $\tau\beta\kappa$ to conclude that

Finally, invoking (B.56) and applying the definition of $\theta$ gives the desired result. ∎

First consider the case when $a_{1}>a_{2}$. In this case, set ${\beta^{\prime}}^{2}=\sqrt{a_{2}/a_{1}}<1$. As $\exp(b)-b-1\leq b^{2}$ for $0\leq b<1$, we have in this case $G(\beta^{\prime})\leq\sqrt{a_{1}a_{2}}$.

Otherwise, consider the case when $a_{2}>a_{1}$ but $\frac{a_{2}}{a_{1}}\leq\log^{2}\frac{a_{2}}{a_{1}}$, $\log^{2}\frac{a_{2}}{a_{1}}<1$, or $\log\frac{a_{2}}{a_{1}}\leq 2\log\log^{2}\frac{a_{2}}{a_{1}}+2$. If any of these three equations hold, then $a_{2}/a_{1}$ is upper bounded by some universal constant, in which case setting $\beta^{\prime}=1$ immediately gives $G(\beta^{\prime})\lesssim\sqrt{a_{1}a_{2}}$.

Otherwise, consider the case when $\frac{a_{2}}{a_{1}}>\log^{2}\frac{a_{2}}{a_{1}}$, $\log^{2}\frac{a_{2}}{a_{1}}\geq 1$, and $\log^{2}\frac{a_{2}}{a_{1}}>2\log\log^{2}\frac{a_{2}}{a_{1}}+2$ all hold. In this case, we set $\beta^{\prime}=\log\left(\frac{a_{2}}{a_{1}\log^{2}\frac{a_{2}}{a_{1}}}\right)$. Note that $\beta^{\prime}>0$ and $\beta^{\prime}=\log\frac{a_{2}}{a_{1}}-\log\log^{2}\frac{a_{2}}{a_{1}}>\frac{1}{2}\log\frac{a_{2}}{a_{1}}+1$ by our conditions on $a_{2}/a_{1}$. Then we have

Combining the three cases gives the desired claim. ∎

B.3 Additional Proofs of Helper Lemmas

We first provide the proof of Lemma B.1. We use $p$ to denote the vector of probabilities predicted by the cross entropy loss, formally defined in Section E. We will rely on the derivatives of the cross entropy loss computed in Section E.

Plugging this back into the previous equation, we obtain

Next, the following statement is useful for Claim B.1.

In the setting of Claim B.1, for any $a,b>0$, we have

First, in the case where $a+b\geq\sigma,a\geq\sigma$, we have $s_{\sigma}(a+b)=s_{\sigma}(a)=0$ so the inequality trivially holds.

Second, in the case where $a+b\geq\sigma,a<\sigma$, we have $s_{\sigma}(a+b)=0$, so the LHS of (B.71) is simply $s_{\sigma}(a)^{2}$. Now we note that $b\geq\sigma-a$, so $b^{2}/\sigma^{2}\geq(1-a/\sigma)^{2}=s_{\sigma}(a)$. Thus, we have $s_{\sigma}(a)^{2}\leq s_{\sigma}(a)b^{2}/\sigma^{2}$, so (B.71) follows.

Third, in the case where $a+b<\sigma,a<\sigma$, we have

B.4 Discussion of Bound

Consider the case where the empirical distribution is only supported on $c^{\prime}\ll c$ tightly clustered classes and all the datapoints $x_{i}$ have norm 1. Suppose that the norms of the rows of $W$ are balanced and concentrated on the $c^{\prime}$ classes in the empirical sample. Let $p$ denote the softmax probability vector defined in (E.3). Consider weights $W$ which are well aligned with the data, so that $1-p(Wx_{i})_{y_{i}}\approx\exp(-\|\vec{1}_{y_{i}}^{\top}W\|_{2})\approx\exp(-\|W\|_{F}/\sqrt{c^{\prime}})$, for every training example $(x_{i},y_{i})$ (this is possible because the softmax probability vector is exponential-tailed).

Appendix C Additional Implementation Details

We implement our code in PyTorch, basing our LSTM implementation on the following code: https://github.com/salesforce/awd-lstm-lm. We base our Transformer-XL implementation on the following code: https://github.com/kimiyoung/transformer-xl. Code for downloading and pre-processing the datasets which we use are also contained in these repositories. We run our code on NVIDIA TitanXp GPUs. We provide detailed descriptions of the algorithms we implement below.

This relationship formally proved in Claim E.1. Note in particular that our regularizer $R_{\textup{approx}}$ does not depend on the true label $y$ since the loss Hessian $H_{\textup{out}}$ is independent of $y$. Algorithm 3 leverages this relationship to compute an unbiased estimator for $R_{\textup{approx}}$ based on (C.1) by sampling a label $\widehat{y}\sim\textup{softmax}(F(x))$ and computing the loss Jacobian for label $\widehat{y}$ instead of the full Hessian matrix $H_{\textup{out}}$.When computing the gradient update, we do not differentiate through the sampling probabilities $\textup{softmax}(F(x))$. Thus, though our loss estimate is unbiased, our estimate of $\nabla_{W}R_{\textup{approx}(}F,x)$ is biased. This does not appear to matter in practice. We formally prove the correctness of Algorithm 3 below.

Algorithm 3 gives an unbiased estimate of $R_{\textup{approx}}$ defined in (4.4).

Summing over $i$ gives the desired result. ∎

Algorithm 3 admits a straightforward extension to the adaptive softmax loss (Grave et al., 2017) which computes the derivative for the loss with respect to a sampled cluster label and sampled word within the cluster.

C.2 Implementing Figure 4 Experiment

Algorithm 4 describes more formally how to implement $\textsc{Dropout}_{k}$ with injection of noise $\xi_{\textup{approx}}$, which was plotted in Figure 4.

C.3 Using Identity Instead of Loss Hessian

It is non-trivial to implement this experiment described in Section 5.2, as the dimensionality of the output is large and naively computing the regularizer $\left\langle{J_{F,{i}}^{\top}J_{F,{i}}},{\textup{diag}(h_{i}^{\odot 2})}\right\rangle$ requires computing the output Jacobian $J_{F,{i}}$ exactly. To circumvent this issue, we use sampling. Letting $\eta_{i}$ be a random vector whose coordinates are independently and uniformly sampled from $\{-1,+1\}$, we have

Now to compute the value $J_{F,{i}}(\eta_{i}\odot h_{i})$ we use the method for computing Jacobian vector products described here: https://j-towns.github.io/2017/06/12/A-new-trick.html.

Appendix D Additional Experimental Results

We use a coefficient of $\lambda_{1}=2/3$ for our explicit regularizer. For our implicit regularizer, for the experiments corresponding to Figure 4, we provide implementation details in Algorithm 4. For the experiments which combine our explicit and implicit regularizers, we implement $\xi_{\textup{approx}}$ by sampling $\eta$ with coordinates independently and uniformly distributed in $\{-1,+1\}$, and computing $\xi_{\textup{approx}}(F,x,\eta)$ with coefficient $\lambda_{2}=\sqrt{2/3}$ in Algorithm 2. This is meant to match the dropout probability of 0.4.

In Tables 3 and 4, we summarize our results across the experiments in Section 5.1.

D.2 Additional Results for Section 5.3

For our Transformer-XL experiments, we use the hyperparameters for the base model on WikiText-103 contained in the following repository: https://github.com/kimiyoung/transformer-xl/. We use an explicit regularization coefficient of $\lambda=0.11$ to match the dropout probability $q=0.1$. Our explicit regularizer takes 3 times longer per iteration than dropout.

We also examine the implicit regularization effect of dropout on the QRNN architecture for WikiText-103. We use a batch size of 15 with the Adam optimizer and an initial learning rate of 5e-4 for all our runs. We chose these parameters to fit the $\textsc{Dropout}_{k}$ updates in memory. The other hyperparameters are set to their defaults in the awd-lstm repository. We chose to use QRNN as they are faster to train than LSTMs (Bradbury et al., 2016; Merity et al., 2018). Table 5 demonstrates that the implicit regularization effect does not appear on the WikiText-103 dataset, which also matches our observations for the Transformer-XL architecture on this same dataset. This suggests that the dataset size, and not architecture, influences whether the implicit regularization effect appears.

Appendix E Useful Properties of Cross-Entropy Loss

The derivatives of the cross-entropy loss are given as follows:

where $p(v)$ is the softmax probability vector given by

Furthermore, we have the following relationship between the first and second derivatives of the loss:

Note that the expectation of $\vec{1}_{\widehat{y}}$ for $\widehat{y}\sim p(v)$ is simply $p(v)$. Thus, the right hand side simplifies to the desired statement. ∎