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: This paper proposes a new type of normalizing flow for incorporating inductive biases into the model architecture. To this end, the authors introduce a so-called “structure layer”, aiming to transform a spherical Gaussian variable into a pre-defined probabilistic program. A modified version, called “gated structured layer” is proposed in order to skip problematic parts of the model. The model is evaluated on a variety of different datasets and outperforms the masked autoregressive flow baseline.