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 of related applications This application claims priority to U.S. Provisional Application No. 63 / 304,507 filed on 28 January 2022, which is incorporated herein by reference in its entirety. [Background technology]

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

[0003] A neural network is a machine learning model that uses one or more layers of nonlinear units to predict an output for a given 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 input to the next layer in the network, i.e., the next hidden layer or output layer. Each layer of the network generates an output from a given input according to the current values ​​of its respective set of parameters. [Overview of the project]

[0004] This specification describes a system implemented as a computer program on one or more computers located in one or more locations, configured to use a neural network to process network inputs 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, each containing a network input or an intermediate representation of a network input, and to generate a block output.

[0005] For example, the first network block in a sequence of network blocks can process the network input and produce a block output, which is an intermediate representation of the network input. Another example is an embedding subnetwork, which can process the network input and produce an embedding of the network input, which is then provided as input to the first network block in the sequence, and process the embedding to produce an intermediate representation of the network input. Each subsequent network block can then process the block output of the previous network block in the sequence. In some embodiments, the network output of the neural network is the block output of the last network block in the sequence. In some other embodiments, the block output of the last network block in the sequence is further processed using one or more output neural network layers to produce the network output of the neural network.

[0006] A sequence of network blocks may include one or more "expert" network blocks. Each expert network block contains multiple different expert subnetworks, each configured to process its respective sub-inputs determined from the block input to the expert network block. Each sub-input of the block input contains a different subset of the elements of the block input.

[0007] Unlike traditional approaches, expert network blocks use “subnetwork selection” or “expert selection” routing to route subinputs to expert subnetworks. That is, for each expert subnetwork, the neural network block independently selects a set of elements from the block input to be processed by the expert subnetwork.

[0008] Certain embodiments of the subject matter described herein may be implemented to achieve one or more of the following advantages:

[0009] By implementing a neural network using 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 during inference time. That is, by selectively activating only a subset of the neural network parameters for each element of the network input, the system can significantly improve the time efficiency 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 contain more network parameters than previously possible, because only a subset of 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, where the neural network block independently selects a set of expert subnetworks to process each element of the block input. However, these existing systems can suffer from load imbalance, where some expert subnetworks process most or all of the elements of the block input, while others process little or none. Such load imbalance can result in suboptimal training, as some 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 using token selection routing may further reduce the computational efficiency of the system by allocating the same amount of computational resources to each element of the network input, ignoring the relative importance of different elements.

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

[0012] In particular, in some embodiments where the system runs different expert subnetworks on different devices, the techniques described herein enable the system to load balance network inputs to neural networks more efficiently across devices compared to existing techniques (e.g., compared to systems implementing token selection routing). The inadequate load balancing suffered by existing systems with token selection routing can result in significantly different loads on different devices running different expert subnetworks, meaning some devices may be underutilized while others may be overloaded, potentially impairing inference performance (e.g., by reducing computational and / or memory efficiency, or by increasing the time required to generate network outputs). Using subnetwork selection routing as described herein allows for "baked-in" full load balancing during inference, and since multiple devices each have a similar or equivalent amount of operation to perform, the system can enjoy significant performance improvements across multiple devices (e.g., improved computational and / or memory efficiency, or reduced time required to generate network outputs). Therefore, compared to other existing approaches, the described approach, by being optimized for distributed hardware implementations, results in a system that can process inputs with higher throughput compared to conventional approaches.

[0013] In some embodiments described herein, each expert subnetwork can be configured through training to process different types of network inputs, enabling the expert subnetwork to "specialize" and further improving the efficiency and performance of the neural network. Since each routing subnetwork corresponding to each expert subnetwork can be trained to select specific types of elements to be processed by the expert subnetwork, "subnetwork selection" routing can further improve the specialization of the expert subnetwork compared to "token selection" routing.

[0014] Details of one or more embodiments of the subject matter of this specification are described in the accompanying drawings and the following description. Other features, aspects, and advantages of the subject matter will become apparent from the description, drawings, and claims. [Brief explanation of the drawing]

[0015] [Figure 1] Figure 1 shows an example of a neural network system. [Figure 2] Figure 2 illustrates the behavior of an expert block using expert-selection routing compared to token-selection routing. [Figure 3] Figure 3 is a flowchart illustrating an exemplary process for processing network inputs and generating network outputs. [Modes for carrying out the invention]

[0016] Similar reference symbols and designations in various drawings refer to the same elements.

[0017] This specification describes a system implemented as a computer program on one or more computers located in one or more locations, which performs a machine learning task on a network input and generates a network output of the machine learning task.

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

[0019] Some examples of machine learning tasks that the system can be configured to perform are shown below.

[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 can 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 can indicate whether a particular word or phrase ("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 can identify the natural language in which that utterance was spoken. In the case of a speech processing task, it will be understood that the input to the neural network may be in the form of a sequence of audio data frames, for example, and may include audio data (e.g., an audio signal), 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 that operates on a text sequence in one or more natural languages and generates a classification output that classifies the text into one or more categories from a set of categories, such as, for example, an entailment task, a paraphrase task, a text similarity task, a sentiment task, a text completion task, a grammar task, etc.

[0022] As another example, the task can be a health prediction task, the input can be a sequence derived from a patient's electronic health record data, and the output can be a prediction related to the patient's future health, such as a predicted treatment to be prescribed to the patient, the likelihood of a harmful health event occurring to the patient, or the predicted diagnosis of the patient.

[0023] As another example, the task can be an agent control task, the input can be an observation sequence or other data characterizing the environmental state, and the output defines an action to be executed by the agent according to the latest data in the sequence. The 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 types of agents.

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

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

[0026] For example, a computer vision task can be a classification task that requires the generation of a classification output. The classification output typically includes a score corresponding to each of several categories. The category score indicates the likelihood that an image belongs to that category. In some cases, a 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 that belongs to the object class corresponding to that category. In some cases, a category may represent a global image characteristic (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 possesses a global characteristic corresponding to that category.

[0027] As another example, a computer vision task can be an object detection task. In an object detection task, the output generated by a neural network identifies the location of a particular type of object within the input image, for example, a bounding box or other region.

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

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

[0030] Another example of a computer vision task could be a depth prediction task. In a depth prediction task, the output generated by a neural network identifies the predicted depth of the scene at each pixel in an image.

[0031] Another example of a computer vision task could be a surface normal prediction task. In a surface normal prediction task, the output generated by a neural network identifies the predicted surface normal of the scene at each pixel in the image.

[0032] When the input is an image or point cloud, the neural network may include an embedding subnetwork that generates an embedding for each of several patches of the image or point cloud, and the input to the first block of the neural network may be a sequence containing each embedding (and optionally one or more additional embeddings at predetermined locations, for example, to be used later to generate an output). Each patch contains the intensity values ​​of pixels from a different region of the input image.

[0033] In some embodiments, the task is a multimodal task requiring processing of both text and image inputs, and therefore the neural network includes both a computer vision neural network and a text processing neural network. That is, the target output produced by the computer vision neural network for a given image depends on one or more outputs produced 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 search, and image-based search.

[0034] In some cases, a machine learning task is a combination of multiple individual machine learning tasks. That is, a system is configured to perform multiple different individual machine learning tasks, for example, two or more of the machine learning tasks described above. For example, a system can be configured to perform multiple individual natural language understanding tasks using a network input that contains identifiers for the individual natural language understanding tasks to be performed on the network input.

[0035] Figure 1 shows an example of a neural network system 100. The neural network system 100 is an example of a system implemented as a computer program on one or more computers located at one or more locations, and can implement the systems, components, and techniques described below.

[0036] System 100 is a system that uses a neural network 110 to process a network input 102 and generates a network output 112 that characterizes the network input 102 for a machine learning task, for example, one of the tasks described above.

[0037] The neural network 110 includes a sequence of one or more network blocks 120, each configured to process a block input, which is either a network input or an intermediate representation of a network input, and to generate a block output.

[0038] As used herein, a “network block” is a collection of one or more neural network layers that receive an input (“block input”), process the input, and produce an output (“block output”).

[0039] For example, the first network block in a sequence of network blocks 120 can process the network input 102, or the embedding of the network input generated by the embedding subnetwork, and produce a block output, which is an intermediate representation of the network input. Then, each subsequent network block 120 can process the block output of the previous network block in the sequence.

[0040] In some embodiments, 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 embodiments, 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] A sequence of network blocks may include one or more "expert" network blocks 130. Each expert network block 130 includes several different expert subnetworks 132, each configured to process its respective sub-inputs 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. In some embodiments, each sub-input contains exactly one of the elements of the block input.

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

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

[0046] More specifically, router 134 implements "expert selection" routing by independently selecting the sub-inputs to be processed by each expert subnetwork 132. This is in contrast to "token selection" routing or "sub-input" selection routing.

[0047] Expert selection routing is explained in more detail below, with reference to Figure 2.

[0048] Next, each expert subnetwork 132 can process only the sub-input(s) selected for the subnetwork 132 and generate a corresponding 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 an appropriate subset of the elements of the block input.

[0049] After each expert subnetwork 132 has been executed, the expert network block 130 can determine the combined sub-outputs corresponding to each sub-input.

[0050] In particular, for each sub-input, the expert network block 130 can combine the sub-outputs generated by each expert subnetwork in accordance with the processing of the sub-input.

[0051] For sub-inputs that were not processed by any of the expert subnetworks 132, the combined sub-output of these sub-inputs can be the same as the sub-input.

[0052] Next, the expert network block 130 can combine its combined sub-outputs 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 (for example, in the same order) as the corresponding sub-inputs of the block input.

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

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

[0055] In some embodiments, each element of a block input resides in exactly one sub-input. In some other embodiments, some or all elements of a block input may reside in multiple different sub-inputs.

[0056] The operations performed by the expert network block 130 will be described in more detail below with reference to Figures 2 and 3.

[0057] In some embodiments, the sequence of network blocks 120 includes expert network blocks 130 interspersed within other types of network blocks, such as self-attention network blocks that apply self-attention, where the expert network blocks do not include routers and expert neural networks; that is, they do not perform conditional calculations and use all the parameters of the network blocks for all inputs to the network blocks. In a particular example, the sequence of network blocks 120 can alternate between expert network blocks and self-attention network blocks. In another particular example, the sequence network block 120 may include a self-attention network block, a feedforward network block that includes a single neural network having the same architecture as the expert neural network 132 and processes all sub-inputs of the block input to the feedforward block, and the expert network block 130. For example, the expert network block 130 may follow immediately after all other self-attention network blocks in the sequence, followed by the feedforward network block for the remainder of the self-attention network blocks.

[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 a sequence of input elements as input and applies an attention mechanism to the sequence of input elements to generate a sequence of layer output elements. Specifically, for each input element, the self-attention neural network layer uses one or more queries derived from the input element to apply an attention mechanism to the sequence of input elements to generate the 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 a sequence of output elements, and then combines multiple sequences of output elements to generate a final sequence of output elements.

[0060] Self-attention will be explained in more detail below.

[0061] Each expert network block 130 can be implemented such that its expert subnetworks 132 run in parallel with respect to a given block input, thereby improving the efficiency of the system. For example, an expert network block 130 can be implemented on a parallel processing device, such as a GPU or TPU, which can run expert subnetworks in separate threads. As another example, at least some of the expert subnetworks 132 of an expert network block 130 can be implemented on separate devices, such as separate devices that are communicably connected and provide the sub-outputs generated by each expert subnetwork to a single device for combining them to produce their combined sub-outputs.

[0062] Therefore, a neural network architecture including an expert network block 130 with multiple expert subnetworks 132 is optimized for efficient execution of the neural network. Such a network architecture allows for the parallelization of the neural network's operation for rapid and low-cost execution, for example, by parallelizing the operation of each expert subnetwork across multiple devices. Thus, the neural network can be implemented in dedicated parallel processing hardware, for example, a network of multiple parallel processing devices that each run on an expert subnetwork of the neural network.

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

[0064] Figure 2 illustrates the operations performed by the routers of the 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 in Figure 2, expert network block 130 processes an intermediate representation of the sequence "We like playing soccer on that field," which contains elements ("tokens") 202-216 corresponding to each word in the sequence, and expert network block 130 contains four expert neural networks 132. In the example in Figure 2, each expert 132 is a feedforward neural network (FFN). Furthermore, in the example in Figure 2, each sub-input contains one different element.

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

[0067] Therefore, when routing token 2, the token selection router 230 does not consider where token 1 was routed, or more generally, the score assigned to FFN1 of other tokens. As a result, this system may suffer from load imbalance, with some expert subnetworks processing most or all of the elements of the block input, while others process little or no elements. Such load imbalance can result in suboptimal training, as some network parameters (corresponding to underutilized expert subnetworks) do not receive meaningful updates during training and therefore do not learn to extract useful information. Furthermore, this system may further reduce the computational efficiency of the system by allocating the same amount of computational resources to each element of the network input, ignoring the relative importance of different elements.

[0068] In expert selection routing 250, the expert selection router 260 instead generates a score distribution for each expert neural network 132. The score distribution includes the respective scores for tokens 202 to 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 can route the token expression "we," "like," "thing (To)," and "do" to FFN1, and can route the token expression "we," "like," "soccer," and "field" to FFN2. Therefore, if a given token is among the top k tokens in the score distribution of multiple different expert neural networks 132, the expert selection router 260 can route the same token to multiple different expert neural networks 132.

[0070] As can be seen in Figure 2, expert selection routing can ensure that expert blocks are fully load-balanced, for example, by selecting the same number of k elements to be processed by each expert subnetwork. Therefore, the computational and time efficiency of training the neural network can be greatly improved so that all parameters of the neural network receive meaningful updates for each network input. In some embodiments, a neural network training system can reduce the time and computational resources required to achieve a given performance by more than double using subnetwork selection routing compared to token selection routing. Furthermore, subnetwork selection routing can allow network blocks to allocate computational resources to each element more flexibly, for example, by routing relatively important elements to more expert subnetworks than relatively less important elements.

[0071] Processing input using expert selection routing is explained in more detail below, with reference to Figure 3.

[0072] Figure 3 is a flowchart 300 of an exemplary process for processing block inputs using expert blocks. For convenience, process 300 is described as being performed by one or more computer systems at one or more locations. For example, one of the expert blocks included in a neural network system, e.g., one of the expert blocks 130 included in the appropriately programmed neural network system 100 of Figure 1, can perform 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 multiple sub-inputs from the block input (step 304). Each sub-input contains a distinct subset of the multiple elements of the block input. In some embodiments, each sub-input contains exactly one of the elements of the block input; that is, each sub-input is a distinct element of the block input.

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

[0076] In some embodiments, each element of a block input resides in exactly one sub-input. In some other embodiments, some or all elements of a block input may reside in multiple different sub-inputs.

[0077] Next, the expert block performs steps 306-310 for each of the expert subnetworks in the expert network block.

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

[0079] For each expert subnetwork, the expert network block, i.e., the expert selection router within the block, can determine the score of the sub-inputs in any appropriate manner. For example, for each expert subnetwork, the expert network block can process each sub-input and generate its respective score using one or more neural network layers, e.g., one or more feedforward neural network layers. As a specific example, the expert network block can compute the following:

number

number

number

[0080] As another example, for each expert subnetwork, the expert network block can process each sub-input using one or more convolutional neural network layers and / or one or more self-attention layers to generate a score for each. Self-attention is explained 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 subnetwork, after the expert network block generates a 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 subnetwork. For example, the expert network block can calculate the following: G,I=TopK(S T ,k) P = Onehot(I) In the formula, TopK(S T ,k) is S T Select the maximum entry k for each row,

number

number

number

[0083] In some embodiments, the expert network block can compel each sub-input to be selected by up to b different expert subnetworks. For example, the expert network block can solve the following entropy regularized linear programming problem:

number

number

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

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

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

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

[0088] In some embodiments, each expert subnetwork selects the same number of sub-inputs; that is, each expert subnetwork can select the k sub-inputs with the highest corresponding score, where k is the same for all expert subnetworks. As a specific example, the same k of sub-inputs processed by each expert subnetwork can be equal to:

number

[0089] In some other embodiments, some expert subnetworks select a different number of sub-inputs. For example, an expert subnetwork may 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 subnetwork to generate its respective sub-output (step 310).

[0091] In other words, after selecting k sub-inputs, each expert subnetwork can process the k sub-inputs using one or more neural network layers to produce its respective sub-output. For example, each expert subnetwork may 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 subnetwork i can generate its own sub-output for each sub-input by calculating the following:

number

number

number

[0093] Once the expert subnetworks have completed processing, for each of the multiple sub-inputs, the system processes the sub-outputs corresponding to the sub-inputs generated by each expert subnetwork to produce a combined sub-output for the sub-input (step 312). In particular, for each sub-input, the expert network block can combine each sub-output generated by each expert subnetwork, depending on the sub-input it has processed. For sub-inputs that have not been processed by any expert subnetwork, the combined sub-output of the sub-input can be the same as the sub-input.

[0094] In some embodiments, an expert network block can generate combined sub-outputs by calculating a sum of the sub-outputs. For example, an expert network block can calculate a weighted sum of sub-outputs, where each sub-output is weighted by the score of the corresponding sub-input of the expert subnetwork that generated the sub-output.

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

number

number

[0096] In some other embodiments, the expert network block combines sub-outputs for each sub-input by processing the sub-outputs using one or more neural network layers, e.g., 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 the combined sub-outputs for multiple sub-inputs (step 314). For example, an expert network block can concatenate 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 sub-outputs, the system may apply one or more additional operations to the concatenation of the combined sub-outputs, for example, the concatenation may be processed by one or more feedforward layers, skip connections, or normalization operations, such as layer normalization, to provide a block output.

[0099] Expert network blocks can be implemented so that their expert subnetworks run in parallel on a given block input, thereby improving system efficiency. For example, an expert network block can be implemented on a parallel processing device, such as a GPU or TPU, which can run expert subnetworks in separate threads. As another example, at least some of the expert subnetworks of an expert network block can be implemented on separate devices, such as separate devices that are communicatively connected and provide the sub-outputs generated by each expert subnetwork to a single device to combine them to produce their combined sub-outputs.

[0100] Therefore, a neural network architecture that includes expert network blocks with multiple expert subnetworks is optimized for the efficient execution of the neural network. Such a network architecture allows for the parallelization of the neural network's operation for rapid and low-cost execution, for example, by parallelizing the operation of each expert subnetwork across multiple devices. Thus, the neural network can be implemented in dedicated parallel processing hardware, for example, a network of multiple parallel processing devices that each run on an expert subnetwork of the neural network.

[0101] In some embodiments, before assigning the sub-inputs of the block input to each expert subnetwork, 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 each expert subnetwork. That is, the input to the expert subnetwork can be a strict subset of the elements of the updated representation of the block input, rather than the block input itself. Generally, this specification refers to the sub-inputs of the block input, but it should be understood that the same technique can be applied to the sub-inputs of the updated representation of the block input. Similarly, these neural network layers preceding the expert subnetwork can be considered as part of the preceding network block in the sequence of network blocks.

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

[0103] Because this system uses expert-selection routing that can "burn in" load balancing, the system does not need to utilize any auxiliary losses to facilitate load balancing among experts during training, improving training stability and efficiency compared to conventional approaches.

[0104] Furthermore, by using expert selection routing and training the routers of each expert block by backpropagating the overall loss gradient, the system allows each expert subnetwork to be trained to handle different types of network inputs, i.e., configured through training, thus enabling the expert subnetworks to "specialize" and 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. Furthermore, as mentioned above, the system can first pre-train the neural network on a large unsupervised dataset by unsupervised learning to minimize, for example, the BERT loss or other unsupervised loss, and then fine-tune the neural network on task-specific training data to optimize the task's loss function.

[0106] As used herein, “embedding” refers to a numerical vector having a predetermined dimensionality, for example, a predetermined number of values, such as a floating-point or other type of numerical vector.

[0107] A self-attention block, as described above, is a neural network layer that includes an attention mechanism that operates on a self-attention block input (or an input derived from the layer input) and generates a self-attention block output. The self-attention mechanism can be causally masked such that any given position in the input sequence does not attend to any position after that given position in the input sequence (e.g., does not use data from there). Many different attention mechanisms are possible.Some examples of self-attention layers including attention mechanisms are found in 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 This is described in chatbot.CoRR,abs / 2001.09977,2020; and in 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] Generally, an attention mechanism maps a query and a set of key-value pairs to an output. Here, the query, key, and value are all vectors. The output is calculated as a weighted sum of the values, and the weight assigned to each value is calculated by a compatibility function of the query with the corresponding key, such as a dot product or a scaled dot product.

[0109] Generally, a self-attention mechanism is configured to associate different positions in the same sequence and determine a transformed version of the sequence as an output. For example, the attention layer input may include vectors for each element of the input sequence. These vectors provide the input to the self-attention mechanism, which is used by the self-attention mechanism to determine a new representation of the same sequence for the attention layer output, and that representation similarly includes vectors for each 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 feed-forward layer, a skip connection, or a normalization operation to provide the attention layer output.

[0110] In some embodiments, the attention mechanism applies each of a query transformation defined by, for example, matrix W Q a key transformation defined by, for example, matrix W k and a value transformation defined by, for example, matrix W V to the attention layer input, which is the input data X to the attention layer and includes a query matrix Q = XW Q that contains a respective query for each vector of the input sequence, a key matrix K = XW K that contains a respective key for each vector of the input sequence, and a value matrix V = XW VThis is used to derive the and determine the attended sequence of the output. For example, the attention mechanism may be a dot product attention mechanism applied by applying each query vector to each key vector to determine the respective weight for each value vector, then combining the value vectors using the respective weights to determine the self-attention layer output for each element of the input sequence. The self-attention layer output can be scaled by a scaling factor, for example by the square root of the dimensions of the query and key to implement a scaled dot product attention. Thus, for example, the output of the attention mechanism is

number

[0111] The attention mechanism may implement multi-head attention; that is, multiple different attention mechanisms may be applied in parallel. These outputs may then be combined, for example, concatenated, with a learned linear transformation that is applied to reduce them to the original dimension as needed.

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

[0113] The subject matter and functional embodiments described herein can be implemented in digital electronic circuits, tangibly embodied computer software or firmware, or computer hardware, including 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 in a tangible non-temporary storage medium, to be executed by or control the operation of a data processing device. The computer storage medium may be a machine-readable storage device, a machine-readable storage board, a random or serial access memory device, or one or more combinations thereof. Alternatively or additionally, the program instructions can be encoded into artificially generated propagating signals, such as mechanically generated electrical signals, optical signals, or electromagnetic signals, which are generated to encode information for transmission to a receiving device suitable for execution by a data processing device.

[0114] The term "data processing device" refers to data processing hardware and encompasses a variety of devices, machines, and equipment for processing data, including, for example, programmable processors, computers, or multiple processors or computers. A device may also be, or further include, a special-purpose logic circuit, such as an FPGA (Field-Programmable Gate Array) or ASIC (Application-Specific Integrated Circuit). In addition to hardware, a device may optionally include code that creates an execution environment for computer programs, such as processor firmware, protocol stacks, database management systems, operating systems, or code that constitutes one or more of these.

[0115] Computer programs, which may also be called or described as programs, software, software applications, apps, modules, software modules, scripts, or code, can be written in any form of programming language, including compiled or interpreted languages, declarative or procedural languages, and can be deployed in any form, including as standalone programs or as modules, components, subroutines, or other units suitable for use in a computing environment. A program may, but may not, correspond to a file in a file system. A program may be stored in a part of a file that holds one or more scripts stored in a markup language document, in a single file dedicated to the program of interest, or in multiple collaborative files, such as a file that holds one or more modules, subprograms, or parts of code. A computer program can be deployed to run on one computer or on multiple computers located in one location or distributed across multiple locations and interconnected by a data communication network.

[0116] In this specification, the term “database” is used broadly to refer to any collection of data. The data does not need to be structured in any particular way, or not structured at all, and can be stored in one or more storage devices. Therefore, 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” as used herein is used broadly to refer to a software-based system, subsystem, or process 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 directed to a particular engine, and in other cases, multiple engines may be installed and run on the same one or more computers.

[0118] The processes and logic flows described herein can be executed by one or more programmable computers running one or more computer programs, performing their functions by manipulating input data to produce outputs. Alternatively, the processes and logic flows can be executed by special-purpose logic circuits, such as FPGAs or ASICs, or by a combination of special-purpose logic circuits and one or more programmed computers.

[0119] A computer suitable for running computer programs can be based on a general-purpose or dedicated microprocessor, or both, or other types of central processing units. Generally, the central processing unit receives instructions and data from read-only memory, random-access memory, or both. The basic elements of a computer are the central processing unit for executing and running instructions, and one or more memory devices for storing instructions and data. The central processing unit and memory can be complemented or incorporated by dedicated logic circuits. Generally, a computer also includes one or more mass storage devices for storing data, such as magnetic disks, magneto-optical disks, or optical disks, or is operablely connected to them to receive data from them, transmit data to them, or both. However, such devices are not necessary for a computer. Furthermore, to give some examples, a computer can be embedded in other devices, such as mobile phones, personal digital assistants (PDAs), mobile audio or video players, game consoles, Global Positioning System (GPS) receivers, or portable storage devices, such as Universal Serial Bus (USB) flash drives.

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

[0121] To interact with a user, embodiments of the subject matter described herein can be implemented in a computer having a display device for displaying information to the user, such as a CRT (cathode ray tube) or LCD (liquid crystal display) monitor, and a keyboard and pointing device, such as a mouse or trackball, through which the user can input to the computer. Other types of devices can also be used to provide interaction with the user. For example, the feedback provided to the user can be any form of sensory feedback, such as visual feedback, auditory feedback, or haptic feedback, and input from the user can be received in any form, such as acoustic, spoken language, or haptic input. Furthermore, the computer can interact with the user by sending and receiving documents to and from the device used by the user, for example, by sending a web page to a web browser on the user's device in response to a request received from a web browser. The computer can also interact with the 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] Data processing equipment for implementing machine learning models may include, for example, dedicated hardware accelerator units for processing machine learning training or machine learning production, i.e., inference, and the general and numerical computation portions of the workload.

[0123] Machine learning models can be implemented and deployed using machine learning frameworks, such as the TensorFlow framework, the Microsoft Cognitive Toolkit framework, the Apache Singa framework, or the Apache MXNet framework.

[0124] Embodiments of the subject matter described herein can be implemented in a computing system including, for example, a backend component as a data server, or a computing system including a middleware component, for example, an application server, or a frontend component, for example, a client computer having a graphical user interface, a web browser, or an application that allows a user to interact with embodiments of the subject matter described herein, or in a computing system including one or more such backend, middleware, or frontend components in any combination. The components of the system can be interconnected by any form or medium of digital data communication, for example, a communication network. Examples of communication networks include local area networks (LANs), wide area networks (WANs), and for example, the Internet.

[0125] A computing system can include clients and servers. Clients and servers are generally geographically distant from each other and typically interact through a communication network. The client-server relationship arises from computer programs running on each computer and from the client-server relationship between them. In some embodiments, the server sends data, such as an HTML page, to a user device for the purpose of displaying data to a user interacting with a device acting as a client and receiving user input from that user. Data generated on the user device, such as the results of user interactions, can be received by the server from the device.

[0126] While this specification includes details of many specific embodiments, these should not be construed as limiting the scope of any invention or claimable scope, but rather as descriptions of features that may be specific to a particular embodiment of a particular invention. Certain features described herein in the context of individual embodiments may also be implemented in combination in a single embodiment. Conversely, various features of the invention described in the context of a single embodiment may be implemented separately or in any preferred secondary combination in multiple embodiments. Furthermore, features may be described above as functioning in a particular combination, and even if initially claimed as such, one or more features from the claimed combination may be removed from the combination, and the claimed combination may cover secondary combinations or variations of secondary combinations.

[0127] Similarly, while the drawings show operations in a specific order and the claims describe operations in a specific order, this should not be understood as requiring that such operations be performed in a specific or sequential order shown, or that all shown operations be performed, in order to obtain the desired results. In certain situations, multitasking and parallel processing may be advantageous. Furthermore, the separation of various system modules and components in the embodiments described above should not be understood as requiring such separation in all embodiments, and the described program components and systems can 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 desired results may still be obtained by performing the actions described in the claims in a different order. As an example, the process shown in the accompanying drawings does not necessarily require that the actions be performed in a specific or sequential order shown in order to obtain the desired results. In some cases, multitasking and parallel processing may be advantageous.

Claims

1. A system including a neural network configured to process a network input and generate a network output for the network input, wherein the neural network includes a sequence of one or more network blocks, the sequence includes one or more expert network blocks configured to perform an operation, the operation is Obtaining a block input representing an intermediate representation of the network input, wherein the block input includes multiple elements, Determining a plurality of sub-inputs from the block input, wherein each sub-input includes a different subset of the plurality of elements of the block input, For each of the multiple expert subnetworks in the aforementioned expert network block, The process involves processing the aforementioned multiple sub-inputs and generating a score for each sub-input, Selecting one or more of the sub-inputs according to each of the aforementioned scores, For each selected sub-input, the expert subnetwork is used to process the selected sub-input and generate its respective sub-output. For each of the aforementioned multiple sub-inputs, the sub-outputs corresponding to the sub-inputs, generated by each expert subnetwork, are processed to generate a sub-output that is combined with the sub-input; A block output is generated by combining the respective combined sub-outputs for the plurality of sub-inputs, A system that includes this.

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

3. The number of sub-inputs k processed by each expert subnetwork is: [Math 1] The system according to claim 2, where l is the number of sub-inputs of the block input, e is the number of expert subnetworks of the expert network block, and c is a hyperparameter of the neural network representing the average number of sub-inputs processed per expert subnetwork.

4. For each expert subnetwork, processing the multiple sub-inputs and generating a score for each sub-input is: [Math 2] This includes calculating, in the formula, [Math 3] l is a matrix containing each 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. [Math 4] The system according to claim 1, wherein is a matrix containing a column corresponding to each expert subnetwork, and e is the number of expert subnetworks in the expert network block.

5. For each expert subnetwork, selecting one or more of the sub-inputs according to their respective scores is: G,I=TopK(S Ｔ ,k) P=Onehot(I) This includes calculating the formula where k is the number of sub-inputs selected by each expert subnetwork, [Math 5] This is a matrix of the (i, j)th element, and the (i, j)th element matrix identifies the sub-input having the j-th maximum score of the i-th expert subnetwork. [Math 6] This is a matrix of the (i, j)th element, where the (i, j)th element matrix represents the score of the sub-input having the jth maximum score of the i-th expert subnetwork. [Number 7] The system according to claim 4, wherein is a one-hot matrix of the (i, j, m)-th element, the one-hot matrix of the (i, j, m)-th element is equal to 1 if the m-th sub-input has the j-th maximum score of the i-th expert subnetwork, and zero otherwise.

6. The system according to claim 1, wherein each sub-input is processed by an expert subnetwork with a maximum of b different threshold values.

7. For each expert subnetwork, selecting one or more of the sub-inputs according to their respective scores is: [Number 8] The calculation is to make it so that, In the formula, <S Ｔ A> is S Ｔ This represents the inner product between and A. [Number 9] Therefore, the calculation described above, G,I=TopK(A,k) P=Onehot(I) This involves calculating the following, where k is the number of sub-inputs selected by each expert subnetwork: [Number 10] This is a matrix of the (i, j)th element, and the (i, j)th element matrix identifies the sub-input having the j-th maximum score of the i-th expert subnetwork. [Math 11] This is a matrix of the (i, j)th element, where the (i, j)th element matrix represents the score of the sub-input having the jth maximum score of the i-th expert subnetwork. [Math 12] The system according to claim 6, wherein is a one-hot matrix of the (i, j, m)-th element, the one-hot matrix of the (i, j, m)-th element is equal to 1 if the m-th sub-input has the j-th maximum score of the i-th expert subnetwork, and zero otherwise.

8. For each expert subnetwork i, processing each selected sub-input using the expert subnetwork to generate each sub-output is: [Number 13] This includes calculating, in the formula, [Number 14] This is the input to the i-th expert subnetwork, W １ [i] and W ２ [i] is the weight matrix of the i-th expert subnetwork, and ACT is the activation function. [Number 15] The system according to claim 5, wherein is a matrix in the j-th row representing the sub-output for the j-th sub-input processed by the expert subnetwork i.

9. For each expert subnetwork i, processing each sub-input selected using the expert subnetwork to generate each sub-output means that one or more bias terms are applied to X. ｅ The system according to claim 8, further comprising adding to [i].

10. For each of the aforementioned multiple sub-inputs, processing the sub-outputs corresponding to the sub-inputs generated by each expert subnetwork to generate a combined sub-output for the sub-inputs is, [Number 16] This includes calculating, in the formula, [Number 17] The system according to claim 8, wherein is a matrix of the mth row representing the combined sub-outputs for the mth sub-input.

11. For each of the aforementioned multiple sub-inputs, processing the sub-outputs corresponding to the sub-inputs generated by each expert subnetwork to generate a combined sub-output for the sub-inputs is, This includes determining the weighted sum of the sub-outputs corresponding to the sub-inputs, The system according to claim 1, wherein each sub-output is weighted by the score of the sub-input to the expert subnetwork that generated the sub-output.

12. The system according to claim 1, wherein at least some of the operations of the plurality of expert subnetworks are performed in parallel.

13. The system according to claim 1, wherein at least some of the operations of the plurality of expert subnetworks are performed by different devices.

14. A method for performing the operation described in any one of claims 1 to 13 using one or more computers.

15. One or more computer storage media, which, when executed by one or more computers, store instructions that cause one or more computers to perform the operations described in any one of claims 1 to 13.

Images