Nonlinear denoising score matching for enhanced learning of structured distributions (2405.15625v1)
Abstract: We present a novel method for training score-based generative models which uses nonlinear noising dynamics to improve learning of structured distributions. Generalizing to a nonlinear drift allows for additional structure to be incorporated into the dynamics, thus making the training better adapted to the data, e.g., in the case of multimodality or (approximate) symmetries. Such structure can be obtained from the data by an inexpensive preprocessing step. The nonlinear dynamics introduces new challenges into training which we address in two ways: 1) we develop a new nonlinear denoising score matching (NDSM) method, 2) we introduce neural control variates in order to reduce the variance of the NDSM training objective. We demonstrate the effectiveness of this method on several examples: a) a collection of low-dimensional examples, motivated by clustering in latent space, b) high-dimensional images, addressing issues with mode collapse, small training sets, and approximate symmetries, the latter being a challenge for methods based on equivariant neural networks, which require exact symmetries.
- S. Asmussen and P. W. Glynn. Stochastic simulation: algorithms and analysis, volume 57. Springer, 2007.
- An optimal control perspective on diffusion-based generative modeling. In NeurIPS 2022 Workshop on Score-Based Methods, 2022.
- Quasi-stationary distributions: Markov chains, diffusions and dynamical systems, volume 1. Springer, 2013.
- Deciphering protein evolution and fitness landscapes with latent space models. Nature communications, 10(1):5644, 2019.
- Score-based generative modeling with critically-damped langevin diffusion. arXiv preprint arXiv:2112.07068, 2021.
- Generative adversarial nets. Advances in neural information processing systems, 27, 2014.
- FFJORD: Free-form continuous dynamics for scalable reversible generative models. arXiv preprint arXiv:1810.01367, 2018.
- Gans trained by a two time-scale update rule converge to a local nash equilibrium. In I. Guyon, U. V. Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 30. Curran Associates, Inc., 2017. URL https://proceedings.neurips.cc/paper_files/paper/2017/file/8a1d694707eb0fefe65871369074926d-Paper.pdf.
- 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, pages 8867–8887. PMLR, 2022.
- Gradient-based learning applied to document recognition. Proceedings of the IEEE, 86(11):2278–2324, 1998.
- Quasi-stationary distribution for the langevin process in cylindrical domains, part i: existence, uniqueness and long-time convergence. Stochastic Processes and their Applications, 144:173–201, 2022.
- Searching for protein variants with desired properties using deep generative models. BMC bioinformatics, 24(1):297, 2023.
- Structure preserving diffusion models. arXiv preprint arXiv:2402.19369, 2024.
- G. A. Pavliotis. Stochastic processes and applications. Texts in Applied Mathematics, 60, 2014.
- Searching for activation functions. arXiv preprint arXiv:1710.05941, 2017.
- High-resolution image synthesis with latent diffusion models. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pages 10684–10695, 2022.
- U-net: Convolutional networks for biomedical image segmentation. In Medical image computing and computer-assisted intervention–MICCAI 2015: 18th international conference, Munich, Germany, October 5-9, 2015, proceedings, part III 18, pages 234–241. Springer, 2015.
- R. Rubinstein. Simulation and the Monte Carlo Method. Wiley Series in Probability and Statistics. Wiley, 2009. ISBN 9780470317228. URL https://books.google.com/books?id=PUpdaQZsCK0C.
- Improved techniques for training gans. Advances in neural information processing systems, 29, 2016.
- Deep unsupervised learning using nonequilibrium thermodynamics. In International conference on machine learning, pages 2256–2265. PMLR, 2015.
- Score-based generative modeling through stochastic differential equations. arXiv preprint arXiv:2011.13456, 2020.
- Fourier features let networks learn high frequency functions in low dimensional domains. Advances in neural information processing systems, 33:7537–7547, 2020.
- Score-based generative modeling in latent space. Advances in Neural Information Processing Systems, 34:11287–11302, 2021.
- P. Vincent. A connection between score matching and denoising autoencoders. Neural computation, 23(7):1661–1674, 2011.
- Tackling the generative learning trilemma with denoising diffusion GANs. arXiv preprint arXiv:2112.07804, 2021.
- A mean-field games laboratory for generative modeling. arXiv preprint arXiv:2304.13534, 2023.
- Wasserstein proximal operators describe score-based generative models and resolve memorization. arXiv preprint arXiv:2402.06162, 2024.