Unified Probabilistic Neural Architecture and Weight Ensembling Improves Model Robustness

TL;DR

This paper introduces UraeNAS, a probabilistic approach combining neural architecture search and Bayesian inference to generate ensembles that improve model robustness and uncertainty calibration in classification tasks.

Contribution

It proposes a unified probabilistic framework for neural architecture and weight ensembling, addressing architecture uncertainty in Bayesian neural networks.

Findings

Significant accuracy and calibration improvements on CIFAR-10 and CIFAR-10-C datasets.

Outperforms baseline deterministic approaches in robustness and uncertainty estimation.

Demonstrates the effectiveness of joint architecture and weight ensembling for model robustness.

Abstract

Robust machine learning models with accurately calibrated uncertainties are crucial for safety-critical applications. Probabilistic machine learning and especially the Bayesian formalism provide a systematic framework to incorporate robustness through the distributional estimates and reason about uncertainty. Recent works have shown that approximate inference approaches that take the weight space uncertainty of neural networks to generate ensemble prediction are the state-of-the-art. However, architecture choices have mostly been ad hoc, which essentially ignores the epistemic uncertainty from the architecture space. To this end, we propose a Unified probabilistic architecture and weight ensembling Neural Architecture Search (UraeNAS) that leverages advances in probabilistic neural architecture search and approximate Bayesian inference to generate ensembles form the joint distribution of neural network architectures and weights. The proposed approach showed a significant improvement both with in-distribution (0.86% in accuracy, 42% in ECE) CIFAR-10 and out-of-distribution (2.43% in accuracy, 30% in ECE) CIFAR-10-C compared to the baseline deterministic approach.