JANA: Jointly Amortized Neural Approximation of Complex Bayesian Models

Neural networks trained on model simulations enable amortized inference: A pre-trained network can be stored and re-used for Bayesian inference on millions of data sets (von This work proposes "jointly amortized neural Krause et al., 2022). Crucially, most previous neural approaches approximation" (JANA) of intractable likelihood have tackled either SM or SBI in isolation, but little functions and posterior densities arising in attention has been paid to learning both tasks simultaneously. Bayesian surrogate modeling and simulation-based To address this gap, we propose JANA ("Jointly Amortized inference. We train three complementary networks Neural Approximation"), a Bayesian neural framework for in an end-to-end fashion: 1) a summary network simultaneously amortized SM and SBI, and show how it enables to compress individual data points, sets, or time novel solutions to challenging downstream tasks like series into informative embedding vectors; 2) a posterior the estimation of marginal and posterior predictive distributions network to learn an amortized approximate (see Figure 1). JANA also presents a major qualitative posterior; and 3) a likelihood network to learn an upgrade to the BayesFlow framework (Radev et al., 2020), amortized approximate likelihood. Their interaction which was originally designed for amortized SBI alone.