Ramblings on Generalization: From Random Noise and Gradient Penalties to Virtual Adversarial Training
Improving a model's generalization performance is one of the goals that machine learning constantly pursues. There are two common approaches to boosting generalization: the first is adding noise—e.g., adding Gaussian noise to the input, inserting Dropout in intermediate layers, and the currently popular adversarial training; data augmentation techniques such as random translation and scaling of images also fall, in a sense, into this category. The second is adding regularization terms to the loss, such as $L_1, L_2$ penalties, gradient penalties, and so on. This post attempts to explore the connections between several common techniques for improving generalization performance.
Random Noise
Let us denote the model as $f(x)$, $\mathcal{D}$ as the training dataset, and $l(f(x), y)$ as the loss for a single sample. Then our optimization objective is
\begin{equation}\mathop{\text{argmin}}_{\theta} L(\theta)=\mathbb{E}_{(x,y)\sim \mathcal{D}}[l(f(x), y)]\end{equation}
where $\theta$ are the trainable parameters inside $f(x)$. If we add noise $\varepsilon$ to the model input, with distribution $q(\varepsilon)$, then the optimization objective becomes
\begin{equation}\mathop{\text{argmin}}_{\theta} L_{\varepsilon}(\theta)=\mathbb{E}_{(x,y)\sim \mathcal{D}, \varepsilon\sim q(\varepsilon)}[l(f(x + \varepsilon), y)]\end{equation}
Of course, noise can be added not only to the input but also to intermediate layers, to the weights $\theta$, or even to the output $y$ (which is equivalent to label smoothing); nor does the noise have to be additive—Dropout, for instance, is multiplicative. For additive noise, the common choice for $q(\varepsilon)$ is a zero-mean Gaussian with fixed variance; for multiplicative noise, common choices are the uniform distribution $U([0,1])$ or the Bernoulli distribution.
The purpose of adding random noise is intuitive: we want the model to learn to withstand random perturbations, thereby reducing its sensitivity to the input or the parameters. Reducing this sensitivity usually means the resulting model relies less heavily on the training set, which helps improve generalization performance. more
Improving Efficiency
Adding random noise is easy to implement and, in many cases, genuinely effective, but it has one obvious drawback: it isn't very "targeted." The noise $\varepsilon$ is random rather than being constructed specifically for $x$, which means that in most cases $x + \varepsilon$ may just end up being a trivial sample—i.e., it doesn't perturb the original model in any noticeable way—so its contribution to improving generalization is limited.
Increasing the Number of Samples
Theoretically, once random noise is introduced, the loss for a single sample becomes
\begin{equation}\tilde{l}(x,y)=\mathbb{E}_{\varepsilon\sim q(\varepsilon)}[l(f(x+\varepsilon),y)]=\int q(\varepsilon) l(f(x+\varepsilon),y) d\varepsilon\label{eq:noisy-loss}\end{equation}
But in practice, for each specific sample $(x,y)$, we typically only draw one noise sample, so the above expression is not well approximated. Of course, we could sample multiple noises $\varepsilon_1,\varepsilon_2,\cdots,\varepsilon_k\sim q(\varepsilon)$ and better approximate
\begin{equation}\tilde{l}(x,y)\approx \frac{1}{k}\sum_{i=1}^k l(f(x+\varepsilon_i),y)\end{equation}
but this is equivalent to expanding the batch size by a factor of $k$, which increases computational cost and is not very friendly.
Approximate Expansion
A direct idea is: if we could compute the integral in equation $\eqref{eq:noisy-loss}$ ahead of time, we would no longer need to inefficiently sample (or equivalently, we'd be sampling infinitely many noises at once). Let's try heading in this direction. Of course, computing the integral exactly in closed form is essentially impossible, but we can perform an approximate expansion:
\begin{equation}l(f(x+\varepsilon),y)\approx l(f(x),y)+(\varepsilon \cdot \nabla_x) l(f(x),y)+\frac{1}{2}(\varepsilon \cdot \nabla_x)^2 l(f(x),y)\end{equation}
Then, multiplying both sides by $q(\varepsilon)$ and integrating—assuming the components of $\varepsilon$ are i.i.d. with mean 0 and variance $\sigma^2$—the result of the integral is
\begin{equation}\int q(\varepsilon)l(f(x+\varepsilon),y)d\varepsilon \approx l(f(x),y)+\frac{1}{2}\sigma^2 \Delta l(f(x),y)\end{equation}
Here $\Delta$ is the Laplace operator, i.e., $\Delta f = \sum\limits_i \frac{\partial^2}{\partial x_i^2} f$. This result is formally simple: it amounts to adding the regularization term $\frac{1}{2}\sigma^2 \Delta l(f(x),y)$ to the loss. However, it is quite difficult in practice, because it requires computing the second-order derivative of $l$, and combined with gradient descent, this means computing third-order derivatives in total—something existing deep learning frameworks struggle to implement efficiently.
Shifting the Objective
Directly simplifying the integral in $l(f(x+\varepsilon),y)$ turns out to be infeasible, but we can instead try replacing the optimization objective with
\begin{equation}l(f(x+\varepsilon),f(x)) + l(f(x),y)\label{eq:loss-2}\end{equation}
That is, we simultaneously shrink the gap between $f(x),y$ and $f(x+\varepsilon),f(x)$, and pursuing both together can, to some extent, also achieve the goal of shrinking the gap in $f(x+\varepsilon),y$. What's important is that this objective leads to more interesting results.
Reasoning Through It
Mathematically speaking, if $l$ is some form of distance metric, then by the triangle inequality we have
\begin{equation}l(f(x+\varepsilon),y) \leq l(f(x+\varepsilon),f(x)) + l(f(x),y)\end{equation}
If $l$ is not a metric, we can usually obtain a similar result via Jensen's inequality instead, for example $l(f(x+\varepsilon),y)=\Vert f(x+\varepsilon) - y\Vert^2$, giving us
\begin{equation}\begin{aligned} \Vert f(x+\varepsilon) - f(x) + f(x) - y\Vert^2 =& \left\Vert \frac{1}{2}\times 2[f(x+\varepsilon) - f(x)] + \frac{1}{2}\times 2[f(x) - y]\right\Vert^2\\ \leq& \frac{1}{2} \Vert 2[f(x+\varepsilon) - f(x)]\Vert^2 + \frac{1}{2} \Vert 2[f(x) - y]\Vert^2\\ =& 2\big(\Vert f(x+\varepsilon) - f(x)\Vert^2 + \Vert f(x) - y\Vert^2\big) \end{aligned}\end{equation}
That is to say, (some multiple of) the objective $\eqref{eq:loss-2}$ can be regarded as an upper bound on $l(f(x+\varepsilon),y)$. Since the original objective is hard to optimize directly, we instead optimize its upper bound.
Notice that of the two terms in the objective $\eqref{eq:loss-2}$, the term $l(f(x+\varepsilon),f(x))$ measures the smoothness of the model itself and has nothing to do with the labels—it can be optimized using unlabeled data. This means it can be combined with labeled data to form a semi-supervised learning pipeline.
Bravely Computing It
For the objective $\eqref{eq:loss-2}$, the result of its integral is:
\begin{equation}\int q(\varepsilon) \big[l(f(x+\varepsilon),f(x)) + l(f(x),y)\big]d\varepsilon = l(f(x),y) + \int q(\varepsilon) l(f(x+\varepsilon),f(x)) d\varepsilon\end{equation}
Following the same old approach, let's do an approximate expansion of $\varepsilon$:
\begin{equation}\begin{aligned}l(f(x+\varepsilon),f(x))\approx &\, l(f(x),f(x)) + \left.\sum_{i,j} \frac{\partial l(F(x),f(x))}{\partial F_i(x)}\frac{\partial f_i(x)}{\partial x_j}\varepsilon_j\right|_{F(x)=f(x)}\\ &\, + \frac{1}{2}\left.\sum_{i,j,k} \frac{\partial l(F(x),f(x))}{\partial F_i(x)}\frac{\partial^2 f_i(x)}{\partial x_j \partial x_k}\varepsilon_j \varepsilon_k\right|_{F(x)=f(x)}\\ &\, + \frac{1}{2}\left.\sum_{i,i',j,k} \frac{\partial^2 l(F(x),f(x))}{\partial F_i(x) \partial F_{i'}(x)}\frac{\partial f_i(x)}{\partial x_j}\frac{\partial f_{i'}(x)}{\partial x_k}\varepsilon_j \varepsilon_k\right|_{F(x)=f(x)} \end{aligned}\label{eq:kongbu}\end{equation}
Looks terrifying? Don't worry—let's recall that a loss function $l$ generally has the following properties:
1. $l$ is smooth;
2. $l(x, x)=0$;
3. $\left.\frac{\partial}{\partial x} l(x,y)\right|_{x=y}=0,\left.\frac{\partial}{\partial y} l(x,y)\right|_{y=x}=0$.
This is really just saying that $l$ is smooth and attains its (minimum) extremum when $x=y$, with the extremum (minimum) value being 0—properties shared by almost every loss function. Based on these properties, the first three terrifying terms in equation $\eqref{eq:kongbu}$ vanish immediately, so the final integral result is:
\begin{equation}\int q(\varepsilon) l(f(x+\varepsilon),f(x)) d\varepsilon \approx \frac{1}{2}\sigma^2\left.\sum_{i,i',j} \frac{\partial^2 l(F(x),f(x))}{\partial F_i(x) \partial F_{i'}(x)}\frac{\partial f_i(x)}{\partial x_j}\frac{\partial f_{i'}(x)}{\partial x_j}\right|_{F(x)=f(x)} \end{equation}
Gradient Penalty
This may still look a bit unsettling, but it's much better than $\eqref{eq:kongbu}$. The above expression is also a regularization term, and its key feature is that it involves only a first-order gradient term. For a specific loss function, $\left.\frac{\partial^2 l(F(x),f(x))}{\partial F_i(x) \partial F_{i'}(x)}\right|_{F(x)=f(x)}$ can be computed in advance; in particular, for several common loss functions, when $i\neq i'$, $\left.\frac{\partial^2 l(F(x),f(x))}{\partial F_i(x) \partial F_{i'}(x)}\right|_{F(x)=f(x)}=0$ holds, so we only need to compute the components of $i=i'$, which we denote $\lambda_{i}(x)$. Then
\begin{equation}\int q(\varepsilon) l(f(x+\varepsilon),f(x)) d\varepsilon \approx \frac{1}{2}\sigma^2 \sum_i \lambda_i(x)\Vert \nabla_x f_i(x)\Vert^2\label{eq:gp}\end{equation}
As can be seen, formally this amounts to computing a gradient penalty term $\Vert \nabla_x f_i(x)\Vert^2$ for each component of every $f(x)$, then taking a weighted sum according to $\lambda_i(x)$.
For example, for MSE, $l(f(x),y)=\Vert f(x) - y\Vert^2$, in which case we can compute $\lambda_i(x)\equiv 2$, so the corresponding regularization term is $\sum\limits_i\Vert \nabla_x f_i(x)\Vert^2$; for KL divergence, $l(f(x),y)=\sum\limits_i y_i \log \frac{y_i}{f_i(x)}$, in which case $\lambda_i(x)=\frac{1}{f_i(x)}$, so the corresponding regularization term is $\sum\limits_i f_i(x) \Vert \nabla_x \log f_i(x)\Vert^2$. Readers may find results similar to these, more or less, in the famous "Flower Book" Deep Learning—these are not new results. A similar derivation can also be found in Training with noise is equivalent to Tikhonov regularization.
Sampling Approximation
Of course, although we've derived the regularization term $\sum\limits_i \lambda_i(x)\Vert \nabla_x f_i(x)\Vert^2$ involving only a first-order gradient, in fact this computation is still not cheap, because we need to compute the gradient for every $f_i(x)$, and if the number of output components is too large, this computational cost remains prohibitive.
At this point we can consider approximating the computation via sampling: assuming $q(\eta)$ is a distribution with mean 0 and variance 1, we have
\begin{equation}\sum\limits_i \Vert \nabla_x f_i(x)\Vert^2=\sum\limits_i \left\Vert \nabla_x f_i(x)\right\Vert^2=\mathbb{E}_{\eta_i\sim q(\eta)}\left[\left\Vert\sum_i \eta_i \nabla_x f_i(x)\right\Vert^2\right]\end{equation}
This way, at each step we only need to compute the gradient of $\sum\limits_i \eta_i f_i(x)$, rather than computing the gradient multiple times. The simplest choice for $q(\eta)$ is a uniform distribution over the space $\{-1,1\}$, i.e., $\eta_i$ picks, with equal probability, one element from $\{-1,1\}$.
Adversarial Training
Looking back over the discussion so far: we first introduced adding random noise as a means of enhancing generalization, then pointed out that random noise addition may not be targeted enough, which led us to try computing the integral first—yielding the results on approximate expansion and gradient penalties derived above. Turning this around, if we could find a way to construct the noise signal more specifically, we could likewise improve training efficiency and enhance generalization performance.
Supervised Adversarial Training
Supervised adversarial training focuses on the original objective $\eqref{eq:noisy-loss}$, and the optimization goal is to make the loss as small as possible. So if we want to choose a more representative noise, we should choose the noise that makes the loss as large as possible, and since
\begin{equation}l(f(x + \varepsilon), y) \approx l(f(x), y) + \varepsilon \cdot \nabla_x l(f(x), y)\end{equation}
making $l(f(x + \varepsilon), y)$ as large as possible means $\varepsilon$ should point in the same direction as $\nabla_x l(f(x), y)$—in other words, the perturbation should move along the direction of gradient ascent, i.e.,
\begin{equation}\varepsilon \sim \nabla_x l(f(x), y)\end{equation}
This constitutes the FGM method in adversarial training, which was already introduced previously in A Brief Discussion on Adversarial Training: Meaning, Methods, and Reflections (with Keras Implementation).
It's worth noting that in A Brief Discussion on Adversarial Training: Meaning, Methods, and Reflections (with Keras Implementation) we also showed that adversarial training is, to some extent, equivalent to adding the gradient penalty term $\left\Vert\nabla_x l(f(x), y)\right\Vert^2$ to the loss—which again resembles the noise-integral result from the previous section. This suggests that the gradient penalty is likely a general-purpose technique for improving model performance.
Virtual Adversarial Training
Earlier we mentioned that the term $l(f(x+\varepsilon),f(x))$ doesn't require label signals, and hence can be used for unsupervised learning; by expanding its Gaussian integral we obtained the gradient penalty $\eqref{eq:gp}$. Following the logic of adversarial training, instead of computing the integral, we can search for the perturbation noise that makes $l(f(x+\varepsilon),f(x))$ as large as possible—this constitutes "Virtual Adversarial Training (VAT)," first introduced in the paper Virtual Adversarial Training: A Regularization Method for Supervised and Semi-Supervised Learning.
Based on our earlier discussion of the properties of the loss function $l$, we know that the first-order gradient of $l(f(x+\varepsilon),f(x))$ with respect to $\varepsilon$ is 0, so to compute the adversarial perturbation we must expand to second order:
\begin{equation}\begin{aligned} l(f(x+\varepsilon),f(x))\approx&\, l(f(x),f(x)) + \varepsilon^{\top} \nabla_x l(f(x),f_{ng}(x)) + \frac{1}{2}\varepsilon^{\top}\nabla_x^2 l(f(x),f_{ng}(x)) \varepsilon\\ =&\, \frac{1}{2}\varepsilon^{\top}\nabla_x^2 l(f(x),f_{ng}(x)) \varepsilon\end{aligned}\end{equation}
Here $f_{ng}(x)$ indicates that we do not need to take the gradient with respect to the $x$ inside it. This leaves us with two problems to solve: 1) how to efficiently compute the Hessian matrix $\mathcal{H}=\nabla_x^2 l(f(x),f_{ng}(x))$; 2) how to find the unit vector $u$ that maximizes $u^{\top}\mathcal{H}u$?
In fact, it's not hard to show that the optimal solution to $u$ is precisely "the eigenvector corresponding to the largest eigenvalue of $\mathcal{H}$," also called the "dominant eigenvector of $\mathcal{H}$." An effective method for approximating the dominant eigenvector is the "power iteration method": starting from a random vector $u_0$, iteratively execute $u_{i+1}=\frac{\mathcal{H}u_i}{\Vert\mathcal{H}u_i\Vert}$. For the relevant derivation, see the "Dominant Eigenvalue" and "Power Iteration" sections of Lipschitz Constraints in Deep Learning: Generalization and Generative Models.
In power iteration, we find that we don't actually need to know the value of $\mathcal{H}$ explicitly—only the value of $\mathcal{H}u$, which can be approximated via finite differences:
\begin{equation}\begin{aligned}\mathcal{H}u =&\, \nabla_x^2 l(f(x),f_{ng}(x)) u\\ =&\, \nabla_x \big(u\cdot\nabla_x l(f(x),f_{ng}(x))\big)\\ \approx&\, \nabla_x \left(\frac{l(f(x + \xi u),f_{ng}(x)) - l(f(x),f_{ng}(x))}{\xi}\right)\\ =&\, \frac{1}{\xi}\nabla_x l(f(x + \xi u),f_{ng}(x))\end{aligned}\end{equation}
where $\xi$ is a scalar constant. Based on this approximation, we arrive at the following VAT procedure:
Initialize the vector $u\sim \mathcal{N}(0,1)$, and scalars $\epsilon$ and $\xi$;
Iterate $r$ times:
$u \leftarrow \frac{u}{\Vert u\Vert}$;
$u \leftarrow \nabla_x l(f(x+\xi u), f_{ng}(x))$
$u \leftarrow \frac{u}{\Vert u\Vert}$;
Use $l(f(x+\epsilon u), f_{ng}(x))$ as the loss and perform ordinary gradient descent.
Experiments show that a single iteration is generally already quite good, while iterating zero times reduces to simply adding Gaussian noise, as mentioned at the start of this post. This shows that virtual adversarial training improves the "specificity" of the noise via $\nabla_x l(f(x+\xi u), f_{ng}(x))$.
Reference Implementation
A Keras implementation of adversarial training was already given in A Brief Discussion on Adversarial Training: Meaning, Methods, and Reflections (with Keras Implementation); here the author provides a reference Keras implementation of virtual adversarial training:
def virtual_adversarial_training(
model, embedding_name, epsilon=1, xi=10, iters=1
):
"""给模型添加虚拟对抗训练
其中model是需要添加对抗训练的keras模型，embedding_name
则是model里边Embedding层的名字。要在模型compile之后使用。
"""
if model.train_function is None: # 如果还没有训练函数
model._make_train_function() # 手动make
old_train_function = model.train_function # 备份旧的训练函数
# 查找Embedding层
for output in model.outputs:
embedding_layer = search_layer(output, embedding_name)
if embedding_layer is not None:
break
if embedding_layer is None:
raise Exception('Embedding layer not found')
# 求Embedding梯度
embeddings = embedding_layer.embeddings # Embedding矩阵
gradients = K.gradients(model.total_loss, [embeddings]) # Embedding梯度
gradients = K.zeros_like(embeddings) + gradients[0] # 转为dense tensor
# 封装为函数
inputs = (
model._feed_inputs + model._feed_targets + model._feed_sample_weights
) # 所有输入层
model_outputs = K.function(
inputs=inputs,
outputs=model.outputs,
name='model_outputs',
) # 模型输出函数
embedding_gradients = K.function(
inputs=inputs,
outputs=[gradients],
name='embedding_gradients',
) # 模型梯度函数
def l2_normalize(x):
return x / (np.sqrt((x**2).sum()) + 1e-8)
def train_function(inputs): # 重新定义训练函数
outputs = model_outputs(inputs)
inputs = inputs[:2] + outputs + inputs[3:]
delta1, delta2 = 0.0, np.random.randn(*K.int_shape(embeddings))
for _ in range(iters): # 迭代求扰动
delta2 = xi * l2_normalize(delta2)
K.set_value(embeddings, K.eval(embeddings) - delta1 + delta2)
delta1 = delta2
delta2 = embedding_gradients(inputs)[0] # Embedding梯度
delta2 = epsilon * l2_normalize(delta2)
K.set_value(embeddings, K.eval(embeddings) - delta1 + delta2)
outputs = old_train_function(inputs) # 梯度下降
K.set_value(embeddings, K.eval(embeddings) - delta2) # 删除扰动
return outputs
model.train_function = train_function # 覆盖原训练函数
# 写好函数后，启用虚拟对抗训练只需要一行代码
virtual_adversarial_training(model_vat, 'Embedding-Token')
For the complete usage script, see: task_sentiment_virtual_adversarial_training.py. Roughly speaking, the model is built twice: one model is trained normally on labeled data, and the other is trained via virtual adversarial training on unlabeled data, with the two alternating. Please make sure you understand the source code before using it—don't just copy it blindly. The experimental task is sentiment classification with roughly 20,000 labeled samples; the first 200 are used as labeled samples, and the rest are treated as unlabeled data. The comparison between VAT and non-VAT performance is as follows (each experiment repeated three times, averaged):
$$\begin{array}{c|cc} \hline & \text{val set} & \text{test set}\\ \hline \text{non-VAT} & 88.93\% & 89.34\%\\ \text{VAT} & 89.83\% & 90.37\%\\ \hline \end{array}$$
Note:
As mentioned earlier, $f_{ng}(x)$ indicates that we do not take the gradient with respect to $x$; however, the gradient with respect to $f$'s own parameters $\theta$ is still required. But readers who work through the code above will notice that this implementation effectively drops the gradient with respect to $x,\theta$ inside $f_{ng}(x)$ entirely, so it isn't theoretically fully equivalent to standard VAT. The issue is that implementing standard VAT in Keras is somewhat cumbersome and increases computational cost; moreover, experiments show that this "knock-off" version already delivers improvements, and the difference from standard VAT is a second-order small quantity, so the discrepancy is minor. The code above basically meets our needs.
Summary
This post first introduced the addition of random noise as a conventional regularization technique, then, through the process of approximate expansion and integration, derived its connection to the gradient penalty, from which we obtained a model-smoothness loss usable for semi-supervised training. We then further connected this to supervised adversarial training and semi-supervised virtual adversarial training, and finally gave a Keras implementation and example of virtual adversarial training.
Translated automatically with claude-sonnet-5; all equations are reproduced verbatim from the source. Copyright remains with the original author.