# MeanFlow in PyTorch and the JVP problem behind one-step generation

> MeanFlow is an unofficial PyTorch implementation of MeanFlow and Improved MeanFlow, two papers on one-step generative modeling, with training configs for MNIST, CIFAR-10 and ImageNet latents. Its most interesting documentation is the known issues section, where the JVP computation and PyTorch native Flash Attention are shown to be incompatible and worked around with two forward passes.

**haidog-yaqub/MeanFlow** — PyTorch implementation of MeanFlow & iMF (one-step generative modeling).

- Repository: https://github.com/haidog-yaqub/MeanFlow
- Stars: 1,213 · Forks: 68
- Language: Python
- License: MIT
- Published: 2026-10-19 · Updated: 2026-10-19 · Language: en
- Canonical page: https://hysenlabs.com/projects/haidog-yaqub-meanflow

## An unofficial port, and why that matters

MeanFlow is a PyTorch implementation of Mean Flows for One-step Generative Modeling, known as MeanFlow, and of Improved Mean Flows, known as iMF. The README states in a note that this is an unofficial implementation, based on the papers and on the official JAX repository imeanflow.

That disclaimer is the first thing to weigh, and it cuts both ways. In favour of the project, a PyTorch implementation means the tooling ecosystem is the one most people already have: Accelerate for distributed training, diffusers for the surrounding model components, and the usual PyTorch checkpointing and visualisation tools. Against it, an unofficial implementation can lag the reference on algorithmic changes, and when a paper is revised, the JAX repo is where the revision lands first.

The subject itself is worth a sentence. These papers are about generating a sample in one step rather than running an iterative denoising loop, and the examples in the repository are the point of the exercise: an MNIST sample after 10k steps, an MNIST sample after 6k steps with classifier-free guidance at a scale of 2.0, and a CIFAR-10 sample after 200k steps with the same guidance scale.

The repository is small and legible, with train.py as the entry point, meanflow.py holding the core, data.py for the data pipeline, a models/ directory and a configs/ directory. The project also says it builds upon Just-a-DiT and EzAudio, and that the iMF support is based on the official JAX repository, with thanks to those authors for releasing their code.

The audience is a researcher or engineer who wants to reproduce or extend one-step generation in PyTorch, and who reads the papers as well as the code.

## Two modes, and the recommendation to use the newer one

The training code supports two algorithms behind one switch, and the README is direct about which to use.

The configuration field meanflow.mode takes the values meanflow and i-meanflow, and the note in the usage section says that i-meanflow is more stable and recommended for your projects, while meanflow mode is kept for reference.

That is a small sentence carrying a lot of information. It tells you the original formulation is retained so that the original results can be reproduced and compared against, and that the improved variant is the one the author considers usable. Anyone starting a new project should take that recommendation at face value, and anyone reproducing a number from the MeanFlow paper should use the meanflow mode deliberately rather than by default.

The update log for 2026.06.13 lists what the current code contains: MeanFlow and iMF training through that mode switch, a dual-head DiT where the u head serves MeanFlow and the v head serves flow matching, classifier-free guidance scale passed as a model input with CFG distillation, the JVP computed under no_grad for dudt only with a separate gradient-enabled forward pass for optimisation, and config-based training for MNIST, CIFAR-10 and ImageNet latents.

The dual-head design is worth understanding, because it explains why the model can serve both objectives. A flow matching model learns a velocity field, which is the v head. MeanFlow adds an average velocity that depends on the time interval, the u head. Keeping both in one network means one set of weights can be trained for either objective, and the switch between them is a configuration value rather than a separate checkpoint.

The guidance scale as a model input is the other piece of that design. By feeding the guidance scale into the network and distilling it during training, the model learns to respond to a guidance weight at sampling time, rather than having guidance applied by an external formula that rescales two predictions.

## The JVP and Flash Attention do not mix

The known issues section is the most technically interesting part of the README, and it deserves to be read before anything is run.

The problem is that jvp is currently incompatible with PyTorch native Flash Attention, that is, with scaled_dot_product_attention. The JVP is the operation MeanFlow needs to compute the derivative of the network output with respect to time, dudt, and it has to be evaluated inside the forward pass, which is exactly where the fused attention kernel is used.

The workaround in this repository is a two-pass structure. The JVP pass runs under no_grad, for the dudt value only, and without Flash Attention. A separate forward pass, with gradients enabled, runs with Flash Attention for the optimisation step. In other words, the model is evaluated twice per training step, once to get the derivative and once to get the gradients.

That has a direct cost in memory and in wall clock time, and it means the two passes are not computing quite the same thing numerically, since the attention implementation differs between them. For a research implementation this is a reasonable trade, and the alternative mentioned in the README is a Triton-based kernel with JVP support, such as the rcm project from NVlabs, which would allow Flash Attention inside the JVP and might offer additional speed and memory benefits. Such kernels are not currently supported in this repository, and the corresponding entry in the TODO list, Triton JVP plus Flash Attention, remains unchecked.

The honest summary is that the project currently cannot have both, and the README says so in plain language rather than leaving you to discover it.

## Three datasets, and how to start a run

The installation is a single pip command with the dependencies the training script needs:

```bash
pip install torch accelerate torchvision einops tqdm diffusers
```

Training is driven by a config file, and the shipped configs cover three datasets: configs/mnist.py for MNIST, configs/cifar10.py for CIFAR-10, and configs/imagenet_latent.py for ImageNet in latent space with a VAE.

The command line is short:

```bash
# single GPU
python train.py --config configs/mnist.py

# custom run name
python train.py --config configs/cifar10.py --run_suffix exp1

# multi-GPU
accelerate launch --num_processes 2 train.py --config configs/mnist.py
```

The run suffix is what keeps experiments apart, and it matters more than it looks, because a training run in this project produces images and checkpoints that you will want to compare across configurations, and the run name is the directory they land in.

The multi-GPU path uses Accelerate rather than a custom launcher, which means the same script runs on one GPU or several, and that the process count is a flag rather than a configuration change.

The common config fields the README lists are n_steps, batch_size, grad_clip, mixed_precision, meanflow.mode and meanflow.cfg_scale. Those six cover the axes a researcher changes most often, the length of the run, the batch size, gradient clipping, whether mixed precision is on, which algorithm to train, and the guidance scale.

## Logs, checkpoints and the shape of an experiment

Training logs are saved to a directory named after the run, and the layout is described in the README:

```
logs/{run_name}/
├── config.py
├── train.log      # loss, FM/MF loss, MF_V_MSE, grad norm, LR
├── images/
└── ckpts/
```

Three things are worth noting in that tree. The config is copied into the run directory as config.py, so a training run carries the exact configuration that produced it and not a reference to a config that may have changed since. The training log records loss, the FM and MF losses, MF_V_MSE, the gradient norm and the learning rate, which is enough to tell a diverging run from a slow one without rerunning anything. And there are separate images/ and ckpts/ directories, so the samples and the weights are kept apart.

The MF_V_MSE entry in the log is the one specific to this project rather than a generic training metric, and it is the number to watch when comparing the two modes, since it measures the discrepancy between the mean flow prediction and the flow matching target.

The samples in the examples are described in the same terms the log uses. The MNIST example is 10k steps with a 1-step sample, so no guidance is involved. The second MNIST example is 6k steps with a 1-step CFG at a scale of 2.0. The CIFAR-10 example is 200k steps with 1-step CFG at a scale of 2.0.

The TODO list gives an honest picture of what is done. Basic training and inference, multi-GPU via Accelerate, classifier-free guidance, latent image training and Improved MeanFlow are all marked complete, with only Triton JVP and Flash Attention outstanding.

## MIT, no releases, and a repository that stopped moving in July

The licence is MIT, with the text in the LICENSE file at the root of the repository, which is the usual arrangement for a research code release and requires keeping the copyright and licence notice.

The maintenance record is the part to note. The repository is not archived, but the last push was on 2026-07-01, more than three months before the latest state of the project, and there are no GitHub releases at all. The only dated record of what changed is the update log entry for 2026.06.13, which lists the current feature set in a single block.

The citation section asks users to cite the original MeanFlow papers as well as the implementation, and it provides three BibTeX entries: the MeanFlow paper with Geng, Shechtman, Kolter and He as authors, the Improved Mean Flows paper with Geng, Lu, Wu, Shechtman, Kolter and He, and the repository itself as a misc entry with the implementation's own author and year.

For a research codebase, that citation request is the part to honour if the work ends up in a paper. A one-step generative model reproduced from an unofficial implementation is the kind of result a reviewer will ask about, and pointing at the official JAX repository is the honest way to describe where the algorithm came from.

The project also asks for contributions, saying that contributions and feedback are welcome and inviting issues or pull requests.

## Conclusion

Adopt this repository if you want to train one-step generative models in PyTorch rather than in JAX, and if the dataset scale you care about is MNIST, CIFAR-10 or ImageNet latents, since those are the configs that exist. Do not adopt it as a production inference library, since the examples are training runs and the project has no releases, and do not expect Flash Attention to work inside the JVP pass until a Triton kernel lands. Verify three things first: that you read the known issues section, because the two-pass workaround changes what you can measure, that your GPU memory budget tolerates running the model twice per step, and that the paper's own repo is where you check for algorithm changes, since this implementation is explicitly unofficial. The last push was on 2026-07-01 and the repository has no GitHub releases.

## FAQ

### What is the MeanFlow repository?

An unofficial PyTorch implementation of the MeanFlow and Improved MeanFlow papers on one-step generative modeling, based on the papers and on the official JAX repository imeanflow. It builds upon Just-a-DiT and EzAudio, and the licence is MIT.

### How do I install the dependencies for MeanFlow training?

Run pip install torch accelerate torchvision einops tqdm diffusers. Training itself is started with train.py and a config file, for example python train.py --config configs/mnist.py.

### Which MeanFlow mode should I use?

The i-meanflow mode, set with meanflow.mode, which the README calls more stable and recommended for your projects. The meanflow mode is kept for reference, so it is the right choice only when reproducing the original formulation.

### Does MeanFlow support multi-GPU training?

Yes, through Accelerate, with the command accelerate launch --num_processes 2 train.py --config configs/mnist.py. Multi-GPU via Accelerate is marked complete in the project TODO list.

### What is the known Flash Attention issue in MeanFlow?

The jvp is incompatible with PyTorch native Flash Attention, so the repository runs the JVP pass under no_grad for dudt only without Flash Attention, and a separate gradient-enabled forward pass for optimisation with Flash Attention. Triton kernels with JVP support would remove the conflict but are not supported yet.

### Which datasets can the MeanFlow implementation train on?

The shipped configs cover MNIST with configs/mnist.py, CIFAR-10 with configs/cifar10.py, and ImageNet latents with configs/imagenet_latent.py. Logs, images and checkpoints are written under logs/{run_name}/, and the last push was on 2026-07-01.

## Sources

- [haidog-yaqub/MeanFlow on GitHub](https://github.com/haidog-yaqub/MeanFlow)
- [Issues](https://github.com/haidog-yaqub/MeanFlow/issues)
- [License: MIT](https://github.com/haidog-yaqub/MeanFlow/blob/main/LICENSE)
- [README](https://github.com/haidog-yaqub/MeanFlow/blob/main/README.md)

---

Hysen Labs editorial analysis, written from the project's own repository and release notes. Cite the canonical page: https://hysenlabs.com/projects/haidog-yaqub-meanflow