Amortized Factor Inference Networks for Posterior Inference

Abstract: 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.

• The paper introduces AFIN, an inference network that generalizes conjugate Bayesian updates to arbitrary models using a dimension-independent encode-merge-decode approach.
• The paper demonstrates that AFIN achieves posterior accuracy comparable to NUTS with 2–4 orders of magnitude less test-time computation across diverse synthetic and real datasets.
• The paper leverages SNIS correction and flexible decoder designs to enhance inference quality, paving the way for scalable and reusable Bayesian inference.

Amortized Factor Inference Networks (AFINs): Dimension-Independent Amortized Bayesian Inference

Overview

"Amortized Factor Inference Networks for Posterior Inference" (2605.26419) presents AFIN, an inference network architecture enabling rapid, zero-shot Bayesian posterior estimation across heterogeneous latent-variable models. The key innovation is a dimension-independent encode-merge-decode pipeline that generalizes conjugate updates, utilizing parameter-tied modules capable of operating across variable numbers of latent dimensions, factor types, and observations. The system maps arbitrary typed model specifications and datasets to variational posterior parameters, achieves accuracy comparable to No-U-Turn Sampler (NUTS) and various variational inference (VI) baselines, and offers dramatic test-time speedup.

Problem Formulation and Motivation

Traditional amortized inference methods are tightly coupled to specific $p(z, y)$, requiring retraining whenever model structure changes. The paper extends amortization to inference over arbitrary prior and likelihood families, arbitrary observation sets, and latent dimensionalities, offering a universal inference interface. The approach is motivated by the structure of conjugate exponential-family inference, where posterior computation reduces to encoding individual natural parameter contributions, merging via addition, and decoding into posterior parameters. AFIN generalizes each step:

1. Encoding: Typed factor adapters and BoxMLPs translate model parameters and observations to per-factor embeddings.
2. Merging: Factor embeddings are aggregated via BoxTransformer blocks, using factor-axis cross-attention to capture nontrivial factor interactions.
3. Decoding: The pooled merged embedding is mapped to parameters for a tractable variational family (e.g., full-rank Gaussian or normalizing flow).

This modular architecture allows AFIN to process arbitrary combinations of factor types, arbitrary latent dimension $d$, and variable $N$ (number of likelihood factors).

Dimension-Independent Neural Modules

Two principal architectural components guarantee dimension-independent operation:

• BoxMLP: An MLP with parameter-tying such that learned weights are independent of input dimension $d$. Local and pooled summaries ensure permutation equivariance and facilitate both coordinate-local and global context sharing.
• BoxTransformer: A transformer analogue in which attention and projections are performed coordinate-wise with BoxMLPs. Attention scores are aggregated along the factor axis, and outputs retain equivariance in latent variable coordinate permutations.

  Figure 1: Architecture of an AFIN, illustrating encoding of typed factors, merging via BoxTransformer blocks, and decoding into variational posterior parameters.

AFIN Architecture and Training Objective

AFIN receives a typed-factor model specification and observations, infers a node-pair embedding for each factor, aggregates factor embeddings with $M$ BoxTransformer blocks, pools over factors, and decodes into $\phi$ for $q_\phi(z)$. The system is trained via simulation-based inference, maximizing the expected log density assigned to true posterior draws, equivalent to minimizing forward KL divergence from the true posterior to the variational approximation.

Empirical Evaluation

AFIN is evaluated on synthetic and real-world tabular regression/classification tasks with variable model structures, latent dimensionality, and observation count. The experiments demonstrate several claims:

1. Posterior Accuracy vs Compute: AFIN achieves NUTS-level posterior accuracy, but requires 2–4 orders of magnitude less test-time computation.

   Figure 2: Synthetic posterior accuracy (mean error, covariance error, sliced Wasserstein-2) for increasing test-time budgets across AFIN+SNIS, NUTS, and FRVI.

2. SNIS Correction: Using AFIN's single-shot posterior as a proposal for self-normalized importance sampling (SNIS) further improves estimates, often exceeding iterative approaches per unit compute.

   Figure 3: SNIS refinement using AFIN and FRVI proposals, showing rapid improvement in posterior metrics and stable importance weights.

3. Difficulty Regimes: Across easy (low-dimensional, large-$N$), medium, and hard (high-dimensional, small-$N$) regimes, AFIN+SNIS consistently achieves competitive accuracy at minimal cost.

   Figure 4: Synthetic posterior accuracy (easy regime, $d=4$, $d$0) indicates single-shot AFIN dominance in low compute budgets.

   

   Figure 5: Medium difficulty ($d$1, $d$2), with AFIN+SNIS outpacing iterative baselines on posterior metrics.

   

   Figure 6: Hard tasks ($d$3, $d$4), highlighting amortized performance in high-dimensional, data-poor settings.

4. UCI Tabular Benchmarks: Zero-shot transfer to UCI datasets confirms generalization across unseen likelihood families and latent dimensions in practical data settings.

   Figure 7: Zero-shot posterior inference on UCI datasets, showing accuracy and compute tradeoffs for AFIN+SNIS relative to NUTS and VI baselines.

   

   Figure 8: Zero-shot posterior mean accuracy on UCI datasets, with AFIN+SNIS matching or exceeding iterative optimization performance.

   

   Figure 9: Posterior covariance accuracy in UCI testing, further demonstrating strong amortized results.

5. Extrapolation: AFIN generalizes outside of its training support (e.g., $d$5, $d$6), maintaining competitive performance, especially when combined with SNIS correction.
6. Decoder Flexibility: Gaussian and flow-based decoders offer similar accuracy, with flows providing marginal gains in challenging settings at higher decoding cost.

Practical and Theoretical Implications

AFIN's architecture enables amortized inference to be reusable across broad classes of Bayesian models. For practical applications, this reduces the need for per-model retraining or problem-specific proposal construction. The forward-KL objective aligns the output for importance sampling correction, further increasing practical utility.

Theoretically, factor-level amortization with dimension-independent neural modules equips AFIN to represent conjugate and non-conjugate interactions, potentially subsuming classical exponential-family inference in the limit.

AFIN's current limitations are bounded by the set of predefined factor types and adapters; generalization to arbitrary probabilistic program factorization would require expansion of the typed adapter interface or further architectural advances. Scalability in high-dimensional settings, particularly addressing the $d$7 pairwise embedding bottleneck, is a natural direction.

Future Directions

Possible avenues for future research include:

• Expanding adapter interfaces to automate recognition and embedding construction for novel factor types.
• Extending the node-pair architecture to richer probabilistic program structures (e.g., deterministic factors, hierarchical compositions).
• Efficient scaling to $d$8 latent dimensions by sparsified or factorized representations.
• Incorporating uncertainty quantification for unseen factor configurations and out-of-support extrapolation.

Conclusion

AFIN introduces a flexible, efficient amortized inference mechanism capable of handling heterogeneous Bayesian model families with variable latent structure and data regimes. Its encode-merge-decode pipeline, grounded in dimension-independent neural modules, achieves strong posterior accuracy, rapid compute, and robust zero-shot generalization. The methodology substantially broadens the scope of amortized inference from fixed programs to typed factor lists, laying groundwork for scalable, reusable Bayesian inference engines.