Technical Summary: Amortized Mean-Shift Interacting Particles

1. Problem Statement

In Bayesian inverse problems, the primary deliverable is not a set of posterior samples for their own sake, but the evaluation of integrals (posterior expectations, tail probabilities, risks) over a stream of observations. The standard approach uses Monte Carlo (MC) averaging of $M$ independent posterior samples. This estimator suffers from a variance that decays only as $O(1/\sqrt{M})$, requiring a prohibitive number of samples when each sample evaluation involves expensive forward models (e.g., partial differential equations).

While "mean-shift interacting particles" (Belhadji et al., 2026a,b) offer a superior alternative by replacing independent samples with a small set of weighted nodes (a quadrature) that minimizes the Maximum Mean Discrepancy (MMD) to the posterior, they suffer from two critical limitations:

1. Per-observation Optimization: Finding the optimal nodes requires solving a non-convex optimization problem from scratch for every new observation.
2. Score/Density Dependency: The most accurate versions of this method require evaluating the posterior score ($\nabla \log p(x|y)$) or density at every step, which is often unavailable in simulator-based settings or requires expensive adjoint solves.

The goal is to achieve the integration accuracy of designed quadrature nodes without the per-observation optimization cost or the need for posterior scores.

2. Methodology: Amortized Mean-Shift Interacting Particles

The authors propose Amortized Mean-Shift Interacting Particles, a learned, set-equivariant map that emits a weighted quadrature in a single forward pass.

Core Mechanism

The method takes an observation $y^\star$ and a "seed" set of $M$ independent samples $\{z^0_j\} \sim p(x|y^\star)$ as input. Instead of learning the nodes directly, the network learns a per-node displacement $\Delta z_j$ such that the final nodes are $z_j = z^0_j + \Delta z_j$.

• Architecture: A set-equivariant network (utilizing self-attention among particles and cross-attention to an observation summary) processes the seed set. It is conditioned on the requested node count $M$ via feature-wise modulation (FiLM), allowing a single trained map to serve any resolution.
• Closed-Form Weights: The network does not output weights. Instead, once the nodes are displaced, the weights are solved in closed form to minimize the MMD to the reference posterior, subject to a unit-sum constraint. This weight solve is a single linear system ($K^{-1}(\mu + \kappa \mathbf{1})$) and is differentiable, allowing end-to-end training.

Training Objective

The network is trained to minimize the expected squared MMD between the emitted quadrature and the reference posterior $\rho(\cdot|y^\star)$:
$$\mathcal{L}(\theta) = \mathbb{E}_{y^\star} \mathbb{E}_{z^0 \sim \rho} \left[ \text{MMD}^2_k \left( \sum \hat{w}_j \delta_{\hat{z}_j}, \rho(\cdot|y^\star) \right) \right]$$
Crucially, the training only requires samples from the posterior. It does not require the posterior density or score. The reference can be a closed-form distribution, a trained conditional flow, a physics-based posterior, or an empirical Nadaraya-Watson conditional measure.

Handling High Dimensions

In high ambient dimensions, isotropic kernels suffer from the "kernel wall," where pairwise distances concentrate, causing the Gram matrix to lose geometric contrast. The method addresses this by:

1. Posterior-Whitened Kernel: Using a Mahalanobis kernel based on the posterior precision (or sample covariance) to normalize the geometry.
2. Informed Subspace: Projecting the problem into a low-dimensional "likelihood-informed" subspace. The authors show that this subspace can be estimated score-free from joint samples (using the covariance of conditional means), recovering the same directions as score-based methods in linear-Gaussian settings.

Optional Refinement

For cases where a score is available and the reference is well-conditioned, the one-pass emission can serve as a "warm start" for a few steps of the original mean-shift iteration, further reducing error. However, the authors note this refinement is fragile on sharp (low-rank) references where the iteration may over-contract.

3. Key Contributions

1. Score-Free Quadrature Below the MC Floor: The method produces a weighted node set that integrates the posterior with lower worst-case error than an equal number of independent samples, without ever evaluating the posterior density or score.
2. Amortization: A single trained network emits the quadrature for any observation and any node budget $M$ in one forward pass, eliminating the need for per-observation optimization.
3. Pareto Improvement: The construction is not a competitor to drawing more samples but a Pareto improvement on the MC estimator itself.
   • Reweighting: Provably no worse than equal weights (Monte Carlo floor) at any budget (Proposition 2).
   • Moving Nodes: Empirically lowers the integration error further by relocating nodes to high-mass regions missed by random seeds (Section 4.4).
4. Generalization: The method generalizes across posterior types (Gaussian, mixtures, learned flows, physics-based) and dimensions (up to 1,024 coefficients in groundwater fields) when combined with the posterior-whitened kernel and informed subspace.

4. Results

The paper validates the method across several benchmarks:

• Closed-Form & Sampled Toys: On Gaussian, Gaussian-mixture, and curved "banana" distributions, the one-pass warm start consistently sits below the independent-sample floor across all node budgets. The "move" lever provides the significant gain, particularly where random samples miss structure (e.g., separated modes or curved ridges).
• Limited-Angle Tomography: On a linear-Gaussian inverse problem, the method integrates the exact posterior below the floor. The per-pixel uncertainty maps are calibrated to the true posterior within the informed subspace.
• Trained Conditional Flow: On a learned posterior where the shape varies with the observation, the amortized map tracks the changing geometry and beats the floor.
• Physics-Based (Groundwater): On a high-dimensional (1,024 coefficients) Darcy flow inverse problem, the method successfully integrates the Nadaraya-Watson conditional posterior in the informed subspace, beating the MC floor where the ambient kernel would otherwise collapse.
• Comparison to Samplers: When compared to tuned Stein Variational Gradient Descent (SVGD), the quadrature (even with just reweighting) achieves lower integration error at larger budgets, as SVGD plateaus at its fixed point while the quadrature continues to descend the discrepancy.

5. Significance and Claims

The paper claims that Amortized Mean-Shift Interacting Particles represent a fundamental shift in how Bayesian integrals are computed in inverse problems.

• Not a Rival to Sampling: The method does not compete with the fidelity of samplers; rather, it improves the estimator fed by those samples. It transforms a set of independent samples into a designed quadrature that is provably better (via reweighting) and empirically superior (via moving).
• Pareto Improvement: The construction offers a "Pareto improvement" on the Monte Carlo estimator: it achieves lower variance for the same computational budget (node count) without requiring the expensive per-observation solves or score evaluations that previously made such quadratures impractical.
• Scope Limitations: The authors are modest about the claims:
  • The guarantee is relative to the user-supplied posterior $\rho$, not the ground truth $\pi$. The fidelity of $\rho$ depends on the user's modeling.
  • The "move" gain is empirical, not theoretically certified for all budgets.
  • The optional score-based refinement is fragile on sharp, low-rank references.
  • The method integrates the posterior, not the truth; if the posterior model is poor, the quadrature will faithfully integrate that poor model.

In summary, the paper presents a learned, score-free mechanism to convert expensive posterior samples into a highly efficient, reusable quadrature rule that outperforms standard Monte Carlo integration across a wide range of inverse problem settings.