Diffusion bridge mixture transports, schrödinger bridge problems and generative modeling
Article No.: 374, Pages 17958 - 18008
Abstract
The dynamic Schrödinger bridge problem seeks a stochastic process that defines a transport between two target probability measures, while optimally satisfying the criteria of being closest, in terms of Kullback-Leibler divergence, to a reference process. We propose a novel sampling-based iterative algorithm, the iterated diffusion bridge mixture (IDBM) procedure, aimed at solving the dynamic Schrödinger bridge problem. The IDBM procedure exhibits the attractive property of realizing a valid transport between the target probability measures at each iteration. We perform an initial theoretical investigation of the IDBM procedure, establishing its convergence properties. The theoretical findings are complemented by numerical experiments illustrating the competitive performance of the IDBM procedure. Recent advancements in generative modeling employ the time-reversal of a diffusion process to define a generative process that approximately transports a simple distribution to the data distribution. As an alternative, we propose utilizing the first iteration of the IDBM procedure as an approximation-free method for realizing this transport. This approach offers greater flexibility in selecting the generative process dynamics and exhibits accelerated training and superior sample quality over larger discretization intervals. In terms of implementation, the necessary modifications are minimally intrusive, being limited to the training loss definition.
References
[1]
Brian D.O. Anderson. Reverse-Time Diffusion Equation Models. Stochastic Processes and their Applications, 12(3):313-326, May 1982.
[2]
Yacine Aït-Sahalia. Closed-form likelihood expansions for multivariate diffusions. The Annals of Statistics, 36(2):906-937, 2008. ISSN 0090-5364, 2168-8966.
[3]
Espen Bernton, Jeremy Heng, Arnaud Doucet, and Pierre E. Jacob. Schrödinger Bridge Samplers, 2019.
[4]
Alexandros Beskos and Gareth O. Roberts. Exact simulation of diffusions. The Annals of Applied Probability, 15(4):2422-2444, 2005. ISSN 1050-5164, 2168-8737.
[5]
Alexandros Beskos, Omiros Papaspiliopoulos, and Gareth O. Roberts. A Factorisation of Diffusion Measure and Finite Sample Path Constructions. Methodology and Computing in Applied Probability, 10(1):85-104, 2008. ISSN 1573-7713.
[6]
Patrick Billingsley. Convergence of Probability Measures. Wiley Series in Probability and Statistics. Probability and Statistics Section. Wiley, 2nd ed edition, 1999.
[7]
Valentin De Bortoli, James Thornton, Jeremy Heng, and Arnaud Doucet. Diffusion Schrödinger Bridge with Applications to Score-Based Generative Modeling. In Thirty-Fifth Conference on Neural Information Processing Systems, 2021.
[8]
Damiano Brigo. The General Mixture-Diffusion SDE and Its Relationship with an Uncertain-Volatility Option Model with Volatility-Asset Decorrelation, December 2002.
[9]
Steve Brooks, Andrew Gelman, Galin Jones, and Xiao-Li Meng. Handbook of Markov Chain Monte Carlo. CRC press, 2011.
[10]
Tianrong Chen, Guan-Horng Liu, and Evangelos Theodorou. Likelihood Training of Schrödinger Bridge using Forward-Backward SDEs Theory. In International Conference on Learning Representations, 2022.
[11]
I. Csiszar. I-Divergence Geometry of Probability Distributions and Minimization Problems. The Annals of Probability, 3(1):146-158, 1975. ISSN 0091-1798.
[12]
Paolo Dai Pra. A stochastic control approach to reciprocal diffusion processes. Applied Mathematics and Optimization, 23(1):313-329, 1991. ISSN 1432-0606.
[13]
Valentin De Bortoli, Arnaud Doucet, Jeremy Heng, and James Thornton. Simulating Diffusion Bridges with Score Matching, 2021.
[14]
W. Edwards Deming and Frederick F. Stephan. On a Least Squares Adjustment of a Sampled Frequency Table When the Expected Marginal Totals are Known. The Annals of Mathematical Statistics, 11(4):427-444, 1940. ISSN 0003-4851.
[15]
Prafulla Dhariwal and Alexander Quinn Nichol. Diffusion Models Beat GANs on Image Synthesis. In A. Beygelzimer, Y. Dauphin, P. Liang, and J. Wortman Vaughan, editors, Advances in Neural Information Processing Systems, 2021.
[16]
Pavel Dvurechensky, Alexander Gasnikov, and Alexey Kroshnin. Computational Optimal Transport: Complexity by Accelerated Gradient Descent Is Better Than by Sinkhorn's Algorithm. In Proceedings of the 35th International Conference on Machine Learning, pages 1367-1376. PMLR, 2018.
[17]
Tarek A. El Moselhy and Youssef M. Marzouk. Bayesian inference with optimal maps. Journal of Computational Physics, 231(23):7815-7850, 2012. ISSN 0021-9991.
[18]
David Lopes Fernandes, Francisco Vargas, Carl Henrik Ek, and Neill D. F. Campbell. Shooting Schrödinger's Cat. In Fourth Symposium on Advances in Approximate Bayesian Inference, 2022.
[19]
Rémi Flamary, Nicolas Courty, Alexandre Gramfort, Mokhtar Z. Alaya, Aurélie Boisbunon, Stanislas Chambon, Laetitia Chapel, Adrien Corenflos, Kilian Fatras, Nemo Fournier, Léo Gautheron, Nathalie T. H. Gayraud, Hicham Janati, Alain Rakotomamonjy, Ievgen Redko, Antoine Rolet, Antony Schutz, Vivien Seguy, Danica J. Sutherland, Romain Tavenard, Alexander Tong, and Titouan Vayer. POT: Python Optimal Transport. Journal of Machine Learning Research, 22(78):1-8, 2021. ISSN 1533-7928.
[20]
Robert Fortet. Résolution d'un systeme d'équations de M. Schrödinger. J. Math. Pure Appl. IX, 1:83-105, 1940.
[21]
Avner Friedman. Stochastic Differential Equations and Applications - Volume 1. Probability and Mathematical Statistics Series ; v. 28. Academic Press, 1975.
[22]
Christiane Fuchs. Inference for Diffusion Processes: With Applications in Life Sciences. Springer Science & Business Media, 2013.
[23]
U. G. Haussmann and E. Pardoux. Time Reversal of Diffusions. The Annals of Probability, 14(4):1188-1205, October 1986.
[24]
Martin Heusel, Hubert Ramsauer, Thomas Unterthiner, Bernhard Nessler, and Sepp Hochreiter. GANs Trained by a Two Time-Scale Update Rule Converge to a Local Nash Equilibrium. In Advances in Neural Information Processing Systems, volume 30. Curran Associates, Inc., 2017.
[25]
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising Diffusion Probabilistic Models. In H. Larochelle, M. Ranzato, R. Hadsell, M. F. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems, volume 33, pages 6840-6851, 2020.
[26]
Emiel Hoogeboom, Jonathan Heek, and Tim Salimans. Simple diffusion: End-to-end diffusion for high resolution images, 2023.
[27]
Jian Huang, Yuling Jiao, Lican Kang, Xu Liao, Jin Liu, and Yanyan Liu. Schrödinger-Föllmer Sampler: Sampling without Ergodicity, 2021.
[28]
A. S. Hurn, J. I. Jeisman, and K. A. Lindsay. Seeing the Wood for the Trees: A Critical Evaluation of Methods to Estimate the Parameters of Stochastic Differential Equations. Journal of Financial Econometrics, 5(3):390-455, 2007. ISSN 1479-8409.
[29]
Aapo Hyvärinen. Estimation of Non-Normalized Statistical Models by Score Matching. Journal of Machine Learning Research, 6(24):695-709, 2005. ISSN 1533-7928.
[30]
Benton Jamison. Reciprocal processes. Zeitschrift für Wahrscheinlichkeitstheorie und Verwandte Gebiete, 30(1):65-86, 1974. ISSN 1432-2064.
[31]
Benton Jamison. The Markov processes of Schrödinger. Zeitschrift für Wahrscheinlichkeitstheorie und Verwandte Gebiete, 32(4):323-331, 1975. ISSN 1432-2064.
[32]
Hicham Janati, Boris Muzellec, Gabriel Peyré, and Marco Cuturi. Entropic Optimal Transport between Unbalanced Gaussian Measures has a Closed Form. In H. Larochelle, M. Ranzato, R. Hadsell, M. F. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems, volume 33, pages 10468-10479. Curran Associates, Inc., 2020.
[33]
Mathieu Kessler, Alexander Lindner, and Michael Sorensen. Statistical Methods for Stochastic Differential Equations. Chapman and Hall/CRC, 2012.
[34]
Peter E. Kloeden and Eckhard Platen. Numerical Solution of Stochastic Differential Equations. Springer Berlin Heidelberg, Berlin, Heidelberg, 1992. ISBN 978-3-642-08107-1 978-3-662-12616-5.
[35]
Yaron Lipman, Ricky T. Q. Chen, Heli Ben-Hamu, Maximilian Nickel, and Matthew Le. Flow Matching for Generative Modeling. In The Eleventh International Conference on Learning Representations, 2023.
[36]
Robert Shevilevich Liptser and Al'bert Nikolaevich Shiriaev. Statistics of Random Processes: General Theory, volume 394. Springer, 1977.
[37]
Guan-Horng Liu, Arash Vahdat, De-An Huang, Evangelos A. Theodorou, Weili Nie, and Anima Anandkumar. I2SB: Image-to-Image Schrödinger Bridge, 2023a.
[38]
Qiang Liu. Rectified Flow: A Marginal Preserving Approach to Optimal Transport, 2022.
[39]
Xingchao Liu, Chengyue Gong, and Qiang Liu. Flow Straight and Fast: Learning to Generate and Transfer Data with Rectified Flow, 2022.
[40]
Xingchao Liu, Lemeng Wu, Mao Ye, and Qiang Liu. Learning Diffusion Bridges on Constrained Domains. In The Eleventh International Conference on Learning Representations, 2023b.
[41]
Christian Léonard. Some Properties of Path Measures. In Catherine Donati-Martin, Antoine Lejay, and Alain Rouault, editors, Séminaire de Probabilités XLVI, Lecture Notes in Mathematics, pages 207-230. Springer International Publishing, 2014a.
[42]
Christian Léonard. A survey of the Schrödinger problem and some of its connections with optimal transport. Discrete & Continuous Dynamical Systems, 34(4):1533, 2014b.
[43]
Anton Mallasto, Augusto Gerolin, and Hà Quang Minh. Entropy-regularized 2-Wasserstein distance between Gaussian measures. Information Geometry, 5(1):289-323, 2022. ISSN 2511-249X.
[44]
Youssef Marzouk, Tarek Moselhy, Matthew Parno, and Alessio Spantini. Sampling via Measure Transport: An Introduction. In Roger Ghanem, David Higdon, and Houman Owhadi, editors, Handbook of Uncertainty Quantification, pages 1-41. Springer International Publishing, 2016.
[45]
Annie Millet, David Nualart, and Marta Sanz. Integration by Parts and Time Reversal for Diffusion Processes. The Annals of Probability, pages 208-238, 1989.
[46]
Alexander Quinn Nichol and Prafulla Dhariwal. Improved Denoising Diffusion Probabilistic Models. In Proceedings of the 38th International Conference on Machine Learning, pages 8162-8171. PMLR, 2021.
[47]
Frank Nielsen. What is an information projection. Notices of the AMS, 65(3):321-324, 2018.
[48]
George Papamakarios, Eric Nalisnick, Danilo Jimenez Rezende, Shakir Mohamed, and Balaji Lakshminarayanan. Normalizing Flows for Probabilistic Modeling and Inference. Journal of Machine Learning Research, 22(57):1-64, 2021. ISSN 1533-7928.
[49]
Michele Pavon, Esteban G. Tabak, and Giulio Trigila. The data-driven Schrödinger bridge, 2018.
[50]
Stefano Peluchetti. Non-Denoising Forward-Time Diffusions. 2021.
[51]
Stefano Peluchetti and Gareth O. Roberts. A Study of the Efficiency of Exact Methods for Diffusion Simulation. In Leszek Plaskota and Henryk Woźniakowski, editors, Monte Carlo and Quasi-Monte Carlo Methods 2010, Springer Proceedings in Mathematics & Statistics, pages 161-187. Springer, 2012.
[52]
Gabriel Peyré and Marco Cuturi. Computational Optimal Transport. 2020.
[53]
L Chris G Rogers and David Williams. Diffusions, Markov Processes and Martingales: Volume 2: Itô Calculus, volume 2. Cambridge university press, 2000.
[54]
Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Björn Ommer. High-Resolution Image Synthesis with Latent Diffusion Models. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition (CVPR), 2022.
[55]
Ludger Ruschendorf. Convergence of the Iterative Proportional Fitting Procedure. The Annals of Statistics, 23(4):1160-1174, 1995. ISSN 0090-5364.
[56]
Hamza Ruzayqat, Alexandros Beskos, Dan Crisan, Ajay Jasra, and Nikolas Kantas. Unbiased estimation using a class of diffusion processes. Journal of Computational Physics, 472: 111643, 2023. ISSN 0021-9991.
[57]
L. Rüschendorf and W. Thomsen. Note on the Schrödinger equation and I-projections. Statistics & Probability Letters, 17(5):369-375, 1993. ISSN 0167-7152.
[58]
Yuyang Shi, Valentin De Bortoli, Andrew Campbell, and Arnaud Doucet. Diffusion Schrödinger Bridge Matching, 2023.
[59]
Vignesh Ram Somnath, Matteo Pariset, Ya-Ping Hsieh, Maria Rodriguez Martinez, Andreas Krause, and Charlotte Bunne. Aligned Diffusion Schrödinger Bridges. In Proceedings of the Thirty-Ninth Conference on Uncertainty in Artificial Intelligence, pages 1985-1995. PMLR, 2023.
[60]
Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-Based Generative Modeling through Stochastic Differential Equations. In International Conference on Learning Representations, 2021.
[61]
Daniel W Stroock and S. R. S Varadhan. Multidimensional Diffusion Processes. 2006.
[62]
Simo Särkkä and Arno Solin. Applied Stochastic Differential Equations. Cambridge University Press, first edition, April 2019. ISBN 978-1-108-18673-5 978-1-316-51008-7 978-1-316- 64946-6.
[63]
James Thornton, Michael Hutchinson, Emile Mathieu, Valentin De Bortoli, Yee Whye Teh, and Arnaud Doucet. Riemannian Diffusion Schrödinger Bridge, 2022.
[64]
Alexander Tong, Nikolay Malkin, Guillaume Huguet, Yanlei Zhang, Jarrid Rector-Brooks, Kilian Fatras, Guy Wolf, and Yoshua Bengio. Conditional Flow Matching: Simulation-Free Dynamic Optimal Transport, 2023.
[65]
Belinda Tzen and Maxim Raginsky. Theoretical guarantees for sampling and inference in generative models with latent diffusions. In Proceedings of the Thirty-Second Conference on Learning Theory, pages 3084-3114. PMLR, 2019.
[66]
Tim van Erven and Peter Harremos. Rényi Divergence and Kullback-Leibler Divergence. IEEE Transactions on Information Theory, 60(7):3797-3820, 2014. ISSN 1557-9654.
[67]
Francisco Vargas, Pierre Thodoroff, Austen Lamacraft, and Neil Lawrence. Solving Schrödinger Bridges via Maximum Likelihood. Entropy, 23(9):1134, September 2021.
[68]
Francisco Vargas, Andrius Ovsianas, David Fernandes, Mark Girolami, Neil D. Lawrence, and Nikolas Nüsken. Bayesian learning via neural Schrödinger-Föllmer flows. Statistics and Computing, 33(1):3, 2022. ISSN 1573-1375.
[69]
Pascal Vincent. A Connection Between Score Matching and Denoising Autoencoders. Neural Computation, 23(7):1661-1674, July 2011.
[70]
Gefei Wang, Yuling Jiao, Qian Xu, Yang Wang, and Can Yang. Deep Generative Learning via Schrödinger Bridge. In Marina Meila and Tong Zhang, editors, Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pages 10794-10804. PMLR, 2021.
[71]
Lemeng Wu, Chengyue Gong, Xingchao Liu, Mao Ye, and Qiang Liu. Diffusion-based Molecule Generation with Informative Prior Bridges. In Advances in Neural Information Processing Systems, 2022.
[72]
Mao Ye, Lemeng Wu, and Qiang Liu. First Hitting Diffusion Models for Generating Manifold, Graph and Categorical Data. In Advances in Neural Information Processing Systems, 2022.
[73]
Qinsheng Zhang and Yongxin Chen. Path Integral Sampler: A Stochastic Control Approach For Sampling. In International Conference on Learning Representations, 2022.
[74]
B. K. Øksendal. Stochastic Differential Equations: An Introduction with Applications. Universitext. Springer, Berlin ; New York, 6th ed. edition, 2003. ISBN 978-3-540-04758-2.
Index Terms
- Diffusion bridge mixture transports, schrödinger bridge problems and generative modeling
Index terms have been assigned to the content through auto-classification.
Recommendations
Diffusion schrödinger bridge matching
NIPS '23: Proceedings of the 37th International Conference on Neural Information Processing SystemsSolving transport problems, i.e. finding a map transporting one given distribution to another, has numerous applications in machine learning. Novel mass transport methods motivated by generative modeling have recently been proposed, e.g. Denoising ...
On the Relation Between Optimal Transport and Schrödinger Bridges: A Stochastic Control Viewpoint
We take a new look at the relation between the optimal transport problem and the Schrödinger bridge problem from a stochastic control perspective. Our aim is to highlight new connections between the two that are richer and deeper than those previously ...
Diffusion schrödinger bridge with applications to score-based generative modeling
NIPS '21: Proceedings of the 35th International Conference on Neural Information Processing SystemsProgressively applying Gaussian noise transforms complex data distributions to approximately Gaussian. Reversing this dynamic defines a generative model. When the forward noising process is given by a Stochastic Differential Equation (SDE), Song et al. (...
Comments
Information & Contributors
Information
Published In
January 2023
18881 pages
Copyright © 2023.
CC-BY 4.0
Publisher
JMLR.org
Publication History
Published: 06 March 2024
Accepted: 01 October 2023
Revised: 01 September 2023
Received: 01 April 2023
Published in JMLR Volume 24, Issue 1
Author Tags
Qualifiers
- Research-article
Contributors
Other Metrics
Bibliometrics & Citations
Bibliometrics
Article Metrics
- 0Total Citations
- 46Total Downloads
- Downloads (Last 12 months)46
- Downloads (Last 6 weeks)13
Reflects downloads up to 03 Oct 2024
Other Metrics
Citations
View Options
Get Access
Login options
Check if you have access through your login credentials or your institution to get full access on this article.
Sign in