Improved Spectral Norm Regularization for Neural Networks

Abstract

We improve on a line of research that seeks to regularize the spectral norm of the Jacobian of the input-output mapping for deep neural networks. While previous work rely on upper bounding techniques, we propose a scheme that targets the exact spectral norm. We evaluate this regularization method empirically with respect to its generalization performance and robustness.

Our results demonstrate that this improved spectral regularization scheme outperforms L2-regularization as well as the previously used upper bounding technique. Moreover, our results suggest that exact spectral norm regularization and exact Frobenius norm regularization have comparable performance. We analyze these empirical findings in the light of the mathematical relations that hold between the spectral and the Frobenius norms. Lastly, in light of our evaluation we revisit an argument concerning the strong adversarial protection that Jacobian regularization provides and show that it can be misleading.

In summary, we propose a new regularization method and contribute to the practical and theoretical understanding of when one regularization method should be preferred over another.

Appendix

1.1 A Experimental details

Network Architectures. We will follow [12] and denote a convolutional-max-pool layer as a tuple (K, $C_{in}\rightarrow C_{out}, S, P, M$) where K is the width of the kernel, $C_{in}$ is the number of in-channels, $C_{out}$ the number of out-channels, S the stride, P the padding of the layer and M the size of the kernel of the max-pool following the convolutional layer. The case $M=1$ can be seen as a convolutional layer followed by an identity function. Linear layers we will denote as the tuple ($N_{in}$, $N_{out}$) where $N_{in}$ is the dimension of the input and $N_{out}$ the size of the output. For KMNIST and FashionMNIST we used the LeNet network which consist of a convolutional-maxpool layer (5, $1 \rightarrow 6$, 1, 2, 2), convolutional-maxpool layer (5, $6 \rightarrow 16$, 1, 0, 2), linear layer (400, 120), linear layer (120, 84) and linear layer (84, 10).

We use the VGG16 network as is available from the torchvision package. For this network we use batch-norm layers directly after every convolutional layers. This network consist of the layers (3, $3 \rightarrow 64$, 1, 1, 1), (3, $64 \rightarrow 64$, 1, 1, 2), (3, $64 \rightarrow 128$, 1, 1, 1), (3, $128 \rightarrow 128$, 1, 1, 1), (3, $128 \rightarrow 256$, 1, 1, 2), (3, $256 \rightarrow 256$, 1, 1, 1), (3, $256 \rightarrow 256$, 1, 1, 1), (3, $256 \rightarrow 512$, 1, 1, 2), (3, $512 \rightarrow 512$, 1, 1, 1), (3, $512 \rightarrow 512$, 1, 1, 1), (3, $512 \rightarrow 512$, 1, 1, 2), (512, 10).

Training Details. We train the LeNet networks for 50 epochs with SGD (with momentum=0.8). For every regularization method we perform a hyperparameter search over these three following parameters and values.

• Learning rate: [0.01, 0.001]

• Batch size: [16, 32]

• Weight factor $\lambda$: [0.0001, 0.001, 0.01, 0.1]

For the VGG16 network we trained the network for 100 epochs with a batch size of 128, SGD with momentum of 0.8 and performed a hyperparameter search over these parameters and values

• Weight factor $\lambda$: [0.00001, 0.0001, 0.001, 0.01, 0.1]

For VGG16 we additionally used a cosine annealing learning rate scheduler with an initial learning rate of 0.1 and the data augmentation techniques of random cropping and horizontal flipping.

For each hyperparameter setting we repeat the training procedure 5 times to be able to obtain mean and standard deviation. We pick the final representative model for each regularization method as the one that achieves the lowest mean validation loss over these 5 training runs.

For the Frobenius regularization we set $n_{proj} = 1$ and for the Spectral-Bound we estimate the spectral norm of the weight matrices through one power iteration.

Details for Figures. Figure 4: The model for each regularization method was chosen randomly among the 5 models from the hyperparameter setting that obtained the best results in Table 1. The distance is only calculated for the points in the validation set that all models predict correctly. In total the distance is predicted for between 8000–9000 validation points on FashionMNIST and KMNIST.

Figure 5 (left): The time for a batch was measured on a computer with NVIDIA K80 GPU as available through Google Colab³. The analytical method works by sequentially calculating $d(x^L \cdot e_i)/dx$ where $e_i$ is a basis-vector for $\mathbb {R}^{n_{out}}$ for $i=1,2,...,n_{out}$. This yields the full Jacobian matrix which we then calculate the singular values of by using inbuilt functions in PyTorch.

Figure 5 (right): The upper bound was evaluated on a network trained with the Spectral-Bound regularization scheme for all data points in the training set. The curve for the spectral method was evaluated on a network trained with the spectral method for all data points in the training set. For the spectral method there was no significant difference in the shape of the curve when using a different network or by working with data points in the validation set.

1.2 Conversion Between Operators

In this section we detail how to convert between the forward F, backward $F^T$ and regular operators G. These can be seen in Table 3 - 5. Other non-linearities such as Dropout can be incorporated identically to ReLU by simply storing the active neurons in a boolean matrix Z.

Skip-Connections. Utilizing networks with skip-connections does not change the forward and backward modes. Simple turn off the bias of all layer transformations and replace the activation functions with the matrices $Z_R^i$ instead. That this is true follows from the definition of a network with skip-connections. For simplicity of presentation, we will assume that the skip-connections only skip one layer. Assume that we have a network with L layers and additionally have skip-connections between layers with indices in the set $\mathcal {S} := \{s_1,s_2,...,s_m\},~1\le s_i \le L$. Then the network $f_{\theta }$ is given recursively as before with

$$ x^l = {\left\{ \begin{array}{ll} f^l(G^l(x^{l-1}) + b^l) &{}\text{ if } l \in \mathcal {S}^C, \\ x^{l-1} + f^l(G^l(x^{l-1}) + b^l) &{} \text{ if } l \in \mathcal {S}. \end{array}\right. } $$

Assuming that we are only working with piecewise linear or linear operators $G^l$, then for $x \in R$ we know that each operator can be represented as a matrix and we can write the derivative of the two cases as

$$ \frac{dx^l}{dx^{l-1}} = {\left\{ \begin{array}{ll} Z^lW^l &{}\text{ if } l \in \mathcal {S}^C, \\ I + Z^lW^l &{} \text{ if } l \in \mathcal {S}, \end{array}\right. } $$

where I denotes a unit-matrix. The Jacobian-vector product $W_Rv$ can thus be obtained as

$$\begin{aligned} W_Rv = \bigg (\prod _{l=1}^L (I - \mathbb {I}\{l \in \mathcal {S}^C\} + Z^lW^l)\bigg )v \end{aligned}$$

(13)

where $\mathbb {I}\{l \in \mathcal {S}^C\}$ is an indicator for the unit-matrix so that we can concisely write the two cases. Thus we see that we can interpret this equation in the same manner as we did for the networks without skip-connections. We simply pass the input v through the network and turn off all the biases and replace the activation functions with $Z^l$. The same is true for the backward mode (Table 4).

Table 3. Conversion table for the linear operator.

Table 4. Conversion table for the convolutional operator.

Table 5. Conversion table for the max-pool operator.

1.3 Time Efficiency and Relative Error

In this section we investigate the difference between targeting the exact spectral norm of the Jacobian compared to working with an upper bound. From Table 1 we saw that this yields an improved generalization performance and from Fig. 3 we observed that the two methods provide a similar protection against noise, with different strengths against different attacks on the two considered data sets.

While an improved generalization performance is beneficial, it cannot come at a too large of a computational cost. Additionally, with approximate methods it is also important to measure the trade-off between computational speed and accuracy of the approximated quantity. We thus analyze the computational overhead that they add to the training routine and the relative error with the analytical spectral norm.

Fig. 5.

Time and error comparison between the Spectral and Spectral-Bound method for the LeNet network. (Left) Time taken to pass over one batch of data points. The Spectral method is slower than the Spectral-Bound method for larger batch sizes but still around two orders of magnitude faster than calculating the exact spectral norm analytically. (Right) The relative error as the number of power iterations is increased. The relative error decreases quickly and is significantly closer to the exact quantity compared to the upper bound $\prod _l ||W^l||_2$.

In Fig. 5 (left) we can thus see the average time taken to optimize over a batch for the Spectral method, the Spectral-Bound method, an analytical method that calculates $||W_R||_2$ exactly and a regular forward-pass. In Fig. 5 (right) the relative error for the power iteration scheme is visible.

From these plots we can see there is a small extra incurred cost of working with our method compared to regularizing with Spectral-Bound, but that our method has a significantly lower relative error while still being orders of magnitude faster than calculating the analytical spectral norm.

1.4 Proof for Extension Scheme

We will denote the directed acyclic graph which when summing the product of every edge element along every path from output to input yields $(df/dx)^T$ as G.

Theorem: Consider the graph F obtained by flipping the direction of all edges of G and adding a node at the end of F with edge elements given by components of v. Summing the product of every edge element along every path from output to input of F yields (df/dx)v.

Proof: We will follow the notation of [6], Theorem 1 and denote the Jacobian between variables $y=f_{\theta }(x)$ and x as the sum of the product of all intermediate Jacobians, meaning

$$\begin{aligned} \frac{dy}{dx} = \sum _{p \in \mathcal {P}(x,y)} \prod _{(a,b) \in p} J^{a\rightarrow b}(\alpha ^b) \end{aligned}$$

(14)

where $\mathcal {P}(x,y)$ is the set of all directed paths between x and y and (a, b) is two successive edges on a given path.

In our scheme we flip the direction of all relevant edges and add a fictitious node at the end of the path the flipped paths. Since we preserve the edge elements, we can realize that flipping the direction of the edges simply transposes the local Jacobian, meaning that $J^{b\rightarrow a}(\alpha ^b) = \big (J^{a\rightarrow b}(\alpha ^b)\big )^T$ with our scheme. Further, our added fictitious node has edge elements given by elements of v, and the Jacobian between that node and the subsequent layer is thus given by $v^T$. For a path $p=[(v_1,v_2), (v_2,v_3),...,(v_{n-1}, v_n)]$ we define the flipped path with the added fictitious node as $p^T$ as $p^T = [(v_n, v_{n-1}), ..., (v_2,v_1), (v_1, v_f)]$ and the reverse-order path $\lnot p$ as $\lnot p = [( v_{n-1}, v_n), ..., (v_2,v_3),(v_1,v_2)]$. For our modified graph we thus have the Jacobian for a path as

$$\begin{aligned} \prod _{(a,b) \in p^T} J^{a\rightarrow b}(\alpha ^b)&= \prod _{(a,b) \in p^T} \big (J^{b\rightarrow a}(\alpha ^b)\big )^T \end{aligned}$$

(15)

$$\begin{aligned}&=v^T\bigg (\prod _{(a,b) \in p} J^{a\rightarrow b}(\alpha ^b)\bigg )^T\end{aligned}$$

(16)

$$\begin{aligned}&=v^T\bigg (\prod _{(a,b) \in \lnot p} J^{a\rightarrow b}(\alpha ^b)^T\bigg ) \end{aligned}$$

(17)

Denoting the fictitious node as $n_f$ and summing over all paths we thus get

$$\begin{aligned}&\sum _{p^T \in \mathcal {P}(y,n_f)} \prod _{(a,b) \in p^t} J^{a\rightarrow b}(\alpha ^b) \end{aligned}$$

(18)

$$\begin{aligned}&= \sum _{p^T \in \mathcal {P}(y,n_f)}v^T\bigg (\prod _{(a,b) \in \lnot p} J^{a\rightarrow b}(\alpha ^b)^T\bigg )\end{aligned}$$

(19)

$$\begin{aligned}&=v^T\sum _{p^T \in \mathcal {P}(y,n_f)}\bigg (\prod _{(a,b) \in \lnot p} J^{a\rightarrow b}(\alpha ^b)^T\bigg )\end{aligned}$$

(20)

$$\begin{aligned}&=v^T\big (\frac{dy}{dx}\big )^T = (\frac{dy}{dx}v)^T \end{aligned}$$

(21)

which proves that working with the modified graph will yield the desired matrix-vector product $\frac{dy}{dx}v$ $\square$.