Reflected Flow Matching (2405.16577v1)
Abstract: Continuous normalizing flows (CNFs) learn an ordinary differential equation to transform prior samples into data. Flow matching (FM) has recently emerged as a simulation-free approach for training CNFs by regressing a velocity model towards the conditional velocity field. However, on constrained domains, the learned velocity model may lead to undesirable flows that result in highly unnatural samples, e.g., oversaturated images, due to both flow matching error and simulation error. To address this, we add a boundary constraint term to CNFs, which leads to reflected CNFs that keep trajectories within the constrained domains. We propose reflected flow matching (RFM) to train the velocity model in reflected CNFs by matching the conditional velocity fields in a simulation-free manner, similar to the vanilla FM. Moreover, the analytical form of conditional velocity fields in RFM avoids potentially biased approximations, making it superior to existing score-based generative models on constrained domains. We demonstrate that RFM achieves comparable or better results on standard image benchmarks and produces high-quality class-conditioned samples under high guidance weight.
- Building normalizing flows with stochastic interpolants. In The Eleventh International Conference on Learning Representations, 2022.
- Wasserstein gan. ArXiv, abs/1701.07875, 2017. URL https://api.semanticscholar.org/CorpusID:13943041.
- Riemannian score-based generative modelling. In Neural Information Processing Systems, 2022.
- Language models are few-shot learners. Advances in neural information processing systems, 33:1877–1901, 2020.
- Cattiaux, P. Time reversal of diffusion processes with a boundary condition. Stochastic processes and their Applications, 28(2):275–292, 1988.
- Neural ordinary differential equations. Advances in Neural Information Processing Systems, 2018.
- Flow matching in latent space. arXiv preprint arXiv:2307.08698, 2023.
- BERT: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv:1810.04805, 2018.
- Diffusion models beat gans on image synthesis. Advances in neural information processing systems, 34:8780–8794, 2021.
- A family of embedded runge-kutta formulae. Journal of computational and applied mathematics, 6(1):19–26, 1980.
- Diffusion models for constrained domains. Transactions on Machine Learning Research, 2023. ISSN 2835-8856.
- Gaussian error linear units (GELUs). arXiv preprint arXiv:1606.08415, 2016.
- Classifier-free diffusion guidance. arXiv preprint arXiv:2207.12598, 2022.
- Denoising diffusion probabilistic models. Advances in Neural Information Processing Systems, 33:6840–6851, 2020.
- Equivariant diffusion for molecule generation in 3d. In International conference on machine learning, pp. 8867–8887. PMLR, 2022.
- Junction tree variational autoencoder for molecular graph generation. In International conference on machine learning, pp. 2323–2332. PMLR, 2018.
- Subspace diffusion generative models. In European Conference on Computer Vision, pp. 274–289. Springer, 2022a.
- Torsional diffusion for molecular conformer generation. In Advances in Neural Information Processing Systems, 2022b.
- Adam: A method for stochastic optimization. In ICLR, 2015.
- Equivariant flow matching. In Thirty-seventh Conference on Neural Information Processing Systems, 2023.
- Numerical solution of stochastic differential equations. Springer, 1992.
- Lamperski, A. Projected stochastic gradient langevin algorithms for constrained sampling and non-convex learning. In Conference on Learning Theory, pp. 2891–2937. PMLR, 2021.
- Flow matching for generative modeling. In The Eleventh International Conference on Learning Representations, 2023. URL https://openreview.net/forum?id=PqvMRDCJT9t.
- Mirror diffusion models for constrained and watermarked generation. In Thirty-seventh Conference on Neural Information Processing Systems, 2023. URL https://openreview.net/forum?id=XPWEtXzlLy.
- Decoupled weight decay regularization. In International Conference on Learning Representations, 2018.
- Reflected diffusion models. arXiv preprint arXiv:2304.04740, 2023.
- Pytorch: An imperative style, high-performance deep learning library. In Neural Information Processing Systems, 2019.
- Pilipenko, A. An introduction to stochastic differential equations with reflection, volume 1. Universitätsverlag Potsdam, 2014.
- Variational inference with normalizing flows. In Proceedings of The 32nd International Conference on Machine Learning, pp. 1530–1538, 2015.
- U-net: Convolutional networks for biomedical image segmentation. In International Conference on Medical image computing and computerassisted intervention. Springer, 2015.
- Schuss, Z. Brownian dynamics at boundaries and interfaces. Springer, 2015.
- Seitzer, M. pytorch-fid: FID Score for PyTorch. https://github.com/mseitzer/pytorch-fid, August 2020. Version 0.3.0.
- Denoising diffusion implicit models. arXiv preprint arXiv:2010.02502, 2020a.
- Score-based generative modeling through stochastic differential equations. arXiv preprint arXiv:2011.13456, 2020b.
- Consistency models. arXiv preprint arXiv:2303.01469, 2023.
- Szabó, Z. Information theoretical estimators toolbox. The Journal of Machine Learning Research, 15(1):283–287, 2014.
- Improving and generalizing flow-based generative models with minibatch optimal transport. In ICML Workshop on New Frontiers in Learning, Control, and Dynamical Systems, 2023.
- Attention is all you need. Advances in neural information processing systems, 30, 2017.
- Vincent, P. A connection between score matching and denoising autoencoders. Neural Computation, 23(7):1661–1674, 2011. doi: 10.1162/NECO˙a˙00142.
- Flow matching for scalable simulation-based inference. In ICML 2023 Workshop on Structured Probabilistic Inference & Generative Modeling, 2023.
- Williams, R. On time-reversal of reflected brownian motions. In Seminar on Stochastic Processes, 1987, pp. 265–276. Springer, 1987.
- Geodiff: A geometric diffusion model for molecular conformation generation. In International Conference on Learning Representations, 2021.
- Opt: Open pre-trained transformer language models. arXiv preprint arXiv:2205.01068, 2022.
- Guided flows for generative modeling and decision making. ArXiv, abs/2311.13443, 2023.