It Just Takes Two: Scaling Amortized Inference to Large Sets

TL;DR

This paper presents a scalable method for neural posterior estimation that efficiently handles large sets by decoupling representation learning from posterior modeling, enabling generalization to arbitrary set sizes.

Contribution

The authors introduce a simple, theoretically grounded approach that trains on small sets and generalizes to larger sets, reducing computational costs significantly.

Findings

Method matches or outperforms baselines on various benchmarks.

Training cost is independent of deployment set size.

Effective across diverse data modalities and large set sizes.

Abstract

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.