Torch-Uncertainty: Deep Learning Uncertainty Quantification

Adrien Lafage, Olivier Laurent, Firas Gabetni, Gianni Franchi

Keywords: Uncertainty Quantification, Bayesian Deep Learning, Model Calibration, Ensemble Methods, Trustworthy AI, PyTorch

TL;DR: This paper introduces torch-uncertainty, a unified PyTorch-based framework that benchmarks state-of-the-art uncertainty quantification methods across multiple deep learning tasks and modalities.

Abstract: Deep Neural Networks (DNNs) have demonstrated remarkable performance across various domains, including computer vision and natural language processing. However, they often struggle to accurately quantify their predictions' uncertainty, limiting their broader adoption in critical industrial applications. Uncertainty Quantification (UQ) for Deep Learning seeks to address this challenge by providing methodologies to improve the reliability of uncertainty estimates. While numerous techniques have been proposed, a unified tool remains lacking that offers a seamless workflow for evaluating and integrating these methods. To bridge this gap, we introduce **Torch-Uncertainty**, a *PyTorch* and *Lightning* framework designed to streamline the training and evaluation of DNNs with UQ techniques. In this paper, we outline the foundational principles of our library and present comprehensive experimental results that benchmark a diverse set of UQ methods across classification, segmentation, and regression tasks. Our library is available at: https://github.com/ENSTA-U2IS-AI/torch-uncertainty.

Code URL: https://github.com/ENSTA-U2IS-AI/torch-uncertainty

Primary Area: Datasets & Benchmarks illustrating Different Deep learning Scenarios (e.g., architectures, generative models, optimization for deep networks, foundation models, LLMs)

Submission Number: 1834