PI3NN UQ Head

[SimonTopp]

This PR includes a prototyping notebook for a new UQ head approach. It's not exactly XAI, but I think it could be the first step to attributing model confidence, specifically in out-of-bounds prediction scenarios.

Prediction Interval from 3 Neural Networks (PI3NN) has some key advantages compared to other UQ methods:

1. It has no apriori assumption about the distribution of uncertainty
2. It's able to detect out-of-bounds predictions and accurately reflect the increase in uncertainty associated with OOB preds.
3. It uses a scaling approach for defining different confidence intervals, so it only needs to be trained once and different CI's can them be calculated post-hoc.

The original implementation is at the bottom the of notebook. It requires training three distinct NN's for the main prediction, upper, and lower bound respectively. Here I adapted this to a multitask modeling approach were we duct tape a UQ head to the main model. I think that this makes it much more modular across models, and it could theoretically be applied to spatially and temporally aware models in order for the UQ to take into account spatial and temporal awareness.

@jdiaz4302, I think this could be a cool UQ approach to add to your already extensive list of distribution based UQ heads. When you get a chance you want to let me know what you think?

[jdiaz4302]

Nice, I'll definitely look into this

Have you used this method with the RGCN on the DRB (with all of its nan observations)? Curious because I'm currently working on the UQ heads in that context (subbasin for the reservoir network forecasting)

[SimonTopp]

I haven't tried it, but I think it'd be relatively straight forward. Do you have a notebook or anything that you used to compare the other UQ heads? I'd be curious to try to plug it into the same pipeline so see both how it affects overall error during training and how it compares to the other heads in similar circumstances.

[jdiaz4302]

Yeah, I have some notebooks that I'll revisit and send along either today or tomorrow. Like their paper talks about, this method has less overhead. By comparison, the UQ heads will add (1) the new model code, (2) separate likelihood loss functions, and (3) sampling/analytical utilities for deriving confidence intervals from the distribution parameters.

[SimonTopp]

For sure, as a UQ head this also needs its own loss function, but it's really just a masked mse function. Plus it needs the root finding code that scales the CI outputs, but similarly that's very straightforward and adds almost no computational overhead, certainly way less than sampling the distributions for the other UQ heads.

[jdiaz4302]

I actually did some work on getting point estimates and CIs without sampling, and they were very promising with speed, memory, and replication of sampling results, but that work isn't well organized/polished/implemented in a pipeline because it isn't compatible with DA since DA adjusts the samples to better match reality rather than a distribution we have the known/estimated parameters for.

When/if other projects want to use the UQ heads and after our operational forecast deadlines, that may not still be the case

[SimonTopp]

Closing this cause I haven't gotten back to the PI3NN heads and am not sure when I will