RADMI: Latent Information Aggregation as a Proxy for Model Uncertainty

TL;DR

RADMI is a single-pass method that estimates segmentation uncertainty by measuring mutual information between decoder layers, offering a scalable alternative to ensemble methods.

Contribution

It introduces RADMI, a novel approach that efficiently estimates uncertainty through inter-layer mutual information without modifying network architecture.

Findings

RADMI achieves the highest correlation with ensemble uncertainty among single-pass methods.

It outperforms baselines by 5.5% in Pearson and 10.7% in Spearman correlation.

RADMI produces sharp, boundary-localized uncertainty maps without additional computational cost.

Abstract

Epistemic uncertainty estimation is essential for identifying regions where deep learning system outputs may be unreliable. However, existing approaches require computationally expensive ensemble methods or multiple stochastic forward passes, limiting their scalability to dense prediction tasks like segmentation. We propose Resolution-Aggregated Decoder Mutual Information (RADMI), a single-pass method that estimates prediction uncertainty by measuring mutual information (MI) between consecutive decoder layers in segmentation networks. We observe that elevated inter-layer MI correlates with prediction uncertainty, as the network must integrate conflicting contextual information at ambiguous regions such as class boundaries. Evaluating on a seismic facies segmentation benchmark, RADMI achieves the highest correlation with deep ensemble uncertainty among all single-pass methods, outperforming the next-best baselines by 5.5% in Pearson and 10.7% in Spearman correlation coefficients. Compared to baselines that either lack spatial precision or demand significant computational overhead, RADMI yields sharp, boundary-localized uncertainty maps without architectural modifications. Our results suggest that linear aggregation of normalized information flow provides a principled and efficient proxy for prediction uncertainty in encoder-decoder architectures.