# It Just Takes Two: Scaling Amortized Inference to Large Sets

arXiv: https://arxiv.org/abs/2605.07972 · HTML: https://mlsanity.com/paper/2605.07972/

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Builds on 24 of its 37 references that are in the map.
- First posted on arXiv 8 May 2026.

## Builds on (most cited first, 24 of 24)

- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [Flow Matching for Generative Modeling](https://mlsanity.com/paper/2210.02747/) (2022)
- [Deep Sets](https://mlsanity.com/paper/1703.06114/) (2017)
- [Normalizing Flows for Probabilistic Modeling and Inference](https://mlsanity.com/paper/1912.02762/) (2019)
- [Perceiver: General Perception with Iterative Attention](https://mlsanity.com/paper/2103.03206/) (2021)
- [Train Short, Test Long: Attention with Linear Biases Enables Input Length Extrapolation](https://mlsanity.com/paper/2108.12409/) (2021)
- [SchNet: A continuous-filter convolutional neural network for modeling quantum interactions](https://mlsanity.com/paper/1706.08566/) (2017)
- [The frontier of simulation-based inference](https://mlsanity.com/paper/1911.01429/) (2019)
- [GEOM: Energy-annotated molecular conformations for property prediction and molecular generation](https://mlsanity.com/paper/2006.05531/) (2020)
- [Sequential Neural Likelihood: Fast Likelihood-free Inference with Autoregressive Flows](https://mlsanity.com/paper/1805.07226/) (2018)
- [Exploring Length Generalization in Large Language Models](https://mlsanity.com/paper/2207.04901/) (2022)
- [Flexible statistical inference for mechanistic models of neural dynamics](https://mlsanity.com/paper/1711.01861/) (2017)
- [Neural Importance Sampling](https://mlsanity.com/paper/1808.03856/) (2018)
- [BayesFlow: Learning complex stochastic models with invertible neural networks](https://mlsanity.com/paper/2003.06281/) (2020)
- [Flow Matching for Scalable Simulation-Based Inference](https://mlsanity.com/paper/2305.17161/) (2023)
- [Unconstrained Monotonic Neural Networks](https://mlsanity.com/paper/1908.05164/) (2019)
- [On the Limitations of Representing Functions on Sets](https://mlsanity.com/paper/1901.09006/) (2019)
- [Neural Approximate Sufficient Statistics for Implicit Models](https://mlsanity.com/paper/2010.10079/) (2020)
- [Universal Approximation of Functions on Sets](https://mlsanity.com/paper/2107.01959/) (2021)
- [Compositional Score Modeling for Simulation-based Inference](https://mlsanity.com/paper/2209.14249/) (2022)
- [Likelihood-Free Parameter Estimation with Neural Bayes Estimators](https://mlsanity.com/paper/2208.12942/) (2022)
- [HNPE: Leveraging Global Parameters for Neural Posterior Estimation](https://mlsanity.com/paper/2102.06477/) (2021)
- [Partially Exchangeable Networks and Architectures for Learning Summary Statistics in Approximate Bayesian Computation](https://mlsanity.com/paper/1901.10230/) (2019)
- [Hierarchical Neural Simulation-Based Inference Over Event Ensembles](https://mlsanity.com/paper/2306.12584/) (2023)

## Built on by (most cited first, 1 of 1)

- [Hybrid Neural Simulation-Based Inference for Robust Applications and Limited-Budget Scenarios](https://mlsanity.com/paper/2609.36044/) (2026)

## Abstract (arXiv)

Neural posterior estimation has emerged as a powerful tool for amortized inference, with growing adoption across scientific and applied domains. In many of these applications, the conditioning variable is a set of observations whose elements depend not only on the target but also on unknown factors shared across the set. Optimal inference therefore requires treating the set jointly, which in turn requires training the estimator at the deployment set size -- a regime where memory and compute quickly become prohibitive. We introduce a simple, theoretically grounded strategy that decouples representation learning from posterior modeling. Our method trains a mean-pool Deep Set on sets of size at most two, producing an encoder that generalizes to arbitrary set sizes. The inference head is then finetuned on pre-aggregated embeddings, making training cost essentially independent of the deployment set size N. Across scalar, image, multi-view 3D, molecular, and high-dimensional conditional generation benchmarks with N in the thousands, our approach matches or outperforms standard baselines at a fraction of the compute.