Amortized Factor Inference Networks for Posterior Inference
This work addresses the need for fast, amortized Bayesian inference that generalizes across different models, reducing test-time computation by orders of magnitude.
AFINs achieve posterior accuracy comparable to NUTS and variational inference methods while requiring 2 to 4 orders of magnitude less test-time compute, enabling generalization across varying priors, likelihoods, and dimensionality without retraining.
Amortized inference promises fast test-time Bayesian inference, but existing methods are inherently tied to fixed models. Extending amortization to unseen models typically requires retraining or costly test-time finetuning. In this paper, we ask: is it possible to build a single inference network capable of generalizing across varying priors, likelihoods, and dimensionality? We introduce Amortized Factor Inference Networks (AFINs), a family of encode-merge-decode inference networks built on dimension-independent modules that map a model specification and its observations to the parameters of a variational posterior. Experimentally, a single trained AFIN achieves posterior accuracy comparable to NUTS and several variational inference methods, while requiring 2 to 4 orders of magnitude less test-time compute. Code is available at https://github.com/joohwanko/AFINs.