Neural Networks with Cheap Differential Operators

Abstract

Gradients of neural networks can be computed efficiently for any architecture, but some applications require differential operators with higher time complexity. We describe a family of restricted neural network architectures that allow efficient computation of a family of differential operators involving dimension-wise derivatives, used in cases such as computing the divergence. Our proposed architecture has a Jacobian matrix composed of diagonal and hollow (non-diagonal) components. We can then modify the backward computation graph to extract dimension-wise derivatives efficiently with automatic differentiation. We demonstrate these cheap differential operators for solving root-finding subproblems in implicit ODE solvers, exact density evaluation for continuous normalizing flows, and evaluating the Fokker--Planck equation for training stochastic differential equation models.

NeurIPS 2019

Cited by in corpus (11)

• Normalizing Flows for Probabilistic Modeling and Inference
• Equivariant Flows: Exact Likelihood Generative Learning for Symmetric Densities
• Deep Learning Theory Review: An Optimal Control and Dynamical Systems Perspective
• Learning Neural Event Functions for Ordinary Differential Equations
• Almost Surely Stable Deep Dynamics
• Neural Spatio-Temporal Point Processes
• Relative gradient optimization of the Jacobian term in unsupervised deep learning
• Climate Modeling with Neural Diffusion Equations
• "Hey, that's not an ODE": Faster ODE Adjoints via Seminorms
• Differentiable Implicit Layers
• Accelerating Continuous Normalizing Flow with Trajectory Polynomial Regularization