Probabilistic Skip Connections for Deterministic Uncertainty Quantification in Deep Neural Networks

2025-01-08

AI Summary

Deep learning models struggle to efficiently quantify uncertainty and detect out-of-distribution (OOD) samples without modifying the training pipeline. Method: This paper proposes Probabilistic Skip Connections (PSC), a plug-and-play, forward-pass uncertainty quantification (UQ) framework. PSC identifies the optimal intermediate feature layer—without retraining—using neural collapse metrics, then decouples aleatoric and epistemic uncertainty via lightweight probabilistic modeling of layer-wise features. Integrating spectral normalization with intermediate-feature probabilistic calibration, PSC yields calibrated uncertainty estimates alongside predictions in a single forward pass. Contribution/Results: Extensive experiments demonstrate that PSC matches or surpasses state-of-the-art single-pass UQ methods requiring retraining across multiple benchmarks. It significantly improves confidence calibration and OOD detection performance, achieving high accuracy, reliability, and computational efficiency—all while preserving the original model architecture and training procedure.

Technology Category

• Machine Learning: Calibration & Uncertainty Quantification
• Reasoning under Uncertainty: Relational Probabilistic Models
• Search and Optimization: Learning to Search

Application Category

• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Web Mining and Content Analysis: Large pretrained models with web data
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs

Abstract

Deterministic uncertainty quantification (UQ) in deep learning aims to estimate uncertainty with a single pass through a network by leveraging outputs from the network's feature extractor. Existing methods require that the feature extractor be both sensitive and smooth, ensuring meaningful input changes produce meaningful changes in feature vectors. Smoothness enables generalization, while sensitivity prevents feature collapse, where distinct inputs are mapped to identical feature vectors. To meet these requirements, current deterministic methods often retrain networks with spectral normalization. Instead of modifying training, we propose using measures of neural collapse to identify an existing intermediate layer that is both sensitive and smooth. We then fit a probabilistic model to the feature vector of this intermediate layer, which we call a probabilistic skip connection (PSC). Through empirical analysis, we explore the impact of spectral normalization on neural collapse and demonstrate that PSCs can effectively disentangle aleatoric and epistemic uncertainty. Additionally, we show that PSCs achieve uncertainty quantification and out-of-distribution (OOD) detection performance that matches or exceeds existing single-pass methods requiring training modifications. By retrofitting existing models, PSCs enable high-quality UQ and OOD capabilities without retraining.

Problem

• Uncertainty Quantification
• Anomaly Detection
• Deep Learning

Innovation

• Probabilistic Skip Connections
• Uncertainty Quantification
• Anomaly Detection