Scaling Structured Inference with Randomization

TL;DR

This paper introduces randomized dynamic programming algorithms that significantly scale structured inference to tens of thousands of states, enabling more complex models and efficient training of structured variational autoencoders.

Contribution

It presents a novel randomized DP method applicable to various graph structures, improving scalability and integration with neural networks for structured inference.

Findings

Achieves orders of magnitude reduction in computation.

Maintains low bias and variance through Rao-Blackwellization.

Improves test likelihood and prevents posterior collapse in VAEs.

Abstract

Deep discrete structured models have seen considerable progress recently, but traditional inference using dynamic programming (DP) typically works with a small number of states (less than hundreds), which severely limits model capacity. At the same time, across machine learning, there is a recent trend of using randomized truncation techniques to accelerate computations involving large sums. Here, we propose a family of randomized dynamic programming (RDP) algorithms for scaling structured models to tens of thousands of latent states. Our method is widely applicable to classical DP-based inference (partition, marginal, reparameterization, entropy) and different graph structures (chains, trees, and more general hypergraphs). It is also compatible with automatic differentiation: it can be integrated with neural networks seamlessly and learned with gradient-based optimizers. Our core technique approximates the sum-product by restricting and reweighting DP on a small subset of nodes, which reduces computation by orders of magnitude. We further achieve low bias and variance via Rao-Blackwellization and importance sampling. Experiments over different graphs demonstrate the accuracy and efficiency of our approach. Furthermore, when using RDP for training a structured variational autoencoder with a scaled inference network, we achieve better test likelihood than baselines and successfully prevent posterior collapse. code at: https://github.com/FranxYao/RDP