Implementing Mean Flows for One-Step Sampling

[amorehead]

Describe the solution you'd like
Adding Mean Flows as an available option would be a nice win for the PyTorch community.

Additional context

[itaigat]

Mean flows is a great work and it would be amazing to have it. However, we currently don't have plans to work on it. Would you be interested in leading it?

[gumran]

Hi @itaigat! The repository lets you build a MeanFlow model via the algorithm from the paper. Would you consider merging a Jupyter notebook like in the examples folder if I were to make one for MeanFlow? Thanks!

[amorehead]

@gumran, did you ever make an example notebook for MeanFlow?

[gumran]

@amorehead, yes, but unfortunately, I've removed it since then. It is fairly straightforward to implement by modifying the 2d_flow_matching.ipynb notebook - just sample two timepoints r and t instead of just one t and use torch.func.jvp and detach to implement the algorithm from the paper.

[itaigat]

@gumran would you like to PR it? we can quickly iterate over it

[gumran]

@itaigat, I'll try to submit a PR today/tomorrow.

[gumran]

@itaigat, feel free to review PR #82, where I wrote a notebook implementing MeanFlows. If the current one-step sampling performance is an issue, we can either further increase the MLP, which will take a long time to train for anyone using Colab, or adjust the architecture otherwise.