We also introduce two coupling approaches advocated this summer to improve FM training: using either very large sharp Sinkhorn couplings (arxiv.org/abs/2506.05526) or, even better, semidiscrete couplings (arxiv.org/abs/2509.25519), as proposed with Alireza Mousavi-Hosseini and @syz.bsky.social

We have been working with Michal Klein on pushing a module to train *flow matching* models using JAX. This is shipped as part of our new release of the OTT-JAX toolbox (github.com/ott-jax/ott) The tutorial to do so is here: ott-jax.readthedocs.io/tutorials/ne...

Wow! Finally OT done on the entire training set to train a diffusion model!

Our two phenomenal interns, Alireza Mousavi-Hosseini and Stephen Zhang @syz.bsky.social have been cooking some really cool work with Michal Klein and me over the summer. Relying on optimal transport couplings (to pick noise and data pairs) should, in principle, be helpful to guide flow matching 🧵

Alireza Mousavi-Hosseini, Stephen Y. Zhang, Michal Klein, Marco Cuturi: Flow Matching with Semidiscrete Couplings https://arxiv.org/abs/2509.25519 https://arxiv.org/pdf/2509.25519 https://arxiv.org/html/2509.25519