STIX: Stochastic Interpolants, A Unifying Framework

👀 Overview

This repository provides a concise implementation and mathematical guide to Stochastic Interpolants, based on the framework introduced in Stochastic Interpolants: A Unifying Framework for Flows and Diffusions by Michael S. Albergo, Nicholas M. Boffi, and Eric Vanden-Eijnden.

The goal of this project is to demonstrate how a single, minimal setup can encompass and unify a wide class of generative models, including but not limited to Flow Matching, EDM-style Diffusion, Bayesian Flow Networks (BFNs) and Discrete Flow Matching.

We also see this library as a stepping stone for future research: every component is highly modular and can be extended or swapped out with ease, especially when it comes to guidance, coupling, or multi-modality.

Written in JAX, stix provides:

• 🌐 Training and sampling of multimodal models
• 🎯 Built-in standard generative frameworks: Flow Matching, Diffusion, and BFN
• 🧭 Guidance for conditional generation
• 🧩 Highly modular components : interpolants, couplings, and guidance are all independently swappable for research liberty
• 📐 Full mathematical infrastructure for stochastic interpolants: interpolant schedules, couplings, and time samplers

Have a look at the Installation section for details on how to install stix. If you want to get started with the library or get a feel of what is possible, you can dive into the introduction and the tutorials notebooks.

Why stix?

A framework for nearly every interpolation scheme: Flow matching, diffusion (e.g. VE and VP), Bayesian Flow Networks, and masked and uniform discrete diffusion all under a single framework. Adding your own is incredibly straightforward.

Discrete and continuous modalities in one framework, by abstracting over the Generator: Build multimodal models with both discrete and continuous data easily. Truly discrete diffusion runs as a continuous-time Markov chain and continuous data as an ODE or SDE, yet both are sampled simultaneously from the same network call.

Sampling decoupled from training: On one-sided paths, target, noise, velocity and score convert in closed form, so a velocity-trained model samples as either an ODE or an SDE with no second head and no retraining.

Different options for handling discrete data: We offer mask and uniform diffusion, as well as methods to learn continuous embeddings of discrete data to be used with continuous interpolants.

Have a look at the Installation section for details on how to install stix. If you want to get started with the library or get a feel of what is possible, you can dive into the introduction and the tutorials notebooks.

We also provide extensive documentation of the different classes and components.

📦 Installation

Install as a dependency (from PyPI)

pip install stix-ml

Or with uv:

uv add stix-ml

Install for development

git clone https://github.com/instadeepai/stix.git && cd stix
uv sync --group dev

Install the hooks once after cloning:

uv run pre-commit install

They will now run automatically on every commit. To run all hooks against every file manually:

uv run pre-commit run --all-files

📓 Tutorials

The tutorials/notebooks directory walks through the library end to end. Each notebook is self-contained and provides a thoroughly documented walkthrough.

Getting started

1. Training and sampling : build the full pipeline to train and sample a generative model.

2. Multimodal data loading with grain : feed real and synthetic data into stix as Batch objects, with checkpointing.

3. Building generative models : write your own GenerativeModel with custom losses and velocity/score conversions.

Going further

4. Conditioning and guidance : conditional sampling via context and intrinsic guidance recipes, and how to write your own.
5. Coupling : pair source and target distributions using methods like product-of-marginals, mini-batch OT, rectified flow.
6. Discrete models : train and sample discrete and mixed continuous-discrete models.

🙏 Acknowledgments

We would like to thank Krisztina Sinkovics (InstaDeep), Bora Guloglu (InstaDeep), Louis Robinson (InstaDeep) and Shaun Kandathil (InstaDeep) for beta-testing and giving feedback on the iterations of this work.

📚 Citing our work

Please cite this repository when using stix in your work.

The BibTeX formatted citation:

@software{stix2026,
  author       = {Simons, Jack and Seince, Maxime and Leach, Adam and
                  Brunken, Christoph and Tilly, Jules and Heyraud, Valentin},
  title        = {{stix}: Stochastic Interpolants, A Unifying Framework},
  year         = {2026},
  version      = {0.1.0},
  organization = {InstaDeep},
  license      = {Apache-2.0},
  url          = {https://github.com/instadeepai/stix},
}