This repository contains the code for the paper Neural Diffusion Processes by Vincent Dutordoir, Alan Saul, Zoubin Ghahramani and Fergus Simpson accepted at ICML 2023.
The code is written in Python and uses the jax framework with the haiku deep learning library. The code has been tested on ubuntu linux with Python 3.9 and 3.10.
Before installing the dependencies, we recommend creating a virtual environment. Then the dependencies and the source (neural_diffusion_processes
) can be installed with the following command using poetry
:
poetry install
This command will install the required dependencies to import the package and run the tests. To verify that the installation was successful, run the following command:
poetry run task test
Instructions to reproduce the paper's experiments can be found in the readme files in each of the experiment directories.
To install the dependencies for the experiments, run the following command:
poetry install --with experiments
This commands installs a lot of extra dependecies that are required for the experiments, such as aim (for experiment tracking), ml-tools, tensorflow-cpu
(for data loading), etc.
Note that we do not install jaxlib
by default, which is required for GPU support. To install jaxlib
with GPU support for your machine following the official instructions here.
@InProceedings{dutordoir2023,
title = {{N}eural {D}iffusion {P}rocesses},
author = {Dutordoir, Vincent and Saul, Alan and Ghahramani, Zoubin and Simpson, Fergus},
url = {https://arxiv.org/abs/2206.03992},
booktitle = {International Conference of Machine Learning},
year = {2023},
}