A Framework for Variational Inference of Lightweight Bayesian Neural Networks with Heteroscedastic Uncertainties

1. Problem Statement

Bayesian Neural Networks (BNNs) are essential for quantifying predictive uncertainty, distinguishing between epistemic uncertainty (model uncertainty due to limited data) and aleatoric uncertainty (data noise).

• The Challenge: In many applications, particularly those requiring heteroscedastic uncertainties (where noise levels vary with the input), standard approaches require the BNN to output the aleatoric variance as an additional network head alongside the predictive mean.
• The Limitation: Adding an output node for variance increases the number of learnable parameters and architectural complexity. This is detrimental for lightweight BNNs intended for deployment on hardware-constrained systems or where computational time is limited.
• The Gap: Existing methods often force a trade-off between model size and the ability to capture complex, input-dependent noise structures.

2. Methodology

The authors propose a novel inference framework that eliminates the need for a separate output node for aleatoric uncertainty by embedding both aleatoric and epistemic uncertainties directly into the variances of the BNN's learnable parameters.

A. Reinterpretation of Parameter Distributions

Instead of treating the variance of Bayesian neuron parameters as solely representing epistemic uncertainty, the authors reinterpret these distributions to encode the total predictive uncertainty (aleatoric + epistemic).

• Standard Approach: $Var(\hat{y}) = \sigma^2_{epistemic}$. Aleatoric variance ($\sigma^2_{aleatoric}$) is predicted as a separate output.
• Proposed Approach: The total variance $\sigma^2_k$ for a given input $k$ is defined as the sum of aleatoric and epistemic variances:
  $$\sigma^2_k = \sigma^2_{a,k} + \sigma^2_{e,k}$$
  During training, the network minimizes the Evidence Lower Bound (ELBO) where the reconstruction loss assumes the output distribution is Gaussian with variance $\sigma^2_k$. The variance $\sigma^2_k$ is computed as the variance of the BNN output derived from the parameter distributions, effectively "absorbing" the aleatoric noise into the parameter variances.

B. Sampling-Free Variational Inference via Moment Propagation

To maintain computational efficiency and avoid the cost of Monte Carlo sampling, the authors employ a moment propagation scheme.

• Mechanism: The mean and variance of the input to each layer are propagated analytically through the network to compute the mean and variance of the output.
• Layer Handling:
  • Fully Connected & Convolutional: Exact analytical propagation of means and variances assuming independence between neurons.
  • Leaky-ReLU: Approximate propagation assuming the input to the activation function is normally distributed. The authors derive closed-form expressions for the mean and variance of the output using the law of total expectation and variance, involving truncated normal distributions.
• Benefit: This allows for deterministic inference without the need for repeated sampling, making the framework suitable for lightweight, real-time applications.

3. Key Contributions

1. Parameter Variance Embedding: A conceptual shift where the variances of BNN weights and biases are trained to represent the sum of aleatoric and epistemic uncertainties, removing the need for a dedicated variance output head.
2. Lightweight Architecture: The framework enables BNNs to capture heteroscedastic uncertainties with fewer learnable parameters compared to standard "learned variance" approaches.
3. Sampling-Free Inference: Integration of a simplified moment propagation technique that provides analytical uncertainty estimates, avoiding the computational overhead of Monte Carlo dropout or sampling.
4. Unified Framework: A complete pipeline combining the reinterpretation of uncertainty sources with efficient moment propagation for training and inference.

4. Experimental Results

The authors evaluated their "Embedded Variance" approach against the standard "Learned Variance" approach (where variance is a separate output) using a polynomial regression task with heteroscedastic Gaussian noise.

• Performance in Lightweight Regimes:
  • With a small number of hidden neurons (e.g., 4), the Embedded Variance model accurately predicted both the trend and the uncertainty for in-distribution data and correctly expressed high uncertainty for out-of-distribution data.
  • The Learned Variance model failed to recover true uncertainty for in-distribution data and significantly underestimated uncertainty for out-of-distribution data with the same small architecture.
• Parameter Efficiency:
  • For a fixed number of neurons, the "Learned Variance" model required ~34% more learnable parameters than the "Embedded Variance" model to achieve comparable performance.
  • The "Embedded Variance" approach achieved lower reconstruction loss (negative log-likelihood) in the low-parameter regime.
• Generalization:
  • The "Embedded Variance" models demonstrated better generalization to out-of-distribution data, showing a more gradual increase in reconstruction loss as data moved away from the training distribution compared to the "Learned Variance" models.
• Scalability: As the network size increased (e.g., 1024 neurons), the performance gap narrowed, and both methods performed similarly, suggesting the proposed method is most critical for resource-constrained, lightweight models.

5. Significance

This work addresses a critical bottleneck in deploying Bayesian Neural Networks for edge computing and safety-critical systems.

• Efficiency: By removing the need for additional output nodes to predict noise, the method significantly reduces model size and computational cost, making BNNs viable for hardware with strict constraints.
• Robustness: The ability to capture heteroscedastic uncertainties without increasing model complexity ensures that lightweight models do not sacrifice reliability or safety assessments.
• Practicality: The sampling-free nature of the inference allows for faster prediction times, which is essential for real-time applications where latency is a concern.

In summary, the paper demonstrates that total uncertainty can be effectively "baked into" the parameter distributions of a BNN, offering a superior alternative to traditional multi-output heads for lightweight, heteroscedastic uncertainty quantification.