Amortized Factor Inference Networks for Posterior Inference
Pith reviewed 2026-06-29 19:02 UTC · model grok-4.3
The pith
A single trained network maps model specifications to accurate variational posteriors for models never seen in training.
A machine-rendered reading of the paper's core claim, the machinery that carries it, and where it could break.
Core claim
The authors show that dimension-independent modules can be composed into an inference network that accepts arbitrary model specifications (priors and likelihoods) together with data and directly produces the parameters of a variational posterior; after training on a distribution of models, the same network achieves posterior accuracy comparable to NUTS and several variational inference baselines while requiring two to four orders of magnitude less test-time compute on previously unseen models.
What carries the argument
Amortized Factor Inference Networks (AFINs): encode-merge-decode networks built from dimension-independent modules that map a model specification and observations to variational posterior parameters.
If this is right
- Bayesian inference on new models no longer requires per-model retraining or test-time optimization.
- A single network can serve an entire family of models that differ in dimensionality, prior, and likelihood.
- Test-time cost drops by 100-10000x compared with NUTS while retaining comparable posterior quality.
- Inference networks become practical for settings where the generative model changes between queries.
Where Pith is reading between the lines
- The same architecture could be applied to families of scientific models that share a common observation format but vary in their parameter priors.
- If the modules truly separate dimension from structure, the approach may extend to non-Euclidean domains such as graphs or point processes without architectural redesign.
- Training data could be generated on the fly from a meta-distribution over models, removing the need for a fixed training corpus.
Load-bearing premise
A fixed collection of dimension-independent modules can faithfully encode arbitrary priors and likelihoods so that the network generalizes accurately to models outside the training distribution.
What would settle it
Train one AFIN on a broad but finite set of models, then evaluate its posterior accuracy on a new model family whose prior or likelihood structure differs qualitatively from everything seen in training; if accuracy collapses relative to per-model baselines, the generalization claim fails.
Figures
read the original 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.
Editorial analysis
A structured set of objections, weighed in public.
Referee Report
Summary. The paper introduces Amortized Factor Inference Networks (AFINs), a family of encode-merge-decode architectures with dimension-independent modules that take a model specification (priors and likelihoods) plus observations as input and output parameters of a variational posterior. It claims that a single trained AFIN generalizes to unseen models with varying priors, likelihoods, and dimensionality, achieving posterior accuracy comparable to NUTS and several variational inference methods while using 2–4 orders of magnitude less test-time compute.
Significance. If the generalization result holds beyond the training distribution, the work would be a notable contribution to amortized inference by removing the need for per-model retraining or finetuning. The release of code supports reproducibility.
major comments (2)
- [§4] §4: The experiments evaluate only held-out instances drawn from the same generative process used to create the training models. This does not establish the headline claim of generalization to arbitrary unseen priors, likelihoods, or dependency structures outside that process, as required by the abstract and the dimension-independent architecture in §3.
- [Abstract, §1] Abstract and §1: No information is supplied on the concrete model families tested, the precise encoding of arbitrary priors/likelihoods into the finite input representation, error bars on accuracy metrics, or data exclusion criteria. These omissions make it impossible to assess whether the reported accuracy is load-bearing for the generalization claim.
minor comments (1)
- [§3] The input representation for priors and likelihoods (e.g., how functional forms or factor graphs are encoded) should be stated explicitly in §3 so that readers can judge coverage of the space of possible models.
Simulated Author's Rebuttal
We thank the referee for the detailed and constructive comments. We address each major point below, clarifying the scope of our generalization results and committing to additions that improve the manuscript's clarity and completeness.
read point-by-point responses
-
Referee: [§4] §4: The experiments evaluate only held-out instances drawn from the same generative process used to create the training models. This does not establish the headline claim of generalization to arbitrary unseen priors, likelihoods, or dependency structures outside that process, as required by the abstract and the dimension-independent architecture in §3.
Authors: We agree that the experiments in §4 evaluate generalization only to held-out models sampled from the same generative process (which itself varies priors, likelihoods, and dimensionalities). This supports the practical claim that one trained AFIN can be applied to new models without retraining or finetuning, but it does not demonstrate performance on models whose dependency structures or families lie entirely outside the training meta-distribution. The abstract and §1 use 'unseen models' to refer to instances not encountered during training but drawn from the same process; we will revise the abstract, §1, and §4 to state this scope explicitly and avoid any implication of extrapolation to arbitrary models beyond the meta-distribution. revision: partial
-
Referee: [Abstract, §1] Abstract and §1: No information is supplied on the concrete model families tested, the precise encoding of arbitrary priors/likelihoods into the finite input representation, error bars on accuracy metrics, or data exclusion criteria. These omissions make it impossible to assess whether the reported accuracy is load-bearing for the generalization claim.
Authors: These details are indeed necessary for proper evaluation. §4 specifies the model families (mixtures of Gaussian, Bernoulli, and Poisson likelihoods with randomly generated priors and sparse dependency graphs) and the generative process used for both training and held-out test models. §3.2 describes the encoding of priors and likelihoods as fixed-length vectors that include type indicators and parameter values. We will add error bars to all accuracy metrics, explicitly document the data exclusion criteria (e.g., rejection of invalid parameter draws), and expand the description of the input encoding in the revised manuscript. revision: yes
Circularity Check
No significant circularity; architecture and empirical claims are independent of fitted inputs
full rationale
The paper defines a new encode-merge-decode architecture for amortized inference over varying model specifications and evaluates it empirically on held-out instances drawn from the same generative process used in training. No derivation step reduces a reported accuracy or generalization result to a quantity defined by the training procedure itself; the central claim rests on direct comparison to NUTS and VI baselines rather than any self-definitional mapping or fitted-parameter renaming. No load-bearing self-citations appear in the provided text, and the method is presented as a constructive proposal rather than a theorem derived from prior author work.
Axiom & Free-Parameter Ledger
Reference graph
Works this paper leans on
-
[1]
Disentangling impact of capacity, objective, batchsize, estimators, and step-size on flow vi.International Conference on Artificial Intelligence and Statistics (AISTATS), 2025
Abhinav Agrawal and Justin Domke. Disentangling impact of capacity, objective, batchsize, estimators, and step-size on flow vi.International Conference on Artificial Intelligence and Statistics (AISTATS), 2025
2025
-
[2]
Stochastic interpolants: A unifying framework for flows and diffusions.Journal of Machine Learning Research, 2025
Michael Albergo, Nicholas M Boffi, and Eric Vanden-Eijnden. Stochastic interpolants: A unifying framework for flows and diffusions.Journal of Machine Learning Research, 2025
2025
-
[3]
Approximate bayesian computation in population genetics.Genetics, 2002
Mark A Beaumont, Wenyang Zhang, and David J Balding. Approximate bayesian computation in population genetics.Genetics, 2002
2002
-
[4]
Variational inference: A review for statisticians.Journal of the American statistical Association, 2017
David M Blei, Alp Kucukelbir, and Jon D McAuliffe. Variational inference: A review for statisticians.Journal of the American statistical Association, 2017
2017
-
[5]
Hoffman, Daniel Lee, Ben Goodrich, Michael Betancourt, Marcus Brubaker, Jiqiang Guo, Peter Li, and Allen Riddell
Bob Carpenter, Andrew Gelman, Matthew D. Hoffman, Daniel Lee, Ben Goodrich, Michael Betancourt, Marcus Brubaker, Jiqiang Guo, Peter Li, and Allen Riddell. Stan: A probabilistic programming language.Journal of Statistical Software, 2017
2017
-
[6]
Amortized probabilistic conditioning for optimization, simulation and inference
Paul Edmund Chang, Nasrulloh Ratu Bagus Satrio Loka, Daolang Huang, Ulpu Remes, Samuel Kaski, and Luigi Acerbi. Amortized probabilistic conditioning for optimization, simulation and inference. InInternational Conference on Artificial Intelligence and Statistics (AISTATS), 2025
2025
-
[7]
Małgorzata Charytanowicz, Jerzy Niewczas, Piotr Kulczycki, Piotr Kowalski, and Szymon Łukasik. Seeds. UCI Machine Learning Repository, 2010
2010
-
[8]
Diffusion posterior sampling for general noisy inverse problems.International Conference on Learning Representations (ICLR), 2023
Hyungjin Chung, Jeongsol Kim, Michael T Mccann, Marc L Klasky, and Jong Chul Ye. Diffusion posterior sampling for general noisy inverse problems.International Conference on Learning Representations (ICLR), 2023
2023
-
[9]
Forest Fires
Paulo Cortez and Aníbal Morais. Forest Fires. UCI Machine Learning Repository, 2007
2007
-
[10]
The frontier of simulation-based inference
Kyle Cranmer, Johann Brehmer, and Gilles Louppe. The frontier of simulation-based inference. Proceedings of the National Academy of Sciences, 2020
2020
-
[11]
Density estimation using real nvp
Laurent Dinh, Jascha Sohl-Dickstein, and Samy Bengio. Density estimation using real nvp. In International Conference on Learning Representations (ICLR), 2017
2017
-
[12]
Neural spline flows
Conor Durkan, Artur Bekasov, Iain Murray, and George Papamakarios. Neural spline flows. Neural Information Processing Systems (NeurIPS), 2019
2019
-
[13]
On contrastive learning for likelihood- free inference
Conor Durkan, Iain Murray, and George Papamakarios. On contrastive learning for likelihood- free inference. InInternational Conference on Machine Learning (ICML), 2020
2020
-
[14]
Least angle regression
Bradley Efron, Trevor Hastie, Iain Johnstone, and Robert Tibshirani. Least angle regression. The Annals of Statistics, 2004
2004
-
[15]
Lasse Elsemüller, Hans Olischläger, Marvin Schmitt, Paul-Christian Bürkner, Ullrich Köthe, and Stefan T. Radev. Sensitivity-aware amortized bayesian inference.Transactions on Machine Learning Research, 2024
2024
-
[16]
Computer Hardware
Jacob Feldmesser. Computer Hardware. UCI Machine Learning Repository, 1987
1987
-
[17]
Rotskoff, and Eric Vanden-Eijnden
Marylou Gabrié, Grant M. Rotskoff, and Eric Vanden-Eijnden. Adaptive monte carlo augmented with normalizing flows.Proceedings of the National Academy of Sciences, 2022
2022
-
[18]
Compositional score modeling for simulation-based inference
Tomas Geffner, George Papamakarios, and Andriy Mnih. Compositional score modeling for simulation-based inference. InInternational Conference on Machine Learning (ICML), 2023
2023
-
[19]
Analytical Methods for Social Research
Andrew Gelman and Jennifer Hill.Data Analysis Using Regression and Multilevel/Hierarchical Models. Analytical Methods for Social Research. Cambridge University Press, Cambridge, 2007
2007
-
[20]
Gerritsma, R
J. Gerritsma, R. Onnink, and A. Versluis. Yacht Hydrodynamics. UCI Machine Learning Repository, 1981. 11
1981
-
[21]
Weilbach, Frank Wood, and Jakob H
Manuel Gloeckler, Michael Deistler, Christian D. Weilbach, Frank Wood, and Jakob H. Macke. All-in-one simulation-based inference. InInternational Conference on Machine Learning (ICML), 2024
2024
-
[22]
Automatic posterior transformation for likelihood-free inference
David Greenberg, Marcel Nonnenmacher, and Jakob Macke. Automatic posterior transformation for likelihood-free inference. InInternational Conference on Machine Learning (ICML), 2019
2019
-
[23]
Haberman
S. Haberman. Haberman’s Survival. UCI Machine Learning Repository, 1976
1976
-
[24]
Likelihood-free MCMC with amortized approximate ratio estimators
Joeri Hermans, V olodimir Begy, and Gilles Louppe. Likelihood-free MCMC with amortized approximate ratio estimators. InInternational Conference on Machine Learning (ICML), 2020
2020
-
[25]
Hoffman and Andrew Gelman
Matthew D. Hoffman and Andrew Gelman. The no-u-turn sampler: Adaptively setting path lengths in hamiltonian monte carlo.Journal of Machine Learning Research, 2014
2014
-
[26]
Stochastic variational inference.Journal of Machine Learning Research, 2013
Matthew D Hoffman, David M Blei, Chong Wang, and John Paisley. Stochastic variational inference.Journal of Machine Learning Research, 2013
2013
-
[27]
Heart Disease
Andras Janosi, William Steinbrunn, Matthias Pfisterer, and Robert Detrano. Heart Disease. UCI Machine Learning Repository, 1989
1989
-
[28]
Heller, Yian Ma, and Michael I
Ghassen Jerfel, Serena Wang, Clara Wong-Fannjiang, Katherine A. Heller, Yian Ma, and Michael I. Jordan. Variational refinement for importance sampling using the forward kullback- leibler divergence. InConference on Uncertainty in Artificial Intelligence (UAI), 2021
2021
-
[29]
The UCI machine learning repository
Markelle Kelly, Rachel Longjohn, and Kolby Nottingham. The UCI machine learning repository. UCI Machine Learning Repository, 2023
2023
-
[30]
Improved variational inference with inverse autoregressive flow
Durk P Kingma, Tim Salimans, Rafal Jozefowicz, Xi Chen, Ilya Sutskever, and Max Welling. Improved variational inference with inverse autoregressive flow. InNeural Information Process- ing Systems (NeurIPS), 2016
2016
-
[31]
Model-informed flows for bayesian inference
Joohwan Ko and Justin Domke. Model-informed flows for bayesian inference. InNeural Information Processing Systems (NeurIPS), 2025
2025
-
[32]
Inference compilation and universal probabilistic programming
Tuan Anh Le, Atilim Gunes Baydin, and Frank Wood. Inference compilation and universal probabilistic programming. InInternational Conference on Artificial Intelligence and Statistics (AISTATS), 2017
2017
-
[33]
Flow matching for generative modeling.International Conference on Learning Representations (ICLR), 2023
Yaron Lipman, Ricky TQ Chen, Heli Ben-Hamu, Maximilian Nickel, and Matt Le. Flow matching for generative modeling.International Conference on Learning Representations (ICLR), 2023
2023
-
[34]
Parkinsons
Max Little. Parkinsons. UCI Machine Learning Repository, 2007
2007
-
[35]
Benchmarking simulation-based inference
Jan-Matthis Lueckmann, Jan Boelts, David Greenberg, Pedro Goncalves, and Jakob Macke. Benchmarking simulation-based inference. InInternational Conference on Artificial Intelligence and Statistics (AISTATS), 2021
2021
-
[36]
Reconstructing the universe with variational self-boosted sampling.Journal of Cosmology and Astroparticle Physics, 2023
Chirag Modi, Yin Li, and David Blei. Reconstructing the universe with variational self-boosted sampling.Journal of Cosmology and Astroparticle Physics, 2023
2023
-
[37]
Transformers can do bayesian inference
Samuel Müller, Noah Hollmann, Sebastian Pineda Arango, Josif Grabocka, and Frank Hutter. Transformers can do bayesian inference. InInternational Conference on Learning Representa- tions (ICLR), 2022
2022
-
[38]
Markovian score climbing: Variational inference withKL(p||q).Neural Information Processing Systems (NeurIPS), 2020
Christian Naesseth, Fredrik Lindsten, and David Blei. Markovian score climbing: Variational inference withKL(p||q).Neural Information Processing Systems (NeurIPS), 2020
2020
-
[39]
Mcmc using hamiltonian dynamics.Handbook of markov chain monte carlo, 2011
Radford M Neal. Mcmc using hamiltonian dynamics.Handbook of markov chain monte carlo, 2011
2011
-
[40]
Fastϵ-free inference of simulation models with bayesian conditional density estimation
George Papamakarios and Iain Murray. Fastϵ-free inference of simulation models with bayesian conditional density estimation. InNeural Information Processing Systems (NeurIPS), 2016. 12
2016
-
[41]
Masked autoregressive flow for density estimation.Neural Information Processing Systems (NeurIPS), 2017
George Papamakarios, Theo Pavlakou, and Iain Murray. Masked autoregressive flow for density estimation.Neural Information Processing Systems (NeurIPS), 2017
2017
-
[42]
Sequential neural likelihood: Fast likelihood-free inference with autoregressive flows
George Papamakarios, David Sterratt, and Iain Murray. Sequential neural likelihood: Fast likelihood-free inference with autoregressive flows. InInternational Conference on Artificial Intelligence and Statistics (AISTATS), 2019
2019
-
[43]
Normalizing flows for probabilistic modeling and inference.Journal of Machine Learning Research, 2021
George Papamakarios, Eric Nalisnick, Danilo Jimenez Rezende, Shakir Mohamed, and Balaji Lakshminarayanan. Normalizing flows for probabilistic modeling and inference.Journal of Machine Learning Research, 2021
2021
-
[44]
R. Quinlan. Auto MPG. UCI Machine Learning Repository, 1993
1993
-
[45]
Radev, Ulf K
Stefan T. Radev, Ulf K. Mertens, Andreas V oss, Lynton Ardizzone, and Ullrich Köthe. Bayesflow: Learning complex stochastic models with invertible neural networks.IEEE Trans- actions on Neural Networks and Learning Systems, 2022
2022
-
[46]
Radev, Marvin Schmitt, Valentin Pratz, Umberto Picchini, Ullrich Köthe, and Paul- Christian Bürkner
Stefan T. Radev, Marvin Schmitt, Valentin Pratz, Umberto Picchini, Ullrich Köthe, and Paul- Christian Bürkner. JANA: Jointly amortized neural approximation of complex bayesian models. InConference on Uncertainty in Artificial Intelligence (UAI), 2023
2023
-
[47]
Black Box Variational Inference
Rajesh Ranganath, Sean Gerrish, and David Blei. Black Box Variational Inference. InInterna- tional Conference on Artificial Intelligence and Statistics (AISTATS), 2014
2014
-
[48]
Arik Reuter, Tim G. J. Rudner, Vincent Fortuin, and David Rügamer. Can transformers learn full bayesian inference in context? InInternational Conference on Machine Learning (ICML), 2025
2025
-
[49]
Variational inference with normalizing flows
Danilo Rezende and Shakir Mohamed. Variational inference with normalizing flows. In International Conference on Machine Learning (ICML), 2015
2015
-
[50]
Terry Sejnowski and R. Gorman. Connectionist Bench (Sonar, Mines vs. Rocks). UCI Machine Learning Repository, 1988
1988
-
[51]
Sequential neural score estima- tion: Likelihood-free inference with conditional score based diffusion models.International Conference on Machine Learning (ICML), 2024
Louis Sharrock, Jack Simons, Song Liu, and Mark Beaumont. Sequential neural score estima- tion: Likelihood-free inference with conditional score based diffusion models.International Conference on Machine Learning (ICML), 2024
2024
-
[52]
Gonçalves, David S
Alvaro Tejero-Cantero, Jan Boelts, Michael Deistler, Jan-Matthis Lueckmann, Conor Durkan, Pedro J. Gonçalves, David S. Greenberg, and Jakob H. Macke. sbi: A toolkit for simulation- based inference.Journal of Open Source Software, 2020
2020
-
[53]
Improving and generalizing flow-based generative models with minibatch optimal transport.Transactions on Machine Learning Research, 2024
Alexander Tong, Kilian Fatras, Nikolay Malkin, Guillaume Huguet, Yanlei Zhang, Jarrid Rector- Brooks, Guy Wolf, and Yoshua Bengio. Improving and generalizing flow-based generative models with minibatch optimal transport.Transactions on Machine Learning Research, 2024
2024
-
[54]
Attention is all you need.Neural Information Processing Systems (NeurIPS), 2017
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin. Attention is all you need.Neural Information Processing Systems (NeurIPS), 2017
2017
-
[55]
Julius Vetter, Manuel Gloeckler, Daniel Gedon, and Jakob H. Macke. Effortless, simulation- efficient bayesian inference using tabular foundation models. InNeural Information Processing Systems (NeurIPS), 2025
2025
-
[56]
George Whittle, Juliusz Ziomek, Jacob Rawling, and Michael A. Osborne. Distribution trans- formers: Fast approximate bayesian inference with on-the-fly prior adaptation, 2026
2026
-
[57]
Flow matching for scalable simulation-based inference.Neural Information Processing Systems (NeurIPS), 2023
Jonas Wildberger, Maximilian Dax, Simon Buchholz, Stephen Green, Jakob H Macke, and Bern- hard Schölkopf. Flow matching for scalable simulation-based inference.Neural Information Processing Systems (NeurIPS), 2023
2023
-
[58]
Foundation posteriors for approximate probabilistic inference
Mike Wu and Noah Goodman. Foundation posteriors for approximate probabilistic inference. InNeural Information Processing Systems (NeurIPS), 2022. 13
2022
-
[59]
Goodman, and Stefano Ermon
Mike Wu, Kristy Choi, Noah D. Goodman, and Stefano Ermon. Meta-amortized variational inference and learning. InProceedings of the AAAI Conference on Artificial Intelligence, 2020
2020
-
[60]
Concrete Slump Test
I-Cheng Yeh. Concrete Slump Test. UCI Machine Learning Repository, 2007
2007
-
[61]
Transport score climbing: Variational inference using forward kl and adaptive neural transport.Transactions on Machine Learning Research, 2023
Liyi Zhang, David M Blei, and Christian A Naesseth. Transport score climbing: Variational inference using forward kl and adaptive neural transport.Transactions on Machine Learning Research, 2023. 14 A Notations Table 4 summarizes the main notation used throughout the paper. We index the prior factor by n= 0 and the likelihood factors by n= 1, . . . , N . ...
2023
discussion (0)
Sign in with ORCID, Apple, or X to comment. Anyone can read and Pith papers without signing in.