Uncertainty Quantification support in TorchGeo

[nilsleh]

Uncertainty Quantification is often deemed essential in prediction tasks that involve Earth Observation data (Gawlikowski et al. 2022). However, it seems that this is still sort of an "afterthought" as most papers and new methodologies focus on point estimates. To this end I think it would be great if TorchGeo would support some UQ approaches "out of the box" to make it easier for users to utilize them in their workflow, because maybe that is what is holding them back. My idea is to begin with the following two frameworks that have the advantage of being "postprocessing steps" to a MAP trained NN and support Classification and Regression Tasks:

1. Laplace Approximation: The core idea is to fit a Gaussian distribution around the trained MAP weights to obtain a measure of uncertainty. Weights can be sampled based on the approximated covariance matrix and Monte Carlo sampling during inference yields predictive uncertainty that is grounded in epistemic uncertainty. There is a nice Laplace library which supports several variants and which I would use (the current requirement is that the last layer is a nn.Linear Layer so support for classification and regression tasks should work)
2. Conformal Prediction: This is a method that recently has gained a lot of interest again with many papers coming out. It is a sort of recalibration method, with the important difference that it has proven marginal coverage guarantees for any desired error rate, for any model and any dataset, given that the assumptions of exchangeability and symmetry of data are fullfilled. Coverage then means that the true class is contained in the prediction set with the desired coverage (1 - error rate). For a nice introduction see (Angelopoulos and Bates 2022). For regression I would support Conformalized Quantile Regression (Romano et al. 2019) and for classification Adaptive Prediction Sets (Romano et al. 2020), as illustrated in the intro linked above.

Although I haven't tried it in details, my idea would be to adapt the Trainer classes for the respective task to make these methods available as an option after someone has called trainer.fit() and then make the uncertainty available to the trainer.predict() method. Additionally it would be nice to support some visualization and metric tools that focus on uncertainty (proper scoring roles, coverage, sharpness, and calibration metrics).

Looking forward to hearing what you think and any suggestions you might have.

[adamjstewart]

Only comment is that this doesn't feel like it's specific to geospatial data and may belong in a different library. That doesn't mean we can't implement it in something like torchvision and call it in our trainers though.

[isaaccorley]

I'd like to see a notebook showing how to combine a classification or segmentation task with this. There are some decent libraries for uncertainty estimation for classification.

I agree with @adamjstewart though. This seems like we could already combine this with other libraries without having to take on maintenance costs.

[nilsleh]

I would not implement either method from scratch but use existing libraries instead. But maybe instead of being part of torchgeo source, I could just write some tutorial notebooks?