Routing to expert subnetworks in mixed-expert neural networks.

Abstract

Methods, systems, and apparatus, such as computer programs encoded on a computer storage medium, for performing machine learning tasks on network inputs to generate network outputs. In one aspect, one of the systems includes a neural network configured to perform the machine learning task, the neural network including one or more expert neural network blocks, each including a router that performs expert selection routing among a plurality of expert neural networks.

Description

[Technical field]

[0001] CROSS-REFERENCE TO RELATED APPLICATIONS This application claims the benefit of priority to U.S. Provisional Application No. 63 / 304,507, filed January 28, 2022, the disclosure of which is incorporated herein by reference in its entirety. [Background technology]

[0002] This specification relates to using neural networks to perform machine learning tasks on network inputs.

[0003] A neural network is a machine learning model that uses one or more layers of nonlinear units to predict an output for a received input. Some neural networks include one or more hidden layers in addition to an output layer. The output of each hidden layer is used as the input to the next layer in the network, i.e., the next hidden layer or the output layer. Each layer of the network generates an output from the received input according to the current values ​​of a respective set of parameters. Summary of the Invention

[0004] Described herein is a system implemented as a computer program on one or more computers at one or more locations that is configured to process network inputs using a neural network and generate network outputs that characterize the network inputs. The neural network includes a sequence of one or more network blocks, each configured to process block inputs that include network inputs or intermediate representations of the network inputs, and generate block outputs.

[0005] For example, a first network block in a sequence of network blocks may process a network input and generate a block output that is an intermediate representation of the network input. As another example, an embedding sub-network may process a network input to generate an embedding of the network input provided as an input to the first network block in the sequence, and process the embedding to generate an intermediate representation of the network input. Each subsequent network block may then process the block output of the previous network block in the sequence. In some implementations, the network output of the neural network is the block output of the last network block in the sequence. In some other implementations, the block output of the last network block in the sequence is further processed using one or more output neural network layers to generate a network output of the neural network.

[0006] A sequence of network blocks may include one or more "expert" network blocks, each of which includes multiple different expert sub-networks, each configured to process a respective sub-input determined from a block input to the expert network block, each sub-input of a block input including a respective different subset of the elements of the block input.

[0007] Unlike conventional approaches, the expert network block routes sub-inputs to expert subnetworks using "sub-network selection" or "expert selection" routing, i.e., for each expert subnetwork, the neural network block independently selects the set of elements of the block inputs that are to be processed by the expert subnetwork.

[0008] Particular embodiments of the subject matter described in this specification can be implemented to realize one or more of the following advantages.

[0009] By implementing a neural network with a network block that includes multiple expert subnetworks, each configured to process a subset of the inputs to the network block, the system can increase the capacity of the neural network without increasing the computational resources required to run the neural network at inference time. That is, by selectively activating only a subset of the neural network's parameters for each element of the network input, the system can significantly improve the time and computational efficiency of the neural network compared to other neural networks with the same number of parameters. Introducing this sparsity allows the neural network to include many more network parameters than was previously possible, since only a subset of the parameters is used to process any given input.

[0010] Some existing systems that implement neural network blocks with multiple expert subnetworks use "token selection" routing, i.e., the neural network block independently selects, for each element of the block input, a set of expert subnetworks to process the element. However, these existing systems may suffer from load imbalance, where some expert subnetworks process most or all of the elements of the block input while other expert subnetworks process few or none of the elements of the block input. Such load imbalance may result in suboptimal training, since some of the network parameters (corresponding to underutilized expert subnetworks) do not receive meaningful updates during training and therefore do not learn to extract useful information. Furthermore, existing systems that use token selection routing may allocate the same amount of computational resources to each element of the network input, ignoring the relative importance of different elements, further reducing the computational efficiency of the system.

[0011] Using the techniques described herein, the system can implement neural network blocks with multiple expert subnetworks using "subnetwork selection" or "expert selection" routing. That is, the neural network block independently selects for each expert subnetwork the set of elements of the block input to be processed by the expert subnetwork. Subnetwork selection routing can ensure that the network block is fully load balanced, for example, by selecting the same number k of elements to be processed by each expert subnetwork. Thus, the computational and time efficiency of training a neural network can be significantly improved, so that all parameters of the neural network receive meaningful updates for each network input. In some implementations, the neural network training system can reduce the time and computational resources required to achieve a given performance using subnetwork selection routing by more than two times compared to token selection routing. Furthermore, subnetwork selection routing can allow the network block to more flexibly allocate computational resources to each element, for example, by routing relatively important elements to more expert subnetworks than relatively unimportant elements.

[0012] In particular, in some implementations where a system runs different expert sub-networks on different devices, the techniques described herein enable the system to load balance network inputs to a neural network more efficiently across devices compared to existing techniques (e.g., compared to systems implementing token selection routing). Poor load balancing suffered by existing systems with token selection routing can impair inference performance (e.g., by reducing computational and / or memory efficiency, or by increasing the time required to generate a network output) because different devices may have significantly different loads running different expert sub-networks, and thus some devices may be underutilized while others may be overloaded. With sub-network selection routing as described herein, perfect load balancing can be "baked in" at inference time, and the system can enjoy significant performance improvements across multiple devices (e.g., improved computational and / or memory efficiency, or reduced time required to generate a network output) because multiple devices each have a similar or equivalent amount of operations to perform. Thus, compared to other existing approaches, the described approach results in a system that can process inputs at a higher throughput compared to conventional approaches by being optimized for distributed hardware implementation.

[0013] In some implementations described herein, each expert subnetwork can be configured via training to process different types of network inputs, allowing the expert subnetworks to "specialize," further improving the efficiency and performance of the neural network. The use of "subnetwork selection" routing can further improve the specialization of the expert subnetworks compared to "token selection" routing, because each routing subnetwork corresponding to each expert subnetwork can learn via training to select specific types of elements to be processed by the expert subnetwork.

[0014] The details of one or more embodiments of the subject matter of this specification are set forth in the accompanying drawings and the description below. Other features, aspects, and advantages of the subject matter will become apparent from the description, drawings, and claims. [Brief description of the drawings]

[0015] [Figure 1] FIG. 1 shows an example of a neural network system. [Diagram 2] FIG. 2 illustrates the operation of the expert block using expert selection routing as compared to token selection routing. [Diagram 3] FIG. 3 is a flow diagram of an exemplary process for processing network inputs and generating network outputs. DETAILED DESCRIPTION OF THE PREFERRED EMBODIMENTS

[0016] Like reference symbols and designations in the various drawings indicate like elements.

[0017] Described herein is a system, implemented as a computer program on one or more computers at one or more locations, that performs a machine learning task on a network input to generate a network output for the machine learning task.

[0018] The machine learning task can be any machine learning task that operates on an input sequence, i.e., a network input that is a set of elements, and produces a network output for the network input.

[0019] Below are some examples of machine learning tasks the system can be configured to perform:

[0020] As another example, the task may be a speech processing task. For example, if the input to the neural network is a sequence representing an utterance, the output may be a classification output that classifies the utterance into one or more categories from a set of categories. For example, if the input to the neural network is a sequence representing an utterance, the output generated by the neural network may indicate whether a particular word or phrase (a "hot word") was spoken within the utterance. As another example, if the input to the neural network is a sequence representing an utterance, the output generated by the neural network may identify the natural language in which the utterance was spoken. It will be appreciated that in the case of a speech processing task, the input to the neural network may include audio data (e.g., an audio signal), for example in the form of a sequence of audio data frames, and the audio data may be processed to perform the speech processing task.

[0021] As another example, the task can be a natural language processing or understanding task, such as an entailment task, a paraphrase task, a text similarity task, a sentiment task, a sentence completion task, a grammar task, etc., that operates on a text sequence in some natural language and generates a classification output that classifies the text into one or more categories from a set of categories.

[0022] As another example, the task can be a health prediction task, where the input is a sequence derived from the patient's electronic health record data and the output is a prediction related to the patient's future health, such as a predicted treatment to be prescribed for the patient, the likelihood that the patient will experience an adverse health event, or a predicted diagnosis for the patient.

[0023] As another example, a task can be an agent control task, where the inputs are sequences of observations or other data characterizing an environmental state, and the outputs define actions to be performed by the agent in response to the most recent data in the sequence. An agent can be, for example, a real-world or simulated robot, a control system for an industrial facility, or a control system that controls different kinds of agents.

[0024] As another example, the task can be a genomics task, where the input is a sequence representing a fragment of a DNA sequence or other molecular sequence, and the output is either an embedding of the fragment for use in a downstream task, e.g., by using unsupervised learning techniques on a dataset of DNA sequence fragments, or the output of a downstream task. Examples of downstream tasks include promoter site prediction, methylation analysis, predicting the functional impact of non-coding variants, etc.

[0025] As another example, the task can be a computer vision task, the input is an image or a point cloud, and the output is a computer vision output of the image or point cloud, It will be appreciated that an image may include pixel data and may be processed to perform a computer vision task.

[0026] For example, the computer vision task can be a classification task that requires the generation of a classification output. The classification output typically includes a respective score corresponding to each of a number of categories. The category score indicates the likelihood that the image belongs to that category. In some cases, the category may be a class of objects (e.g., dogs, cats, people, etc.), and an image may belong to a category if it depicts an object included in the object class that corresponds to the category. In some cases, the category may represent global image characteristics (e.g., whether the image depicts a daytime or nighttime scene, or whether the image depicts a summer or winter scene), and an image may belong to a category if it has global characteristics that correspond to the category.

[0027] As another example, the computer vision task can be an object detection task, in which the output generated by a neural network identifies locations, e.g., bounding boxes or other regions, within an input image where a particular type of object is depicted.

[0028] As another example, the computer vision task can be an instance segmentation task, in which the output generated by the neural network identifies, for each pixel of the image that belongs to a particular object type, the object instance to which the pixel corresponds.

[0029] As another example, the computer vision task can be a semantic segmentation task, in which the output generated by the neural network identifies, for each pixel of an image, to which of multiple categories the pixel belongs.

[0030] As another example, the computer vision task can be a depth prediction task, in which the output generated by the neural network identifies, for each pixel of an image, a predicted depth of the scene at the pixel.

[0031] As another example, the computer vision task can be a surface normal prediction task, in which the output generated by the neural network identifies, for each pixel of the image, a predicted surface normal of the scene at the pixel.

[0032] When the input is an image or a point cloud, the neural network may include an embedding sub-network that generates a respective embedding for each of multiple patches of the image or point cloud, and the input to a first block of the neural network may be a sequence including the respective embeddings (and optionally, one or more additional embeddings at predetermined locations that are subsequently used, e.g., to generate the output), each patch including intensity values ​​of pixels in a different region of the input image.

[0033] In some implementations, the neural network includes both a computer vision neural network and a text processing neural network, since the task is a multimodal task that requires processing of both text inputs and image inputs. That is, the target output generated by the computer vision neural network for a given image depends on one or more outputs generated by the text processing neural network for one or more corresponding text inputs, and vice versa. Examples of such tasks include open vocabulary image classification, open vocabulary object detection, image captioning, text-based image retrieval, image-based search, etc.

[0034] In some cases, the machine learning task is a combination of multiple individual machine learning tasks. That is, the system is configured to perform multiple different individual machine learning tasks, e.g., two or more of the machine learning tasks listed above. For example, the system can be configured to perform multiple individual natural language understanding tasks with a network input that includes an identifier of the individual natural language understanding task to be performed on the network input.

[0035] 1 illustrates an example of a neural network system 100. Neural network system 100 is an example of a system implemented as a computer program on one or more computers at one or more locations, which may implement the systems, components, and techniques described below.

[0036] The system 100 is a system that processes network inputs 102 using a neural network 110 to generate network outputs 112 that characterize the network inputs 102 for a machine learning task, such as one of the tasks listed above.

[0037] Neural network 110 includes a sequence of one or more network blocks 120, each configured to process block inputs that include network inputs or intermediate representations of network inputs, and generate block outputs.

[0038] As used herein, a "network block" is a collection of one or more neural network layers that receives an input (a "block input"), processes the input, and produces an output (a "block output").

[0039] For example, the first network block in a sequence of network blocks 120 may process the network input 102, or an embedding of the network input generated by an embedding sub-network, and generate a block output that is an intermediate representation of the network input. Each subsequent network block 120 may then process the block output of the previous network block in the sequence.

[0040] In some implementations, the network output 112 of the neural network 110 is the block output of the last network block 120 in the sequence.

[0041] In some other implementations, the block output of the last network block 120 in the sequence is further processed using one or more output neural network layers to generate the network output 112 of the neural network 110.

[0042] The sequence of network blocks may include one or more “expert” network blocks 130. Each expert network block 130 includes multiple different expert sub-networks 132, each configured to process a respective sub-input determined from the block input to the expert network block 132.

[0043] Each sub-input of a block input contains a different subset of the elements of the block input, hi some implementations, each sub-input contains exactly one of the elements of the block input.

[0044] For each expert sub-network 132 of the expert network block 130 , an expert selection router 134 within the expert network block 130 is configured to select one or more of the sub-inputs to be processed by the expert sub-network 132 .

[0045] For example, the expert selection router 134 may be configured to generate a respective score for each sub-input for each expert sub-network 132 and then select the sub-input with the highest corresponding score for processing by the expert sub-network 132.

[0046] More specifically, the router 134 performs "expert selection" routing by independently selecting the sub-inputs for each expert subnetwork 132 to process, as opposed to "token selection" routing or "sub-input" selection routing.

[0047] Expert selection routing is described in more detail below with reference to FIG.

[0048] Each expert subnetwork 132 can then process only the sub-input(s) selected for the subnetwork 132 and generate a respective sub-output for each processed sub-input. The expert subnetwork 132 does not process any other sub-inputs that were not selected for the expert subnetwork. In other words, each expert subnetwork 132 is configured to process only a proper subset of the elements of the block inputs.

[0049] After each expert sub-network 132 is executed, the expert network block 130 can determine the combined sub-outputs that correspond to each sub-input.

[0050] In particular, for each sub-input, the expert network block 130 can combine each sub-output generated by the respective expert sub-network in response to processing the sub-input.

[0051] For sub-inputs that have not been processed by any of the expert sub-networks 132, the combined sub-outputs of the sub-inputs may be the same as the sub-inputs.

[0052] The expert network block 130 can then combine each combined sub-output to generate a block output of the expert network block. For example, the expert network block 130 can concatenate the combined sub-outputs in the same configuration (e.g., in the same order) as the corresponding sub-inputs of the block input.

[0053] The sub-inputs may be any suitable subset of the elements of the block inputs. For example, if the neural network is configured to process an input sequence (e.g., an input sequence representing images, text data, or audio data), each block input may be an intermediate sequence that is an intermediate representation of the input sequence, and the sub-inputs may be subsequences of the intermediate sequence.

[0054] In some implementations, each sub-input is the same size, i.e., contains the same number of elements. For example, each sub-input can be a different one of the elements of the block input. In some other implementations, different sub-inputs can be different sizes, i.e., contain different numbers of elements.

[0055] In some implementations, each element of a block input is in exactly one sub-input, while in some other implementations, the elements of some or all of the block inputs may be in multiple different sub-inputs.

[0056] The operations performed by expert network block 130 are described in more detail below with reference to FIGS.

[0057] In some implementations, the sequence of network blocks 120 includes expert network blocks 130 interspersed among other types of network blocks, e.g., self-attention network blocks that apply self-attention, which do not include routers and expert neural networks, i.e., do not perform conditional calculations and use all parameters of the network blocks for all inputs to the network blocks. As a particular example, the sequence of network blocks 120 can alternate between expert network blocks and self-attention network blocks. As another particular example, a sequence of network blocks 120 can include a self-attention network block, a feedforward network block that includes a single neural network with the same architecture as the expert neural network 132 and processes all sub-inputs of the block input to the feedforward block, and an expert network block 130. For example, every other self-attention network block in the sequence can be immediately followed by an expert network block 130, and the rest of the self-attention network blocks can be followed by a feedforward network block.

[0058] Each self-attention network block is configured to process the block input using one or more self-attention neural network layers.

[0059] A self-attention neural network layer receives as input a sequence of input elements and applies an attention mechanism to the sequence of input elements to generate a sequence of layer output elements. In particular, for each input element, the self-attention neural network layer applies an attention mechanism to the sequence of input elements using one or more queries derived from the input elements to generate a respective output element. Some self-attention neural network layers are multi-head self-attention neural network layers. A multi-head self-attention neural network layer applies h different attention mechanisms in parallel to generate respective sequences of output elements, and then combines the multiple sequences of output elements to generate a final sequence of output elements.

[0060] Self-attention is described in more detail below.

[0061] Each expert network block 130 can be implemented such that the expert sub-networks 132 of the expert network block 130 run in parallel for a given block input, thereby improving system efficiency. For example, the expert network block 130 can be implemented in a parallel processing device, e.g., a GPU or a TPU, capable of executing the expert sub-networks in different threads. As another example, at least some of the expert sub-networks 132 of the expert network block 130 can be implemented in different devices, e.g., different devices that are communicatively connected and provide the sub-outputs generated by the respective expert sub-networks to a single device for combining to generate respective combined sub-outputs.

[0062] Thus, the network architecture of the neural network, including the expert network block 130 with multiple expert sub-networks 132, is optimized for efficient execution of the neural network. Such a network architecture allows for parallelizing the operation of the neural network for rapid and low-cost execution, e.g., by parallelizing the operation of each expert sub-network across multiple devices. Thus, the neural network can be implemented with dedicated parallel processing hardware, e.g., a network of multiple parallel processing devices each running a respective expert sub-network of the neural network.

[0063] As described below, by implementing expert selection routing, system 100 optimizes the parallelism of the processing performed for each network input.

[0064] FIG. 2 illustrates the operations performed by routers of expert network block 130 under token selection routing 200 used by some conventional systems, and under expert selection routing 250 as described herein.

[0065] In particular, in the example of Figure 2, the expert network block 130 is processing an intermediate representation of the sequence "We like to play soccer on that field", which includes respective elements ("tokens") 202-216 corresponding to each word in the sequence, and the expert network block 130 includes four expert neural networks 132. In the example of Figure 2, each expert 132 is a feed-forward neural network (FFN). Moreover, in the example of Figure 2, each sub-input includes one different element.

[0066] In token selection routing 210, the expert block 130 processes each token 202-216 independently. As shown in Figure 2, for token 1 202, the token selection router 230 generates a respective score for each expert network 132 and routes token 1 to "FFN2" since FFN2 has the highest score. For token 2 204, the token selection router 230 generates a respective score for each expert network 132 and routes token 2 to "FFN1" since FFN1 has the highest score.

[0067] Thus, when routing token 2, the token selection router 230 does not consider where token 1 was routed, or more generally, the scores assigned to FFN1 of other tokens. As a result, the system may suffer from load imbalance, with some expert sub-networks processing most or all of the elements of the block input, while other expert sub-networks processing few or none of the elements of the block input. Such load imbalance may result in suboptimal training, since some of the network parameters (corresponding to underutilized expert sub-networks) do not receive meaningful updates during training and therefore do not learn to extract useful information. Moreover, the system may allocate the same amount of computational resources to each element of the network input, ignoring the relative importance of different elements, further reducing the computational efficiency of the system.

[0068] In expert selection routing 250, the expert selection router 260 instead generates a respective score distribution for each expert neural network 132. The score distribution includes respective scores for each token 202-216. The expert selection router 260 then selects the top k scoring tokens for each expert neural network 132 according to their respective scores in the score distribution of the expert neural network 132.

[0069] For example, when k is equal to 4, the expert selection router 260 may route representations of the tokens "we" "like" "to" "do" to FFN1 and may route representations of the tokens "we" "like" "soccer" "field" to FFN2. Thus, the expert selection router 260 may route the same token to multiple different expert neural networks 132 if the given token is within the top k of the score distributions of the multiple different expert neural networks 132.

[0070] As can be seen from FIG. 2, expert selection routing can ensure that the expert block is fully load-balanced, for example by selecting the same number k of elements to be processed by each expert sub-network. Thus, the computational and time efficiency of training a neural network can be significantly improved, so that all parameters of the neural network receive meaningful updates for each network input. In some implementations, the neural network training system can reduce the time and computational resources required to achieve a given performance using sub-network selection routing by more than two times compared to token selection routing. Furthermore, sub-network selection routing can allow the network block to more flexibly allocate computational resources to each element, for example by routing relatively important elements to more expert sub-networks than relatively unimportant elements.

[0071] Processing inputs using expert selection routing is described in more detail below with reference to FIG.

[0072] 3 is a flow diagram 300 of an exemplary process for processing block inputs using an expert block. For convenience, the process 300 is described as being performed by one or more computer systems at one or more locations. For example, an expert block included in a neural network system, such as one of the expert blocks 130 included in the neural network system 100 of FIG. 1, when appropriately programmed, can perform the process 200.

[0073] The expert block obtains a block input that represents an intermediate representation of the network input (step 302).

[0074] The expert block determines 304 a number of sub-inputs from the block input, each of which includes a different subset of the elements of the block input. In some implementations, each sub-input includes exactly one of the elements of the block input. That is, each sub-input is a different one of the elements of the block input.

[0075] More generally, in some implementations, each sub-input is the same size, i.e., contains the same number of elements, while in some other implementations, different sub-inputs can be different sizes, i.e., contain different numbers of elements.

[0076] In some implementations, each element of a block input is in exactly one sub-input, while in some other implementations, the elements of some or all of the block inputs may be in multiple different sub-inputs.

[0077] Next, the expert block executes steps 306 to 310 for each of the multiple expert sub-networks of the expert network block.

[0078] The expert block processes multiple sub-inputs and generates a respective score for each sub-input (step 306).

[0079] For each expert sub-network, the expert network block, i.e., the expert selection router within the block, may determine the scores of the sub-inputs in any suitable manner. For example, for each expert sub-network, the expert network block may use one or more neural network layers, e.g., one or more feedforward neural network layers, to process each sub-input and generate a respective score. As a specific example, the expert network block may calculate:

number

number

number

[0080] As another example, for each expert sub-network, the expert network block may process each sub-input using one or more convolutional neural network layers and / or one or more self-attention layers to generate a respective score, with self-attention being described in more detail below.

[0081] The system selects one or more of the sub-inputs according to their respective scores (step 308).

[0082] For each expert sub-network, after the expert network block generates a respective score for each sub-input, the expert network block, i.e., the expert selection router within the block, can select and process the sub-input with the highest score k for the expert sub-network. For example, the expert network block can calculate: G,I = TopK(S T ,k) P=Onehot(I) In the formula, TopK(S T ,k) is S T For each row of, choose the kth largest entry,

number

number

number

[0083] In some implementations, the expert network block can enforce that each sub-input is selected by at most b different expert sub-networks. For example, the expert network block can solve the entropy regularized linear programming problem:

number

number

[0084] Adding a small entropy term, i.e., (λH(A)), can give near-integer solutions, and if an upper bound is imposed during processing, the expert network block can be more efficiently executed by a parallel processing device, e.g., a graphics processing unit (GPU) or a tensor processing unit (TPU). Deploying neural networks on one or more parallel processing hardware is described in more detail below.

[0085] As a specific example, the expert network block can use Dijkstra's algorithm to solve the above problem.

[0086] After computing the matrix A, the network block can generate the matrices I and G by determining G,I=TopK(A,k) as above.

[0087] Optionally, during training, for each expert sub-network, the expert network block randomly samples a noise value for each sub-input and adds the noise value to the corresponding score before determining the highest score k of the expert sub-network. Alternatively, or in addition, the expert network block can apply a non-linear activation function, such as a softmax function, a Tanh function, or a ReLU function, to the scores before determining the highest score k. Performing either or both of these operations can aid in exploration during training.

[0088] In some implementations, each expert sub-network selects the same number of sub-inputs. That is, each expert sub-network can select the k sub-inputs with the highest corresponding scores, where k is the same for all expert sub-networks. As a specific example, the same k of sub-inputs processed by each expert sub-network can be equal to:

number

[0089] In some other implementations, some expert sub-networks select a different number of sub-inputs. For example, an expert sub-network can select a sub-input if its score meets a threshold.

[0090] For each selected sub-input, the expert block processes the selected sub-input using the expert sub-network to generate a respective sub-output (step 310).

[0091] In other words, after selecting the k sub-inputs, each expert sub-network can process the k sub-inputs using one or more neural network layers to generate respective sub-outputs. For example, each expert sub-network can include one or more feedforward neural network layers, one or more convolutional neural network layers, one or more recurrent neural network layers, and / or one or more self-attention neural network layers.

[0092] As a specific example, each expert sub-network i can generate each sub-output for each sub-input by computing:

number

number

number

[0093] Once the expert sub-networks have completed processing, for each of the multiple sub-inputs, the system processes the sub-outputs corresponding to the sub-inputs generated by the respective expert sub-networks to generate combined sub-outputs for the sub-inputs (step 312). In particular, for each sub-input, the expert network block can combine each sub-output generated by the respective expert sub-networks according to having processed the sub-input. For sub-inputs that were not processed by any of the expert sub-networks, the combined sub-outputs of the sub-inputs can be the same as the sub-input.

[0094] In some implementations, the expert network block can generate a combined sub-output by computing a sum of the sub-outputs, for example, the expert network block can compute a weighted sum of the sub-outputs, where each sub-output is weighted by the score of the corresponding sub-input of the expert sub-network that generated the sub-output.

[0095] In particular, the expert network block can compute:

number

number

[0096] In some other implementations, the expert network block combines the sub-outputs for each sub-input by processing the sub-outputs using one or more neural network layers, such as one or more self-attention layers, one or more convolutional neural network layers, and / or one or more recurrent neural network layers.

[0097] The system generates a block output by combining each combined sub-output for the multiple sub-inputs (step 314). For example, the expert network block may concatenate the combined sub-outputs in the same configuration (e.g., in the same order) as the corresponding sub-inputs of the block input.

[0098] Optionally, as part of combining the sub-outputs, the system may apply one or more additional operations to the concatenation of the combined sub-outputs, e.g., the concatenation may be processed by one or more feedforward layers, skip connections, or normalization operations, e.g., layer normalization, to provide a block output.

[0099] The expert network block can be implemented such that the expert sub-networks of the expert network block run in parallel for a given block input, thereby improving system efficiency. For example, the expert network block can be implemented in a parallel processing device, e.g., a GPU or a TPU, that can execute the expert sub-networks in different threads. As another example, at least some of the expert sub-networks of the expert network block can be implemented in different devices, e.g., different devices that are communicatively connected and provide sub-outputs generated by the respective expert sub-networks to a single device for combination to generate respective combined sub-outputs.

[0100] Thus, the network architecture of a neural network including an expert network block with multiple expert sub-networks is optimized for efficient execution of the neural network. Such a network architecture allows for parallelizing the operation of the neural network for rapid and low-cost execution, e.g., by parallelizing the operation of each expert sub-network across multiple devices. Thus, the neural network can be implemented with dedicated parallel processing hardware, e.g., a network of multiple parallel processing devices each executing a respective expert sub-network of the neural network.

[0101] In some implementations, before assigning the sub-inputs of a block input to their respective expert sub-networks, the expert network block first processes the block input using one or more neural network layers to generate an updated representation of the block input, and then assigns the sub-inputs of the updated representation of the block input to their respective expert sub-networks. That is, the input to the expert sub-network can be a strict subset of the elements of the updated representation of the block input, rather than the block input itself. In general, this specification refers to the sub-inputs of the block input, but it should be understood that the same techniques can be applied to the sub-inputs of the updated representation of the block input. Similarly, these neural network layers preceding the expert sub-network can be considered part of the previous network block in the sequence of network blocks.

[0102] Before using the neural network to perform a machine learning task, the training system trains the neural network to perform the task, i.e., to determine the parameters of the neural network, i.e., trained values ​​of the parameters of the blocks in the sequence, and optionally trains an embedding sub-network used to generate the input to the first block in the sequence, an output sub-network that generates the network output from the output of the last block in the sequence, or both. For example, the training system can use traditional machine learning techniques to train the neural network from scratch on the training data of the task to minimize the task loss, e.g., a function cross-entropy loss, a negative log-likelihood loss. As another example, the training system can first pre-train the neural network on an unsupervised objective and then fine-tune the neural network on the training data of the task. As yet another example, the training system can train the neural network on both unlabeled data and the training data of the task by semi-supervised learning.

[0103] Because the system uses expert selection routing that allows load balancing to be "baked in," the system does not need to utilize any auxiliary loss to facilitate load balancing among experts during training, improving training stability and efficiency compared to traditional approaches.

[0104] Furthermore, by using expert selection routing and training the routers of each expert block by backpropagating the gradient of the global loss, the system allows each expert sub-network to be trained, i.e., configured through training, to handle different types of network inputs, thereby allowing the expert sub-networks to "specialize," further improving the efficiency and performance of the neural network.

[0105] During training, the training system can incorporate any number of techniques to improve the speed, effectiveness, or both of the training process. For example, the system can use dropout, label smoothing, or both to reduce overfitting. As another example, the system can perform training using a distributed architecture that trains multiple instances of the neural network in parallel. Additionally, as described above, the system can first pre-train the neural network on a large unsupervised dataset by unsupervised learning, e.g., to minimize the BERT loss or other unsupervised loss, and then fine-tune the neural network on task-specific training data to optimize the loss function for the task.

[0106] An "embedding" as used herein is a vector of numeric values, eg, floating point or other types of numeric values, having a given dimensionality, eg, having a given number of values.

[0107] A self-attention block as described above is a neural network layer that includes an attention mechanism that operates on the self-attention block inputs (or inputs derived from the layer inputs) and produces a self-attention block output. The self-attention mechanism may be causally masked such that any given position in the input sequence does not attend (e.g., does not use data from) any positions after the given position in the input sequence. There are many different attention mechanisms that can be implemented.Some examples of self-attention layers that include attention mechanisms are Vaswani et al. “Attention is all you need”, 31st Conference on Neural Information Processing Systems (NIPS2017), Long Beach, CA, USA; Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li and Peter J Liu, Exploring the limits of transfer learning with a unified text-to-text transformer. arXiv preprint arXiv:1910.10683, 2019; Daniel Adiwardana, Minh-Thang Luong, David R. So, Jamie Hall, Noah Fiedel, Romal Thoppilan, Zi Yang, Apoorv Kulshreshtha, Gaurav Nemade, Yifeng Lu and Quoc V. Le, Towards a human-like open-domain chatbot. CoRR, abs / 2001.09977, 2020; and Tom B Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. arXiv preprint arXiv:2005.14165, 2020.

[0108] In general, an attention mechanism maps a query and a set of key-value pairs to an output, where the query, key, and value are all vectors. The output is computed as a weighted sum of the values, where the weight assigned to each value is computed by a compatibility function of the query with the corresponding key, e.g., the dot product or scaled dot product.

[0109] In general, the self-attention mechanism is configured to associate different positions in the same sequence to determine a transformed version of the sequence as an output. For example, the attention layer input may include vectors per element of the input sequence. These vectors provide input to the self-attention mechanism and are used by the self-attention mechanism to determine a new representation of the same sequence for the attention layer output, which also includes vectors per element of the input sequence. The output of the self-attention mechanism may be used as the attention layer output or may be processed by one or more of a feedforward layer, a skip connection, or a normalization operation to provide the attention layer output.

[0110] In some implementations, the attention mechanism may be, for example, a matrix W Q A query transformation defined by, say, the matrix W k and the key transformation defined by, for example, the matrix W V to the attention layer input, where the attention layer input is the input data X to the attention layer, and a query matrix Q=XW containing a respective query for each vector of the input sequence. Q and a key matrix K=XW that contains each key for each vector in the input sequence. K and a value matrix V=XW that contains the respective values ​​for each vector in the input sequence. Vand used to determine an attended sequence of outputs. For example, the attention mechanism may be a dot-product attention mechanism applied by applying each query vector to each key vector to determine a respective weight for each value vector, then combining the value vectors using the respective weights to determine a self-attention layer output for each element of the input sequence. The self-attention layer output may be scaled by a scaling factor, e.g., by the square root of the dimensions of the query and keys, to implement scaled dot-product attention. Thus, for example, the output of the attention mechanism may be

number

[0111] The attention mechanism may implement multi-head attention, i.e., multiple different attention mechanisms may be applied in parallel. These outputs may then be combined, e.g., concatenated, with a learned linear transformation applied to reduce them back to the original dimensionality, if necessary.

[0112] The term "configured" is used herein in connection with systems and computer program components. A system of one or more computers configured to perform a particular operation or action means that the system has installed thereon software, firmware, hardware, or a combination thereof that causes the system to perform an operation or action. A computer program or programs configured to perform a particular operation or action means that the program or programs contain instructions that, when executed by a data processing device, cause the device to perform an operation or action.

[0113] Embodiments of the subject matter and functional operations described herein can be implemented in digital electronic circuitry, tangibly embodied computer software or firmware, computer hardware, such as the structures disclosed herein and their structural equivalents, or one or more combinations thereof. Embodiments of the subject matter described herein can be implemented as one or more computer programs, i.e., one or more modules of computer program instructions, encoded on a tangible, non-transitory storage medium for execution by or to control the operation of a data processing apparatus. The computer storage medium may be a machine-readable storage device, a machine-readable storage substrate, a random or serial access memory device, or one or more combinations thereof. Alternatively, or additionally, the program instructions can be encoded in an artificially generated propagated signal, e.g., a mechanically generated electrical, optical, or electromagnetic signal, that is generated to encode information for transmission to a receiving device suitable for execution by the data processing apparatus.

[0114] The term "data processing apparatus" refers to data processing hardware and encompasses various apparatus, devices, and machines for processing data, including, by way of example, a programmable processor, a computer, or multiple processors or computers. An apparatus can also be or further include special purpose logic circuitry, such as an FPGA (field programmable gate array) or an ASIC (application specific integrated circuit). In addition to hardware, an apparatus can also optionally include code that creates an execution environment for a computer program, such as code that constitutes processor firmware, a protocol stack, a database management system, an operating system, or one or more combinations thereof.

[0115] A computer program, which may also be referred to or described as a program, software, software application, app, module, software module, script, or code, may be written in any form of programming language, including compiled or interpreted languages, declarative or procedural languages, and may be deployed in any form, including as a stand-alone program or as a module, component, subroutine, or other unit suitable for use in a computing environment. A program may correspond to a file in a file system, but this is not necessarily the case. A program may be stored in a portion of a file that holds other programs or data, e.g., one or more scripts stored in a markup language document, in a single file dedicated to the program in question, or in multiple associated files, e.g., files that store one or more modules, subprograms, or portions of code. A computer program may be deployed to be executed on one computer, or on multiple computers located at one site or distributed across multiple sites and interconnected by a data communication network.

[0116] As used herein, the term "database" is used broadly to refer to any collection of data. The data need not be structured in any particular way, or at all, and may be stored on storage devices in one or more locations. Thus, for example, an index database may contain multiple collections of data, each of which may be organized and accessed in a different way.

[0117] Similarly, the term "engine" is used broadly herein to refer to a software-based system, subsystem, or process that is programmed to perform one or more specific functions. Generally, an engine is implemented as one or more software modules or components and installed on one or more computers in one or more locations. In some cases, one or more computers are dedicated to a particular engine, and in other cases, multiple engines may be installed and executed on the same computer or computers.

[0118] The processes and logic flows described herein may be performed by one or more programmable computers executing one or more computer programs to perform functions by manipulating input data to generate output. The processes and logic flows may also be performed by special purpose logic circuitry, such as an FPGA or an ASIC, or a combination of special purpose logic circuitry and one or more programmed computers.

[0119] A computer suitable for executing a computer program may be based on a general-purpose or special-purpose microprocessor, or on both, or on other types of central processing units. Typically, the central processing unit receives instructions and data from a read-only memory, a random access memory, or on both. The basic elements of a computer are a central processing unit for implementing and executing instructions, and one or more memory devices for storing instructions and data. The central processing unit and memory may be supplemented by, or incorporated with, special-purpose logic circuitry. Typically, a computer also includes one or more mass storage devices, such as magnetic disks, magneto-optical disks, or optical disks, for storing data, or is operatively connected to receive data from or transmit data to them, or both. However, such devices are not required for a computer. Furthermore, a computer may be embedded in other devices, such as a mobile phone, a personal digital assistant (PDA), a mobile audio or video player, a game console, a global positioning system (GPS) receiver, or a portable storage device, such as a universal serial bus (USB) flash drive, to name a few examples.

[0120] Computer-readable media suitable for storing computer program instructions and data include all types of non-volatile memory, media, and memory devices, such as, for example, semiconductor memory devices, e.g., EPROM, EEPROM, and flash memory devices, magnetic disks, e.g., internal hard disks or removable disks, magneto-optical disks, and CD-ROM and DVD-ROM disks.

[0121] To interact with a user, embodiments of the subject matter described herein can be implemented in a computer having a display device, such as a CRT (cathode ray tube) or LCD (liquid crystal display) monitor, for displaying information to the user, and a keyboard and pointing device, such as a mouse or trackball, by which the user can provide input to the computer. Other types of devices can also be used to provide interaction with the user. For example, feedback provided to the user can be any form of sensory feedback, such as visual feedback, auditory feedback, or tactile feedback, and input from the user can be received in any form, such as acoustic, spoken language, or tactile input. Additionally, a computer can interact with a user by sending and receiving documents to a device used by the user, such as by sending a web page to a web browser on a user's device in response to a request received from the web browser. A computer can also interact with a user by sending text messages or other forms of messages to a personal device, such as a smartphone running a messaging application, and receiving a response message from the user in return.

[0122] A data processing apparatus for implementing machine learning models may also include dedicated hardware accelerator units for handling the general and numerically-intensive parts of, for example, machine learning training or machine learning production, i.e., inference, workload.

[0123] The machine learning model can be implemented and deployed using a machine learning framework, for example, the TensorFlow framework, the Microsoft Cognitive Toolkit framework, the Apache Singa framework, or the Apache MXNet framework.

[0124] An embodiment of the subject matter described herein can be implemented in a computing system including a back-end component, e.g., a data server, or a middleware component, e.g., an application server, or a front-end component, e.g., a client computer having a graphical user interface, a web browser, or an app through which a user can interact with an embodiment of the subject matter described herein, or any combination of one or more such back-end, middleware, or front-end components. The components of the system can be interconnected by any form or medium of digital data communication, e.g., a communication network. Examples of communication networks include a local area network (LAN), a wide area network (WAN), e.g., the Internet.

[0125] A computing system may include clients and servers. Clients and servers are generally remote from each other and typically interact through a communication network. The relationship of client and server arises by virtue of computer programs running on the respective computers and by virtue of having a client-server relationship to each other. In some embodiments, the server sends data, e.g., HTML pages, to a user device, e.g., for the purpose of displaying the data to and receiving user input from a user interacting with the device acting as a client. Data generated at the user device, e.g., results of user interaction, can be received from the device at the server.

[0126] Although the present specification contains many specific implementation details, these should not be construed as limiting the scope or scope of any invention or patentable scope, but rather as descriptions of features that may be inherent to a particular embodiment of a particular invention. Certain features described herein in the context of separate embodiments can also be implemented in combination in a single embodiment. Conversely, various features of the invention described in the context of a single embodiment can also be implemented in multiple embodiments separately or in any suitable subcombination. Furthermore, even if features may be described above as functioning in a particular combination and originally claimed as such, one or more features from the claimed combination may in some cases be deleted from the combination, and the claimed combination may be directed to a subcombination or a variation of the subcombination.

[0127] Similarly, although the figures show operations in a particular order, and the claims describe operations in a particular order, this should not be understood as requiring that such operations be performed in the particular order or sequential order shown, or that all of the operations shown be performed, to achieve desirable results. In certain situations, multitasking and parallel processing may be advantageous. Furthermore, the separation of various system modules and components in the above embodiments should not be understood as requiring such separation in all embodiments, and it should be understood that the program components and systems described may generally be integrated into a single software product or packaged into multiple software products.

[0128] Specific embodiments of the subject matter have been described. Other embodiments are within the scope of the following claims. For example, the actions recited in the claims may be performed in a different order and still achieve desirable results. By way of example, the processes depicted in the accompanying figures do not necessarily require the particular order shown, or sequential order, to achieve desirable results. In some cases, multitasking and parallel processing may be advantageous.

Claims

1. 1. A system including a neural network configured to process network inputs and generate network outputs for the network inputs, the neural network including a sequence of one or more network blocks, the sequence including one or more expert network blocks configured to perform operations, the operations comprising: obtaining a block input representing an intermediate representation of the network input, the block input including a plurality of elements; determining a plurality of sub-inputs from the block input, each sub-input comprising a respective different subset of the plurality of elements of the block input; For each of a plurality of expert sub-networks of the expert network block, processing said plurality of sub-inputs to generate a respective score for each sub-input; selecting one or more of said sub-inputs according to their respective scores; for each selected sub-input, processing the selected sub-input using the expert sub-network to generate a respective sub-output; for each of said plurality of sub-inputs, processing said sub-outputs generated by respective expert sub-networks corresponding to said sub-inputs to generate combined sub-outputs for said sub-inputs; generating a block output by combining the respective combined sub-outputs for the plurality of sub-inputs; Including, the system.

2. The system of claim 1 , wherein each expert sub-network is configured to process the same number of sub-inputs.

3. The number of equal sub-inputs, k, processed by each expert sub-network is [0010] 3. The system of claim 2, wherein l is the number of sub-inputs of the block input, e is the number of expert sub-networks of the expert network block, and c is a hyper-parameter of the neural network representing the average number of sub-inputs processed per expert sub-network.

4. processing, for each expert sub-network, the plurality of sub-inputs to generate a respective score for each sub-input; [0025] wherein: [0030] is a matrix with a respective row corresponding to each sub-input, l is the number of sub-inputs of the block input, and d is the dimension of each sub-input; [0045] The system of any one of claims 1 to 3, wherein x is a matrix with a respective column corresponding to each expert sub-network, and e is the number of expert sub-networks in the expert network block.

5. selecting, for each expert sub-network, one or more of the sub-inputs according to their respective scores, G,I=TopK(S Ｔ ,k) P=Onehot(I) where k is the number of sub-inputs selected by each expert sub-network; [0050] is a matrix of (i,j) elements that identifies the sub-input with the jth highest score of the i-th expert sub-network; [006] is a matrix with (i,j) elements, which represents the scores of the jth highest scoring sub-input of the ith expert sub-network; [0070] 5. The system of claim 4, wherein: m is a one-hot matrix of (i,j,m) elements, the one-hot matrix of (i,j,m) elements being equal to one if the m sub-input has the j maximum score of the i expert sub-network, and zero otherwise.

6. A system according to any one of claims 1 to 5, wherein each sub-input is processed by up to a threshold number b of different expert sub-networks.

7. selecting, for each expert sub-network, one or more of the sub-inputs according to their respective scores, [0080] The calculation is carried out so that In the formula, <S Ｔ , A> is S Ｔ represents the dot product between A and A, [0090] and G,I=TopK(A,k) P=Onehot(I) where k is the number of sub-inputs selected by each expert sub-network; [0010] is a matrix of (i,j) elements that identifies the sub-input with the jth highest score of the i-th expert sub-network; ##EQU00011## is a matrix with (i,j) elements, which represents the scores of the jth highest scoring sub-input of the ith expert sub-network; ##EQU00012## 7. The system of claim 6, wherein: m is a one-hot matrix of (i,j,m) elements, the one-hot matrix of (i,j,m) elements being equal to one if the m sub-input has the j maximum score of the i expert sub-network, and zero otherwise.

8. For each expert sub-network i, processing each selected sub-input using said expert sub-network to generate a respective sub-output, ##EQU00013## wherein: ##EQU14## is the input to the i-th expert sub-network, and W １ [i] and W ２ [i] is the weight matrix of the i-th expert sub-network, ACT is the activation function, ##EQU00015## 8. The system of claim 5, wherein i is a matrix with j-th row representing the sub-output for the j-th sub-input processed by the expert sub-network i.

9. For each expert sub-network i, processing each selected sub-input using said expert sub-network to generate a respective sub-output includes adjusting one or more bias terms by X ｅ The system of claim 8 , further comprising: adding to [i].

10. for each of the plurality of sub-inputs, processing the sub-outputs generated by a respective expert sub-network corresponding to the sub-input to generate combined sub-outputs for the sub-inputs; ##EQU00016## wherein: ##EQU00017## 10. The system of claim 8, wherein m is a matrix with an m-th row representing the combined sub-output for the m-th sub-input.

11. for each of the plurality of sub-inputs, processing the sub-outputs generated by a respective expert sub-network corresponding to the sub-input to generate combined sub-outputs for the sub-inputs; determining a weighted sum of the sub-outputs corresponding to the sub-inputs; A system according to any preceding claim, wherein each sub-output is weighted by the score of the sub-input for the expert sub-network that generated it.

12. A system according to any preceding claim, wherein the operations of at least some of the plurality of expert sub-networks are performed in parallel.

13. A system according to any preceding claim, wherein the operations of at least some of the plurality of expert sub-networks are performed by respective different devices.

14. A method comprising the acts of any one of claims 1 to 13.

15. One or more computer storage media having stored thereon instructions that, when executed by one or more computers, cause the one or more computers to perform the operations recited in any one of claims 1 to 13.

Images