Flow Matching Update

What happened

My implementation of flow matching in tinygrad has been an awesome intellectual challenge.

With limited time to work on personal projects, at the beginning of this year I decided to technically challenge myself by implementing a very interesting mathematical framework in a very modern, lightweight machine learning framework.

Flow matching takes advantage of the concept of flow, a solution to an ordinary differential equation. This flow goes from a starting point (initial condition) to a specified moment in time following a velocity law given by an ordinary differential equation. This is a very loose explanation.

For a machine learning setting, imagine data slowly shifting from noise (initial point, our start at t0) to an image of a cat (our flow reached time t1).

There are many technical difficulties that can arise during that process. From tractability to computational efficiency in ODE solving, we take careful steps in designing such systems.

After a few months of working through Flow Matching Guide and Code, I reached a stage where my model trained very well on a simple dataset (sklearn.datasets.moons) but not on more complex items like MNIST.

What I did

After months of debating, I went back and refactored the code significantly.

I split training examples, reorganized and modularized final data generation logic, and finally debugged my ODE solvers as much as I could.

Aside from refactoring, debugging, simplifying, and modularizing the code, one thing that helped was going from a multi-layer perceptron architecture to a convolutional UNet-like structure.

After all that, the examples, which you are more than welcome to run yourself, generate figures like these:

Next steps

There were many more types of flow matching algorithms listed in the paper.

I am still hoping to get to them (for example, really interested in non-Euclidean flow matching and Markov Chain Models in flow matching!) but the timing for this is now more difficult.

My previous Flow Matching Posts: See all my previous posts on Flow Matching.