🤖 machine learning

Amortized Factor Inference Networks for Posterior Inference

This paper introduces Amortized Factor Inference Networks (AFINs), a novel architecture that enables a single trained inference network to generalize across varying priors, likelihoods, and dimensionalities, achieving posterior accuracy comparable to state-of-the-art methods while reducing test-time compute by 2 to 4 orders of magnitude.

Original authors: Joohwan Ko, Justin Domke

Published 2026-05-27

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Joohwan Ko, Justin Domke

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

The Big Problem: The "One-Size-Fits-None" Dilemma

Imagine you are a detective trying to solve a mystery. In the world of statistics, this mystery is called Bayesian Inference. You have some clues (data) and a theory about how the world works (a model), and you want to figure out the most likely explanation for what happened.

Traditionally, solving this mystery is slow. It's like trying to solve a new puzzle every single time by hand, testing millions of different pieces until you find the right fit. This is accurate but takes forever.

To speed things up, scientists invented "amortized inference." Think of this as training a super-smart AI detective. You show it thousands of puzzles, and it learns to guess the answer instantly.

But here's the catch: The old AI detectives were trained on one specific type of puzzle. If you gave them a puzzle with a different number of clues, a different type of clue, or a different size of the mystery, they would get confused. You'd have to fire them and hire a new one, or spend days retraining them.

The Solution: The "Universal Detective" (AFIN)

The authors of this paper, Joohwan Ko and Justin Domke, asked: Can we build one single AI detective that can solve ANY mystery, no matter the clues, the number of clues, or the size of the case?

They say yes, and they built it. They call it AFIN (Amortized Factor Inference Network).

How It Works: The LEGO Analogy

Imagine you have a box of LEGO bricks.

Old AI: You built a specific castle out of those bricks. If you wanted to build a spaceship, you had to take the whole castle apart and start over.
AFIN: Instead of building a fixed castle, AFIN is like a master builder who looks at the instructions (the model) and the bricks you have (the data) and instantly knows how to assemble the perfect structure.

AFIN works in three steps, like a factory assembly line:

The Encoder (The Translator):
Every mystery has different parts. Some parts are "priors" (what we think before seeing clues), and some are "likelihoods" (the clues themselves). These parts can be very different—some are numbers, some are categories, some are huge, some are tiny.
AFIN has a special translator that looks at each part and turns it into a standard "LEGO brick" (an embedding). It doesn't matter if the original part was a giant boulder or a tiny pebble; the translator turns them all into the same standard shape so the machine can understand them.
The Merger (The Team Huddle):
Once all the parts are translated into standard bricks, AFIN puts them in a circle. It uses a special attention mechanism (like a team huddle) where every brick talks to every other brick. They share information to figure out how they fit together.
- The Magic: This step doesn't care if there are 5 bricks or 500 bricks. The "huddle" works the same way regardless of the group size.
The Decoder (The Architect):
Finally, the merged information is passed to an architect who draws the final blueprint. This blueprint is the "answer" to the mystery (the posterior distribution). It tells you exactly what the solution looks like.

Why Is This a Big Deal?

The paper claims three major wins for this new "Universal Detective":

It's Fast:
The old way of solving these mysteries (using methods like NUTS) is like walking through a maze blindfolded, feeling every wall. It's accurate but slow. AFIN is like having a map. It solves the problem 100 to 10,000 times faster than the old methods.
- Analogy: If the old method takes 100 seconds to solve a puzzle, AFIN might take 0.01 seconds.
It's Flexible (Zero-Shot):
You can train AFIN once on a bunch of made-up puzzles. Then, you can hand it a completely new puzzle it has never seen before—maybe with more clues, fewer clues, or different types of clues—and it solves it instantly without needing to be retrained.
- Analogy: You learn to drive a car on a training course. With AFIN, you can hop into a truck, a motorcycle, or a boat, and it knows how to drive them all immediately.
It's Accurate:
Despite being incredibly fast and flexible, AFIN is just as good at finding the right answer as the slow, careful methods. In fact, if you use AFIN's quick guess as a starting point for a more careful check, it becomes even more accurate.

The "Box" Secret

The paper mentions "dimension-independent modules" and "BoxMLPs." Here is the simple version:

Usually, computer brains are built with fixed sizes. If you want to process a list of 10 items, you build a brain for 10 items. If you have 100 items, you need a different brain.

AFIN uses a special kind of brain cell (a "BoxMLP") that works like a stencil.

Imagine a stencil that says "Process this number."
You can hold it up to one number, or a thousand numbers. The stencil doesn't change; you just move it down the line.
Because the "stencil" doesn't change size, the AI doesn't need to learn new rules for bigger or smaller problems. It just applies the same rule over and over.

Summary

The paper introduces AFIN, a new type of AI that acts as a universal solver for statistical mysteries.

Old way: Train a specific solver for every specific problem. Slow and rigid.
AFIN way: Train one universal solver that understands the structure of any problem, regardless of size or type. Fast, flexible, and accurate.

The authors tested this on synthetic puzzles and real-world datasets (like predicting car fuel efficiency or diagnosing diseases from data) and found that AFIN could solve them instantly with high accuracy, beating traditional methods by orders of magnitude in speed.

Technical Summary: Amortized Factor Inference Networks (AFINs)

1. Problem Statement

Amortized inference aims to learn a single neural network that maps observations to posterior distributions, offering fast test-time inference after a one-time training cost. However, existing amortized methods are typically tied to fixed probabilistic models. If the prior, likelihood, latent dimensionality ( $d$ ), or number of observations ( $N$ ) changes, the inference network usually requires retraining or costly test-time fine-tuning. While some recent works generalize to varying $N$ or specific parameter changes, they often lack the ability to handle arbitrary combinations of varying prior/likelihood families, dimensionality, and observation counts simultaneously without sacrificing accuracy or requiring model-specific adaptation.

The central question posed by this work is: Can a single inference network generalize across an entire family of distinct Bayesian inference problems, handling varying priors, likelihoods, dimensions, and observation counts, while matching the accuracy of non-amortized methods like NUTS?

2. Methodology: Amortized Factor Inference Networks (AFINs)

The authors introduce Amortized Factor Inference Networks (AFINs), an encode-merge-decode architecture designed to map a typed representation of a model specification and its observations directly to the parameters of a variational posterior.

A. Typed-Factor Representation
Instead of treating the model as a monolithic block, AFIN represents the joint density $p(z, y_{1:N})$ as a collection of $N+1$ factors: one prior factor $p_0(z)$ and $N$ likelihood factors $p_n(y_n|z)$ . Each factor is represented by:

Type ( $t_n$ ): A label indicating the distribution family (e.g., "gaussian likelihood," "student-t likelihood").
Parameters ( $\theta_n$ ): Factor-specific parameters (e.g., scale, degrees of freedom) required to instantiate the factor.
Observations ( $y_n$ ): The data associated with the likelihood factors.

This representation allows the network to handle heterogeneous models where different observations may follow different likelihood families.

B. Dimension-Independent Architecture
To support varying latent dimensions ( $d$ ) and numbers of factors ( $N$ ), AFIN employs modules that are independent of $d$ and equivariant to permutations of latent coordinates:

BoxMLP: A parameter-tied Multilayer Perceptron where learned weights do not depend on $d$ . It processes coordinate-wise features and pooled summaries (e.g., mean over dimensions) to ensure the output is equivariant to permutations of the latent vector $z \in \mathbb{R}^d$ .
BoxTransformer: A dimension-independent analogue of the standard Transformer. It operates on a sequence of factor embeddings. Crucially, the projection maps (Query, Key, Value, Output, FFN) are implemented as BoxMLPs applied to coordinate-indexed tokens. This allows the network to model interactions between factors without flattening the latent dimension, preserving the $O(d)$ or $O(d^2)$ structure of the parameters.

C. The Encode-Merge-Decode Pipeline

Encoding: Each factor is processed independently. A type-specific adapter maps the factor's parameters and observations to a common embedding space, followed by a shared BoxMLP encoder. The embedding includes both node features (coordinate-wise) and pair features (coordinate-pair interactions, useful for covariance/precision matrices).
Merging: The sequence of factor embeddings is processed by $M$ BoxTransformer blocks. These blocks use attention across the factor axis (treating factors as the sequence dimension) to aggregate information. The attention mechanism combines node and pair scores to update the representations.
Decoding: The merged representations are sum-pooled over the factor axis to produce a single task-level embedding. A final BoxMLP decoder maps this embedding to the variational parameters $\phi$ (e.g., mean and precision for a Gaussian, or flow parameters for a normalizing flow).

D. Training Objective
AFIN is trained via simulation-based inference. The network is trained to maximize the log density assigned to latent draws $z$ sampled from the joint model $p(z, y)$ , conditioned on the generated observations $y$ . This is equivalent to minimizing the expected forward KL divergence ( $KL(p(z|y) || q_\phi(z))$ ) between the true posterior and the variational approximation. The forward KL objective is noted to be beneficial for using the approximation as a proposal for importance sampling.

3. Key Contributions

Generalization: A single trained AFIN can perform posterior inference across varying prior families, likelihood families, latent dimensions ( $d$ ), and numbers of observations ( $N$ ) without retraining.
Dimension Independence: The introduction of BoxMLP and BoxTransformer modules allows the network to operate on latent spaces of arbitrary dimension $d$ while maintaining a fixed parameter count, a significant departure from standard amortized networks that require fixed input/output sizes.
Factor-Level Abstraction: By operating on typed factors rather than fixed model graphs, AFIN supports heterogeneous likelihoods (e.g., mixing Gaussian and Bernoulli observations) within a single inference run.
Efficiency: The method achieves high accuracy with significantly lower test-time compute compared to iterative methods.

4. Experimental Results

The authors evaluated a single 5M-parameter AFIN checkpoint on synthetic and real-world tasks without fine-tuning.

Synthetic Performance: Across 48 settings (16 prior-likelihood combinations $\times$ $\times$ 3 difficulty levels), AFIN achieved posterior accuracy comparable to NUTS and various Variational Inference (VI) baselines (Full-Rank Gaussian, IAF, MAF, NSF).
- Speed: AFIN single-shot inference is 2 to 4 orders of magnitude faster than NUTS and flow-based VI methods.
- Refinement: When used as a proposal for Self-Normalized Importance Sampling (SNIS), AFIN+SNIS matched or exceeded NUTS accuracy with substantially less wall-clock time.
Extrapolation: AFIN demonstrated zero-shot generalization to out-of-distribution (OOD) settings, including latent dimensions ( $d=32$ ) and observation counts ( $N=512$ ) larger than those seen during training ( $d \le 16, N \le 256$ ).
Real-World Transfer: On 12 UCI datasets (including heterogeneous likelihood settings), AFIN+SNIS reached competitive accuracy with long-run NUTS references almost immediately, whereas iterative baselines required significant optimization time to converge.
Decoder Variants: Both Gaussian and Flow decoders were tested. The Flow decoder offered slight accuracy improvements over the Gaussian decoder while maintaining the speed advantage of amortization.

5. Significance and Claims

The paper claims that AFIN successfully bridges the gap between the speed of amortized inference and the flexibility of non-amortized methods. By decoupling the inference network from fixed model specifications through a typed-factor interface and dimension-independent modules, AFIN enables reusable amortized inference.

The authors position AFIN as a step toward a "foundation model" for Bayesian inference, capable of handling diverse probabilistic programs. They acknowledge limitations, such as the requirement for a predefined set of factor types and the $O(d^2)$ scaling of the pair-component embeddings, but emphasize that the approach allows for a single network to solve a broad family of inference problems that previously required distinct, retrained models. The work suggests that amortization can match the quality of state-of-the-art sampling methods (NUTS) while offering massive computational savings at test time.