用两个样本训练,就能高效推断上千样本的后验分布。
It Just Takes Two: Scaling Amortized Inference to Large Sets

- 先用最多两样本训练深度集合模型,学习通用表示
- 部署时处理千级样本集,性能媲美甚至超越传统方法
- 适合需要快速推理大规模集合数据的科研与工程场景
神经后验估计已成为一种强大的通用推理工具,广泛应用于科学和实际领域。在许多应用中,条件变量是一组观测值,其元素不仅依赖于目标变量,还共享未知的全局因素。最优推理需联合处理整个集合,但传统方法要求在部署规模下训练,导致内存和计算成本急剧上升。本文提出一种简单且理论完备的策略:将表征学习与后验建模解耦。方法在最大大小为2的集合上训练均值池化深度集合模型,得到可泛化至任意集合规模的编码器;随后在预聚合嵌入上微调推理头,使训练开销几乎与部署规模N无关。在标量、图像、多视角3D、分子及高维条件生成等基准测试中,当N达数千时,本方法以极低算力消耗达到或超过标准基线性能。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。