Discrete flow posteriors for variational inference in discrete dynamical systems

Each training step for a variational autoencoder ( VAE ) requires us to sample from the approximate posterior , so we usually choose simple ( e.g . factorised ) approximate posteriors in which sampling is an efficient computation that fully exploits GPU parallelism . However , such simple approximate posteriors are often insufficient , as they eliminate statistical dependencies in the posterior . While it is possible to use normalizing flow approximate posteriors for continuous latents , there is nothing analogous for discrete latents . The most natural approach to model discrete dependencies is an autoregressive distribution , but sampling from such distributions is inherently sequential and thus slow . We develop a fast , parallel sampling procedure for autoregressive distributions based on fixed-point iterations which enables efficient and accurate variational inference in discrete state-space models . To optimize the variational bound , we considered two ways to evaluate probabilities : inserting the relaxed samples directly into the pmf for the discrete distribution , or converting to continuous logistic latent variables and interpreting the K-step fixed-point iterations as a normalizing flow . We found that converting to continuous latent variables gave considerable additional scope for mismatch between the true and approximate posteriors , which resulted in biased inferences , we thus used the former approach . We tested our approach on the neuroscience problem of inferring discrete spiking activity from noisy calcium-imaging data , and found that it gave accurate connectivity estimates in an order of magnitude less time . The development of variational auto-encoders ( VAE ) ( Kingma & Welling , 2014 ; Rezende et al. , 2014 ) has enabled Bayesian methods to be applied to very large-scale data , by leveraging neural networks , trained by stochastic gradient ascent , to approximate the posterior . Importantly , to perform each stochastic gradient ascent step , we need to sample from the current approximate posterior , restricting us to approximate posteriors in which sampling can be performed rapidly by leveraging parallel GPU computations . As such , most work on VAE ’ s has used factorised approximate posteriors ( e.g . Kingma & Welling , 2014 ; Rezende et al. , 2014 ; Blundell et al. , 2015 ; Higgins et al. , 2017 ; Aitchison et al. , 2017 ) , but in many domains of interest we expect the posterior over latents to be highly correlated , not only because the posterior inherits correlations from the prior ( e.g . as we might find in a dynamical system ) , but also because the likelihood itself can induce correlations due to effects such as explaining away ( Pearl , 1988a ; b ) . One approach to introducing correlations into the approximate posterior is normalizing flows ( Rezende & Mohamed , 2015 ; Kingma et al. , 2016 ) which transforms variables generated from a simple , often factorised distribution into a complex correlated distribution , in such a way that the determinant of the Jacobian ( and thus the probability of the transformed variables ) can easily be computed . However , the normalizing flow approach can only be applied to continuous latents , and there are important problems which require discrete latent variables and correlated posteriors , making efficient and accurate stochastic variational inference challenging . In particular , we consider the neuroscience problem of inferring the correlated spiking activity of neural populations recorded by calcium imaging . Due to the indirect nature of calcium imaging , spike inference algorithms must be used to infer the underlying neural spiking activity ( Friedrich et al. , 2017 ; Speiser et al. , 2017 ) . Not only does the data naturally give strong explaining away induced anticorrelations ( as a spike in nearby timebins produces very similar data ) , but there are prior correlations induced by synaptic connectivity which induces similar correlations in the approximate posterior . To address these challenging tasks with discrete latent variables , and correlated approximate posteriors , we considered two approaches . First , we considered applying normalizing flows by transforming our discrete latents into continuous latents , which are thresholded to recover the original discrete variables ( Maddison et al. , 2016 ) . However , we found that working with continuous variables gave rise to far more scope for mismatch between the approximate and true posteriors than working with discrete variables , and that this mismatch resulted in biased inferences . Instead , we developed a fast-sampling procedure for discrete autoregressive posteriors . In particular , we considered an autoregressive approximate posterior and circumvent the requirement for slow sequential sampling by developing flow-like fixed-point iterations that are guaranteed to sample the true posterior after T iterations , but in practice converge much more rapidly ( in our simulations , ∼ 5 iterations ) , and efficiently exploit GPU parallelism . Applying the flow-like fixed point iterations to simulated neuroscience problems , we were able to sample from autoregressive approximate posteriors in almost the same time required for factorised posteriors , and at least an order of magnitude faster than sequentially sampling from the underlying autoregressive process , allowing us to realize the benefits of correlated posteriors in large-scale settings . 1 BACKGROUND . The evidence lower bound objective ( ELBO ) , takes the form of an expectation computed over the approximate posterior , Qφ ( z ) , L = EQφ ( z ) [ log Pθ ( x|z ) + log Pθ ( z ) − log Qφ ( z ) ] ( 1 ) Optimizing this objective with respect to the generative model parameters , θ , is straightforward , as we can push the derivative inside the expectation , and perform stochastic gradient ascent . However , we can not do the same with the recognition parameters , φ , as they control the distribution over which the expectation is taken . To solve this issue , the usual approach is the reparameterisation trick ( Kingma & Welling , 2014 ; Rezende et al. , 2014 ) which performs the expectation over IID random noise , and transforms this noise into samples from the approximate posterior , L = E [ log Pθ ( x , z ( ; φ ) ) − log Qφ ( z ( ; φ ) ) ] . ( 2 ) As the recognition parameters φ no longer appear in the distribution over which the expectation is taken , we can again optimize this expression using stochastic gradient descent . While the reparameterisation trick is extremely effective for continuous latent variables , it can not be used for discrete latents , as we can not back-propagate gradients through discrete z ( ; φ ) . To rectify this issue , one approach is to relax the discrete variables , z̄ , to form an approximately equivalent model with continuous random variables , ẑ , through which gradients can be propagated ( Maddison et al. , 2016 ; Jang et al. , 2016 ) . To consider the simplest possible case , we take a single binary variable , z̄ , drawn from a Bernoulli distribution with log-odds ratio u , and probability p = σ ( u ) = 1/ ( 1 + e−u ) , P̄ ( z̄ ) = Bernoulli ( σ ( u ) ) . ( 3 ) Instead of sampling z̄ directly from a Bernoulli , we can obtain samples of z̄ from the same distribution by first sampling from a Logistic , and thresholding that sample using , l = Logistic ( u , 1 ) , z̄ = Θ ( l ) = lim β→∞ σ ( βl ) , ẑ = σ ( βl ) , ( 4 ) where Θ ( l ) is the Heaviside step function , which is 0 for negative inputs , and 1 for positive inputs , and ẑ is a relaxed version of the original discrete latent variable that lies between 0 and 1 and becomes equal to z̄ in the limit that the inverse-temperature , β , goes to infinity . We now have two options for how we compute probaiblities in the VAE setting for relaxed discrete latent variables . 1.1 EVALUATING PROBABILITIES UNDER THE DISCRETE MODEL . The most straightforward approach is to simply insert the relaxed variables , ẑ into the original probability mass function for the discrete model , P̄ ( ẑ ) , and Q̄ ( ẑ ) ( Jang et al. , 2016 ) . Taking a univariate example , this gives , log P̄ ( ẑ ) = ẑ log p+ ( 1− ẑ ) log ( 1− p ) . ( 5 ) However , this immediately highlights a key issue : to obtain a valid variational bound , we need to evaluate the probability density of samples from the approximate posterior . In our case , samples from the approximate posterior are the relaxed variables , ẑ , so we need to evaluate Q̂ ( ẑ ) . However , we are actually using a different expression , Q̄ ( ẑ ) , and while this may in practice be an effective approximation , it can not give us a valid variational bound . To obtain a valid variational bound , we must use Q̄ ( z̄ ) , which we can compute but not differentiate . 1.2 EVALUATING PROBABILITIES UNDER THE CONTINUOUS MODEL . To obtain a valid variational bound , we need to evaluate the actual probability of the relaxed variable , i.e . we need to compute Q̂ ( ẑ ) , and P̂ ( ẑ ) . While we could work with the ẑ directly , this is known to be numerically unstable so instead we work in terms of the logistic variables , l ( Maddison et al. , 2016 ) . These two approaches are equivalent , in the sense that the ratio of prior and approximate posterior probabilities is the same , because the gradient terms introduced by the change of variables cancel . The generative model is described above ( Eq . 4 ) , but we have not yet specified the approximate posterior , Q̂ ( l ) . Following usual practice Maddison et al . ( 2016 ) , we use a Logistic for the approximate posterior , with learned mean and scale parameter 1 ( so as to avoid a problem described below in Fig . 1E ) . Further , note that we have not specified whether these distributions correspond to the discrete or relaxed model , as we can obtain either ẑ or z̄ from exactly the same l , by using different inverse-temperatures for the transformation ( Eq . 4 ) . 2 RESULTS . We are interested in discrete dynamical systems with autoregressive generative models and approximate posteriors given by , P̄ ( z̄t|z̄1 : t−1 ) = Bernoulli ( σ ( u ( z̄1 : t−1 ) ) ) , ( 6a ) Q̄ ( z̄t|x , z̄1 : t−1 ) = Bernoulli ( σ ( v ( x , z̄1 : t−1 ) ) ) . ( 6b ) These equations can be rewritten using continuous Logistic variables , P̂ ( lt|l1 : t−1 ) = Logistic ( u ( σ ( βl1 : t−1 ) ) , 1 ) ( 7a ) Q̂ ( lt|l1 : t−1 ) = Logistic ( v ( x , σ ( βl1 : t−1 ) ) , 1 ) ( 7b ) where u and v are map from past activations to the logits value for the next time-step . Note that it is not possible to use classical message passing techniques to perform inference in this model due to the relatively long temporal dependencies . 2.1 A FLOW-LIKE SAMPLING PROCEDURE FOR DISCRETE AUTOREGRESSIVE DISTRIBUTIONS . One approach to sampling the sequential autoregressive posterior in Eq . ( 7b ) is to use , lt = ηt + v ( x , σ ( βl1 : t−1 ) ) = ft ( l1 : t−1 ) , ( 8 ) where ηt is Logistic ( 0 , 1 ) noise , lt and ηt are length N vectors , and we define ft so as to highlight the temporal dependencies ( cf . Eq.10 ) . However , in practice these iterations can be extremely slow as the sequential structure ( Fig . 1A top ) fails to fully exploit the parallelism available on today ’ s GPU hardware . As such , we considered fixed-point iterations that do fully utilize GPU parallelism . To construct these iterations , we simply apply the autoregressive updates at all time-steps in parallel , ( rather than sequentially , see Fig . 1A ) , lk+1 ( η , lk ) = η + v ( x , σ ( βlk ) ) . ( 9 ) where all the variables are N × T matrices , and where v ( x , σ ( βlk ) ) is simply Eq . ( 8 ) computed in parallel across time for an externally specified input . Making the time dependencies explicit , lk+1t = ηt + vt ( x , σ ( βlk1 : t−1 ) ) = ft ( l k 1 : t−1 ) . ( 10 ) To show that this will become equivalent to the autoregressive model after K = T iterations , note that after the first fixed-point iteration , l11 is exactly equal to the sequential autoregressive result , l1 ( Eq . 8 ) , as l11 = f1 ( ∅ ) = l1 , because there is no history to take into account ( which we represent , in a slight abuse of notation , by the empty set ) . After the second fixed-point iteration , l21 and l 2 2 are equal to the corresponding sequential autoregressive values , as l21 is still equal to the correct , autoregressive value , ( i.e . l21 = f1 ( ∅ ) = l1 ) , and l 2 2 = f2 ( l 1 1 ) = f2 ( l1 ) = l2 . By induction , after t iterations , the first t time-points lt1 : t are now equal to their sequential autoregressive values , l t 1 : t = l1 : t until after T parallel iterations we have equality for the whole time-series . Thus , the worst-case is extremely slow , converting a O ( T ) computation to an order O ( T 2 ) computation . However , in practice the iterations reach steady-state rapidly so we are able to use K T , giving an order of magnitude improvement in efficiency by making improved use of GPU parallelism . Finally , note that this procedure generalizes straightforwardly to the case of Categorical variables : we convert the sigmoid to a softmax , and convert the Logistic to a Gumbel distribution . While this section gives a fast procedure for sampling from the approximate posterior , it is still unclear how we should evaluate probabilities . We have two options , which mirror the two options introduced above ( Sec . 1.1 and Sec . 1.2 respectively ) .

Summary: Normalizing flows provide a way for variational posteriors to go beyond the mean-field assumption and introduce correlations in the posterior for latent variables. Typically, normalizing flows are only defined for continuous distributions, and the authors tackle the issue of creating flexible variational posteriors for discrete latent models. They posit a general autoregressive posterior family for discrete variables or their continuous relaxations. In order to perform variational inference with reparameterization gradients, one needs to have a sample from their variational family and be able to evaluate the density at the sample. The obvious way to sample from this autoregressive variational family is O(T) for T the number of autoregressive time-steps. Instead, they propose a method based off fixed point iterations to compute logits in parallel based off the results of previous iterations to generate an approximate sample. Moreover, they can interpret each iteration as a volume-preserving flow, so they don't have to add any terms to the density. They also use a continuous relaxation of Bernoulli random variables so they can back-propagate through a recognition network. They use their method on synthetic calcium spike data, and show that the correlated posterior is better-suited to handle uncertainty.