📊 statistics

It Just Takes Two: Scaling Amortized Inference to Large Sets

This paper introduces a scalable amortized inference method that decouples representation learning from posterior modeling by training a mean-pool Deep Set on sets of size at most two, enabling efficient inference on arbitrarily large sets with compute costs independent of the deployment set size while matching or outperforming standard baselines.

Original authors: Antoine Wehenkel, Michael Kagan, Lukas Heinrich, Chris Pollard

Published 2026-05-11

📖 6 min read🧠 Deep dive

CC BY 4.0

Original authors: Antoine Wehenkel, Michael Kagan, Lukas Heinrich, Chris Pollard

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 "Too Many Cooks" Dilemma

Imagine you are a detective trying to solve a mystery (finding a hidden truth, or $\theta$ ). You have a huge pile of clues (a set of observations).

In many real-world scenarios, these clues aren't independent. They are all affected by a hidden, shared factor (a nuisance variable, or $\psi$ ).

Example: Imagine trying to guess the weight of a specific apple ( $\theta$ ) by weighing it on 1,000 different scales. But all the scales are slightly broken in the exact same way because they were all calibrated by the same faulty technician ( $\psi$ ).

To get the right answer, you can't just look at one scale at a time. You have to look at all 1,000 scales together to figure out how the faulty calibration is skewing the results and correct for it.

The Catch:

Method A (The Lazy Detective): Look at each scale individually and average the results. This is fast and easy, but it fails because it ignores the fact that all the scales are broken in the same way. You get a wrong answer.
Method B (The Perfect Detective): Look at all 1,000 scales at once to solve the puzzle perfectly. This works, but it requires a supercomputer. If you try to train a neural network to do this with 1,000 clues at once, the computer runs out of memory and crashes. It's too expensive.

The Solution: PAIRS (Pretraining Aggregators for Inference at aRbitrary Set-sizes)

The authors introduce a clever trick called PAIRS. Their philosophy is simple: "Train small, deploy large."

They realized that to learn how to handle a group of clues that share a hidden flaw, you don't need to see the whole group at once. You only need to see two clues at a time.

Think of it like learning a language:

If you want to learn how a specific accent affects a sentence, you don't need to listen to a whole choir of 1,000 people singing at once.
You just need to listen to two people singing together. Once you hear how their voices blend and how the shared accent changes the sound, you understand the rule.
Once you know the rule, you can apply it to a choir of 1,000 people without needing to re-learn the rule.

How PAIRS Works (The Three-Step Recipe)

The paper proposes a three-stage process:

Stage 1: The "Two-Person" Training (Pretraining)
The AI is trained on tiny sets of data, containing only 1 or 2 observations at a time. It learns to recognize the pattern of the shared hidden flaw (the nuisance) by comparing pairs. It learns a "summary rule" (an encoder) that can compress any single clue into a useful piece of information.
- Analogy: The detective studies pairs of scales to figure out exactly how the faulty technician messed up the readings.
Stage 2: The "Freezing" Step
Once the AI has learned the rule from the pairs, the "brain" that learned the rule (the encoder) is frozen. It is locked in place. It will never change again.
Stage 3: The "Big Group" Finetuning
Now, the AI is given the huge sets of data (1,000 clues). Because the "brain" is frozen, the computer doesn't have to do the heavy lifting of processing 1,000 items simultaneously. It just uses the frozen brain to quickly summarize each clue into a small note, adds all the notes together (mean pooling), and then trains a simple "inference head" to read the final summary.
- Analogy: The detective now looks at the 1,000 scales. Instead of analyzing all 1,000 at once, they use their pre-learned rule to quickly jot down a note for each scale, sum up the notes, and then make a final decision. This is fast and doesn't crash the computer.

Why "Two" is the Magic Number

The paper proves mathematically that you never need more than two to learn the rule.

If you train on 1 item, you only see the individual clue, not the shared flaw.
If you train on 2 items, you see the shared flaw in action.
If you train on 3, 4, or 1,000 items, you aren't learning anything new about the rule that you didn't already learn from the pair. The extra items just add more of the same information.

Therefore, training on sets of size 2 is just as good as training on sets of size 1,000, but it costs a fraction of the computing power.

What They Tested

The authors tested this on several real-world scenarios where the "shared flaw" makes things hard:

Particle Physics: Finding a signal amidst background noise where the noise is shared across events.
Images: Identifying the size of a circle in a picture when the picture is rotated (the rotation is the shared flaw).
3D Objects: Guessing the volume of a 3D object from multiple 2D photos taken from different angles.
Molecules: Predicting chemical properties from different 3D shapes of the same molecule.
Image Generation: Creating a new view of a 3D scene based on a few existing photos.

The Results

Performance: PAIRS performed just as well as (or better than) the expensive methods that try to train on huge sets all at once.
Cost: It was dramatically cheaper. In some cases, the "expensive" method required 100 times more computing power to get the same result.
Scalability: While other methods crashed or became impossible to train when the set size got huge (thousands of items), PAIRS handled them easily because it only had to train on pairs.

Summary

The paper says: Don't try to eat the whole elephant at once. If you want to understand how a group of things work together, just look at two of them. Once you understand the relationship between a pair, you can apply that knowledge to the whole herd without needing a bigger brain or a bigger computer.

Technical Summary: It Just Takes Two: Scaling Amortized Inference to Large Sets

Problem Statement

Neural Posterior Estimation (NPE) has become a standard tool for amortized probabilistic inference, particularly in scientific domains where one must infer a target parameter $\theta$ from a set of observations $X = \{x_i\}_{i=1}^N$ . A critical challenge arises when the observations in the set are coupled by shared, unknown nuisance factors $\psi$ (e.g., common detector calibration, scene illumination, or molecular identity) that are not of direct interest but impact the entire set collectively.

In such scenarios, optimal inference requires processing the set jointly to account for these shared factors. However, existing approaches face a scalability trade-off:

Per-observation strategies: Train on individual measurements and recombine scores at test time. While scalable in $N$ , this fails when shared nuisance factors couple the observations, leading to incorrect inference.
Joint set strategies: Use permutation-invariant architectures (e.g., Deep Sets) trained end-to-end on sets of size $N$ . While principled, the memory and compute costs of training scale linearly with the target deployment cardinality $N$ . This makes training on large sets (e.g., $N$ in the thousands) computationally prohibitive.

Practitioners are thus forced to choose between a scalable but sub-optimal estimator or a correct but unscalable one.

Methodology: PAIRS

The authors introduce PAIRS (Pretraining Aggregators for Inference at aRbitrary Set-sizes), a three-stage procedure that decouples representation learning from posterior modeling to achieve scalability without sacrificing correctness.

Theoretical Insight

The core theoretical contribution is the proof that learning a permutation-invariant data representation (the encoder) does not require training on large sets. Under the assumption that the data admits a finite-dimensional mean-pool sufficient statistic, the authors prove that training on sets of size $N \le 2$ is sufficient to recover the same encoder (up to affine reparameterization) as training on any larger cardinality. This relies on the fact that pairs are the smallest set size that exposes the coupling induced by shared nuisance factors, while larger sets provide no new information about the underlying generative process beyond what pairs reveal.

The PAIRS Procedure

The method consists of three stages:

Pretraining (Small Cardinality): A Deep Set encoder (mean-pool) and an inference head (conditional generative model) are jointly trained on observation sets of size $N \in \{1, 2\}$ . This stage learns the representation $t_\omega(x)$ and the initial posterior mapping.
Embedding Caching: The trained encoder is frozen. For the finetuning dataset (which may contain sets of arbitrary large size $N$ ), individual observations are passed through the encoder, and the resulting embeddings are pre-aggregated via mean-pooling.
Finetuning (Large Cardinality): The inference head is finetuned on the pre-aggregated embeddings. Since the encoder is frozen and embeddings are pre-computed, the memory and compute cost of this stage is independent of the deployment set size $N$ .

Key Contributions

Theoretical Guarantee (Theorem 1): The authors prove that under mild regularity conditions, a mean-pool Deep Set encoder trained on sets of size at most two recovers a sufficient statistic for $\theta$ at any cardinality $N \ge 1$ . This implies that the representation learning cost is decoupled from the deployment set size.
The PAIRS Algorithm: A practical, three-stage pipeline that allows training on small sets while enabling inference on sets of thousands of elements. The training cost is essentially independent of the target $N$ .
Empirical Validation: Extensive experiments across scalar, image, multi-view 3D, molecular, and high-dimensional conditional generation benchmarks with $N$ ranging up to thousands.

Experimental Results

The paper evaluates PAIRS against several baselines, including per-observation marginalization ( $N=1$ ), extended pretraining ( $N=1\text{--}10$ ), Gaussian regression, and end-to-end training at large $N$ (where feasible).

Performance: PAIRS matches or outperforms standard baselines across all tasks. Notably, it tracks the analytical posterior of a bivariate Gaussian model up to $N=10^5$ and matches MCMC estimates on a "bump hunt" particle physics task with shared nuisance factors.
Scalability: In tasks where end-to-end training at $N=1000$ is computationally feasible, PAIRS achieves comparable performance at a fraction of the compute cost (approx. 100x cheaper). In high-dimensional generative tasks (novel-view synthesis), end-to-end training at large $N$ is infeasible, whereas PAIRS scales successfully.
Embedding Dimension: Experiments show that performance saturates once the embedding dimension $l$ is sufficiently large to capture the problem's complexity, confirming that $l$ does not need to grow with $N$ .
Cost-Performance Trade-off: PAIRS lies on the Pareto frontier, offering the best balance between posterior quality and computational cost. Strategies pretraining on $N=1\text{--}10$ were found to be significantly more expensive and often performed worse than PAIRS.

Significance and Claims

The paper claims to resolve the scalability vs. correctness trade-off in amortized inference for set-valued data. By demonstrating that "training small, deploying large" is theoretically grounded for mean-pool architectures, PAIRS enables the application of neural posterior estimation to problems with large observation sets that were previously intractable.

The authors emphasize that their approach is not merely a heuristic but is supported by a proof showing that pairs are sufficient to learn the necessary additive structure of the sufficient statistic. They note that while the theory assumes an exponential-family marginal likelihood (a restriction inherent to mean-pool Deep Sets), empirical results suggest this idealization is mild enough to generalize well to realistic, non-EF problems. The work opens avenues for applying amortized inference to large-scale scientific datasets without the prohibitive compute costs of end-to-end training at deployment cardinalities.