Light Schrödinger Bridge (2310.01174v3)
Abstract: Despite the recent advances in the field of computational Schr\"odinger Bridges (SB), most existing SB solvers are still heavy-weighted and require complex optimization of several neural networks. It turns out that there is no principal solver which plays the role of simple-yet-effective baseline for SB just like, e.g., $k$-means method in clustering, logistic regression in classification or Sinkhorn algorithm in discrete optimal transport. We address this issue and propose a novel fast and simple SB solver. Our development is a smart combination of two ideas which recently appeared in the field: (a) parameterization of the Schr\"odinger potentials with sum-exp quadratic functions and (b) viewing the log-Schr\"odinger potentials as the energy functions. We show that combined together these ideas yield a lightweight, simulation-free and theoretically justified SB solver with a simple straightforward optimization objective. As a result, it allows solving SB in moderate dimensions in a matter of minutes on CPU without a painful hyperparameter selection. Our light solver resembles the Gaussian mixture model which is widely used for density estimation. Inspired by this similarity, we also prove an important theoretical result showing that our light solver is a universal approximator of SBs. Furthemore, we conduct the analysis of the generalization error of our light solver. The code for our solver can be found at https://github.com/ngushchin/LightSB
- Brandon Amos. On amortizing convex conjugates for optimal transport. In The Eleventh International Conference on Learning Representations, 2022.
- Neural optimal transport with general cost functionals. In The Twelfth International Conference on Learning Representations, 2024. URL https://openreview.net/forum?id=gIiz7tBtYZ.
- Existence, duality, and cyclical monotonicity for weak transport costs. Calculus of Variations and Partial Differential Equations, 58(6):1–28, 2019.
- Learning single-cell perturbation responses using neural optimal transport. bioRxiv, pp. 2021–12, 2021.
- Proximal optimal transport modeling of population dynamics. In International Conference on Artificial Intelligence and Statistics, pp. 6511–6528. PMLR, 2022.
- The schrödinger bridge between gaussian measures has a closed form. In International Conference on Artificial Intelligence and Statistics, pp. 5802–5833. PMLR, 2023.
- Likelihood training of schrödinger bridge using forward-backward sdes theory. In International Conference on Learning Representations, 2021a.
- Optimal steering of a linear stochastic system to a final probability distribution, part i. IEEE Transactions on Automatic Control, 61(5):1158–1169, 2015.
- On the relation between optimal transport and schrödinger bridges: A stochastic control viewpoint. Journal of Optimization Theory and Applications, 169:671–691, 2016.
- Stochastic control liaisons: Richard sinkhorn meets gaspard monge on a schrodinger bridge. SIAM Review, 63(2):249–313, 2021b.
- Provably convergent schrödinger bridge with applications to probabilistic time series imputation. 2023.
- Generative modeling through the semi-dual formulation of unbalanced optimal transport. arXiv preprint arXiv:2305.14777, 2023.
- Marco Cuturi. Sinkhorn distances: Lightspeed computation of optimal transport. Advances in neural information processing systems, 26, 2013.
- Score-based generative neural networks for large-scale optimal transport. Advances in neural information processing systems, 34:12955–12965, 2021.
- Diffusion schrödinger bridge with applications to score-based generative modeling. Advances in Neural Information Processing Systems, 34:17695–17709, 2021.
- Rates of estimation of optimal transport maps using plug-in estimators via barycentric projections. Advances in Neural Information Processing Systems, 34:29736–29753, 2021.
- Maximum likelihood from incomplete data via the em algorithm. Journal of the royal statistical society: series B (methodological), 39(1):1–22, 1977.
- Computational optimal transport: Complexity by accelerated gradient descent is better than by sinkhorn’s algorithm. In International conference on machine learning, pp. 1367–1376. PMLR, 2018.
- Neural monge map estimation and its applications. Transactions on Machine Learning Research, 2023. ISSN 2835-8856. URL https://openreview.net/forum?id=2mZSlQscj3. Featured Certification.
- Learning with minibatch wasserstein: asymptotic and gradient properties. In International Conference on Artificial Intelligence and Statistics, pp. 2131–2141. PMLR, 2020.
- Pot: Python optimal transport. Journal of Machine Learning Research, 22(78):1–8, 2021. URL http://jmlr.org/papers/v22/20-451.html.
- Robert Fortet. Résolution d’un système d’équations de m. schrödinger. Journal de Mathématiques Pures et Appliquées, 19(1-4):83–105, 1940.
- Extremal domain translation with neural optimal transport. In Thirty-seventh Conference on Neural Information Processing Systems, 2023. URL https://openreview.net/forum?id=vZRiMjo826.
- Aude Genevay. Entropy-regularized optimal transport for machine learning. PhD thesis, Paris Sciences et Lettres (ComUE), 2019.
- Stochastic optimization for large-scale optimal transport. In Advances in neural information processing systems, pp. 3440–3448, 2016.
- Kantorovich duality for general transport costs and applications. Journal of Functional Analysis, 273(11):3327–3405, 2017.
- Entropic neural optimal transport via diffusion processes. In Advances in Neural Information Processing Systems, 2023a.
- Building the bridge of schr\\\backslash\” odinger: A continuous entropic optimal transport benchmark. In Thirty-seventh Conference on Neural Information Processing Systems Datasets and Benchmarks Track, 2023b.
- Pierre Henry-Labordere. (martingale) optimal transport and anomaly detection with neural networks: A primal-dual algorithm. Available at SSRN 3370910, 2019.
- Denoising diffusion probabilistic models. Advances in Neural Information Processing Systems, 33:6840–6851, 2020.
- Minimax estimation of smooth optimal transport maps. 2021.
- Entropic optimal transport between unbalanced gaussian measures has a closed form. Advances in neural information processing systems, 33:10468–10479, 2020.
- A style-based generator architecture for generative adversarial networks. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp. 4401–4410, 2019.
- Unpaired image-to-image translation via neural schrödinger bridge. In The Twelfth International Conference on Learning Representations, 2024. URL https://openreview.net/forum?id=uQBW7ELXfO.
- Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6980, 2014.
- Wasserstein-2 generative networks. In International Conference on Learning Representations, 2021a. URL https://openreview.net/forum?id=bEoxzW_EXsa.
- Do neural optimal transport solvers work? a continuous wasserstein-2 benchmark. Advances in Neural Information Processing Systems, 34:14593–14605, 2021b.
- Wasserstein iterative networks for barycenter estimation. In Thirty-Sixth Conference on Neural Information Processing Systems, 2022a. URL https://openreview.net/forum?id=GiEnzxTnaMN.
- Kantorovich strikes back! wasserstein GANs are not optimal transport? In Thirty-sixth Conference on Neural Information Processing Systems Datasets and Benchmarks Track, 2022b. URL https://openreview.net/forum?id=VtEEpi-dGlt.
- Kernel neural optimal transport. In International Conference on Learning Representations, 2023a. URL https://openreview.net/forum?id=Zuc_MHtUma4.
- Neural optimal transport. In International Conference on Learning Representations, 2023b. URL https://openreview.net/forum?id=d8CBRlWNkqH.
- Neural lagrangian schr\\\backslash\”{{\{{o}}\}} dinger bridge: Diffusion modeling for population dynamics. In The Eleventh International Conference on Learning Representations, 2022.
- Solomon Kullback. Probability densities with given marginals. The Annals of Mathematical Statistics, 39(4):1236–1243, 1968.
- The effect of the intrinsic dimension on the generalization of quadratic classifiers. In A. Beygelzimer, Y. Dauphin, P. Liang, and J. Wortman Vaughan (eds.), Advances in Neural Information Processing Systems, 2021. URL https://openreview.net/forum?id=_hKvtsqItc.
- A tutorial on energy-based learning. Predicting structured data, 1(0), 2006.
- Christian Léonard. A survey of the schr\\\backslash\” odinger problem and some of its connections with optimal transport. arXiv preprint arXiv:1308.0215, 2013.
- I22{}^{2}start_FLOATSUPERSCRIPT 2 end_FLOATSUPERSCRIPTsb: Image-to-image schr\\\backslash\” odinger bridge. arXiv preprint arXiv:2302.05872, 2023.
- Flow straight and fast: Learning to generate and transfer data with rectified flow. In The Eleventh International Conference on Learning Representations, 2022.
- Optimal transport mapping via input convex neural networks. In International Conference on Machine Learning, pp. 6672–6681. PMLR, 2020.
- Entropy-regularized 2-wasserstein distance between gaussian measures. Information Geometry, 5(1):289–323, 2022.
- Plugin estimation of smooth optimal transport maps. arXiv preprint arXiv:2107.12364, 2021.
- Foundations of machine learning. MIT press, 2018.
- Energy-guided entropic neural optimal transport. In The Twelfth International Conference on Learning Representations, 2024. URL https://openreview.net/forum?id=d6tUsZeVs7.
- On unbalanced optimal transport: Gradient methods, sparsity and approximation error. arXiv preprint arXiv:2202.03618, 2022.
- Approximation by finite mixtures of continuous density functions that vanish at infinity. Cogent Mathematics & Statistics, 7(1):1750861, 2020.
- Unbalanced diffusion schr\\\backslash\” odinger bridge. arXiv preprint arXiv:2306.09099, 2023.
- The matrix cookbook. Technical University of Denmark, 7(15):510, 2008.
- Computational optimal transport. Foundations and Trends® in Machine Learning, 11(5-6):355–607, 2019.
- Adversarial latent autoencoders. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 14104–14113, 2020.
- 8 - brownian motion and related processes. In Mark A. Pinsky and Samuel Karlin (eds.), An Introduction to Stochastic Modeling (Fourth Edition), pp. 391–446. Academic Press, Boston, fourth edition edition, 2011. ISBN 978-0-12-381416-6. doi: https://doi.org/10.1016/B978-0-12-381416-6.00008-3. URL https://www.sciencedirect.com/science/article/pii/B9780123814166000083.
- Entropic estimation of optimal transport maps. arXiv preprint arXiv:2109.12004, 2021.
- Variational inference with normalizing flows. In International conference on machine learning, pp. 1530–1538. PMLR, 2015.
- On the sample complexity of entropic optimal transport. arXiv preprint arXiv:2206.13472, 2022.
- Energy distance. wiley interdisciplinary reviews: Computational statistics, 8(1):27–38, 2016.
- High-resolution image synthesis with latent diffusion models. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 10684–10695, 2022.
- Generative modeling with optimal transport maps. In International Conference on Learning Representations, 2021.
- Ludger Ruschendorf. Convergence of the iterative proportional fitting procedure. The Annals of Statistics, pp. 1160–1174, 1995.
- Erwin Schrödinger. Über die umkehrung der naturgesetze. Verlag der Akademie der Wissenschaften in Kommission bei Walter De Gruyter u …, 1931.
- Erwin Schrödinger. Sur la théorie relativiste de l’électron et l’interprétation de la mécanique quantique. In Annales de l’institut Henri Poincaré, volume 2, pp. 269–310, 1932.
- Large scale optimal transport and mapping estimation. In International Conference on Learning Representations, 2018.
- Understanding machine learning: From theory to algorithms. Cambridge university press, 2014.
- Diffusion schrödinger bridge matching. In Thirty-seventh Conference on Neural Information Processing Systems, 2023. URL https://openreview.net/forum?id=qy07OHsJT5.
- Aligned diffusion schr\\\backslash\” odinger bridges. arXiv preprint arXiv:2302.11419, 2023.
- How to train your energy-based models. arXiv preprint arXiv:2101.03288, 2021.
- Austin Stromme. Sampling from a schrödinger bridge. In International Conference on Artificial Intelligence and Statistics, pp. 4058–4067. PMLR, 2023.
- 2-Wasserstein approximation via restricted convex potentials with application to improved training for GANs. arXiv preprint arXiv:1902.07197, 2019.
- Trajectorynet: A dynamic optimal transport network for modeling cellular dynamics. In International conference on machine learning, pp. 9526–9536. PMLR, 2020.
- Simulation-free schr\\\backslash\” odinger bridges via score and flow matching. arXiv preprint arXiv:2307.03672, 2023.
- Transport, variational inference and diffusions: with applications to annealed flows and schr\\\backslash\” odinger bridges. arXiv preprint arXiv:2307.01050, 2023.
- Solving schrödinger bridges via maximum likelihood. Entropy, 23(9):1134, 2021.
- Deep generative learning via schrödinger bridge. In International Conference on Machine Learning, pp. 10794–10804. PMLR, 2021.
- An accelerated stochastic algorithm for solving the optimal transport problem. arXiv preprint arXiv:2203.00813, 2022.
- Scalable unbalanced optimal transport using generative adversarial networks. In International Conference on Learning Representations, 2018.
- Unpaired image-to-image translation using cycle-consistent adversarial networks. In Proceedings of the IEEE international conference on computer vision, pp. 2223–2232, 2017.