Rectified Flow Matching

Cutting-edge and feature-rich implementation of Rectified Flows from Flow Straight and Fast: Learning to Generate and Transfer Data with Rectified Flows in jax and equinox.

Features

• Deterministic and stochastic sampling of the associated ODE and SDE respectively,
• Mixed precision optimisation,
• Array-typed to-the-teeth for dependable execution with jaxtyping and beartype.

To implement:

• Guidance by score of conditioning
• Mixed precision
• EMA
• AdaLayerNorm
• Stochastic sampling
• ODE Sampling
• Likelihoods
• DiT
• Hyperparameter/model saving

Usage

pip install rectified_flows

uv run --all-extras python examples/cifar10.py

Citations

@misc{liu2022flowstraightfastlearning,
      title={Flow Straight and Fast: Learning to Generate and Transfer Data with Rectified Flow}, 
      author={Xingchao Liu and Chengyue Gong and Qiang Liu},
      year={2022},
      eprint={2209.03003},
      archivePrefix={arXiv},
      primaryClass={cs.LG},
      url={https://arxiv.org/abs/2209.03003}, 
}

@misc{lipman2023flowmatchinggenerativemodeling,
    title={Flow Matching for Generative Modeling}, 
    author={Yaron Lipman and Ricky T. Q. Chen and Heli Ben-Hamu and Maximilian Nickel and Matt Le},
    year={2023},
    eprint={2210.02747},
    archivePrefix={arXiv},
    primaryClass={cs.LG},
    url={https://arxiv.org/abs/2210.02747}, 
}

@misc{singh2024stochasticsamplingdeterministicflow,
      title={Stochastic Sampling from Deterministic Flow Models}, 
      author={Saurabh Singh and Ian Fischer},
      year={2024},
      eprint={2410.02217},
      archivePrefix={arXiv},
      primaryClass={cs.LG},
      url={https://arxiv.org/abs/2410.02217}, 
}

@misc{yang2024consistencyflowmatchingdefining,
      title={Consistency Flow Matching: Defining Straight Flows with Velocity Consistency}, 
      author={Ling Yang and Zixiang Zhang and Zhilong Zhang and Xingchao Liu and Minkai Xu and Wentao Zhang and Chenlin Meng and Stefano Ermon and Bin Cui},
      year={2024},
      eprint={2407.02398},
      archivePrefix={arXiv},
      primaryClass={cs.CV},
      url={https://arxiv.org/abs/2407.02398}, 
}

@misc{peebles2023scalablediffusionmodelstransformers,
      title={Scalable Diffusion Models with Transformers}, 
      author={William Peebles and Saining Xie},
      year={2023},
      eprint={2212.09748},
      archivePrefix={arXiv},
      primaryClass={cs.CV},
      url={https://arxiv.org/abs/2212.09748}, 
}