Testing and Improving the Robustness of Amortized Bayesian Inference for Cognitive Models

TL;DR

This study test and improve the robustness of parameter estimation using amortized Bayesian inference (ABI) with neural networks and proposes a data augmentation or noise injection approach that incorporates a contamination distribution into the data-generating process during training.

Abstract

Contaminant observations and outliers often cause problems when estimating the parameters of cognitive models, which are statistical models representing cognitive processes. In this study, we test and improve the robustness of parameter estimation using amortized Bayesian inference (ABI) with neural networks. To this end, we conduct systematic analyses on a toy example and analyze both synthetic and real data using a popular cognitive model, the Drift Diffusion Models (DDM). First, we study the sensitivity of ABI to contaminants with tools from robust statistics: the empirical influence function and the breakdown point. Next, we propose a data augmentation or noise injection approach that incorporates a contamination distribution into the data-generating process during training. We examine several candidate distributions and evaluate their performance and cost in terms of accuracy and efficiency loss relative to a standard estimator. Introducing contaminants from a Cauchy distribution during training considerably increases the robustness of the neural density estimator as measured by bounded influence functions and a much higher breakdown point. Overall, the proposed method is straightforward and practical to implement and has a broad applicability in fields where outlier detection or removal is challenging.

Figures (20)

• Figure 1: A basic amortized Bayesian workflow. Parameters and data are simulated from a prior and an observation model. The simulations are used as training data for the summary and inference networks that jointly learn the posterior. Once the networks are trained, they can instantly sample from the approximate posterior given any new data.
• Figure 2: A graphical illustration of the drift diffusion process and the resulting reaction time data in a hypothetical visual recognition memory task. Participants view an image and judge whether it is old (i.e., previously seen) or new (i.e., previously unseen). The rows in the resulting data table correspond to individual trials of the experiment, with conditions and responses coded as 0 (old image) and 1 (new image).
• Figure 3: Predicted sufficient summary statistics $\hat{\boldsymbol{s}}_{EZ}$ from $\boldsymbol{s}_{B}$. The x-axis are the observed values of $M_{RT}$, $V_{RT}$, and $P_c$, and the y-axis is the $\hat{M}_{RT}$, $\hat{V}_{RT}$ and $\hat{P}_c$ predicted by $\boldsymbol{s}_{B}$
• Figure 4: The SSC and BP for the $\mu$ estimator when $n \in \{10, 20, 100\}$. The left panel displays the SSC for estimating $\mu$. For reasons of comparison, both the theoretical and the ABI-based SSC are shown. The shaded areas around the ABI-based SSC lines indicate the 2.5% to 97.5% percentiles of the influence. The right panel shows the average estimate of $\mu$ with a fraction of $p$ outliers (i.e., copies of $x^c=-100$) inserted into the sample, that is, $\overline{\hat{\mu}}(\boldsymbol{x}^{c},p)$.
• Figure 5: SSC and BP of amortized DDM estimation. Row (a) shows the SSC when $x^c$ ranges from 0.01s to 25s. Row (b) zooms in on the SSC when $x^c = 0.01 \sim 1s$, with a higher resolution. Row (c) shows the BP of the estimator when $x^c=0.01s$, while row (d) shows that with large outliers, namely $x^c=20s$. The true posterior mean across 1000 data sets is the prior mean, which is indicated with the red dashed line. The grey area marks the prior density for this parameter. The deviation between the estimates of perturbed data sets (in the blue line) and the red dashed line represents the systematic bias introduced by a certain fraction of outliers.