PaperlayerView exact revision on arXiv

Source-linked AI summary

Training independent subnetworks for robust prediction

Marton Havasi, Rodolphe Jenatton, Stanislav Fort, Jeremiah Zhe Liu, Jasper Snoek, Balaji Lakshminarayanan, Andrew M. Dai, Dustin Tran

arXiv:2010.06610v2cs.LGcs.CVstat.ML

TL;DR

Existing efficient ensembles improve robustness and uncertainty but still require multiple forward passes. MIMO trains independent subnetworks inside one network and ensembles their predictions in a single pass, with reported gains across standard and out-of-distribution benchmarks.

  • Problem

    Efficient ensemble methods can provide strong robustness and uncertainty performance with few extra parameters, but typically require 4-10 forward passes for prediction.

  • Method

    MIMO changes the input and output layers to train multiple subnetworks jointly, then repeats the test input so their predictions form a single-pass ensemble.

  • Results

    MIMO achieves new state-of-the-art after wall-clock adjustment on CIFAR10, CIFAR100, ImageNet, and their out-of-distribution variants.

  • Takeaways & Limitations

    MIMO provides multiple diverse predictions with minor parameter and compute changes while improving uncertainty estimation and robustness.

Abstract

from arXiv · show

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.

1 INTRODUCTION

MIMO trains multiple independent subnetworks within one network and ensembles their predictions in a single forward pass, targeting robust and well-calibrated prediction without substantial computational cost.

  • Motivation: Multiple forward passes remain a runtime cost in efficient ensemble methods, motivating single-pass ensembling.These methods can add few parameters while typically requiring 4-10 prediction passes.
  • Method: MIMO uses one network to concurrently train multiple subnetworks that co-habit without explicit separation.The approach is motivated by evidence that heavily overparameterized networks can retain performance after pruning 70-80% of connections.
  • Method: During training, independently sampled inputs are matched to separate output heads; during testing, repeating one input produces an ensemble in one forward pass.The heads learn to ignore features from nonmatching inputs because those features are not useful for the corresponding prediction.
  • Method: MIMO subnetworks use disjoint network parts and have diversity comparable to independently trained neural networks.This diversity is presented as essential because ensembles are more robust when members are less likely to make the same mistake.
  • Contributions: MIMO enables multiple independent predictions in a single forward pass, improving uncertainty estimation and robustness with minor parameter and compute changes.The contribution summary describes the benefits as achieved “for free.”
  • Results: MIMO achieves new state-of-the-art after wall-clock adjustment on CIFAR10, CIFAR100, ImageNet, and their out-of-distribution variants.The reported comparison spans both standard and distribution-shifted benchmark variants.

2 MULTI-INPUT MULTI-OUTPUT NETWORKS

MIMO trains multiple subnetworks from independently sampled input-output pairs and evaluates them jointly in one forward pass. Its ensemble can trade a slight bias increase for lower variance and improved generalization when the network has sufficient capacity.

  • Architecture and training: MIMO concatenates M inputs and produces M corresponding predictive distributions using small input- and output-layer changes.During training, each input is sampled independently and each output predicts its matching input.
  • Architecture and training: The training objective sums the negative log-likelihoods of the M predictions with a regularization term and is optimized using stochastic gradient descent.Because the input-output pairs are independent, this is equivalent to training on M-tuples of independently sampled datapoints.
  • Prediction and computation: At evaluation, one input is tiled M times, producing M predictions that are averaged as an ensemble.All subnetworks are evaluated simultaneously, so MIMO’s wall-clock time is almost equivalent to that of a standard neural network.
  • Subnetwork diversity: MIMO subnetworks converge to distinct low-loss basins, whereas naive multihead subnetworks converge to the same mode and are essentially equivalent.The comparison examines test accuracy and subnetwork disagreement across a weight-space section containing M = 3 subnetworks and the origin.
  • Synthetic regression example: In the synthetic regression experiment, expected mean squared error is evaluated across ensemble sizes M ∈ {1, . . . , 5}.The experiment uses a two-hidden-layer multilayer perceptron and training and test sets of 64 and 3000 observations, respectively.
  • Synthetic regression example: MIMO learns diverse predictors whose lower ensemble variance can offset the slight bias increase associated with M > 1.The method improves when the model has sufficient capacity to fit multiple diverse, well-performing ensemble members.

3 UNDERSTANDING THE SUBNETWORKS

MIMO subnetworks converge to distinct, independently behaving solutions because separate inputs encourage diverse feature usage. The ensemble benefits most at moderate M, while capacity limits and input correlation constrain performance.

  • Subnetwork independence: MIMO subnetworks converge to distinct local optima and functionally behave like independently trained neural networks.
  • Subnetwork independence: Separate inputs promote diversity because features derived from each input are useful primarily for its corresponding output.A naive multiheaded model with shared inputs instead produces nearly identical predictions.
  • Diversity comparison: MIMO matches independently trained neural networks in prediction diversity and exceeds prior efficient ensemble approaches.
  • Robustness mechanism: Ensembling diverse subnetwork predictions is identified as the source of MIMO’s robustness, replicating Deep Ensemble behavior within one network.
  • Subnetwork separation: Conditional-variance analysis finds activations responsive to one input and close to zero with respect to the others, supporting functional subnetwork separation.The analysis measures how pre-activations change when individual inputs vary.
  • Optimal number of subnetworks: The ensemble peaks between M = 2 and M = 4 because ensembling benefits outweigh declining individual-subnetwork performance.Accuracy peaks earlier than log-likelihood, suggesting greater ensemble benefit for log-likelihood.

4 BENCHMARKS

The benchmarks evaluate MIMO across CIFAR10, CIFAR100, and ImageNet using predictive quality, robustness, calibration, and computational-cost metrics. MIMO outperforms other single-forward-pass methods and approaches Deep Ensemble robustness without added computational cost.

  • MIMO is evaluated with ResNet28-10 on CIFAR10 and CIFAR100, and ResNet50 on ImageNet.
  • The baselines include deterministic networks, MC-Dropout, BatchEnsemble, ensembles, Naive multihead, and TreeNet.
  • The evaluation measures accuracy, negative log-likelihood, and expected calibration error on IID and corrupted test sets.Computational cost is measured as TPUv2 evaluation time in milliseconds per example.
  • MIMO significantly outperforms other single forward pass methods on all three benchmarks.
  • MIMO approaches the robustness of a Deep Ensemble without increasing computational costs.

5 RELATED WORK

Related efficient ensemble methods reduce memory or computation through parameter sharing, multiheaded architectures, or training heuristics. MIMO differs by targeting ensemble diversity while retaining single-forward-pass evaluation.

  • Earlier multiheaded networks share inputs and network layers while separating final layers and outputs, reducing computation but providing limited ensemble diversity.
  • Aggregated Learning uses multiple inputs to create an information bottleneck and regularizes mutual information between inputs and that representation.
  • BatchEnsemble, Rank-1 BNNs, and hyper batch ensembles share most parameters but still require multiple forward passes at evaluation.
  • Data augmentation, temperature scaling, label smoothing, and contrastive training are compute-efficient heuristics that are orthogonal to MIMO.

6 CONCLUSIONS

The paper presents MIMO as a simple, efficient method for training independent subnetworks within one network. Its empirical results support improved performance and robustness with minor parameter and compute changes.

  • MIMO trains multiple independent subnetworks within a single network and evaluates them in one forward pass.
  • The subnetworks separate within the model and behave like independently trained neural networks.
  • MIMO requires few hyperparameters and no significant network-architecture modifications.
  • MIMO improves performance and robustness with minor changes to parameter count and compute cost.

A PSEUDOCODE

The pseudocode trains MIMO by sampling multiple inputs and labels, producing one prediction per input, optimizing their summed loss, and updating parameters. At inference, one repeated input produces multiple predictions for averaging.

  • Each training iteration samples M input-label pairs independently from the training set.
  • MIMO produces M predictions from the multi-input batch, with each prediction matched to its corresponding input.
  • The training objective sums the negative log-probabilities for the M matched predictions with a regularizer before a gradient update.
  • At inference, the same input is supplied across the M input slots to obtain M predictions for aggregation.

B HYPERPARAMETERS

The experiments use distinct training settings for CIFAR and ImageNet, while baseline and MIMO configurations differ in training duration, regularization, and batch repetition.

  • Baseline settings: CIFAR baseline models use batch-size 512, learning-rate decay, L2 regularization 2e-4, and 200–250 training epochs.Deterministic, Dropout, and Ensemble models train for 200 epochs; BatchEnsemble, Naive multihead, and TreeNet train for 250.
  • Baseline settings: ImageNet baseline models use batch-size 4096, learning-rate decay, L2 regularization 1e-4, and 90–150 training epochs.Training duration varies by model, from 90 epochs for Deterministic, Dropout, and Ensemble to 150 for Naive multihead and TreeNet.
  • Baseline settings: Ensemble and BatchEnsemble use M = 4 members, while Dropout uses a 10% rate and one evaluation-time forward pass.The four-member choice is intended to provide most ensemble benefits without significantly increasing computational costs.
  • MIMO settings: MIMO uses batch repetition 4 on CIFAR and 2 on ImageNet, with 250 and 150 training epochs respectively.Its L2 regularization is 3e-4 for CIFAR and 1e-4 for ImageNet.

C MIMO BETTER EXPLOITS THE NETWORK CAPACITY: PERFORMANCE

The capacity analysis tests whether MIMO’s performance becomes more sensitive as regularization constrains the network. Results support this prediction, with stronger effects for larger ensembles.

  • Experimental rationale: Higher regularization is treated as greater capacity constraint, while larger MIMO ensembles are expected to exploit more capacity.The protocol predicts stronger sensitivity when both regularization and ensemble size increase.
  • Experimental setup: MIMO is evaluated on ResNet28-10 with CIFAR10 and CIFAR100 while varying L1 or L2 regularization paths.Accuracy and log-likelihood are averaged over three repetitions.
  • Findings: MIMO becomes more sensitive to increasing regularization, and the effect strengthens with larger ensemble size.This matches the experimental hypothesis about capacity exploitation.

D ADDITIONAL IMAGENET OOD RESULTS

Additional ImageNet evaluations examine multiple out-of-distribution datasets using standard and dataset-specific robustness metrics. The reported results are consistent with the main benchmarks across distribution shifts.

  • Evaluation scope: The additional ImageNet evaluation covers ImageNet, ImageNet-C, ImageNet-A, ImageNetV2, ImageNet-Vid-Robust, YTTB-Robust, and ObjectNet.The first four datasets use negative log-likelihood, accuracy, and expected calibration error; ImageNet-C also reports mCE.
  • Evaluation scope: ImageNet-Vid-Robust, YTTB-Robust, and ObjectNet use their own predefined stability metrics.These metrics differ from the three metrics used on the other listed datasets.
  • Results: MIMO consistently outperforms methods on corrupted images and across distribution shifts in the additional experiments.These results are reported as consistent with the main-text benchmarks.
Loading 2010.06610v2…