Iterative Amortized Inference

Joseph Marino; Yisong Yue; Stephan Mandt

Source-linked AI summary

Iterative Amortized Inference

Joseph Marino, Yisong Yue, Stephan Mandt

arXiv:1807.09356v1Jul 24, 2018cs.LGstat.ML

TL;DR

Standard inference models use direct mappings that may fail to reach optimized approximate posteriors, creating an amortization gap. The paper introduces iterative inference models that repeatedly encode approximate-posterior gradients to refine estimates, and reports advantages over standard inference models across image and text benchmarks.

Problem
Standard inference models use direct data-to-posterior mappings that may fail to reach fully optimized approximate posterior estimates, creating an amortization gap.
Method
The paper introduces iterative inference models that encode approximate posterior gradients or errors to iteratively update inference estimates, generalizing standard inference models.
Results
Iterative inference models outperform standard inference models on several benchmark data sets of images and text.
Takeaways & Limitations
Iterative inference models provide theoretical justification for top-down inference techniques while learning to perform variational inference optimization.
Takeaways & Limitations
Iterative inference models require additional computation and may require many inference iterations because their formulation is restricted to first-order approximate posterior derivatives.

Abstract

from arXiv · show

Inference models are a key component in scaling variational inference to deep latent variable models, most notably as encoder networks in variational auto-encoders (VAEs). By replacing conventional optimization-based inference with a learned model, inference is amortized over data examples and therefore more computationally efficient. However, standard inference models are restricted to direct mappings from data to approximate posterior estimates. The failure of these models to reach fully optimized approximate posterior estimates results in an amortization gap. We aim toward closing this gap by proposing iterative inference models, which learn to perform inference optimization through repeatedly encoding gradients. Our approach generalizes standard inference models in VAEs and provides insight into several empirical findings, including top-down inference techniques. We demonstrate the inference optimization capabilities of iterative inference models and show that they outperform standard inference models on several benchmark data sets of images and text.

1. Introduction

Standard inference models make variational inference more efficient by amortizing inference across examples, but direct mappings can leave approximate posteriors under-optimized. Iterative inference models address this gap by repeatedly encoding optimization information and are evaluated as extensions of standard inference models.

1. Introduction: Inference models replace per-example optimization with shared amortized mappings from observations to approximate posterior estimates.This improves computational efficiency, but standard models use direct encoding.
1. Introduction: Direct bottom-up inference models cannot account for top-down priors in hierarchical latent variable models, motivating top-down techniques without rigorous theoretical justification.
1. Introduction: Iterative inference models encode approximate posterior gradients or errors to learn repeated updates, reframing inference as optimization learned through an inference model.The approach is related to learning to learn applied to variational inference optimization.
1. Introduction: The paper introduces a family of iterative inference models that generalize standard inference models and provides a theoretical justification for top-down inference techniques.
1. Introduction: Iterative inference models outperform standard inference models on several image and text data sets.

2. Background

Variational inference turns otherwise intractable posterior inference into optimization of an evidence lower bound. Conventional gradient optimization is expensive for each example, so amortized inference models learn shared mappings from observations to posterior estimates.

2. Background: Latent variable models use local latent variables z to model observations x with global parameters θ through a joint distribution.
2. Background: Variational inference introduces a tractable approximate posterior q(z|x) and minimizes its KL-divergence from the true posterior.The posterior itself prevents direct minimization.
2. Background: Maximizing the ELBO with respect to q(z|x) performs approximate inference because the KL decomposition makes the ELBO a lower bound on log pθ(x).The ELBO combines reconstruction quality with divergence from the prior.
2. Background: Variational inference and learning alternate expectation and maximization steps, typically using stochastic variational inference in batches.
2. Background: Conventional optimization repeatedly estimates stochastic gradients for each example's posterior parameters, making inference computationally expensive and requiring step-size hyperparameters.
2. Background: Amortized inference replaces per-example posterior-parameter optimization with shared parameters in an inference model, as exemplified by VAEs.

3. Iterative Amortized Inference

Iterative inference models address the amortization gap by replacing a single direct mapping with learned, repeated updates based on approximate posterior gradients or errors. They retain amortized inference while enabling optimization-like estimation and are shown to generalize standard inference models and justify top-down inference.

3.1. Standard Inference Models & Amortization Gaps: Standard inference models approximate a difficult non-convex variational optimization through a direct mapping from observations to posterior parameters.This direct mapping can produce sub-optimal approximate posteriors and an amortization gap.
3.1. Standard Inference Models & Amortization Gaps: In a 2-D latent Gaussian MNIST example, gradient ascent reaches a better final estimate than a standard inference model, although it requires many iterations and is step-size sensitive.The standard model produces a reasonable estimate in one pass but does not reach the optimum.
3.2. Learning to Iteratively Optimize: Iterative inference models learn to update approximate posterior estimates by encoding gradients or errors, combining amortized computation with iterative estimation.The approach is motivated as learning to optimize inference while improving efficiency over conventional optimizers.
3.2. Learning to Iteratively Optimize: Iterative inference models generalize standard inference models and provide theoretical justification for top-down inference techniques in hierarchical models.The paper also reports empirical evaluations showing that iterative models outperform standard inference models on image and text data sets.
3.2. Learning to Iteratively Optimize: At each iteration, the model uses approximate posterior gradients from the ELBO to produce updated posterior parameters, with model parameters trained through accumulated gradients.The computation graph initializes λ, samples z, evaluates the ELBO, backpropagates gradients, and repeats the update process.

4. Iterative Inference in Latent Gaussian Models

For latent Gaussian models, iterative inference uses posterior gradients containing bottom-up data errors and top-down prior errors to update posterior parameters. The framework generalizes standard inference under stated assumptions, while derivative approximation can reduce the iterations needed for convergence.

4. Iterative Inference in Latent Gaussian Models: Latent Gaussian models use Gaussian priors and commonly Gaussian approximate posteriors for continuous-valued latent variables.The posterior parameters include the mean and diagonal variance terms.
4.2. Approximating Approximate Posterior Derivatives: Iterative models update the posterior mean and variance through separate learned functions that can share inputs and parameters.The update functions use current posterior estimates together with bottom-up and top-down error signals.
4. Iterative Inference in Latent Gaussian Models: Approximate posterior gradients combine distribution-mismatch errors with Jacobian matrices that invert generative mappings.For Gaussian likelihoods, the gradients can be expressed using bottom-up errors from the data and top-down errors from the prior.
4.1. Top-Down Inference: The gradient terms push the approximate posterior toward agreement with both the data and the prior, making inference simultaneously bottom-up and top-down.In hierarchical models, the top-down prior varies across examples, which can challenge purely bottom-up procedures.
4.2. Approximating Approximate Posterior Derivatives: The basic formulation uses only first-order approximate posterior derivatives, which may require many inference iterations to reach reasonable estimates.The paper motivates learned derivative approximations as a way to reduce iterations without calculating costly higher-order derivatives.
4.3. Generalizing Standard Inference Models: Under assumptions for single-level latent Gaussian models, a one-step iterative inference model is equivalent to a standard inference model.Iterative models instead learn to traverse the optimization landscape beyond the local curvature captured by the one-step special case.

5. Experiments

Experiments show that iterative inference models learn adaptive, multi-step approximate-posterior optimization across image and text tasks. They converge quickly, remain stable beyond training horizons, improve with additional iterations or samples, and outperform standard inference models.

Experimental setup: Experiments evaluate iterative inference models on MNIST, Omniglot, SVHN, CIFAR-10, and RCV1 using latent Gaussian models.Image likelihoods use Bernoulli or Gaussian outputs, while RCV1 word counts use a multinomial output.
Approximate inference optimization: Iterative inference models adaptively adjust update step sizes and quickly reach near-optimal approximate-posterior estimates.A 16-step trajectory on a 2D latent Gaussian MNIST model demonstrates this optimization behavior.
Approximate inference optimization: Iterative inference models converge faster than conventional optimizers to better ELBO estimates and remain stable over hundreds of iterations.They use only local gradients and were trained with 16 inference iterations, unlike optimizers using momentum or similar terms.
Additional inference iterations and samples: Reconstructions become gradually sharper across inference iterations for MNIST, Omniglot, SVHN, and CIFAR-10, then remain stable.The reconstruction term is typically dominant in the objective, linking visual improvement to inference optimization.
Additional inference iterations and samples: Approximate-posterior mean gradient magnitudes decrease over inference iterations on RCV1 as estimates approach local maxima.This pattern is observed during training for a model encoding gradients over 16 inference iterations.
Additional inference iterations and samples: Increasing inference iterations improves performance, all iterative models outperform standard models, and additional samples improve iterative models by more than 1 nat.Encoding errors to approximate higher-order derivatives helps with fewer iterations, while additional samples widen the improvement over similar standard models.
Comparison with standard inference models: Across MNIST, CIFAR-10, and RCV1, iterative inference models outperform standard models for both single-level and hierarchical models.Larger improvements occur on high-dimensional RCV1; fixed generative architectures indicate the gains arise from inference optimization.

6. Conclusion

Iterative inference models refine approximate posterior estimates by encoding gradients or errors, accounting for priors during inference and extending standard inference models. Experiments show optimization advantages over current inference techniques, although the approach requires additional computation.

6. Conclusion: Iterative inference models refine inference estimates by encoding approximate posterior gradients or errors.They learn to update approximate posterior estimates repeatedly rather than relying only on a direct inference mapping.
6. Conclusion: The models naturally account for priors during inference and provide insight into top-down inference techniques.In hierarchical models, gradients include bottom-up reconstruction terms and top-down prior terms.
6. Conclusion: Iterative inference models generalize and extend standard inference models while learning to perform variational inference optimization.The inference process uses approximate-posterior gradients to update estimates iteratively.
6. Conclusion: The approach outperforms current inference techniques on several benchmark data sets of images and text.The paper reports advantages over standard inference models across multiple image and text benchmarks.
6. Conclusion: Additional computation is required compared with similar standard inference models.The authors suggest online filtering as a future setting where fewer inference iterations may reduce this cost.

B.3. Training

Iterative inference models are trained by accumulating stochastic gradient estimates during inference and jointly updating inference and generative parameters.

B.3. Training: Stochastic reparameterization gradients are accumulated during inference before jointly updating φ and θ.Training uses a fixed number of inference iterations.

C. Experiment Details

The models use joint Adam optimization for inference and generative parameters, with shared learning-rate settings and exponential decay.

C. Experiment Details: Inference and generative parameters are jointly trained with Adam using a learning rate of 0.0002.Learning rates decay exponentially by a factor of 0.999 each epoch, and models use ELU activations unless otherwise noted.

C.1. Two-Dimensional Latent Gaussian Models

The two-dimensional latent Gaussian experiments evaluate inference on binarized MNIST using small neural-network generative models and a point-estimate posterior.

C.1. Two-Dimensional Latent Gaussian Models: Binarized MNIST experiments use a two-dimensional Dirac-delta posterior and generative networks with two hidden layers of 512 units.The optimization surface is evaluated on a grid from -5 to 5 in increments of 0.05 for each latent dimension.

C.2. L During Inference

Iterative inference models were evaluated by encoding approximate posterior gradients across multiple inference iterations and compared with conventional optimization techniques. Despite taking more time per iteration, the iterative model outperformed those techniques.

Experimental setup: The MNIST experiment encoded approximate posterior gradients for 16 inference iterations and compared iterative inference with SGD, momentum, RMSProp, and Adam.Learning rates for conventional optimizers were selected from {0.5, 0.4, 0.3, 0.2, 0.1, 0.01, 0.001}, using the best result.
Results: Despite requiring more time per inference iteration, the iterative inference model outperformed conventional optimization techniques.
Experimental setup: Iterative inference models were trained on MNIST, Omniglot, and SVHN with 16 encoded-gradient iterations, while CIFAR-10 used encoded errors for 10 iterations.
Inference behavior: Approximate posterior gradient magnitudes generally decreased during inference optimization, especially for gradients of the approximate posterior mean.

C.5. Additional Inference Iterations

Additional inference-iteration experiments used fixed network and latent-variable architectures across 1,500 training epochs. Multiple trials were unavailable, but observed variation in a subset of runs was smaller than differences associated with inference-iteration counts.

Experimental limitation: Multiple trials could not be run for each experimental setup, limiting direct assessment of performance variability.In a subset of standard and iterative inference runs, final-performance standard deviation was less than 0.1 nats.
Experimental limitation: The observed standard deviation below 0.1 nats was smaller than the performance difference between models trained with different numbers of inference iterations.

C.6. Additional Latent Samples

Additional experiments varied latent-variable model settings and posterior-sample usage across image and text datasets. The models used fixed architectures and training procedures, with text performance evaluated using test-set perplexity.

Experimental setup: The experiment used two hidden layers of 512 units, a 64-dimensional latent variable, five inference iterations, and 1,500 training epochs.Iterative models encoded data and approximate posterior gradients to obtain ELBO estimates and gradients.
Image models: MNIST models used a 64-dimensional one-level latent variable or hierarchical latent variables of sizes 64 and 32.Inference and generative networks used two hidden layers, with 512 units at the first level and 256 at the second hierarchical level.
Image models: CIFAR-10 one-level models used a 1,024-dimensional latent variable, a three-layer encoder with 2,048-unit layers, and a one-layer decoder with 2,048 units.This asymmetric encoder-decoder setup performed better than a symmetric architecture for both standard and iterative inference models.
Text models: The text experiment used normalized TF-IDF features, two-layer 2,048-unit encoder and decoder networks, a 1,024-dimensional latent variable, and 10 gradient-encoding steps.Both models used five approximate posterior samples at each iteration and were evaluated with test-set perplexity.
Evaluation: Test perplexity was estimated by evaluating each log p(x(i)) with 5,000 importance-weighted samples, alongside an upper bound on perplexity.