PaperlayerView exact revision on arXiv

Source-linked AI summary

Improving Variational Inference with Inverse Autoregressive Flow

Diederik P. Kingma, Tim Salimans, Rafal Jozefowicz, Xi Chen, Ilya Sutskever, Max Welling

arXiv:1606.04934v2cs.LGstat.ML

TL;DR

Flexible posterior inference is difficult when high-dimensional latent variables require efficient density computation and sampling, often favoring limited diagonal Gaussian approximations. The paper introduces inverse autoregressive flow, which improves inference quality and achieves near-state-of-the-art image log-likelihood with faster sampling.

  • Problem

    Efficient high-dimensional variational inference restricts practical approximate posteriors, often leading to diagonal Gaussian models because density computation, differentiation, sampling, and parallelization must remain tractable.

  • Method

    Inverse autoregressive flow chains invertible transformations based on autoregressive neural networks to construct flexible approximate posteriors that scale to high-dimensional latent spaces.

  • Results

    IAF yields significant gains over factorized Gaussian posteriors and reaches 3.11 bits per dimension on CIFAR-10, close to state-of-the-art log-likelihood with much faster sampling.

  • Takeaways & Limitations

    IAF provides a high-dimensional inference approach that combines improved performance with substantially faster synthesis than comparable high-likelihood models.

  • Takeaways & Limitations

    Earlier bottleneck-based flow transformations may require long chains to capture high-dimensional dependencies.

Abstract

from arXiv · show

The framework of normalizing flows provides a general strategy for flexible variational inference of posteriors over latent variables. We propose a new type of normalizing flow, inverse autoregressive flow (IAF), that, in contrast to earlier published flows, scales well to high-dimensional latent spaces. The proposed flow consists of a chain of invertible transformations, where each transformation is based on an autoregressive neural network. In experiments, we show that IAF significantly improves upon diagonal Gaussian approximate posteriors. In addition, we demonstrate that a novel type of variational autoencoder, coupled with IAF, is competitive with neural autoregressive models in terms of attained log-likelihood on natural images, while allowing significantly faster synthesis.

1 Introduction

The paper introduces inverse autoregressive flow as a flexible inference transformation built from Gaussian autoregressive functions, designed for high-dimensional tensor-valued latent variables. It demonstrates the method by improving inference networks in deep variational auto-encoders with hierarchical three-dimensional tensor latent variables.

  • 1 Introduction: Stochastic variational inference becomes especially efficient for continuous latent variables through reparameterization and inference networks that amortize computational cost.This produces a highly scalable learning procedure for large datasets.
  • 1 Introduction: Inverse autoregressive flow uses Gaussian autoregressive functions that output elementwise means and standard deviations conditioned on preceding elements under a specified ordering.Examples of related autoregressive density estimators include RNNs, MADE, PixelCNN, and WaveNet.
  • 1 Introduction: The transformation is well suited to high-dimensional tensor variables, including variables organized spatio-temporally.This suitability distinguishes it from previously used normalizing flows discussed in the introduction.
  • 1 Introduction: The method improves inference networks for deep variational auto-encoders with latent variables at multiple hierarchy levels.The stochastic variables are three-dimensional tensors, described as stacks of featuremaps.
  • 1 Introduction: The resulting deep variational auto-encoders demonstrate improved performance when each hierarchical stochastic variable is a three-dimensional tensor.The experiments specifically use latent variables represented as stacks of featuremaps.

2 Variational Inference and Learning

Variational inference introduces a parametric inference model q(z|x) and optimizes a lower bound when marginal likelihood optimization is intractable. Effective inference models must support efficient density computation, differentiation, and sampling, motivating flexible normalizing flows while exposing limitations in prior transformations.

  • Variational inference: When marginal likelihood is intractable, variational inference introduces q(z|x) and optimizes a variational lower bound on each observation’s marginal log-likelihood.This difficulty arises in flexible generative models, including those whose components are parameterized by neural networks.
  • Optimization: For continuous z, re-parameterization of q(z|x) enables efficient optimization of the variational lower bound.The section cites Kingma and Welling (2013) and Rezende et al. (2014) for this approach.
  • Variational inference: Maximizing L(x; θ) simultaneously increases log p(x) and decreases DKL(q(z|x)||p(z|x)), tightening the bound as the divergence approaches 0.The divergence may also be the direct objective when q(z|x) is used for inference after optimization.
  • Inference-model requirements: Inference models must make q(z|x) density computation, differentiation, and sampling computationally efficient for every datapoint in each minibatch iteration.For high-dimensional z, parallelizability is also needed to use GPU resources efficiently.
  • Normalizing flows: Normalizing flows build flexible posterior distributions by applying a chain of invertible parameterized transformations to an initial variable with a simple, tractable density.The final density remains computable when each transformation’s Jacobian determinant can be computed.
  • Normalizing flows: Earlier transformations using a single-unit bottleneck require long chains to capture high-dimensional dependencies because information passes through that bottleneck.This limitation motivates seeking more expressive invertible transformations than the restricted family studied by Rezende and Mohamed (2015).

3 Inverse Autoregressive Transformations

Inverse autoregressive transformations invert sequential autoregressive sampling into a parallelizable operation with a simple Jacobian determinant. Their flexibility, dimension-wise parallelizability, and tractable log-determinant make them suitable for normalizing flows in high-dimensional latent spaces.

  • Motivation: Autoregressive sampling is sequential, with each y_i depending on previously generated elements, so its computation scales with dimensionality D.The forward transformation is therefore unsuitable for direct variational-inference sampling despite computation proportional to D.
  • Inverse transformation: The inverse transformation is one-to-one when every σ_i > 0 and can be parallelized because each ϵ_i computation is independent of the others.This reverses the computational disadvantage of autoregressive sampling for normalizing flows.
  • Jacobian: Autoregressive structure yields a lower-triangular Jacobian whose diagonal satisfies ∂ϵ_i/∂y_i = σ_i, making the log-determinant simple to compute.The determinant is the product of the diagonal terms.
  • Design rationale: Model flexibility, parallelizability across dimensions, and a simple log-determinant motivate using the transformation as a normalizing flow over high-dimensional latent space.These properties jointly address the requirements of flexible variational inference in large latent spaces.

4 Inverse Autoregressive Flow (IAF)

IAF constructs flexible approximate posteriors through a chain of autoregressive, invertible transformations whose tractable Jacobians enable simple density computation. Its flexibility grows with autoregressive-model expressivity and chain depth, with implementation choices improving stability and results.

  • Flow construction: IAF chains T autoregressive transformations, conditioning each step on the previous latent state and an encoder-produced auxiliary input h.The initial encoder outputs μ0, σ0, and h; subsequent autoregressive networks output μt and σt.
  • Flow construction: Autoregressive structure makes each transformation Jacobian triangular, with σt on the diagonal, yielding a tractable determinant and final log-density.The determinant is the product of the diagonal scale terms, ∏i σt,i.
  • Expressivity: The final iterate becomes more flexible as the autoregressive models become more expressive and the transformation chain becomes deeper.This flexibility is intended to help the approximate posterior closely fit the true posterior.
  • Implementation: A numerically stable IAF update uses σt = sigmoid(st) and interpolates between zt−1 and an autoregressive vector mt.This LSTM-inspired formulation remains compatible with simple final log-density computation.
  • Implementation: Results improved when variables were reordered after each IAF step, while masked or convolutional autoregressive networks supported nonconvolutional and CIFAR-10 experiments.The reordering is volume-preserving, so the simple log-density expression remains unchanged.

5 Related work

IAF belongs to the normalizing-flow family and is distinguished from earlier approaches by its focus on scaling beyond relatively low-dimensional latent spaces. Related alternatives include NICE transformations, Hamiltonian flows, and auxiliary-variable methods, while the experiments combine deep latent-variable models with IAF.

  • Normalizing flows: IAF extends the normalizing-flow family introduced for stochastic variational inference, where planar and radial flows were effective in latent spaces of at most a few hundred dimensions.The supplied passage notes that scaling those earlier flows to much higher-dimensional latent spaces was unclear.
  • NICE: NICE updates only half the latent variables per step using a neural-network function of the remaining variables, making inversion cheap but typically requiring longer chains.The updated variables are z1:D/2, while the conditioning variables are zD/2+1:D.
  • Hamiltonian flows: Hamiltonian Variational Inference generates transformations by simulating a Hamiltonian system over latent variables and auxiliary momenta.The passage also states that this transformation is guided by the exact posterior and leaves it invariant for small step sizes.
  • Auxiliary-variable methods: Auxiliary latent variables and corresponding inference models provide another route to more flexible variational inference, with multilayer stochastic latent-variable models often equivalent to such methods.The experiments combine deep latent-variable models with IAF to benefit from both approaches.

6 Experiments

Experiments show that increasingly expressive IAF posteriors improve variational autoencoder performance on MNIST, while a ResNet VAE with IAF achieves strong CIFAR-10 likelihood and much faster synthesis than PixelCNN.

  • MNIST: Increasing posterior expressiveness improves generative modeling performance and tightens variational lower bounds on dynamically binarized MNIST.The tighter bounds reduce the gap between variational lower bounds and marginal likelihoods.
  • MNIST: A deep-IAF VAE achieves a log-likelihood of -79.88 on statically binarized MNIST, slightly worse than PixelCNN’s -79.2.These results compare the reported VAE result with the best reported PixelCNN result.
  • CIFAR-10: On CIFAR-10, the ResNet VAE with IAF achieves 3.11 bits per dimension, outperforming other published latent-variable models and nearly matching PixelCNN.The authors suggest that additional flow steps could further improve the result.
  • CIFAR-10: Sampling takes 0.05 seconds per image with the ResNet VAE versus 52.0 seconds per image with PixelCNN on an NVIDIA Titan X GPU.PixelCNN was sampled naïvely by generating pixels sequentially with the full generative model at each iteration.
  • CIFAR-10: ResNet VAE synthesis is a parallel computation that requires no custom code, whereas PixelCNN’s sequential sampling limits speedups on parallel hardware.Custom code can accelerate PixelCNN by evaluating only the relevant network portion, but parallel hardware limits the benefit.

7 Conclusion · A Linear IAF · B MNIST

IAF is presented as a scalable normalizing flow that improves variational posteriors and achieves near-state-of-the-art CIFAR-10 log-likelihood with faster sampling. Its linear special case connects diagonal Gaussian posteriors to full-covariance Gaussians through inverse Cholesky structure, while the MNIST VAE uses a 32-dimensional stochastic layer and ResNet-based inference network.

  • 7 Conclusion: IAF scales well to high-dimensional latent spaces and yields significant gains over factorized Gaussian approximate posteriors.The experiments also report close to state-of-the-art CIFAR-10 log-likelihood with much faster sampling.
  • 7 Conclusion: IAF achieves close to state-of-the-art log-likelihood results on CIFAR-10 while allowing much faster sampling.
  • A Linear IAF: The simplest linear IAF transforms a diagonal-covariance Gaussian into one with linear dependencies.
  • A Linear IAF: Any full-covariance Gaussian with mean m and covariance C can be expressed as an autoregressive model with conditional means defined by preceding variables.
  • A Linear IAF: The conditional variance σi(y1:i−1) is given by the covariance Schur-complement expression involving C[i, i] and preceding-variable blocks.
  • A Linear IAF: Inverting the autoregressive model yields ϵ = (y −µ(y))/σ(y) = L(y −m), where L is the inverse Cholesky factor of C.
  • A Linear IAF: The variational encoder produces µ(x), σ(x), and a lower triangular inverse Cholesky matrix L(x) after an initially factorized Gaussian posterior.
  • B MNIST: The MNIST convolutional VAE uses a single Gaussian stochastic layer with dimension 32 and a ResNet-based inference network with three stride-2 blocks.The blocks use 3x3 filters and [16,32,32] feature maps, with additional stride-1 ResNet blocks between alternating strided convolutions.

C ResNet VAE · C.1 Bottom-Up versus Bidirectional inference · C.2 Inference with stochastic ResNet

The ResNet VAE combines stacked residual blocks and stochastic latent layers with an autoregressive prior, improving posterior flexibility while retaining straightforward sampling. Its inference variants differ in latent-variable ordering and conditioning, with bottom-up inference sampling in reverse generative order and bidirectional inference incorporating both bottom-up and top-down information.

  • C ResNet VAE: The CIFAR-10 ResNet VAE uses L stacked blocks combining bottom-up residual inference units with top-down residual units for inference and generation.Each block produces bottom-up activations h(q) and top-down activations h(p).
  • C ResNet VAE: Each generative residual function combines deterministic hidden units with a small number of stochastic units distributed by a heteroscedastic diagonal Gaussian.The stochastic distribution is p(z_l|h^(p)_l), followed by a nonlinearity.
  • C ResNet VAE: The latent prior factorizes autoregressively as p(z_1:L) = p(z_L) product_{l=1}^{L-1} p(z_l|z_{l+1:L}), increasing posterior flexibility and yielding a tighter variational bound.The stated mechanism is that a flexible prior makes matching approximate and true posteriors easier without sacrificing generative-model flexibility.
  • C.1 Bottom-Up versus Bidirectional inference: Bottom-up inference reverses the generative model’s topological ordering: q(z_1,z_2,z_3|x) = q(z_1|x)q(z_2|z_1,x)q(z_3|z_2,z_1,x).For the example, the generative ordering is z_1 → z_2 → z_3 → x, while inference proceeds in the reverse ordering.
  • C.1 Bottom-Up versus Bidirectional inference: Bidirectional inference first performs a fully deterministic bottom-up pass, then samples posterior variables in the generative model’s top-down topological order.This contrasts with bottom-up inference, whose latent-variable ordering is reversed relative to generation.
  • C.2 Inference with stochastic ResNet: Both inference variants are implemented with ResNet blocks, and bottom-up inference samples latent variables in reverse generative order.Its residual functions compute q(z_l|h^(p)_{l+1}) from the bottom-up residual unit’s input.
  • C.2 Inference with stochastic ResNet: In bottom-up inference, each sampled latent is transformed by a nonlinearity and incorporated into the residual hidden layer for upstream computation.The sample is therefore used upstream in subsequent bottom-up residual processing.
  • C.2 Inference with stochastic ResNet: In bidirectional inference, each approximate posterior is conditioned on both bottom-up and top-down inputs, and the transformed sample is used downstream.The posterior is written as q(z_l|h^(q)_l), with the sample incorporated into the top-down residual hidden layer after the nonlinearity.

C.3 Approximate posterior … C.7 Nonlinearity

The paper specifies alternative approximate posteriors and implementation choices for the ResNet VAE, including IAF structure, bottom-layer operations, likelihood modeling, normalization, and nonlinearities. Reported experiments favor ELU and weight normalization, while additional spatial resampling and learned IAF rescaling provide limited benefit.

  • C.3 Approximate posterior: Approximate posteriors use either diagonal Gaussians or IAF transformations, with IAF contexts supplied by inference-direction-dependent encoder states.The supplied passage introduces the two posterior families and indicates that the context depends on inference direction.
  • C.3 Approximate posterior: IAF uses a single masked PixelCNN step with zero, one, or two hidden layers; zero hidden layers reduce it to a linear transformation with off-diagonal covariance.The hidden layers use ELU nonlinearities.
  • C.3 Approximate posterior: The performance difference between learned dynamic σt(.) rescaling and fixed σt(.) = 1 was almost negligible on CIFAR-10 test-set bpp.The comparison concerns location-only versus full location-and-scale IAF perturbations.
  • C.4 Bottom layer: The encoder begins with 2 × 2 spatial subsampling and the decoder ends with matching 2 × 2 spatial upsampling; additional ResNet resampling did not improve empirical results.These operations define the bottom-layer input and output projections.
  • C.5 Discretized Logistic Likelihood: Pixel data are scaled to [0, 1], and pixel likelihoods use probability mass under a logistic distribution computed from sigmoid-based CDF differences.The supplied formula evaluates the mass over a one-pixel interval using the logistic CDF.
  • C.6 Weight initialisation and normalization: Batch-normalization noise hurt performance, so the experiments used weight normalization with data-dependent parameter initialization.The initialization technique follows Salimans and Kingma (2016).
  • C.7 Nonlinearity: Among ReLU, softplus, and ELU, ELU produced significantly better empirical results and was used throughout the inference and generative models.ELU was used for all reported experiments.

C.8 Objective with Free Bits

Free bits modifies the variational objective with a constant per-group information constraint, avoiding undesirable equilibria that can trap optimization. On CIFAR-10, λ values from 0.125 to 2 improved bits/pixel by more than 0.1 nats.

  • Motivation: The unmodified lower-bound objective often gets stuck in an undesirable stable equilibrium during stochastic optimization.Early in training, weak log p(x|z) makes q(z|x) ≈p(z) attractive, leaving encoder gradients with low signal-to-noise ratio.
  • Method: Free bits replaces annealing with a constant modified objective that partitions latent dimensions into K groups or subsets with shared parameters.The grouping can use latent feature maps or individual dimensions when parameters are not shared across dimensions.
  • Method: The objective ensures that using less than λ nats of information per latent subset is not penalized on average per minibatch M.Because increasing latent information benefits the unaffected negative reconstruction-error term, the average KL divergence satisfies Ex∼M [DKL(q(zj|x)||p(zj))] ≥λ in practice.
  • Results: More than 0.1 nats improvement in bits/pixel on CIFAR-10 resulted from λ ∈ [0.125, 0.25, 0.5, 1, 2].The tested values were λ ∈[0, 0.125, 0.25, 0.5, 1, 2, 4, 8].

C.9 IAF architectural comparison · D Equivalence with Autoregressive Priors

The paper compares IAF architectures within ResNet VAEs on CIFAR-10 and reports a best result of 3.11 bits/dim. It also explains that IAF posterior improvement is equivalent to using an autoregressive prior under an alternative latent-variable representation.

  • C.9 IAF architectural comparison: On CIFAR-10, the study evaluates ResNet VAE depths, approximate posteriors, and inference directions.The comparison covers various architectural and inference choices.
  • C.9 IAF architectural comparison: 3.11 bits/dim is the best CIFAR-10 result, obtained with a depth-20 ResNet VAE.This result is reported in table 2.
  • C.9 IAF architectural comparison: The best configuration uses bidirectional inference with nonlinear IAF, two hidden layers, and one step.These architectural choices produced the reported best CIFAR-10 result.
  • D Equivalence with Autoregressive Priors: The method improves the fully-factorized posterior rather than directly improving the latent-variable prior.The paper contrasts this approach with earlier work focused on improving p(z).
  • D Equivalence with Autoregressive Priors: When latent variables y are treated as the inference target, the corresponding prior is autoregressive.This equivalence follows from the analysis in section 3.
  • D Equivalence with Autoregressive Priors: When whitened variables z are treated as the variables of interest, the prior is fully-factorized and the posterior uses an IAF.The IAF whitens the data according to equation 7.
Loading 1606.04934v2…