https://github.com/lebellig/jax-fm

Minimal implementation of flow matching with JAX
https://github.com/lebellig/jax-fm

deep-learning flow-matching generative-ai jax

Last synced: over 1 year ago
JSON representation

Minimal implementation of flow matching with JAX

• Host: GitHub
• URL: https://github.com/lebellig/jax-fm
• Owner: lebellig
• Created: 2024-09-19T08:41:41.000Z (almost 2 years ago)
• Default Branch: main
• Last Pushed: 2024-10-04T13:10:06.000Z (almost 2 years ago)
• Last Synced: 2025-02-06T15:53:21.560Z (over 1 year ago)
• Topics: deep-learning, flow-matching, generative-ai, jax
• Language: Jupyter Notebook
• Homepage:
• Size: 3.56 MB
• Stars: 5
• Watchers: 1
• Forks: 0
• Open Issues: 0
• Metadata Files:
  • Readme: README.md

Awesome Lists containing this project

README

          
# Flow Matching with Jax

_notebook by Georges Le Bellier_  - [Twitter](https://twitter.com/_lebellig), [Website](https://gle-bellier.github.io)

![flow(1)](https://github.com/gle-bellier/jax_fm/blob/main/swiss_1000.gif)

This notebook proposes a minimal implementation of _Flow Matching_ using Jax, Flax and JaxTyping. I mostly wanted to define the conditional velocity as the time derivative of the interpolant (i.e. conditional path):

```python

def interpolant(x0: Float[Array, "N"] , x1: Float[Array, "N"], t: float) -> Float[Array, "N"]:

    return x0 + (x1 - x0) * t

velocity = jax.jacrev(interpolant, argnums=2)

```

## References:

📄 [1] **Flow Matching for Generative Modeling** by Yaron Lipman, Ricky T. Q. Chen, Heli Ben-Hamu, Maximilian Nickel, Matt Le - [Article](https://arxiv.org/abs/2210.02747)

🐍 [2] **Jax** - [Code](https://github.com/jax-ml/jax)

🐍 [3] **Flax** - [Code](https://github.com/google/flax)

🐍 [4] **JaxTyping** by Patrick Kidger - [Doc](https://github.com/patrick-kidger/jaxtyping)

🐍 [5] **Introduction to Flow Matching** by Georges Le Bellier - [Notebook](https://github.com/gle-bellier/flow-matching/tree/main)