L2 Regularization May Not Be as Good as We Thought — "Weight Scale Shift" Might Be to Blame
L2 regularization is a commonly used technique in machine learning for preventing overfitting (and probably a frequently asked interview question too). Roughly speaking, it aims to keep the norm of the weights as small as possible, so that the model can better withstand perturbations, ultimately improving generalization. But readers may also have noticed that L2 regularization often doesn't perform as well in practice as the theory suggests — in many cases, adding it can even have a negative effect. A recent paper, Improve Generalization and Robustness of Neural Networks via Weight Scale Shifting Invariant Regularizations, analyzes the shortcomings of L2 regularization from the angle of "weight scale shift" and proposes a new regularization term called WEISSI. The whole analysis is quite interesting, so let me share it here.
Background
In this section we'll briefly review L2 regularization, and then introduce its connection to weight decay and the related AdamW optimizer.
Understanding L2 Regularization
Why do we add L2 regularization? There may be multiple answers to this question. Some answer it from the perspective of Ridge regression, others from Bayesian inference. Here I'll give an understanding based on perturbation sensitivity. more
For two (column) vectors $\boldsymbol{w},\boldsymbol{x}$, we have the Cauchy inequality $\left|\boldsymbol{w}^{\top}\boldsymbol{x}\right|\leq \Vert\boldsymbol{w}\Vert_2\Vert\boldsymbol{x}\Vert_2$. Based on this result, we can prove that
\begin{equation}\Vert\boldsymbol{W}(\boldsymbol{x}_2 - \boldsymbol{x}_1)\Vert_2\leq \Vert\boldsymbol{W}\Vert_2\Vert\boldsymbol{x}_2 - \boldsymbol{x}_1\Vert_2\end{equation}
Here, $\Vert\boldsymbol{W}\Vert_2^2$ equals the sum of the squares of all the elements of matrix $\boldsymbol{W}$. The proof isn't difficult, and interested readers can work it out themselves. This result tells us: the change in $\boldsymbol{W}\boldsymbol{x}$ can be controlled by $\Vert\boldsymbol{W}\Vert_2$ and $\Vert\boldsymbol{x}_2 - \boldsymbol{x}_1\Vert_2$, so if we want the change in $\boldsymbol{W}\boldsymbol{x}$ to be as small as possible when $\Vert\boldsymbol{x}_2 - \boldsymbol{x}_1\Vert_2$ is small, we can reduce $\Vert\boldsymbol{W}\Vert_2$. At this point, we can add a regularization term $\mathcal{L}_{reg}=\Vert\boldsymbol{W}\Vert_2^2$ to the task objective $\mathcal{L}_{task}$. It's not hard to see that this is exactly L2 regularization. For related discussion from this angle, see also Lipschitz Constraints in Deep Learning: Generalization and Generative Models (note that the notation between the two articles differs slightly).
The AdamW Optimizer
When optimizing with SGD, suppose the original update is $\boldsymbol{\theta}_{t}=\boldsymbol{\theta}_{t-1} - \varepsilon\boldsymbol{g}_{t}$. It's not hard to show that adding L2 regularization $\Vert\boldsymbol{\theta}\Vert_2^2$ turns this into
\begin{equation}\boldsymbol{\theta}_{t}=(1-\varepsilon\lambda)\boldsymbol{\theta}_{t-1} - \varepsilon\boldsymbol{g}_{t}\end{equation}
Since $0 < 1-\varepsilon\lambda < 1$, this causes the parameters $\boldsymbol{\theta}$ to have a tendency to "shrink" toward zero throughout the optimization process, and this modification is called "weight decay."
However, the equivalence between L2 regularization and weight decay only holds under the SGD optimizer. If we use an adaptive-learning-rate optimizer like Adagrad or Adam, the two are no longer equivalent. In adaptive-learning-rate optimizers, the effect of L2 regularization is roughly equivalent to adding $-\varepsilon\lambda\text{sign}(\boldsymbol{\theta}_{t-1})$ rather than $-\varepsilon\lambda\boldsymbol{\theta}_{t-1}$ to the optimization process — that is, the penalty on each element becomes fairly uniform, rather than penalizing elements with larger absolute values more heavily, which partially cancels out the effect of L2 regularization. The paper Decoupled Weight Decay Regularization was the first to highlight this issue, and it proposed the improved AdamW optimizer.
A New Regularizer
In this section, we'll point out that a phenomenon called "weight scale shift" often exists in common deep learning models, and this phenomenon may be the reason L2 regularization's effect isn't as pronounced as expected. Furthermore, we can construct a new regularization term that has a similar effect to L2 but is more compatible with the weight scale shift phenomenon, and is therefore theoretically more effective.
Weight Scale Shift
We know that the basic structure of deep learning models is "linear transformation + nonlinear activation function," and one of the most commonly used activation functions today is $\text{relu}(x)=\max(x,0)$. Interestingly, both of these satisfy "positive homogeneity" — that is, for $\varepsilon \geq 0$, we always have $\varepsilon\phi(x)=\phi(\varepsilon x)$. As for other activation functions such as SoftPlus, GELU, and Swish, they are all smooth approximations of $\text{relu}$, so we can consider them to approximately satisfy "positive homogeneity" as well.
"Positive homogeneity" gives deep learning models a certain invariance to weight scale shift. Specifically, suppose we have an $L$-layer model:
\begin{equation}\begin{aligned} \boldsymbol{h}_L =& \phi(\boldsymbol{W}_L \boldsymbol{h}_{L-1} + \boldsymbol{b}_L) \\ =& \phi(\boldsymbol{W}_L \phi(\boldsymbol{W}_{L-1} \boldsymbol{h}_{L-2} + \boldsymbol{b}_{L-1}) + \boldsymbol{b}_L) \\ =& \cdots\\ =& \phi(\boldsymbol{W}_L \phi(\boldsymbol{W}_{L-1} \phi(\cdots\phi(\boldsymbol{W}_1\boldsymbol{x} + \boldsymbol{b}_1)\cdots) + \boldsymbol{b}_{L-1}) + \boldsymbol{b}_L) \end{aligned}\end{equation}
Suppose each parameter is shifted by $\boldsymbol{W}_l = \gamma_l\tilde{\boldsymbol{W}}_l,\boldsymbol{b}_l = \gamma_l\tilde{\boldsymbol{b}}_l$. Then, by positive homogeneity, we get
\begin{equation}\begin{aligned} \boldsymbol{h}_L =& \left(\prod_{l=1}^L \gamma_l\right)\phi(\tilde{\boldsymbol{W}}_L \boldsymbol{h}_{L-1} + \tilde{\boldsymbol{b}}_L) \\ =& \cdots\\ =& \left(\prod_{l=1}^L \gamma_l\right) \phi(\tilde{\boldsymbol{W}}_L \phi(\tilde{\boldsymbol{W}}_{L-1} \phi(\cdots\phi(\tilde{\boldsymbol{W}}_1\boldsymbol{x} + \tilde{\boldsymbol{b}}_1)\cdots) + \tilde{\boldsymbol{b}}_{L-1}) + \tilde{\boldsymbol{b}}_L) \end{aligned}\end{equation}
If $\prod\limits_{l=1}^L \gamma_l = 1$, then the model with parameters $\{\boldsymbol{W}_l,\boldsymbol{b}_l\}$ is exactly equivalent to the model with parameters $\{\tilde{\boldsymbol{W}}_l,\tilde{\boldsymbol{b}}_l\}$. In other words, the model is invariant to a weight scale shift of $\prod\limits_{l=1}^L \gamma_l = 1$ (WEIght-Scale-Shift-Invariance, WEISSI).
Incompatibility with L2 Regularization
We just said that as long as the scale shift satisfies $\prod\limits_{l=1}^L \gamma_l = 1$, the two sets of parameters correspond to equivalent models. But the problem is that their corresponding L2 regularization values are not equivalent:
\begin{equation}\sum_{l=1}^L \Vert\boldsymbol{W}_l\Vert_2^2=\sum_{l=1}^L \gamma_l^2\Vert\tilde{\boldsymbol{W}}_l\Vert_2^2\neq \sum_{l=1}^L \Vert\tilde{\boldsymbol{W}}_l\Vert_2^2\end{equation}
Furthermore, it can be shown that if we fix $\Vert\boldsymbol{W}_1\Vert_2,\Vert\boldsymbol{W}_2\Vert_2,\dots,\Vert\boldsymbol{W}_L\Vert_2$ and maintain the constraint $\prod\limits_{l=1}^L \gamma_l = 1$, then the minimum of $\sum\limits_{l=1}^L
\Vert\tilde{\boldsymbol{W}}_l\Vert_2^2$ is attained at
\begin{equation}\Vert\tilde{\boldsymbol{W}_1}\Vert_2=\Vert\tilde{\boldsymbol{W}}_2\Vert_2=\dots=\Vert\tilde{\boldsymbol{W}}_L\Vert_2=\left(\prod_{l=1}^L \Vert\boldsymbol{W}_l\Vert_2\right)^{1/L}\end{equation}
In fact, this reveals the inefficiency of L2 regularization. Imagine we've already trained a set of parameters $\{\boldsymbol{W}_l,\boldsymbol{b}_l\}$, and this set of parameters might not generalize particularly well. So we hope that L2 regularization can help the optimizer find a better set of parameters (sacrificing a bit of $\mathcal{L}_{task}$ to reduce $\mathcal{L}_{reg}$ somewhat). But the result above tells us that, due to the invariance under weight scale shift, the model can simply find a new set of parameters $\{\tilde{\boldsymbol{W}}_l,\tilde{\boldsymbol{b}}_l\}$ that is completely equivalent to the original model (with no improvement in generalization), yet has a smaller L2 regularization value (so L2 regularization is indeed "doing something"). In other words, L2 regularization really is having an effect, but it's not improving the model's generalization — which defeats the whole purpose of using L2 regularization in the first place.
The WEISSI Regularizer
The root of the above problem is that the model is invariant to weight scale shift, but L2 regularization is not. If we can find a new regularization term that has a similar effect while also being invariant to weight scale shift, we can solve this problem. Personally, I feel the original paper's explanation of this part isn't very clear, so the derivation below reflects my own understanding.
Let's consider a regularization term of the following general form:
\begin{equation}\mathcal{L}_{reg}=\sum_{l=1}^L \varphi(\Vert\boldsymbol{W}_l\Vert_2)\end{equation}
For L2 regularization, $\varphi(x)=x^2$; as long as $\varphi(x)$ is a monotonically increasing function of $x$ over $[0,+\infty)$, this ensures that the optimization objective shrinks $\Vert\boldsymbol{W}_l\Vert$. Note that for the regularization term to have scale-shift invariance, we don't actually need $\varphi(\gamma x) = \varphi(x)$ — we only need
\begin{equation}\frac{d}{dx}\varphi(\gamma x)=\frac{d}{dx}\varphi(x)\label{eq:varphi}\end{equation}
because the optimization process only uses its gradient. Some readers may already be able to spot a solution directly — it's simply the logarithmic function $\varphi(x) = \log x$. So the newly proposed regularization term is
\begin{equation}\mathcal{L}_{reg}=\sum_{l=1}^L \log\Vert\boldsymbol{W}_l\Vert_2=\log \left(\prod_{l=1}^L \Vert\boldsymbol{W}_l\Vert_2\right)\end{equation}
In addition, the original paper, perhaps worried that the above regularization term isn't strong enough, also adds an L1 penalty on the direction of the parameters, giving the overall form:
\begin{equation}\mathcal{L}_{reg}=\lambda_1\sum_{l=1}^L \log\Vert\boldsymbol{W}_l\Vert_2 + \lambda_2\sum_{l=1}^L \big\Vert\boldsymbol{W}_l\big/\Vert\boldsymbol{W}_l\Vert_2\big\Vert_1\end{equation}
Brief Summary of Experimental Results
As usual, let's show the experimental results from the original paper — and of course, since the authors bothered to write it all up, the results are naturally positive:
One of the original paper's experimental results for WEISSI regularization
For us, the main takeaway is simply knowing that this new option exists — one more thing to try when doing model training. After all, when it comes to regularization terms, there's no theoretical guarantee that any given one will actually work; you only find out by trying it yourself. No matter how convincing someone else's argument sounds, it may not actually pan out for you.
Summary
This article introduced the phenomenon of weight scale shift invariance in neural network models, pointed out its incompatibility with L2 regularization, and then proposed a regularization term that has a similar effect but resolves this incompatibility.
Translated automatically with claude-sonnet-5; all equations are reproduced verbatim from the source. Copyright remains with the original author.