Embedded-model flows: Combining the inductive biases of model-free deep learning and explicit probabilistic modeling

1 INTRODUCTION . Normalizing flows have emerged in recent years as a strong framework for high-dimensional density estimation , a core task in modern machine learning ( Papamakarios et al. , 2021 ; Kobyzev et al. , 2020 ; Kingma & Dhariwal , 2018 ) . As with other deep generative architectures such as GANs and VAEs , normalizing flows are commonly designed with the flexibility to model a very general class of distributions . This works well for large naturalistic datasets , such as those common in computer vision . However , general-purpose flows are less appealing when datasets are small or when we possess strong a priori knowledge regarding the relationships between variables . On the other hand , differentiable ( or ‘ deep ’ ) probabilistic programming is a powerful and general framework for explicit stochastic modeling that can express strong assumptions about the data- generating process ( Tran et al. , 2017 ; 2016 ; Bingham et al. , 2019 ; Dillon et al. , 2017 ; Piponi et al. , 2020 ; Kucukelbir et al. , 2017 ; Ambrogioni et al. , 2021a ) . For example , the user may know that the data represents a process evolving over time or in space , or is subject to known physical laws , or that there is an unobserved common cause or other shared factor behind a set of observations . Such structured prior information can enable the user to form effective inferences using much less data than would otherwise be required . For example , in the field of astrophysics differentiable probabilistic programs have been successfully used to study strong gravitational lensing ( Chianese et al. , 2020 ) and dark matter substructures ( Coogan et al. , 2020 ; Varma et al. , 2020 ) . All of these applications entail the translation of sophisticated physical theories into highly parameterized and structured probabilistic programs . However , it remains the case that “ all models are wrong ” ( Box , 1976 ) , and it is rare that a modeler can be fully confident in the strong assumptions ( e.g. , Gaussianity , statistical independence ) expressed by typical parametric probability models . In this paper we bridge these two worlds , providing an automatic technique to convert differentiable probabilistic programs into equivalent normalizing flow layers with domain-specific inductive biases . We call these architectures embedded-model flows ( EMF ) . EMFs can be use to integrate the structure of “ wrong but useful ” models within generic deep learning architectures capable of correcting the deviation of the model from the data . Furthermore , in the context of Bayesian inference , we show that the EMF architecture can be used to define a powerful class of approximate posterior distributions that embed the structure of the prior . 2 PRELIMINARIES . Normalizing flows define a random variable as an invertible differentiable transformation x = fϕ ( ϵ ) of a base variable ϵ ∼ p0 ( ϵ ) . In these expressions , ϕ are parameters that control the form of the transformation and consequently the distribution of x . The transformed ( log- ) density can be given explicitly using the change of variables formula from probability theory , log p ( x ) = log p0 ( f −1 ϕ ( x ) ) − log |det Jϕ ( x ) | , ( 1 ) where Jϕ ( x ) is the Jacobian matrix of fϕ . The base distribution p0 is often taken to be standard normal ( ϵ ∼ N ( 0 , I ) ) . In typical applications , the functional form of fϕ is specified by a deep architecture engineered to guarantee stable and tractable invertibility and tractable ( log- ) Jacobian computations ( Dinh et al. , 2014 ; Rezende & Mohamed , 2015 ) . The parameters ϕ are usually trained by stochastic gradient descent to maximize the log-likelihood ( given observed data ) or evidence lower bound ( in the variational inference setting ) . 2.1 PROBABILISTIC PROGRAMS AND GRAPHICAL MODELS . Probabilistic programming allows a joint probability distribution over a set of random variables to be specified intuitively as a program that includes random sampling steps ( van de Meent et al. , 2018 ) . Computations on this distribution , including the joint log-density function , structured inference algorithms , and the bijective transformations described in this paper , may then be derived automatically as program transformations . Recently frameworks have emerged for ‘ deep ’ probabilistic programming , which exploit the automatic differentiation and hardware acceleration provided by deep learning frameworks ; examples include Pyro ( Bingham et al. , 2019 ) , PyMC ( Salvatier et al. , 2016 ) , Edward ( Tran et al. , 2017 ) , and TensorFlow Probability ( Dillon et al. , 2017 ; Piponi et al. , 2020 ) . These allow for straightforward implementation of gradient-based inference and parameter learning , and enabling methods such as ours that integrate probabilistic programs with standard ‘ deep learning ’ components such as flows . Probabilistic programs that contain only deterministic control flow may equivalently be viewed as directed graphical models , in which each random variable xi is a graph node connected by incoming edges to some subset of the previous variables x < i . This leads to the following general expression for the joint density of a set of vector-valued variables x1 : n = { xj } nj=1 : p ( x0 : n ; ϕ ) = p0 ( x0 ; θ0 ( ϕ ) ) n∏ j=1 pj ( xj |θj ( x < j ; ϕ ) ) ( 2 ) Here θj ( x < j ; ϕ ) denotes the ‘ link function ’ associated with each node that takes as input the model parameters ϕ and the values of all parent variables ( excepting θ0 , where there are no parents ) , and outputs the parameters of its respective density function . As a practical matter , flow architectures generally assume layers of fixed size , so the experiments in this paper focus on programs of fixed structure and we use graphical-model notation for simplicity . However , the methods we describe do not require access to an explicit graph structure and could be straightforwardly applied to programs with stochastic control flow . 3 CONVERTING PROBABILISTIC PROGRAMS INTO NORMALIZING FLOWS . In this section , we will show that a large class of differentiable probabilistic programs can be converted into equivalent normalizing flow layers . More precisely , we will construct an invertible function Fϕ that maps spherical normal random inputs to the joint distribution of the probabilistic program : x = Fϕ ( ϵ ) ∼ p ( x ; ϕ ) . ( ϵ ∼ N ( 0 , I ) ) ( 3 ) This is a special case of the so called Rosenblatt transform ( Rosenblatt , 1952 ) . We interpret this transform as a flow layer and we extend it by including trainable gate variables that can be use to ” disconnect ” a set of variables from its parents so to correct model miss-specification . In our applications , the input to the bijective transformation may be arbitrary real-valued vectors , for example , the output of a previous flow layer , which likely will not be normally distributed . In such cases the transformed values will not be samples from the original probabilistic program , but we expect the transformation to capture some of its properties , such as multimodality and the dependence structure of its random variables . 3.1 INDIVIDUAL RANDOM VARIABLES AS BIJECTIVE TRANSFORMATIONS . We first consider the case of a univariate random variable x ∼ pϕ ( x ) , which we can express as a continuous bijective transformation fϕ of a standard normal variable ϵ via a slight extension of inverse transform sampling . The probability integral transform theorem states that the cumulative distribution function ( CDF ) of pϕ , Cϕ ( x ) = ∫ x a pϕ ( y ) dy , with pϕ defined over a generic interval ( a , b ) and a < x < b , is uniformly distributed , u = Cϕ ( x ) ∼ U ( 0 , 1 ) ( Angus , 1994 ) . If the CDF is invertible , i.e. , strictly increasing within its domain , it follows that we can sample x = C−1ϕ ( u ) ∼ pϕ given only a uniform random variable u ∼ U ( 0 , 1 ) . We obtain u through a second application of the probability integral transform theorem to the standard normal CDF Φ0,1 . Our univariate ‘ flow ’ is therefore given by x = fϕ ( ϵ ) = C −1 ϕ ( Φ0,1 ( ϵ ) ) ( ϵ ∼ N ( 0 , 1 ) ) ( 4 ) Closed-form expressions for fϕ may be available in some cases ( see below ) . In the general univariate setting where the inverse distribution function C−1ϕ ( u ) is not available in closed form , it may be evaluated efficiently using numerical root-finding methods such as the method of bisection ( see Appendix A.6 ) , with derivatives obtained via the inverse function theorem and the implicit reparameterization trick ( Figurnov et al. , 2018 ) . However , this is not needed for maximum likelihood training of the parameters ϕ , since eq . ( 1 ) only involves f−1ϕ and its derivatives . Example : Gaussian variables . The inverse CDF of a Gaussian variable x ∼ N ( µ , σ2 ) can be expressed as a scale-shift transformation of the inverse CDF of a standard normal variable : Φ−1µ , σ ( u ) = µ+σΦ −1 0,1 ( u ) . We can obtain the forward flow transformation by composing this function with the CDF of a standard Gaussian : fµ , σ ( ϵ ) = Φ−1µ , σ ( Φ0,1 ( ϵ ) ) = µ + σϵ . This is the famous reparameterization formula used to train VAEs and other variational methods ( Kingma & Welling , 2013 ) . The inverse function is f−1µ , σ ( x ) = ( x− µ ) /σ while the log-Jacobian is equal to log σ . It is not generally straightforward to evaluate or invert the CDF of a multivariate distribution . However , the most popular multivariate distributions , such as the multivariate Gaussian , can be expressed as transformations of univariate standard normal variables , and so may be treated using the methods of this section . More generally , the next section presents a simple approach to construct structured bijections ( i.e. , flows ) from multivariate distributions expressed as probabilistic programs . 3.2 FROM DIFFERENTIABLE PROBABILISTIC PROGRAMS TO STRUCTURED LAYERS . A probabilistic program samples a sequence of variables in turn , where ( as in eq . ( 2 ) ) the distribution of each variable is conditioned on the previously-sampled variables . At the sampling site of the j-th random variable xj , we therefore have all the information necessary to represent its conditional distribution as an invertible transformation fj , ϕ ( ϵj ; x < j ) of a standard normal input ϵj . Feeding the result back into the model defines a sequence of local transformations , which together form a joint flow layer Fϕ ( ϵ1 : n ) ( algorithm 1 ) . This layer transforms a spherical Gaussian variable ϵ1 : n into a variable that follows the same distribution of the input program . As with autoregressive flows , the forward transformation of the j-th input requires us to have computed the preceding outputs x < j = F j ϕ ( ϵ < j ) . This may be expressed recursively as follows : [ Fϕ ( ϵ1 : n ) ] j = fj , ϕ ( ϵj ; F j ϕ ( ϵ < j ) ) , ( 5 ) This recursive expression terminates since the parents have an acyclic-directed-graph structure . For variables without parents , F jϕ reduces to a trivial map between empty sets . In the reverse direction , the variables x1 : n are given , so all dimensions may be evaluated in parallel : [ F−1ϕ ( x1 : n ) ] j = f−1j , ϕ ( xj ; x < j ) ( 6 ) This is a generalization of the inversion formula used in autoregressive flows ( Kingma et al. , 2016 ; Papamakarios et al. , 2017 ) . The Jacobian of the transformation is block triangular ( up to a permutation of the variables ) and , consequently , the log determinant is just the sum of the log determinants of the block diagonal entries : log ∣∣∣det J−1ϕ ( x1 : n ) ∣∣∣ = ∑ j log ∣∣∣det J−1ϕ , j ( xj ; x < j ) ∣∣∣ ( 7 ) where J−1ϕ , j ( xj ; x < j ) is the Jacobian of the local inverse transformation f −1 ϕ ( xj ; x < j ) . When used in a normalizing flow , we denote the transformation Fϕ ( ϵ1 : n ) as a structured layer . Algorithm 1 Gated forward transformation : ϕ : Model parameters ; { ϵk } nk=1 : Input variables sorted respecting parenting structure ; { λk } nk=1 : Gate parameters ; n Number of variables . procedure GATEDFORWARDSAMPLER ( { ϵk } nk=1 ) State log |det J | ← 0 for j in Range ( n ) do { pi } mj i=1 ← ϵj .parents idx ▷ { pi } mj i=1 : indices of parents of j-th variable . ▷ mj : number of parents of j-th variable . if { pi } mj i=1 is not empty then x̃j ← fj , ϕ ( ϵj ; { xpi } mj i=1 ) xj ← λjx̃j + ( 1− λj ) ϵj log |det J | ← log |det J |+ log ( det |Jϕ , j ( { xk } nk=1 ) + ( 1− λj ) I| ) else x̃j ← fj , ϕ ( ϵj ; ϕ0j ) ▷ ϕ0j : parameters of the j-th root distribution . xj ← λjx̃j + ( 1− λj ) ϵj log |det J | ← log |det J |+ log ( det |Jϕj0 , j ( xj ) + ( 1− λj ) I| ) return { xk } nk=1 , log |det J | Gated structured layers Most user-specified models offer a simplified and imperfect description of the data . In order to allow the flow to skip the problematic parts of the model , we consider gated layers . In a gated layer , we take each local bijective transformation gj , ϕ to be a convex combination of the original transformation fj , ϕ and the identity mapping : gj , ϕ ( ϵj ; x < j , λj ) = λjfj , ϕ ( ϵj ; x < j ) + ( 1− λj ) ϵj ( 8 ) This gating technique is conceptually similar to the method used in highway networks ( Srivastava et al. , 2015 ) , recurrent highway networks ( Zilly et al. , 2017 ) and , more recently , highway flows ( Ambrogioni et al. , 2021b ) . Here the gates λj ∈ ( 0 , 1 ) are capable of decoupling each node from its parents , bypassing part of the original structure of the program . As before , the local inverses g−1ϕ ( xj ; x < j , λj ) can in general be computed by numeric root search , but the Gaussian case admits a closed form in both directions : gµ , σ ( ϵ , λ ) = ( 1 + λ ( σ − 1 ) ) ϵ+ λµ ; g−1µ , σ ( x , λ ) = x− λµ 1 + λ ( σ − 1 ) ( 9 ) Analogues of eq . ( 5 ) , eq . ( 6 ) , and eq . ( 7 ) then define the gated joint transformation Gλ , ϕ , its inverse and log Jacobian determinant , respectively . Note that the gating affects only the diagonal of the Jacobian , so the overall block-triangular structure is unchanged .

Summary: The authors describe how to translate a probabilistic program into a normalising flow layer which maps samples from a unit Gaussian into samples from the probabilistic program. They propose a "gating" mechanism for this layer so that it can learn to interpolate between this transformation and the identity transform. They then show how this layer can be combined with generic normalizing flow layers, and in this way add an inductive bias to a normalising flow without sacrificing its expressivity. In the experiments, they provide various examples using probabilistic programs in this way to construct normalising flows with helpful inductive biases.