Plug-in Losses for Evidential Deep Learning: A Simplified Framework for Uncertainty Estimation that Includes the Softmax Classifier

Technical Summary: Plug-in Losses for Evidential Deep Learning

Problem Statement
Real-world sensor-based learning systems, particularly in safety-critical and resource-constrained environments (e.g., embedded sensing), require uncertainty estimation that is both reliable and computationally efficient. While methods like Monte Carlo dropout, Bayesian neural networks, and deep ensembles provide uncertainty quantification, they often incur significant computational overhead or require complex training procedures unsuitable for real-time inference.

Evidential Deep Learning (EDL) addresses this by enabling single-pass uncertainty estimation. Instead of predicting a single class probability, EDL models class probabilities via a Dirichlet distribution, with the network predicting the distribution's parameters (evidence). However, the standard EDL approach relies on minimizing the expected loss over the unknown Dirichlet distribution. This results in complex objective functions (e.g., involving digamma functions for cross-entropy or variance terms for MSE) that depend on both projected probabilities and total Dirichlet concentration. These complexities can lead to difficult optimization landscapes, sensitivity to hyperparameters, and implementation challenges, creating a gap between theoretical foundations and algorithmic practicality.

Methodology
The authors propose a simplified framework that approximates the EDL objective using plug-in losses. The core theoretical insight is derived from a first-order empirical risk minimization (ERM) perspective:

1. Plug-in Approximation: The authors approximate the Dirichlet-expected loss, $\mathbb{E}_{\pi \sim \text{Dir}(\alpha)}[\ell(\pi, y)]$, by evaluating the loss function directly at the Dirichlet mean, $\hat{p} = \Pi(\alpha) = \alpha / \alpha_0$.

   • For smooth losses (e.g., Mean-Squared Error), a Taylor expansion around the Dirichlet mean shows that the approximation error decays as $O((\alpha_0 + 1)^{-1})$ as the evidence (concentration parameter $\alpha_0$) grows.
   • For cross-entropy loss, which is not globally smooth, a separate analysis using Dirichlet logarithmic moments demonstrates that the error is $O(\alpha_0^{-1})$ under an interiority condition (where the predicted probability is bounded away from zero).

2. Simplified Evidential Classifiers: The framework defines a "simplified evidential classifier" where the optimization depends solely on the projected probabilities $\hat{p}(z)$ rather than the intermediate evidence parameters $\alpha(z)$ directly.

   • Theoretical Unification: A key theoretical result (Theorem 1) establishes that any neural network with a standard softmax output layer is a specific instance of a simplified evidential classifier. This occurs when the evidence mapping is exponential ($e_i = \exp(z_i)$) and the Dirichlet parameter mapping is identity ($\alpha_i = e_i$).

3. Role of KL Regularization: The paper analyzes the role of the Kullback-Leibler (KL) divergence regularizer, commonly used in classical EDL to penalize evidence assigned to incorrect classes. The authors interpret this as a deterministic regularization function that influences the learned evidence scale and concentration, thereby affecting uncertainty scores like "vacuity" (inverse concentration), but noting it is not strictly required to define the simplified classifier itself.

Key Contributions

• Theoretical Derivation: The paper derives that the EDL framework can be simplified by using plug-in losses evaluated at the Dirichlet mean, rigorously showing that the approximation error vanishes with growing evidence for a broad class of losses (including MSE and cross-entropy).
• Softmax Justification: The analysis provides a theoretical justification for the use of standard softmax classifiers within the EDL framework, identifying them as a special case of simplified evidential classifiers.
• Empirical Validation: The proposed simplified objectives are validated on the Google Speech Commands v1 dataset. The study compares classical EDL (with Dirichlet-expected losses) against simplified variants (plug-in losses) and standard softmax models.
• Operational Analysis: The work presents the first uncertainty analysis in the spirit of EDL for speech recognition, evaluating coverage-accuracy trade-offs and uncertainty distributions for correct versus incorrect predictions.

Results
Experiments on the Google Speech Commands v1 dataset using the MatchboxNet architecture yield the following findings:

• Predictive Performance: Simplified evidential classifiers trained with plug-in losses (e.g., Plug-in CE, Plug-in MSE) achieve predictive accuracy and selective prediction performance comparable to classical EDL.
• Uncertainty Scores:
  • Entropy: Entropy-based selection remains a strong metric. Simplified classifiers (including standard Softmax) perform competitively, often achieving the best or near-best total accuracy at high-reliability operating points (e.g., 99.9% thresholded accuracy).
  • Vacuity: Dirichlet-expected losses alone do not make vacuity (inverse concentration) competitive with entropy. However, adding KL regularization significantly improves the reliability ranking induced by vacuity. For instance, a Softmax model with KL regularization achieved 80.36% total accuracy at a 99.9% thresholded accuracy target under vacuity-based selection, outperforming non-regularized variants.
• Implementation: The simplified objectives are easier to implement using standard deep learning losses and training pipelines, removing the need for complex closed-form expressions of Dirichlet expectations.

Significance and Claims
The paper claims that the proposed simplified framework offers a practical remedy to the computational and implementation complexities of classical EDL without sacrificing predictive or selective prediction performance. By demonstrating that plug-in losses approximate the EDL objective well under high-evidence regimes, the authors suggest that the complex Dirichlet-expected structure is often unnecessary for achieving effective uncertainty estimation.

Crucially, the work bridges the gap between standard softmax classifiers and evidential learning, suggesting that softmax networks can be viewed as simplified evidential classifiers. This provides a foundation for understanding the performance of softmax in uncertainty estimation contexts. The authors note that their empirical analysis is the first to obtain coverage-accuracy trade-offs for speech recognition tasks through EDL.

Limitations and Future Directions
The authors acknowledge that the empirical evaluation is currently restricted to a single keyword-spotting benchmark (Google Speech Commands v1) and focuses on in-distribution selective prediction. They suggest that future work should evaluate the framework across a wider range of datasets, architectures, and distribution-shift settings to fully assess its deployment reliability. Additionally, understanding the local and global agreement between uncertainty orderings induced by vacuity and entropy remains an open question.