Attention for Inference Compilation

TL;DR

This paper introduces an attention mechanism in inference compilation for probabilistic programming, improving the modeling of long-range dependencies and enhancing posterior approximation accuracy.

Contribution

It proposes a novel attention-based neural network architecture for inference compilation, addressing limitations of previous models in capturing long-range dependencies.

Findings

Attention improves posterior approximation in probabilistic programs.

Enhanced inference performance over existing methods.

Better modeling of long-range dependencies in latent variables.

Abstract

We present a new approach to automatic amortized inference in universal probabilistic programs which improves performance compared to current methods. Our approach is a variation of inference compilation (IC) which leverages deep neural networks to approximate a posterior distribution over latent variables in a probabilistic program. A challenge with existing IC network architectures is that they can fail to model long-range dependencies between latent variables. To address this, we introduce an attention mechanism that attends to the most salient variables previously sampled in the execution of a probabilistic program. We demonstrate that the addition of attention allows the proposal distributions to better match the true posterior, enhancing inference about latent variables in simulators.