Uncertainty in Real-Time Semantic Segmentation on Embedded Systems

1. Problem Statement

Semantic segmentation is critical for safety-critical applications like autonomous vehicles and robotics, where real-time performance on resource-constrained embedded hardware (e.g., NVIDIA Jetson) is mandatory. While deep learning models have improved in accuracy, they often lack the ability to quantify uncertainty (epistemic and aleatoric) during inference.

• The Challenge: Traditional methods for uncertainty quantification, such as Monte Carlo Dropout or complete Bayesian inference, require multiple forward passes or expensive sampling, making them computationally intractable for real-time embedded systems.
• The Gap: Existing real-time models prioritize speed over uncertainty reasoning, while existing uncertainty methods prioritize accuracy over speed. There is a lack of solutions that provide meaningful uncertainty estimates in real-time on edge devices without significant performance degradation.

2. Methodology

The authors propose a lightweight, analytic probabilistic framework that combines deterministic feature extraction with Bayesian regression and moment propagation. The core architecture consists of three stages:

A. Hybrid Architecture (Deterministic + Probabilistic)

• Feature Extraction: The model utilizes a pre-trained, deterministic neural network (e.g., ENet, BiSeNet, PIDNet) with fixed parameters ($\theta$) to extract features. This ensures the heavy lifting of feature extraction remains fast and compatible with existing real-time models.
• Probabilistic Classification: The final linear/convolutional layer of the network is replaced with a probabilistic layer. Instead of fixed weights, this layer uses latent variables ($\omega$) treated as random variables.
• Logit Space Inference: Inference is performed in the logit space rather than the categorical space to simplify the math. The model assumes a linear relationship between the fixed features and the probabilistic weights.

B. Efficient Inference via Diagonal SWAG

To avoid expensive Monte Carlo sampling, the authors approximate the posterior distribution of the latent variables $\omega$:

• Prior: A conjugate Gaussian prior is placed on $\omega$.
• Posterior Approximation: The authors use a modified SWAG (Stochastic Weight Averaging-Gaussian) method. Instead of computing the empirical mean from training iterates, they use the pre-trained weights as the mean ($\bar{\omega}$) and compute the empirical variance ($\Sigma_{diag}$) using the diagonal approximation of the covariance matrix.
• Benefit: This results in a factorized Gaussian distribution, reducing the covariance matrix to a diagonal matrix. This drastically lowers memory usage and computational complexity, enabling analytic inference without multiple forward passes.

C. Moment Propagation (EPSoftmax)

Since the final output requires categorical probabilities (via Softmax), and the exact analytical solution for the distribution of a Softmax applied to Gaussian logits is intractable, the authors employ Expectation Propagation (EP):

• Log-Normal Approximation: They approximate the exponential of the Gaussian logits as Log-Normal distributions.
• Ratio Distribution: Using Taylor expansion and properties of Log-Normal distributions, they approximate the mean and variance of the final softmax output ($t_j$).
• Result: This yields an EPSoftmax layer that outputs a valid categorical distribution with associated mean and variance, allowing for analytic uncertainty calculation.

D. Uncertainty Metrics

The framework distinguishes between two types of uncertainty:

1. Epistemic Uncertainty (Model Uncertainty): Calculated using the entropy of the Gaussian distribution over the logits (Eq. 18). This represents reducible uncertainty due to the model's lack of knowledge.
2. Aleatoric Uncertainty (Data Uncertainty): Approximated using the entropy of the final categorical distribution (Eq. 19). This represents irreducible noise inherent in the data (e.g., occlusion, blur).
3. Class-Conditional Uncertainty: The variance for specific classes can be extracted to identify uncertainty for safety-critical objects (e.g., pedestrians).

3. Key Contributions

1. Lightweight Uncertainty Method: A novel combination of Bayesian regression and moment propagation that enables uncertainty-aware predictions on embedded hardware with minimal computational overhead.
2. Real-Time Metrics: Development of meaningful real-time metrics for both aleatoric and epistemic uncertainty, demonstrating how they can inform end-user decision protocols.
3. Adaptability: A demonstration that the probabilistic module can be easily adapted to various pre-trained real-time segmentation models (ENet, BiSeNet, PIDNet) without retraining the entire feature extractor.

4. Experimental Results

The method was evaluated on the NVIDIA Jetson AGX Xavier embedded device using the CityScapes, CamVid, ADE20k, and CoCoStuff datasets.

• Performance vs. Speed:

  • The proposed "Bayes" models achieved comparable predictive accuracy (mIOU, Macro-F1) to their deterministic point-estimate counterparts.
  • Inference Speed: While there was a slight decrease in Frames Per Second (FPS) compared to standard models, the models maintained real-time performance (e.g., Bayes-PIDNet achieved ~69 FPS on CityScapes, Bayes-BiSeNetV1 ~57 FPS).
  • Overhead: The computational overhead was significantly lower than Monte Carlo methods, which are impossible to run in real-time on this hardware.

• Qualitative Analysis:

  • Epistemic Uncertainty: Visualized as high entropy in regions where the model is unsure (e.g., object boundaries, unfamiliar objects).
  • Aleatoric Uncertainty: Concentrated around the edges of objects and in noisy regions, effectively capturing data ambiguity.
  • Class-Conditional Uncertainty: Successfully highlighted uncertainty for specific classes (e.g., "truck" or "person"), showing higher variance at object boundaries.

5. Significance and Future Work

• Safety Critical Applications: This work bridges the gap between high-performance real-time segmentation and the safety requirements of autonomous systems. By providing uncertainty estimates in real-time, systems can make safer decisions (e.g., slowing down when uncertainty is high).
• Edge AI: It proves that complex probabilistic reasoning does not require massive cloud compute; it can be efficiently executed on resource-constrained edge devices.
• Future Directions: The authors suggest integrating aleatoric uncertainty directly into the training loss of the base network and exploring how to utilize these uncertainty metrics for optimal decision-making protocols in automated systems.

In conclusion, the paper presents a practical, analytically derived solution for embedding uncertainty quantification into real-time semantic segmentation, making it viable for deployment in safety-critical, resource-constrained environments.