[Paper Review] Multi-Loss Sub-Ensembles for Accurate Classification with Uncertainty Estimation

Omer Achrack, Raizy Kellerman+1 more|arXiv (Cornell University)|Oct 5, 2020

Adversarial Robustness in Machine Learning45 references4 citations

TL;DR

This paper proposes a multi-loss sub-ensemble method that trains multiple heads with different loss functions on shared features, enabling end-to-end training for accurate classification with uncertainty estimation. The approach achieves state-of-the-art accuracy on SVHN, CIFAR10, CIFAR100, and ImageNet while maintaining low inference time and competitive uncertainty calibration.

ABSTRACT

Deep neural networks (DNNs) have made a revolution in numerous fields during the last decade. However, in tasks with high safety requirements, such as medical or autonomous driving applications, providing an assessment of the models reliability can be vital. Uncertainty estimation for DNNs has been addressed using Bayesian methods, providing mathematically founded models for reliability assessment. These model are computationally expensive and generally impractical for many real-time use cases. Recently, non-Bayesian methods were proposed to tackle uncertainty estimation more efficiently. We propose an efficient method for uncertainty estimation in DNNs achieving high accuracy. We simulate the notion of multi-task learning on single-task problems by producing parallel predictions from similar models differing by their loss. This multi-loss approach allows one-phase training for single-task learning with uncertainty estimation. We keep our inference time relatively low by leveraging the advantage proposed by the Deep-Sub-Ensembles method. The novelty of this work resides in the proposed accurate variational inference with a simple and convenient training procedure, while remaining competitive in terms of computational time. We conduct experiments on SVHN, CIFAR10, CIFAR100 as well as Image-Net using different architectures. Our results show improved accuracy on the classification task and competitive results on several uncertainty measures.

Motivation & Objective

• Address the need for reliable uncertainty estimation in safety-critical deep learning applications such as autonomous driving and medical diagnosis.
• Overcome the computational inefficiency of Bayesian methods and the high inference cost of existing non-Bayesian uncertainty estimation techniques.
• Improve model accuracy and uncertainty calibration without sacrificing inference speed, especially for real-time systems.
• Enable end-to-end training of diverse sub-ensembles through loss-specific head optimization, avoiding separate training phases.
• Provide a practical, scalable alternative to deep ensembles and MC-dropout with competitive performance on standard benchmarks.

Proposed method

• Train multiple sub-ensemble heads on shared feature extractors, each optimized with a distinct loss function (e.g., cross-entropy, MSE, MAE).
• Use a single forward pass during inference, leveraging the diversity of predictions from different loss-optimized heads to estimate uncertainty.
• Apply variational inference principles by treating each head as a sample from a posterior distribution over weights, enabling epistemic uncertainty estimation.
• Maintain low inference cost by reusing the shared trunk, similar to Deep Sub-Ensembles (DSE), while enhancing diversity through loss-based head specialization.
• Integrate uncertainty estimation into standard classification training without requiring post-hoc modeling or multiple forward passes.
• Use the ensemble of heads to compute uncertainty via prediction variance or entropy of softmax outputs, enabling detection of mispredictions on in-distribution data.

Experimental results

Research questions

• RQ1 Can training multiple heads with different loss functions on shared features improve classification accuracy while enabling uncertainty estimation?
• RQ2 Does the proposed multi-loss sub-ensemble method achieve better uncertainty calibration than existing non-Bayesian methods like MC-dropout and Deep Ensembles?
• RQ3 Can the method maintain low inference time while improving model accuracy and uncertainty estimation performance?
• RQ4 How does the diversity induced by different loss functions affect uncertainty estimation and model robustness?
• RQ5 Can uncertainty measures derived from the ensemble effectively detect mispredictions even on in-distribution data?

Key findings

• The proposed method achieves state-of-the-art accuracy on SVHN, CIFAR10, CIFAR100, and ImageNet, outperforming MC-dropout and DSE on these benchmarks.
• On CIFAR10, the method achieves a test accuracy of 95.2% with an expected calibration error (ECE) of 0.113%, competitive with the best-performing methods.
• On ImageNet using Xception, the method achieves 82.3% top-1 accuracy and an ECE of 0.129%, demonstrating strong performance on large-scale datasets.
• The method reduces ECE compared to prior work, with ECE values consistently below 5% across all datasets, indicating strong calibration.
• Uncertainty measures from the ensemble effectively separate correct from incorrect predictions, enabling threshold-based detection of model mispredictions.
• The method maintains low inference cost—significantly lower than Deep Ensembles—while achieving competitive uncertainty estimation and improved accuracy.