An axiomatized PDE model of deep neural networks
Inspired by the relation between deep neural network (DNN) and partial
differential equations (PDEs), we study the general form of the PDE models of
deep neural networks. To achieve this goal, we formulate DNN as an evolution
operator from a simple base model. Based on several reasonable assumptions, we
prove that the evolution operator is actually determined by
convection-diffusion equation. This convection-diffusion equation model gives
mathematical explanation for several effective networks. Moreover, we show that
the convection-diffusion model improves the robustness and reduces the
Rademacher complexity. Based on the convection-diffusion equation, we design a
new training method for ResNets. Experiments validate the performance of the
proposed method.
Tangjun Wang, Wenqi Tao, Chenglong Bao, Zuoqiang Shi
# An Axiomatized PDE Model of Deep Neural Networks +

###### Abstract

Inspired by the relation between deep neural network (DNN) and partial differential equations (PDEs), we study the general form of the PDE models of deep neural networks. To achieve this goal, we formulate DNN as an evolution operator from a simple base model. Based on several reasonable assumptions, we prove that the evolution operator is actually determined by convection-diffusion equation. This convection-diffusion equation model gives mathematical explanation for several effective networks. Moreover, we show that the convection-diffusion model improves the robustness and reduces the Rademacher complexity. Based on the convection-diffusion equation, we design a new training method for ResNets. Experiments validate the performance of the proposed method.

R esidual network, axiomatization, convection-diffusion equation

35K57, 93B35

## 1 Introduction

Deep neural networks (DNN) have achieved success in tasks such as image classification [33], speech recognition [7], video analysis [3], and action recognition [39]. Among these networks, residual networks (ResNets) are important architectures, making it practical to train ultra-deep DNN, and has ability to avoid gradient vanishing [12, 13]. Also, the idea of ResNets has motivated the development of many other DNNs including WideResNet [43], ResNeXt [41], and DenseNet [15].

In recent years, understanding the ResNets from the dynamical perspective has become a promising approach [8, 11]. More specifically, assume \(\mathbf{x}_{0}\in\mathbb{R}^{d}\) as the input of ResNet [12] and define \(\mathcal{F}\) to be a mapping, then the \(l\)-th residual block can be realized by

\[\mathbf{x}_{l+1}=\mathbf{x}_{l}+\mathcal{F}(\mathbf{x}_{l},\mathbf{w}_{l}) \tag{1}\]

where \(\mathbf{x}_{l}\) and \(\mathbf{x}_{l+1}\) are the input and output tensors of the residual mapping, and \(\mathbf{w}_{l}\) are parameters of \(l\)-th layer that are learned by minimizing the training error. Define \(\mathbf{x}_{L}\) as the output of a ResNet with \(L\) layers, then the classification score is determined by \(\mathbf{y}=\text{softmax}(\mathbf{w}_{\text{fc}}\mathbf{x}_{L})\), where \(\mathbf{w}_{\text{fc}}\) are also learnable parameters of the final linear layer.

For any \(T>0\), introducing a temporal partition \(\Delta t=T/L\), the ResNet represented by (1) is the explicit Euler discretization with time step \(\Delta t\) of the following differential equation:

\[\frac{\mathrm{d}\mathbf{x}(t)}{\mathrm{d}t}=v(\mathbf{x}(t),t),\quad\mathbf{x}(0)=\mathbf{x}_ {0}\quad t\in[0,T], \tag{2}\]

where \(v(\mathbf{x}(t),t)\) is a velocity field such that \(\Delta tv(\mathbf{x}(t),t)=F(\mathbf{x}(t),\mathbf{w}(t))\). The above ordinary differential equation (ODE) interpretation of ResNet provides new perspective and has inspired many networks. As shown in [25], by applying stable and different numerical methods for (2), it leads to PolyNet [45] and FractalNet [21]. Besides,
the other direction is to consider the the continuous form of (1.2), representing the velocity \(v(\mathbf{x},t)\) by a deep neural network. One typical method is the Neural ODE [4] in which \(v\) is updated by the adjoint state method. Similar extensions along this direction include neural stochastic differential equation (SDE) [17] and neural jump SDE [18]. The theoretical property of these continuous models has been analyzed in [37, 2, 44].

The connection between ODE and partial differential equation (PDE) through the well-known characteristics method has motivated the analysis of ResNet from PDE perspective, including theoretical analysis [34], novel training algorithms [36] and improvement of adversarial robustness [38] for DNNs. To be specific, from the PDE theory, the ODE (1.2) is the characteristic curve of the transport equation:

\[\frac{\partial u}{\partial t}(\mathbf{x},t)=-v(\mathbf{x},t)\nabla u(\mathbf{x},t),\ (\mathbf{x},t)\in\mathbb{R}^{d}\times[0,T]. \tag{1.3}\]

The method of characteristics tells us that, along the curve \((\mathbf{x},t)\) defined by (1.2), the function value \(u(\mathbf{x},t)\) remains unchanged. Assume at \(t=T\), \(u(\mathbf{x},T)=f(\mathbf{x})\coloneqq\mathrm{softmax}(\mathbf{w}_{\mathrm{fc}}\mathbf{x})\) is the linear classifier, then

\[u(\mathbf{x}(0),0)=u(\mathbf{x}(T),T)=f(\mathbf{x}(T))=f\circ k(\mathbf{x}(0))\]

where \(k\) represents the mapping from \(\mathbf{x}(0)\) to \(\mathbf{x}(T)\), which is the continuous form of feature extraction in ResNet. Thus at \(t=0\), \(u(\cdot,0)\) is the composition of a feature extractor and a classifier, which is analogous to ResNet. Nonetheless, since the transport equation (1.3) is reversible in time, and initial value problem is more common than terminal value problem in PDE, we assume \(u(\mathbf{x},0)=f(\mathbf{x})\) in our paper. Consequently, the direction of solving ODE (1.2) needs to be reversed, but its connection to ResNet remains consistent. In one word, the transport equation (1.3) can describe the evolution from a linear classifier to ResNet.

Suppose we fix the initial condition \(u(\cdot,0)\) as the linear classifier \(f\). DNN can be seen as a map between two functions: \(u(\cdot,0)\) and \(u(\cdot,T)\). In ResNets, this map is

Figure 1: \(\mathcal{T}_{T}\) is the map between two colored boxes. The left represents the base linear classifier. Right represents feature extractor + linear classifier, which together form a typical neural network.
formulated as a convection equation. A natural question is: _Is convection equation is the only PDE to formulate this map? If not, can we derive a general form of PDE to formulate the map?_ In this paper, we try to answer above questions from mathematical point of view. First we construct a continuous flow \(\mathcal{T}_{t}\) which maps a simple linear classifier to a more complicate function, as illustrated in Figure 1,

\[\mathcal{T}_{t}:f\mapsto u(\cdot,t),\quad t\in[0,T], \tag{4}\]

The idea in this paper is also classical in mathematics. First, we extract some _basic properties_\(\mathcal{T}_{t}\) should satisfies. Then based on these basis properties, a general form of \(\mathcal{T}_{t}\) can be derived rigorously. More specifically, inspired by the scale space theory, we prove that under several reasonable assumptions on \(\mathcal{T}_{t}\), \(u(\mathbf{x},t)=\mathcal{T}_{t}f(\mathbf{x})\) is the solution of a second order convection-diffusion equation. This theoretical result provides a unified framework which covers transport equation and some existing works including Gaussian noise injection [38, 24], dropout techniques [36, 35] and randomized smoothing [5, 23, 31]. It also illuminates new thinking for designing networks. In summary, we list the main contributions as follows.

* We establish several basic assumptions on operator \(\mathcal{T}_{t}\), and prove the sufficiency of these assumptions for generalizing ResNet and beyond. To the best of our knowledge, this is the first theoretical attempt for establishing a sufficient condition for designing the variants of ResNet from the PDE perspective, which may provide some insights when considering the search space in neural architecture search (NAS).
* Inspired by our theoretical analysis, we propose an isotropic model by adding isotropic diffusion to (3). Compared to the linear classifier \(f\), we prove that the proposed model has lower Rademacher complexity and larger region with certified robustness. Moreover, we design a training method by applying the operator splitting scheme for solving the proposed convection-diffusion equation.

_Notations._ We denote scalars, vectors, and matrices by lowercase and uppercase letters, where vectors and matrices are bolded. We denote the \(\ell_{2}\) and \(\ell_{\infty}\) norms of the vector \(\mathbf{x}=(x_{1},\cdots,x_{d})\in\mathbb{R}^{d}\) by \(\|\mathbf{x}\|_{2}=(\sum_{i=1}^{d}|x_{i}|^{2})^{1/2}\) and \(\|\mathbf{x}\|_{\infty}=\max_{i=1}^{d}|x_{i}|\), respectively. We denote the gradient and Laplace operators by \(\nabla=(\frac{\partial}{\partial x_{1}},\cdots,\frac{\partial}{\partial x_{d}})\) and \(\Delta=\sum_{i=1}^{d}\frac{\partial^{2}}{\partial x_{i}^{2}}\), respectively. For a function \(f:\mathbb{R}^{d}\to\mathbb{R}\), \(D^{\alpha}f\) denotes its \(\alpha\)-order derivative, \(\|f(\mathbf{x})\|_{L^{\infty}}=\sup_{\mathbf{x}\in\mathbb{R}^{d}}|f(\mathbf{x})|\) its \(L^{\infty}\) norm. \(\mathcal{N}(\mathbf{0},\sigma^{2}\mathbf{I})\) denotes Gaussian noise with mean \(\mathbf{0}\) and variance \(\sigma^{2}\). \(C_{b}^{\infty}\) is the space of bounded functions which have bounded derivatives at any order.

## 2 General PDE model

In this section, we show under several reasonable assumptions, the sequence of operator images \(u(\mathbf{x},t)=\mathcal{T}_{t}f(\mathbf{x})\) is the solution of the convection-diffusion equation. Then we show that our convection-diffusion equation model can naturally cover various existing effective models.

### The characterization of \(\mathcal{T}_{t}\)

Throughout this section, we assume \(\mathcal{T}_{t}\) is well defined on \(C_{b}^{\infty}\), and \(\mathcal{T}_{t}f\) is a bounded continuous function. The assumption is reasonable, since typical classifiers \(f\) like a linear classifier is indeed bounded (between 0 and 1) and has bounded derivatives. The operator image \(\mathcal{T}_{t}f\), which we hope to be a neural network, is obviously bounded and continuous. To get the expression of the evolution operator \(\mathcal{T}_{t}\), we assume it has some fundamental properties, which fall into two categories: deep neural network type and partial differential equation type.
#### 2.1.1 DNN-type assumptions

Suppose we are given two classifiers \(f\) and \(g\) such that \(f(\mathbf{x})\geq g(\mathbf{x})\) for all data point \(\mathbf{x}\in\mathbb{R}^{d}\). Then \(f\circ k(\mathbf{x})\geq g\circ k(\mathbf{x})\) if we replace the data points with extracted features. Recall that for ResNet, \(f\circ k=\mathcal{T}_{T}(f)\), which implies \(\mathcal{T}_{T}\left(f\right)\geq\mathcal{T}_{T}\left(g\right)\). Since the order-preserving property holds both at initial time step \(t=0\) and final time step \(t=T\), it is reasonable to make the following assumption **[Comparison Principle]** For all \(t\geq 0\) and \(f,g\in C_{b}^{\infty}\), if \(f\leq g\), then \(\mathcal{T}_{t}(f)\leq\mathcal{T}_{t}(g)\).

The prediction of a deep neural network is computed using forward propagation, i.e. the network uses output of former layer as input of current layer. Thus, for a DNN model, it's natural that the output of a DNN can be deduced from the output of intermediate \(l\)-th layer without any information depending upon the original data point \(\mathbf{x}\) and output of \(m\)-th layer (\(m<l\)). Regarding the evolution of operator \(\mathcal{T}_{t}\) as stacking layers in the neural network, we should require that \(\mathcal{T}_{t+s}\) can be computed from \(\mathcal{T}_{t}\) for any \(s\geq 0\), and \(T_{0}\) is of course the identity, which implies **[Markov Property]**\(\mathcal{T}_{t+s}=T_{t}\circ T_{t+s,t}\), for all \(s,t\geq 0\) and \(t+s\leq T\). \(\mathcal{T}_{t+s,t}\) denotes the flow from time \(t\) to time \(t+s\).

Linearity is also an intrinsic property of deep neural networks. Notice that we are not referring to a single DNN's output v.s. input linearity, which is obviously wrong because of the activation function. Rather, we are stating that two different DNN with the same feature extractor can be merged in to a new DNN with a new classifier composed with the shared extractor, i.e.

\[(\beta_{1}f+\beta_{2}g)\circ k=\beta_{1}f\circ k+\beta_{2}g\circ k\]

This is linearity at \(t=T\), and for \(t=0\) it is trivial. Thus we assume,

**[Linearity]** For any \(f,g\in C_{b}^{\infty}\), and real constants \(\beta_{1},\beta_{2}\), we have

\[\mathcal{T}_{t}(\beta_{1}f+\beta_{2}g)=\beta_{1}\mathcal{T}_{t}(f)+\beta_{2} \mathcal{T}_{t}(g)\]

if \(C\) is a constant function, then \(\mathcal{T}_{t}(C)=C\).

#### 2.1.2 PDE-type assumptions

First of all, we need an assumption to ensure the existence of a differential equation. If two classifiers \(f\) and \(g\) have the same derivatives of any order at some point, then we should assume same evolution at this point when \(t\) is small. If we unrigorously define \(\partial T_{t}(f)/\partial t=\left(\mathcal{T}_{t}(f)-f\right)/t\) when \(t\to 0^{+}\) (or infinitesimal generator in our proof), then \(\partial T_{t}(f)/\partial t\) should equal to \(\partial T_{t}(g)/\partial t\). Thus, we give the following assumption concerning the local character of the operator \(\mathcal{T}_{t}\) for \(t\) small.

**[Locality]** For all fixed \(\mathbf{x}\), if \(f,g\in C_{b}^{\infty}\) satisfy \(D^{\alpha}f(\mathbf{x})=D^{\alpha}g(\mathbf{x})\) for all \(|\alpha|\geq 0\), then

\[\lim_{t\to 0^{+}}\frac{(\mathcal{T}_{t}(f)-\mathcal{T}_{t}(g))(\mathbf{x})}{t}=0\]

Regularity is an essential component in PDE theory. Thus, when considering PDE-type assumptions on \(\mathcal{T}_{t}\), it is necessary to study its regularity. We separate the regularity requirements into spatial and temporal. First, spatial regularity means that if we add a perturbation \(\mathbf{h}\) to data point \(\mathbf{x}\), the output \(\mathcal{T}_{t}(f)(\mathbf{x}+\mathbf{h})\) will not be much different from adding the same perturbation to the output \(\mathcal{T}_{t}(f)(\mathbf{x})\). One can relate it to the well-known translation invariance in image processing, but our assumption is weaker, as we allow small difference rather than require strict equivalence,

**[Spatial Regularity]** There exist a positive constant \(C\) depending on \(f\) such that

\[\|\mathcal{T}_{t}(\tau_{\mathbf{h}}f)-\tau_{\mathbf{h}}(\mathcal{T}_{t}f)\|_{L^{ \infty}}\leq Cht\]
for all \(f\in C_{b}^{\infty},\mathbf{h}\in\mathbb{R}^{d},t\geq 0\), where \((\tau_{\mathbf{h}}f)(\mathbf{x})=f(\mathbf{x}+\mathbf{h})\) and \(\|\mathbf{h}\|_{2}=h\).

_Remark 2.1_.: Spatial regularity is also beneficial for adversarial robustness. DNN have been shown to be vulnerable to some well-designed input samples (adversarial examples) [10, 20]. These adversarial examples are produced by adding carefully hand-crafted perturbations to the inputs of the targeted model. Although these perturbations are imperceptible to human eyes, they can fool DNN to make wrong prediction. In some sense, the existence of these adversarial examples is due to spatial unstability of DNN. So in our method, we hope the new model \(\mathcal{T}_{t}(f)\) to be spatially stable.

Secondly, temporal stability requires that in any small time interval, the evolution process will not be rapid. We want a smooth operator \(\mathcal{T}_{t}\) in time. Our assumption goes

**[Temporal Regularity]** For all \(t,s,t+s\in[0,T]\) and all \(f\in C_{b}^{\infty}\), there exist a constant \(C\geq 0\) depending on \(f\) such that

\[\|\mathcal{T}_{t+s,s}(f)-f\|_{L^{\infty}} \leq Ct\] \[\|\mathcal{T}_{t+s,s}(f)-\mathcal{T}_{t}(f)\|_{L^{\infty}} \leq Cst\]

Finally, combine all the assumptions on \(\mathcal{T}_{t}\), we can derive the following theorem, emphasizing that the output value of neural network \(T_{t}(f)\) with time evolution satisfies a convection-diffusion equation,

**Theorem 2.2**.: _Under the above assumptions, there exists Lipschitz continuous function \(v:\mathbb{R}^{d}\times[0,T]\to\mathbb{R}^{d}\) and Lipschitz continuous positive function \(\sigma:\mathbb{R}^{d}\times[0,T]\to\mathbb{R}^{d\times d}\) such that for any bounded and uniformly continuous base classifier \(f(\mathbf{x})\), \(u(\mathbf{x},t)=\mathcal{T}_{t}(f)(\mathbf{x})\) is the unique solution of the following convection-diffusion equation:_

\[\begin{cases}\frac{\partial u(\mathbf{x},t)}{\partial t}=v(\mathbf{x},t)\cdot\nabla u (\mathbf{x},t)+\sum_{i,j}\sigma_{i,j}\frac{\partial^{2}u}{\partial x_{i}\partial x _{j}}(\mathbf{x},t),\\ u(\mathbf{x},0)=f(\mathbf{x}),\end{cases} \tag{1}\]

_where \(\mathbf{x}\in\mathbb{R}^{d},t\in[0,T]\). Here \(\sigma_{i,j}\) is the \(i,j\)-th element of matrix function \(\sigma(\mathbf{x},t)\)._

_Remark 2.3_.: The right hands of the differential equation in (1) consist of two terms, the first order term \(v(\mathbf{x},t)\cdot\nabla u(\mathbf{x},t)\) called convection term and the second order term \(\sum_{i,j}\sigma_{i,j}\frac{\partial^{2}u}{\partial x_{i}\partial x_{j}}(\mathbf{ x},t)\) called diffusion term.

_Remark 2.4_.: When \(\sigma(\mathbf{x},t)=\sigma^{2}\mathbf{I}\), we call these type equations isotropic equations and the corresponding models isotropic models. When \(\sigma(\mathbf{x},t)\) is a diagonal matrix and \(\sigma(\mathbf{x},t)\neq\sigma^{2}\mathbf{I}\), we call these type equations anisotropic equations that lead to anisotropic models.

We will provide the proof of Theorem 2.2 in Appendix A. In this subsection, we have introduced a convection-diffusion equation framework for ResNets. The framework is quite general, as many existing models with residual connections can be interpreted as special cases in our framework.

### Examples Convection-Diffusion Model

Under the convection-diffusion framework, we can give interpretation to several regularization mechanisms including Gaussian noise injection [38, 24], ResNet with stochastic dropping out the hidden state of residual block [35, 36] and randomized smoothing [5, 23, 31]. Actually, they can be seen as convection-diffusion model with different diffusion term. Corresponding diffusion terms of these models are listed in Table 1.
**Gaussian noise injection:** Gaussian noise injection is an effective regularization mechanism for a DNN model. For a vanilla ResNet with \(L\) residual mapping, the \(n\)-th residual mapping with Gaussian noise injected can be written as

\[\mathbf{x}_{n+1}=\mathbf{x}_{n}+\mathcal{F}(\mathbf{x}_{n},\mathbf{w}_{n})+a\mathbf{\varepsilon}_{n },\ \mathbf{\varepsilon}_{n}\sim\mathcal{N}(0,\mathbf{I})\]

where the parameter \(a\) is a noise coefficient. By introducing a temporal partition: \(t_{n}=nT/L\), for \(n=0,1,..,L\) with the time interval \(\Delta t=T/L\) and let \(\mathbf{x}(t_{n})=\mathbf{x}_{n}\) and \(\mathbf{w}(t_{n})=\mathbf{w}_{n}\). And let \(a=\sigma\sqrt{\Delta t}\) and \(\mathcal{F}(\mathbf{x}_{n},\mathbf{w}_{n})/\Delta t=v(\mathbf{x},t)\). This noise injection technique in a discrete neural network can be viewed as the approximation of continuous dynamic

\[d\mathbf{x}(t)=v(\mathbf{x}_{n},t)dt+\sigma d\mathbf{B}(t) \tag{2.2}\]

where \(\mathbf{B}(t)\) is multidimensional Brownian motion. The output of \(L\)-th residual mapping can be written as \(\hat{\text{I}}\hat{\text{$o$}}\) process (2.2) at terminal time \(T\), \(\mathbf{x}(T)\). So, an ensemble prediction over all the possible sub-networks with shared parameters can be written as

\[\hat{y}=\mathbb{E}(\text{softmax}(\mathbf{w}_{\text{fc}}\mathbf{x}(T))|\mathbf{x}(0)=\mathbf{x }_{0}). \tag{2.3}\]

According to Feynman-Kac formula [26], Equation (2.3) is known to solve the following convection-diffusion equation

\[\begin{cases}\frac{\partial u(\mathbf{x},t)}{\partial t}=v(\mathbf{x},t)\cdot\nabla u +\sigma^{2}\Delta u,\quad\mathbf{x}\in\mathbb{R}^{d},t\in[0,T]\\ u(\mathbf{x},0)=\text{softmax}(\mathbf{w}_{\text{fc}}\mathbf{x}).\end{cases}\]

**Dropout of Hidden Units:** Consider the case that we disable every hidden units independently from a Bernoulli distribution \(\mathcal{B}(1,p)\) with \(p\in(0,1)\) in each residual mapping

\[\mathbf{x}_{n+1} =\mathbf{x}_{n}+\mathcal{F}(\mathbf{x}_{n},\mathbf{w}_{n})\odot\frac{\mathbf{z}_ {n}}{p}\] \[=\mathbf{x}_{n}+\mathcal{F}(\mathbf{x}_{n},\mathbf{w}_{n})+\mathcal{F}(\mathbf{x} _{n},\mathbf{w}_{n})\odot(\frac{\mathbf{z}_{n}}{p}-\mathbf{I})\]

where \(\mathbf{z}_{n}\sim\mathcal{B}(1,p)\) namely \(\mathbb{P}(\mathbf{z}_{n}=0)=1-p\), \(\mathbb{P}(\mathbf{z}_{n}=1)=p\) and \(\odot\) indicates the Hadamard product. If the number of the ensemble is large enough, according to Central Limit Theorem, we have

\[\mathcal{F}(\mathbf{x}_{n},\mathbf{w}_{n})\odot(\frac{\mathbf{z}_{n}}{p}-\mathbf{I})\approx \mathcal{F}(\mathbf{x}_{n},\mathbf{w}_{n})\odot\mathcal{N}(0,\frac{1-p}{p})\]

\begin{table}
\begin{tabular}{c c} \hline Models & Diffusion terms \\ \hline ResNet & 0 \\ Gaussian noise injection & \(\sigma^{2}\Delta u\) \\ Dropout of Hidden Units & \(\frac{1-p}{2p}\sum_{i}(v^{T}v)_{i,i}\frac{\partial^{2}u}{\partial x_{i}^{2}}\) \\ Randomized smoothing & \(\sigma^{2}\Delta u\) \\ \hline \end{tabular}
\end{table}
Table 1: Examples of networks under our proposed framework
The similar way with Gaussian noise injection, the ensemble prediction \(\hat{y}\) can be viewed as the solution \(u(\mathbf{x},T)\) of following equation:

\[\begin{cases}\frac{\partial u(\mathbf{x},t)}{\partial t}=v(\mathbf{x},t)\cdot\nabla u( \mathbf{x},t)+\frac{1-p}{2p}\sum_{i}(v^{T}v)_{i,i}\frac{\partial^{2}u}{\partial x_{ i}^{2}}(\mathbf{x},t),\quad\mathbf{x}\in\mathbb{R}^{d},t\in[0,T]\\ u(\mathbf{x},0)=\operatorname{softmax}(\mathbf{w}_{\text{fc}}\mathbf{x}),\end{cases}\]

_Remark 2.5_.: In fact, similar to dropout, shake-shake regularization [9, 14] and ResNet with stochastic depth [16] can also be interpreted by our convection-diffusion equation model.

**Randomized Smoothing:** Consider to transform a trained classifier into a new smoothed classifier by adding Gaussian noise to the input when inference time. If we denote the trained classifier by \(f(\mathbf{x})\) and denote the new smoothed classifier by \(g(\mathbf{x})\). Then \(f(\mathbf{x})\) and \(g(\mathbf{x})\) have the following relation:

\[g(\mathbf{x})=\frac{1}{N}\sum_{i=1}^{N}f(\mathbf{x}+\mathbf{\varepsilon}_{i})\approx \mathbb{E}_{\mathbf{e}\sim\mathcal{N}(0,\sigma^{2}I)}[f(\mathbf{x}+\mathbf{\varepsilon})]\]

where \(\mathbf{\varepsilon}_{i}\sim\mathcal{N}(0,\sigma^{2}I)\). According to Feynman-Kac formula, \(g(\mathbf{x})\) can be viewed as the solution of the following PDEs

\[\begin{cases}\frac{\partial u(\mathbf{x},t)}{\partial t}=\frac{1}{2}\sigma^{2} \Delta u,t\in[0,1]\\ u(\mathbf{x},0)=f(\mathbf{x}).\end{cases} \tag{4}\]

Especially, when \(f(\mathbf{x})\) is ResNet, the smoothed classifier \(g(\mathbf{x})=u(\mathbf{x},T+1)\) can be expressed as

\[\begin{cases}\frac{\partial u(\mathbf{x},t)}{\partial t}=v(\mathbf{x},t)\cdot\nabla u (\mathbf{x},t),\mathbf{x}\in\mathbb{R}^{d},t\in[0,T]\\ \frac{\partial u(\mathbf{x},t)}{\partial t}=\frac{1}{2}\sigma^{2}\Delta u,t\in[T,T +1]\\ u(\mathbf{x},0)=\operatorname{softmax}(\mathbf{w}_{\text{fc}}\mathbf{x}).\end{cases}\]

The differential equation formulation of randomized smoothing is similar to our method presented in the next section. However, randomized smoothing is a postprocessing step which ensures certified robustness. It does not involve the training of velocity field \(v\), which is parametrized by a neural network. Moreover, our method adds regularization term in each time step \(t\), while randomized smoothing only adds Gaussian noise on initial time step \(t=0\).

Obviously, ResNet is no diffusion model, Gaussian noise injection and randomized smoothing are isotropic models and dropout of hidden units is anisotropic models.

The viewpoint that forward propagation of a ResNet is the solving process of a TE enables us to interpret poor robustness and generalizability as the irregularity of the solution. Moreover, we can also interpret the effectiveness of the models above-mentioned in improving generalizability as the action of diffusion term. Next, we focus on the case of isotropic models, because anisotropic models are difficult for theoretical analysis and practical experiment and we take it as our future work.

## 3 Theoretical Analysis

In this section, under our convection-diffusion model, we focus on isotropic-type equation. Furthermore, for the sake of simplicity and practical computation, we consider the model with split convection and diffusion in
time.

\[\begin{cases}\frac{\partial u(\mathbf{x},t)}{\partial t}=v(\mathbf{x},t)\cdot\nabla u(\mathbf{ x},t),&\mathbf{x}\in\mathbb{R}^{d},\quad t\in[0,T-1]\\ \frac{\partial u(\mathbf{x},t)}{\partial t}=\sigma^{2}\Delta u,&\mathbf{x}\in\mathbb{R }^{d},\quad t\in[T-1,T]\\ u(\mathbf{x},0)=f(\mathbf{x})\end{cases} \tag{1}\]

In the following subsections, we will theoretically illustrate its benefits for improving the robustness and reducing the Rademacher complexity of neural networks.

### Robustness guarantee

Assume the data point lies in a bounded domain \(\mathcal{D}\). Consider a multi-class classification problem from \(\mathcal{D}\subset\mathbb{R}^{d}\) to label class \(\mathcal{Y}=\{1,\cdots,k\}\). Let \(G\) be a prediction function defined by \(G(\mathbf{x})=\operatorname*{arg\,max}_{i\in\mathcal{Y}}u^{i}(\mathbf{x},T)\), where \(u^{i}(\mathbf{x},T)\) is the \(i\)-th element of \(u(\mathbf{x},T)\). Suppose that DNN classifies \(\mathbf{x}\) the most probable class \(c_{A}\) is returned with probability \(p_{A}=u^{c_{A}}(\mathbf{x},T)\), and the "runner-up" class is returned with probability \(p_{B}\). Our main result of this subsection is to estimate the area around the data point \(\mathbf{x}\) in which the prediction \(G(\mathbf{x})\) is robust,

Suppose the velocity \(v(\mathbf{x},t)\) is a continuous function \(\mathbb{R}^{d}\times[0,T]\) which satisfies the Lipschitz condition, i.e. there exists a given constant \(L>0\), such that

\[\|v(\mathbf{x}_{1},t)-v(\mathbf{x}_{2},t)\|\leq L\|\mathbf{x}_{1}-\mathbf{x}_{2}\|,\quad\forall (\mathbf{x}_{1},t),(\mathbf{x}_{2},t)\in\mathbb{R}^{d}\times[0,T]\]

Let \(u(\mathbf{x},T)\) be the bounded solution of Equation (1). Suppose \(c_{A}\in\mathcal{Y}\) and

\[u^{c_{A}}(\mathbf{x},T)=p_{A}\geq p_{B}=\max_{i\neq c_{A}}u^{i}(\mathbf{x},T),\]

then \(G(\mathbf{x}+\mathbf{\delta})=c_{A}\) for all \(\|\mathbf{\delta}\|_{2}\leq R\), where

\[R=\frac{\sigma}{\sqrt{2d}}(p_{A}-p_{B})\]

We assume \(v\) as Lipschitz continuous because \(v(\mathbf{x},t)=\frac{1}{\Delta t}F(\mathbf{x},\mathbf{w})\) corresponds to a residual block, so its gradient can be bounded easily using the network parameters \(\mathbf{w}\).

We provide the proof of Theorem 3.1 in Appendix B. According to Theorem 3.1, we can obtain the certified radius \(R\) is large when the diffusion coefficient \(\sigma^{2}\) and the probability of the top class is high. We will also consider the generalization ability of model Equation (1) in the next subsection.

### Rademacher complexity

For simplicity, we consider binary classification problems. Assume that data points are drawn from the underlying distribution \(\mathcal{D}\). The training set \(S_{N}=\{\mathbf{x}_{i}\}_{i=1}^{N}\) is composed of \(N\) samples drawn i.i.d. from \(\mathcal{D}\). Rademacher complexity is one of the classic measures of generalization error. We first recap on the definition of empirical Rademacher complexity.

[[22]] Let \(\mathcal{H}:X\to\mathbb{R}\) be the space of real-valued functions on the space \(X\). For a given sample set \(S_{N}=\{\mathbf{x}_{i}\}_{i=1}^{N}\), the empirical Rademacher complexity of \(\mathcal{H}\) is defined as

\[R_{S_{N}}(\mathcal{H}):=\frac{1}{N}\mathbb{E}_{\sigma}[\sup_{h\in\mathcal{H}} \sum_{i=1}^{N}\sigma_{i}h(\mathbf{x}_{i})]\]

where \(\sigma_{1},\sigma_{2},\cdots,\sigma_{N}\) are i.i.d. Rademacher random variables with \(\mathbb{P}(\sigma_{i}=1)=\mathbb{P}(\sigma_{i}=-1)=\frac{1}{2}\).
Rademacher complexity is a tool to bound the generalization error. The smaller the generalization gap is, the less overfitting the model is. We are interested in the empirical Rademacher complexity of the following function class:

\[\mathcal{G}_{\sigma}:=\left\{g(\mathbf{x})=\mathcal{T}_{T}(f)|f\in\mathcal{F}\right\}\]

where

\[\mathcal{F}:=\left\{f:\mathbf{x}\mapsto\phi(\mathbf{w}_{\mathrm{fc}}\mathbf{x})|\left\|\mathbf{ w}_{\mathrm{fc}}\right\|_{1}\leq W\right\}\]

Here \(\mathcal{T}_{t}\) is the solution operator of (1), \(\phi\) is sigmoid activation function. The function class \(\mathcal{F}\) represents the hypothesis linear classifiers function class, where we assume the \(\ell^{1}\) norm of the weight of fully-connected layer is bounded by \(W\). The function class \(\mathcal{G}_{\sigma}\) includes the evolved residual neural networks from base classifiers. We can also assume the data points are bounded by \(R\), which is reasonable, e.g. in CIFAR-10 dataset [19], the pixel values lie in \([0,1]\). Then we have the following theorem.

**Theorem 3.1**: _Given a data set \(S_{N}=\{\mathbf{x}_{i}\}_{i=1}^{N}\). Suppose the data points \(\|\mathbf{x}_{i}\|_{\infty}\leq R\) for \(\mathbf{x}_{i}\in S_{N}\), then_

\[R_{S_{N}}(\mathcal{G}_{\sigma})\leq R_{S_{N}}(\mathcal{F})\leq\inf_{\epsilon} \left(\sqrt{\frac{2d\log(3WR/\epsilon)}{N}}+\epsilon\right)\]

We provide the proof of Theorem 3.1 in Appendix C. According to Theorem 3.1, we can obtain that the empirical Rademacher complexity of \(\mathcal{G}_{\sigma}\) can be upper bounded by that of a simple linear classifier. This can help bounding the generalization error of the evolved neural network. In the next section, we will present a training method for ResNet to verify our results.

## 4 Training Algorithm

To verify the effectiveness of the convection-diffusion model, we design a training method for ResNet. We assume that the data point lies in an unknown domain \(\mathcal{D}\subset\mathbb{R}^{d}\) and the label function \(l(\mathbf{x})\) is defined on \(\mathcal{D}\). Let \(\mathcal{S}=\{(\mathbf{x}_{i},y_{i})\}_{i=1}^{N}\subset\mathcal{D}\) be the training set, where \(\mathbf{x}_{i}\) is a data point sampled from \(\mathcal{D}\) and \(y_{i}=l(\mathbf{x}_{i})\) is the corresponding label.

As stated in the introduction, the forward propagation of ResNet corresponds to the transport equation. Denote ResNet as \(g_{\theta}\) with trainable parameters \(\theta\). Then the process from \(0\) to \(T-1\) in model (1) corresponds to the forward propagation of \(g_{\theta}\). In other words, the transport equation part in the convection-diffusion model (1) is already inherently included in \(g_{\theta}\). To incorporate the diffusion part, we enforce the network \(g_{\theta}\) to satisfy the following constraint

\[\begin{cases}\frac{\partial g_{\theta}(\mathbf{x},t)}{\partial t}=\sigma^{2}\Delta g _{\theta}(\mathbf{x},t),&\mathbf{x}\in\mathcal{D},t\in[0,1]\\ g_{\theta}(\mathbf{x},t)=l(\mathbf{x}),&\mathbf{x}\in\mathcal{S},t\in[0,1]\end{cases} \tag{1}\]

where we set \(T-1\) as the new initial time. The first equation is the diffusion part in convection-diffusion model (1). The second equation is the boundary condition for points in the training set. It is natural to require the network to classify training data correctly at any time \(t\in[0,1]\). To impose the constraints, we follow the techniques in PINN [30] and design two loss terms for the boundary condition and differential equation respectively.
First, to fit the boundary condition, we use the following loss

\[L_{1}=\int_{\mathcal{S}}\int_{0}^{1}l^{\mathrm{CE}}\left(g_{\theta}(\mathbf{x},s),l( \mathbf{x})\right)\mathrm{d}s\mathrm{d}\mathbf{x}\approx\frac{1}{MN}\sum_{i=1}^{N}\sum_ {k=1}^{M}l^{\mathrm{CE}}(g_{\theta}(\mathbf{x}_{i},t_{k}),y_{i})\]

where \(l^{\mathrm{CE}}\) denotes cross-entropy loss. \(N\) is the number of data points in \(\mathcal{S}\), and \(M\) is the number of time steps. In practice, we evenly choose \(t_{k}=(k-1)/(M-1)\) to discretize time.

Then, to fit the differential equation, we use the following mean square error loss

\[L_{2}=\int_{\mathcal{D}}\int_{0}^{1}\left(\frac{\partial g_{\theta}}{\partial t }(\mathbf{x},s)-\sigma^{2}\Delta g_{\theta}(\mathbf{x},s)\right)^{2}\mathrm{d}s \mathrm{d}\mathbf{x}\]

In practice, we only have access to the training set \(\mathcal{S}\) and do not know the underlying domain \(\mathcal{D}\), thus we treat the neighborhood of \(\mathcal{S}\) as the domain. The neighborhood is obtained by adding a uniform noise \(\epsilon\) to \(\mathcal{S}\). Then the integral term is approximated by

\[L_{2} \approx\frac{1}{N}\sum_{i=1}^{N}\int_{0}^{1}\left(\frac{\partial g _{\theta}}{\partial t}(\mathbf{x}_{i}+\epsilon,s)-\sigma^{2}\Delta g_{\theta}( \mathbf{x}_{i}+\epsilon,s)\right)^{2}\mathrm{d}s\] \[\approx\frac{1}{MN}\sum_{i=1}^{N}\sum_{k=1}^{M}\left(\frac{ \partial g_{\theta}}{\partial t}(\mathbf{x}_{i}+\epsilon,t_{k})-\sigma^{2}\Delta g _{\theta}(\mathbf{x}_{i}+\epsilon,t_{k})\right)^{2}.\]

Nonetheless, in neural network, the exact computation of Laplace of output w.r.t the input is computationally intractable [27] because of the high dimension inputs, e.g. for CIFAR-10 pictures the dimension is 3072. Thus, we use finite difference method to approximate. Denote difference operator \(\Delta_{h,\mathbf{v}}\) by

\[\Delta_{h,\mathbf{v}}g_{\theta}(\mathbf{x}_{i},s)=\frac{g_{\theta}(\mathbf{x}_{i}+h\mathbf{v},s)+g_{\theta}(\mathbf{x}_{i}-h\mathbf{v},s)-2g_{\theta}(\mathbf{x}_{i},s)}{h^{2}}\]

Using Taylor formula and law of large numbers, we have

\[\Delta g_{\theta}(\mathbf{x}_{i},s)=\mathbb{E}_{\mathbf{v}\sim\mathcal{N}(0,I)}\left( \Delta_{h,\mathbf{v}}g_{\theta}(\mathbf{x}_{i},s)\right)+O(h^{2})\approx\frac{1}{K} \sum_{j=1}^{K}\Delta_{h,\mathbf{v}_{i,j}}g_{\theta}(\mathbf{x}_{i},s)\]

where \(\{\mathbf{v}_{i,j}\}\) is i.i.d and standard normal distributed, \(K\) is the average number. Unless otherwise specified, we set \(K\) equals \(1\) in order to reduce the computation cost. Similarly, we introduce the time differential operator \(\mathbf{dt}_{\tau}\),

\[\mathbf{dt}_{\tau}g_{\theta}(\mathbf{x}_{i},s)=\frac{g_{\theta}(\mathbf{x}_{i},s+\tau) -g_{\theta}(\mathbf{x}_{i},s-\tau)}{2\tau}\]

to substitute for \(\frac{\partial g_{\theta}}{\partial t}\). Include these difference operators into \(L_{2}\), now the loss term for differential equation becomes

\[L_{2}=\frac{1}{MN}\sum_{i=1}^{N}\sum_{k=1}^{M}\left(\mathbf{dt}_{\tau}g_{ \theta}(\mathbf{x}_{i}+\epsilon,t_{k})-\sigma^{2}\frac{1}{K}\sum_{j=1}^{K}\Delta_ {h,\mathbf{v}_{i,j}}g_{\theta}(\mathbf{x}_{i}+\epsilon,t_{k})\right)^{2}\]
The final loss function for training \(g_{\theta}\) is a weighted sum of \(L_{1}\) and \(L_{2}\), \(L=L_{1}+\lambda L_{2}\).

Since the input of ResNet \(g_{\theta}\) is \((\mathbf{x},t)\), which now contains an additional time variable \(t\), we modify the network input dimension from \([C\times H\times W]\) to \([(C+1)\times H\times W]\), where the additional channal represents \(t\). Accordingly, we slightly change the structure of ResNet by adding an additional channel in the first convolution layer. Other than that, our ResNet structure is the same as vanilla ResNet. Obviously, when \(t=0\), the additional channel has no impact on the network output, and it functions as same as vanilla ResNet.

## 5 Experiments

In this section, we will numerically verify the performance of our training method for ResNets.

### Preliminaries

_Datasets._ In our experiments, we consider both synthetic half-moon dataset and real-world CIFAR-10 [19], Fashion-MNIST [40] and SVHN [28] datasets. Half-moon dataset is a randomly generated 2d synthetic dataset in which we randomly generate 500 points and 1000 points with a standard deviation of 0.3 as training set and testing set, respectively. CIFAR-10 contains 60K \(3\times 32\times 32\) color images from 10 different classes with 50K and 10K of them used for training and testing, respectively. SVHN is a 10-class house number classification dataset which contains 73257 training images and 26032 testing images, each of size \(3\times 32\times 32\). Fashion-MNIST is a 10-class greyscale dataset which contains 60K training images and 10K testing images. The size of each image is \(28\times 28\).

_Performance evaluations._ We evaluate the performance by both natural accuracy on original test samples and robust accuracy on adversarial test samples within a perturbation range \(\epsilon\). To craft adversarial samples, we use Project Gradient Descent (PGD) and AutoAttack [6]. PGD first adds random uniform noise from \((-\epsilon,\epsilon)\) to clean samples, and iterates Fast Gradient Sign Method (FGSM) with fixed step size \(\alpha\) for \(m\) steps,

\[\mathbf{x}^{(m)}=\text{Clip}_{\mathbf{x},\epsilon}\{\mathbf{x}^{(m-1)}+\alpha\text{sign}( \nabla\mathcal{L}(\mathbf{x}^{(m-1)},y))\}\]

AutoAttack [6] is a more reliable benchmark for evaluating robustness, which ensembles four strong parameter-free attacks. It is reported that some defense methods which claim to be robust against PGD attack are destroyed by AutoAttack. Thus, we report classification accuracy on adversarial examples crafted by AutoAttack to evaluate the robustness of our method.

For half-moon dataset, we apply PGD[20] (step number \(m=20\)) attack with \(\alpha=0.01\) and \(\epsilon=0.2\) under \(l^{\infty}\) norm. For CIFAR-10 and SVHN, we apply PGD[20] attack with \(\alpha=2/255\) and \(\epsilon=8/255\). For Fashion-MNIST, we apply PGD[20] attack with \(\alpha=0.01\) and \(\epsilon=0.1\). AutoAttack is parameter-free and thus we only set \(\epsilon\) as same as that of PGD attack.

### Experiments on synthetic dataset

In this subsection, we numerically verify the efficacy of our training method in improving the robustness of ResNet on the half-moon dataset.

We first vary hyperparameters \(\lambda\) and \(\sigma^{2}\) and present the classification accuracy on adversarial examples in Figure 2. The robust accuracy increase when \(\lambda\) and \(\sigma^{2}\) increase, which indicates that the introduced regularizer is helpful for improving the performance of ResNet. Moreover, we plot the decision boundary of naturally trained ResNet and ResNet trained with different \(\lambda,\sigma^{2}\) in Figure 3. We can observe that
the decision boundary of natural training is irregular, while that of our models is smoother. These experimental results are consistent with our theory.

### Experiment on benchmarks

We further test the performance of our method on CIFAR-10, SVHN and Fashion-MNIST datasets. We choose ResNet18 as the backbone model, where our method will add an additional channel in the

\begin{table}
\begin{tabular}{c c c c c} \hline \hline Dataset & Methods & Natural & PGD[20] & AutoAttack \\ \hline \multirow{2}{*}{CIFAR-10} & ResNet18 & 95.18 & 0.0 & 0.0 \\  & Ours & 88.54 & 20.02 & 17.64 \\ \hline \multirow{2}{*}{SVHN} & ResNet18 & 96.58 & 0.40 & 0.02 \\  & Ours & 94.16 & 19.05 & 15.48 \\ \hline \multirow{2}{*}{Fashion-MNIST} & ResNet18 & 93.95 & 0.0 & 0.0 \\  & Ours & 92.01 & 35.27 & 23.48 \\ \hline \hline \end{tabular}
\end{table}
Table 2: Natural accuracy and robust accuracy under PGD[20] and AutoAttack on different datasets (%).

Figure 3: Decision boundary of natural trained ResNet and ResNet trained by our method with different hyperparameters \(\lambda\) and \(\sigma^{2}\).

Figure 2: Robust accuracy of ResNet trained by our model with different hyperparameters \(\lambda\) and \(\sigma^{2}\) on half-moon dataset.
input convolutional layer. During training, we apply standard data augmentation techniques including random crops and horizontal flips. The batch size is 128. We run 200 epochs with initial learning rate of 0.1, which decays by a factor of 10 at the 80th, 120th and 160th epochs. We use stochastic gradient descent optimizer with momentum of 0.9 and weight decay of \(5\times 10^{-4}\).

There are several hyperparameters in our algorithm that need to be determined. \(h\) and \(\tau\) are spatial and temporal discretization parameters. We choose \(\tau=10^{-4}\) and \(h=0.1\). We choose a fairly large \(h\) because neural networks are highly unstable on the spatial dimension and cannot converge when \(h\) is too small. To reduce the computation cost, we choose \(K=1\) when computing the Laplacian \(\Delta\), and only consider starting and final time step \(t_{1}=0,t_{2}=1\) when computing the regularization loss. The uniform noise \(\epsilon\) which we add to training set \(\mathcal{S}\) to fit the underlying domain \(\mathcal{D}\) is chosen the same as the attack range of each dataset, i.e. 8/255 for CIFAR-10 and SVHN, and 0.1 for Fashion-MNIST. \(\lambda\) and \(\sigma^{2}\) both affect the trade-off between natural accuracy and robust accuracy. With larger \(\lambda\) or \(\sigma^{2}\), the natural accuracy decreases, while robust accuracy increases. We report the result of \(\lambda=0.005\) and \(\sigma^{2}=0.2\) in Table 2, and an ablation study on the two parameters can be found in Table 3. From Table 2, we can see that our method, which introduces a regularization term based on convection-diffusion differential equation, raises the classification accuracy of ResNet against adversarial samples by a large margin. We want to emphasize that our method does not include any adversarial training techniques, which trains the model using adversarial samples crafted by PGD attacks. From Table 3, we observe that with fixed parameter \(\sigma^{2}\), the natural accuracy decreases and robust accuracy increases as the increasing of parameters \(\lambda\) and vice versa.

## 6 Conclusion

In this paper, we theoretically prove that under reasonable assumptions, the evolution from a base linear classifier to residual neural networks should be modeled by a convection-diffusion equation. Motivated by PDE theory, we analyze the robustness and Rademacher complexity of the proposed isotropic models. Based on these theoretical results, we develop a training method for ResNet and verify its effectiveness through experiments. We are aware that modeling the convection-diffusion equation through introducing a regularization term is one of the

\begin{table}
\begin{tabular}{c c c c c} \hline \hline \(\lambda\) & \(\sigma^{2}\) & Natural & PGD[20] & AutoAttack \\ \hline
0.2 & 0.2 & 80.33 & 24.91 & 23.89 \\
0.1 & 0.2 & 82.64 & 24.04 & 22.97 \\
0.05 & 0.2 & 84.51 & 23.58 & 22.53 \\
0.01 & 0.2 & 87.45 & 20.63 & 18.37 \\
0.005 & 0.2 & 88.54 & 20.02 & 17.64 \\ \hline
0.2 & 0.1 & 81.79 & 20.56 & 19.74 \\
0.1 & 0.1 & 84.73 & 20.02 & 18.87 \\
0.05 & 0.1 & 85.95 & 19.07 & 17.76 \\
0.01 & 0.1 & 89.52 & 18.60 & 14.12 \\
0.005 & 0.1 & 89.95 & 15.22 & 11.11 \\ \hline \hline \end{tabular}
\end{table}
Table 3: Natural accuracy and robust accuracy under PGD[20] and AutoAttack on CIFAR-10, with varying parameter \(\lambda\) and \(\sigma^{2}\) (%).
many possible approaches. We are looking forward to explore other paths in the future work.
