Papers
Topics
Authors
Recent
Gemini 2.5 Flash
Gemini 2.5 Flash
194 tokens/sec
GPT-4o
7 tokens/sec
Gemini 2.5 Pro Pro
46 tokens/sec
o3 Pro
4 tokens/sec
GPT-4.1 Pro
38 tokens/sec
DeepSeek R1 via Azure Pro
28 tokens/sec
2000 character limit reached

Equivariant flow matching (2306.15030v2)

Published 26 Jun 2023 in stat.ML, cs.LG, physics.chem-ph, and physics.comp-ph

Abstract: Normalizing flows are a class of deep generative models that are especially interesting for modeling probability distributions in physics, where the exact likelihood of flows allows reweighting to known target energy functions and computing unbiased observables. For instance, Boltzmann generators tackle the long-standing sampling problem in statistical physics by training flows to produce equilibrium samples of many-body systems such as small molecules and proteins. To build effective models for such systems, it is crucial to incorporate the symmetries of the target energy into the model, which can be achieved by equivariant continuous normalizing flows (CNFs). However, CNFs can be computationally expensive to train and generate samples from, which has hampered their scalability and practical application. In this paper, we introduce equivariant flow matching, a new training objective for equivariant CNFs that is based on the recently proposed optimal transport flow matching. Equivariant flow matching exploits the physical symmetries of the target energy for efficient, simulation-free training of equivariant CNFs. We demonstrate the effectiveness of flow matching on rotation and permutation invariant many-particle systems and a small molecule, alanine dipeptide, where for the first time we obtain a Boltzmann generator with significant sampling efficiency without relying on tailored internal coordinate featurization. Our results show that the equivariant flow matching objective yields flows with shorter integration paths, improved sampling efficiency, and higher scalability compared to existing methods.

Definition Search Book Streamline Icon: https://streamlinehq.com
References (74)
  1. Semi-supervised learning with deep generative models. Advances in neural information processing systems, 27, 2014.
  2. Density estimation using real NVP. In International Conference on Learning Representations, 2017.
  3. Glow: Generative flow with invertible 1x1 convolutions. In S. Bengio, H. Wallach, H. Larochelle, K. Grauman, N. Cesa-Bianchi, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 31. Curran Associates, Inc., 2018.
  4. Analyzing and improving the image quality of stylegan. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pages 8110–8119, 2020.
  5. Attention is all you need. Advances in Neural Information Processing Systems, 30, 2017.
  6. BERT: Pre-training of deep bidirectional transformers for language understanding. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume 1 (Long and Short Papers), pages 4171–4186, Minneapolis, Minnesota, June 2019. Association for Computational Linguistics.
  7. Gpt-3: Its nature, scope, limits, and consequences. Minds and Machines, 30:681–694, 2020.
  8. Boltzmann generators — sampling equilibrium states of many-body systems with deep learning. Science, 365:eaaw1147, 2019.
  9. Highly accurate protein structure prediction with alphafold. Nature, 596(7873):583–589, 2021.
  10. Broadly applicable and accurate protein design by integrating structure prediction networks and diffusion generative models. bioRxiv, 2022.
  11. Density estimation by dual ascent of the log-likelihood. Communications in Mathematical Sciences, 8(1):217–233, 2010.
  12. A family of nonparametric density estimation algorithms. Communications on Pure and Applied Mathematics, 66(2):145–164, 2013.
  13. Normalizing flows for probabilistic modeling and inference. CoRR, 2019.
  14. Variational inference with normalizing flows. In International conference on machine learning, pages 1530–1538. PMLR, 2015.
  15. Equivariant flow-based sampling for lattice gauge theory. Phys. Rev. Lett., 125:121601, Sep 2020.
  16. Equivariant manifold flows. In M. Ranzato, A. Beygelzimer, Y. Dauphin, P.S. Liang, and J. Wortman Vaughan, editors, Advances in Neural Information Processing Systems, volume 34, pages 10600–10612. Curran Associates, Inc., 2021.
  17. Equivariant flows: exact likelihood generative learning for symmetric densities. In International conference on machine learning, pages 5361–5370. PMLR, 2020.
  18. E(n) equivariant normalizing flows. In M. Ranzato, A. Beygelzimer, Y. Dauphin, P.S. Liang, and J. Wortman Vaughan, editors, Advances in Neural Information Processing Systems, volume 34, pages 4181–4192. Curran Associates, Inc., 2021.
  19. Rigid body flows for sampling molecular crystal structures. In International Conference on Machine Learning, ICML 2023, volume 202 of Proceedings of Machine Learning Research, pages 17301–17326. PMLR, 2023.
  20. Smooth normalizing flows. In M. Ranzato, A. Beygelzimer, Y. Dauphin, P.S. Liang, and J. Wortman Vaughan, editors, Advances in Neural Information Processing Systems, volume 34, pages 2796–2809. Curran Associates, Inc., 2021.
  21. Flow matching for generative modeling. In The Eleventh International Conference on Learning Representations, 2023.
  22. Conditional flow matching: Simulation-free dynamic optimal transport. arXiv preprint arXiv:2302.00482, 2023.
  23. Equivariant flow matching with hybrid probability transport for 3d molecule generation. In Thirty-seventh Conference on Neural Information Processing Systems, 2023.
  24. SE(3) equivariant augmented coupling flows. In Thirty-seventh Conference on Neural Information Processing Systems, 2023.
  25. Temperature steerable flows and Boltzmann generators. Phys. Rev. Res., 4:L042005, Oct 2022.
  26. Targeted free energy estimation via learned mappings. The Journal of Chemical Physics, 153(14):144112, 2020.
  27. Flow annealed importance sampling bootstrap. In The Eleventh International Conference on Learning Representations, 2023.
  28. Deepbar: A fast and exact method for binding free energy computation. Journal of Physical Chemistry Letters, 12:2509–2515, 3 2021.
  29. Stochastic normalizing flows. In H. Larochelle, M. Ranzato, R. Hadsell, M. F. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems, volume 33, pages 5933–5944. Curran Associates, Inc., 2020.
  30. Computing absolute free energy with deep generative models. Biophysical Journal, 120(3):195a, 2021.
  31. Multimap targeted free energy estimation, 2023.
  32. Targeted free energy perturbation revisited: Accurate free energies from mapped reference potentials. Journal of Physical Chemistry Letters, 12:9449–9454, 2021.
  33. Skipping the replica exchange ladder with normalizing flows. The Journal of Physical Chemistry Letters, 13:11643–11649, 2022.
  34. Timewarp: Transferable acceleration of molecular dynamics by learning time-coarsened dynamics. In Thirty-seventh Conference on Neural Information Processing Systems, 2023.
  35. Equivariant diffusion for molecule generation in 3D. In Kamalika Chaudhuri, Stefanie Jegelka, Le Song, Csaba Szepesvari, Gang Niu, and Sivan Sabato, editors, Proceedings of the 39th International Conference on Machine Learning, volume 162 of Proceedings of Machine Learning Research, pages 8867–8887. PMLR, 17–23 Jul 2022.
  36. Geodiff: A geometric diffusion model for molecular conformation generation. In International Conference on Learning Representations, 2022.
  37. Torsional diffusion for molecular conformer generation. In S. Koyejo, S. Mohamed, A. Agarwal, D. Belgrave, K. Cho, and A. Oh, editors, Advances in Neural Information Processing Systems, volume 35, pages 24240–24253. Curran Associates, Inc., 2022.
  38. Implicit transfer operator learning: Multiple time-resolution models for molecular dynamics. In Thirty-seventh Conference on Neural Information Processing Systems, 2023.
  39. How to train your neural ode: the world of jacobian and kinetic regularization. In International conference on machine learning, pages 3154–3164. PMLR, 2020.
  40. Trajectorynet: A dynamic optimal transport network for modeling cellular dynamics. In International conference on machine learning, pages 9526–9536. PMLR, 2020.
  41. Ot-flow: Fast and accurate continuous normalizing flows via optimal transport. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 35, pages 9223–9232, 2021.
  42. Building normalizing flows with stochastic interpolants. In The Eleventh International Conference on Learning Representations, 2023.
  43. Flow straight and fast: Learning to generate and transfer data with rectified flow. In The Eleventh International Conference on Learning Representations, 2023.
  44. Multisample flow matching: Straightening flows with minibatch couplings. arXiv preprint arXiv:2304.14772, 2023.
  45. Towards optimal transport with global invariances. In Kamalika Chaudhuri and Masashi Sugiyama, editors, Proceedings of the Twenty-Second International Conference on Artificial Intelligence and Statistics, volume 89 of Proceedings of Machine Learning Research, pages 1870–1879. PMLR, 16–18 Apr 2019.
  46. Flow-matching: Efficient coarse-graining of molecular dynamics without forces. Journal of Chemical Theory and Computation, 19(3):942–952, 2023.
  47. Normalizing flows for probabilistic modeling and inference. J. Mach. Learn. Res., 22(57):1–64, 2021.
  48. Neural ordinary differential equations. In Advances in neural information processing systems, pages 6571–6583, 2018.
  49. Scalable reversible generative models with free-form continuous dynamics. In International Conference on Learning Representations, 2019.
  50. Equivariant Hamiltonian flows. arXiv preprint arXiv:1909.13739, 2019.
  51. Minibatch optimal transport distances; analysis and applications. arXiv preprint arXiv:2101.01792, 2021.
  52. Harold W Kuhn. The hungarian method for the assignment problem. Naval research logistics quarterly, 2(1-2):83–97, 1955.
  53. Wolfgang Kabsch. A solution for the best rotation to relate two sets of vectors. Acta Crystallographica Section A: Crystal Physics, Diffraction, Theoretical and General Crystallography, 32(5):922–923, 1976.
  54. E (n) equivariant graph neural networks. In International conference on machine learning, pages 9323–9332. PMLR, 2021.
  55. A family of embedded runge-kutta formulae. Journal of computational and applied mathematics, 6(1):19–26, 1980.
  56. Gfn2-xtb—an accurate and broadly parametrized self-consistent tight-binding quantum chemical method with multipole electrostatics and density-dependent dispersion contributions. Journal of Chemical Theory and Computation, 15(3):1652–1671, 2019. PMID: 30741547.
  57. Jerome M Kurtzberg. On approximation methods for the assignment problem. Journal of the ACM (JACM), 9(4):419–439, 1962.
  58. Learning from protein structure with geometric vector perceptrons. In International Conference on Learning Representations, 2021.
  59. Equivariant message passing for the prediction of tensorial properties and molecular spectra. In International Conference on Machine Learning, pages 9377–9388. PMLR, 2021.
  60. Equiformer: Equivariant graph attention transformer for 3d atomistic graphs. In The Eleventh International Conference on Learning Representations, 2023.
  61. Gemnet: Universal directional graph neural networks for molecules. Advances in Neural Information Processing Systems, 34:6790–6802, 2021.
  62. A new perspective on building efficient and expressive 3d equivariant graph neural networks. arXiv preprint arXiv:2304.04757, 2023.
  63. Conditional permutation invariant flows. arXiv preprint arXiv:2206.09021, 2022.
  64. Equivariant normalizing flows for point processes and sets, 2021.
  65. Exchangeable neural ode for set modeling. Advances in Neural Information Processing Systems, 33:6936–6946, 2020.
  66. Equivariant manifold flows. Advances in Neural Information Processing Systems, 34:10600–10612, 2021.
  67. Alleviating label switching with optimal transport. Advances in Neural Information Processing Systems, 32, 2019.
  68. Yann Brenier. Polar factorization and monotone rearrangement of vector-valued functions. Communications on pure and applied mathematics, 44(4):375–417, 1991.
  69. Pytorch: An imperative style, high-performance deep learning library. In Advances in Neural Information Processing Systems 32, pages 8024–8035. Curran Associates, Inc., 2019.
  70. Torchdyn: Implicit models and neural numerical methods in pytorch. In Neural Information Processing Systems, Workshop on Physical Reasoning and Inductive Biases for the Real World, volume 2, 2021.
  71. Pot: Python optimal transport. Journal of Machine Learning Research, 22(78):1–8, 2021.
  72. Openmm 7: Rapid development of high performance algorithms for molecular dynamics. PLoS computational biology, 13(7):e1005659, 2017.
  73. The atomic simulation environment—a python library for working with atoms. Journal of Physics: Condensed Matter, 29(27):273002, 2017.
  74. H Wiegand. Kish, l.: Survey sampling. john wiley & sons, inc., new york, london 1965, ix+ 643 s., 31 abb., 56 tab., preis 83 s., 1968.
Citations (42)

Summary

We haven't generated a summary for this paper yet.