Neural Networks with Cheap Differential OperatorsRicky T. Q. Chen and David Duvenaud
University of Toronto, Vector Institute
Overview
Given a function f : Rd → Rd , we seek to obtain a vectorcontaining its dimension-wise k-th order derivatives,
Dkdimf (x) :=
[∂kf1(x)
∂xk1· · · ∂
kfd(x)
∂xkd
]T∈ Rd (1)
using only k evaluations of automatic differentiation regardlessof the dimension d .
Full Jacobianof f (x)
Ddimf (x)(diag of Jac)
This has applications in:
I. Solving differential equations.
II. Continuous Normalizing Flows.
III. Learningstochastic differential equations.
Reverse-mode Automatic Differentiation
Deep learning software rely on automatic differentiation (AD)to compute gradients. More generally, AD computesvector-Jacobian products.
vT∂f (x)
∂x=
d∑i=1
vi∂fi(x)
∂x(2)
However, computing the Jacobian trace (ie.∑Ddim)
is as expensive as the full Jacobian. One evaluation ofAD can only compute one row of the Jacobian.
Forward: Network Structure
Our approach:Dimension-wise derivatives can be efficiently computed forrestricted architectures with simple modifications to the ADprocedure in the backward pass.
Conditioner hi = ci(x−i). The i -th hidden dimensiondepends on all inputs except the i -th input dimension. Can becomputed in parallel using masked neural networks (e.g.MADE, PixelCNN).
Transformer fi(x) = τi(xi , hi). τi : Rd → R outputs the i -thdimension given the concatenated vector. Can be computed inparallel if composed of matrix multiplication and element-wiseoperations.
Backward: Modified Computation Graph
x1
x2
xd
...
h1
h2
hd
...
f1
f2
fd
...
Forward computation graph
x1
x2
xd
...
h1
h2
hd
...
f1
f2
fd
...
Backward computation graph
Let h be stop_gradient(h) and f = τ (x , h), so
∂ f i(x)
∂xj=∂τi(xi , hi)
∂xj=
{∂fi(x)∂xi
if i = j
0 if i 6= j(3)
Dimension-wise derivatives.
1T ∂ f (x)
∂x= Ddimf =
[∂f1(x)∂x1· · · ∂fd(x)
∂xd
]T= Ddimf (4)
Higher orders.
1T ∂Dk−1dim f (x)
∂x= Dk
dimf (x) = Dkdimf (x) (5)
Backpropagating through dim-wise derivatives.
∂Dkdimf
∂w+∂Dk
dimf
∂h
∂h
∂w=∂Dk
dimf
∂w(6)
App I: Linear Multistep ODE Solvers
Implicit ODE solvers need to solve an optimizationsub-problem in the inner loop.
General idea: replace Newton-Raphson
y (k+1) = y (k) −
[∂F (y (k))
∂y (k)
]−1
F (y (k)) (7)
with Jacobi-Newton
y (k+1) = y (k) − [DdimF (y)]−1 � F (y (k)) (8)
0 2000 4000 6000 8000 10000 12000
Training Iteration
0
500
1000
1500
2000
Num
.E
valu
atio
ns
RK4(5)
ABM
ABM-Jacobi
App II: Continuous Normalizing Flows
If dxdt = f (t, x), then ∂ log p(x)
∂t = − tr(∂f∂x
). (See ”Neural ODEs”)
ModelMNIST Omniglot
ELBO ↑ NLL ↓ ELBO ↑ NLL ↓VAE -86.55 82.14 -104.28 97.25
Planar -86.06 81.91 -102.65 96.04
IAF -84.20 80.79 -102.41 96.08
Sylvester -83.32 80.22 -99.00 93.77
FFJORD -82.82 − -98.33 −DiffOp-CNF -82.37 80.22 -97.42 93.90
Table: Evidence lower bound and negative log-likelihood for staticMNIST and Omniglot.
Exact trace vs stochastic trace. Converges faster andresults in easier to solve dynamical systems.
App III: Stochastic Differential Eqs
dx(t) = f (x(t), t)dt + g(x(t), t)dW (9)
Data Density Learned
Idea: match left- and right-hand-side of the Fokker-Planckequation, which describes the change in density.
∂p(t, x)
∂t= −
d∑i=1
∂
∂xi[fi(t, x)p(t, x)] +
1
2
d∑i=1
∂2
∂x2i
[g 2ii (t, x)p(t, x)
](10)
References
• Germain et al. “MADE: Masked autoencoder for distributionestimation.” (2015)
• Huang et al. “Neural Autoregressive Flows.” (2018)
• Chen et al. “Neural ODEs.” (2018)
