Modular Networks: Learning to Decompose Neural Computation
Scaling model capacity has been vital in the success of deep learning. For a typical network, necessary compute resources and training time grow dramatically with model size. Conditional computation is a promising way to increase the number of parameters with a relatively small increase in resources. We propose a training algorithm that flexibly chooses neural modules based on the data to be processed. Both the decomposition and modules are learned end-to-end. In contrast to existing approaches, training does not rely on regularization to enforce diversity in module use. We apply modular networks both to image recognition and language modeling tasks, where we achieve superior performance compared to several baselines. Introspection reveals that modules specialize in interpretable contexts.
The paper is concerned with conditional computation, which is an interesting topic yet at early stages of research, and as such one that requires much research and investigation. The paper proposes a latent-variable approach to constructing modular networks, modeling the choice of processing modules in a layer as a discrete latent variable.  A modular network is composed of L modular layers, each comprised of M modules and a controller. Each module is a function (standard layer) f_i(x; \theta_i). The controller accepts the input, chooses K of the M modules to process the input, and outputs the  as the module output. Modular layers can be stacked, or placed anywhere inside a standard network. The input/output of a modular layer is a tensor, and as such the layer is easily integrated with standard networks. Interestingly, the paper proposes a controller that outputs a distribution over compositions of modules (i.e., assigns a probability to each of the (M choose K) possible compositions). This is different from Shazeer et al., who assign a probability to each module and choose the K most likely.  The key issue is a training/inference procedure, which is complicated due to the nature of the discrete latent variable. The authors propose trianing the model with Viterbi EM, optimizing a lower bound by introducing a q distribution over $a$, the choice of modules. The authors also mention REINFORCE and noisy top-K as options for training the model, though note that these require additional regularization in order to not collapse modules.   Pros: - Key contribution: conditional computation layer that selects K of M modules per layer by assigning probability to the complete set (more general than greedy assignment as in Shazeer, et al., 2017), and training with EM explicitly leverages sparse computation (as opposed to MOE). - Simple, but powerful, modular layer for conditional computation that can be stacked and assembled in a straight-forward manner. - Well justified training procedure using EM to optimize a lower-bound over the marginal likelihood of y. - Experiments on Penn TreeBank and Cifar100 imply that the method is scalable to larger('ish) datasets.  Cons: ----- - Contribution is somewhat incremental in that the the core idea was introduced by Shazeer et al. 2017 - Related work is somewhat sparse. - Biggest issue is the experimental results: the model does not work well on either benchmark at all: test accuracy was very low and perplexity very high compared to standard-compute models, though these are not shown in the paper. As the authors mention, there is yet a lot of work to make conditional-compute models viable.  My overall impression is that this is a good paper that should be accepted at NIPS. The paper is clearly written, and easy to read. It links nicely to existing work in the field, clearly stating its contributions. The model design and derivation of the training procedure are clearly explained and justified. The paper makes a reasonable contribution in the area of conditional-compute by suggesting distributions over K modules per layer rather than softmax with top-K. This enables the use of a reasonable q distribution and generlized EM for training. Despite the weak experimental results, the paper represents a step in the direction for conditional neural computation.