Split-Ensemble: Efficient OOD-aware Ensemble via Task and Model Splitting

TL;DR

This work tackles the challenge of robust uncertainty estimation and OOD detection without relying on external OOD data or increased inference costs. It introduces Split-Ensemble, which partitions a multiclass task into complementary subtasks learned by submodels sharing a common backbone, forming a tree-like ensemble that estimates uncertainty via OOD-aware subtask objectives. The method combines a subtask-splitting training objective with correlation-based automated splitting and Hessian-based global pruning to produce an efficient architecture that preserves submodel diversity while sharing low-level features. Empirical results on CIFAR-10/100, Tiny ImageNet, and ImageNet-1K show improved accuracy and OOD AUROC over single-model baselines and larger naive ensembles, demonstrating practical gains for reliable deployment in real-world scenarios.

Abstract

Uncertainty estimation is crucial for machine learning models to detect out-of-distribution (OOD) inputs. However, the conventional discriminative deep learning classifiers produce uncalibrated closed-set predictions for OOD data. A more robust classifiers with the uncertainty estimation typically require a potentially unavailable OOD dataset for outlier exposure training, or a considerable amount of additional memory and compute to build ensemble models. In this work, we improve on uncertainty estimation without extra OOD data or additional inference costs using an alternative Split-Ensemble method. Specifically, we propose a novel subtask-splitting ensemble training objective, where a common multiclass classification task is split into several complementary subtasks. Then, each subtask's training data can be considered as OOD to the other subtasks. Diverse submodels can therefore be trained on each subtask with OOD-aware objectives. The subtask-splitting objective enables us to share low-level features across submodels to avoid parameter and computational overheads. In particular, we build a tree-like Split-Ensemble architecture by performing iterative splitting and pruning from a shared backbone model, where each branch serves as a submodel corresponding to a subtask. This leads to improved accuracy and uncertainty estimation across submodels under a fixed ensemble computation budget. Empirical study with ResNet-18 backbone shows Split-Ensemble, without additional computation cost, improves accuracy over a single model by 0.8%, 1.8%, and 25.5% on CIFAR-10, CIFAR-100, and Tiny-ImageNet, respectively. OOD detection for the same backbone and in-distribution datasets surpasses a single model baseline by, correspondingly, 2.2%, 8.1%, and 29.6% mean AUROC.

Figures (6)

• Figure 1: Overview. We split an original task into complementary subtasks to create objectives for submodel training. All submodels form an ensemble to perform the original task, and, importantly, each submodel can be trained with OOD-aware objectives of the subtask as proposed in Section \ref{['sec:split']}. To implement an efficient Split-Ensemble architecture, we start with a shared backbone and iteratively perform splitting and pruning based on subtask similarity and sensitivity described in Section \ref{['sec:prune']}.
• Figure 2: Subtask splitting. Each submodel learns its subtask using a subset of the original training data. OOD detection by outlier exposure training is realized using other subtasks' examples. Concatenated ID logits from all submodels implement the original multiclass classification task.
• Figure 3: Iterative splitting and pruning. Starting from a shared backbone, we compute the layer-wise sensitivity mask $\mathcal{M}$ for each subtask loss, and calculate pair-wise IoU score $J$ across different subtasks for the layer-wise correlation graph. Model is split at the layer with a small minimal cutting threshold (MCT), and, then, is pruned globally. Applying splitting and pruning in an iterative fashion leads to the final Split-Ensemble architecture that satisfies computational cost constraints.
• Figure 4: TSNE visualization of learned features from ResNet-18 with CIFAR10 test set. Samples with different class labels are marked in different colors. The adjacent features within the dashed lines are to be grouped into one subtask.
• Figure 5: Visualization of Split-Ensemble's learned features using Score-CAM. The number of splits is set to 8 and the model is trained on ImageNet1K with ResNet-18 as backbone. The feature maps within the dashed lines across the layers indicate shared representations. The input image's class is 'Angora', targeted by submodel 2.