Learning Differential Equations that are Easy to Solve

Code for the paper:

Jacob Kelly*, Jesse Bettencourt*, Matthew James Johnson, David Duvenaud. "Learning Differential Equations that are Easy to Solve" arXiv preprint (2020). [arxiv] [bibtex]

*Equal Contribution

Requirements

JAX

Follow installation instructions here.

Haiku

Follow installation instructions here.

Tensorflow Datasets

For using the MNIST dataset, follow installation instructions here.

Usage

Different scripts are provided for each task and dataset.

MNIST Classification

python mnist.py --reg r3 --lam 6e-5

Latent ODEs

python latent_ode.py --reg r3 --lam 1e-2

FFJORD (Tabular)

python ffjord_tabular.py --reg r2 --lam 1e-2

FFJORD (MNIST)

python ffjord_mnist.py --reg r2 --lam 3e-4

Datasets

MNIST

tensorflow-datasets (instructions for installing above) will download the data when called from the training script.

Physionet

The file physionet_data.py, adapted from Latent ODEs for Irregularly-Sampled Time Series will download and process the data when called from the training script. A preprocessed version is available in releases.

Tabular (FFJORD)

Data must be downloaded following instructions from gpapamak/maf and placed in data/. Only MINIBOONE is needed for experiments in the paper.

Code in datasets/, adapted from Free-form Jacobian of Reversible Dynamics (FFJORD), will create an interface for the MINIBOONE dataset once it's downloaded. It is called from the training script.

Acknowledgements

Code in lib is modified from google/jax under the license.

BibTeX

@article{kelly2020easynode,
  title={Learning Differential Equations that are Easy to Solve},
  author={Kelly, Jacob and Bettencourt, Jesse and Johnson, Matthew James and Duvenaud, David},
  journal={arXiv preprint arXiv:2007.04504},
  year={2020}
}