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

**Authors:** Antoine Wehenkel, Michael Kagan, Lukas Heinrich, Chris Pollard  
**arXiv:** [2605.07972](https://arxiv.org/abs/2605.07972)  
**Date:** 2026-05-08  
**Summary page:** https://www.summarizepaper.com/en/paper/2605.07972v1/

## Key points

- Authors introduce PAIRS strategy for scaling neural posterior estimation to large sets
- Traditional amortized inference involves joint conditioning on a set of observations with shared factors
- PAIRS method decouples representation learning from posterior modeling by training mean-pool Deep Set on sets of size at most two
- Encoder can generalize to arbitrary set sizes, enabling efficient inference
- Inference head fine-tuned on pre-aggregated embeddings, reducing training costs independent of deployment set size N
- Experiments validate effectiveness of PAIRS in diverse real-world scenarios with superior performance over baseline methods
- Demonstrated improved predictive performance in particle physics-inspired toy problem and low-dimensional parameter estimation tasks across different observation modalities
- Applied to novel-view synthesis and high-dimensional conditional generative modeling tasks, outperforming standard baselines while requiring less computational resources

## Summary

In the paper titled "It Just Takes Two: Scaling Amortized Inference to Large Sets," authors Antoine Wehenkel, Michael Kagan, Lukas Heinrich, and Chris Pollard introduce a novel strategy called PAIRS for scaling neural posterior estimation to large sets. The traditional approach to amortized inference involves conditioning on a set of observations where elements depend not only on the target but also on unknown shared factors across the set. This necessitates treating the set jointly for optimal inference, which typically requires training the estimator at the deployment set size - a scenario that quickly becomes computationally prohibitive due to memory and compute constraints. The PAIRS method decouples representation learning from posterior modeling by training a mean-pool Deep Set on sets of size at most two. This produces an encoder that can generalize to arbitrary set sizes, allowing for efficient inference. The inference head is then fine-tuned on pre-aggregated embeddings, making training costs essentially independent of the deployment set size N.

The authors conducted experiments in diverse scenarios to validate the effectiveness of PAIRS in real-world applications where traditional assumptions may not hold. They first present a toy problem inspired by particle physics, demonstrating how joint processing of observations can lead to improved predictive performance. The proposed strategy is then evaluated on four low-dimensional parameter-estimation tasks across different observation modalities, showcasing its superiority over baseline methods. Furthermore, PAIRS is applied to novel-view synthesis from a set of images and high-dimensional conditional generative modeling tasks where end-to-end training at the deployment cardinality is computationally challenging. Despite these complex scenarios, PAIRS consistently outperforms standard baselines while requiring significantly less computational resources. Overall,"It Just Takes Two" presents a groundbreaking approach that addresses the challenges of scaling amortized inference to large sets by decoupling representation learning and posterior modeling, offering an efficient and effective solution for various scientific and applied domains.
