We present a method that trains large capacity neural networks with significantly improved accuracy and lower dynamic computational cost. This is achieved 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. Almost all deep neural networks have a prior that seems suboptimal: All features are calculated all the time. Both from a generalization and an inference-time perspective, this is superfluous. For example, there is no reason to compute features that help differentiate between several dog breeds, if there is no dog to be seen in the image. The necessity of specific features for classification performance depends on other features. We can make use of this natural prior information, to improve our neural networks. We can also exploit this to spend less computational power on simple and more on complicated examples. This general idea is commonly encapsulated in the terms conditional computing (Bengio, 2013) or gating architectures (Sigaud et al., 2016) . It is known that models with increased capacity, for example increased model depth or width (Zagoruyko & Komodakis, 2016) , generally help increase model performance when properly regularized. However, as models increase in size, so do their training and inference times. This often limits practical applications of deep learning. By conditionally turning parts of the architecture on or off we can train networks with very large capacity while keeping the computational overhead small. The hypothesis is that when training conditionally gated networks, we can train models with a better accuracy/computation cost trade-off than their fully feed-forward counterparts. In addition, conditional computation with channel gating can potentially increase the interpretability of models. For example, gating can allow us to identify input patterns that trigger the response of a single unit in the network. Several works in recent literature that successfully learn conditional architectures, for example, ConvNet-AIG (Veit & Belongie, 2018) and dynamic channel pruning . However, the conditionality is often very coarse as in ConvNet-AIG (Veit & Belongie, 2018) , or the amount of actual conditional features learned is very minimal as in Gaternet . We attempt to solve both. Our contributions are as follows: • We propose a fine-grained gating architecture that turns individual input and output convolutional maps on or off, leading to features that are individually gated. This allows for a better trade-off between computation cost and accuracy than previous work. • We propose a generally applicable tool named batch-shaping that matches the marginal aggregated posterior of a feature in the network to a specified prior. Depending on the chosen prior, networks can match activation distributions to e.g. the uniform distribution for better quantization, or the Gaussian to enforce behavior similar to batch-normalization (Ioffe & Szegedy, 2015) . Specifically, in this paper, we apply batch-shaping to help the network learn conditional features. We show that this helps performance by controlling the gates to be more conditional on the input data at the end of training. • We show state-of-the-art results compared to other conditional computing architectures such as Convnet-AIG (Veit & Belongie, 2018) , SkipNet (Wang et al., 2018a) , Dynamic Channel Pruning , and soft-guided adaptively-dropped neural network (Wang et al., 2018b) . In this paper, we presented a fine-grained gating architecture that enables conditional computation in deep networks. Our gating network achieves state-of-the-art accuracy among competing conditional computation architectures on CIFAR10 and ImageNet datasets. In both datasets, given a model with large capacity, our gating method was the only approach that could reduce the inference computation to a value equal or lower than that of a lower capacity base network, while obtaining a higher accuracy. On ImageNet, our ResNet50-BAS and ResNet34-BAS improve the accuracy by more than 4.8% and 2.8% over a ResNet18 model at the same computation cost. We also proposed a novel batch-shaping loss that can match the marginal aggregated posterior of a feature in the network to any prior PDF. We use it to enforce each gate in the network to be more conditionally activated at the start of training and improve performance significantly. We look forward to seeing many novel applications for this loss in the future, for e.g., autoencoders, better quantized models, and as an alternative to batch-normalization. Another important future research direction that can benefit from our fine-grained gating architecture is continual learning. Designing gating mechanisms that can dynamically decide to allow or prevent the flow of gradients through certain parts of the network could potentially mitigate catastrophic forgetting. Finally, with our gating setup, we could distill smaller sub-networks that work on only a subset of the trained classes.
A method that trains large capacity neural networks with significantly improved accuracy and lower dynamic computational cost A method to train a network with large capacity, only parts of which are used at inference time dependent on input, using fine-grained conditional selection and a new method of regularization, "batch shaping."