Optimal transport based simulation methods for deep probabilistic models

Thesis

Deep probabilistic models have emerged as state-of-the-art for high-dimensional, multi-modal data synthesis and density estimation tasks. By combining abstract probabilistic formulations with the expressivity and scalability of neural networks, deep probabilistic models have become a fundamental component of the machine learning toolbox. Such models still have a number of limitations however. For example, deep probabilistic models are often limited to gradient based training and hence struggle to incorporate non-differentiable operations; they are expensive to train and sample from; and often deep probabilistic models do not leverage prior geometric and problem-specific structural knowledge.

This thesis consists of four contributing pieces of work and advances the field of deep probabilistic models through optimal transport based simulation methods. First, by using regularized optimal transport via the Sinkhorn algorithm, we provide a theoretically grounded and differentiable approximation to resampling within particle filtering. This allows one to perform gradient based training of state space models, a class of sequential probabilistic model, with end-to-end differentiable particle filtering. Next, we explore initialization strategies for the Sinkhorn algorithm to address speed issues. We show that careful initializations result in dramatic acceleration of the Sinkhorn algorithm. This has applications in differentiable sorting; clustering within the latent space of a variational autoencoder; and within particle filtering. The remaining two works contribute to the field of diffusion based generative modelling through the Schrödinger Bridge. First, we connect diffusion models to the Schrödinger Bridge, coined the Diffusion Schrödinger Bridge. This methodology enables accelerated sampling; data-to-data simulation, and a novel way to compute regularized optimal transport for high dimensional, continuous state-space problems. Finally, we extend the Diffusion Schrödinger Bridge to the Riemannian manifold setting. This allows one to incorporate prior geometric knowledge and hence enable more efficient training and inference for diffusion models on Riemannian manifold valued data. This has applications in climate and Earth science.

Authors: Thornton, J

• DOI: 10.5287/ora-j0onmxmok
• Type of award: DPhil
• Level of award: Doctoral
• Awarding institution: University of Oxford
• Language: English
• Keywords: score-based generative models; Schrodinger Bridge; sinkhorn; diffusion models; optimal transport; generative modeling
• Subjects: Computational statistics; Optimal transport; Machine learning