Robust Neural Networks inspired by Strong Stability Preserving Runge-Kutta methods

Neural networks such as ResNet [15], PolyNet [40] and recurrent neural networks share a common operation represented as $x_{t+1}=x_{t}+F(x_{t};\Theta_{t})$. Interestingly, a sequence of the operations (or equivalently the network architectures) can be interpreted as an explicit Euler method for numerical discretization [6, 7, 20, 24, 25]. For instance, ResNet can be written mathematically as

where $A$ denotes the number of layers in the network.

If we multiply the function $F$ by $\Delta t$, i.e., $x_{k+1}=x_{k}+\Delta tF(x_{k};\Theta_{k})$, then ResNet can be seen as the explicit Euler numerical scheme discretization with an initial condition, $x(0)$, to solve the initial value problem given as

The explicit Euler method is the simplest Runge-Kutta method and often suffers from low accuracy because it is a first-order method. In this regard, higher-order numerical methods are natural candidates to obtain a more precise numerical solution, but the higher accuracy from higher-order methods may come with the cost of instability, e.g., poor convergence behaviour on stiff differential equations compared to the first order Euler method [4]. Therefore, it is important to understand the trade-off between accuracy and stability when considering a numerical method.

Recently, some network architectures inspired by the computational similarity between ResNet and Euler discretization have been proposed, e.g., NeuralODE and FFJORD [6, 11]. Unlike ResNet, which requires the discretization of observation/emission intervals to be represented by a finite number of hidden layers, NeuralODE and FFJORD use numerical discretization methods in the forward propagation to define continuous-depth and continuous-time latent variable models. These require ODE solvers for training and inference, unlike our implementation of SSP networks. Since we changed only computational graphs and coefficients based on the numerical discretization theory, our methods perform the standard forward/backward propagation in the discrete space as ResNet.

Another approach to design new blocks/layers of neural networks is to make them have operations similar to advanced numerical discretization techniques that possess desirable properties [20, 25]. From the partial differential equation perspective, analysis on numerical stability of conventional residual connections lead to the development of new architectures: parabolic/hyperbolic CNNs to achieve better stability as parabolic/hyperbolic PDEs [25]. The models use theoretical assumptions on the function to achieve stability with a positive semi-definite Jacobian of the function resulting in constraints on convolutional kernels; alternatively, our networks do not require such constraints.

Stability and robustness of neural networks have been studied in the context of adversarial attacks after the success of deep learning [2, 30, 36]. Gradient-based adversarial attacks create adversarial examples solving optimization problems. One example is the maximization of loss against ground truth labels within a small ball, e.g., $\max_{\delta}\mathcal{L}(h_{\theta}(x+\delta),y),\textit{ s.t. }\|\delta\|_{\infty}\leq\epsilon$, where $h_{\theta}$ is a model parameterized by $\theta$, $x$, $y$ are the input (natural sample) and its target label respectively, and $\mathcal{L}$ is a loss function. The simplest procedure to approximate the solution is to use the fast gradient sign method (FGSM) [8]. It can be seen as an optimal solution to a linearized loss function, i.e., $\arg\max_{\|v\|_{\infty}\leq\alpha}v^{T}\nabla_{\delta}\mathcal{L}(h_{\theta}(x+\delta),y)=\alpha\cdot\text{sign}(\nabla_{\delta}\mathcal{L}(h_{\theta}(x+\delta),y))$. Furthermore, the FGSM can be more powerful when it is used with iterative methods such as the projected gradient descent (PGD). PGD has been used in both untargeted and targeted attacks [21, 5].

One of the early attempts to defend against adversarial attacks is adversarial training using FGSM, a single-step method [8]. After that, various defensive techniques have been proposed [3, 22, 26, 29, 31]. Many of them were defeated by iterative attack methods [5] and Backward Pass Differentiable Approximation [1]. Adversarial training with stronger multi-step attack methods is still promising and shows state-of-the-art performance [21, 32, 35]. More recently, provably robust neural networks have been successfully trained by minimizing the lower bound of risk based on convex duality and convex relaxation [33, 34]. Most adversarial training methods above assume that attack methods are known a priori, i.e., a white-box attack, and generate augmented samples using the attacks. Another defensive technique is to alleviate the effect of perturbation by augmentation and reconstruction [23, 37], or denoising [35]. These methods alongside adversarial training achieved comparable robustness. Similarly, in this work we will introduce our approach and evaluate it with adversarial training.

Our objective is to solve the non-autonomous differential equation given as

where $t_{0},t_{N}$ are the initial and terminal time state respectively; a non-autonomous system permits a time varying solution, e.g., the learned function varies as the depth of the network increases. The function $L$ is a linear (or nonlinear) function and $u(t_{0})$ is given by the initial condition. The objective is to figure out the terminal state of the function $u$, i.e., $u(t_{N})$.

A general high-order Runge-Kutta time discretization for solving the initial value problem (3) introduced in [28] is given as

where $\sum^{i-1}_{k=0}\alpha_{i,k}=1$ and $\alpha_{i,k}\geq 0$. For example, if $m=1$, it becomes the first-order Euler method as in Equation (1) with $\alpha_{1,0}=\beta_{1,0}=1$.

Shu et al. [27, 28] propose a TVD time discretization method that is called the SSP time discretization method; for more discussion on the TVD method, we refer the reader to [12, 13]. The procedure of TVD time discretization is to take the high-order method to decrease the local truncation error and maintain the stability under a suitable restriction on the time step. While applying the TVD scheme into the explicit high-order Runge-Kutta methods, there needs the assumption to hold it: The first-order Euler method in time is strongly stable under a certain (semi) norm when the time step $\Delta t$ is suitably restricted [10]. More precisely, if we assume that the forward Euler time discretization is stable under a certain norm, the SSP methods find a higher-order time discretization that maintains strong stability for the same norm; improving accuracy.

Followed by this assumption, for a sufficiently small time step known as Courant-Friedrichs-Lewy (CFL) condition $\Delta t\leq\Delta t_{CFL}$, the total variation semi-norm of the numerical scheme does not increase in time, that is,

where the total variation is defined by

where $j$ is the spatial discretization. The explicit high-order Runge-Kutta discretization with the SSP property maintains a higher order accuracy with a modified CFL condition $\Delta t\leq c\Delta t_{CFL}$. In other words, the high-order SSP Runge-Kutta scheme improves accuracy while retaining its stability. This has been theoretically studied by the following Lemma 1.

We provide a sketch of the proof of Lemma 1 in the supplement. The full proof of the Lemma 1 can be found in [28]. Following this representation, we can figure out the specific coefficients $\alpha_{i,k}$ and $\beta_{i,k}$ in equation (4). In particular, the second and third order nonlinear SSP Runge-Kutta method was studied in [28].

A sketch of the proof for Lemma 2 can be found in the supplement and for the detailed proof, we refer the reader to [9, 10, 28].

Next, we show how to incorporate the explicit SSP Runge-Kutta method into neural networks. Equation (19) and (20) can be implemented with standard residual blocks and simple operations, as shown in Figure 2.

Let ResBlock denote a standard residual block written as $\textit{ResBlock}(x(t_{k});\Theta(t_{k}))$ $=x(t_{k})+F(x(t_{k});\Theta(t_{k})),$ where $\Theta(t_{k})$ are the parameters of ResBlock($\cdot$;$\Theta(t_{k})$). The function $F$ is typically composed of two or three sets of normalization, activation, and convolutional layers, e.g., Figure 2 and [15, 16]. When the numbers of input and output channels differ, we use the expansive residual block ResBlock-E; this can be implemented with a $1\times 1$ convolutional filter to expand the number of channels.

Using the standard modules in ResNet (ResBlock and ResBlock-E), SSPNets can be constructed. First, SSP blocks can be implemented using linear combinations of ResBlocks. As the Euler method interpretation of ResNet requires $\Delta t=1$, we assume $\Delta t=1$ in Equation (19), then the SSP2-block is given by,

Similarly, the third order SSP in Equation (20) (SSP3-block) is written as

The SSP block schematic is presented in Figure 2 and SSP blocks are only used when the number of channels does not change.

The explicit SSP Runge-Kutta methods in Equation (19) and (20) use the same function $L$ multiple times. Similarly, SSP blocks in Equation (9) and (10) apply the same ResBlock multiple times. Using the same ResBlock multiple times can be viewed as parameter sharing, which is a kind of regularization. In other words, without increasing the number of parameters, a SSP block implementation improves the robustness of neural networks by utilizing higher-order schemes.

The variance of SSP blocks is smaller than that of ResBlock. The variance adds to our argument that the SSP property is the reason for improved robustness; for more detailed derivation and proof, see the supplement.

Also, we generalize Equation (4) with the second-order Adaptive Runge-Kutta block (ArkBlock) that has the SSP property by construction. These novel computational blocks slightly increase the number of parameters compared to ResBlock but also provide greater robustness and natural accuracy than SSP2-Block or SSP3-Block. Finally, we explore different computational architectures within each group to retain natural accuracy and further improve robustness.

A naive implementation of Equation (4) yields 5 additional parameters. We can retain the SSP property in ArkBlocks by reducing the number of parameters with Ralston’s method [9]. Thus, the number of additional learned parameters per block, when compared with ResBlock, is 2 and is defined as

We further improve performance by reducing the number of parameters by fixing $\alpha_{2,1}$ and simply learning $\beta_{1,0}$ in each block.

Adaptive SSP networks still maintain the same architecture, as in Figure 3, but are comprised of blocks that have the form

We implement ArkBlocks with,

The ArkBlocks are inspired by the generalized Runge-Kutta method in (14). However, the numerical scheme in Equation (14), keeps $\alpha_{2,1}$ and $\beta_{1,0}$ constant in all blocks, while ArkBlocks set those parameters as learnable; varying in each block. Such an adaptivity based on data and architectures cannot be obtained by mathematically derived coefficients. To our knowledge, this is the first attempt.

We demonstrate that SSPNets are more robust than ResNet with standard training. Since MNIST has relatively low-resolution images compared to CIFAR10, we used a smaller architecture by skipping group 1 and 2 in Figure 3.

We analyze the robustness of SSP networks, on the CIFAR10 dataset. Our preliminary experiments show that all the models, e.g., ResNet, SSP-2, SSP-3, and SSP-adap trained without adversarial training are easily fooled by PGD attacks, but more analysis is needed on a more challenging dataset. For this reason, we focus on the adversarial training setting for CIFAR10. Please see supplementary materials for more analysis on SSP networks with adversarial training.

To summarize, our adversarial training procedure works as follows: First, randomly perturb the image within the allowed perturbation range $\epsilon$. Next, generate the candidate adversarial example by PGD attack. Finally, take the gradient descent step on a minibatch composed of only candidate adversarial examples. The adversarial training is closely related to the Frank-Wolfe Algorithm and two projections in the original adversarial training can be simplified to one projection to the intersection of two convex sets. The pseudocode of adversarial training and a detailed discussion of implementation are provided in the supplement.

Before discussing about the effectiveness of SSP, we briefly show the relationship among robustness, the amount of model parameters, and the strength of attacks used in adversarial training. As shown in Table 2, the robustness of all the models is improved by stronger attacks during training (e.g., larger $K$ in PGD). The same observation is reported in [32]. For instance, SSP-3 (N=6, K=12) shows higher accuracy than SSP-3 (N=6, K=7) against all the attacks and especially the improvement is about 7% against PGD with $20$ iterations. Also, a bigger model size (e.g., larger $N$) increases the robustness against adversarial examples. This is closely related to the finding in [21] that increasing the number of channels in hidden layers often improves the robustness. Our experiments show that increasing the model size by adding more layers improves the robustness. For example, when $K=7$, SSP-3 with $N=10$ blocks show overall higher accuracy than SSP-3 with $N=6$ blocks, the gain is about $2\%$. From the numerical discretization perspective, more blocks can be seen as a finer time discretization that leads to a more accurate numerical solution (or prediction).

All SSPNets, SSP-2, SSP-3 and SSP-adap, consistently outperform ResNet by, roughly, $1\sim 3.9\%$ when ResNet and SSPNets have the same number of blocks, $N$, and iterations, $K$, in adversarial training. Note that we compare SSP-2 (and SSP-3) with ResNet, which has the same amount of parameters, and this is important to assure that the gain is not from an increased the amount of model parameters. Also, SSP-2, SSP-3 and SSP-adap have the same time discretization as ResNet. So, we conclude that the improvement in robustness against adversarial attacks solely comes from the strength of a higher-order numerical discretization. Table 2 shows one more interesting property of SSP networks. Unlike adversarial training and defensive methods that usually cause the label leaking effects [31, 38], SSP-2, SSP-3 and SSP-adap (our architectural changes) do not bring any additional loss of accuracy on natural samples.

On the other hand, Table 2 also shows that the mid-RK2 architecture does not outperform ResNet, SSP-2, SSP-3 or SSP-adap even though the mid-RK2 is derived from the second order numerical scheme. This gives credence to the implementation of SSPNets and implies that the robust performance is not a result of arbitrary high-order methods. In addition, SSP-adap achieves comparable natural accuracy as ResNet and improves robustness. Table 2 demonstrates the consistent improvement across various settings. For example, SSP-adap achieves nearly 4% absolute performance improvement for $N=10,K=7$. The improvement by SSP networks compared to ResNet and the performance difference between different SSP networks are relatively smaller than Table 1. Our conjecture is that this is due to the improvement by adversarial training. We believe that the Strong Stability Preserving property imposed by our architectural change allows the SSPNets to improve the robustness against adversarial attacks.

Figure 4 presents the perturbation growth ratio when $N=6,K=7$ and $N=10,K=12$ at each group in the models. Since the adversarial examples change the final predictions, the perturbation growth ratio increases in all the models. However, for SSPNets, the perturbation growth ratio is significantly lower than ResNet. This result supports that the proposed SSP blocks improve robustness of networks against adversarial attacks when compared to ResNet. We also conducted an experiment when $x^{\prime}$ is corrupted by adding a random perturbation to $x$ and the result is consistent with Figure 4. For full version of Figure 4 and more discussion, see the supplement.