It Just Takes Two: Scaling Amortized Inference to Large Sets

• URL: http://arxiv.org/abs/2605.07972v1
• Date: Fri, 08 May 2026 16:32:34 GMT
• Title: It Just Takes Two: Scaling Amortized Inference to Large Sets
• Abstract summary: We introduce a strategy that decouples representation learning from posterior modeling.<n>Our method trains a mean-pool Deep Set on sets of size at most two, producing an encoder that generalizes to arbitrary set sizes.<n>Our approach matches or outperforms standard baselines at a fraction of the compute.
• Score: 1.2879523047871226
• License: http://creativecommons.org/licenses/by/4.0/
• 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.

Related papers

• A Layer Importance Metric for Quantization Accounting for the Speed-Quality Trade-off in Autoregressive Models [0.0]
  Small language models (sLLMs) are nowadays hosted on devices with limited memory and computational budget.<n>We propose a composite metric that combines information retention (measured in terms of a normalized SQNR-based coefficient) and throughput gains.
  arXiv  Detail & Related papers  (2026-08-27T10:23:01Z)
• CRUMB: Efficient Prior Fitted Network Inference via Distributionally Matched Context Batching [1.4211459456330362]
  Prior-fitted networks (PFNs) are a promising class of foundation models that perform in-context learning.<n>We propose CRUMB (Clustered Retrieval Using Minimised-MMD ), a three-stage inference wrapper that clusters the test queries.<n>CRUMB selects a small, distributionally matched training subset for each cluster by greedily minimising the maximum mean discrepancy.
  arXiv  Detail & Related papers  (2026-06-09T22:07:04Z)
• Making Foundation Models Probabilistic via Singular Value Ensembles [56.4174499669573]
  Foundation models have become a dominant paradigm in machine learning, achieving remarkable performance across diverse tasks through large-scale pretraining.<n>Standard approach to quantifying uncertainty, training an ensemble of independent models, incurs prohibitive computational costs that scale linearly with ensemble size.<n>We propose Singular Value Ensemble (SVE), a parameter-efficient implicit ensemble method that builds on a simple, but powerful core assumption.<n>We show that SVE uncertainty quantification achieves comparable to explicit deep ensembles while increasing the parameter count of the base model by less than 1%.
  arXiv  Detail & Related papers  (2026-01-29T18:07:18Z)
• Tuning-Free Structured Sparse Recovery of Multiple Measurement Vectors using Implicit Regularization [13.378211527081582]
  We introduce a tuning-free framework to recover sparse signals in multiple measurement vectors.<n>We show that the optimization dynamics exhibit a "momentum-like" effect, causing the norms of rows in the true support to grow significantly faster than others.
  arXiv  Detail & Related papers  (2025-12-03T02:53:11Z)
• OPUS: Occupancy Prediction Using a Sparse Set [64.60854562502523]
  We present a framework to simultaneously predict occupied locations and classes using a set of learnable queries. OPUS incorporates a suite of non-trivial strategies to enhance model performance. Our lightest model achieves superior RayIoU on the Occ3D-nuScenes dataset at near 2x FPS, while our heaviest model surpasses previous best results by 6.1 RayIoU.
  arXiv  Detail & Related papers  (2024-09-14T07:44:22Z)
• LoRA-Ensemble: Efficient Uncertainty Modelling for Self-Attention Networks [52.46420522934253]
  We introduce LoRA-Ensemble, a parameter-efficient ensembling method for self-attention networks.<n>The method not only outperforms state-of-the-art implicit techniques like BatchEnsemble, but even matches or exceeds the accuracy of an Explicit Ensemble.
  arXiv  Detail & Related papers  (2024-05-23T11:10:32Z)
• Likelihood Ratio Confidence Sets for Sequential Decision Making [51.66638486226482]
  We revisit the likelihood-based inference principle and propose to use likelihood ratios to construct valid confidence sequences. Our method is especially suitable for problems with well-specified likelihoods. We show how to provably choose the best sequence of estimators and shed light on connections to online convex optimization.
  arXiv  Detail & Related papers  (2023-11-08T00:10:21Z)
• RGM: A Robust Generalizable Matching Model [49.60975442871967]
  We propose a deep model for sparse and dense matching, termed RGM (Robust Generalist Matching) To narrow the gap between synthetic training samples and real-world scenarios, we build a new, large-scale dataset with sparse correspondence ground truth. We are able to mix up various dense and sparse matching datasets, significantly improving the training diversity.
  arXiv  Detail & Related papers  (2023-10-18T07:30:08Z)
• SqueezeLLM: Dense-and-Sparse Quantization [80.32162537942138]
  Main bottleneck for generative inference with LLMs is memory bandwidth, rather than compute, for single batch inference. We introduce SqueezeLLM, a post-training quantization framework that enables lossless compression to ultra-low precisions of up to 3-bit. Our framework incorporates two novel ideas: (i) sensitivity-based non-uniform quantization, which searches for the optimal bit precision assignment based on second-order information; and (ii) the Dense-and-Sparse decomposition that stores outliers and sensitive weight values in an efficient sparse format.
  arXiv  Detail & Related papers  (2023-06-13T08:57:54Z)
• Scaling Structured Inference with Randomization [64.18063627155128]
  We propose a family of dynamic programming (RDP) randomized for scaling structured models to tens of thousands of latent states. Our method is widely applicable to classical DP-based inference. It is also compatible with automatic differentiation so can be integrated with neural networks seamlessly.
  arXiv  Detail & Related papers  (2021-12-07T11:26:41Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.