An Inequality Between Input Gradient Penalty and Parameter Gradient Penalty
In this blog, gradient penalties have already come up many times. In terms of form, gradient penalty terms come in two kinds: one is the gradient penalty with respect to the input $\Vert\nabla_{\boldsymbol{x}} f(\boldsymbol{x};\boldsymbol{\theta})\Vert^2$, discussed in posts such as "A Brief Discussion of Adversarial Training: Meaning, Methods, and Thoughts (with Keras Implementation)" and "Random Thoughts on Generalization: From Random Noise and Gradient Penalty to Virtual Adversarial Training"; the other is the gradient penalty with respect to the parameters $\Vert\nabla_{\boldsymbol{\theta}} f(\boldsymbol{x};\boldsymbol{\theta})\Vert^2$, discussed in posts such as "A Dynamical Perspective on Optimizers (V): Why Shouldn't the Learning Rate Be Too Small?" and "Do We Really Need to Reduce the Training Loss to Zero?".
In these related posts, both kinds of gradient penalty are claimed to improve the generalization performance of the model, so is there any connection between the two? I learned of an inequality connecting them from a recent Google paper, "The Geometric Occam's Razor Implicit in Deep Learning", which partially answers this question, and since it feels like something that might come in handy later, I'm writing it down here as a note.
The Final Result
Suppose we have an $l$-layer MLP model, written as
\begin{equation}\boldsymbol{h}^{(t+1)} = g^{(t)}(\boldsymbol{W}^{(t)}\boldsymbol{h}^{(t)}+\boldsymbol{b}^{(t)})\end{equation}
where $g^{(t)}$ is the activation function of the current layer, $t\in\{1,2,\cdots,l\}$, and we denote $\boldsymbol{h}^{(1)}$ as $\boldsymbol{x}$, i.e., the original input to the model; for convenience in the derivation below, we write $\boldsymbol{z}^{(t+1)}=\boldsymbol{W}^{(t)}\boldsymbol{h}^{(t)}+\boldsymbol{b}^{(t)}$; the full set of parameters is $\boldsymbol{\theta}=\{\boldsymbol{W}^{(1)},\boldsymbol{b}^{(1)},\boldsymbol{W}^{(2)},\boldsymbol{b}^{(2)},\cdots,\boldsymbol{W}^{(l)},\boldsymbol{b}^{(l)}\}$. Let $f$ be any scalar function of $\boldsymbol{h}^{(l+1)}$; then the following inequality holds:
\begin{equation}\Vert\nabla_{\boldsymbol{x}} f\Vert^2\left(\frac{1 + \Vert \boldsymbol{h}^{(1)}\Vert^2}{\Vert\boldsymbol{W}^{(1)}\Vert^2 \Vert\nabla_{\boldsymbol{x}}\boldsymbol{h}^{(1)}\Vert^2}+\cdots+\frac{1 + \Vert \boldsymbol{h}^{(l)}\Vert^2}{\Vert\boldsymbol{W}^{(l)}\Vert^2 \Vert\nabla_{\boldsymbol{x}}\boldsymbol{h}^{(l)}\Vert^2}\right)\leq \Vert\nabla_{\boldsymbol{\theta}} f\Vert^2\label{eq:f}\end{equation}more
where in the above, $\Vert\nabla_{\boldsymbol{x}} f\Vert$, $\Vert\nabla_{\boldsymbol{\theta}} f\Vert^2$, and $\Vert \boldsymbol{h}^{(i)}\Vert$ use the ordinary $l_2$ norm—i.e., the square root of the sum of squares of each element—while $\Vert\boldsymbol{W}^{(1)}\Vert$ and $\Vert\nabla_{\boldsymbol{x}}\boldsymbol{h}^{(1)}\Vert$ use the "spectral norm" of the matrix (see Lipschitz Constraints in Deep Learning: Generalization and Generative Models). This inequality shows that the gradient penalty on the parameters, to some extent, subsumes the gradient penalty on the input.
Derivation
Clearly, to prove inequality $\eqref{eq:f}$, it suffices to prove, for each individual parameter:
\begin{align}\Vert\nabla_{\boldsymbol{x}} f\Vert^2\left(\frac{\Vert \boldsymbol{h}^{(t)}\Vert^2}{\Vert\boldsymbol{W}^{(t)}\Vert^2 \Vert\nabla_{\boldsymbol{x}}\boldsymbol{h}^{(t)}\Vert^2}\right)\leq&\, \Vert\nabla_{\boldsymbol{W}^{(t)}} f\Vert^2 \label{eq:w}\\ \Vert\nabla_{\boldsymbol{x}} f\Vert^2\left(\frac{1}{\Vert\boldsymbol{W}^{(t)}\Vert^2 \Vert\nabla_{\boldsymbol{x}}\boldsymbol{h}^{(t)}\Vert^2}\right)\leq&\, \Vert\nabla_{\boldsymbol{b}^{(t)}} f\Vert^2 \label{eq:b} \end{align}
and then sum both sides over all $t$. The proof of these two inequalities is essentially a matrix-calculus problem, but most readers, like myself, may not be very comfortable with matrix calculus. In that case, the best approach is to write everything out in component form, turning it into a problem of scalar-valued derivatives.
Specifically, we write $\boldsymbol{z}^{(t+1)}=\boldsymbol{W}^{(t)}\boldsymbol{h}^{(t)}+\boldsymbol{b}^{(t)}$ in component form as:
\begin{equation}z^{(t+1)}_i = \sum_j w^{(t)}_{i,j} h_j^{(t)} + b^{(t)}_i\end{equation}
Then, by the chain rule:
\begin{equation}\frac{\partial f}{\partial x_i} = \sum_{j,k} \frac{\partial f}{\partial z^{(t+1)}_j} \frac{\partial z^{(t+1)}_j}{\partial h^{(t)}_k} \frac{\partial h^{(t)}_k}{\partial x_i} = \sum_{j,k} \frac{\partial f}{\partial z^{(t+1)}_j} w^{(t)}_{j,k} \frac{\partial h^{(t)}_k}{\partial x_i}\label{eq:l}\end{equation}
and then
\begin{equation}\frac{\partial z^{(t+1)}_j}{\partial w^{(t)}_{m,n}} = \delta_{j,m}h^{(t)}_n\end{equation}
Here $\delta_{j,m}$ is the Kronecker delta. We can now write
\begin{equation}w^{(t)}_{j,k} = \sum_m \delta_{j,m}w^{(t)}_{m,k} = \sum_m \frac{\partial z^{(t+1)}_j}{\partial w^{(t)}_{m,n}} (h^{(t)}_n)^{-1} w^{(t)}_{m,k}\end{equation}
Substituting into $\eqref{eq:l}$ gives
\begin{equation}\frac{\partial f}{\partial x_i} = \sum_{j,k,m} \frac{\partial f}{\partial z^{(t+1)}_j} \frac{\partial z^{(t+1)}_j}{\partial w^{(t)}_{m,n}} (h^{(t)}_n)^{-1} w^{(t)}_{m,k} \frac{\partial h^{(t)}_k}{\partial x_i}=\sum_{k,m} \frac{\partial f}{\partial w^{(t)}_{m,n}} (h^{(t)}_n)^{-1} w^{(t)}_{m,k} \frac{\partial h^{(t)}_k}{\partial x_i}\end{equation}
Multiplying both sides by $h^{(t)}_n$ gives
\begin{equation}h^{(t)}_n\frac{\partial f}{\partial x_i} = \sum_{k,m} \frac{\partial f}{\partial w^{(t)}_{m,n}} w^{(t)}_{m,k} \frac{\partial h^{(t)}_k}{\partial x_i}\end{equation}
If we agree that the original vectors are column vectors, so that taking a gradient reverses the shape of the resulting matrix, then the above can be written in matrix form as:
\begin{equation}\boldsymbol{h}^{(t)}(\nabla_{\boldsymbol{x}} f) = (\nabla_{\boldsymbol{W}^{(t)}} f )\boldsymbol{W}^{(t)}(\nabla_{\boldsymbol{x}}\boldsymbol{h}^{(t)})\end{equation}
Left-multiplying both sides by $(\boldsymbol{h}^{(t)})^{\top}$ gives
\begin{equation}\Vert\boldsymbol{h}^{(t)}\Vert^2(\nabla_{\boldsymbol{x}} f) = (\boldsymbol{h}^{(t)})^{\top}(\nabla_{\boldsymbol{W}^{(t)}} f )\boldsymbol{W}^{(t)}(\nabla_{\boldsymbol{x}}\boldsymbol{h}^{(t)})\end{equation}
Taking norms on both sides gives
\begin{equation}\Vert\boldsymbol{h}^{(t)}\Vert^2 \Vert\nabla_{\boldsymbol{x}} f\Vert = \Vert (\boldsymbol{h}^{(t)})^{\top}(\nabla_{\boldsymbol{W}^{(t)}} f )\boldsymbol{W}^{(t)}(\nabla_{\boldsymbol{x}}\boldsymbol{h}^{(t)})\Vert \leq \Vert\boldsymbol{h}^{(t)}\Vert \Vert\nabla_{\boldsymbol{W}^{(t)}} f \Vert \Vert \boldsymbol{W}^{(t)}\Vert \Vert\nabla_{\boldsymbol{x}}\boldsymbol{h}^{(t)}\Vert\end{equation}
Regarding the second inequality: it holds whether we use the $l_2$ norm or the spectral norm for the matrix. So once we choose the appropriate norm and rearrange, we obtain formula $\eqref{eq:w}$; the proof of formula $\eqref{eq:b}$ is similar, so we won't repeat it here.
A Brief Commentary
Some readers might ask how exactly one should interpret formula $\eqref{eq:f}$. Honestly, I mainly find formula $\eqref{eq:f}$ interesting in its own right—it might come in useful in some setting down the road—so this post is really more of a "note" on it, and I don't have a particularly satisfying interpretation to offer.
As for the logical thread of the original paper, it goes like this: in "A Dynamical Perspective on Optimizers (V): Why Shouldn't the Learning Rate Be Too Small?", we discussed "Implicit Gradient Regularization" (by the same authors as this paper), which points out that SGD implicitly contains a gradient penalty term on the parameters. Formula $\eqref{eq:f}$ then shows that the parameter gradient penalty implicitly contains the input gradient penalty, and the input gradient penalty is in turn related to the Dirichlet energy, which can serve as a measure of model complexity. So, chaining the whole argument together, the conclusion is: SGD itself tends to favor models of lower complexity.
However, the original paper makes a small mistake when interpreting formula $\eqref{eq:f}$. It claims that at initialization, $\Vert \boldsymbol{W}^{(t)}\Vert$ is very close to 0, so the bracketed term in formula $\eqref{eq:f}$ would be very large; hence, in order to reduce the parameter gradient penalty on the right-hand side of formula $\eqref{eq:f}$, the input gradient penalty on the left-hand side of formula $\eqref{eq:f}$ would need to be sufficiently small. However, from "Understanding Parameter Initialization Strategies from a Geometric Perspective", we know that commonly used initialization schemes are actually close to orthogonal initialization, and the spectral norm of an orthogonal matrix is in fact 1; if we further account for the activation function, the spectral norm at initialization is actually even greater than 1. So the claim that $\Vert \boldsymbol{W}^{(t)}\Vert$ is close to 0 at initialization does not hold.
In fact, for a network that hasn't collapsed during training, both the model's parameters and the inputs/outputs of each layer will generally remain in a fairly stable state throughout training. So in practice, $\Vert \boldsymbol{h}^{(t)}\Vert$, $\Vert\boldsymbol{W}^{(t)}\Vert$, and $\Vert\nabla_{\boldsymbol{x}}\boldsymbol{h}^{(t)}\Vert$ don't fluctuate much over the course of training, which means the parameter gradient penalty on the right-hand side is approximately equivalent to the input multiplicative penalty on the left-hand side. This is my own reading of it, and it doesn't require the assumption that "$\Vert \boldsymbol{W}^{(t)}\Vert$ is close to 0."
Summary
This post introduced an inequality connecting two kinds of gradient penalty terms, along with my own proof and a brief commentary on it.
Translated automatically with claude-sonnet-5; all equations are reproduced verbatim from the source. Copyright remains with the original author.