Minimum Distance Summaries for Robust Neural Posterior Estimation

2026-02-09

AI Summary

This work addresses the fragility of neural posterior estimation under test-time distribution shifts, where model misspecification can severely degrade performance. To enhance robustness without retraining, the authors propose a plug-and-play test-time summary adaptation method that optimizes summary statistics solely at inference by minimizing the Maximum Mean Discrepancy (MMD) between observed data and the summary-conditioned predictive distribution. This approach decouples robustness from model training, preserving the reusability and modularity of pretrained inference networks. Leveraging random Fourier features enables efficient, model-free optimization with minimal computational overhead. Empirical evaluations across multiple synthetic and real-world tasks demonstrate substantial improvements in inference robustness while maintaining negligible additional cost.

Technology Category

• Machine Learning: Adversarial Learning & Robustness
• Computer Vision: Adversarial Attacks & Robustness
• Natural Language Processing: Safety and Robustness

Application Category

• User Modeling, Personalization and Recommendation: Attacks and countermeasures in recommendation systems
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Systems and Infrastructure for Web, Mobile and WoT: Applied ML and AI for Web-based mobile applications

Abstract

Simulation-based inference (SBI) enables amortized Bayesian inference by first training a neural posterior estimator (NPE) on prior-simulator pairs, typically through low-dimensional summary statistics, which can then be cheaply reused for fast inference by querying it on new test observations. Because NPE is estimated under the training data distribution, it is susceptible to misspecification when observations deviate from the training distribution. Many robust SBI approaches address this by modifying NPE training or introducing error models, coupling robustness to the inference network and compromising amortization and modularity. We introduce minimum-distance summaries, a plug-in robust NPE method that adapts queried test-time summaries independently of the pretrained NPE. Leveraging the maximum mean discrepancy (MMD) as a distance between observed data and a summary-conditional predictive distribution, the adapted summary inherits strong robustness properties from the MMD. We demonstrate that the algorithm can be implemented efficiently with random Fourier feature approximations, yielding a lightweight, model-free test-time adaptation procedure. We provide theoretical guarantees for the robustness of our algorithm and empirically evaluate it on a range of synthetic and real-world tasks, demonstrating substantial robustness gains with minimal additional overhead.

Problem

• neural posterior estimation
• simulation-based inference
• robustness
• model misspecification
• amortized inference

Innovation

• minimum-distance summaries
• neural posterior estimation
• maximum mean discrepancy
• simulation-based inference
• test-time adaptation