2f:null 30:["$","title",null,{"children":"Neural networks & partial differential equations | Yao Liu | Observable"}] 35:T81d,md`TL;DR * If "Neural ODE" and "infinitely wide neural nets" are to turn the depth and width dimensions into continuous, it stands to reason that we shall gain some insights viewing the two image dimensions as continuous. Here's a quick summary in the form of a dictionary: convolutional neural nets | partial differential equations :----:|:-------: input layer | initial condition feed forward | solving the equation hidden layers | solution at intermediate times output layer | solution at final time **convolution** with 3×3 kernel | **differential operator** of order ≤ 2 weights | coefficients boundary handling
(padding) | boundary condition multiple channels
[e.g. 16×16×3×3 kernel;
16×16×1×1 kernel] | system of (coupled) PDEs
[16×16 matrix of differential operators;
16×16 matrix of constants] groups (=g) | matrix is block-diagonal (direct sum of g blocks) * It is then natural to consider *differential operators with variable coefficients*. With the "variables" simply *linear* (degree 1 in the image coordinates), the neural net already has the capacity to **rotate** and **scale** the image. Though simple to implement, it is not intuitive, if not impossible, to motivate or justify without taking the *continuous* perspective. * The training of a neural net would be the inverse problem to a PDE: given the solutions (or rather the solutions at initial and final times), what are the coefficients of the PDE that fits best? One might consider the backpropagation as another PDE (on those coefficients) in disguise that "solves" a certain **variational problem** numerically. Applied mathematicians have taken note; see the brief survey *[Machine Learning from a Continuous Viewpoint](https://arxiv.org/abs/1912.12777)* and references therein. * One might regard it as part of the general scheme of bringing the language and reasoning of continuous mathematics into traditionally discrete-based computer science, while being quite distinct from the language of probability and statistics (inference, distribution, etc.). `36:T5fd,md`Specifically, I'd like to highlight one natural **modification of ResNet** that comes out of this perspective (which seems to have escaped the other authors). For simplicity, consider a ResNet with a fixed number of channels throughout, and here are a few simple things that one could try out: 1. Convolution with ${tex`3\times 3`} kernel as the "residual" at *every* layer. 2. If we wish the model to be rotation- and/or scale-invariant, for each layer also include two other branches that get multiplied by the *coordinate functions* (i.e., ${tex`x`} and ${tex`y`}). That is, have three separate ${tex`3\times 3`} kernels for each layer, and take ${tex.block`I_{n+1}=I_n + K^{(0)} * I_n + x\, K^{(1)} * I_n + y\, K^{(2)} * I_n`} where ${tex`x`} and ${tex`y`} are (entry-wise multiplication by) the matrices \`\`np.indices((h,w))[1]\`\` and \`\`np.indices((h,w))[0]\`\`, if you will. 3. If that does not seem to make the neural net learn rotation, "symmetrize" the weights of ${tex`K^{(1)}`} and ${tex`K^{(2)}`} so that opposite corners (of the ${tex`3\times 3`} square) are equal and opposite. Furthermore, we may force ${tex`K^{(2)}`} to be a ${tex`90^\circ`} rotation of ${tex`K^{(1)}`}. 4. Instead of L1 or L2 regularization, we may impose a restriction or cost on the **variation** of the weights across layers (which one might call the **energy** of the neural net), or even freeze the kernels for several layers (weight sharing). Equivalently, one might do "weight smoothing" instead of, or in addition to, weight decay. `37:T566,md`Now, as it happens, deep neural nets may involve an astronomical number of parameters, aka weights, to optimize, far exceeding the dataset we have. (Remember, in passing to the continuous, the dimension of the space ${tex`X`} is simply the number of weights in the model.) Partly for this reason, the spectacular success of deep neural nets came as a surprise to many experts in AI, and the mystery is a big unsolved problem to this day. The conventional wisdom is that the "true" parameter space must be a much smaller submanifold of ${tex`X`}, though we can't know a priori what the submanifold is. In my opinion, the key to a better understanding of deep learning may lie in going a step *further*, to an optimization problem over an *infinite*-dimensional space ${tex`X`}, typically as the *space of functions* over some finite-dimensional ${tex`\Omega`}. It may seem a big conceptual leap, but keep in mind that it too is meant to make things *simpler* to reason about. This is squarely in the realm of the [calculus of variations](https://en.wikipedia.org/wiki/Calculus_of_variations), a subject as old as calculus itself but turned out to be essential for the fundamental laws of physics, cf. the [principle of least action](https://en.wikipedia.org/wiki/Principle_of_least_action) and other [variational principles](https://en.wikipedia.org/wiki/Variational_principle).`38:T895,md`## The heat equation There are many classical problems of the calculus of variations, but I think the [heat equation](https://en.wikipedia.org/wiki/Heat_equation) offers the best starting point. In addition, as one of the main characters for the story, I'd need to introduce it (as briefly as I can) anyways. Let ${tex`\Omega`} be a region in the 2-dimensional plane, and the heat inside is described by a function ${tex`u: \Omega\times\mathbb R_{\geq 0}\to\mathbb R`} that evolves over time (${tex`t\in\mathbb R_{\geq 0}`}) according to the **heat equation**: ${tex.block`\frac{\partial u}{\partial t} = \Delta u`} starting from some given initial data ${tex`u_0`} at time ${tex`t=0`}. Here, ${tex`\Delta`} is the Laplace's operator, or the Laplacian: ${tex.block`\Delta:=\frac{\partial^2}{\partial x^2}+\frac{\partial^2}{\partial y^2}`} (In words, the equation says that: where the heat is "concave down ${tex`\cap`}", the second derivative is negative, so it decreases in time, i.e., it should cool down; and where the heat is "concave up ${tex`\cup`}" it should heat up.) For a finite region ${tex`\Omega`} (as opposed to the entire ${tex`\mathbb R^2`}), we would also need to impose some boundary condition, e.g., ${tex`u=0`} at every point of the boundary of ${tex`\Omega`}, for *all* time. The problem then is to *solve the heat equation, inside ${tex`\Omega`} and for all ${tex`t>0`}, under these initial and boundary conditions*. It is a rather non-trivial theorem that the solution exists and is unique, even though it *appears* rather obvious that the heat must flow in some specific, deterministic way, reaching a thermal equilibrium. ![](https://upload.wikimedia.org/wikipedia/commons/a/a9/Heat_eqn.gif) > (Well, I suppose it's impossible to do better than 3blue1brown. Watch his recent [video](https://www.youtube.com/watch?v=ly4S0oi3Yz8) for an excellent visual introduction of the heat equation; and the [following one](https://www.youtube.com/watch?v=ToIXSwZ1pJU) on solving it by Fourier series, even though it's not so relevant to our story.) In passing, I should note that the heat equation also works in any (spatial) dimension, with much the same properties below. `39:T5ab,md`### Variational perspective As time ${tex`t`} increases, the function ${tex`u_t`} is approaching the minimum point ${tex`u^*`} of the **energy functional** ${tex.block` \begin{aligned} E[u] :=&\, \frac{1}{2}\int_{\Omega} |\nabla u|^2 \,da \\ =&\, \frac{1}{2} \int_{\Omega} \left( \bigl(\frac{\partial u}{\partial x}\bigr)^2 + \bigl(\frac{\partial u}{\partial y}\bigr)^2 \right) da \end{aligned} `} (${tex`da = dx\,dy`}, the ${tex`\frac{1}{2}`} is not important). Let's spell things out: ${tex`E`} is a function from ${tex`X\to\mathbb R`} where ${tex`X`} is the space of *all* (differentiable) functions ${tex`u:\Omega\to\mathbb R`} satisfying the *same* boundary condition, and ${tex.block`u^* = \argmin_{u\in X} E[u]`} Note that ${tex`X`} is an *infinite*-dimensional space. We'll come to this in a moment, but simply by virtue of the heat equation, the *terminal* or *equilibrium state* must satisfy ${tex`\Delta u=0`} in ${tex`\Omega`}. This equation with no time variable is known as **Laplace's equation**, and it is also very important in physics (as well as in pure mathematics, especially in complex analysis), and the intuition to keep in mind is that the *graph* of any solution ${tex`u`} would look like a rubber sheet that is taut and fixed along the boundary; in particular, it should *not* have any local maxima or minima in the *interior* of ${tex`\Omega`} — any local extrema would have been "ironed out" by the heat equation. `3a:T715,md` In fact, the Laplace's equation can be *derived* from the requirement that ${tex`u^*`} is a (local) minimizer of ${tex`E`}, by "differentiating ${tex`E`} with respect to ${tex`u`}": ${tex.block` \begin{aligned} \delta E &= \frac{1}{2} \int_\Omega \biggl( |\nabla(u+\delta u)|^2 - |\nabla u|^2 \biggr) da \\ &= \frac{1}{2} \int_\Omega 2\nabla u\cdot \nabla(\delta u)\, da \\ &=-\int_\Omega \Delta u\, \delta u \,da \end{aligned} `} which reads: a change in ${tex`u`} by an *arbitrary* infinitesimal ${tex`\delta u:\Omega\to\mathbb R`} results in a change in ${tex`E`} by ${tex`\delta E`}, which can be written as ${tex`\delta E =\langle -\Delta u, \delta u\rangle`} using the "dot product" of two functions. This means that ${tex`-\Delta u`} is *literally* the gradient of ${tex`E`}. [It's worth pausing for a moment. Recall that if ${tex`\Omega`} consists of a finite (but large) number of points, ${tex`u`} is equivalent to a finite-dimensional vector. The gradient of ${tex`E`} would be a vector of the same length (=dimension), and the (total) "differential" of ${tex`E`} would be given by a dot product: ${tex`dE=\langle\nabla E, du\rangle`}. In the infinite-dimensional case, the "gradient" should then be an infinite-dimensional vector of the same type as ${tex`u`}, i.e., a function of the type ${tex`\Omega\to\mathbb R`}.] Thus, we can solve this optimization problem either analytically, * Set the gradient to zero: ${tex`\Delta u=0`} (Laplace's equation); or "numerically", * Start with any ${tex`u_0`} and always take the steepest descent in the direction of ${tex`\Delta u`}, i.e., ${tex`u`} "flows" per the equation ${tex`\dfrac{\partial u}{\partial t}=\Delta u`} (the heat equation). That is, the **heat equation** = **gradient descent** for an *infinite-dimensional* optimization problem!`3b:T461,md`## ResNet What seemed too simple an idea, ResNet or Residual Net introduced by He Kaiming 何恺明, et al. (2015) quickly took over the field of computer vision and became the "backbone" of almost all deep neural nets. Let's leave aside the nonlinear activation, and only consider the "pure" ResNet, with a single channel throughout, i.e., ${tex.block`I_{n+1} = I_n + K_n * I_n`} (The input and the output are by necessity the same size.) As many have noticed, it is structually identical to the discretized heat equation we just saw ${tex.block`u_{t+\epsilon} = u_t + \epsilon\Delta u_t`} where the discrete Laplacian is also a convolution operator, namely ${tex.block`\Delta = \frac{1}{4} \begin{bmatrix} 0 & 1 & 0 \\ 1 & -4 & 1 \\ 0 & 1 & 0 \end{bmatrix} * `} That is, the layers or depth of the neural net correspond to the time variable, and the input is being *locally* deformed from one layer to the next (this is the intuitive picture that most people have for CNNs, if not put in these terms). The difference is that ${tex`K_n`} contains learnable parameters (the weights), while ${tex`\Delta`} is fixed. `3c:T686,md`Now, we can interpret the ResNet (without activation) as a linear partial differential equation of the form ${tex.block`\frac{\partial u}{\partial t} = \alpha u + \beta \frac{\partial u}{\partial x} + \gamma\frac{\partial u}{\partial y} + \delta\Delta u + \cdots`} where *the coefficients ${tex`\alpha, \beta, \ldots`} are functions of ${tex`t`}*, and these are what the neural net is actually learning. In fact, we have simply "re-packaged" the weights into functions on the time interval. Each of the coefficients has a "physical" interpretation: * ${tex`\alpha`} is how fast the signal intensifies or diminishes (according to the sign); * ${tex`\beta`} is how fast it shifts in the ${tex`x`} direction; * ${tex`\gamma`} is how fast it shifts in the ${tex`y`} direction; * ${tex`\delta`} is how fast the signal diffuses (if ${tex`\delta>0`}) in *all* directions, etc. It is not difficult to fill up the following correspondence neural network | heat equation :----:|:-------: input layer | initial condition "feed forward" | solving the equation hidden layers | solution at intermediate times output layer | solution at final time convolution with 3×3 kernel | differential operator of order ≤ 2 weights | coefficients boundary handling in convolution | boundary condition Note that, contrary to what the study of PDEs is concerned with, here we don't know a priori what the equation is, but we have solutions (or rather, the solutions at final time) for a variety of initial conditions, and we are trying to find the PDE that best fits these solutions. In other words, it is the inverse problem to solving a PDE. All this is in Ruthotto & Haber (2018). `3d:T432,md`Clearly, all this would be very helpful in **interpreting** a neural net. Should one as simple as this be successfully trained, we can convert the weights into coefficients for the spatial derivatives, which tell us what the neural net is doing to the image at each and every layer. More interestingly, this makes it clear why convolutional neural nets are *translation-invariant*, but not rotation- or scale-invariant, as usually understood. The standard thing to do is to augment the dataset to force the model to learn about these transformations. There are also other attempts at addressing this problem, including RotationNet, the Spatial Transformer Networks, and works that use fancy words like *group equivariance* and even *gauge theory* (or so claimed), but here we have a simple and natural solution to make the neural net "innately" rotation- and scale-invariant: incorporate the differential operators for these transformations into the neural net from the get-go. (Those who are familiar with the relevant mathematics may already see what I'm getting at.)`3e:T54f,md`It is well-known that rotation is likewise generated by an operator, namely ${tex.block`-y\frac{\partial}{\partial x} + x\frac{\partial}{\partial y}`} and that scaling is generated by the (so-called Euler) operator ${tex.block`x\frac{\partial}{\partial x} + y\frac{\partial}{\partial y}`} [Instead of expanding the exponential, one may simply write down a solution that rotates or scales the initial data ${tex`u_0`} as ${tex`t`} increases, and check by hand that it does satisfy the differential equation.] These are first-order differential operators but *with variable coefficients*, and they can also be interpreted (and visualized) as vector fields in the ${tex`xy`}-plane. Incidentally, the expressions above are actually the standard notation for vector fields in modern differential geometry. On a related note, infinitesimal generator and the **exponential map** are fundamental concepts in the theory of Lie groups and Lie algebras. It's also easy to write down the vector field for rotation around an arbitrary point: ${tex.block`-(y-y_0)\frac{\partial}{\partial x} + (x-x_0)\frac{\partial}{\partial y}`} In general, any first-order differential operator (with variable coefficients) is a vector field that generates a **flow** (a technical term, but just what you'd expect). In PDE, it is commonly known as the **method of characteristics**. `3f:T5d8,md`Higher-order operators are more complicated. The heat equation, from this perspective, is simply solved by (formally) exponentiating the Laplacian: ${tex.block`u_t = e^{t\Delta} u_0`} and one can write the operator ${tex`e^{t\Delta}`} more explicitly, at least for the rectangular domains ${tex`\Omega`}. Let me just add that, in the finite-dimensional case, we know to first diagonalize the matrix ${tex`A`}, then exponentiate; here we can use the same strategy: diagonalizing a differential operator is nothing but taking the Fourier series or Fourier transform. One surprising and little-known fact is that the Fourier transform can also be thought of as an exponential — namely, of the (Hermite) operator ${tex.block`-\frac{1}{2}\Delta + \frac{1}{2}(x^2+y^2) `} but for *imaginary* ${tex`t=\frac{i\pi}{2}`}. (For other values of ${tex`t`}, we obtain what's known as the [fractional Fourier transform](https://en.wikipedia.org/wiki/Fractional_Fourier_transform).) I'm not sure where one would typically learn about this stuff, but I encountered it in the context of the metaplectic (or Weil) representation of ${tex`\widetilde{SL}_2(\mathbb R)`} or of Lie algebra ${tex`\mathfrak{sl}_2`}. Very meta, indeed! With the appearance of ${tex`i`}, this looks eerily similar to the Schödinger's equation for the (quantum) harmonic oscillator... All this is to say that when you combine two differential operators (that do not commute), the exponential could be something entirely different. `40:T49b,md`I do not want to give the impression that this is the ultimate theory that *explains* neural nets as they are currently practiced. Yes, the simple proposal should expand the types of transformations that a neural net is capable of learning, but whether the neural net *actually* learns is yet to be seen. The full theory of neural nets is almost surely a combination of many explanations that all work at once, as with the oft-quoted analogy of the working of an airplane. For instance, *stochastic* gradient descent (SGD) is routinely used not only to speed up the training process, but to "get around" the local minima of the loss function. Furthermore, this is after all still a *linear* partial differential equation, and would respond the same way (shifts, rotates, diffuses) no matter what the input is. The nonlinear activation would probably be crucial to "separate" out the different classes. [The simplest approach is to add these nonlinear activations only once every few layers; that is, we'd have *blocks* of "pure ResNet", each of which is a linear PDE, connected by the usual activation functions and possibly some encoder/decoder net that change the width.] `41:[] 33:{"id":"c52dbcf9e1d9b456","slug":"neural-networks-and-partial-differential-equations","title":"Neural networks & partial differential equations","thumbnail":"75d41989e6a63c2450211dec83d49a6413c082194390fff5cd1d23a15d3bcc20","access_level":"public","create_time":"$D2019-02-17T01:05:13.701Z","update_time":"$D2021-03-16T14:07:53.213Z","view_time":null,"first_public_time":"$D2019-05-20T22:14:07.353Z","has_unpublished_edits":null,"fork":false,"stars":15,"starred":false,"owner":{"type":"team","id":"490427fcc4ec2c1d","login":"liuyao12","name":"Yao Liu","avatar_url":"https://avatars-next.observableusercontent.com/895964f8b86922c183ea757cd86944e8e9d463ef5ba08171fab539d4dc087c54","tier":"starter_2024"},"version":3061,"publish_version":null,"role":null,"authored":false,"authors":[{"type":"individual","id":"ae908de8dd38723a","login":"liuyao12","name":"Yao Liu","avatar_url":"https://avatars-next.observableusercontent.com/895964f8b86922c183ea757cd86944e8e9d463ef5ba08171fab539d4dc087c54","tier":null}],"members":[{"type":"individual","id":"ae908de8dd38723a","login":"liuyao12","name":"Yao Liu","avatar_url":"https://avatars-next.observableusercontent.com/895964f8b86922c183ea757cd86944e8e9d463ef5ba08171fab539d4dc087c54","tier":null,"role":"editor"}],"comments":[],"custom_thumbnail":false,"trashed":false,"lapsed":null,"subscribed":null,"forks":9,"collection":null,"type":"notebook","body":{"title":"Neural networks & partial differential equations","cells":[{"id":0,"value":"md`# Neural networks & partial differential equations\n\n> *Keywords*: ConvNet, ResNet, PDE, heat equation, optimal control, flow of vector field, rotation- & scale-invariant\n\n> (**Update**: 3/24/2020) *For a more practical approach and implementation in PyTorch, see [this Jupyter notebook](https://liuyao12.github.io/blog/research/2020/03/07/Conv-Twist.html).*\n`","pinned":false,"hidden":false,"mode":"ojs","author":"$undefined","summary":"$undefined","output":"$undefined","database":"$undefined"},{"id":3054,"value":"{\n replay0\n await visibility()\n yield* rotate()\n}","pinned":false,"hidden":false,"mode":"ojs","author":"$undefined","summary":"$undefined","output":"$undefined","database":"$undefined"},{"id":3057,"value":"viewof replay0 = html`