Training independent subnetworks for robust prediction

TL;DR

This paper introduces a method to train multiple independent subnetworks within a single model using a multi-input multi-output configuration, enhancing robustness and uncertainty estimation without additional computational cost.

Contribution

It presents a novel MIMO-based training approach that enables independent subnetworks to be trained within one model, improving robustness and uncertainty estimation efficiently.

Findings

Improved negative log-likelihood, accuracy, and calibration on CIFAR10, CIFAR100, ImageNet.

Achieved robustness benefits with a single forward pass, reducing computational costs.

Enhanced out-of-distribution detection performance.

Abstract

Recent approaches to efficiently ensemble neural networks have shown that strong robustness and uncertainty performance can be achieved with a negligible gain in parameters over the original network. However, these methods still require multiple forward passes for prediction, leading to a significant computational cost. In this work, we show a surprising result: the benefits of using multiple predictions can be achieved `for free' under a single model's forward pass. In particular, we show that, using a multi-input multi-output (MIMO) configuration, one can utilize a single model's capacity to train multiple subnetworks that independently learn the task at hand. By ensembling the predictions made by the subnetworks, we improve model robustness without increasing compute. We observe a significant improvement in negative log-likelihood, accuracy, and calibration error on CIFAR10, CIFAR100, ImageNet, and their out-of-distribution variants compared to previous methods.