# A Neural-preconditioned Poisson Solver for Mixed Dirichlet and Neumann
Boundary Conditions

Kai Weixian Lan &Elias Gueidon &Ayano Kaneda &Julian Panetta 11footnotemark: 1
&Joseph Teran 11footnotemark: 1 University of California, Davis University of
California, Los AngelesWaseda University

###### Abstract

We introduce a neural-preconditioned iterative solver for Poisson equations
with mixed boundary conditions. The Poisson equation is ubiquitous in
scientific computing: it governs a wide array of physical phenomena, arises as
a subproblem in many numerical algorithms, and serves as a model problem for
the broader class of elliptic PDEs. The most popular Poisson discretizations
yield large sparse linear systems. At high resolution, and for performance-
critical applications, iterative solvers can be advantageous for these—but
only when paired with powerful preconditioners. The core of our solver is a
neural network trained to approximate the inverse of a discrete structured-
grid Laplace operator for a domain of arbitrary shape and with mixed boundary
conditions. The structure of this problem motivates a novel network
architecture that we demonstrate is highly effective as a preconditioner even
for boundary conditions outside the training set. We show that on challenging
test cases arising from an incompressible fluid simulation, our method
outperforms state-of-the-art solvers like algebraic multigrid as well as some
recent neural preconditioners.

##  1 Introduction

††footnotetext: Correspondence to Kai Weixian Lan<EMAIL_ADDRESS>⟩

The solution of linear systems of equations involving discrete Laplace
operators is the bottleneck in many engineering and scientific applications.
These large, symmetric positive definite and sparse systems of equations are
notoriously ill-conditioned. Fast Fourier Transforms (Cooley & Tukey, 1965)
are optimal for these problems when discretized over trivial geometric
domains, however they are not applicable for practical domain shapes. Direct
methods like Cholesky factorization (Golub & Loan, 2012) resolve conditioning
issues, but suffer from loss of sparsity/fill-in and are prohibitively costly
in practice when per-time-step refactoring is necessary (_e.g._ , with
changing domain shape). Iterative methods like preconditioned conjugate
gradient (PCG) (Saad, 2003) and multigrid (Brandt, 1977) can achieve good
performance, however an optimal preconditioning strategy is not generally
available, and though multigrid can guarantee modest iteration counts,
computational overhead associated with solver creation and other per-iteration
costs can dominate runtimes in practice. Unfortunately, there is no clear
algorithmic solution.

Recently, machine learning techniques have shown promise for these problems.
Tompson et al. (2017) showed that a network (FluidNet) can be used to generate
an approximate inverse across domain shapes, albeit only with Neumann boundary
conditions. Kaneda et al. (2023) developed DCDM, which improves on this
approach by using a similar network structure and an iterative technique where
gradient descent in the matrix-norm of the error is preconditioned with a
neural network. While their approach is similar to PCG, the nonlinearity of
their approximate inverse required a generalization of the PCG method which
proved effective. We build on this approach and generalize it to domains with
mixed Dirichlet and Neumann boundary conditions. Notably, these problems arise
in simulating free-surface liquid flows. The DCDM approach cannot handle these
cases, however we show that a novel, more lightweight network structure can be
used in DCDM’s iterative formalism that is both linear and capable of handling
mixed boundary conditions over time-varying fluid domains. Furthermore, we
show that this structure drastically improves performance over that in DCDM.
We design our network structure to represent the dense nature of the inverse
of a discrete Laplacian matrix. That is, the inverse matrix for a discrete
Laplace operator has the property that local perturbations anywhere in the
domain have non-negligible effects at all other points in the domain. Our
network structure uses a hierarchy of grid scales to improve the resolution of
this behavior over what is possible with the DCDM structure. In effect, the
process of transferring information across the hierarchy from fine grid to
increasingly coarse grids and back again facilitates rapid propagation of
information across the domain. This structure has similarities with multigrid,
however there are some important differences. We incorporate the effects of
the Dirichlet and Neumann conditions at irregular boundaries with a novel
convolution design. Specifically, we use stencils that learn spatially varying
weights based on a voxel’s proximity to the boundary and the boundary
condition types encoded there.

Although our approximate inverses are linear (unlike the DCDM preconditioner)
we still adopt the DCDM iterative formalism. We do this because we cannot
guarantee that our neural network produces a symmetric and positive definite
approximate inverse as required for standard PCG. It is possible to use a
flexible PCG technique (Golub & Ye, 1999) in this case though (as in
(Bouwmeester et al., 2015)), however we show that the matrix-orthogonal
gradient descent iteration in DCDM provides superior results. We show that our
network outperforms state-of-the-art preconditioning strategies, including
DCDM, FluidNet, algebraic multigrid and incomplete Cholesky. We perform our
comparison across a number of representative free-surface liquid and fluid
flow problems. To promote reproducibility we have released our full code and a
link to our pretrained model at https://anonymous.4open.science/r/MLPCG-2102.

##  2 Related Work

Many recent approaches leverage machine learning techniques to accelerate
numerical linear algebra computations. Ackmann et al. (2020) use supervised
learning to compute preconditioners from fully-connected feed-forward networks
in semi-implicit time stepping for weather and climate models. Sappl et al.
(2019) use convolutional neural networks (CNNs) to learn banded approximate
inverses for discrete Poisson equations arising in incompressible flows
discretized over voxelized spatial domains. However, their loss functions is
the condition number of the preconditioned operator which is prohibitively
costly at high resolution. Özbay et al. (2021) also use CNN to approximate
solutions to Poisson problems arising in incompressible flow discretized over
voxelized domains, however they do not learn a preconditioner and their
approach only supports two-dimensional square domains. Our approach is most
similar to those of Tompson et al. (2017) and Kaneda et al. (2023) who also
consider discrete Poisson equations over voxelized fluid domains, however our
lighter-weight network outperforms them and generalizes to a wider class of
boundary conditions. Li et al. (2023) build on the approach of Sappl et al.
(2019), but use a more practical loss function based on the supervised
difference between the inverse of their preconditioner times a vector and its
image under the matrix under consideration. Their preconditioner is the
product of easily invertible, sparse lower triangular matrices. Notably, their
approach works on disretizations over unstructured meshes. Götz & Anzt (2018)
learn Block-Jacobi preconditioners using deep CNNs. The choice of optimal
blocking is unclear for unstructured discretizations, and they use machine
learning techniques to improve upon the selection.

Various works use hybrid deep learning/multigrid techniques. For example, the
UNet (Ronneberger et al., 2015) and MSNet architectures (Mathieu et al., 2016)
are similar to a multigrid V-cycle in terms of data flow, as noted by Cheng et
al. (2021) and Azulay & Treister (2023). Cheng et al. (2021) use the multi-
scale network architecture MSNet to approximate the solution of Poisson
equations arising in plasma flow problems. However, they only consider flows
over a square domain in 2D. Azulay & Treister (2023) note the similarity
between the multi-scale UNet architecture and a multigrid V-cycle. They use
this structure to learn preconditioners for the solution of heterogeneous
Helmholtz equations. Eliasof et al. (2023) also use a multigrid-like
architecture for a general class of problems. Huang et al. (2023) use deep
learning to generate multigrid smoothers at each grid resolution that
effectively smooth high frequencies: CNNs generate the smoothing stencils from
matrix entries at each level in the multigrid hierarchy. This is similar to
our boundary-condition-dependent stencils, however we note that our network is
lighter-weight and allowed to vary at a larger scale during learning.
Furthermore, optimal stencils are known for the problems considered in this
work, and we provide evidence that our solvers outperforms them.

##  3 Motivation: Incompressible Fluids With Mixed B.C.s

While our solver architecture can be applied to any Poisson equation
discretized on a structured grid, our original motivation was to accelerate a
popular method for incompressible inviscid fluid simulation based on the
splitting scheme introduced by Chorin (1967). The fluid’s velocity ${\bf
u}({\bf x},t)$ is governed by the incompressible Euler equations:

$\rho\left(\frac{\partial{\bf u}}{\partial t}+({\bf u}\cdot\nabla){\bf
u}\right)+\nabla p={\bf f}^{\text{ext}}\quad\text{s.t.}\quad\nabla\cdot{\bf
u}=0\quad\text{in }\Omega,$ (1)

where $\Omega$ is the domain occupied by fluid, pressure $p$ is the Lagrange
multiplier for the incompressibility constraint $\nabla\cdot{\bf u}=0$, $\rho$
is the mass density, and ${\bf f}^{\text{ext}}$ accounts for external forces
like gravity. These equations are augmented with initial conditions ${{\bf
u}({\bf x},0)={\bf u}^{0}({\bf x})}$ and ${\rho({\bf x},0)=\rho^{0}}$ as well
as the boundary conditions discussed in Section 3.1. Incompressibility implies
that the initial homogeneous mass density is conserved throughout the
simulation ($\rho\equiv\rho^{0}$).

Chorin’s scheme employs finite differences in time and splits the integration
from time $t^{n}$ to ${t^{n+1}=t^{n}+\Delta t}$ into two steps. First, a
provisional velocity field ${\bf u}^{*}$ is obtained by an _advection step_
that neglects the pressure and incompressibility constraint:

$\frac{{\bf u}^{*}-{\bf u}^{n}}{\Delta t}+({\bf u}^{n}\cdot\nabla){\bf
u}^{n}=\frac{1}{\rho^{0}}{\bf f}^{\text{ext}}.$ (2)

Second, a _projection step_ obtains ${\bf u}^{n+1}$ by eliminating divergence
from ${\bf u}^{*}$:

$\displaystyle-\nabla\cdot\frac{1}{\rho^{0}}\nabla p^{n+1}$
$\displaystyle=-\frac{1}{\Delta t}\nabla\cdot{\bf u}^{*},$ (3)
$\displaystyle\frac{{\bf u}^{n+1}-{\bf u}^{*}}{\Delta t}$
$\displaystyle=-\frac{1}{\rho^{0}}\nabla p^{n+1}.$ (4)

Equations 2-4 hold inside $\Omega$, and we have deferred discussion of
boundary conditions to Section 3.1. The bottleneck of this full process is
(3), which is a Poisson equation since $\rho^{0}$ is spatially constant.

###  3.1 Boundary Conditions

Our primary contribution is handling both Neumann and Dirichlet boundary
conditions for the Poisson equation. We assume the computational domain
$\mathcal{D}$ is decomposed into
$\mathcal{D}=\Omega\cup\Omega_{a}\cup\Omega_{s}$, as sketched in the inset,
where $\Omega_{a}$ denotes free space and $\Omega_{s}$ the region filled with
solid. This decomposition induces a partition of the fluid boundary
$\partial\Omega=\Gamma_{n}\cup\Gamma_{d}$. Boundary $\Gamma_{n}$ represents
the fluid-solid interface as well as the intersection
$\partial\Omega\cap\partial\mathcal{D}$ (_i.e._ , the region outside
$\mathcal{D}$ is treated as solid); on it a free-slip boundary condition is
imposed: (1), ${{\bf u}({\bf x},t)\cdot\hat{{\bf n}}({\bf
x})=u^{\Gamma}_{n}({\bf x},t)}$, where $\hat{{\bf n}}$ denotes the outward-
pointing unit surface normal. This condition on ${\bf u}$ translates via (4)
into a Neumann condition on (3):

$\hat{{\bf n}}\cdot\nabla p^{n+1}=\frac{\rho_{0}}{\Delta t}(\hat{{\bf
n}}\cdot{\bf u}^{*}-u^{\Gamma}_{n})\quad\text{on }\Gamma_{n}.$ (5)

Free-surface boundary $\Gamma_{d}$ represents the interface between the fluid
and free space. Ambient pressure $p_{a}$ then imposes on (3) a Dirichlet
condition $p^{n+1}=p_{a}\text{ on }\Gamma_{d}$. In our examples, we set
$p_{a}=0$.

The Dirichlet conditions turn out to make solving (3) fundamentally more
difficult: while the DCDM paper Kaneda et al. (2023) discovered that a
preconditioner blind to the domain geometry and trained solely on an empty box
is highly effective for simulations featuring pure Neumann conditions, the
same is not true for Dirichlet (see Figure 5).

###  3.2 Spatial Discretization

We discretize the full domain $\mathcal{D}$ using a regular marker-and-cell
(MAC) staggered grid with $n_{c}$ cubic elements Harlow (1964). The disjoint
subdomains $\Omega$, $\Omega_{a}$, and $\Omega_{s}$ are each represented by a
per-cell rasterized indicator field; these are collected into a 3-channel
image, stored as a tensor $\bm{\mathcal{I}}$. In the case of a 2D square with
$n_{c}=N^{2}$, this tensor is of shape $(3,N,N)$, and summing along the first
index yields a single-channel image filled with ones.

Velocities and forces are represented at the _corners_ of this grid, and for
smoke simulations the advection step (2) is implemented using an explicit
semi-Lagrangian method (Stam, 1999; Robert, 1981). For free-surface
simulations, advection is performed by interpolating fluid velocities from the
grid onto particles responsible for tracking the fluid state, advecting those
particles, and then transferring their velocities back to the grid. In our
examples, we use a PIC/FLIP blend transfer scheme with a 0.99 ratio (Zhu &
Bridson, 2005).

Pressure values are stored at element _centers_ , and the Laplace operator in
(3) is discretized into a sparse symmetric matrix
${A^{\bm{\mathcal{I}}}\in\mathbb{R}^{n_{c}\times n_{c}}}$ using the standard
second-order accurate finite difference stencil (with 5 points in 2D and 7 in
3D) but with modifications to account for Dirichlet and Neumann boundary
conditions: stencil points falling outside $\Omega$ are dropped, and the
central value (_i.e._ , the diagonal matrix entry) is determined as the number
of neighboring cells belonging to either $\Omega$ or $\Omega_{a}$. Examples of
these stencils are visualized in 2D in the inset. Rows and columns
corresponding to cells outside $\Omega$ are left empty, meaning
$A^{\bm{\mathcal{I}}}$ typically has a high-dimensional nullspace. These empty
rows and columns are removed before solving, obtaining a smaller positive
definite matrix ${\tilde{A}}^{\bm{\mathcal{I}}}\in\mathbb{R}^{n_{f}\times
n_{f}}$, where $n_{f}$ is the number of fluid cells.

The right-hand side of (3) is discretized using the standard MAC divergence
finite difference stencil into a vector ${\bf b}\in\mathbb{R}^{n_{c}}$, which
also receives contributions from the Neumann boundary. Entries of this vector
corresponding to cells outside $\Omega$ are removed to form right-hand side
vector $\tilde{\bf b}\in\mathbb{R}^{n_{f}}$ of the reduced linear system
representing the discrete Poisson equation:

${\tilde{A}}^{\bm{\mathcal{I}}}\tilde{\bf x}=\tilde{\bf b},$ (6)

where $\tilde{\bf x}\in\mathbb{R}^{n_{f}}$ collects the fluid cells’ unknown
pressure values (a discretization of $p^{n+1}$).

The constantly changing domains and boundary conditions of a typical fluid
simulation mean traditional preconditioners for (6) like multigrid or
incomplete Cholesky, as well as direct sparse Cholesky factorizations, need to
be _rebuilt at every frame_. This prevents their high fixed costs from being
amortized across frames and means they struggle to outperform a highly tuned
GPU implementation of unpreconditioned CG. This motivates our neural-
preconditioned solver which, after training, instantly adapts to arbitrary
subdomain shapes encoded in $\bm{\mathcal{I}}$.

##  4 Neural-preconditioned Steepest Descent with Orthogonalization

Our neural-preconditioned solver combines a carefully chosen iterative method
(Section 4.1) with a preconditioner based on a novel neural network
architecture (Section 4.2.1) inspired by multigrid.

###  4.1 Algorithm

For symmetric positive definite matrices $A$ (like the discrete Laplacian
${\tilde{A}}^{\bm{\mathcal{I}}}$ from (6)), the preconditioned conjugate
gradient (PCG) algorithm (Shewchuk, 1994) is by far the most efficient
iterative method for solving linear systems $A{\bf x}={\bf b}$ when an
effective preconditioner is available. Unfortunately, its convergence rate is
known to degrade when the preconditioner itself fails to be symmetric, as is
the case for our neural preconditioner. Bouwmeester et al. (2015) have shown
that good convergence can be recovered for nonsymmetric multigrid
preconditioners using the “flexible PCG” variant at the expense of an
additional dot product. However, this variant turns out to perform
suboptimally with our neural preconditioner, as shown in Table 1. Instead, we
adopt the preconditioned steepest descent with orthogonalization (PSDO) method
proposed in Kaneda et al. (2023), which was shown to perform well even for
their nonlinear preconditioning operator.

The PSDO algorithm can be understood as a modification of standard CG that
replaces the residual with the preconditioned residual as the starting point
for generating search directions and, consequently, cannot enjoy many of the
simplifications baked into the traditional algorithm. Most seriously,
$A$-orthogonalizing against only the previous search direction no longer
suffices to achieve $A$-orthogonality to all past steps. Therefore, iteration
$k$ of PSDO obtains its step direction ${\bf d}_{k}$ by explicitly
$A$-orthogonalizing the preconditioned residual against the last
$n_{\textrm{ortho}}$ directions (where $n_{\textrm{ortho}}$ is a tunable
parameter) and determines step length $\alpha_{k}$ with an exact line search.
PSDO reduces to standard preconditioned steepest descent (PSD) when
$n_{\textrm{ortho}}=0$, and it is mathematically equivalent to
unpreconditioned CG when $n_{\textrm{ortho}}\geq 1$ and the identity operator
is used as the preconditioner. In the case of a symmetric preconditioner
$P=LL^{\top}$, PSDO differs from PCG by taking steps that are $A$-orthogonal
rather than $L^{-1}AL^{-\top}$-orthogonal. When combined with our neural
preconditioner, we call this algorithm NPSDO, presented formally in Algorithm
1 in the appendix. We empirically determined $n_{\textrm{ortho}}=2$ to perform
well, and we use this value in all reported experiments.

###  4.2 Neural Preconditioner

The ideal preconditioner for all iterative methods described in Section 4.1 is
the exact inverse $A^{-1}$; with it, each method would converge to the exact
solution in a single step. Of course, the motivation for using an iterative
solver is that inverting or factorizing $A$ is too costly (Figure 6), and
instead we must seek an inexpensive approximation of $A^{-1}$. Examples are
incomplete Cholesky, which does its best to factorize $A$ with a limited
computational budget, and multigrid, which applies one or more iterations of a
multigrid solver.

Our method approximates the map ${\bf r}\mapsto A^{-1}{\bf r}$ by our neural
network $\bm{\mathcal{P}}^{\textrm{net}}(\bm{\mathcal{I}},{\bf r})$. Departing
from recent works like Kaneda et al. (2023), we use a novel architecture that
both substantially boosts performance on pure-Neumann problems and generalizes
to the broader class of Poisson equations with mixed boundary conditions by
considering geometric information from $\bm{\mathcal{I}}$. The network
performs well on 2D or 3D Poisson equations of varying sizes, but to simplify
the exposition, our figures and notation describe the method on small square
grids of size $N\times N$.

We note that Algorithm 1 runs on linear system
${\tilde{A}}^{\bm{\mathcal{I}}}\tilde{\bf x}=\tilde{\bf b}$, featuring vectors
of smaller size $n_{f}$, but the network always operates on input vectors of
full size $n_{c}$, reshaped into $(N,N)$ tensors. Therefore, to evaluate
$\tilde{\bf d}=\bm{\mathcal{P}}^{\textrm{net}}(\bm{\mathcal{I}},\tilde{\bf
r})$, $\tilde{\bf r}$ is first padded by inserting zeros into locations
corresponding to cells in $\Omega_{a}$ and $\Omega_{s}$, and then those
locations of the output are removed to obtain $\tilde{\bf
d}\in\mathbb{R}^{n_{f}}$.

Figure 1: Our network architecture sketched for a 2D preconditioner with
$\mathcal{L}=3$ levels.

####  4.2.1 Architecture

Our neural network architecture (Figure 1) is inspired by geometric multigrid,
aiming to propagate information across the computational grid faster than the
one-cell-per-iteration of unpreconditioned CG. The architecture is constructed
recursively, consisting of levels $1\leq\ell\leq\mathcal{L}$. A given level
$\ell$ operates on an input image $\bm{\mathcal{I}}^{(\ell)}$ and input vector
${\bf r}^{(\ell)}$. It performs a special image-dependent convolution
operation on ${\bf r}^{(\ell)}$ and then downsamples the resulting vector
${\bf y}^{(\ell)}$, as well as $\bm{\mathcal{I}}^{(\ell)}$, to the next-
coarser level $\ell+1$ using average pooling (analogous to restriction in
multigrid). The output of the level $\ell+1$ subnetwork is then upsampled
(analogous to prolongation), run through another convolution stage, and
finally linearly combined with ${\bf y}^{(\ell)}$ to obtain the output. At the
finest level, $\bm{\mathcal{I}}^{(1)}=\bm{\mathcal{I}}$ and ${\bf
r}^{(1)}={\bf r}$, while at the coarsest level only a single convolution
operation is performed.

One crucial difference between our network and existing neural solvers like
FluidNet (Tompson et al., 2017) is how geometric information from
$\bm{\mathcal{I}}$ is incorporated. Past architectures treat this geometric
data on the same footing as input tensor ${\bf r}$, _e.g._ feeding both into
standard multi-channel convolution blocks. However, we note that
$\bm{\mathcal{I}}$ determines the entries of $A^{\bm{\mathcal{I}}}$, and so if
the convolutions are to act analogously to the smoothing operations of
multigrid, really this geometry information should inform the weights of
convolutions applied to ${\bf r}$. This motivates our use of custom
convolutional blocks whose _spatially varying kernels_ depend on local
information from $\bm{\mathcal{I}}$.

Each custom convolutional block (at the right corner in Figure 1) at level
$\ell$ learns an affine map from a $3\times 3$ sliding window in
$\bm{\mathcal{I}}^{(\ell)}$ to a $3\times 3$ kernel
$\bm{\mathcal{K}}^{(i,j)}$. This affine map is parametrized by a weights
tensor $\bm{\mathcal{W}}$ of shape $(3^{2},3,3,3)$ and a bias vector
$\bm{\mathcal{B}}\in\mathbb{R}^{3^{2}}$. Entry $y_{i,j}$ of the block’s output
is computed as:

$\displaystyle y_{i,j}$
$\displaystyle=\sum_{a,b=-1}^{1}\mathcal{K}_{a,b}^{(i,j)}x_{i+a,j+b},\quad\quad\mathcal{K}_{a,b}^{(i,j)}\vcentcolon=\sum_{c=0}^{2}\sum_{l,m=-1}^{1}\mathcal{W}_{3a+b,c,l,m}\mathcal{I}_{c,i+l,j+m}^{(\ell)}+\mathcal{B}_{3a+b}.$

Out-of-bounds accesses in these formulas are avoided by padding
$\bm{\mathcal{I}}^{(\ell)}$ with solid pixels (_i.e._ , the values assigned to
cells in $\Omega_{s}$) and ${\bf x}$ with zeros.

In multigrid, the solutions obtained on the coarser grids of the hierarchy are
_corrections_ that are added to the finer grids’ solutions; likewise, our
network includes connections labeled “linear combination” in Figure 1 that mix
in upsampled data from the lower level. Our network determines each of the two
coefficients in this combination by learning affine functions of the input
image defined by (i) convolving $\bm{\mathcal{I}}^{(\ell)}$ with a (spatially
constant) kernel $\overline{\bm{\mathcal{K}}}$ of shape $(3,3,3)$; (ii)
averaging to produce a scalar; and (iii) adding a scalar bias
$\overline{\mathcal{B}}$. For efficiency, these evaluation steps are fused
into a custom linear block (indicated by blue arrows in Figure 1) that
implements the formula:

$\displaystyle
z=\overline{\mathcal{B}}+\frac{1}{3^{2}n_{c}}\sum_{i,j=0}^{N-1}\sum_{c=0}^{2}\sum_{l,m=-1}^{1}{\overline{\mathcal{K}}}_{c,l,m}\mathcal{I}_{c,i+l,j+m}^{(\ell)}.$

Our custom network architecture has numerous advantages. Its output is a
linear function of the input vector (unlike the nonlinear map learned by
Kaneda et al. (2023)), making it easier to interpret as a preconditioner. The
architecture is also very lightweight: a model with $\mathcal{L}=4$ coarsening
levels has only $\sim\\!25$k parameters. Its simplicity accelerates network
evaluations at solve time, critical to make NPSDO competitive with the state-
of-the-art solvers used in practice.

We note that our solver is fully matrix free, with
$\bm{\mathcal{P}}^{\textrm{net}}$ relying only on the image $\bm{\mathcal{I}}$
of the simulation scene to infer information about $A^{\bm{\mathcal{I}}}$.
Furthermore, since all network operations are formulated in terms of local
windows into $\bm{\mathcal{I}}$ and ${\bf r}$, it can train and run on
_problems of any size divisible by_ $2^{\mathcal{L}}$.

The 3D version of our architecture is a straightforward extension of the 2D
formulas above, simply using larger tensors with additional indices to account
for the extra dimension, as well as extending the sums to run over these
indices.

####  4.2.2 Training

We train our network $\bm{\mathcal{P}}^{\textrm{net}}$ to approximate
${A^{\bm{\mathcal{I}}}}\big{\backslash}{\bf b}$ when presented with image
$\bm{\mathcal{I}}$ and input vector ${\bf b}$. We calculate the loss for an
example $(\bm{\mathcal{I}},A^{\bm{\mathcal{I}}},{\bf r})$ from our training
dataset as the residual norm:

$\displaystyle Loss=\left\lVert{\bf
b}-A^{\bm{\mathcal{I}}}\bm{\mathcal{P}}^{\textrm{net}}(\bm{\mathcal{I}},{\bf
b})\right\rVert_{2}.$

We found the more involved loss function used in Kaneda et al. (2023) not to
benefit our network.

Our training data set consists of 183 matrices collected from 10 different
simulation scenes, some of domain shape (128, 128, 128) and others (256, 128,
128). For each matrix, we generate 800 right-hand side vectors using a similar
approach to Kaneda et al. (2023), but with far fewer Rayley-Ritz vectors. We
first compute 1600 Ritz vectors using Lanczos iterations (Lanczos, 1950) and
then generate from them 800 random linear combinations. These linear
combinations are finally normalized and added to the training set. To
accelerate data generation, we create the right-hand sides for different
matrices in parallel; it takes between 0.5 and 3 hours to generate the data
for each scene. Since Ritz vector calculation is expensive, we experimented
with other approaches, like picking random vectors or constructing analytical
eigenmodes for the Laplacian on $\mathcal{D}$ and masking out entries outside
$\Omega$. Unfortunately these cheaper generation techniques led to degraded
performance.

In each epoch of training, we loop over the matrices of our dataset in
shuffled order. For each matrix, we process all of its 800 right-hand sides in
batches of 128, repeating five times. The full training process takes 5-7 days
on an AMD EPYC 9554P 64-Core Processor with an NVIDIA RTX 6000 GPU. The
training and validation losses are computed every five epochs, and we found it
beneficial to terminate after 50 epochs.

Figure 2: Renderings of some benchmark scenes.

####  4.2.3 Implementation

We built our network using PyTorch (Paszke et al. (2019)), but implemented our
custom convolutional and linear blocks as custom CUDA extensions. The neural
network was trained using single precision floating point.

##  5 Results and Analysis

We evaluate the effectiveness and efficiency of our neural preconditioned
solver by comparing it to high-performance state-of-the-art implementations of
several baseline methods: unpreconditioned CG provided by the CuPy library
(Okuta et al., 2017), as well as CG preconditioned by the algebraic multigrid
(AMG) and incomplete Cholesky (IC) implementations from the AMGCL library
(Demidov, 2020). All of these baseline methods are accelerated by CUDA
backends running on the GPU, with the underlying IC implementation coming from
NVIDIA’s cuSparse library. Where appropriate, we also compared against past
neural preconditioners FluidNet (Tompson et al., 2017) and DCDM (Kaneda et
al., 2023). Finally, we included characteristic performance statistics of a
popular sparse Cholesky solver CHOLMOD (Chen et al., 2008). In all cases, our
method outperforms these baselines, often dramatically.

We executed all benchmarks on a workstation featuring an AMD Ryzen 9 5950X
16-Core Processor and an NVIDIA GeForce RTX 3080 GPU. We used as our
convergence criterion for all methods a reduction of the residual norm by a
factor of $10^{6}$, which is sufficiently accurate to eliminate visible
simulation artifacts. We evaluate our neural preconditioner in single
precision floating point but implement the rest of the NPSDO algorithm in
double precision for numerical stability.

We benchmarked on twelve simulation scenes with various
shapes—$(128,128,128)$, $(256,128,128)$, and $(256,256,256)$—each providing
200 linear systems to solve. For each solve, we recorded the number of
iterations and runtime taken by each solver. These performance statistics are
summarized visually in Figures 5-6 and in tabular form in Appendix A.2.

Figure 3(a) summarizes timings from all solves in our benchmark suite: for
each system, we divide the unpreconditioned CG solve time by the other
methods’ solve times to calculate their speedups and plot a histogram. We note
that our method significantly outperforms the others on a majority of solves:
ours is fastest on 94% of the linear systems, which account for 98% of our
total solve time.

Our improvements are more substantial on larger problems, (Figures 3(b) and
3(c)) for two reasons. First, condition numbers increase with size, impeding
solvers without effective preconditioners; this is seen clearly by comparing
results from two different resolutions (Figures 3(d) and 3(e)). Second, the
small matrices ${\tilde{A}}^{\bm{\mathcal{I}}}$ correspond to simulation grids
with mostly non-fluid cells. While CG, AMGCL and IC timings shrink
significantly as fluid cells are removed, our network’s evaluation cost does
not: it always processes all of $\mathcal{D}$ regardless of occupancy. This
scaling behavior is visible in Figure 6.

Our speedups are also greater for examples with $\Gamma_{d}=\emptyset$. DCDM
is applicable for these, and so we included in it Figure 3(d) (but not in
Figure 3(e) due to the network overspilling GPU RAM). DCDM’s failure to
outperform CG and IC in these results, contrary to (Kaneda et al., 2023), can
be attributed to the higher-performance CUDA-accelerated implementations of
those baselines used in this work. With Dirichlet conditions (Figure 3(f)),
our preconditioner is less effective, and yet we still outperform the rest on
91% of the frames, which account for 96% of our total solve time. Statistics
are not reported in this setting for DCDM and FluidNet, which struggle to
reduce the residual (Figure 5).

(a) All systems

(b) Smallest 25% of systems

(c) Largest 75% of systems

(d) Neumann only, $128^{3}$

(e) Neumann only, $256^{3}$

(f) Dirichlet present, $128^{3}$ and $256^{3}$

Figure 3: Histograms of solution speedup vs. a baseline of unpreconditioned CG
(a) for all solves; and (b-f) for certain subsets of the systems to help tease
apart the different modes of the distribution.

Figure 4: Comparisons among AMG, IC, CG and NSPDO (Ours) on a single frame at
$256^{3}$ with Neumann only BC (left two) and mixed BC (right two).

Figure 5: Comparisons among AMG, IC, CG, DCDM, FluidNet (FN) and NSPDO (Ours)
on a single frame at $128^{3}$ with Neumann only BC (left two) and mixed BC
(right two).

Figure 6: Solver scaling for mixed BC system matrices originating from a
fixed-resolution domain $(n_{c}=256^{3})$; matrix row/col size $n_{f}$ is
determined by the proportion of cells occupied by fluid. The vast majority of
total solve time is contributed by the high-occupancy systems clustered to the
right, where our method outperforms the rest.

Further insights can be obtained by consulting Figures 5 and 5, which show the
convergence behavior of each iterative solver on characteristic example
problems. AMG is clearly the most effective preconditioner, but this comes at
the high cost of rebuilding the multigrid hierarchy before each solve: its
iterations cannot even start until long after our solver already converged.
Our preconditioner is the second most effective and, due to its lightweight
architecture, achieves the fastest solves. DCDM is also quite effective at
preconditioning for Neumann-only problems, but its iterations are slowed by
costly network evaluations. IC’s setup time is shorter than AMG but still
substantial, and it is much less effective as a preconditioner.

We note that the smoke example (Figure 5) also includes a comparison to
FluidNet _applied as a preconditioner_ for PSDO. In the original paper,
FluidNet was presented as a standalone solver, to be run just once per
simulation frame. However, in this form it cannot produce highly accurate
solutions. Incorporating it as a preconditioner as we do here in theory allows
the system to be solved to controlled accuracy, but this solver ended up
stalling before reaching a $10^{6}$ reduction in our experiments; for this
reason it was omitted from Figure 5.

On average, our solver spends 81% of its time evaluating
$\bm{\mathcal{P}}^{\textrm{net}}$, 4% of its time in orthogonalization, and
the remaining 15% in other CG operations. In contrast, AMG takes a full 90% of
its time in its setup stage. IC’s quicker construction and slower convergence
mean it takes only 23% in setup. Our architecture also confers GPU memory
usage benefits: for $128^{3}$ grids, our solver uses 1.5GiB of RAM, while
FluidNet and DCDM consume 5GiB and 8.3GiB, respectively (Appendix A.2).

##  6 Conclusions

The neural-preconditioned solver we propose not only addresses more general
boundary conditions than past machine learning approaches for the Poisson
equation (Tompson et al., 2017; Kaneda et al., 2023) but also dramatically
outperforms these solvers. It even surpasses state-of-the art high-performance
implementations of standard methods like algebraic multigrid and incomplete
Cholesky. It achieves this through a combination of its strong efficacy as a
preconditioner and its fast evaluations enabled by our novel lightweight
architecture.

Nevertheless, we see several opportunities to improve and extend our solver in
future work. First, although we implemented our spatially-varying convolution
block in CUDA, it remains the computational bottleneck of the network
evaluation and is not yet fully optimized. We are also excited to try porting
our architecture to special-purpose acceleration hardware like Apple’s Neural
Engine; not only could this offer further speedups, but also it would free up
GPU cycles for rendering the results in real-time applications like visual
effects and games. Second, we would like to explore was to explicitely enforce
symmetry and even positive definiteness of our preconditioning operator so
that the less expensive PCG algorithm could be used rather than PSDO. Third,
for applications where fluid occupies only a small portion of the
computational domain, we would like to develop techniques to exploit sparsity
for better scaling (Figure 6). Finally, we look forward to extending our ideas
to achieve competitive performance for problems posed on unstructured grids as
well as equations with non-constant coefficients, vector-valued unknowns
(_e.g._ , elasticity), and nonlinearities.

## References

  * Ackmann et al. (2020) J. Ackmann, P. D. Düben, T. N. Palmer, and P. K. Smolarkiewicz.  Machine-learned preconditioners for linear solvers in geophysical fluid flows.  _arXiv preprint arXiv:2010.02866_ , 2020.
  * Azulay & Treister (2023) Y. Azulay and R. Treister.  Multigrid-augmented deep learning preconditioners for the helmholtz equation.  _SIAM Journal on Scientific Computing_ , 45(3):S127–S151, 2023.  doi: 10.1137/21M1433514.  URL https://doi.org/10.1137/21M1433514.
  * Bouwmeester et al. (2015) H. Bouwmeester, A. Dougherty, and A.V. Knyazev.  Nonsymmetric preconditioning for conjugate gradient and steepest descent methods.  _Procedia Computer Science_ , 51:276–285, 2015.  ISSN 1877-0509.  doi: https://doi.org/10.1016/j.procs.2015.05.241.  URL https://www.sciencedirect.com/science/article/pii/S1877050915010492.  International Conference On Computational Science, ICCS 2015.
  * Brandt (1977) A. Brandt.  Multi-level adaptive solutions to boundary-value problems.  _Math Comp_ , 31(138):333–390, 1977.
  * Chen et al. (2008) Y. Chen, T.A. Davis, W.W. Hager, and S. Rajamanickam.  Algorithm 887: Cholmod, supernodal sparse cholesky factorization and update/downdate.  _ACM Trans. Math. Softw._ , 35(3), oct 2008.  ISSN 0098-3500.  doi: 10.1145/1391989.1391995.  URL https://doi.org/10.1145/1391989.1391995.
  * Cheng et al. (2021) L. Cheng, E.A. Illarramendi, G. Bogopolsky, M. Bauerheim, and B. Cuenot.  Using neural networks to solve the 2d poisson equation for electric field computation in plasma fluid simulations.  _arXiv preprint arXiv:2109.13076_ , 2021.
  * Chorin (1967) A. Chorin.  A numerical method for solving incompressible viscous flow problems.  _J Comp Phys_ , 2(1):12–26, 1967.
  * Cooley & Tukey (1965) J. Cooley and J. Tukey.  An algorithm for the machine calculation of complex fourier series.  _Math Comp_ , 19(90):297–301, 1965.
  * Demidov (2020) D. Demidov.  Amgcl —a c++ library for efficient solution of large sparse linear systems.  _Software Impacts_ , 6:100037, 2020.  ISSN 2665-9638.  doi: https://doi.org/10.1016/j.simpa.2020.100037.  URL https://www.sciencedirect.com/science/article/pii/S2665963820300282.
  * Eliasof et al. (2023) M. Eliasof, J. Ephrath, L. Ruthotto, and E. Treister.  Mgic: Multigrid-in-channels neural network architectures.  _SIAM Journal on Scientific Computing_ , 45(3):S307–S328, 2023.  doi: 10.1137/21M1430194.  URL https://doi.org/10.1137/21M1430194.
  * Golub & Loan (2012) G. Golub and C. Van Loan.  _Matrix computations_ , volume 3.  JHU Press, 2012.
  * Golub & Ye (1999) G. Golub and Q. Ye.  Inexact preconditioned conjugate gradient method with inner-outer iteration.  _SIAM J Sci Comp_ , 21(4):1305–1320, 1999.  doi: 10.1137/S1064827597323415.
  * Götz & Anzt (2018) M. Götz and H. Anzt.  Machine learning-aided numerical linear algebra: Convolutional neural networks for the efficient preconditioner generation.  In _2018 IEEE/ACM 9th Workshop on Latest Advances in Scalable Algorithms for Large-Scale Systems (scalA)_ , pp. 49–56, 2018.  doi: 10.1109/ScalA.2018.00010.
  * Harlow (1964) F. Harlow.  The particle-in-cell method for numerical solution of problems in fluid dynamics.  _Meth Comp Phys_ , 3:319–343, 1964.
  * Huang et al. (2023) R. Huang, R. Li, and Y. Xi.  Learning optimal multigrid smoothers via neural networks.  _SIAM Journal on Scientific Computing_ , 45(3):S199–S225, 2023.  doi: 10.1137/21M1430030.  URL https://doi.org/10.1137/21M1430030.
  * Kaneda et al. (2023) A. Kaneda, O. Akar, J. Chen, V.A.T. Kala, D. Hyde, and J. Teran.  A deep conjugate direction method for iteratively solving linear systems.  In Andreas Krause, Emma Brunskill, Kyunghyun Cho, Barbara Engelhardt, Sivan Sabato, and Jonathan Scarlett (eds.), _Proceedings of the 40th International Conference on Machine Learning_ , volume 202 of _Proceedings of Machine Learning Research_ , pp. 15720–15736. PMLR, 23–29 Jul 2023.  URL https://proceedings.mlr.press/v202/kaneda23a.html.
  * Lanczos (1950) C. Lanczos.  An iteration method for the solution of the eigenvalue problem of linear differential and integral operators.  1950\.
  * Li et al. (2023) Y. Li, P.Y. Chen, T. Du, and W. Matusik.  Learning preconditioners for conjugate gradient PDE solvers.  In Andreas Krause, Emma Brunskill, Kyunghyun Cho, Barbara Engelhardt, Sivan Sabato, and Jonathan Scarlett (eds.), _Proceedings of the 40th International Conference on Machine Learning_ , volume 202 of _Proceedings of Machine Learning Research_ , pp. 19425–19439. PMLR, 23–29 Jul 2023.  URL https://proceedings.mlr.press/v202/li23e.html.
  * Mathieu et al. (2016) M. Mathieu, C. Couprie, and Y. LeCun.  Deep multi-scale video prediction beyond mean square error, 2016.
  * Okuta et al. (2017) R. Okuta, Y. Unno, D. Nishino, S. Hido, and C. Loomis.  Cupy: A numpy-compatible library for nvidia gpu calculations.  In _Proceedings of Workshop on Machine Learning Systems (LearningSys) in The Thirty-first Annual Conference on Neural Information Processing Systems (NIPS)_ , 2017.  URL http://learningsys.org/nips17/assets/papers/paper_16.pdf.
  * Paszke et al. (2019) A. Paszke, S. Gross, F. Massa, A. Lerer, J. Bradbury, G. Chanan, T. Killeen, Z. Lin, N. Gimelshein, L. Antiga, A. Desmaison, A. Kopf, E. Yang, Z. DeVito, M. Raison, A. Tejani, S. Chilamkurthy, B. Steiner, L. Fang, J. Bai, and S. Chintala.  Pytorch: An imperative style, high-performance deep learning library.  In H. Wallach, H. Larochelle, A. Beygelzimer, F. d'Alché-Buc, E. Fox, and R. Garnett (eds.), _Advances in Neural Information Processing Systems 32_ , pp. 8024–8035. Curran Associates, Inc., 2019.  URL http://papers.neurips.cc/paper/9015-pytorch-an-imperative-style-high-performance-deep-learning-library.pdf.
  * Robert (1981) A. Robert.  A stable numerical integration scheme for the primitive meteorological equations.  _Atm Ocean_ , 19(1):35–46, 1981.
  * Ronneberger et al. (2015) O. Ronneberger, P. Fischer, and T. Brox.  U-net: Convolutional networks for biomedical image segmentation, 2015\.
  * Saad (2003) Y. Saad.  _Iterative Methods for Sparse Linear Systems_.  Society for Industrial and Applied Mathematics, USA, 2nd edition, 2003\.  ISBN 0898715342.
  * Sappl et al. (2019) J. Sappl, L. Seiler, M. Harders, and W. Rauch.  Deep learning of preconditioners for conjugate gradient solvers in urban water related problems, 2019.  URL https://arxiv.org/abs/1906.06925.
  * Shewchuk (1994) J.R. Shewchuk.  An introduction to the conjugate gradient method without the agonizing pain.  Technical report, USA, 1994.
  * Stam (1999) J. Stam.  Stable fluids.  In _Siggraph_ , volume 99, pp. 121–128, 1999.
  * Tompson et al. (2017) J. Tompson, K. Schlachter, P. Sprechmann, and K. Perlin.  Accelerating Eulerian fluid simulation with convolutional networks.  In D. Precup and Y. Teh (eds.), _Proc 34th Int Conf Mach Learn_ , volume 70 of _Proc Mach Learn Res_ , pp. 3424–3433. PMLR, 06–11 Aug 2017\.
  * Zhu & Bridson (2005) Y. Zhu and R. Bridson.  Animating sand as a fluid.  _ACM Trans Graph_ , 24(3):965–972, 2005.
  * Özbay et al. (2021) A.G. Özbay, A. Hamzehloo, S. Laizet, P. Tzirakis, G. Rizos, and B. Schuller.  Poisson cnn: Convolutional neural networks for the solution of the poisson equation on a cartesian mesh.  _Data-Centric Engineering_ , 2:e6, 2021.  doi: 10.1017/dce.2021.7.

##  Appendix A Appendix

Table 1: Performance statistics for several solver variants averaged over a 30 frame simulation. PSD is implemented by Algorithm 1 with $n_{\textrm{ortho}}=0$. Each solver except unpreconditioned CG (left) was limited to 200 iterations per frame. | CG | PCG | Flexible PCG | PSD | PSDO
---|---|---|---|---|---
Iterations | 457.93 | 29.6667 | 24.6333 | 82.0333 | 20.7333
Time | 0.5253 | 0.2295 | 0.1911 | 0.6299 | 0.1614

###  A.1 Algorithm

For completeness, we provide the following pseudocode for the NPSDO algorithm,
which is based on the same iterative solver as DCDM Kaneda et al. (2023) but
with our new geometry-aware preconditioner $\bm{\mathcal{P}}^{\textrm{net}}$.

Algorithm 1 Neural-preconditioned Steepest Descent with $A$-Orthogonalization
(NPSDO)

Given matrix $A$, right-hand side $\mathbf{b}$, image $\bm{\mathcal{I}}$, and
trained network $\bm{\mathcal{P}}^{\textrm{net}}$

${\bf r}_{0}\leftarrow{\bf b}-A{\bf x}_{0}$

$k\leftarrow 0$

while $\left\lVert{\bf r}_{k}\right\rVert\geq\epsilon$ do

$k\leftarrow k+1$

${\bf
d}_{k}\leftarrow\bm{\mathcal{P}}^{\textrm{net}}\left(\bm{\mathcal{I}},\frac{{\bf
r}_{k-1}}{\left\lVert{\bf r}_{k-1}\right\rVert}\right)$

for $k-n_{\textrm{ortho}}\leq i<k$ do

$\mathbf{d}_{k}\leftarrow\mathbf{d}_{k}-\frac{{\bf d}_{k}^{\top}A{\bf
d}_{i}}{{\bf d}_{i}^{\top}A{\bf d}_{i}}\mathbf{d}_{i}$

end for

$\alpha_{k}\leftarrow\frac{{\bf r}_{k-1}^{\top}{\bf d}_{k}}{{\bf
d}_{k}^{\top}A{\bf d}_{k}}$

${\bf x}_{k}\leftarrow{\bf x}_{k-1}+\alpha_{k}{\bf d}_{k}$

${\bf r}_{k}\leftarrow{\bf b}-A{\bf x}_{k}$

end while

return ${\bf x}_{k}$

###  A.2 Benchamrking results

This section lists the average iteration count and runtime among all the test
cases.

Table 2: Average iteration count and runtime among all frames for each test suite | AMG | IC | CG | Ours
---|---|---|---|---
Examples | Iteration | Time | Iteration | Time | Iteration | Time | Iteration | Time
Smoke solid $128^{3}$ | 11.2 | 0.696 | 196.1 | 0.852 | 612.5 | 1.050 | 20.4 | 0.124
Smoke solid $256^{3}$ | 14.3 | 6.433 | 343.7 | 11.1 | 1076.8 | 13.245 | 34.8 | 1.187
Smoke bunny $128^{3}$ | 11.3 | 0.713 | 200.5 | 0.862 | 615.8 | 1.048 | 20.0 | 0.122
Smoke bunny $256^{3}$ | 14.3 | 6.046 | 345.7 | 10.9 | 1069.7 | 13.042 | 33.3 | 1.130
Scooping $128^{3}$ | 12.2 | 0.170 | 128.1 | 0.270 | 392.8 | 0.234 | 14.6 | 0.073
Scooping $256^{3}$ | 12.2 | 1.643 | 246.0 | 1.968 | 749.2 | 2.593 | 23.2 | 0.569
Waterflow torus $128^{3}$ | 9.4 | 0.086 | 66.9 | 0.129 | 208.6 | 0.095 | 11.1 | 0.053
Waterflow torus $256^{3}$ | 10.4 | 1.024 | 130.0 | 0.946 | 401.2 | 1.121 | 16.3 | 0.382
Waterflow ball $128^{3}$ | 11.2 | 0.173 | 136.5 | 0.275 | 405.4 | 0.257 | 17.7 | 0.088
Waterflow ball $256^{3}$ | 12.1 | 2.442 | 328.4 | 3.559 | 969.7 | 4.875 | 37.2 | 0.967
Dambreak pillars $256\cdot 128^{2}$ | 11.2 | 0.476 | 167.7 | 0.625 | 523.2 | 0.704 | 18.0 | 0.144
Dambreak bunny $256\cdot 128^{2}$ | 11.1 | 0.395 | 148.2 | 0.514 | 452.1 | 0.522 | 20.4 | 0.159
Average | 11.7 | 1.691 | 203.2 | 2.667 | 623.1 | 3.232 | 22.3 | 0.417

Peak memory usage was recorded with the command `nvidia-smi`.

Table 3: Peak memory usage on a smoke example at $128^{3}$ and $256^{3}$ Method | AMG | IC | CG | NPSDO | DCDM | FN
---|---|---|---|---|---|---
$128^{3}$ | 1248 MiB | 1668 MiB | 1418 MiB | 1548 MiB | 8532 MiB | 5170 MiB
$256^{3}$ | 5032 MiB | 8214 MiB | 3716 MiB | 4776 MiB | NA | NA


# Solving Image PDEs with a Shallow Network

Pascal Getreuer Google Research, Mountain View CA, USA Peyman
Milanfar††footnotemark:  Xiyang Luo††footnotemark:

###### Abstract

Partial differential equations (PDEs) are typically used as models of physical
processes but are also of great interest in PDE-based image processing.
However, when it comes to their use in imaging, conventional numerical methods
for solving PDEs tend to require very fine grid resolution for stability, and
as a result have impractically high computational cost. This work applies
BLADE (Best Linear Adaptive Enhancement), a shallow learnable filtering
framework, to PDE solving, and shows that the resulting approach is efficient
and accurate, operating more reliably at coarse grid resolutions than
classical methods. As such, the model can be flexibly used for a wide variety
of problems in imaging.

###### keywords:

Filter learning, machine learning, partial differential equations, anisotropic
diffusion.

##  1 Introduction

Classically, finite element numerical schemes for solving PDEs are derived
from local Taylor series analysis, essentially approximating the solution to
be locally a polynomial (see for instance [16]). An important requirement is
that the scheme is stable, meaning that the approximation error accumulated
over multiple successive time steps stays under control instead of exploding.
Classical schemes tend to require that both the time step and grid step are
very small in order to ensure stability. There are numerous techniques to do
it better, and accuracy and stability requirements depend a lot on the
particular PDE at hand, but generally the trade-off between stability and step
size is the key difficulty.

BLADE is a trainable adaptive filtering framework that is simple, fast, and
useful for a wide range of imaging problems. In contrast to classic finite
difference solvers, BLADE is a generic data-driven method. A training set is
first formed from a set of input images and corresponding target outputs,
which can be computed once by a more computationally intensive reference
method. BLADE is then trained to approximate the input-target relationship. We
previously showed [12] that BLADE is capable of approximating a variety of
operators. In this work we show that compelling quality and efficiency is
broadly possible on PDEs.

Previous work. There are two camps of related literature at the intersection
of PDE and ML. The first is incorporating ML-models as a part of a PDE solver,
which is directly relevant to this paper. The second is to use PDE methods in
designing deep neural nets.

The interplay between PDEs and machine learning has lead to many interesting
works [24, 30, 23, 6]. On the one hand, classical numerical schemes for PDEs
have guided and provided more interpretability for the success of certain
neural network architectures. For example, several works have formulated
ResNets as a discretization of a certain differential equation [29, 18], which
allowed techniques for improving stability from classical numerical schemes to
be applied to deep networks as well [23, 5]. Chen et al. [6] formalized this
connection by connecting depth to the time horizon of a dynamical system. On
the other hand, the representational power of deep neural networks (DNN) makes
it an excellent candidate for replacing certain components of a PDE solver.
For example Weinan et al. used DNNs for general variational problems [30] as
well as for high-dimensional parabolic PDEs [13]. Such an approach has also
been successfully applied to time-dependent PDEs, such as wave propagation
[24], weather forecasting [20], as well as generic non-linear evolution PDEs
[17].

In contrast to most existing methods which rely on the use of large deep
neural nets, our method uses only a shallow network with about 50K parameters,
making it an extremely efficient solution while maintaining a surprisingly
high fidelity. This is made possible by the use of BLADE, an efficient
trainable and adaptive filtering framework that is applicable to a wide range
of imaging applications.

We are interested particularly in time-dependent PDEs appearing in image
processing, such as total variation (TV) flow,

(1)
$\left\\{\begin{array}[]{l}\partial_{t}u=\operatorname{div}\bigl{(}\frac{\nabla
u}{|\nabla u|}\bigr{)},\\\ u(x,y,0)=f(x,y).\end{array}\right.$

Above, $u(x,y,t)$ is the image, $(x,y)$ are the spatial coordinates, and $t$
is the evolution time. Such PDEs appear routinely in variational methods as
the gradient flow for the minimization of some objective [2].

Consider an equation of the form

(2) $\partial_{t}u=F(u,\partial_{x}u,\partial_{y}u,\ldots),$

where $F$ is a function of $u$, its gradient, and possibly higher spatial
derivatives. Let $\mathbf{u}$ denote the image in discrete space with pixels
$u_{m,n}$. A straightforward conventional approach to solving (2) is to use
finite differences to approximate spatial derivatives and apply explicit Euler
to integrate in time,

(3)
$\mathbf{u}^{(k+1)}=\mathbf{u}^{(k)}+\mathit{dt}F(\mathbf{u}^{(k)},D_{x}\mathbf{u}^{(k)},D_{y}\mathbf{u}^{(k)},\ldots).$

The finite differences ($D_{x},D_{y}$) are usually chosen with a localized
stencil, e.g. the forward finite difference with a $2\times 1$ footprint:

(4)
$(D_{x}\mathbf{u})_{m,n}=\tfrac{1}{\mathit{dx}}(u_{m+1,n}-u_{m,n})=\partial_{x}u+O(\mathit{dx}).$

Higher-order-accurate finite differences have limited use in image processing
as their use assumes a several-times continuously differentiable function
(required to expand higher terms of the Taylor series). However, natural
images have relatively low regularity as they contain edges and texture. As a
result, low-order, localized differences are most appropriate and common
practice.

On the other hand, this means that each pixel of $\mathbf{u}^{(k+1)}$ has a
small, localized domain of dependence in $\mathbf{u}^{(k)}$. So information
propagates slowly, and a small $\mathit{dt}$ is necessary for reliable
solution (the Courant–Friedrichs–Lewy condition [9]).

Additionally, although numerically convenient, it is well known that explicit
Euler becomes unstable if $\mathit{dt}$ is too large. This problem tends to be
more severe if $F$ includes higher-order spatial derivatives.

We train BLADE to approximate $M$ timesteps of a reference solution. To get
high-quality training data, we apply a reference method, then subsample both
temporally and spatially. With this approach, BLADE is trained to capture
longer time dependency, addressing the small timestep issue. Additionally,
since the training data is created at higher spatial resolution and then
coarsened, BLADE is trained to super-resolve spatial derivatives in the
equation.

$\operatorname{BLADE}$$\vphantom{x}\smash{\mathbf{z}}$$\vphantom{x}\smash{\mathbf{\hat{u}}}$$s(i)$
$\mathbf{R}_{i}\mathbf{z}$ input patch
$\mathbf{h}^{0}$$\mathbf{h}^{1}$$\mathbf{h}^{2}$ $\hat{\mathbf{u}}_{i}$ output
pixel

Figure 1: BLADE, a shallow 2-layer structure, taking image $\mathbf{z}$ as
input and producing image $\hat{\mathbf{u}}$ as output.

###  1.1 Notation

We denote the PDE solution by $u(x,y,t)$, a function of continuous space and
time, where $(x,y)\in\Omega\subset\mathbb{R}^{2}$ is the spatial domain
(usually a bounded rectangle). We follow the conventional notations
$\partial_{t}u$ for partial derivative of $u$ with respect to $t$, $\nabla$
for spatial gradient, $\operatorname{div}$ for spatial divergence, and
$\Delta:=\operatorname{div}\nabla$ for spatial Laplacian.

In numerical implementation, we consider an $M\times N$ discrete image as a
vector $\mathbf{f}\in\mathbb{R}^{M\cdot N}$. Subscripts $f_{m,n}$ or $f_{i}$,
$i=(m,n)\in\mathbb{Z}^{2}$, denote the sample at $x=m\mathit{dx}$,
$y=n\mathit{dx}$. We denote by $\mathbf{u}^{(k)}$ the discrete solution at
time $t=k\mathit{dt}$.

##  2 BLADE

Best linear adaptive enhancement (BLADE) is a trainable adaptive filtering
framework that is simple, fast, and useful for a wide range of imaging
problems. It is the generalization of the Rapid and Accurate Image Super-
Resolution (RAISR) method [21] to tasks other than image upscaling.

The BLADE network takes an input image $\mathbf{z}$ as input and computes an
output image $\hat{\mathbf{u}}$. The network has a set of learnable, locally
linear filters with small footprint: $\mathbf{h}^{0},\mathbf{h}^{1},\ldots$
and a filter selection mechanism $s(i)$ that decides for each output pixel
which filter to apply. The $i$th output pixel is computed by applying filter
$\mathbf{h}^{s(i)}$:

(5) $\hat{u}_{i}=\sum_{j\in F}h_{j}^{s(i)}z_{i+j},$

where $F\subset\mathbb{Z}^{2}$ denote the footprint of the filter. In
experiments, we will set $F$ to a $5\times 5$ filter footprint. To filter
pixels near the image boundaries, We extend $\mathbf{z}$ by replicating border
pixels. BLADE may be seen as a type of spatially-varying convolution. With a
spatially constant selection $s(i)\equiv s$, (5) is the cross-correlation of
$\mathbf{h}^{s}$ with $\mathbf{z}$.

Another way to express (5) is as extracting an input patch and computing its
dot product with (only) one selected filter. Denoting patch extraction by
$(\mathbf{R}_{i}\mathbf{z})_{j}=z_{i+j}$, $j\in F$, the $i$th output pixel is

(6) $\hat{u}_{i}=(\mathbf{h}^{s(i)})^{T}\mathbf{R}_{i}\mathbf{z}.$

BLADE can be seen as a shallow two-layer structure, where the first layer
selects the filter, and the second layer applies the filter. Notably, only the
selected filter at each pixel is evaluated. Therefore, in contrast to
conventional convolutional layers, BLADE’s inference computation cost is
independent of the number of filters. BLADE is successfully used in consumer
products where the number of filters is in the hundreds to low thousands [21,
12, 7].

###  2.1 Filter selection

Input | Orientation | Strength | Coherence
---|---|---|---
|  |  |
Figure 2: Example of quantized structure tensor features used for filter
selection, distinguishing 24 orientations, 3 strength values, and 3 coherence
values.

An effective choice for the selection mechanism $s(i)$ is to use features of
the $2\times 2$ image structure tensor. Define the structure tensor

(7) $J(\nabla
u)=\left(\begin{array}[]{cc}(\partial_{x}u)^{2}&(\partial_{x}u)\,(\partial_{y}u)\\\
(\partial_{x}u)\,(\partial_{y}u)&(\partial_{y}u)^{2}\end{array}\right),$

which is a symmetric $2\times 2$ matrix at each pixel. The tensor is smoothed
with a Gaussian $G_{\rho}$ of standard deviation $\rho$,

(8) $J_{\rho}(\nabla u)=G_{\rho}*J(\nabla u),$

where convolution is applied spatially to each component. At each pixel, we
compute the eigenvalues $\lambda_{1}\geq\lambda_{2}$ and dominant eigenvector
$\mathbf{w}_{1}$ of $J_{\rho}(\nabla u)$ and use them to define three
features:

  * •

orientation $=\angle\mathbf{w}_{1}$, gradient orientation,

  * •

strength $=\sqrt{\lambda_{1}}$, gradient magnitude,

  * •

coherence
$=\frac{\sqrt{\lambda_{1}}-\sqrt{\lambda_{2}}}{\sqrt{\lambda_{1}}+\sqrt{\lambda_{2}}}$,
local anisotropy.

Theses features are quantized to a small number of possible values. Figure 2
shows these three quantized features for an example image. We then consider
the quantized values as a three-dimensional index into the bank of filters. A
typical configuration is 24 possible orientations, 3 strengths, and 3
coherences for a total of $24\cdot 3\cdot 3=216$ filters, and flattened to a
filter index $\\{0,\ldots,215\\}$ as

(9)
$s(i)=\mathit{ori}(i)+24\cdot\bigl{(}\mathit{str}(i)+3\cdot\mathit{coh}(i)\bigr{)}.$

These structure tensor features enable BLADE to perform robust edge-adaptive
filtering. Other features can be used for selection. When approximating the
Cahn–Hilliard equation for instance, we will use the input intensity as an
additional selection feature.

###  2.2 BLADE for PDEs

We now develop how BLADE can be applied to the solution of hyperbolic PDEs,
for instance the form (2),

$\partial_{t}u=F(u,\partial_{x}u,\partial_{y}u,\ldots).$

$\operatorname{BLADE}$$\smash{\hat{\partial_{t}u}}\vphantom{x}$$\mathit{dt}$+$\smash{\hat{\mathbf{u}}^{(k)}}\vphantom{x}$$\smash{\hat{\mathbf{u}}^{(k+1)}}\vphantom{x}$

Figure 3: BLADE sequence model based on explicit Euler time integration.

Our approach is to estimate the time derivative $\partial_{t}u$ from
$\mathbf{u}$; that is, $\operatorname{BLADE}(\mathbf{u})=\hat{\partial_{t}u}$,
where $\operatorname{BLADE}(\mathbf{u})$ denotes the application of BLADE to
$\mathbf{u}$ as in equation (5) and $\hat{\cdot}$ denotes that the quantity is
an estimate. We then use explicit Euler integration to advance to the next
time step,

(10)
$\hat{\mathbf{u}}^{(k+1)}=\hat{\mathbf{u}}^{(k)}+\mathit{dt}\,\operatorname{BLADE}(\hat{\mathbf{u}}^{(k)}).$

The above equation is initialized with
$\hat{\mathbf{u}}^{(0)}=\mathbf{u}^{(0)}$ and applied repeatedly to predict
the evolution $\mathbf{u}^{(1)},\mathbf{u}^{(2)},\ldots$ over multiple time
steps.

$\smash{\mathbf{u}^{(0)}}\vphantom{x}$$\smash{\hat{\mathbf{u}}^{(N)}}\vphantom{x}$$\operatorname{BLADE}$$\mathit{dt}$+$\operatorname{BLADE}$$\mathit{dt}$+$\operatorname{BLADE}$$\mathit{dt}$+

Figure 4: BLADE estimation of a stopping time $T=N\mathit{dt}$ viewed as an
$N$-layer deep network with weights shared across layers.

The structure of (10) may be seen as a spatially-varying convolutional layer
with a skip connection, forming what is called a “residual” block in the
learning literature [14] (Figure 3). Its repeated application may be seen as a
deeper model composed of relatively simple blocks (Figure 4). If we are
interested only in some stopping time $T=N\mathit{dt}$, the model could be
considered as an $N$-layer deep network with weight sharing across layers that
produces $u^{(N)}$ as the final output.

$\operatorname{BLADE}$$\tfrac{\mathit{dt}}{2}$+$\smash{\mathbf{u}^{(k)}}\vphantom{x}$$\operatorname{BLADE}$$\mathit{dt}$+$\smash{\hat{\mathbf{u}}^{(k+1)}}\vphantom{x}$

Figure 5: BLADE with explicit midpoint time integration.

In the time dimension, explicit Euler integration could be substituted with a
higher-order method at the cost of more network evaluations per time step. For
instance explicit midpoint integration could be done as

(11) $\displaystyle\hat{\mathbf{u}}^{(k+1)}$
$\displaystyle=\hat{\mathbf{u}}^{(k)}+\mathit{dt}\,\operatorname{BLADE}\bigl{(}$
$\displaystyle\quad\hat{\mathbf{u}}^{(k)}+\tfrac{\mathit{dt}}{2}\,\operatorname{BLADE}(\hat{\mathbf{u}}^{(k)})\bigr{)}.$

This may be seen as a network structure with two skip connections (Figure 5).

###  2.3 BLADE training

Reference solver High-res initial condition
$\mathbf{u}^{(2)}$$\mathbf{u}^{(1)}$$\mathbf{u}^{(0)}$Subsample timeCoarsen
space

Figure 6: Preparing a target frame sequence for training.

Next we describe how we train the BLADE sequence model for approximating the
solution of the described anisotropic diffusion PDEs. To create high-quality
data for training, we perform the following steps, as illustrated in Figure 6:

  1. 1.

We begin with an initial condition that is, say, $4\times$ higher spatial
resolution than the operating resolution at which we intend to run BLADE. A
reference implementation of the PDE is executed on this initial condition with
fine time step $\mathit{dt}_{\mathit{HR}}$. This is such that quantization
both spatially and temporally is fine, so that the computed reference solution
is high quality.

  2. 2.

We subsample temporally, keeping every $M$th frame, to coarsen time resolution
to $\mathit{dt}=M\mathit{dt}_{\mathit{HR}}$.

  3. 3.

The frames are spatially downscaled using an area-averaging kernel, reducing
space to the target operating resolution. The resulting sequence of frames is
the training target.

We can generate target frame sequences of high quality in this manner,
possibly exceeding the quality of simply running the reference method directly
at the operating resolution.

In training, we unroll ten time steps of (10). Training examples are formed by
partitioning a target frame sequence into windows of 11 consecutive frames,
where the first frame is the input to the model, and the later 10 frames
$\mathbf{u}^{(k)}$, $k=1,\ldots,10$, are compared against model predictions
$\hat{\mathbf{u}}^{(k)}$ with a summed squared $L^{2}$ training loss:

(12)
$L(\hat{\mathbf{u}})=\sum_{k=1}^{10}\|\hat{\mathbf{u}}^{(k)}-\mathbf{u}^{(k)}\|_{2}^{2}.$

We remark that instead of $L^{2}$, other differentiable metrics could be used
as the loss without impact on cost at inference time. For PDEs on natural
images, mean structural similarity (SSIM) [26], difference in VGG activations
[15], or neural image assessment (NIMA) [25] could be used to optimize for
perceptual quality of the approximation.

We use TensorFlow [1] in this work to train BLADE. The BLADE filtering formula
(5) is implemented as an op, a function taking the filters
$(\mathbf{h}^{0},\mathbf{h}^{1},\ldots)$, selection $s$, and image
$\mathbf{z}$ as inputs and producing a filtered image $\hat{u}$ as output. We
backpropagate the training loss gradient $\tfrac{\partial
L}{\partial\hat{\mathbf{u}}}$ as

(13) $\left\\{\begin{aligned} \frac{\partial L}{\partial
h^{k}_{j}}&=\sum_{\mathchoice{\hbox
to0.0pt{\hss$\displaystyle{i:s(i)=k}$\hss}}{\hbox
to0.0pt{\hss$\textstyle{i:s(i)=k}$\hss}}{\hbox
to0.0pt{\hss$\scriptstyle{i:s(i)=k}$\hss}}{\hbox
to0.0pt{\hss$\scriptscriptstyle{i:s(i)=k}$\hss}}}z_{i+j}\frac{\partial
L}{\partial\hat{u}_{i}},\\\ \frac{\partial L}{\partial z_{i}}&=\sum_{j\in
F}h^{s(i-j)}_{j}\frac{\partial L}{\partial\hat{u}_{i+j}}.\end{aligned}\right.$

Note that (5) is discontinuous with respect to selection, due to the filter
lookup. Therefore, no $\tfrac{\partial L}{\partial s}$ gradient is
backpropagated. This is suboptimal, since training ignores how selection could
be changed (i.e. due to changes in preceding operations) to improve the
training objective. Nevertheless, we find in experiments that training rate
and convergence are satisfactory with gradient equation (13).

###  2.4 Properties

The proposed structure is quite flexible, even after training. It is
reasonable to run the time step prediction (10) with a smaller $\mathit{dt}$
than used in training if it is desired to estimate $u$ at a specific point in
time. Moreover, it is reasonable to add other terms to the right hand side of
(2) not seen during training, simply by adding them in the explicit Euler
step. We will show for example that BLADE trained to perform TV flow
$\partial_{t}u=\operatorname{div}(\nabla u/|\nabla u|)$ can be applied as

(14)
$\hat{\mathbf{u}}^{(k+1)}=\mathbf{u}^{(k)}+\mathit{dt}\bigl{(}\operatorname{BLADE}(\mathbf{u}^{(t)})+\lambda
A^{T}(u^{(0)}-Au)\bigr{)},$

where $A$ is a blurring operator to approximate

(15) $\partial_{t}u=\operatorname{div}\bigl{(}\tfrac{\nabla u}{|\nabla
u|}\bigr{)}+\lambda A^{T}(f-Au)$

whose steady state solution is TV-regularized deblurring. So once trained, a
BLADE network is potentially useful for multiple applications without
retraining.

We remark on several other attractive properties of this approach:

  * •

BLADE learns to approximate spatial derivatives in the equation. With high-
quality training examples created at finer resolution, BLADE’s approximation
is possibly superior to finite differences.

  * •

Training may help accuracy in the temporal dimension as well, since the
explicit Euler time integration is incorporated in the model and trained end-
to-end.

  * •

Training the network unrolled over multiple time steps encourages BLADE to be
stable, or in the very least to have controlled error growth over the unrolled
steps.

##  3 Anisotropic diffusion PDEs

In this section, we investigate using BLADE to approximate PDEs for several
image processing tasks: TV flow, Perona–Malik, coherence enhancing diffusion,
and Cahn–Hilliard.

#### Total variation flow

Total variation (TV) flow is the edge-preserving diffusion

(16) $\partial_{t}u=\operatorname{div}\Bigl{(}\frac{\nabla u}{|\nabla
u|}\Bigr{)}.$

TV flow arises from gradient descent of the TV seminorm $\int|\nabla u|\,dx$
[22]. This makes it of interest in implementing TV-regularized variational
methods. We use as reference implementation the explicit scheme suggested by
Rudin, Osher, and Fatemi [22]:

(17) $\displaystyle u_{m,n}^{(k+1)}=u^{(k)}_{m,n}+\mathit{dt}\Biggl{[}$
$\displaystyle
D_{-}^{x}\Biggl{(}\frac{D_{+}^{x}u_{m,n}^{(k)}}{\sqrt{(D_{+}^{x}u_{m,n}^{(k)})^{2}+\min(D_{+}^{y}u_{m,n}^{(k)},D_{-}^{y}u_{m,n}^{(k)})^{2}}}\Biggr{)}$
$\displaystyle+D_{-}^{y}\Biggl{(}\frac{D_{+}^{y}u_{m,n}^{(k)}}{\sqrt{(D_{+}^{y}u_{m,n}^{(k)})^{2}+\min(D_{+}^{x}u_{m,n}^{(k)},D_{-}^{x}u_{m,n}^{(k)})^{2}}}\Biggr{)}\Biggr{]}.$

where $D_{+}^{x}u_{m,n}:=u_{m+1,n}-u_{m,n}$ and
$D_{-}^{x}u_{m,n}:=u_{m,n}-u_{m-1,n}$ denote forward and backward finite
differences in the x direction, and similarly $D_{+}^{y}$ and $D_{-}^{y}$
denote finite differences in the y direction.

#### Perona–Malik

The Perona–Malik equation [19] is another edge-preserving diffusion,111As
noted by Esedoḡlu [8], the Perona–Malik equation is often thought of as
$\partial_{t}u=\operatorname{div}\bigl{(}g(|\nabla u|)\nabla u\bigr{)}$, but
more precisely, Perona and Malik’s intention and discretization are suggestive
of the anisotropic form written above.

(18)
$\partial_{t}u=\partial_{x}\bigl{(}g(|u_{x}|^{2})u_{x}\bigr{)}+\partial_{y}\bigl{(}g(|u_{y}|^{2})u_{y}\bigr{)},$

in which we set $g(s)=1/(1+s/c^{2})$.

We use the original scheme suggested by Perona and Malik [19] as reference
implementation:

(19) $\displaystyle
u_{m,n}^{(k+1)}=u^{(k)}_{m,n}+\frac{\mathit{dt}}{\mathit{dx}}\Bigl{[}$
$\displaystyle g(|D_{+}^{x}u_{m,n}^{(k)}|^{2})D_{+}^{x}u_{m,n}^{(k)}$
$\displaystyle{-}$ $\displaystyle
g(|D_{-}^{x}u_{m,n}^{(k)}|^{2})D_{-}^{x}u_{m,n}^{(k)}$ $\displaystyle{+}$
$\displaystyle g(|D_{+}^{y}u_{m,n}^{(k)}|^{2})D_{+}^{y}u_{m,n}^{(k)}$
$\displaystyle{-}$ $\displaystyle
g(|D_{-}^{y}u_{m,n}^{(k)}|^{2})D_{-}^{y}u_{m,n}^{(k)}\Bigr{]}.$

#### Coherence enhancing diffusion

Coherence enhancing diffusion (CED) is an anisotropic diffusion introduced by
Weickert [27],

(20) $\partial_{t}u=\operatorname{div}(D\nabla u),$

where $D$ is a tensor, a $2\times 2$ matrix at each pixel, constructed from
the structure tensor of $u$ as

(21) $\displaystyle\mu_{1}$ $\displaystyle=\alpha,$ (22)
$\displaystyle\mu_{2}$
$\displaystyle=\alpha+(1-\alpha)\exp\bigl{(}-C/(\lambda_{1}-\lambda_{2})^{2}\bigr{)},$
(23) $\displaystyle D$
$\displaystyle=\mu_{1}\mathbf{w}_{1}\mathbf{w}_{1}^{T}+\mu_{2}\mathbf{w}_{2}\mathbf{w}_{2}^{T},$

where $\alpha$ and $C$ are positive parameters, and as described previously in
§2.1, $\lambda_{1},\lambda_{2},\mathbf{w}_{1}$, $\mathbf{w}_{2}$ are the
eigenvalues and eigenvectors of $J_{\rho}(\nabla u)$. The CED equation has an
oil painting like effect, diffusing the image along oriented features in a
flowing self-reinforcing manner.

We use the scheme developed by Weickert and Scharr [28] as reference
implementation, in which spatial derivatives in (20) are discretized with the
rotation-optimized $3\times 3$ filters

(24) $\frac{1}{32}\begin{pmatrix}-3&0&3\\\ -10&0&10\\\
-3&0&3\end{pmatrix},\;\frac{1}{32}\begin{pmatrix}3&10&3\\\ 0&0&0\\\
-3&-10&-3\end{pmatrix}.$

#### Cahn–Hilliard

The Cahn–Hilliard equation is

(25) $\partial_{t}u=\Delta\bigl{(}W^{\prime}(u)-\gamma\Delta u\bigr{)}$

where $\gamma$ is a positive parameter and $W^{\prime}$ denotes the derivative
of the double-well potential $W(u)=u^{2}(u-1)^{2}$.

This equation describes the process of phase separation of a binary fluid. If
applied to a grayscale image, the equation drives the image toward binary 0/1
values (the minima of $W$). The equation conserves the total mass
$\int_{\Omega}u\,dx$, so regions of darker graylevels tend to evolve toward
islands of white surrounded by black, and symmetrically for lighter grays,
creating a dithering-like effect.

Bertozzi, Esedoḡlu and Gillette [3] showed that the Cahn–Hilliard equation
when applied on binary images has the effect of bridging small gaps, which
makes it useful for inpainting binary images.

Following Bertozzi et al. [3], we use the following semi-implicit scheme as
reference implementation

(26)
$\mathbf{u}^{(k+1)}+\mathit{dt}\gamma\Delta^{2}\mathbf{u}^{(k+1)}=\mathbf{u}^{(k)}+\mathit{dt}\Delta
W^{\prime}(\mathbf{u}^{(k)}).$

The equation is solved for $\mathbf{u}^{(k+1)}$ by inverting
$(I+\mathit{dt}\gamma\Delta^{2})$ in the Fourier domain.

###  3.1 Conservative Model

Several of the anisotropic diffusion PDEs in the previous section have the
form of a conservation law:

(27) $\partial_{t}u=\operatorname{div}G(u),$

where $G$ is

  * •

TV flow (16): $G(u)=\nabla u/|\nabla u|$,

  * •

Coherence enhancing diffusion (20): $G(u)=D\nabla u)$,

  * •

Cahn–Hilliard (25): $G(u)=\nabla\bigl{(}W^{\prime}(u)-\gamma\Delta u\bigr{)}$.

The form (27) implies that total mass $\int u\,dx$ is conserved. We can
integrate over a control area $A$ and apply divergence theorem to express the
right hand side of (27) as fluxes across its boundary $S$:

(28)
$\partial_{t}\Bigl{(}\int_{A}u\,dx\Bigr{)}=\int_{A}\operatorname{div}G(u)\,dx=\oint_{S}G(u)\cdot\mathbf{n}\,dS.$

Provided no flux crosses the image boundaries, flux exiting one element enters
a neighboring element so that total mass is conserved.

$\operatorname{BLADE}^{x}$$\operatorname{BLADE}^{y}$$\smash{\hat{\mathbf{g}}^{x}}\vphantom{x}$$\smash{\hat{\mathbf{g}}^{y}}\vphantom{x}$
Eq. (31)
$\mathit{dt}$+$\smash{\hat{\mathbf{u}}^{(k)}}\vphantom{x}$$\smash{\hat{\mathbf{u}}^{(k+1)}}\vphantom{x}$

Figure 7: BLADE flux model for conservative laws.

Figure 8:  Input images to BLADE in Fig. 10, $4\times$ coarsened grayscale versions of Kodak images #5 and #7. | TV flow | Perona–Malik | CED | Cahn–Hilliard
---|---|---|---|---
Reference |  |  |  |
BLADE |  |  |  |
Reference |  |  |  |
BLADE |  |  |  |
Figure 9:  First and third rows: Results computed by reference implementation at high resolution. Second and fourth rows: Approximations with BLADE. The input image to BLADE is shown in Fig. 8. Color input | BLADE TV flow | BLADE Perona–Malik | BLADE CED
---|---|---|---
|  |  |
Figure 10:  BLADE PDE approximations applied to a color input image.  Figure
11:  BLADE filters for TV flow. [Different orientation bins along columns, and
strength and then coherence along rows. There are 24 orientations, 3
strengths, and 3 coherences, and $\rho=1.0$.] Figure 12:  BLADE filters for
CED. Figure 13:  BLADE filters for Perona–Malik. Figure 14:  BLADE filters for
Cahn–Hilliard. Different orientation bins along rows, and strength and then
intensity along columns. There are 8 orientations, 5 strengths, and 6
intensities.

Finite volume methods partition the domain into volume elements and implement
the right hand side of (28) as a sum of fluxes between adjacent elements,
ensuring that the discrete scheme conserves mass. We can follow this approach
to make a BLADE-based finite volume method that conserves mass. Let
$\operatorname{BLADE}^{x}$ and $\operatorname{BLADE}^{y}$ denote two BLADE
networks (with independent filters and selection rule) that estimate
respectively $G(u)^{x}_{m+1/2,n}$ and $G(u)^{y}_{m,n+1/2}$ (the $x$ and $y$
components of $G(u)$ at boundary midpoints),

(29) $\displaystyle\operatorname{BLADE}^{x}(u)_{m,n}$
$\displaystyle=\hat{G}^{x}_{m+1/2,n}\approx G(u)^{x}_{m+1/2,n},$ (30)
$\displaystyle\operatorname{BLADE}^{y}(u)_{m,n}$
$\displaystyle=\hat{G}^{y}_{m,n+1/2}\approx G(u)^{y}_{m,n+1/2}.$

Fluxes across the image boundaries are set to zero (Neumann boundary
condition). The flux estimates $\hat{G}^{x}$ and $\hat{G}^{y}$ are then summed
as follows to estimate $\partial_{t}u$:

(31) $\displaystyle\hat{\partial_{t}u}_{m,n}=$
$\displaystyle\hat{G}^{x}_{m+1/2,n}-\hat{G}^{x}_{m-1/2,n}$
$\displaystyle{+}\,\,\hat{G}^{y}_{m,n+1/2}-\hat{G}^{y}_{m,n-1/2}.$

The BLADE model with this time derivative conserves mass.

###  3.2 Results

With all above PDEs, we set parameters and a large enough stopping time to
produce a moderate effect, and approximate them with BLADE model (10) with ten
time steps.

###  3.3 Evaluation

We compare the image at the stopping time between the reference implementation
and the BLADE approximation. Figure 9 shows reference vs. BLADE for Kodak
images #5 and #7.

The following table lists average PSNR and SSIM over the Kodak Image Suite
(higher is better).

| Average PNSR | Average SSIM
---|---|---
TV flow | 33.43 | 0.9553
Perona–Malik | 35.68 | 0.9821
CED | 35.54 | 0.9663
Cahn–Hilliard | 12.92 | 0.7737

The BLADE Cahn–Hilliard approximation is quantitatively poor, yet it is
visually close (Fig. 9).

###  3.4 Color Images

Even though we train BLADE on grayscale images, the resulting model extends
readily to color images by two minor modifications: first, selection is
performed on the input’s luma channel222Or as described in [7], a more robust
structure tensor analysis can be done using jointly all color channels, with
only a small increase in computation cost., and second, the BLADE filtering
equation (5) is performed independently on the R, G, B channels.

##  4 Combining with Other Terms

Ground Truth | Input | Deconvolution
---|---|---
|  |
Figure 15: BLADE CED-regularized deconvolution with (36). Input has PSNR 23.04 dB and SSIM 0.6424. The deconvolved image has PSNR 24.66 dB and SSIM 0.7436. Input | Upscaled $4\times$
---|---
|
Figure 16:  BLADE TV-regularized upscaling by factor 4.

TV flow, Perona–Malik, and CED are all denoising or regularizing processes in
the sense that they tend to remove noise while preserving image content. They
may be used as regularizers to solve deblurring, inpainting, image upscaling,
and other inverse problems by adding a $+\lambda A^{T}(f-Au)$ term to the PDE.

Consider generically a degradation (observation) model of the form
$f=Au+\mathit{noise}$ where $A$ is a linear operator and the noise is white
Gaussian, then restoration of $u$ for instance with CED is

(32) $\partial_{t}u=\operatorname{div}(D\nabla u)+\lambda A^{T}(f-Au),$

in which $A^{T}$ denotes the transpose (adjoint) of $A$ and $\lambda$ is a
parameter balancing between matching the degradation model and regularization.

Using the BLADE CED approximation from the previous section, we can implement
CED-regularized restoration (32), without needing to retrain BLADE, as

(33)
$\hat{\mathbf{u}}^{(k+1)}=\hat{\mathbf{u}}^{(k)}+\mathit{dt}\bigl{(}\operatorname{BLADE}(\hat{\mathbf{u}}^{(k)})+\lambda
A^{T}(\mathbf{f}-A\hat{\mathbf{u}}^{(k)})\bigr{)}$

initialized with $\hat{\mathbf{u}}^{(0)}=\mathbf{f}$. BLADE TV flow and BLADE
Perona–Malik approximations can be applied similarly.

###  4.1 Nonblind deconvolution

In the case of nonblind deconvolution with blur kernel $\varphi$, we have
$Au=\varphi*u$ and

$\displaystyle\lambda A^{T}(\mathbf{f}-A\hat{\mathbf{u}})$
$\displaystyle=\lambda\tilde{\varphi}*(\mathbf{f}-\varphi*\hat{\mathbf{u}})$
(34)
$\displaystyle=-\lambda\tilde{\varphi}*\varphi*\hat{\mathbf{u}}+\lambda\tilde{\varphi}*\mathbf{f},$

where $\tilde{\varphi}$ denotes spatial reversal of $\varphi$. The
$-\lambda\tilde{\varphi}*\varphi*\hat{\mathbf{u}}$ term can be absorbed into
the BLADE filters, still without needing to retrain. Let $\mathbf{h}^{k}$
denote the BLADE CED filters, then we create BLADE CED deconvolution filters

(35)
$\mathbf{h}^{k}_{\mathit{deconv}}=\mathbf{h}^{k}-\lambda\tilde{\varphi}*\varphi,$

and (33) with these filters becomes

(36)
$\hat{\mathbf{u}}^{(k+1)}=\hat{\mathbf{u}}^{(k)}+\mathit{dt}\bigl{(}\operatorname{BLADE}(\hat{\mathbf{u}}^{(k)})+\lambda\tilde{\varphi}*\mathbf{f}\bigr{)}.$

The parameter $\lambda$ balances between deconvolution vs. denoising strength.
In Fig. 15, we blurred the input image with a Gaussian with standard deviation
of 1 pixel and added noise of standard deviation 5. The CED-regularized
deconvolution is computed with ten time steps (36) and $\lambda=2$.

###  4.2 Image upscaling

For image upscaling, we consider the degradation model
$f={\downarrow}(\varphi*u)+\mathit{noise}$, where $\varphi$ is the point
spread function and $\downarrow$ denotes subsampling. We perform upscaling as

(37)
$\begin{array}[]{@{}r@{}l@{}}\hat{\mathbf{u}}^{(k+1)}=\hat{\mathbf{u}}^{(k)}+\mathit{dt}\bigl{[}&\operatorname{BLADE}(\hat{\mathbf{u}}^{(k)})\\\
&+\lambda\tilde{\varphi}*{\uparrow}\bigl{(}\mathbf{f}-{\downarrow}(\varphi*\hat{\mathbf{u}}^{(k)})\bigr{)}\bigr{]}\end{array}$

where $\uparrow$ denotes the transpose operation of $\downarrow$, which is
upsampling by inserting zeros, and $\hat{\mathbf{u}}^{(0)}$ is initialized
with Lanczos interpolation of $\mathbf{f}$. Figure 16 shows factor 4 upscaling
of a crop of Kodak image #14 using in this example BLADE TV flow as the
regularizer and ten time steps. The point spread function $\varphi$ is a
Gaussian with standard deviation $0.4$ high-res pixels, and $\lambda=0.35$.

Image | Initialization
---|---
|
Reference implementation | BLADE implementation
|
Figure 17: Example Chan–Vese segmentation using BLADE TV approximation. Top
left: input image. Top right: initialization. Bottom left: segmentation with
reference implementation, $\mu=0.2$. Bottom right: segmentation with BLADE TV
approximation (41), $\mu=0.04$.

If $f$ is assumed to have no noise, another way to perform upscaling with a
trained BLADE regularizer is to project the PDE as

(38)
$\hat{\mathbf{u}}^{(k+1)}=\hat{\mathbf{u}}^{(k)}+\mathit{dt}P_{0}\bigl{(}\operatorname{BLADE}(\hat{\mathbf{u}}^{(k)})\bigr{)},$

where, as detailed for instance in [10], $\hat{\mathbf{u}}^{(0)}$ is
constructed in the Fourier domain to satisfy
$\mathbf{f}={\downarrow}(\varphi*\hat{\mathbf{u}}^{(0)})$ and $P_{0}$ denotes
orthogonal projection onto the subspace
$\\{\mathbf{u}:{\downarrow}(\varphi*\mathbf{u})=0\\}$.

###  4.3 Segmentation

BLADE TV is also useful for image segmentation. In the Chan–Vese “active
contours without edges” segmentation method [4], the segmentation contour is
found as the zero level set of a function $\varphi$ that minimizes

(39)
$\displaystyle\mu\int\delta_{\epsilon}\bigl{(}\varphi(x)\bigr{)}\lvert\nabla\varphi\rvert+\nu\int
H_{\epsilon}\bigl{(}\varphi(x)\bigr{)}$
$\displaystyle{+}\>\lambda_{1}\int\lvert
f(x)-c_{1}\rvert^{2}H_{\epsilon}\bigl{(}\varphi(x)\bigr{)}$
$\displaystyle{+}\>\lambda_{2}\int\lvert
f(x)-c_{2}\rvert^{2}\bigl{(}1-H_{\epsilon}\bigl{(}\varphi(x)\bigr{)}\bigr{)},$

where $c_{1}$, $c_{2}$ are scalars that are simultaneously optimized,
$H_{\epsilon}$ denotes a smoothed version of the Heaviside step function,
$\delta_{\epsilon}$ is its derivative, and $\mu$, $\nu$, $\lambda_{1}$,
$\lambda_{2}$ are constant parameters. In the gradient descent equation for
$\varphi$, a scaled total variation flow appears in the first term:

(40) $\displaystyle\partial_{t}\varphi$
$\displaystyle=\delta_{\epsilon}(\varphi)\bigl{[}\mu\operatorname{div}\bigl{(}\tfrac{\nabla\varphi}{\lvert\nabla\varphi\rvert}\bigr{)}$
$\displaystyle\hphantom{=}{-}\>\nu-\lambda_{1}(f-c_{1})^{2}+\lambda_{2}(f-c_{2})^{2}\bigr{]}.$

We insert BLADE TV flow from the previous section to approximate Chan–Vese
gradient descent as

(41) $\displaystyle\varphi^{(k+1)}$
$\displaystyle=\varphi^{(k)}+\mathit{dt}\,\delta_{\epsilon}(\varphi^{(k)})\bigl{[}\mu\operatorname{BLADE}(\varphi^{(k)})$
$\displaystyle\hphantom{=}{-}\>\nu-\lambda_{1}(f-c_{1}^{(k)})^{2}+\lambda_{2}(f-c_{2}^{(k)})^{2}\bigr{]}.$

Figure 17 shows an example segmentation with (41) with $\mu=0.04$, $\nu=0$,
$\lambda_{1}=\lambda_{2}=1$, $\mathit{dt}=0.5$ at convergence after 300
iterations, and compares with a reference implementation [11] of Chan–Vese,
starting from a checkerboard initialization as suggested in [4]. When evolved
for such a long time, the BLADE approximation is well beyond the 10 frames of
unrolling done during training, and the approximation is blurrier than the
reference. To counteract this effect, the BLADE result shown was performed
with $\mu=0.04$ while the reference was with $\mu=0.2$ and otherwise same
parameters. With this adjustment, the BLADE-based segmentation is
qualitatively similar, capturing nearly the same boundary for around the
wrench, including the small hole in the handle.

Reference implementation | BLADE implementation
---|---
|
Figure 18: Example Chan–Vese–Sandberg segmentation. Left: segmentation with
reference implementation, $\mu=0.2$. Right: segmentation with BLADE TV
approximation, $\mu=0.04$.

For color image segmentation, the Chan–Vese–Sandberg method extends (40) to

(42) $\displaystyle\partial_{t}\varphi$
$\displaystyle=\delta_{\epsilon}(\varphi)\bigl{[}\mu\operatorname{div}\bigl{(}\tfrac{\nabla\varphi}{\lvert\nabla\varphi\rvert}\bigr{)}$
$\displaystyle\hphantom{=}{-}\>\nu-\lambda_{1}\lVert
f-c_{1}\rVert^{2}+\lambda_{2}\lVert f-c_{2}\rVert^{2}\bigr{]},$

in which $f$ is now a color image and $c_{1}$, $c_{2}$ are color vectors, yet
$\varphi$ and the TV flow term have a single channel as before. We approximate
Chan–Vese–Sandberg segmentation by again replacing the TV flow term with BLADE
TV. Figure 18 shows an example segmentation.

##  5 Optical Flow

$\operatorname{BLADE}^{x}$$\operatorname{BLADE}^{y}$$\delta^{x}$$\delta^{y}$+$\smash{\mathbf{u}}\vphantom{x}$$\smash{\tilde{\mathbf{u}}}\vphantom{x}$

Figure 19:  Network for resampling by a fractional displacement, based on
BLADE approximation of a version of the brightness constancy equation (46).

HR image Translate $\smash{\varphi}\vphantom{x}$
$\smash{\downarrow}\vphantom{x}$ $\smash{\varphi}\vphantom{x}$
$\smash{\downarrow}\vphantom{x}$ LR image Translated
LR image

Figure 20: Preparing training data for BLADE-based resampling. Input (frame 0) | Exact (frame 2)
---|---
|
Bicubic resampling | BLADE resampling
|
Figure 21: Comparison of bicubic and BLADE for resampling frame 0 to frame 1, then resampling that result from frame 1 to frame 2. Input (frame 0) | Exact (frame 2)
---|---
|
Bicubic resampling | BLADE resampling
|
Figure 22: Another example of resampling frame 0 to frame 1, then resampling
frame 1 to frame 2.

In this section, we describe a method of resampling an image at displaced
positions using BLADE. Given an image $\mathbf{u}$ and displacement field
$(\mathbf{v}^{x},\mathbf{v}^{y})$, we seek to compute a resampled image whose
$(m,n)$th pixel is sampled at the displaced location
$(m+v^{x}_{m,n},n+v^{y}_{m,n})$. Such resampling is needed for instance in
optical flow or image registration when mapping a moving image to a fixed
image. This BLADE-based resampling is a variation of the RAISR super-
resolution method of Romano et al. [21]. Super-resolution can be seen as a
case of resampling at regularly-spaced sample locations; the new difficulty
here is that the sampling locations are typically irregular as determined by
the displacement field $(\mathbf{v}^{x},\mathbf{v}^{y})$.

We model the observed image $\mathbf{u}$ as having been sampled from a
convolution of an underlying continuous domain image $U(x,y)$ with a point
spread function $\varphi(x,y)$, plus noise,

(43) $u_{m,n}=(\varphi*U)(m,n)+\mathit{noise}_{m,n}.$

The desired resampled image, displaced by $(\mathbf{v}^{x},\mathbf{v}^{y})$,
is

(44) $\tilde{u}_{m,n}=(\varphi*U)(m+v^{x}_{m,n},n+v^{y}_{m,n}).$

Considering a fixed pixel $(m,n)$, we decompose the displacement vector
$(v^{x}_{m,n},v^{y}_{m,n})$ into integer and fractional parts,

(45) $v^{x}=[v^{x}]+\delta^{x},\quad v^{y}=[v^{y}]+\delta^{y}$

where $[\cdot]$ denotes rounding to the nearest integer and
$\delta^{x},\delta^{y}$ are fractional pixel displacements in $[-1/2,+1/2]$.
For the integer part, (44) reduces to indexing a denoised estimate of
$\mathbf{u}$ at a position shifted by a whole number of pixels. Therefore, we
focus on the fractional part.

For a fractional displacement, we have by Taylor expansion

$\displaystyle\tilde{u}_{m,n}$
$\displaystyle=(\varphi*U)(m+\delta^{x},n+\delta^{y})$
$\displaystyle\approx(\varphi*U)(m,n)+\delta^{x}(\partial_{x}\varphi*U)+\delta^{y}(\partial_{x}\varphi*U)$
$\displaystyle=\bigl{(}(\varphi+\delta^{x}\partial_{x}\varphi+\delta^{y}\partial_{y}\varphi)*U\bigr{)}(m,n)$
(46) $\displaystyle=\bigl{(}(I+\delta\cdot\nabla)\varphi*U)\bigr{)}(m,n).$

Therefore in this sense, resampling (44) corresponds to applying the
differential filter $(I+\delta\cdot\nabla)$. This is essentially linearization
of a version of the brightness constancy equation or the optical flow
constraint equation including a point spread function.

We approximate (46) with BLADE as

(47)
$\displaystyle\tilde{\mathbf{u}}=\mathbf{u}+\delta^{x}\operatorname{BLADE}^{x}(\mathbf{u})+\delta^{y}\operatorname{BLADE}^{y}(\mathbf{u}),$

in which $\operatorname{BLADE}^{x}$ and $\operatorname{BLADE}^{y}$ have
independent filters (Figure 19). We prepare data for training by beginning
with a high-resolution image, translating, convolving with $\varphi$, and
downsampling (Figure 20). By this process we create example pairs of an
observed (low resolution) image and a target image translated by a subpixel
displacement. For the point spread function $\varphi$, we use a Gaussian with
standard deviation of $0.5$ pixels.

We test the method on image bursts from a handheld camera so that there is
significant motion (Figures 21 and 22). The flow field between frames was
estimated by a pyramid-based block matching alignment algorithm as described
in [31]. To emphasize the effect of resampling, we show the outcome after two
steps of resampling: frame 0 is resampled to frame 1, then that result is
resampled again to frame 2. Resampling with standard bicubic is also shown. In
comparison, the BLADE result is sharper and captures the structure better.

##  6 Conclusions and Future Work

We have shown that BLADE, a shallow 2-layer network, is reliable and efficient
in approximating several hyperbolic PDEs in image processing. The use of
machine learning for approximating PDEs raises a number of questions. For
instance, stability is an essential property in PDE methods. An interesting
question is whether it is possible to develop learning-based PDE methods that
are both provably stable and have flexible capacity to learn.

## References

  * [1] M. Abadi, P. Barham, J. Chen, Z. Chen, A. Davis, J. Dean, M. Devin, S. Ghemawat, G. Irving, M. Isard, et al., Tensorflow: a system for large-scale machine learning, in 12th USENIX Symposium on Operating Systems Design and Implementation (OSDI 16), 2016, pp. 265–283.
  * [2] F. Andreu, C. Ballester, V. Caselles, J. M. Mazón, et al., Minimizing total variation flow, Differential and integral equations, 14 (2001), pp. 321–360.
  * [3] A. L. Bertozzi, S. Esedoḡlu, and A. Gillette, Inpainting of binary images using the Cahn–Hilliard equation, IEEE Transactions on Image Processing, 16 (2007), pp. 285–291.
  * [4] T. F. Chan and L. A. Vese, Active contours without edges, IEEE Transactions on Image Processing, 10 (2001), pp. 266–277.
  * [5] B. Chang, L. Meng, E. Haber, F. Tung, and D. Begert, Multi-level residual networks from dynamical systems view, arXiv preprint arXiv:1710.10348, (2017).
  * [6] R. T. Chen, Y. Rubanova, J. Bettencourt, and D. Duvenaud, Neural ordinary differential equations, arXiv preprint arXiv:1806.07366, (2018).
  * [7] S. Choi, J. Isidoro, P. Getreuer, and P. Milanfar, Fast, trainable, multiscale denoising, in 2018 25th IEEE International Conference on Image Processing (ICIP), IEEE, 2018, pp. 963–967.
  * [8] S. Esedoḡlu, An analysis of the Perona–Malik scheme, Communications on Pure and Applied Mathematics: A Journal Issued by the Courant Institute of Mathematical Sciences, 54 (2001), pp. 1442–1487.
  * [9] L. C. Evans, Partial differential equations, American Mathematical Society, Providence, R.I., 2010.
  * [10] P. Getreuer, Roussos–Maragos tensor-driven diffusion for image interpolation, Image Processing On Line, 1 (2011), pp. 178–186, https://doi.org/10.5201/ipol.2011.g_rmdi.
  * [11] P. Getreuer, Chan–Vese segmentation, Image Processing On Line, 2 (2012), pp. 214–224.  https://doi.org/10.5201/ipol.2012.g-cv.
  * [12] P. Getreuer, I. Garcia-Dorado, J. Isidoro, S. Choi, F. Ong, and P. Milanfar, BLADE: filter learning for general purpose computational photography, in 2018 IEEE International Conference on Computational Photography (ICCP), IEEE, 2018, pp. 1–11, https://doi.org/10.1109/ICCPHOT.2018.8368476.
  * [13] J. Han, A. Jentzen, and E. Weinan, Solving high-dimensional partial differential equations using deep learning, Proceedings of the National Academy of Sciences, 115 (2018), pp. 8505–8510.
  * [14] K. He, X. Zhang, S. Ren, and J. Sun, Deep residual learning for image recognition, in Proceedings of the IEEE conference on computer vision and pattern recognition, 2016, pp. 770–778.
  * [15] J. Johnson, A. Alahi, and L. Fei-Fei, Perceptual losses for real-time style transfer and super-resolution, in European Conference on Computer Vision, Springer, 2016, pp. 694–711.
  * [16] S. Lele, Compact finite difference schemes with spectral-like resolution., Journal of Computational Physics, 103 (1992), pp. 16–42, https://doi.org/10.1016/0021-9991(92)90324-R.
  * [17] Z. Long, Y. Lu, X. Ma, and B. Dong, Pde-net: Learning pdes from data, in International Conference on Machine Learning, PMLR, 2018, pp. 3208–3216.
  * [18] Y. Lu, A. Zhong, Q. Li, and B. Dong, Beyond finite layer neural networks: Bridging deep architectures and numerical differential equations, in International Conference on Machine Learning, PMLR, 2018, pp. 3276–3285.
  * [19] P. Perona and J. Malik, Scale-space and edge detection using anisotropic diffusion, IEEE Transactions on Pattern Analysis and Machine Intelligence, 12 (1990), pp. 629–639.
  * [20] E. R. Rodrigues, I. Oliveira, R. L. Cunha, and M. A. Netto, DeepDownscale: a deep learning strategy for high-resolution weather forecast, arXiv preprint arXiv:1808.05264, (2018).
  * [21] Y. Romano, J. Isidoro, and P. Milanfar, RAISR: rapid and accurate image super resolution, IEEE Transactions on Computational Imaging, 3 (2017), pp. 110–125.
  * [22] L. I. Rudin, S. Osher, and E. Fatemi, Nonlinear total variation based noise removal algorithms, Physica D: Nonlinear Phenomena, 60 (1992), pp. 259–268.
  * [23] L. Ruthotto and E. Haber, Deep neural networks motivated by partial differential equations, Journal of Mathematical Imaging and Vision, (2019), pp. 1–13.
  * [24] W. E. Sorteberg, S. Garasto, A. S. Pouplin, C. D. Cantwell, and A. A. Bharath, Approximating the solution to wave propagation using deep neural networks, arXiv preprint arXiv:1812.01609, (2018).
  * [25] H. Talebi and P. Milanfar, Learned perceptual image enhancement, in Computational Photography (ICCP), 2018 IEEE International Conference on, IEEE, 2018, pp. 1–13.
  * [26] Z. Wang, A. C. Bovik, H. R. Sheikh, and E. P. Simoncelli, Image quality assessment: from error visibility to structural similarity, IEEE Transactions on Image Processing, 13 (2004), pp. 600–612.
  * [27] J. Weickert, Anisotropic diffusion in image processing, vol. 1, Teubner Stuttgart, 1998.
  * [28] J. Weickert and H. Scharr, A scheme for coherence-enhancing diffusion filtering with optimized rotation invariance, Journal of Visual Communication and Image Representation, 13 (2002), pp. 103–118.
  * [29] E. Weinan, A proposal on machine learning via dynamical systems, Communications in Mathematics and Statistics, 5 (2017), pp. 1–11.
  * [30] E. Weinan and B. Yu, The deep Ritz method: a deep learning-based numerical algorithm for solving variational problems, Communications in Mathematics and Statistics, 6 (2018), pp. 1–12.
  * [31] B. Wronski, I. Garcia-Dorado, M. Ernst, D. Kelly, M. Krainin, C.-K. Liang, M. Levoy, and P. Milanfar, Handheld multi-frame super-resolution, ACM Transactions on Graphics (TOG), 38 (2019), pp. 1–18.