Year: 2018

Abstract:
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.

text: 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.

heading: 1 BACKGROUND
text: 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 cannot 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 cannot be used for discrete latents, as we cannot 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.

heading: 1.1 EVALUATING PROBABILITIES UNDER THE DISCRETE MODEL
text: 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 cannot give us a valid variational bound. To obtain a valid variational bound, we must use Q̄ (z̄), which we can compute but not differentiate.

heading: 1.2 EVALUATING PROBABILITIES UNDER THE CONTINUOUS MODEL
text: 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).

heading: 2 RESULTS
text: 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.

heading: 2.1 A FLOW-LIKE SAMPLING PROCEDURE FOR DISCRETE AUTOREGRESSIVE DISTRIBUTIONS
text: 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).

heading: 2.2 EVALUATING PROBABILITIES UNDER THE DISCRETE MODEL
text: First, we can compute relaxed Bernoulli variables using ẑ = σ ( βlK ) , and insert them into the probability mass function for the autoregressive discrete approximate posterior (Eq. 8). While this can be computed efficiently in parallel, it introduces another level of mismatch between the distribution we sample and the distribution under which we evaluate the probability, in the sense that we sample using the fixed point iterations, but evaluate probabilities under the sequential autoregressive process (Eq. 8). Remarkably, this is often not an issue for the discrete model, as we can simply iterate until convergence, and convergence is very well-defined as the latents are either 0 or 1. However, for the relaxed model it is more difficult to define convergence as the latents can lie anywhere between 0 and 1, so in practice we used a fixed number of iterations (in particular, K = 5, which was generally sufficient for the discrete model to converge).

heading: 2.3 EVALUATING PROBABILITIES UNDER THE CONTINUOUS MODEL
text: To evaluate the probability of the continuous variables under the fixed point iterations, we interpret the iterations as constituting a normalizing flow. Normalizing flows exploit the fact that we can compute the probability density of a random variable, lk(η), generated by transforming, via a oneto-one function, a sample η,
P ( lk(η) ) = P (η) ∣∣∣∣∣∂lk(η)∂η ∣∣∣∣∣ −1
(11)
where |∂lk/∂η| is the absolute value of the determinant of the Jacobian of lk(η). While the determinant of the Jacobian is often very difficult to compute, we can ensure the Jacobian is 1 by using a restricted family of transformations, under which the value of lkt depends only on the current value, ηt via simple addition, and on past values of η via an arbitrary function,
lk+1t = ηt + g k+1 ( x,η1:t−1 ) . (12)
And we know that our fixed-point iterations indeed lie in this family of functions, as Eq. (10) mirrors the above definition, with,
gk+1 ( x,η1:t−1 ) = vt ( x,σ ( βlk1:t−1 )) (13)
where lk1:t−1 indeed depends only on η1:t−1 (see Eq. 10).

heading: 2.4 ISSUES WITH CONTINUOUS MODEL
text: Using continuous latent variables might seem appealing, because we can compute the exact approximate posterior probability even when the fixed point iterations have not converged, and because there is additional flexibility in that we are not restricted to approximate posteriors which have an interpretation as a discrete dynamical system. However, moving from a discrete latent variable to a continuous latent variable introduces scope for considerable mismatch between the true and approximate posterior (Fig. 1B), which simply is not possible for the binary model where the posterior remains Bernoulli. This mismatch between the approximate and true posterior implies a discrepancy between the ELBO and the true model evidence (Fig. 1C), and this discrepancy grows as the evidence in favour of z̄ = 0 or z̄ = 1 increases. This mismatch has a potentially large magnitude (compare the scale on Fig. 1B to that in Fig. 2F), and thus can dramatically modify the variational inference objective function, introducing the possibility for biased inferences. In fact, this is exactly what we see (Fig. 1D), with the approximate posterior underestimating the evidence in favour of z̄ = 0 or z̄ = 1.
Further, the Logistic approximate posterior can introduce additional biases when optimized in combination with a relaxed binary variables (Figs. 1B–D use a hard-thresholding or equivalently, an infinite inverse-temperature). In particular, we consider a Gaussian likelihood, P (x = 0.5|ẑ) = N (x = 0.5; ẑ, σ), and thus, as σ decreases, there is increasing evidence that ẑ is close to 0.5. Under the true model, with a hard-threshold, this can never be achieved: z̄ must be either 0 or 1. However, if we combine a relaxation (Eq. 4) with a Logistic approximate posterior, it is possible to reduce the variance of the logistic sufficiently such that the relaxed Bernoulli variable, ẑ, is indeed close to 0.5. To quantify this effect, we consider how the standard deviation of ẑ (which should be 0.5 under the true posterior), varies as we change the standard deviation of the likelihood, σ (Fig. 1E).

heading: 2.5 EXPERIMENTS
text: One case where binary latent variables are essential is that of calcium spike deconvolution: inferring latent binary variables representing the presence or absence of a neural spike in a small time bin, based on noisy optical observations (Fig. 2C).
We take the binary variables, zit as representing whether neuron i spiked in time-bin t, and as such, we can interpret uit as the corresponding synaptic input (in contrast, as vt is not part of the generative model, it does not have a specific a biological interpretation). For the generative model, we take a weighted sum of inputs from past spikes, filtered by a temporal kernel κu and κv, which
is mirrored by the recognition model, to ensure that the recognition model can capture any priorinduced statistical structure,
ut = b u + Wu τ∑ t′=1 κut′ zt−t′ vt = bv(x) + Wv τ∑ t′=1 κvt′ zt−t′ . (14)
where is the Hadamard product, so r = g h implies ri = gihi. In the first experiment, we considered a single neuron, whose spiking was IID (Poisson) with firing rate 0.25 Hz (Fig. 2A). Fluorescence data was simulated by convolving spikes with a doubleexponential temporal kernel with rise time 0.3 s and decay time 1 s and adding noise with standard deviation e−1.5. We learned the recognition model, which consisted of two components. The first is a neural network mapping data to spike inferences, bv(x), which consisted of two hidden layers, with 20 units per cell per time bin, where the first layer takes input from 200 time points from a single cell’s fluorescence trace, and the second layer takes 5 time points from the previous hidden layer, and we use Elu nonlinearities (Clevert et al., 2015) (for further details see Speiser et al.,
2017; Aitchison et al., 2017). The second is the recognition temporal kernel κv, which captures the anticorrelations induced by explaining away (Fig. 2BC).
All strategies, including factorised and autoregressive (flow) VAE’s, and supervised training give roughly similar reconstructions (Fig. 2C top). Thus, to understand how the autoregressive posterior is superior, we need to look in more depth at the posteriors themselves. In particular, the factorised VAE had very narrow posteriors, spuriously indicating a very high degree of certainty in the spike timing (Fig. 2C middle), whereas both the supervised and autoregressive VAE and also the true posterior (estimate by importance sampling) indicate a higher level of temporal uncertainty. These differences are even more evident if we consider spike trains sampled from the approximate posterior (Fig. 2C bottom), or if we consider calibration: the probability of there actually being a spike in the underlying data, when the inference method under consideration indicates a particular probability of spiking (Fig. 2D). However, the sampled spike trains indicate another issue: while the true and VAE posteriors generally have one spike corresponding to each ground-truth spike, supervised training produces considerable uncertainty about the spike-count (Fig. 2E). Thus, the VAE with an autoregressive (flow) posterior combines the best of both worlds: achieving reasonable timing uncertainty (unlike the factorised VAE), whilst achieving reasonable spike counts (unlike supervised training).
As such, the autoregressive VAE performs more effectively than the factorised methods considering the IWAE objective (with 10 samples) (Burda et al., 2015), under which the models are trained
(Fig. 2F). As expected given Sec. 2.4, to get good performance, it is important to compute probabilities by putting the relaxed variables into the discrete pmf (binary, as opposed to logistic), and to use the parallel fixed point iteration based flow, as opposed to the autoregressive distribution (note however, that we did not thoroughly explore the space of normalizing flows). The considerable differences in speed between the approximate parallel fixed point iterations, and the exact autoregressive computation arises from an order-of-magnitude difference in the time required for an individual iteration (Fig. 2G).
In our second experiment, we considered inferring synaptic connectivity between cells, based on noisy observations. Here, we considered a network of 100 cells with no self-connectivity, and with weights, Wu that are sparse (probability of 0.1 of being non-zero), with the non-zero weights drawn from a Gaussian with standard deviation 5 (Fig. 3A), and with a 200 ms temporal decay. We used an autoregressive, fully connected recognition model (Fig. 3B). We use the same parameters as in the above simulation, except that we use a somewhat more realistic fluorescence rise-time of 100 ms (previously, we used 300 ms so as to highlight uncertainty in timing).
We inferred weights under three methods: a factorised VAE, and an autoregressive posterior where we use the fast-flow like sampling procedure (flow), and where we use the slow sequential sampling procedure (autoregressive) (Fig. 3C). The flow posterior gave a considerable advantage over the other methods in terms of the correlation between ground truth and inferred weights (Fig. 3D), and in terms of the slope (Fig. 3E), indicating a reduction in the bias towards underestimating the magnitude of weights, while the bias (i.e. the additive offset in Fig. 3C) remained small for all the methods. As such, the autoregressive posterior with flow-based sampling increased the ELBO considerably over the factorised model, or the autoregressive model with the slow sequential sampling, (Fig. 3G), and these differences again arise because of large differences in the time required for a single training iteration (Fig. 3H).

heading: 3 DISCUSSION
text: We have described an approach to sampling from a discrete autoregressive distribution using a parallel, flow-like procedure, derived by considering fixed-point iterations that converge to a sample from the underlying autoregressive process. We applied this procedure to speed up sampling from autoregressive approximate posteriors in the variational inference training loop. This allowed us to rapidly learn autoregressive posteriors in the context of neural data analysis, allowing us to realise the benefits of autoregressive approximate posteriors for single and multi cell data in reasonable timescales.
It is important to remember that while we can sample using K fixed-point iterations, we can only evaluate the probability of a sample once it has converged. This mismatch introduces a level of approximation in addition to those that are typical when relaxing discrete distributions (Jang et al., 2016; Maddison et al., 2016), but we can deal with the additional approximation error in the same way: by evaluating the model using samples drawn from the underlying discrete, autoregressive approximate posterior.
Past work has used similar properties of the underlying generative model, to speed up messagepassing based inference algorithms (Gonzalez et al., 2009; Domke, 2011). It is likely that their approach will be preferable when exact inference is possible albeit costly due to large tree-width/timecourses, whereas our approach will be preferable when exact inference is not possible due to longrange temporal dependencies.
Finally, our work suggests two directions for future work. First, while it is possible to use normalizing flows to define approximate posteriors for continuous state-space models, it may be difficult to know exactly which normalizing flow will prove most effective. In this context, our procedure of using fixed-point iterations may be a useful starting point. Second, we showed that while it may be possible to convert a discrete latent variable model to an equivalent model with continuous latents, this typically introduces considerable scope for mismatch between the prior and approximate posterior. However, the actual approximate posterior is relatively simple, a mixture of truncated Logistics, and as such, it may be possible to design approximate posteriors or even whole relaxation schemes that more closely match the true posterior, and indeed this may underlie the gains shown by (Vahdat et al., 2018).

Review summaries:
This paper uses an autoregressive filtering variational approximation for parameter estimation in discrete dynamical systems. One issue that crops up with this *particular choice* of variational distribution is that (a) inference proceeds sequentially (by definition) and (b) this does not make use of parallelism in modern hardware. To mitigate this, the paper proposes using fixed point iterations. After the first iteration, the approximate posterior for each latent variable corresponds to a random draw from a logistic distribution. Each subsequent application of the fixed point iteration modifies the posterior distribution by incorporating information from the previous latent state. After T applications of the fixed point equation, the procedure approximates an auto-regressive variational posterior distribution. Its possible I've misunderstood the point being made in Sections 2.2 - 2.4, but the paper points out that the choice of iterations "looks like" a normalizing flow (with Jacobian 1).

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.