pnkraemer / mini-stochastic-interpolants

Stochastic interpolants

This repository contains a minimal implementation of some concepts related to stochastic interpolants in JAX, based on this paper by Michael S. Albergo, Nicholas M. Boffi, Eric Vanden-Eijnden.

Disclaimer: This implementation is meant to be didactic. For a more functional version (in Pytorch), see the repository published by the authors of the paper here.

Installation

Before installing this project, and after creating & activating your virtual environment, you must install JAX yourself because CPU and GPU backends require different installation commands. See here for instructions. For the small examples, pip install jax[cpu] will suffice. For the bigger demos, a GPU is helpful.

Then, move to the root of the directory and run

pip install .

This command installs all requirements (Flax, Optax, etc.).

Then, find the content as

from stochint import *

Demonstrations

Find the demos in demos/.

Acknowledgements

Thanks to Paul Jeha (@pablo2909) for teaching us how to write a name with 2d samples.

mini-stochastic-interpolants's Issues

Get rid of write-word.py or move it to a test directory

Consider moving the content from setup.sh to the pyproject.toml respectively README.md

Move all training-related functions from losses.py to util_train.py

Move all plotting-related functions from losses.py to util_plot.py

Move dependencies from requirements.txt to pyproject.toml

Move all simulation-related functions from losses.py to util_simulate.py

Write module-level docstrings for the demos

Implement antithetic sampling

Rename `b_*` to velocity

Write module-level docstrings for the source code

Delete interpolate.py

Implement the loss function for the decoder

Fix and rename the 2d-Gaussian-mixture example

Fix and rename the name-sampling example

Rename `s_*` to score