BI-EqNO: Generalized Approximate Bayesian Inference with an Equivariant Neural Operator Framework

TL;DR

BI-EqNO introduces an equivariant neural operator framework that enhances approximate Bayesian inference, improving flexibility, accuracy, and computational efficiency for both deterministic and stochastic models.

Contribution

The paper presents BI-EqNO, a novel neural operator framework that generalizes Bayesian inference with permutation equivariance and invariance, applicable to diverse prior and posterior representations.

Findings

Outperforms traditional Gaussian processes in regression tasks.

EnNF surpasses ensemble Kalman filter in small-ensemble scenarios.

Potential to serve as a 'super' ensemble filter for improved data assimilation.

Abstract

Bayesian inference offers a robust framework for updating prior beliefs based on new data using Bayes' theorem, but exact inference is often computationally infeasible, necessitating approximate methods. Though widely used, these methods struggle to estimate marginal likelihoods accurately, particularly due to the rigid functional structures of deterministic models like Gaussian processes and the limitations of small sample sizes in stochastic models like the ensemble Kalman method. In this work, we introduce BI-EqNO, an equivariant neural operator framework for generalized approximate Bayesian inference, designed to enhance both deterministic and stochastic approaches. BI-EqNO transforms priors into posteriors conditioned on observation data through data-driven training. The framework is flexible, supporting diverse prior and posterior representations with arbitrary discretizations and varying numbers of observations. Crucially, BI-EqNO's architecture ensures (1) permutation equivariance between prior and posterior representations, and (2) permutation invariance with respect to observational data. We demonstrate BI-EqNO's utility through two examples: (1) as a generalized Gaussian process (gGP) for regression, and (2) as an ensemble neural filter (EnNF) for sequential data assimilation. Results show that gGP outperforms traditional Gaussian processes by offering a more flexible representation of covariance functions. Additionally, EnNF not only outperforms the ensemble Kalman filter in small-ensemble settings but also has the potential to function as a "super" ensemble filter, capable of representing and integrating multiple ensemble filters for enhanced assimilation performance. This study highlights BI-EqNO's versatility and effectiveness, improving Bayesian inference through data-driven training while reducing computational costs across various applications.