Batch-Shaping for Learning Conditional Channel Gated Networks

TL;DR

This paper introduces a novel gating mechanism and batch-shaping technique for neural networks, enabling conditional computation that improves accuracy while reducing average computational cost on image classification and segmentation tasks.

Contribution

It proposes a new residual block architecture with fine-grained channel gating and a batch-shaping method to align feature posteriors with prior distributions, enhancing conditional computation.

Findings

Gated networks achieve higher accuracy with similar or lower computational cost.

Networks learn to allocate more features to difficult examples and fewer to simple ones.

Significant accuracy improvements on ImageNet and CIFAR-10 datasets.

Abstract

We present a method that trains large capacity neural networks with significantly improved accuracy and lower dynamic computational cost. We achieve this by gating the deep-learning architecture on a fine-grained-level. Individual convolutional maps are turned on/off conditionally on features in the network. To achieve this, we introduce a new residual block architecture that gates convolutional channels in a fine-grained manner. We also introduce a generally applicable tool $batch$-$shaping$ that matches the marginal aggregate posteriors of features in a neural network to a pre-specified prior distribution. We use this novel technique to force gates to be more conditional on the data. We present results on CIFAR-10 and ImageNet datasets for image classification, and Cityscapes for semantic segmentation. Our results show that our method can slim down large architectures conditionally, such that the average computational cost on the data is on par with a smaller architecture, but with higher accuracy. In particular, on ImageNet, our ResNet50 and ResNet34 gated networks obtain 74.60% and 72.55% top-1 accuracy compared to the 69.76% accuracy of the baseline ResNet18 model, for similar complexity. We also show that the resulting networks automatically learn to use more features for difficult examples and fewer features for simple examples.