[Paper Review] Batch-Shaping for Learning Conditional Channel Gated Networks

Babak Ehteshami Bejnordi, Tijmen Blankevoort+1 more|arXiv (Cornell University)|Jul 15, 2019

Advanced Neural Network Applications43 references38 citations

TL;DR

The paper introduces fine-grained channel gating in ResNet-based architectures plus a batch-shaping loss to align feature marginals to a prior, enabling conditional computation with higher accuracy at similar or lower runtime cost, demonstrated on CIFAR-10, ImageNet, and Cityscapes.

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.

Motivation & Objective

• Motivate conditional computation to reduce unnecessary feature computation during inference.
• Propose a fine-grained channel gating mechanism within residual blocks to selectively activate convolutional channels.
• Introduce batch-shaping to align feature distributions to a chosen prior and promote data-conditioned gates.
• Combine batch-shaping with an L0-style sparsity loss to control sparsity and trade off accuracy vs. computation.

Proposed method

• Introduce a gated ResNet block where a gating vector G(x_l) multiplies the intermediate activations after the first convolution and before the second convolution.
• Gate outputs g_c are produced by a lightweight gating module that uses channel-wise global average pooling followed by a small two-layer fully connected network and a Gumbel-Softmax/BinConcrete relaxation for differentiable binary decisions.
• Use a batch-shaping loss based on the Cramér-von-Mises criterion to match the marginal posterior of gate activations to a chosen prior (e.g., Beta distribution) to encourage conditionality.
• Apply an L0-like sparsity loss on gates to regularize the overall gating and control the number of active channels, with a delayed introduction to avoid premature capacity loss.
• Train end-to-end on CIFAR-10, ImageNet, and Cityscapes with a curriculum of losses (batch-shaping first, then sparsity) and varying lambda, gamma to explore trade-offs.

Experimental results

Research questions

• RQ1 Can fine-grained channel-level gating provide better accuracy vs. computation trade-offs than gating at coarser granularity (e.g., whole blocks) or network pruning approaches?
• RQ2 Does batch-shaping effectively induce gates to be more conditional on the input, and how does this affect accuracy and MAC savings?
• RQ3 What is the impact of combining batch-shaping with an L0-like sparsity loss on performance and compute in large-scale datasets?
• RQ4 How do gated networks compare to existing conditional-computation methods (ConvNet-AIG, SkipNet, dynamic channel pruning) in terms of accuracy at similar compute?

Key findings

• On ImageNet, ResNet50-BAS and ResNet34-BAS achieve 74.60% and 72.55% top-1 accuracy at comparable compute to ResNet18 baseline (69.76%), demonstrating improved accuracy with conditional computation.
• Batch-shaping guides gates to be more data-dependent; models using batch-shaping plus L0-like sparsity achieve higher accuracy at similar or reduced MACs compared to baselines.
• Gated ResNet variants consistently outperform ConvNet-AIG and ConvNet-FBS at similar accuracy-cost points across CIFAR-10 and ImageNet.
• Semantics segmentation on Cityscapes shows PSPNet with gating reaching IoU 0.719 and pixel accuracy 0.935 using 76.3% MAC of the original PSPNet, with improvements when pretrained on ImageNet (IoU 0.747, acc 0.948 at 95% MAC).
• Gates distribute into always-on, conditionally-on, and always-off categories, with the BAS variant favoring conditional sparsity over complete gating-off for better utilization of capacity.