Random Thoughts on Generalization: From Random Noise and Gradient Penalty to Virtual Adversarial Training

Increasing Sampling

Theoretically, after adding random noise, the loss for a single sample becomes: \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} However, in practice, for each specific sample (x,y), we usually only sample one instance of noise, so we do not approximate the above equation well. We could, of course, sample multiple noises \varepsilon_1, \varepsilon_2, \cdots, \varepsilon_k \sim q(\varepsilon) to better approximate it: \tilde{l}(x,y)\approx \frac{1}{k}\sum_{i=1}^k l(f(x+\varepsilon_i),y) But this is equivalent to expanding the batch size by k times, which increases computational costs and is not very friendly.

Approximate Expansion

A direct idea is that if we could calculate the integral in Eq. [eq:noisy-loss] beforehand, we wouldn’t need to sample inefficiently (or it would be equivalent to sampling infinite noise at once). Let’s try moving in this direction. Of course, an exact explicit integral is generally impossible, but we can perform an approximate expansion: 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) Then, multiply both sides by q(\varepsilon) and integrate. Assuming the components of \varepsilon are independent and identically distributed (i.i.d.) with mean 0 and variance \sigma^2, the result of the integral is: \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) Here \Delta is the Laplacian operator, i.e., \Delta f = \sum_i \frac{\partial^2}{\partial x_i^2} f. This result is simple in form, equivalent to adding a 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 calculating the second-order derivative of l, and combined with gradient descent, it would require calculating third-order derivatives, which is hard to implement efficiently in existing deep learning frameworks.

Logic Analysis

In mathematical terms, if l is some form of distance metric, then by the triangle inequality: l(f(x+\varepsilon),y) \leq l(f(x+\varepsilon),f(x)) + l(f(x),y) If l is not a metric, a similar result can usually be obtained via Jensen’s inequality. For example, if l(f(x+\varepsilon),y)=\Vert f(x+\varepsilon) - y\Vert^2, then we have: \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} This means that the target in Eq. [eq:loss-2] (multiplied by some constant) can be considered an upper bound of l(f(x+\varepsilon),y). Since the original target is hard to optimize, we optimize its upper bound instead.

Note that among the two terms in target [eq:loss-2], l(f(x+\varepsilon),f(x)) measures the smoothness of the model itself and has nothing to do with labels. It can be optimized using unlabeled data, which means it can be combined with labeled data to form a semi-supervised learning process.

Brave Calculation

For the target in Eq. [eq:loss-2], the integral result is: \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 Following the same path, we approximate the expansion of \varepsilon: \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} Looks terrifying? Don’t worry. Let’s recall that as a loss function, l generally has the following characteristics:

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 essentially says that l is smooth and reaches a local (minimum) value at x=y, and the minimum value is 0. These characteristics are common to almost all loss functions. Based on these, the first three terms of the terrifying Eq. [eq:kongbu] become zero, so the final integral result is: \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)}

Gradient Penalty

It still looks a bit daunting, but it’s much better than Eq. [eq:kongbu]. This is also a regularization term, characterized by containing only first-order gradient terms. For specific loss functions, \left.\frac{\partial^2 l(F(x),f(x))}{\partial F_i(x) \partial F_{i'}(x)}\right|_{F(x)=f(x)} can be calculated 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, so only the i=i' components need to be calculated. Let’s denote it as \lambda_{i}(x), then: \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} As we can see, the form is equivalent to calculating a gradient penalty term \Vert \nabla_x f_i(x)\Vert^2 for each component of f(x) and then summing them weighted by \lambda_i(x).

For example, for MSE, l(f(x),y)=\Vert f(x) - y\Vert^2, we find \lambda_i(x)\equiv 2, so the corresponding regularization term is \sum_i\Vert \nabla_x f_i(x)\Vert^2. For KL divergence, l(f(x),y)=\sum_i y_i \log \frac{y_i}{f_i(x)}, then \lambda_i(x)=\frac{1}{f_i(x)}, and the corresponding regularization term is \sum_i f_i(x) \Vert \nabla_x \log f_i(x)\Vert^2. These results can more or less be found in the famous "Deep Learning" book (the "Flower Book"), and are not new. Similar derivations can also be found in the reference "Training with noise is equivalent to Tikhonov regularization".

Sampling Approximation

Of course, although we can derive a regularization term \sum_i \lambda_i(x)\Vert \nabla_x f_i(x)\Vert^2 that only involves first-order gradients, the computational cost is still not low because it requires calculating the gradient for each f_i(x). If the number of output components is too large, this cost remains unbearable.

In this case, a sampling approximation can be considered: assuming q(\eta) is a distribution with mean 0 and variance 1, we have: \sum_i \Vert \nabla_x f_i(x)\Vert^2=\sum_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] In this way, in each step, we only need to calculate the gradient of \sum_i \eta_i f_i(x), rather than multiple gradients. One of the simplest ways to choose q(\eta) is a uniform distribution over \{-1,1\}, where each \eta_i is chosen from \{-1,1\} with equal probability.

Supervised Adversarial

Supervised adversarial training focuses on the original target [eq:noisy-loss]. The goal is to minimize the loss, so if we want to choose more representative noise, we should choose noise that makes the loss larger. Since: l(f(x + \varepsilon), y) \approx l(f(x), y) + \varepsilon \cdot \nabla_x l(f(x), y) making l(f(x + \varepsilon), y) as large as possible means \varepsilon should be in the same direction as \nabla_x l(f(x), y). In other words, the perturbation should move in the direction of the gradient ascent: \varepsilon \sim \nabla_x l(f(x), y) This constitutes the FGM method in adversarial training, which was previously introduced in "A Brief Talk on Adversarial Training: Significance, Methods, and Reflections (with Keras Implementation)".

It is worth noting that in that previous article, we also derived that adversarial training is, to some extent, equivalent to adding a gradient penalty term \left\Vert\nabla_x l(f(x), y)\right\Vert^2 to the loss, which is similar to the results of the noise integral in the previous section. This suggests that gradient penalty should be one of the universal means to improve model performance.

Virtual Adversarial

As mentioned earlier, the term l(f(x+\varepsilon),f(x)) does not require label signals and can therefore be used for unsupervised learning. From its Gaussian integral expansion, we obtained the gradient penalty [eq:gp]. If we follow the idea of adversarial training, instead of calculating the integral, we look for the perturbation noise that maximizes l(f(x+\varepsilon),f(x)). This constitutes "Virtual Adversarial Training (VAT)," which first appeared in the paper "Virtual Adversarial Training: A Regularization Method for Supervised and Semi-Supervised Learning".

Based on our previous 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 zero. Therefore, to calculate the adversarial perturbation, we must expand it to the second order: \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} Here f_{ng}(x) indicates that we do not need to take the gradient with respect to the x inside it. Thus, we need to solve two problems: 1. How to efficiently calculate 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 is not difficult to prove that the optimal solution for u is actually the "eigenvector corresponding to the largest eigenvalue of \mathcal{H}," also known as the "principal eigenvector of \mathcal{H}." An effective method to approximate the principal 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}. Related derivations can be found in the "Principal Eigenvalue" and "Power Iteration" sections of "Lipschitz Constraints in Deep Learning: Generalization and Generative Models".

In power iteration, we find that we don’t need to know the specific value of \mathcal{H}, only the value of \mathcal{H}u, which can be approximated by finite differences: \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} where \xi is a scalar constant. Based on this approximation, we can obtain the following VAT process:

Initialize vector u\sim \mathcal{N}(0,1), 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 to perform regular gradient descent.

Experiments show that iterating once is usually sufficient. If we iterate zero times, it is equivalent to adding Gaussian noise as mentioned at the beginning of this article. This shows that virtual adversarial training improves the "specificity" of the noise through \nabla_x l(f(x+\xi u), f_{ng}(x)).

Reference Implementation

Regarding the Keras implementation of adversarial training, it was already provided in the previous article. Here, I provide a reference implementation of virtual adversarial training in Keras:

def virtual_adversarial_training(
    model, embedding_name, epsilon=1, xi=10, iters=1
):
    """Add virtual adversarial training to the model.
    model: the keras model to which VAT is added.
    embedding_name: the name of the Embedding layer in the model.
    Should be used after model.compile().
    """
    if model.train_function is None:  # If no train function yet
        model._make_train_function()  # Manually make it
    old_train_function = model.train_function  # Backup old train function

    # Find Embedding layer
    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')

    # Calculate Embedding gradients
    embeddings = embedding_layer.embeddings  # Embedding matrix
    gradients = K.gradients(model.total_loss, [embeddings])  # Embedding gradients
    gradients = K.zeros_like(embeddings) + gradients[0]  # Convert to dense tensor

    # Wrap as functions
    inputs = (
        model._feed_inputs + model._feed_targets + model._feed_sample_weights
    )  # All input layers
    model_outputs = K.function(
        inputs=inputs,
        outputs=model.outputs,
        name='model_outputs',
    )  # Model output function
    embedding_gradients = K.function(
        inputs=inputs,
        outputs=[gradients],
        name='embedding_gradients',
    )  # Model gradient function

    def l2_normalize(x):
        return x / (np.sqrt((x**2).sum()) + 1e-8)

    def train_function(inputs):  # Redefine train function
        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):  # Iteratively find perturbation
            delta2 = xi * l2_normalize(delta2)
            K.set_value(embeddings, K.eval(embeddings) - delta1 + delta2)
            delta1 = delta2
            delta2 = embedding_gradients(inputs)[0]  # Embedding gradient
        delta2 = epsilon * l2_normalize(delta2)
        K.set_value(embeddings, K.eval(embeddings) - delta1 + delta2)
        outputs = old_train_function(inputs)  # Gradient descent
        K.set_value(embeddings, K.eval(embeddings) - delta2)  # Remove perturbation
        return outputs

    model.train_function = train_function  # Override original train function

# After defining the function, enabling VAT takes only one line
virtual_adversarial_training(model_vat, 'Embedding-Token')

For the complete script, please refer to: task_sentiment_virtual_adversarial_training.py. Roughly, the model is built twice: one model is trained normally with labeled data, and the other is trained with unlabeled data using virtual adversarial training, executing alternately. Please understand the source code before using it; do not just copy and paste. The experimental task is sentiment classification with about 20,000 labeled data points. We take the first 200 as labeled samples and the rest as unlabeled data. The performance comparison between VAT and non-VAT is as follows (each experiment was repeated three times, and the average was taken):

Note: As mentioned earlier, f_{ng}(x) means not taking the gradient with respect to x, but the gradient of f’s own parameters \theta still needs to be calculated. However, readers who understand the above code will find that the implementation effectively removes gradients for both x and \theta in f_{ng}(x). Theoretically, this is not perfectly equivalent to standard VAT. The issue is that implementing standard VAT in Keras is somewhat cumbersome and increases computational load. Furthermore, experiments found that this "bootleg" version already brings improvements. Standard VAT differs from it only by second-order small terms, so the difference is not significant, and the code above basically meets the requirements.