Generalised Flow Maps for Few-Step Generative Modelling on Riemannian Manifolds

Oscar Davis, Michael S. Albergo, Nicholas M. Boffi, Michael M. Bronstein, Avishek Joey Bose.

Official implementation.

About

This repository includes all of the code necessary to run Generalised Flow Maps on the datasets provided in the paper, with Eulerian, Lagrangian and Progressive self-distillations, Generalised Mean-Flows and Riemannian Flow Matching all implemented, here. Having used Lightning for all our modules, it is easy to add more datasets and to reuse this code in another repository. 💫

• All of the Generalised Flow Map-specific code is located in gfm/models/flowmap_module.py 🧠.

Running the code

1. Install the dependencies:

mamba env create -f environment.yaml

2. Activate the environment:

mamba activate gfm

3. For the experimental data, please follow the instructions given in the Riemannian Flow Matching repository. The structure should be:

gfm/
├── configs
├── data
├── gfm
├── res

4. Run the experiment you want! 💥 For example,

python -m gfm.train experiment=euler/euler_amino_general trainer=gpu

For wandb logging, add logger=wandb as an argument.

Citation

For citing the paper or the code, please use the following BibTeX citation:

@misc{davis2025generalisedflowmapsfewstep,
    title={Generalised Flow Maps for Few-Step Generative Modelling on Riemannian Manifolds}, 
    author={Oscar Davis and Michael S. Albergo and Nicholas M. Boffi and Michael M. Bronstein and Avishek Joey Bose},
    year={2025},
    eprint={2510.21608},
    archivePrefix={arXiv},
    primaryClass={cs.LG},
    url={https://arxiv.org/abs/2510.21608}, 
}