In simulation-based inference, you train a neural network with pairs of model parameters θ and correspondingly data D, (θ, D). In the case of NRE, the network is a classifier, which predicts the probability that the data came from the (θ, D) stream corresponding to the joint p(θ,D) vs. a scrambled-up (θ, D) stream, corresponding to the respective priors p(θ)*p(D). With some algebra, this gives you p(D|θ), i.e., the likelihood. In the case of NPE, the network predicts p(θ|D) using a normalising flow, by optimizing the probability function to output most probability weight where θ is for the given D.
The stream of (θ, D) is created by defining a prior p(θ), and a generator function D|θ which simulates data. Easy. Standard “forward-folding” (astronomer-speak) or generative modelling (statistician-speak).
This scheme works well in a ML setting, where you have many tuples (θ, D), ideally, billions. Yes, you need a lot of data. This training cost hinders iterative model building.
You can reuse standard industry network architectures to handle image data easily.
Limitation 1: Varying observing conditions
The scheme does not fit if you have variations in observing conditions. For example, if your instrument behaves slightly differently for each observation.
Varying observing conditions are pervasive in X-ray astronomy because the object position on the focal plane changes the instrument response, both in the spectral and imaging sense. In optical astronomy, the equivalent is that you have different seeing conditions (filters are not considered time-variable). So we need to deal with this.
[Aside on backgrounds: The background can also vary with position and time, which matters unless you are studying X-ray binaries and feel brave. The background can be simulated as well, with extending both θ and D, to background parameters and background region data, so this is not a deal-breaker.]
The standard method then would be to either sample a random observing condition, and generate (θ, D), which would marginalise over the observing condition. Not wrong, but would lose information. The standard optical method is to derive the seeing conditions in the data analysis step from point sources in the image (e.g., SourceExtractor, photutils). This is reasonably fast. For X-ray astronomy, this is not possible, because the response is not location-invariant, so you cannot transplant information from nearby to the source.
The proper way would be to sample not just p(θ, D), but p(θ, D | C) where C are the observing conditions, and define a conditional SBI (NRE or NPE) that gives its results conditional on an input C. If C was just the exposure time, and PSF FWHM, then we could make it part of D and sample it, and at inference time condition on it. However, in X-ray astronomy, C is a response matrix, and it is not trivial to smoothly sample from it, or to inform a SBI neural network about how to use it, because the SBI does not forward-fold. I suppose one could project C into a low-dimensional space.
In any case, the samplings of C would likely be incomplete, and a new observation may not lie inside, erasing much of the cost savings at inference time that SBI should bring.
Limitation 2: Poisson count data
Why are there so few SBI – and, more generally – neural network papers with X-ray data? I think it is because neural networks suck with count data.
You would think that you take a standard network that works with continuous data trained and knowing the L2 loss is proportional to a negative log-likelihood of a Gaussian switch for the loss for a negative log-likelihood of a Poisson and your done. But try to apply it to count data, and no, it doesn’t work. Why not? That’s a interesting ML research question. Correctly specified loss does not mean a trainable loss.
So X-ray ML astronomers are hacking around by making X-ray count data Gaussian-like by smoothing, losing information in the process, or testing with trivial high-count situations.
But once you bin finely enough in time or spectral channels (think XRISM, X-IFU), you are in the low count regime at least in those bins and have to deal with it.
In my personal opinion, SBI is worth pursuing, but there are alternative approaches as well that scale to high dimensions that are not losing information. In any case, neural networks play a role, but it matters how you use them. Yes, I am being coy here on what is better, I’d like to be funded to resolve these issues.
Leave a Reply