Generalized Bayesian Inference for Scientific Simulators via Amortized Cost Estimation

Summary

The authors present a novel approach for SBI by utilizing neural networks (NN) to estimate a generalized cost in GBI. Their proposed method, ACE, demonstrates superior computational efficiency compared to previous approaches, without compromising competitive performance across various evaluation metrics. The authors conducted extensive benchmarking of their method under diverse experimental settings, ranging from 1D to 10D. Furthermore, they conducted experiments using real intracellular recordings, adding an additional dimension to their research.

Strengths

The paper is well written and maintains a smooth flow throughout. Furthermore, the contribution made by the authors appears to be a novel and valuable improvement upon existing ideas in the field. Moreover, the authors have thoughtfully acknowledged and addressed the limitations of their work, placing their findings within the broader context of the field. It is worth noting that their code repository is well-organized, adding to the overall credibility and reproducibility of their research.

Weaknesses

1. To enhance clarity, it would be beneficial to improve the visual presentation of Figure 2. When referring to subfigure C, the arrangement could be adjusted to prevent the automatic focus on the third pane. 2. Propositions 1 and 2 appear somewhat forced, as their results do not seem entirely compelling or noteworthy. It may be worth revisiting these propositions. In the appendix the authors state, "The proofs closely follow standard proofs that regression converges to the conditional expectation." Thus highlighting the lack of importance of these propositions. 3. A minor typo on line 116, "parameter- or".

Questions

nothing, great work!

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

4 excellent

Contribution

3 good

Limitations

The authors do address the limitations of their work.

Summary

This paper studies the problem of simulation-based inference - which can is encountered in a wide range of scientific problems - where one is interested in performing Bayesian inference using simulators with implicit likelihoods. Scientific problems can have two unique properties - a) the predictive quality of the simulation is more important than the posterior b) the model can be misspecified. These can be handled though Generalized Bayesian Inference - which generalizes traditional Bayesian Inference by admitting generalized likelihood functions which are defined by some general loss / cost function defined on the parameters. A challenge with employing Generalized Bayesian Inference in the simulation-based case is that estimating the cost function can be computationally expensive. The authors propose learning an amortized neural network estimator for the cost function to reduce the overhead for enabling GBI in simulation-based inference. This amortized cost estimator can be combined with standard inference algorithms (MCMC) to sample from the Generalized Bayesian posterior. The overall proposed approach consists of 3 stages: 1) Collecting data 2) Training ACE 3) Sampling using ACE. The authors present several experiments on a variety of benchmark simulation-based inference tasks and present a case study with the Hodgkin-Huxley simulator.

Strengths

* The problem of simulation-based inference is an important one for various scientific problems and misspecification is an often under-studied but important aspect in practice. The paper attempts to address this important issue relevant to various communities. * The proposed approach is novel within the context of the simulation based inference. (Although as I mention below, similar ideas have been employed is some related areas) * Even despite some caveats that I discuss below, the idea is quite neat, and relatively simple (which is a good thing!) Replacing the estimation of the cost with an amortized predictor seems like a natural extension to enable generalized posteriors in the simulation-based inference setting. The simpler algorithm also makes it easier to adopt by practitioners with little overhead. * The experiments cover a fairly wide variety of tasks used in prior work in the SBI setting. The case study on the HH model was a nice example of leveraging the method on realistic tasks. * The presentation in the paper was generally quite clear. The paper is well written and easy to follow. * I also appreciate the authors open-sourcing the code to aid reproducibility.

Weaknesses

* A fundamental weakness I see with the method is that it relies on the neural estimator to generalize well from finitely many simulations. This might not be true in practice (certainly isn’t guaranteed). Put differently, an implicit assumption for the approach to work is that the ACE is trained with _enough_ data to be able to provide informative guidance to the sampler learned in the next stages. As with other work relying on NN estimators this is a fundamental failure mode. * On a similar note, the method appears to be practical only for relatively low-dimensional problems. The NN learning issues would crop up if the dimensionality of the problem is increased. * Similar ideas of learning amortized neural estimators have been recently explored in the uncertainty estimation literature [1,2]. In particular, the secondary predictor learned in [1] amortizes a similar quantity. [1] DEUP: Direct Epistemic Uncertainty Prediction. Lahlou et al., TMLR 2023. [2] Epistemic Neural Networks. Osband et al., 2023.

Questions

* There are already some results with different simulation budgets - but a more careful study of the effect of the number of samples available to train the ACE might be quite useful. * Could you comment on the connections to the work on uncertainty estimation mentioned above? * Do the authors have any thoughts on extending the approach to higher dimensional settings?

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

3 good

Contribution

2 fair

Limitations

The authors already discuss several limitations in the paper including a) need for another method to sample from the posterior induced by ACE, b) accommodation of arbitrary distances, c) additional hyper parameters. In addition to these I would also add access to sufficient samples for training ACE and being limited to low dimensional settings as limitations.

Summary

This paper proposes amortized cost estimation (ACE) for generalized Bayesian inference (GBI) for SBI. The paper trains a neural network to approximate the cost function. The paper demonstrates results on baseline synthetic SBI examples followed by a real-world application using experimental data from the Allen Cell Types Database.

Strengths

* The paper is written clearly and is structured well. * The idea appears novel. While, GBI was previously defined in [13], its application within the SBI literature is interesting and is therefore of value to the community. * Algorithm box makes the implementation of the algorithm clear to the reader.

Weaknesses

* The main weakness of the paper seems to be the experimentation. While the figures are neatly presented, it is not that clear what the reader is supposed to conclude from the results: * For example, it was not clearly defined what the difference between specified and misspecified samples means within Figure 3. It would be helpful to mathematically define these differences rather than describe them within the text. * What is the definition of the GT-GBI? It first appears on line 173 and it is not clear what the true distance is and how it is obtained. This misunderstanding makes it further difficult to understand Figure 3. * Another question is why are the first three rows only showing MSE and the last row showing MMD? This is confusing and makes one question why the Figure is not consistent with the metric. * Additionally, for the MSE/MMD it seems that as $\beta$ is increased, the MSE/MMD goes down, but the C2ST seems to go up. Why is this the case? What happens in the limit that $\beta$ goes to infinity?

Questions

* Compared to BGI, equation (2) seems to not include the KL divergence term. Can this approach really be called GBI if it is not being regularized by some prior? Why was the KL term not included? * Following on from the question regarding the prior, is the addition of the noise ($\epsilon$) implicitly adding a prior? It might be the case that adding this noise is also enforcing some implicit distance metric between nearby parameters. I.e. does it assume that there is a local Euclidean distance metric? * Why does this approach require multi-chain slice-sampling? Is random-walk Metropolis-Hastings not sufficient?

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

* The idea behind the paper is interesting, but the experimental results are difficult to interpret. Further clarifications on the above would help increase the score.

Summary

The paper presents a new technique - amortized cost estimation (ACE) - that, as stated on the can, amortizes a broad class of loss functions used in generalized Bayesian inference (GBI) in place of the (log) likelihood. After training on a moderate-to-large number of model simulations (10K-100K in the examples of the paper), the amortized loss can then be used in place of the real loss to perform GBI (e.g., via MCMC). Crucially, this step does not require further simulations from the model. The authors show that the amortized loss generally matches the ground-truth loss, and show that their amortized method inherits the advantages of GBI, in that it is typically more robust to misspecification than standard and amortized Bayesian inference, as in neural likelihood estimation and neural posterior estimation

Strengths

- **Quality:** The general polish and quality of the paper is high. - **Clarity:** The paper is well written and generally clear. The aim and methodology are very well explained and motivated, with some points that would benefit from further expansion. - **Significance:** The idea of the paper (amortizing GBI for simulator-based inference, with the additional goal of dealing with misspecification) is very timely and there are surely many applications. - **Originality:** The idea at the core is not particularly original - amortizing the loss in GBI seems a natural direction for the field at the moment - but clearly worth pursuing. Similarly, the execution seems fairly straightforward (which is not a bad thing). ### Post-rebuttal: Thanks for having addressed most of my comments. I appreciated the additional experiments with tempered NLE, with varying noise augmentation, and the recommendations for the choice of $\beta$. The proposed heuristic seems both reasonable and practical. I understand that given the limited time, it was not feasible to conduct an additional set of experiments. Overall, I am satisfied with the changes and I increased my score accordingly from 6 to 7.

Weaknesses

Most of my concerns have to do with the treatment of the hyperparameter $\beta$ (inverse temperature), a key element of GBI. The role of $\beta$ could be explained and discussed more (it is mentioned in the **Limitations**). Many experiments show the results for a range of $\beta$ (which also changes from experiment to experiment), and it is unclear how the practitioner should choose its value. For example, $\beta \rightarrow \infty$, (generalized) Bayesian inference tends to maximum-a-posteriori (or minimum-generalized-loss) point estimation. The paper is ambiguous on what the user should be trying to achieve. I appreciate that some of these issues are not of ACE (the authors' method) but of GBI, but the paper would benefit from directly addressing these points. The comparison between GBI and Bayesian inference is somewhat unfair in that Bayesian inference is not allowed an inverse temperature hyperparameter (i.e., likelihood tempering), but in principle it could be easily applied (at least for NLE). Is ACE+GBI truly better here, or is the advantage just given by the fact that ACE/GBI has an extra free hyperparameter? In some cases, standard Bayesian inference can do better with large $\beta$. We know this can be the case for example with neural network posteriors (possibly a case of prior misspecification). - With NLE, just run MCMC with a scaling factor $\beta$ on the amortized log likelihood. - Amortizing different temperatures with NPE would be more complex. Naively, one would need to retrain the network for different values of $\beta$. One could also subtract the log prior from the NPE posterior to get the log likelihood, and then rescale it, but that might lead to instabilities. Similarly to $\beta$ (but less importantly), the paper could explore more the role of the noise $\sigma$ added to $S$ observations. This parameter seems to be less relevant though - it is very nice that the authors use a fixed value throughout the paper, and this value can be determined from prior-predictive checks. Nonetheless, a bit of exploration / insight could help. For the rest, the empirical evaluation is acceptable, but it would have been nice to see at least a couple of applications to real-world data (even another simple one), since this is arguably (also according to the authors) where GBI shines.

Questions

- Can you expand on the role of $\beta$, e.g. how you chose (and a practitioner should choose) reasonable ranges of the hyperparameter? - What about having an inverse temperature to the log likelihood of standard Bayesian inference (as per standard tempering, also seen helping in Bayesian deep learning applications)? Can you show that in your examples (for NLE, and possibly NPE)? - Consider adding a lesion study or analysis about the role of $\sigma$ and $S$. - Not strictly necessary, but it would be nice to see the method at work on another (simple) example with real data.

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

The authors address the limitations of their method (and I already asked above to expand on the role of $\beta$). The work has no particular potential negative social impact (at least not above any other general method for statistical inference).

Response to the rebuttal

Thanks for having addressed most of my comments. I appreciated the additional experiments with tempered NLE, with varying noise augmentation, and the recommendations for the choice of $\beta$. The proposed heuristic seems both reasonable and practical. I understand that given the limited time, it was not feasible to conduct an additional set of experiments. Overall, I am satisfied with the changes and I will argue for acceptance of the paper. I will increase my score accordingly.

Response

Thanks for the detailed response. Please do include those clarifications in the main paper/supplementary materials. I will happily increase my score.

Author-reviewer discussion

Dear all, The author-reviewer discussion period has now started. It will continue for one more week, until August 21. @authors: Please respond to the comments or questions reviewers may further have. Remain short and to the point. @reviewers: Please read the author's responses and ask any further questions you may have. To facilitate the decision by the end of the process, please also acknowledge that you have read the responses and indicate whether you want to update your evaluation. - You can update your evaluation positively (if you are satisfied with the responses) or negatively (if you are not satisfied with the responses or share other reviewers' concerns). Please note that major changes are a reason for rejection. - You can also keep your evaluation unchanged. In this case, please indicate that you have read the responses and that you do not have any further comments. Best regards, The AC

Response to author rebuttal

Thank you for the response and sorry for the late response. > Relations to work on uncertainty estimation Thanks for the elaborate explanation. To clarify, my comment on the uncertainty estimation methods was merely to highlight the connection between the two, not necessarily as a shortcoming. Apologies for the miscommunication. I think a concise version of this summary would be good to have in the paper. I appreciate the authors responses which partially address the concerns I raised about the reliability of NNs trained on finite data and high-dimensional problems. But these fundamental limitations still remain so I will keep my score, recommending acceptance.