Diversity and Depth in Per-Example Routing Models

1 INTRODUCTION . Modern neural networks process each input in the exact same way . This static paradigm is rigid compared to how brains process sensory inputs . Brains can utilize different subnetworks to process different categories of objects , such as face-specific processing in the fusiform face area ( Kanwisher et al. , 1997 ; Grill-Spector & Weiner , 2014 ) . While static neural networks are empirically effective , it remains an open question whether neural networks with input-dependent processing can improve performance . Input-dependent processing holds the promise of offering better parameter efficiency and reduced computation due to the specialization of processing . Input-dependent processing has been underexplored in comparison with the wealth of work on static networks . Much of the work exploring input-dependent processing has taken the form of per-example routing within a network ( Shazeer et al. , 2017 ; Fernando et al. , 2017 ; McGill & Perona , 2017 ; Rosenbaum et al. , 2018 ) , which is form of conditional computation ( Bengio et al. , 2013 ) . In per-example routing , different examples are processed by different subcomponents , or experts ( Jacobs et al. , 1991 ) , inside a larger model , or supernetwork ( Fernando et al. , 2017 ) . Only a subset of experts in the supernetwork are active for any given example . This paradigm enables the experts , each of which has its own set of parameters , to specialize to subsets of the input domain . The process of routing each example , which determines the experts that are used , is learned jointly with the parameters of the experts . Routing models to date have been relatively small , homogeneous networks . Typically , the same architectural unit ( e.g. , a fully connected layer of the same width ) is used for every expert . The experts differ only in the parameters . Intuitively , the diversity of input examples is best handled by a diversity of architectural units with varying properties , implying that the usage of homogeneous experts is limiting . Furthermore , the number of routing decisions made in prior routing network works has typically been five or fewer . More routing decisions increase the number of distinct paths in the network , which may increase representational power . Making static networks deeper reliably improves performance , so we suspect that the representational power of routing networks is limited when only a few routing decisions are made . In this work , we address these two deficiencies in routing models . Since we aim to increase the representational capacities of routing models , we first introduce a simple trick that reduces overfitting . We then show how routing models with architectural diversity represent a broad family of models that generalize a number of powerful models . We also discuss the tradeoffs of scaling up the number of routing decisions with respect to optimization difficulty . In our experiments , we demonstrate that architecturally diverse routing models beat the best baselines with a 35 % improvement in error rate on an Omniglot setup . By ablating the architectural diversity , we show that diversity plays a key role in achieving strong performance . We then scale up the number of decisions in routing models on CIFAR-10 and demonstrate that while competitive performance can be achieved , the accuracy drops as the number of decisions increase due to optimization challenges . Finally , we discuss our both our positive and negative findings and suggest future research directions for routing models . 2 METHODS . 2.1 ROUTING . Static neural network architectures apply the same function to every example . In contrast , inputdependent models attempt to tailor the function to each example . While it is straightforward for a human to manually specify a single static architecture , it is infeasible to specify every input-dependent function by hand . Instead , the input-dependent function must be automatically inferred by the model , which introduces an extra level of complexity in optimization . Given the need to automatically infer architectures for each example , a natural solution is to define a single large model ( supernetwork ) with a numerous subnetworks ( experts ) , and route examples through a path in the supernetwork ( Fernando et al. , 2017 ; Rosenbaum et al. , 2018 ; Pham et al. , 2018 ; Bender et al. , 2018 ) . Figure 1 visualizes this framework . Intuitively , similar examples can be routed through similar paths and dissimilar examples can be routed through different paths . The example-dependent routing also encourages expert specialization , in which experts devote their representational capacity to transforming a chosen subset of examples . Learning to route examples to well-matched experts is critical for good performance . Effective routing can be achieved by training another small neural network ( router ) that learns to route examples through the supernetwork . The router takes the example as input and outputs the next expert to use . The router can take advantage of the intermediate representations of the example produced in the supernetwork . Given that the router is characterized as neural network , it must be trained . One option is to use reinforcement learning ( Rosenbaum et al. , 2018 ) . In this work , we use the noisy top-k gating technique of Shazeer et al . ( 2017 ) that enables the learning of the router directly by gradient descent . Noisy top-k gating is used to choose k experts out of a list of n experts in a way that enables gradients to flow to the router . The noisy top-k gating technique is defined as : y = n∑ i G ( x ) iEi ( x ) G ( x ) = softmax ( keep-top-k ( H ( x ) , k ) ) H ( x ) i = ( x ·Wg ) i + · log ( 1 + exp ( ( x ·Wnoise ) i ) ) keep-top-k ( v , k ) i = { vi if vi is in the top k elements of v −∞ otherwise ∼ N ( 0 , 1 ) where y is the output , G ( x ) are the expert weights , and E ( x ) i is the ith expert . Both Wg and Wnoise can be learned using the gradients that flow through G ( x ) . While the keep-top-k operation introduces discontinuities in the output , gradient descent learning with noisy top-k gating works well in practice . The induced sparsity of choosing k components significantly reduces computational cost , in contrast to soft gating mechanisms where the entire supernetwork is activated ( Eigen et al. , 2013 ) . The number of routes taken by the examples is controlled by the hyperparameter k. The choice of k plays an important role in controlling the tradeoff between weight sharing and specialization . We find that using a small fixed k throughout the network leads to overfitting . To combat overfitting , we introduce a simple trick which we call k-annealing . In k-annealing , instead of using the same value of k for every layer , k is annealed downwards over the layers of the routing model . That is , initial layers will have high k ( more weight sharing ) and layers closer to the classification layer have smaller k ( more specialization ) . We found in early experimentation that k-annealing works well in practice , and we use it for all experiments in this work . If noisy top-k gating is used without modifications , the model tends to immediately collapse to using a small number of experts instead of using all the capacity . The collapse happens in the beginning of training . Some experts tend to perform slightly better due to random initialization , and gradient descent will accordingly route more examples to the best few experts , further reinforcing their dominance . To combat this problem , Shazeer et al . ( 2017 ) used additional losses that balance expert utilization , such as the importance loss : Limportance ( X ) = cv ( ∑ x∈X G ( x ) ) 2 where cv is the coefficient of variation . We use the importance loss and the load loss , which is defined in detail in Shazeer et al . ( 2017 ) , in all our models . 2.2 DIVERSITY . Previous works on routing use a homogeneous set of experts , where all experts share the same architecture and only differ in the parameter values ( Shazeer et al. , 2017 ; Rosenbaum et al. , 2018 ) . Intuitively , this pattern is suboptimal because different examples may need to be processed by architectural units with different properties . For example , Figure 2 shows two images from Imagenet ( Russakovsky et al. , 2015 ) of borzoi dogs from the same class . The borzoi on the left fills the image and may require layers that can combine information across the entirety of the image , such as convolutions with large spatial extent , whereas the borzoi on the right is less prominent and may require layers , such as max pooling , that can ignore the distracting background information . We attempt to remedy this problem by introducing architectural diversity in routing models , with different experts having different architectures and parameters . The simplest example of this pattern is a feedforward network with multiple layers . In a standard static neural network , an example would be processed as : hl+1i = φ ( conv-3x3 ( h l i ) ) where φ denotes a nonlinearity . In contrast , a model with diverse routing could be specified as : hl+1i = gatingk ( [ h l i , φ ( conv-3x3 ( h l i ) ) , φ ( conv-7x7 ( h l i ) ) , max-pool-3x3 ( h l i ) ] ) Even this simple structure can represent a number of computational models . If the example requires gathering information from a large receptive field , the conv-7x7 can be used . If skip connections ( He et al. , 2016 ) are beneficial , the router can choose the identity operation . Interestingly , introducing architectural diversity into routing models can be viewed as a way to generalize a number of recently proposed models . One-shot neural architecture search models ( Brock et al. , 2017 ; Pham et al. , 2018 ; Bender et al. , 2018 ) jointly perform architecture learning and parameter learning in the same supernetwork . These one-shot models are just a special case of diverse routing models when the input to the router is a constant across examples , so every example follows the same path inside the supernetwork . Diverse routing models can be extended to the multi-task setting by simply concatenating a task embedding to the input of the routing model ( Johnson et al. , 2016 ) , with the task embedding being shared across all examples of the same task . If only the task embedding is fed as input to the router , then all examples from the same task will be routed through the same path , just as in task-based routing models such as Pathnet ( Fernando et al. , 2017 ) . Thus , routing models with diversity can implement complex architectures with powerful representational capacity . We take the example of implementing a neural architecture search cell from the one-shot architecture search model of Bender et al . ( 2018 ) , which is depicted in Figure 3 . The model in Bender et al . ( 2018 ) consists of stacking the same cell structure multiple times . The inputs to each cell are the outputs of the previous two cells . The cell is built by stacking numerous choice blocks . Each choice block takes as input the cell inputs and the outputs of all previous choice blocks in the cell . Within each choice block , two decisions are made : what inputs to use , and what convolutional operations to apply . After choosing inputs , the aggregated input representation is passed to a router that selects between the different convolution , pooling , and identity operations . The input representation is fed to the selected operations , and the outputs of each operation are aggregated . This complex architecture can be readily implemented as a diverse routing model . The main differences compared to Bender et al . ( 2018 ) are that ( a ) in the diverse routing models , each example is routed differently , whereas Bender et al . ( 2018 ) uses a single path for all examples , and ( b ) for diverse routing models , each cell is different , whereas for Bender et al . ( 2018 ) , a single cell architecture is chosen and replicated across the network .

Summary: Overall, this is a valuable read. Authors tackle the head on problem of what is a good architecture where we can having routing with diverse models. The papers is written well, with comparisons to mixture of experts, other models that tackled this problem with either homogeneous architectures or static architectures. Below is my assessment on various axis: