DISCOVERING OBJECTS AND THEIR RELATIONS FROM ENTANGLED SCENE REPRESENTATIONS
Our world can be succinctly and compactly described as structured scenes of objects and relations. A typical room, for example, contains salient objects such as tables, chairs and books, and these objects typically relate to each other by virtue of their correlated features, such as position, function and shape. Humans exploit knowledge of objects and their relations for learning a wide spectrum of tasks, and more generally when learning the structure underlying observed data. In this work, we introduce relation networks (RNs) a general purpose neural network architecture for object-relation reasoning. We show that RNs are capable of learning object relations from scene description data. Furthermore, we show that RNs can act as a bottleneck that induces the factorization of objects from entangled scene description inputs, and from distributed deep representations of scene images provided by a variational autoencoder. The model can also be used in conjunction with differentiable memory mechanisms for implicit relation discovery in one-shot learning tasks. Our results suggest that relation networks are a powerful architecture for solving a variety of problems that require object relation reasoning.
1 INTRODUCTION
The ability to reason about objects and relations is important for solving a wide variety of tasks (Spelke et al., 1992; Lake et al., 2016). For example, object relations enable the transfer of learned knowledge across superficial (dis)similarities (Tenenbaum et al., 2011): the predator-prey relationship between a lion and a zebra is knowledge that is similarly useful when applied to a bear and a salmon, even though many features of these animals are very different.
In this work, we introduce a neural network architecture for learning to reason about – or model – objects and their relations, which we call relation networks (RNs). RNs accomplish their goal by adhering to a few fundamental design principles. Firstly, RNs are designed to be invariant to permutations of object descriptions in their input. For example, RN representations of the object set {table, chair, book} will be identical for arbitrary re-orderings of the elements of the set. Secondly, RNs are designed to learn relations across multiple objects rather than within a single object – a basic defining property of object relations. This design principle manifests itself in the use of shared computations across groups of objects. In designing the RN architecture, we took inspiration from the recently developed Interaction Network (IN) (Battaglia et al., 2016) which is more generally concerned with modelling temporal interactions.
In principle, a deep network with a sufficiently large number of parameters and a large enough training set should be capable of matching the performance of a RN. However, such networks would have to learn both the permutation invariance of objects and the relational structure of the objects in the execution of a desired computation. This quickly becomes unfeasible as the number of objects and relations increase.
∗Denotes equal contribution.
To test the ability of RNs to discover relations between objects, we turned to the classification of scenes, wherein classification boundaries were defined by the relational structure of the objects in the scenes. There are various ways of encoding scenes as observable data; for example, scene description data can consist of sets of co-occurring objects and their respective features (location, size, color, shape, etc.). A typical room scene, then, might consist of the object set {table, chair, lamp} and their respective descriptions (e.g., the table is large and red). Similar datasets have been used for many decades in cognitive psychology to study human relation learning (Palmer, 1975). Here we consider synthetic scene description data generated by hierarchical probabilistic generative models of objects and object features.
We begin by motivating RNs as a general purpose architecture for reasoning about object relations. Critically, we describe how RNs implement a permutation invariant computation on implicit groups of factored “objects.” We then demonstrate the utility of RNs for classification of static scenes, where classification boundaries are defined by the relations between objects in the scenes. Next, we exploit RNs’ implicit use of factored object representations to demonstrate that RNs can induce the factorization of objects from entangled scene inputs. Finally, we combine RNs with memoryaugmented neural networks (MANNs) (Santoro et al., 2016) to solve a difficult one-shot learning task, demonstrating the ability of RNs’ to act in conjunction with other neural network architectures to rapidly discover new object relations from entirely new scenes.
2 MODEL
2.1 DESCRIPTION AND IMPLEMENTATION
RNs are inspired by Interaction Networks (INs) (Battaglia et al., 2016), and therefore share similar functional insights. Both operate under the assumption that permutation invariance is a necessary requirement for problems with inputs that can be described with a graphical structure. However, INs use relations between objects as input to determine object interactions, mainly for the purpose of reasoning about dynamics. RNs compute object relations, and hence aim to determine objectrelational structure from static inputs.
Suppose we have an object, o = (o1, o2, ..., on), represented as a vector of n features encoding properties such as the object’s type, color, size, position, etc. A collection of m objects can be gathered into a m× n matrix D, called the scene description. Although the term scene description alludes to visual information, this need not be the case; scenes can be entirely abstract, as can the objects that constitute the scene, and the features that define the objects.
We can imagine tasks (see section 3) that depend on the relations, r, between objects. One such task is the discovery of the relations themselves. For example, returning to the predator and prey analogy, the predator-prey relation can be determined from relative features between two animals – such as their relative sizes, perhaps. Observation of the size of a single animal, then, does not inform whether this animal is a predator or prey to any other given animal, since its size necessarily needs to be compared to the sizes of other animals.
There are a number of possible functions that could allow for the discovery of object relations (figure 1). Consider a function g, with parameters ψ. The function gψ can be defined to operate on a particular factorization of D; for example, gψ(D), or gψ(D1,2, ..., Di,j,, ..., Dm,n). We are interested in models defined by the composite function f ◦ g, where f is a function that returns a prediction r̃.
One implementation of gψ would process the entire contents of D without exploiting knowledge that the features in a particular row are related through their description of a common object, and would instead have to learn the appropriate parsing of inputs and any necessary sub-functions: r̃ = fφ(gψ(o 1 1, o 2 1, ..., o n m)). An alternative approach would be to impose a prior on the parsing of the input space, such that g operates on objects directly: r̃ = fφ(gψ(o1), gψ(o2), ..., gψ(om)). A third middle-ground approach – which is the approach taken by RNs – recognizes that relations necessarily exist in the context of a set of objects that have some capacity to be related. Thus, the computation of relations should entail some common function across these sets. For example, g may operate on pairs of objects: r̃ = fφ(gψ(o1, o2), gψ(o1, o3), ..., gψ(om−1, om)).
For RNs, g is implemented as a multi-layered perceptron (MLP) that operates on pairs of objects. The same MLP operates on all possible pairings of objects from D, to produce summaries si,j . Depending on the complexity of the problem, any aggregation function a can be chosen to operate on si,j . However, to ensure the input to f is order invariant, the aggregation function should be commutative and associative, with summation being a natural choice. Thus, r̃ = fφ(a(s1,2, s1,3, ..., sm−1,m)) = fφ( ∑ i,j si,j) = fφ( ∑ i,j gψ(oi, oj)). Training therefore entails the optimization of two MLPs: fφ and gψ .
The order invariance of the aggregation function is a critical feature of the model, since without this invariance, the model would be unable to take advantage of the shared gψ function that operates on all permuted pairs of objects.
3 EXPERIMENTAL TASKS AND DATA
3.1 DATASETS
To probe a model’s ability to both infer relations from scene descriptions and implicitly use relations to solve more difficult tasks – such as one-shot learning – we first developed datasets of scene descriptions and their associated images. To generate a scene description, we first defined a graph of object relations (see figure 2). For example, suppose there are four types of squares, with each type being identified by its color. A graph description of the relations between each colored square could identify the blue square as being a parent of the orange square. If the type of relation is “position,” then this particular relation would manifest as blue squares being independently positioned in the scene, and orange squares positioned in close proximity to blue squares. Similarly, suppose we have triangles and circles as the object types, with color as the relation. If triangles are parents to circles, then the color of triangles in the scene will be randomly sampled, while the color of a circle will be derived from the color of the parent triangle. Datasets generated from graphs, then, impose solutions that depend on relative object features. That is, no information can be used from within object features alone – such as a particular coordinate position, or RGB color value.
Graphs define generative models that can be used to produce scenes. For scenes with position relations, root node coordinates (opx, o p y), were first randomly chosen in a bounded space. Children were then randomly assigned to a particular parent object, and their position was determined as a function of the parent’s position: (ocx, o c y) = (o p x + d cos(θ c), opy + d sin(θ c)). Here, θc ∼ U(θp − π/3, θp + π/3), where θp is the angle computed for the parent. For root nodes, then θp ∼ U(0, 2π). d is a computed distance: d = d0+d1, where d0 is a minimum distance to prevent significant object overlap, and d1 is sampled from a half-normal distribution. This same pattern – inherit features from parents and apply noise – is used to generate scenes from graphs that define other relations, such as color. For the case of color, the inherited features are RGB values. Ultimately, scene descriptions consisted of matrices with 16 rows, with each row describing the object type (four rows for each of four types), position, color, and size.
Custom datasets were required to both explicitly test solutions to the task of inferring object relations, and to actively control for solutions that do not depend on object relations. For example,
consider the common scenario of child objects positioned close to a parent object, analogous to chairs positioned around a table. In this scenario, the count information of objects (i.e., that there are more child objects than parent objects) is non-relational information that can nonetheless be used to constrain the solution space; the prediction of the relation between objects doesn’t entirely depend on explicitly computed relations, such as the relative distance between the child objects and the parent objects. To return to the kitchen scene analogy, one wouldn’t need to know the relative distance of the chairs to the table; instead, knowing the number of chairs and number of tables could inform the relation, if it is known that less frequently occurring objects tend to be parents to more frequently occurring objects. Although information like object count can be important for solving many tasks, here we explicitly sought to test models’ abilities to compute and operate on object relations.
The datasets will be made freely available.
3.2 TASK DESCRIPTIONS
The tasks on which we assessed the RN’s performance fell into three categories. The first category involved the classification of scenes. In this task, the network was given a scene description as input, with the target being a binary vector description of the edges between object types (which constitute the nodes of the graph – see figure 2). Training data consisted of 5000 samples derived from 5, 10, or 20 unique classes (i.e., graphs), with testing data comprising withheld within-class samples. Although the target was a vector description of the generating graph, this task was fundamentally one of classification, and the use of one-hot vector labels could have constituted equivalent training. Nonetheless, since class membership can only be determined from the relational structure of the objects within a particular sample, the ability to classify subsumes the ability to infer the relations.
The second category of tasks tested the ability of the RN to classify scenes – and hence, reason about object relations – from unstructured input domains (see figure 3). Since RNs operate on factored object representations, this task specifically probed the ability of RNs to induce the learning of object factorizations from entangled scene descriptions. In the first set of experiments we broke the highly structured scene description input by passing it through a fixed permutation matrix. To decode the entangled state we used a linear layer whose output provided the input to the RN. We then asked the network to classify scenes, as in the previous tasks. In the second set of experiments we pushed
this idea further, and tested the ability of the RN to operate on distributed representations of image depictions of scenes. The images were passed through a variational autoencoder (VAE), and the latent variables were provided as input to a RN.
The final category of tasks tested the implicit use of discovered relations to solve difficult overarching problem: one-shot learning (Vinyals et al., 2016; Santoro et al., 2016; Lake et al., 2015). In the one-shot learning task, sequences of samples were fed to a memory-augmented neural network (MANN) with a relational network pre-processor. Sequences – or episodes – consisted of 50 random samples generated from five unique graphs, from a pool of 1900 total classes, presented jointly with time-offset label identifiers, as per Hochreiter et al. (2001) and Santoro et al. (2016). Critically, the labels associated with particular classes change from episode-to-episode. So, the network must depend on within-episode knowledge to perform the task; it must learn the particular label assigned to a class within an episode, and learn to assign this label to future samples of the same class, within the same episode. Labels are presented as input in a time-offset manner (that is, the correct label for the sample presented at time t is given as input to the network at time t+1) to enable learning of an arbitrary binding procedure. Unique information from a sample – which is necessarily something pertaining to the relations of the contained objects – must be extracted, bound to its associated label, and stored in memory. Upon subsequent presentations of samples from this same class, the network must query its memory, and use stored information to infer class membership. There is a critical difference in this task compared to the first. In this task, identifying class labels constantly change episode-to-episode. So, the network cannot simply encode mappings from certain learned relational structures to class labels. Instead, the only way the network can solve the task is to develop an ability to compare and contrast extracted relational structures between samples as they occur within an episode. Please see the appendix for more details on the task setup.
4 ADDITIONAL MODEL COMPONENTS AND TRAINING DETAILS
4.1 VARIATIONAL AUTOENCODER
For inferring relations from latent representations of pixels we used a variational autoencoder (VAE) (Kingma & Welling, 2013) with a convolutional neural network (CNN) as the feature encoder and deconvolutional network as the decoder (see figure 3b). The CNN consisted of two processing blocks. The input to each block was sent through four dimension preserving parallel convolution streams using 8 kernels of size 1x1, 3x3, 5x5, and 7x7, respectively. The outputs from these convolutions were passed through a batch normalization layer, rectified linear layer and concatenated. This was then convolved again with a down-sampling kernel of size 3x3, halving the dimension size of the input, and, except for the final output layer, again passed through batch normalization and rectified linear layers. The entire CNN consisted of two of these blocks positioned serially. There-
fore, input images of size 32x32 were convolved to feature-maps of size 8x8. The feature decoder consisted of these same blocks, except convolution operations were replaced with deconvolution operations.
The final feature maps provided by the CNN feature encoder were then passed to a linear layer, whose outputs constituted the observed variables x for the VAE. x, which was decomposed into µ and σ, was then used with an auxiliary Gaussian noise variable to infer the latent variables z = µ + σ. These latent variables were then decoded to generate the reconstruction x̂, as per the conventional implementation of VAEs. However, z was also fed as input to a linear layer, which projected it to a higher dimensional space – the scene description D (see figure 3b). Importantly, this connection – from z to D – did not permit the backward flow of gradients to the VAE. This prevented the VAE architecture from contributing to the RN’s disentangling solution.
4.2 MEMORY-AUGMENTED NEURAL NETWORK
For implicit discovery of relations from scene descriptions, the RN was used as a pre-processor for a memory-augmented neural network (MANN). The MANN was implemented as in Santoro et al. (2016), and the reader is directed here for full details on using networks augmented with external memories. Briefly, the core MANN module consists of a controller – a long-short term memory (LSTM) (Hochreiter & Schmidhuber, 1997) – that interacts with read and write heads, which in turn interact with an external memory store. The external memory store consists of a number of memory slots, each of which contains a vector “memory.” During reading, the LSTM takes in an input and produces a query vector, which the read head uses to query the external memory store using a cosine distance across the vectors stored in the memory slots, and returns a weighted sum of these vectors based on the cosine distance. During writing, the LSTM outputs a vector that the write head uses to write into the memory store using a least recently used memory access mechanism (Santoro et al., 2016).
4.3 TRAINING DETAILS
The sizes of the RN – in terms of number of layers and number of units for both fφ and gψ – were {200, 200}, {500, 500}, {1000, 1000}, or {200, 200, 200}. The MLP baseline models used equivalent sizes. We experimented with different sizes for the summary vector si,j , which is the output from the RN. Performance is generally robust to the choice of size, with similar results emerging for 100, 200, or 500. The MANN used a LSTM controller size of 200, 128 memory slots, 40 for the memory size, and 4 read and write heads.
The Adam optimizer was used for optimization (Kingma & Ba, 2014), with learning rate of 1e−4 for the scene description tasks, and a learning rate of 1e−5 for the one-shot learning task. The number of iterations varied for each experiment, and are indicated in the relevant figures. All figures show performance on a withheld test-set, constituting 2-5% of the size of the training set. The number of training samples was 5000 per class for scene description tasks, and 200 per class (for 100 classes) for the pixel disentangling experiment. We used minibatch training, with batch-sizes of 100 for the scene description experiments, and 16 (with sequences of length 50) for the one-shot learning task.
5 RESULTS
5.1 SUPERVISED LEARNING OF RELATIONAL STRUCTURE FROM SCENE DESCRIPTIONS
Here, we trained RNs on variations of the classification task described in section 3.2 and contrasted their performance with that of MLPs of different sizes and depths. First, we compared the performance of these models on scenes where the relational structure was defined by position (figure 4a). After 200,000 iterations the RNs reached a cross entropy loss of 0.01 on a withheld test set, with MLPs of similar size failing to match (figure 4a, top). In fact, the smallest RN performed significantly better than the largest MLP. The performance of the MLPs remained poor even after 1 million iterations (cross entropy loss above 0.2 – not shown). This result was consistent for datasets with 5, 10 and 20 scene classes (figure 4a, bottom).
These results were obtained using targets that explicitly described the relations between objects types (binary vector indicating the edges of the graph). Using one-hot vectors as output targets resulted in very similar results, with the RNs’ accuracy reaching 97% (see figure 7 in appendix). Our reasoning for using graph descriptions as targets relates to the potential of the RN model to represent relations between individual pairs of object types independent of other relations that might be present in a scene. Explicitly targeting the individual relations could in principle allow the model to learn the particular components that form the overall scene structure. Indeed, when trained to classify hundreds of scene classes, the RN was then able to generalize to unobserved classes (see figure 8 in appendix). (Note: this generalization to unseen classes would be impossible with the use of one-hot labels, since there is no training information that may inform the vector label a particular unseen class). Additionally, the ability to generalize to unobserved classes suggests that the RN is able to generalize in a combinatorially complex object-relation space because of its ability to learn compositional structure. It is able to use pieces (i.e., specific relations) of learned information and combine them in unique, never-before-seen ways, which is a hallmark feature of compositional learning.
We repeated this experiment using relational structure defined by the color of the objects, with position being randomly sampled for each object (figure 4b). Again, RNs reached a low classification error on withheld data (below 0.01), whereas the MLP models did not (error above 0.2; figure 4b, top). This was observed across datasets with 5, 10 and 20 classes (figure 4b, bottom).
5.2 INFERRING OBJECT RELATIONS FROM NON-SCENE DESCRIPTION INPUTS
There are a number of difficulties that present when inferring object relations. Firstly, one cannot naively apply a deep neural network to infer relations; instead, a particular type of architecture – possibly one that is permutation invariant to some component of its input space – may be necessary.
Secondly, the discovery of the very things that can possess relations – which was assumed knowledge in previous experiments – is difficult.
We showed the ability of a RN to solve the first problem, and here, we gain insight on the second problem. A particular feature – or constraint – of RNs is their computation on implicit groups of factored objects. We tested, then, whether this constraint could act as a learning signal to induce the factorization of objects from entangled scenes. Specifically, we tested whether RNs can learn with two types of inputs that do not have nicely factored object representations: entangled scene descriptions and pixel images.
5.2.1 INFERRING RELATIONS FROM ENTANGLED SCENE DESCRIPTIONS
In this task we probed the network’s ability to classify scenes from entangled scene descriptions. Intuitively, the RN should behave as an architectural bottleneck that can aid the disentangling of objects by a perceptual model to produce factored object representations on which it can operate. To entangle data we started from the scene description D, and reshaped it into a vector of size mn (i.e., a concatenated vector of objects [o1; o2; o3; ..; om]). This vector was then projected using a random permutation matrix B of size mn × mn. We chose a permutation matrix for two reasons. First, because it preserved all the information of the input without adding any additional scaling between the factors. Second, because it was invertable using a matrix multiplication, and hence produced an interpretable solution found by the model. We used a learnable matrix U of size mn × mn – implemented as a linear layer without biases, positioned after the permutation and before the RN (see figure 3). Thus, this layer provided the m objects to the RN. Figure 5a shows the performance of the model on this task.
During learning we visualized the absolute value of BU (figure 5a, inset). We noticed a block structure emerging, illustrating the identification of different objects. Note, white pixels denote a value of 0, and black denote a value of 1. BU was of size mn ×mn, making the multiplication of a flattened D with BU produce a vector of size mn. This vector was then reshaped into a matrix of size m× n, and provided as input to the RN. The block structure of BU suggests that the object k, as perceived by the network, was a linear transformation (given by the block) of a different object l from the ground truth input in D. Of particular interest is the gradual discovery of new objects over time. Because the model is order invariant, there is no pressure to recover the ground truth order of the objects. Similarly, the particular order of features that define the object were not disentangled, since there is no real pressure for this to happen. Therefore, RNs managed to successfully force the identification of different objects without imposing a particular form, or order of the features within objects.
5.2.2 INFERRING RELATIONS FROM PIXELS
Image depictions from 100 unique classes (using the position-relation dataset; see figure 11 in appendix) were passed through a VAE, whose latent variables were provided as input to a linear layer, which acted as a disentangler, and which sent its output to a RN. Both the VAE and RN were trained concurrently – however, gradients from the RN were not propagated to the VAE portion of the model so as to prevent any of the VAE components from contributing to the solution of the RN’s task. Thus, this experiment explicitly tested the ability of the RN to operate on distributed representations of scenes. It should be noted that this task is unique from the previous entangling tasks, as the VAE may be capable of providing both disentangled object representations, as well as relational information, in its compressed latent code. Nonetheless, the results from this task suggest a capacity for RNs to operate on distributed representations, which opens the door to the possible conjunction of RNs with perceptual neural network modules (figure 5b).
5.3 IMPLICIT USE OF RELATIONS FOR ONE-SHOT LEARNING
Here, we assess the potential to use RNs in conjunction with memory-augmented neural networks to quickly – and implicitly – discover object relations and use that knowledge for one-shot learning.
We trained a MANN with a RN pre-processor to do one-shot classification of scenes, as described in section 3.2. In order to solve this task, the network must store representations of the scenes (which, if class representations are to be unique, should necessarily contain relational information), and the
episode-unique label associated with the scene. Once a new sample of the same class is observed, it must use inferred relational information from this sample to query its memory and retrieve the appropriate label.
During training, a sequence of 50 random samples from 5 random classes (out of a pool of 1900) were picked to constitute an episode. The test phase consisted of episodes with scenes from 5 randomly selected and previously unobserved classes (out of a pool of 100). After 500,000 episodes the network showed high classification accuracy when presented with just the second instance of a class (76%) and performance reached 93% and 96% by the 5th and 10th instance, respectively (figure 6a). As expected, since class labels change from episode-to-episode, performance is at chance for the first instance of a particular class.
Although samples from a class are visually very dissimilar, a RN coupled to an external memory was able to do one-shot classification. This suggests that the network has the capacity to quickly extract information about the relational structure of objects, which is what defines class membership, and does so without explicit supervision about the object-relation structure. Replacing the RN preprocessor with an MLP resulted in chance performance across all instances (figure 6b), suggesting that the RN is a critical component of this memory-augmented model.
6 CONCLUSIONS
RNs are a powerful architecture for reasoning about object-relations. They operate under the assumption that relations exist between pairs of objects, or entities. This assumption affords the ability to classify scenes wherein class boundaries are defined by object relations, as well as the ability to induce factored object representations from entangled scene descriptions.
RNs have the potential to successfully operate with perceptual modules to extract object representations on which to operate. This ability to operate with different neural network modules also extends to the domain of one-shot learning, wherein they are able to pair with MANNs to perform one-shot structure learning.
The utility of RNs as a relation-reasoning module suggests that they have the potential to be useful for solving tasks that require reasoning not only about object-object relations, but also about verb-
object relations, as in human-object interaction datasets (Chao et al., 2015) or question-answering tasks that involve reasoning between multiple objects (Krishna et al., 2016).
The ability of RNs to induce a disentangling of their input potentially allows for great flexibility in the discovery of the “objects” that are to be factored. In our case, objects were defined a priori, but this needn’t necessarily be the case. Depending on the task, they may not even constitute physically consistent groups of matter. Combinations of objects can be grouped to become one “cluster” object, or, the notion of what constitutes an object can instead be determined by a deep neural network, which may come up with a very flexible interpretation for “object-ness” in a scene. Developing a way to process various types of input to produce “objects” with which RNs can operate is an exciting direction of future research.
ACKNOWLEDGMENTS
We would like to thank Scott Reed, Daan Wierstra, Nando de Freitas, James Kirkpatrick, and many others on the DeepMind team.

+ Understanding relations between objects is an important task in domains like vision, language and robotics. However, models trained on real-life datasets can often exploit simple object properties (not relation-based) to identify relations (eg: animals of bigger size are typically predators and small-size animals are preys). Such models can predict relations without necessarily understanding them. Given the difficulty of the task, a controlled setting is required to investigate if neural networks can be designed to actually understand pairwise object relations. The current paper takes a significant step in answering this question through a controlled dataset. Also, multiple experiments are presented to validate the "relation learning" ability of proposed Relation Networks (RN).
+ The dataset proposed in the paper ensures that relation classification models can succeed only by learning the relations between objects and not by exploiting "predator-prey" like object properties.
+ The paper presents very thorough experiments to validate the claim that "RNs" truly learn the relation between objects.
  1. In particular, the ability of the RN to force a simple linear layer to disentangle scene description from VAE latent space and permuted description is very interesting. This clearly demonstrates that the RN learns object relations.
  2. The one-shot experiments again demonstrate this ability in a convincing manner. This requires the model to understand relations in each run, represent them through an abstract label and assign the label to future samples from the relationship graph.
Some suggestions:
- Is g_{\psi}(.) permutation invariant as well. Since it works on pairs of objects, how did you ensure that the MLP is invariant to the order of the objects in the pair?
- The RNs need to operate over pairs of objects in order to identify pairwise interactions. However, in practical applications there are more complicated group interactions. (eg. ternary interaction: "person" riding a "bike" wears "helmet"). Would this require g(.) of RN to not just operate on pairs but on every possible subset of objects in the scene? More generally, is such a pairwise edge-based approach scalable to larger number of objects?
- The authors mention that " a deep network with a sufficiently large number of parameters and a large enough training set should be capable of matching the performance of a RN". This is an interesting point, and could be true in practice. Have the authors investigated this effect by trying to identify the minimum model capacity and/or training examples required by a MLP to match the performance of RN for the provided setup? This would help in quantifying the significance of RN for practical applications with limited examples. In other words, the task in Sec. 5.1 could benefit from another plot: the performance of MLP and RN at different amounts of training samples.
- While the simulation setup in the current paper is a great first-step towards analyzing the "relation-learning" ability of RNs, it is still not clear if this would transfer to real-life datasets. I strongly encourage the authors to experiment on real-life datasets like Coco, visual genome or HICO as stated in the pre-review stage.
- Minor: Some terminologies in the paper such as "objects" and "scene descriptions" used to refer to abstract entities can be misleading for readers from the object detection domain in computer vision. This could be clarified early on in the introduction.
- Minor: Some results like Fig. 8 which shows the ability of RN to generalize to unseen categories are quite interesting and could be moved to the main draft for completeness.
The paper proposes a network which is capable of understanding relationships between objects in a scene. This ability of the RN is thoroughly investigated through a series of experiments on a controlled dataset. While, the model is currently evaluated only on a simulated dataset, the results are quite promising and could translate to real-life datasets as well.


OUTRAGEOUSLY LARGE NEURAL NETWORKS: THE SPARSELY-GATED MIXTURE-OF-EXPERTS LAYER
The capacity of a neural network to absorb information is limited by its number of parameters. Conditional computation, where parts of the network are active on a per-example basis, has been proposed in theory as a way of dramatically increasing model capacity without a proportional increase in computation. In practice, however, there are significant algorithmic and performance challenges. In this work, we address these challenges and finally realize the promise of conditional computation, achieving greater than 1000x improvements in model capacity with only minor losses in computational efficiency on modern GPU clusters. We introduce a Sparsely-Gated Mixture-of-Experts layer (MoE), consisting of up to thousands of feed-forward sub-networks. A trainable gating network determines a sparse combination of these experts to use for each example. We apply the MoE to the tasks of language modeling and machine translation, where model capacity is critical for absorbing the vast quantities of knowledge available in the training corpora. We present model architectures in which a MoE with up to 137 billion parameters is applied convolutionally between stacked LSTM layers. On large language modeling and machine translation benchmarks, these models achieve significantly better results than state-of-the-art at lower computational cost.
The capacity of a neural network to absorb information is limited by its number of parameters. Conditional computation, where parts of the network are active on a per-example basis, has been proposed in theory as a way of dramatically increasing model capacity without a proportional increase in computation. In practice, however, there are significant algorithmic and performance challenges. In this work, we address these challenges and finally realize the promise of conditional computation, achieving greater than 1000x improvements in model capacity with only minor losses in computational efficiency on modern GPU clusters. We introduce a Sparsely-Gated Mixture-of-Experts layer (MoE), consisting of up to thousands of feed-forward sub-networks. A trainable gating network determines a sparse combination of these experts to use for each example. We apply the MoE to the tasks of language modeling and machine translation, where model capacity is critical for absorbing the vast quantities of knowledge available in the training corpora. We present model architectures in which a MoE with up to 137 billion parameters is applied convolutionally between stacked LSTM layers. On large language modeling and machine translation benchmarks, these models achieve significantly better results than state-of-the-art at lower computational cost.
1 INTRODUCTION AND RELATED WORK
1.1 CONDITIONAL COMPUTATION
Exploiting scale in both training data and model size has been central to the success of deep learning. When datasets are sufficiently large, increasing the capacity (number of parameters) of neural networks can give much better prediction accuracy. This has been shown in domains such as text (Sutskever et al., 2014; Bahdanau et al., 2014; Jozefowicz et al., 2016; Wu et al., 2016), images (Krizhevsky et al., 2012; Le et al., 2012), and audio (Hinton et al., 2012; Amodei et al., 2015). For typical deep learning models, where the entire model is activated for every example, this leads to a roughly quadratic blow-up in training costs, as both the model size and the number of training examples increase. Unfortunately, the advances in computing power and distributed computation fall short of meeting such demand.
Various forms of conditional computation have been proposed as a way to increase model capacity without a proportional increase in computational costs (Davis & Arel, 2013; Bengio et al., 2013; Eigen et al., 2013; Ludovic Denoyer, 2014; Cho & Bengio, 2014; Bengio et al., 2015; Almahairi et al., 2015). In these schemes, large parts of a network are active or inactive on a per-example basis. The gating decisions may be binary or sparse and continuous, stochastic or deterministic. Various forms of reinforcement learning and back-propagation are proposed for trarining the gating decisions.
∗Equally major contributors †Work done as a member of the Google Brain Residency program (g.co/brainresidency)
While these ideas are promising in theory, no work to date has yet demonstrated massive improvements in model capacity, training time, or model quality. We blame this on a combination of the following challenges:
• Modern computing devices, especially GPUs, are much faster at arithmetic than at branching. Most of the works above recognize this and propose turning on/off large chunks of the network with each gating decision.
• Large batch sizes are critical for performance, as they amortize the costs of parameter transfers and updates. Conditional computation reduces the batch sizes for the conditionally active chunks of the network.
• Network bandwidth can be a bottleneck. A cluster of GPUs may have computational power thousands of times greater than the aggregate inter-device network bandwidth. To be computationally efficient, the relative computational versus network demands of an algorithm must exceed this ratio. Embedding layers, which can be seen as a form of conditional computation, are handicapped by this very problem. Since the embeddings generally need to be sent across the network, the number of (example, parameter) interactions is limited by network bandwidth instead of computational capacity.
• Depending on the scheme, loss terms may be necessary to achieve the desired level of sparsity per-chunk and/or per example. Bengio et al. (2015) use three such terms. These issues can affect both model quality and load-balancing.
• Model capacity is most critical for very large data sets. The existing literature on conditional computation deals with relatively small image recognition data sets consisting of up to 600,000 images. It is hard to imagine that the labels of these images provide a sufficient signal to adequately train a model with millions, let alone billions of parameters.
In this work, we for the first time address all of the above challenges and finally realize the promise of conditional computation. We obtain greater than 1000x improvements in model capacity with only minor losses in computational efficiency and significantly advance the state-of-the-art results on public language modeling and translation data sets.
1.2 OUR APPROACH: THE SPARSELY-GATED MIXTURE-OF-EXPERTS LAYER
Our approach to conditional computation is to introduce a new type of general purpose neural network component: a Sparsely-Gated Mixture-of-Experts Layer (MoE). The MoE consists of a number of experts, each a simple feed-forward neural network, and a trainable gating network which selects a sparse combination of the experts to process each input (see Figure 1). All parts of the network are trained jointly by back-propagation.
While the introduced technique is generic, in this paper we focus on language modeling and machine translation tasks, which are known to benefit from very large models. In particular, we apply a MoE convolutionally between stacked LSTM layers (Hochreiter & Schmidhuber, 1997), as in Figure 1. The MoE is called once for each position in the text, selecting a potentially different combination of experts at each position. The different experts tend to become highly specialized based on syntax and semantics (see Appendix E Table 9). On both language modeling and machine translation benchmarks, we improve on best published results at a fraction of the computational cost.
1.3 RELATED WORK ON MIXTURES OF EXPERTS
Since its introduction more than two decades ago (Jacobs et al., 1991; Jordan & Jacobs, 1994), the mixture-of-experts approach has been the subject of much research. Different types of expert architectures hae been proposed such as SVMs (Collobert et al., 2002), Gaussian Processes (Tresp, 2001; Theis & Bethge, 2015; Deisenroth & Ng, 2015), Dirichlet Processes (Shahbaba & Neal, 2009), and deep networks. Other work has focused on different expert configurations such as a hierarchical structure (Yao et al., 2009), infinite numbers of experts (Rasmussen & Ghahramani, 2002), and adding experts sequentially (Aljundi et al., 2016). Garmash & Monz (2016) suggest an ensemble model in the format of mixture of experts for machine translation. The gating network is trained on a pre-trained ensemble NMT model.
The works above concern top-level mixtures of experts. The mixture of experts is the whole model. Eigen et al. (2013) introduce the idea of using multiple MoEs with their own gating networks as parts of a deep model. It is intuitive that the latter approach is more powerful, since complex problems may contain many sub-problems each requiring different experts. They also allude in their conclusion to the potential to introduce sparsity, turning MoEs into a vehicle for computational computation.
Our work builds on this use of MoEs as a general purpose neural network component. While Eigen et al. (2013) uses two stacked MoEs allowing for two sets of gating decisions, our convolutional application of the MoE allows for different gating decisions at each position in the text. We also realize sparse gating and demonstrate its use as a practical way to massively increase model capacity.
2 THE STRUCTURE OF THE MIXTURE-OF-EXPERTS LAYER
The Mixture-of-Experts (MoE) layer consists of a set of n “expert networks" E1, · · · , En, and a “gating network" G whose output is a sparse n-dimensional vector. Figure 1 shows an overview of the MoE module. The experts are themselves neural networks, each with their own parameters. Although in principle we only require that the experts accept the same sized inputs and produce the same-sized outputs, in our initial investigations in this paper, we restrict ourselves to the case where the models are feed-forward networks with identical architectures, but with separate parameters.
Let us denote by G(x) and Ei(x) the output of the gating network and the output of the i-th expert network for a given input x. The output y of the MoE module can be written as follows:
y = n∑ i=1 G(x)iEi(x) (1)
We save computation based on the sparsity of the output ofG(x). WhereverG(x)i = 0, we need not compute Ei(x). In our experiments, we have up to thousands of experts, but only need to evaluate a handful of them for every example. If the number of experts is very large, we can reduce the branching factor by using a two-level hierarchical MoE. In a hierarchical MoE, a primary gating network chooses a sparse weighted combination of “experts", each of which is itself a secondary mixture-of-experts with its own gating network. In the following we focus on ordinary MoEs. We provide more details on hierarchical MoEs in Appendix B.
Our implementation is related to other models of conditional computation. A MoE whose experts are simple weight matrices is similar to the parameterized weight matrix proposed in (Cho & Bengio, 2014). A MoE whose experts have one hidden layer is similar to the block-wise dropout described in (Bengio et al., 2015), where the dropped-out layer is sandwiched between fully-activated layers.
2.1 GATING NETWORK
Softmax Gating: A simple choice of non-sparse gating function (Jordan & Jacobs, 1994) is to multiply the input by a trainable weight matrix Wg and then apply the Softmax function.
Gσ(x) = Softmax(x ·Wg) (2)
Noisy Top-K Gating: We add two components to the Softmax gating network: sparsity and noise. Before taking the softmax function, we add tunable Gaussian noise, then keep only the top k values, setting the rest to −∞ (which causes the corresponding gate values to equal 0). The sparsity serves to save computation, as described above. While this form of sparsity creates some theoretically scary discontinuities in the output of gating function, we have not yet observed this to be a problem in practice. The noise term helps with load balancing, as will be discussed in Appendix A. The amount of noise per component is controlled by a second trainable weight matrix Wnoise.
G(x) = Softmax(KeepTopK(H(x), k)) (3)
H(x)i = (x ·Wg)i + StandardNormal() · Softplus((x ·Wnoise)i) (4)
KeepTopK(v, k)i = { vi if vi is in the top k elements of v. −∞ otherwise. (5)
Training the Gating Network We train the gating network by simple back-propagation, along with the rest of the model. If we choose k > 1, the gate values for the top k experts have nonzero derivatives with respect to the weights of the gating network. This type of occasionally-sensitive behavior is described in (Bengio et al., 2013) with respect to noisy rectifiers. Gradients also backpropagate through the gating network to its inputs. Our method differs here from (Bengio et al., 2015) who use boolean gates and a REINFORCE-style approach to train the gating network.
3 ADDRESSING PERFORMANCE CHALLENGES
3.1 THE SHRINKING BATCH PROBLEM
On modern CPUs and GPUs, large batch sizes are necessary for computational efficiency, so as to amortize the overhead of parameter loads and updates. If the gating network chooses k out of n experts for each example, then for a batch of b examples, each expert receives a much smaller batch of approximately kbn b examples. This causes a naive MoE implementation to become very inefficient as the number of experts increases. The solution to this shrinking batch problem is to make the original batch size as large as possible. However, batch size tends to be limited by the memory necessary to store activations between the forwards and backwards passes. We propose the following techniques for increasing the batch size:
Mixing Data Parallelism and Model Parallelism: In a conventional distributed training setting, multiple copies of the model on different devices asynchronously process distinct batches of data, and parameters are synchronized through a set of parameter servers. In our technique, these different batches run synchronously so that they can be combined for the MoE layer. We distribute the standard layers of the model and the gating network according to conventional data-parallel schemes, but keep only one shared copy of each expert. Each expert in the MoE layer receives a combined batch consisting of the relevant examples from all of the data-parallel input batches. The same set of devices function as data-parallel replicas (for the standard layers and the gating networks) and as model-parallel shards (each hosting a subset of the experts). If the model is distributed over d devices, and each device processes a batch of size b, each expert receives a batch of approximately kbd n examples. Thus, we achieve a factor of d improvement in expert batch size.
In the case of a hierarchical MoE (Section B), the primary gating network employs data parallelism, and the secondary MoEs employ model parallelism. Each secondary MoE resides on one device.
This technique allows us to increase the number of experts (and hence the number of parameters) by proportionally increasing the number of devices in the training cluster. The total batch size increases, keeping the batch size per expert constant. The memory and bandwidth requirements per device also remain constant, as do the step times, as does the amount of time necessary to process a number of training examples equal to the number of parameters in the model. It is our goal to train a trillionparameter model on a trillion-word corpus. We have not scaled our systems this far as of the writing of this paper, but it should be possible by adding more hardware.
Taking Advantage of Convolutionality: In our language models, we apply the same MoE to each time step of the previous layer. If we wait for the previous layer to finish, we can apply the MoE to all the time steps together as one big batch. Doing so increases the size of the input batch to the MoE layer by a factor of the number of unrolled time steps.
Increasing Batch Size for a Recurrent MoE: We suspect that even more powerful models may involve applying a MoE recurrently. For example, the weight matrices of a LSTM or other RNN could be replaced by a MoE. Sadly, such models break the convolutional trick from the last paragraph, since the input to the MoE at one timestep depends on the output of the MoE at the previous timestep. Gruslys et al. (2016) describe a technique for drastically reducing the number of stored activations in an unrolled RNN, at the cost of recomputing forward activations. This would allow for a large increase in batch size.
3.2 NETWORK BANDWIDTH
Another major performance concern in distributed computing is network bandwidth. Since the experts are stationary (see above) and the number of gating parameters is small, most of the communication involves sending the inputs and outputs of the experts across the network. To maintain computational efficiency, the ratio of an expert’s computation to the size of its input and output must exceed the ratio of computational to network capacity of the computing device. For GPUs, this may be thousands to one. In our experiments, we use experts with one hidden layer containing thousands of RELU-activated units. Since the weight matrices in the expert have sizes input_size×hidden_size and hidden_size× output_size, the ratio of computation to input and output is equal to the size of the hidden layer. Conveniently, we can increase computational efficiency simply by using a larger hidden layer, or more hidden layers.
4 BALANCING EXPERT UTILIZATION
We have observed that the gating network tends to converge to a state where it always produces large weights for the same few experts. This imbalance is self-reinforcing, as the favored experts are trained more rapidly and thus are selected even more by the gating network. Eigen et al. (2013) describe the same phenomenon, and use a hard constraint at the beginning of training to avoid this local minimum. Bengio et al. (2015) include a soft constraint on the batch-wise average of each gate.1
We take a soft constraint approach. We define the importance of an expert relative to a batch of training examples to be the batchwise sum of the gate values for that expert. We define an additional loss Limportance, which is added to the overall loss function for the model. This loss is equal to the square of the coefficient of variation of the set of importance values, multiplied by a hand-tuned scaling factor wimportance. This additional loss encourages all experts to have equal importance.
Importance(X) = ∑ x∈X G(x) (6)
Limportance(X) = wimportance · CV (Importance(X))2 (7) 1Bengio et al. (2015) also include two additional losses. One controls per-example sparsity, which we do not need since it is enforced by the fixed value of k. A third loss encourages diversity of gate values. In our experiments, we find that the gate values naturally diversify as the experts specialize (in a virtuous cycle), and we do not need to enforce diversity of gate values.
While this loss function can ensure equal importance, experts may still receive very different numbers of examples. For example, one expert may receive a few examples with large weights, and another may receive many examples with small weights. This can cause memory and performance problems on distributed hardware. To solve this problem, we introduce a second loss function, Lload , which ensures balanced loads. Appendix A contains the definition of this function, along with experimental results.
5 EXPERIMENTS
5.1 1 BILLION WORD LANGUAGE MODELING BENCHMARK
Dataset: This dataset, introduced by (Chelba et al., 2013) consists of shuffled unique sentences from news articles, totaling approximately 829 million words, with a vocabulary of 793,471 words.
Previous State-of-the-Art: The best previously published results (Jozefowicz et al., 2016) use models consisting of one or more stacked Long Short-Term Memory (LSTM) layers (Hochreiter & Schmidhuber, 1997; Gers et al., 2000). The number of parameters in the LSTM layers of these models vary from 2 million to 151 million. Quality increases greatly with parameter count, as do computational costs. Results for these models form the top line of Figure 2-right.
MoE Models: Our models consist of two stacked LSTM layers with a MoE layer between them (see Figure 1). We vary the sizes of the layers and the number of experts. For full details on model architecture, training regimen, additional baselines and results, see Appendix C.
Low Computation, Varied Capacity: To investigate the effects of adding capacity, we trained a series of MoE models all with roughly equal computational costs: about 8 million multiply-andadds per training example per timestep in the forwards pass, excluding the softmax layer. We call this metric (ops/timestep). We trained models with flat MoEs containing 4, 32, and 256 experts, and models with hierarchical MoEs containing 256, 1024, and 4096 experts. Each expert had about 1 million parameters. For all the MoE layers, 4 experts were active per input.
The results of these models are shown in Figure 2-left. The model with 4 always-active experts performed (unsurprisingly) similarly to the computationally-matched baseline models, while the largest of the models (4096 experts) achieved an impressive 24% lower perplexity on the test set.
Varied Computation, High Capacity: In addition to the largest model from the previous section, we trained two more MoE models with similarly high capacity (4 billion parameters), but higher computation budgets. These models had larger LSTMs, and fewer but larger experts. Details can
be found in Appendix C.2. Results of these three models form the bottom line of Figure 2-right. Table 1 compares the results of these models to the best previously-published result on this dataset . Even the fastest of these models beats the best published result (when controlling for the number of training epochs), despite requiring only 6% of the computation.
Computational Efficiency: We trained our models using TensorFlow (Abadi et al., 2016) on clusters containing 16-32 Tesla K40 GPUs. For each of our models, we determine computational efficiency in TFLOPS/GPU by dividing the number of floating point operations required to process one training batch by the observed step time and the number of GPUs in the cluster. The operation counts used here are higher than the ones we report in our ops/timestep numbers in that we include the backwards pass, we include the importance-sampling-based training of the softmax layer, and we count a multiply-and-add as two separate operations. For all of our MoE models, the floating point operations involved in the experts represent between 37% and 46% of the total.
For our baseline models wtih no MoE, observed computational efficiency ranged from 1.07-1.29 TFLOPS/GPU. For our low-computation MoE models, computation efficiency ranged from 0.74- 0.90 TFLOPS/GPU, except for the 4-expert model which did not make full use of the available parallelism. Our highest-computation MoE model was more efficient at 1.56 TFLOPS/GPU, likely due to the larger matrices. These numbers represent a significant fraction of the theoretical maximum of 4.29 TFLOPS/GPU claimed by NVIDIA. Detailed results are in Appendix C, Table 7.
5.2 100 BILLION WORD GOOGLE NEWS CORPUS
On the 1-billion-word corpus, adding additional capacity seems to produce diminishing returns as the number of parameters in the MoE layer exceeds 1 billion, as can be seen in Figure 2-left. We hypothesized that for a larger training set, even higher capacities would produce significant quality improvements.
We constructed a similar training set consisting of shuffled unique sentences from Google’s internal news corpus, totalling roughly 100 billion words. Similarly to the previous section, we tested a series of models with similar computational costs of about 8 million ops/timestep. In addition to a baseline LSTM model, we trained models augmented with MoE layers containing 32, 256, 1024,
4096, 16384, 65536, and 131072 experts. This corresponds to up to 137 billion parameters in the MoE layer. Details on architecture, training, and results are given in Appendix D.
Results: Figure 3 shows test perplexity as a function of capacity after training on 10 billion words (top line) and 100 billion words (bottom line). When training over the full 100 billion words, test perplexity improves significantly up to 65536 experts (68 billion parameters), dropping 39% lower than the computationally matched baseline, but degrades at 131072 experts, possibly a result of too much sparsity. The widening gap between the two lines demonstrates (unsurprisingly) that increased model capacity helps more on larger training sets.
Even at 65536 experts (99.994% layer sparsity), computational efficiency for the model stays at a respectable 0.72 TFLOPS/GPU.
5.3 MACHINE TRANSLATION (SINGLE LANGUAGE PAIR)
Model Architecture: Our model was a modified version of the GNMT model described in (Wu et al., 2016). To reduce computation, we decreased the number of LSTM layers in the encoder and decoder from 9 and 8 to 3 and 2 respectively. We inserted MoE layers in both the encoder (between layers 2 and 3) and the decoder (between layers 1 and 2). Each MoE layer contained up to 2048 experts each with about two million parameters, adding a total of about 8 billion parameters to the models. Further details on model architecture, testing procedure and results can be found in Appendix E.
Datasets: We benchmarked our method on the WMT’14 En→Fr and En→De corpora, whose training sets have 36M sentence pairs and 5M sentence pairs, respectively. The experimental protocols were also similar to those in (Wu et al., 2016): newstest2014 was used as the test set to compare against previous work (Luong et al., 2015a; Zhou et al., 2016; Wu et al., 2016), while the combination of newstest2012 and newstest2013 was used as the development set. We also tested the same model on Google’s Production English to French data.
Results: Tables 2, 3, and 4 show the results of our largest models, compared with published results. Our approach achieved BLEU scores of 40.56 and 26.03 on the WMT’14 En→Fr and En→De benchmarks. As our models did not use RL refinement, these results constitute significant gains of 1.34 and 1.12 BLEU score on top of the strong baselines in (Wu et al., 2016). The perplexity scores are also better.2 On the Google Production dataset, our model achieved 1.01 higher test BLEU score even after training for only one sixth of the time.
5.4 MULTILINGUAL MACHINE TRANSLATION
Dataset: (Johnson et al., 2016) train a single GNMT (Wu et al., 2016) model on a very large combined dataset of twelve language pairs. Results are somewhat worse than those for 12 separately trained single-pair GNMT models. This is not surprising, given that the twelve models have 12 times the capacity and twelve times the aggregate training of the one model. We repeat this experiment with a single MoE-augmented model. See Appendix E for details on model architecture. We train our model on the same dataset as (Johnson et al., 2016) and process the same number of training examples (about 3 billion sentence pairs). Our training time was shorter due to the lower computational budget of our model.
Results: Results for the single-pair GNMT models, the multilingual GNMT model and the multilingual MoE model are given in Table 5. The MoE model achieves 19% lower perplexity on the dev set than the multilingual GNMT model. On BLEU score, the MoE model significantly beats the multilingual GNMT model on 11 of the 12 language pairs (by as much as 5.84 points), and even beats the monolingual GNMT models on 8 of 12 language pairs. The poor performance on English → Korean seems to be a result of severe overtraining, as for the rarer language pairs a small number of real examples were highly oversampled in the training corpus.
6 CONCLUSION
This work is the first to demonstrate major wins from conditional computation in deep networks. We carefully identified the design considerations and challenges of conditional computing and addressed them with a combination of algorithmic and engineering solutions. While we focused on text, conditional computation may help in other domains as well, provided sufficiently large training sets. We look forward to seeing many novel implementations and applications of conditional computation in the years to come.
ACKNOWLEDGMENTS
We would like to thank all of the members of the Google Brain and Google Translate teams who helped us with this project, in particular Zhifeng Chen, Yonghui Wu, and Melvin Johnson. Thanks also to our anonymous ICLR reviewers for the helpful suggestions on making this paper better.
2Reported perplexities relative to the tokenization used by both our models and GNMT.
A LOAD-BALANCING LOSS
As discussed in section 4, for load-balancing purposes, we want to define an additional loss function to encourage experts to receive roughly equal numbers of training examples. Unfortunately, the number of examples received by an expert is a discrete quantity, so it can not be used in backpropagation. Instead, we define a smooth estimator Load(X) of the number of examples assigned to each expert for a batch X of inputs. The smoothness allows us to back-propagate gradients through the estimator. This is the purpose of the noise term in the gating function. We define P (x, i) as the probability that G(x)i is nonzero, given a new random choice of noise on element i, but keeping the already-sampled choices of noise on the other elements. To compute P (x, i), we note that the G(x)i is nonzero if and only if H(x)i is greater than the kth-greatest element of H(x) excluding itself. The probability works out to be:
P (x, i) = Pr ( (x ·Wg)i + StandardNormal() · Softplus((x ·Wnoise)i)
> kth_excluding(H(x), k, i) ) (8)
Where kth_excluding(v, k, i) means the kth highest component of v, excluding component i. Simplifying, we get:
P (x, i) = Φ ( (x ·Wg)i − kth_excluding(H(x), k, i)
Softplus((x ·Wnoise)i)
) (9)
Where Φ is the CDF of the standard normal distribution.
Load(X)i = ∑ x∈X P (x, i) (10)
We can now define the load loss to be the square of the coefficient of variation of the load vector, multiplied by a hand-tuned scaling factor wload.
Lload(X) = wload · CV (Load(X))2 (11)
Initial Load Imbalance: To avoid out-of-memory errors, we need to initialize the network in a state of approximately equal expert load (since the soft constraints need some time to work). To accomplish this, we initialize the matrices Wg and Wnoise to all zeros, which yields no signal and some noise.
Experiments: We trained a set of models with identical architecture (the MoE-256 model described in Appendix C), using different values of wimportance and wload. We trained each model for 10 epochs, then measured perplexity on the test set. We also measured the coefficients of variation in Importance and Load, as well as ratio of the load on the most overloaded expert to the average load. This last value is significant for load balancing purposes on distributed hardware. All of these metrics were averaged over several training batches.
Results: Results are reported in Table 6. All the combinations containing at least one the two losses led to very similar model quality, where having no loss was much worse. Models with higher values of wload had lower loads on the most overloaded expert.
B HIERACHICAL MIXTURE OF EXPERTS
If the number of experts is very large, we can reduce the branching factor by using a two-level hierarchical MoE. In a hierarchical MoE, a primary gating network chooses a sparse weighted combination of “experts", each of which is itself a secondary mixture-of-experts with its own gating network.3 If the hierarchical MoE consists of a groups of b experts each, we denote the primary gating network by Gprimary, the secondary gating networks by (G1, G2..Ga), and the expert networks by (E0,0, E0,1..Ea,b). The output of the MoE is given by:
yH = a∑ i=1 b∑ j=1 Gprimary(x)i ·Gi(x)j · Ei,j(x) (12)
Our metrics of expert utilization change to the following:
ImportanceH(X)i,j = ∑ x∈X Gprimary(x)i ·Gi(x)j (13)
LoadH(X)i,j = Loadprimary(X)i · Loadi(X(i))j
|X(i)| (14)
Loadprimary and Loadi deonte the Load functions for the primary gating network and ith secondary gating network respectively. X(i) denotes the subset of X for which Gprimary(x)i > 0.
It would seem simpler to let LoadH(X)i,j = Loadi(Xi)j , but this would not have a gradient with respect to the primary gating network, so we use the formulation above.
C 1 BILLION WORD LANGUAGE MODELING BENCHMARK - EXPERIMENTAL DETAILS
C.1 8-MILLION-OPERATIONS-PER-TIMESTEP MODELS
Model Architecture: Our model consists of five layers: a word embedding layer, a recurrent Long Short-Term Memory (LSTM) layer (Hochreiter & Schmidhuber, 1997; Gers et al., 2000), a MoE layer, a second LSTM layer, and a softmax layer. The dimensionality of the embedding layer, the number of units in each LSTM layer, and the input and output dimensionality of the MoE layer are all equal to 512. For every layer other than the softmax, we apply dropout (Zaremba et al., 2014) to the layer output, dropping each activation with probability DropProb, otherwise dividing by (1 − DropProb). After dropout, the output of the previous layer is added to the layer output. This residual connection encourages gradient flow (He et al., 2015).
MoE Layer Architecture: Each expert in the MoE layer is a feed forward network with one ReLU-activated hidden layer of size 1024 and an output layer of size 512. Thus, each expert contains [512 ∗ 1024] + [1024 ∗ 512] = 1M parameters. The output of the MoE layer is passed through a sigmoid function before dropout. We varied the number of experts between models, using ordinary MoE layers with 4, 32 and 256 experts and hierarchical MoE layers with 256, 1024 and 4096 experts. We call the resulting models MoE-4, MoE-32, MoE-256, MoE-256-h, MoE-1024-h and MoE-4096h. For the hierarchical MoE layers, the first level branching factor was 16, corresponding to the number of GPUs in our cluster. We use Noisy-Top-K Gating (see Section 2.1) with k = 4 for the ordinary MoE layers and k = 2 at each level of the hierarchical MoE layers. Thus, each example is processed by exactly 4 experts for a total of 4M ops/timestep. The two LSTM layers contribute 2M ops/timestep each for the desired total of 8M.
3We have not found the need for deeper hierarchies.
Computationally-Matched Baselines: The MoE-4 model does not employ sparsity, since all 4 experts are always used. In addition, we trained four more computationally-matched baseline models with no sparsity:
• MoE-1-Wide: The MoE layer consists of a single "expert" containing one ReLU-activated hidden layer of size 4096.
• MoE-1-Deep: The MoE layer consists of a single "expert" containing four ReLU-activated hidden layers, each with size 1024.
• 4xLSTM-512: We replace the MoE layer with two additional 512-unit LSTM layers.
• LSTM-2048-512: The model contains one 2048-unit LSTM layer (and no MoE). The output of the LSTM is projected down to 512 dimensions (Sak et al., 2014). The next timestep of the LSTM receives the projected output. This is identical to one of the models published in (Jozefowicz et al., 2016). We re-ran it to account for differences in training regimen, and obtained results very similar to the published ones.
Training: The models were trained on a cluster of 16 K40 GPUs using the synchronous method described in Section 3. Each batch consisted of a set of sentences totaling roughly 300,000 words. In the interest of time, we limited training to 10 epochs, (27,000 steps). Training took 12-16 hours for all models, except for MoE-4, which took 18 hours (since all the expert computation was performed on only 4 of 16 GPUs). We used the Adam optimizer (Kingma & Ba, 2015). The base learning rate was increased linearly for the first 1000 training steps, and decreased after that so as to be proportional to the inverse square root of the step number. The Softmax output layer was trained efficiently using importance sampling similarly to the models in (Jozefowicz et al., 2016). For each model, we performed a hyper-parmeter search to find the best dropout probability, in increments of 0.1.
To ensure balanced expert utilization we set wimportance = 0.1 and wload = 0.1, as described in Section 4 and Appendix A.
Results: We evaluate our model using perplexity on the holdout dataset, used by (Chelba et al., 2013; Jozefowicz et al., 2016). We follow the standard procedure and sum over all the words including the end of sentence symbol. Results are reported in Table 7. For each model, we report the test perplexity, the computational budget, the parameter counts, the value of DropProb, and the computational efficiency.
C.2 MORE EXPENSIVE MODELS
We ran two additional models (MoE-34M and MoE-143M) to investigate the effects of adding more computation in the presence of a large MoE layer. These models have computation budgets of 34M and 143M ops/timestep. Similar to the models above, these models use a MoE layer between two LSTM layers. The dimensionality of the embedding layer, and the input and output dimensionality of the MoE layer are set to 1024 instead of 512. For MoE-34M, the LSTM layers have 1024 units. For MoE-143M, the LSTM layers have 4096 units and an output projection of size 1024 (Sak et al., 2014). MoE-34M uses a hierarchical MoE layer with 1024 experts, each with a hidden layer of size 2048. MoE-143M uses a hierarchical MoE layer with 256 experts, each with a hidden layer of size 8192. Both models have 4B parameters in the MoE layers. We searched for the best DropProb for each model, and trained each model for 10 epochs.
The two models achieved test perplexity of 31.3 and 28.0 respectively, showing that even in the presence of a large MoE, more computation is still useful. Results are reported at the bottom of Table 7. The larger of the two models has a similar computational budget to the best published model from the literature, and training times are similar. Comparing after 10 epochs, our model has a lower test perplexity by 18%.
D 100 BILLION WORD GOOGLE NEWS CORPUS - EXPERIMENTAL DETAILS
Model Architecture: The models are similar in structure to the 8-million-operations-per-timestep models described in the previous section. We vary the number of experts between models, using an ordinary MoE layer with 32 experts and hierarchical MoE layers with 256, 1024, 4096, 16384, 65536 and 131072 experts. For the hierarchical MoE layers, the first level branching factors are 32, 32, 64, 128, 256 and 256, respectively.
Training: Models are trained on a cluster of 32 Tesla K40 GPUs, except for the last two models, which are trained on clusters of 64 and 128 GPUs so as to have enough memory for all the parameters. For all models, training batch sizes are approximately 2.5 million words. Models are trained once-through over about 100 billion words.
We implement several memory optimizations in order to fit up to 1 billion parameters per GPU. First, we do not store the activations of the hidden layers of the experts, but instead recompute them on the backwards pass. Secondly, we modify the optimizer on the expert parameters to require less auxiliary storage:
The Adam optimizer (Kingma & Ba, 2015) keeps first and second moment estimates of the perparameter gradients. This triples the required memory. To avoid keeping a first-moment estimator, we set β1 = 0. To reduce the size of the second moment estimator, we replace it with a factored approximation. For a matrix of parameters, instead of maintaining a full matrix of second-moment estimators, we maintain vectors of row-wise and column-wise averages of that matrix. At each step, the matrix of estimators is taken to be the outer product of those two vectors divided by the mean of either one. This technique could similarly be applied to Adagrad (Duchi et al., 2010).
Results: We evaluate our model using perplexity on a holdout dataset. Results are reported in Table 8. Perplexity after 100 billion training words is 39% lower for the 68-billion-parameter MoE
model than for the baseline model. It is notable that the measured computational efficiency of the largest model (0.30 TFLOPS/GPU) is very low compared to the other models. This is likely a result of the fact that, for purposes of comparison to the other models, we did not increase the training batch size proportionally to the number of GPUs. For comparison, we include results for a computationally matched baseline model consisting of 4 LSTMs, and for an unpruned 5-gram model with Kneser-Ney smoothing (Kneser & Ney, 1995).4
E MACHINE TRANSLATION - EXPERIMENTAL DETAILS
Model Architecture for Single Language Pair MoE Models: Our model is a modified version of the GNMT model described in (Wu et al., 2016). To reduce computation, we decrease the number of LSTM layers in the encoder and decoder from 9 and 8 to 3 and 2 respectively. We insert MoE layers in both the encoder (between layers 2 and 3) and the decoder (between layers 1 and 2). We use an attention mechanism between the encoder and decoder, with the first decoder LSTM receiving output from and providing input for the attention 5. All of the layers in our model have input and output dimensionality of 512. Our LSTM layers have 2048 hidden units, with a 512-dimensional output projection. We add residual connections around all LSTM and MoE layers to encourage gradient flow (He et al., 2015). Similar to GNMT, to effectively deal with rare words, we used subword units (also known as “wordpieces") (Schuster & Nakajima, 2012) for inputs and outputs in our system.
We use a shared source and target vocabulary of 32K wordpieces. We also used the same beam search technique as proposed in (Wu et al., 2016).
We train models with different numbers of experts in the MoE layers. In addition to a baseline model with no MoE layers, we train models with flat MoE layers containing 32 experts, and models with hierarchical MoE layers containing 512 and 2048 experts. The flat MoE layers use k = 4 and the hierarchical MoE models use k = 2 at each level of the gating network. Thus, each input is processed by exactly 4 experts in each MoE layer. Each expert in the MoE layer is a feed forward network with one hidden layer of size 2048 and ReLU activation. Thus, each expert contains [512 ∗ 2048] + [2048 ∗ 512] = 2M parameters. The output of the MoE layer is passed through a sigmoid function. We use the strictly-balanced gating function described in Appendix F.
Model Architecture for Multilingual MoE Model: We used the same model architecture as for the single-language-pair models, with the following exceptions: We used noisy-top-k gating as described in Section 2.1, not the scheme from Appendix F. The MoE layers in the encoder and decoder are non-hierarchical MoEs with n = 512 experts, and k = 2. Each expert has a larger hidden layer of size 8192. This doubles the amount of computation in the MoE layers, raising the computational budget of the entire model from 85M to 102M ops/timestep.
Training: We trained our networks using the Adam optimizer (Kingma & Ba, 2015). The base learning rate was increased linearly for the first 2000 training steps, held constant for an additional 8000 steps, and decreased after that so as to be proportional to the inverse square root of the step number. For the single-language-pair models, similarly to (Wu et al., 2016), we applied dropout (Zaremba et al., 2014) to the output of all embedding, LSTM and MoE layers, using DropProb = 0.4. Training was done synchronously on a cluster of up to 64 GPUs as described in section 3. Each training batch consisted of a set of sentence pairs containing roughly 16000 words per GPU.
To ensure balanced expert utilization we set wimportance = 0.01 and wload = 0.01, as described in Section 4 and Appendix A.
Metrics: We evaluated our models using the perplexity and the standard BLEU score metric. We reported tokenized BLEU score as computed by the multi-bleu.pl script, downloaded from the public implementation of Moses (on Github), which was also used in (Luong et al., 2015a).
4While the original size of the corpus was 130 billion words, the neural models were trained for a maximum of 100 billion words. The reported Kneser-Ney 5-gram models were trained over 13 billion and 130 billion words respectively, giving them a slight advantage over the other reported results.
5For performance reasons, we use a slightly different attention function from the one described in (Wu et al., 2016) - See Appendix G
Results: Tables 2, 3 and 4 in Section 5.3 show comparisons of our results to other published methods. Figure 4 shows test perplexity as a function of number of words in the (training data’s) source sentences processed for models with different numbers of experts. As can be seen from the Figure, as we increased the number of experts to approach 2048, the test perplexity of our model continued to improve.
We found that the experts indeed become highly specialized by syntax and/or semantics, as can be seen in Table 9. For example, one expert is used when the indefinite article “a" introduces the direct object in a verb phrase indicating importance or leadership.
F STRICTLY BALANCED GATING
Due to some peculiarities in our infrastructure which have since been fixed, at the time we ran some of the machine translation experiments, our models ran faster if every expert received exactly the same batch size. To accommodate this, we used a different gating function which we describe below.
Recall that we define the softmax gating function to be:
Gσ(x) = Softmax(x ·Wg) (15)
Sparse Gating (alternate formulation): To obtain a sparse gating vector, we multiply Gσ(x) component-wise with a sparse mask M(Gσ(x)) and normalize the output. The mask itself is a function of Gσ(x) and specifies which experts are assigned to each input example:
G(x)i = Gσ(x)iM(Gσ(x))i∑n j=1Gσ(x)jM(Gσ(x))j
(16)
Top-K Mask: To implement top-k gating in this formulation, we would let M(v) = TopK(v, k), where:
TopK(v, k)i = { 1 if vi is in the top k elements of v. 0 otherwise.
(17)
Batchwise Mask: To force each expert to receive the exact same number of examples, we introduce an alternative mask function, Mbatchwise(X,m), which operates over batches of input vectors. Instead of keeping the top k values per example, we keep the top m values per expert across the training batch, where m = k|X|n , so that each example is sent to an average of k experts.
Mbatchwise(X,m)j,i = { 1 if Xj,i is in the top m values for to expert i 0 otherwise
(18)
As our experiments suggest and also observed in (Ioffe & Szegedy, 2015), using a batchwise function during training (such as Mbatchwise) requires modifications to the inference when we may not have a large batch of examples. Our solution to this is to train a vector T of per-expert threshold values to approximate the effects of the batchwise mask. We use the following mask at inference time:
Mthreshold(x, T )i = { 1 if xi > Ti 0 otherwise
(19)
To learn the threshold values, we apply an additional loss at training time which is minimized when the batchwise mask and the threshold mask are identical.
Lbatchwise(X,T,m) = |X|∑ j=1 n∑ i=1 (Mthreshold(x, T )i −Mbatchwise(X,m)j,i)(Xj,i − Ti) (20)
G ATTENTION FUNCTION
The attention mechanism described in GNMT (Wu et al., 2016) involves a learned “Attention Function" A(xi, yj) which takes a “source vector" xi and a “target vector" yj , and must be computed for every source time step i and target time step j. In GNMT, the attention function is implemented as a feed forward neural network with a hidden layer of size n. It can be expressed as:
AGNMT (xi, yj) = n∑ d=1 Vdtanh((xiU)d + (yjW )d) (21)
Where U and W are trainable weight matrices and V is a trainable weight vector.
For performance reasons, in our models, we used a slightly different attention function:
A(xi, yj) = n∑ d=1 Vdtanh((xiU)d)tanh((yjW )d) (22)
With our attention function, we can simultaneously compute the attention function on multiple source time steps and multiple target time steps using optimized matrix multiplications. We found little difference in quality between the two functions.

This paper proposes a method for significantly increasing the number of parameters in a single layer while keeping computation in par with (or even less than) current SOTA models. The idea is based on using a large mixture of experts (MoE) (i.e. small networks), where only a few of them are adaptively activated via a gating network. While the idea seems intuitive, the main novelty in the paper is in designing the gating network which is encouraged to achieve two objectives: utilizing all available experts (aka importance), and distributing computation fairly across them (aka load).
Additionally, the paper introduces two techniques for increasing the batch-size passed to each expert, and hence maximizing parallelization in GPUs.
Experiments applying the proposed approach on RNNs in language modelling task show that it can beat SOTA results with significantly less computation, which is a result of selectively using much more parameters. Results on machine translation show that a model with more than 30x number of parameters can beat SOTA while incurring half of the effective computation.
I have the several comments on the paper:
- I believe that the authors can do a better job in their presentation. The paper currently is at 11 pages (which is too long in my opinion), but I find that Section 3.2 (the crux of the paper) needs better motivation and intuitive explanation. For example, equation 8 deserves more description than currently devoted to it. Additional space can be easily regained by moving details in the experiments section (e.g. architecture and training details) to the appendix for the curious readers. Experiment section can be better organized by finishing on experiment completely before moving to the other one. There are also some glitches in the writing, e.g. the end of Section 3.1.
- The paper is missing some important references in conditional computation (e.g.