Nonlinear denoising score matching for enhanced learning of structured distributions

2024-05-24

arXiv.org

AI Summary

To address the limitations of score-based generative models in modeling structured distributions—such as multimodal or approximately symmetric ones—and their poor generalization under small-sample regimes, this paper proposes the Nonlinear Denoising Score Matching (NDSM) framework. NDSM introduces learnable nonlinear drift into score matching for the first time, leveraging nonlinear stochastic differential equations (SDEs) to enhance structural representation capability. It further incorporates a neural control variate technique to substantially reduce gradient estimation variance. Crucially, NDSM enables data-driven structural embedding without requiring explicit symmetry priors. Experiments demonstrate that NDSM effectively mitigates mode collapse on both low-dimensional structured distributions and high-dimensional image data, improves small-sample generalization, and accurately captures approximate symmetries—outperforming equivariant networks and linear score matching baselines.

Technology Category

• Machine Learning: Structured Learning
• Computer Vision: Diffusion Models for Vision
• Search and Optimization: Non-convex Optimization

Application Category

• Graph Algorithms and Modeling for the Web: Graph embeddings and representation learning for Web-related graphs
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Economics, Online Markets and Human Computation: Economic ramifications for generative AI infrastructure and applications

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.

Problem

• Improves learning of structured distributions using nonlinear noising dynamics
• Addresses training challenges with nonlinear denoising score matching
• Enhances performance in high-dimensional data with reduced computational cost

Innovation

• Nonlinear noising dynamics enhance structured distribution learning
• Neural control variates reduce training objective variance
• Inexpensive preprocessing extracts data structure for dynamics