More Like this

1. Soft-constrained Schrödinger Bridge: a Stochastic Control Approach

   Garg, Jhanvi; Zhang, Xianyang; Zhou, Quan (May 2024, Proceedings of The 27th International Conference on Artificial Intelligence and Statistics)

   Dasgupta, Sanjoy; Mandt, Stephan; Li, Yingzhen (Ed.)

   Schrödinger bridge can be viewed as a continuous-time stochastic control problem where the goal is to find an optimally controlled diffusion process whose terminal distribution coincides with a pre-specified target distribution. We propose to generalize this problem by allowing the terminal distribution to differ from the target but penalizing the Kullback-Leibler divergence between the two distributions. We call this new control problem soft-constrained Schrödinger bridge (SSB). The main contribution of this work is a theoretical derivation of the solution to SSB, which shows that the terminal distribution of the optimally controlled process is a geometric mixture of the target and some other distribution. This result is further extended to a time series setting. One application is the development of robust generative diffusion models. We propose a score matching-based algorithm for sampling from geometric mixtures and showcase its use via a numerical example for the MNIST data set. 

   more » « less
2. Boomerang: Local Sampling on Image Manifolds using Diffusion Models

   Luzi, L; Siahkoohi, A; Mayer, P M; Casco-Rodriguez, J; Baraniuk, R G (November 2023, Transactions on machine learning research)

   The inference stage of diffusion models involves running a reverse-time diffusion stochastic differential equation, transforming samples from a Gaussian latent distribution into samples from a target distribution on a low-dimensional manifold. The intermediate values can be interpreted as noisy images, with the amount of noise determined by the forward diffusion process noise schedule. Boomerang is an approach for local sampling of image manifolds, which involves adding noise to an input image, moving it closer to the latent space, and mapping it back to the image manifold through a partial reverse diffusion process. Boomerang can be used with any pretrained diffusion model without adjustments to the reverse diffusion process, and we present three applications: constructing privacy-preserving datasets with controllable anonymity, increasing generalization performance with Boomerang for data augmentation, and enhancing resolution with a perceptual image enhancement framework. 

   more » « less
3. Particle Deposition Driven by Evaporation in Membrane Pores and Droplets

   https://doi.org/10.1137/22S1532937  

   Jiang, Juliet; Zhang, Ruohan; Jeong, Dominic (January 2023, SIAM Undergraduate Research Online)

   This paper investigates particle deposition driven by fluid evaporation in a single pore channel representative of those found in porous membranes. A moving boundary problem for the 2D heat equation is coupled with an evolution equation for the pore radius, and describes the physical processes of fluid evaporation, diffusion of the particle concentration, and deposition on the pore channel wall. Furthermore, a stochastic differential equation (SDE) approach based on a Brownian motion particle-level description of diffusion is used as a similar phenomenological representation to the partial differential equation (PDE) model. Sensitivity analysis reveals trends in dominant model parameters such as evaporation rate, deposition rate, the volume scaling coefficient, and investigates the monotonicity of concentration. Evaluations of the asymptotically reduced model and the SDE model against the 2D PDE model are done in terms of the pore radius and solute concentration over time. For further exploration, we apply the model to a 2D droplet as well with both deterministic and stochastic approaches. 

   more » « less
4. A deep learning method for solving Fokker-Planck equations

   Jiayu Zhai, Matthew Dobson (January 2022, Proceedings of Machine Learning Research)

   Joan Bruna, Jan S (Ed.)

   The time evolution of the probability distribution of a stochastic differential equation follows the Fokker-Planck equation, which usually has an unbounded, high-dimensional domain. Inspired by Li (2019), we propose a mesh-free Fokker-Planck solver, in which the solution to the Fokker-Planck equation is now represented by a neural network. The presence of the differential operator in the loss function improves the accuracy of the neural network representation and reduces the demand of data in the training process. Several high dimensional numerical examples are demonstrated. 

   more » « less
5. Fit Like You Sample: Sample-Efficient Generalized Score Matching from Fast Mixing Diffusions

   Qin, Y; Risteski, A (June 2024, Conference on Learning Theory (COLT), 2024)

   Score matching is an approach to learning probability distributions parametrized up to a constant of proportionality (e.g. Energy-Based Models). The idea is to fit the score of the distribution, rather than the likelihood, thus avoiding the need to evaluate the constant of proportionality. While there's a clear algorithmic benefit, the statistical "cost'' can be steep: recent work by Koehler et al. 2022 showed that for distributions that have poor isoperimetric properties (a large Poincaré or log-Sobolev constant), score matching is substantially statistically less efficient than maximum likelihood. However, many natural realistic distributions, e.g. multimodal distributions as simple as a mixture of two Gaussians in one dimension -- have a poor Poincaré constant. In this paper, we show a close connection between the mixing time of a broad class of Markov processes with generator  and an appropriately chosen generalized score matching loss that tries to fit pp. This allows us to adapt techniques to speed up Markov chains to construct better score-matching losses. In particular, ``preconditioning'' the diffusion can be translated to an appropriate ``preconditioning'' of the score loss. Lifting the chain by adding a temperature like in simulated tempering can be shown to result in a Gaussian-convolution annealed score matching loss, similar to Song and Ermon, 2019. Moreover, we show that if the distribution being learned is a finite mixture of Gaussians in d dimensions with a shared covariance, the sample complexity of annealed score matching is polynomial in the ambient dimension, the diameter of the means, and the smallest and largest eigenvalues of the covariance -- obviating the Poincaré constant-based lower bounds of the basic score matching loss shown in Koehler et al. 2022. 

   more » « less