MSc Artificial Intelligence Master Thesis Continuous normalizing flows on manifolds by Luca Falorsi 11426322 2020 36 ECTS Supervisor: Assessor: Patrick Forre´, PhD Max Welling arXiv:2104.14959v1 [stat.ML] 14 Mar 2021 Informatics Institute Acknowledgements I would first like to thank my thesis supervisor Patrick Forré for helping me not to lose my way during these two years. I would also like to thank Nicola de Cao, Tim R. Davidson, Fabrizio Ambrogi, Pim de Haan, Jonas Köhler, Marco Federici and Luca Simonetto, for the wonderful years spent together in Amsterdam. Finally, I thank my parents and my brother for their love and support. i Abstract Normalizing flows are a powerful technique for obtaining repa- rameterizable samples from complex multimodal distributions. Unfortunately, current approaches are only available for the most basic geometries and fall short when the underlying space has a nontrivial topology, limiting their applicability for most real- world data. Using fundamental ideas from differential geometry and geometric control theory, we describe how the recently in- troduced Neural ODEs and continuous normalizing flows can be extended to arbitrary smooth manifolds. We propose a general methodology for parameterizing vector fields on these spaces and demonstrate how gradient-based learning can be performed. Ad- ditionally, we provide a scalable unbiased estimator for the di- vergence in this generalized setting. Experiments on a diverse selection of spaces empirically showcase the defined framework’s ability to obtain reparameterizable samples from complex distri- butions. ii Contents Abstracti Contents iii 1 Introduction1 1.1 Summary . .2 2 Preliminaries4 2.1 Reparameterization trick . .4 2.2 Normalizing flows . .6 2.3 Neural ODEs and continuous normalizing flows . .8 3 Densities and measure pushforwards 12 3.1 Reparameterization trick as measure pushforward . 16 3.1.1 Differential geometric prospective: pullback of inverse . 17 3.1.2 Computing the local volume change . 18 3.1.3 Computing the volume change: practical advice . 20 3.1.4 Summary . 21 4 CNF on Manifolds 23 4.1 Integral curves and flows . 23 4.1.1 Vector fields . 23 4.1.2 Flows . 26 iii CONTENTS iv 4.2 Measure pushforward induced by the flow . 29 4.2.1 Divergence of a vector field . 29 4.2.2 Continuity equation on Manifolds . 32 4.2.3 Continuity equation on characteristic curves . 33 5 Parameterizing vector fields 36 5.0.1 Local frames and global constraints . 37 5.0.2 Lie groups . 38 5.1 Generators of vector fields . 39 5.1.1 Time dependent vector fields . 42 5.2 Divergence using a generating set . 42 5.3 Homogeneous spaces . 44 5.4 Embedded submanifolds of Rm ................. 46 5.4.1 Embedded Riemannian Submanifolds . 48 5.4.2 Generators defined by gradients of laplacian eigenfunc- tions . 49 5.4.3 Isometrically embedded Submanifolds . 50 6 Backpropagation through flows 53 6.1 A short primer on symplectic geometry . 53 6.2 Cotangent lift . 56 6.2.1 Using local coordinates . 57 6.2.2 Using a generating set . 58 6.2.3 Using a local frame . 59 6.2.4 Parallelizable manifolds . 61 6.2.5 Cotangent lift on Lie Groups . 62 6.2.6 Cotangent lift on embedded submanifolds . 63 7 Experiments 66 7.1 General setup . 66 CONTENTS v 7.1.1 Density matching on manifolds . 66 7.1.2 Mixture of densities . 67 7.1.3 Implementation details . 67 7.2 Manifold specific setup and results . 68 7.2.1 Special Orthonormal matrices . 68 7.2.2 Stiefel Manifold . 69 7.2.3 Hypersphere . 69 7.2.4 Unitary matrices . 72 7.2.5 Special Unitary matrices . 73 7.2.6 Positive definite symmetric matrices . 75 7.3 Summary of results and comment . 77 8 Conclusion 79 A Densities and induced measure 81 A.1 volume forms . 81 A.2 Densities . 83 A.2.1 Induced measure . 86 Bibliography 89 Chapter 1 Introduction Normalizing flows (NF) are a family of methods for defining flexible reparam- eterizable probability distributions on high dimensional data. They accom- plish this by mapping a sample from a simple base distribution into a complex one, using a series of invertible mappings that are usually parameterized by invertible neural networks. The probability density of the final distribution is then given by the change of variable formula. As such, normalizing flows allow for fast and efficient sampling as well as density evaluation. Due to their desirable features, flow based model have been subject to an ever growing interest in the machine learning community since their intro- duction (Rezende and Mohamed, 2015), and have been successfully used in the context of generative modelling, variational inference and density esti- mation1, with a growing number of applications in different fields of science. (Noé et al., 2019; Kanwar et al., 2020; Wirnsberger et al., 2020). As many real word problems in robotics, physics, chemistry, and the earth sciences are naturally defined on spaces with a non-trivial topology, recent work has focused on building probabilistic deep learning frameworks that can work on manifolds different from the Euclidean space (Davidson et al., 2018, 2019; Falorsi et al., 2018, 2019; Pérez Rey et al., 2019; Nagano et al., 2019). For these type of data the possibility of defining complex reparameterizable densities on manifolds through normalizing flows is of central importance. However, as of today there exist few alternatives, mostly limited to the most basic and simple topologies. The main obstacle for defining normalizing flows on manifolds is that cur- 1See Papamakarios et al.(2019); Kobyzev et al.(2020) for a comprehensive review of NF. 1 CHAPTER 1. INTRODUCTION 2 rently there is no general methodology for parameterizing maps F : M ! N between two manifolds. Neural networks can only accomplish this for the Euclidean space, Rn. In this work we propose to use vector fields on a mani- fold M as a flexible way to express diffeomorphic maps from the manifold to itself. As vector fields define an infinitesimal displacement on the manifold for every point, they naturally give rise to diffeomorphisms without needing to impose further constraints. Furthermore, vector fields are significantly easier to parameterize using neural architectures, as they form a free module over the ring of functions on the manifold. In doing so we make it possible to define NF on manifolds. Furthermore, there exists decades old research on how to numerically integrate ODEs on manifolds.2 Recently, normalizing flows build using differential equations have proven successful in Euclidean space (Chen et al., 2018; Grathwohl et al., 2019) taking advantage of unrestricted neural network architectures. Using ideas from differential geometry and building on the concepts first introduced in Chen et al.(2018), this work continues this line of research by defining a flexible framework for constructing normalizing flows on manifolds that is trivially extendable to any manifold of interest. 1.1 Summary Chapter 2: We start by reviewing the fundamental Machine Learning con- cepts on which we build upon in the rest of the thesis: reparameterizable distributions, normalizing flows and neural ODEs. We additionally review the main works that tried to extend these frameworks to manifold setting. Chapter 3 and Appendix A: Here we take care of carefully describ- ing the mathematical constructs needed for building a coherent theory of reparameterizable distributions on smooth manifolds. We first define volume forms and densities on manifolds, which are the fundamental objects used for measuring volumes and integrating. We then outline how, using the Riesz extension theorem, a smooth nonnegative density uniquely defines a Radon measure on the space, providing a general methodology for instantiating a base measure on a manifold. After this, we illustrate how the concept of repa- rameterizable distribution can be generalized to abstract measure spaces via measure pushforward. We conclude by describing how a measure pushfor- ward given by a diffeomorphism operates between densities, giving a change 2See Hairer et al.(2006) for a review of the main methods. CHAPTER 1. INTRODUCTION 3 of variables formula to use when building normalizing flows on manifolds. Chapter 4: In this chapter, we first delineate how vector fields and ODEs on a manifold M can be defined in the context of differential geometry, and explain how they give rise to diffeomorphisms on M through their associated flow. We then apply the notions developed in Chapter 3 to the specific flow case, demonstrating how flows defined by vector fields allow defining continuous normalizing flows on manifolds. As in the Euclidean setting, the change of density is given by integrating the divergence on the ODE solutions, where now the divergence is a generalized quantity that depends on the geometry of the space. Chapter 5: We describe a general methodology for parameterizing vector fields on smooth manifolds. Using module theory, we show that all vector fields on a manifold can be obtained by linear combinations of a finite set of generating vector fields, with coefficients given by functions on the manifold. We then outline how generating sets of vector fields can be built on embedded submanifolds of Rm and homogeneous spaces. We conclude by describing how the divergence can be computed when employing a generating set, and give an unbiased Monte Carlo estimator in this context. Chapter 6: Here we show how the adjoint sensitivity method can be gener- alized to vector fields on manifolds in the context of geometric control theory (Agrachev and Sachkov, 2013), highlighting important connections with sym- plectic geometry and the Hamiltonian formalism. Similarly, as in the adjoint method in the Euclidean space, to backpropagate through the flow defined by a vector field we have to solve an ODE in an augmented space. In this case, the ODE is given by a vector field on the cotangent bundle T ∗M, called cotangent lift, which is a lift of the original vector field on M.