Image Processing Using Mixture of Experts

Abstract

Methods, systems, and apparatuses, including computer programs encoded on a computer storage medium, for generating predictions for an image. One of the systems includes a neural network including a sequence of one or more network blocks each configured to perform operations including obtaining a block input representing an intermediate representation of an input image, determining a plurality of patches of the block input or of an updated representation of the block input, each patch including a different subset of elements of the block input or of the updated representation of the block input, assigning each patch to one or more respective expert modules of a plurality of expert modules of the network block, for each patch of the plurality of patches, processing the patch using a corresponding expert module to generate a respective module output, and combining the module outputs to generate a block output.

Description

[Technical field]

[0001] CROSS-REFERENCE TO RELATED APPLICATIONS This application claims the benefit of priority under 35 U.S.C. § 119 of U.S. Provisional Application No. 63 / 194,926, filed May 28, 2021, the entirety of which is incorporated herein by reference.

[0002] This specification relates to processing images using neural networks. [Background technology]

[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 an 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 [Means for solving the problem]

[0004] This specification describes a system implemented as a computer program on one or more computers at one or more locations that is configured to process an image using a neural network and generate a network output that characterizes the image. The neural network includes a sequence of one or more network blocks, each configured to process block inputs that include an image or an intermediate representation of the image, and generate a block output.

[0005] At least one of the network blocks is an expert network block that includes a collection of multiple different expert modules, each configured to process a proper subset of the elements of the block input. In particular, the expert network block is configured to identify multiple "patches" of the elements of the block input, and to process the patch using the routing module to identify, for each patch, a proper subset of expert modules to route the patch to. Each expert module then processes the patch to generate, for each patch routed to the expert module, a respective expert module output for the routed patch. The expert network block can, for each patch, combine corresponding expert module outputs generated by the respective expert modules in response to processing the patch to generate a respective patch output. The expert network block can then compose together the patch outputs corresponding to each patch of the block input to generate a block output, the patch outputs having the same spatial relationship to one another in the block output as the corresponding patches had in the block input.

[0006] The subject matter described herein can be implemented in particular embodiments to realize one or more of the following advantages.

[0007] Some existing systems process images using "dense" feed-forward neural networks, i.e., feed-forward neural networks that, for each neural network layer of the feed-forward neural network, use every network parameter of the neural network layer to process every element of the layer input to the neural network layer. With large models configured to process high-resolution images, such dense neural networks may require significant time and computational resources to process a single image. Using the techniques described herein, a system can process images using a feed-forward neural network by selectively activating a subset of the neural network's parameters based on the network inputs, greatly improving the time and computational efficiency of processing images. Furthermore, introducing this sparsity may allow feed-forward neural networks to include many more network parameters than was previously feasible, since only a subset of the parameters is used to process any given input.

[0008] Using the techniques described herein, a system can implement neural networks that include expert network blocks with multiple expert modules, where each expert module is configured through training to process a different type of image (or different type of image patch), allowing the modules to "specialize" and further improving the efficiency and performance of the neural network.

[0009] As described herein, a self-attention based neural network configured to process images may require much less computation to achieve the same performance as a state-of-the-art convolutional neural network. In other words, for a given compute budget, a self-attention based neural network will perform better (e.g., as measured by test accuracy, precision, recall, etc.) than a convolutional neural network configured to perform the same machine learning task on images. This is because applying self-attention is generally more computationally efficient than convolving a kernel over the entire image, since the self-attention mechanism can direct attention to different regions of the image with fewer computations than convolution.

[0010] The details of one or more embodiments of the subject matter herein 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]

[0011] [Figure 1A] FIG. 2 is a diagram of an example neural network configured to generate predictions for images. [Figure 1B] FIG. 2 is a diagram of an example neural network block including multiple expert modules. [Diagram 2] FIG. 2 illustrates an example image segmented into image patches. [Diagram 3] 1 is a flow diagram of an example process for executing a network block that includes multiple expert modules. DETAILED DESCRIPTION OF THE PREFERRED EMBODIMENTS

[0012] Like reference numbers and designations in the various drawings indicate like elements.

[0013] This specification describes a system, implemented as a computer program on one or more computers at one or more locations, that is configured to compile computer program modules using a just-in-time compiler and a compilation cache.

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

[0015] The neural network 100 is configured to process the input image 102 according to a set of trained network parameters to generate a network output 108 that represents an embedding of or a prediction about the input image 102. The neural network 100 may be configured to perform any suitable machine learning task using the input image 102. Exemplary machine learning tasks are discussed below. As used herein, processing an image refers to processing intensity values ​​of pixels of the image.

[0016] Neural network 100 includes a sequence of L network blocks 110a-l, where L > 1. Each network block 110a-l may be configured to process a block input that includes an image or an intermediate representation of an image, and to generate a block output.

[0017] 1A, a first network block 110a in the sequence of network blocks 110a-l is configured to process the image 102 itself to generate a first block output 104 that is an intermediate representation of the image 102. Each subsequent network block 110b-l may then process the block output of the previous network block in the sequence, i.e., the block input of the (i+1)th network block 110a-n is the block output of the ith network block 110a-l.

[0018] In some other implementations, the first network block 110a is configured to process an embedding of the input image 102 instead of the input image 102 itself. For example, although the following description generally refers to a neural network processing the image 102, in some implementations the neural network 100 may be configured to receive as input an embedding of the image 102, e.g., an embedding generated by a different neural network in response to processing the image 102.

[0019] In some implementations, the network output 108 of the neural network 100 is a block output of the last network block 110l in the sequence of network blocks 110a-l. For example, as described in more detail below, the block output of the last network block 110l may include a respective updated representation of each of a plurality of patches of the input image 102, and thus the block outputs may collectively represent an embedding of the input image 102 that may be stored or further processed, for example, by one or more downstream machine learning models.

[0020] In some other implementations, the block output of the last network block 110a-l in the sequence is further processed using one or more neural network layers to generate a network output 108 of the neural network 100. For example, if the neural network 100 is configured to classify the input image 102, the neural network 100 may process the block output of the last network block 110a-l in the sequence using one or more neural network layers, including a final logistic or softmax layer, to generate a network output 108 that identifies, for each class in a set of candidate classes, a value representing the likelihood that the input image 102 belongs to the class.

[0021] As a particular example, neural network 100 may include one or more "head" sub-networks, each configured to process the block output of the last network block 110a-l of a sequence using one or more neural network layers to generate predictions for the input image 102 for a respective different machine learning task. That is, after neural network 100 is trained, the output of the last network block 110a-l of a sequence (which may be viewed as an embedding of input image 102) may encode information useful for multiple different machine learning tasks.

[0022] In some implementations, the neural network 100 first processes the input image 102 to determine a set of patches of the input image, e.g., to determine an input sequence including a plurality of elements representing the determined patches of the input image 102. The neural network 100 can then use a sequence of network blocks 110a-l to process the determined patches (or embeddings of the determined patches generated by one or more initial neural network layers of the neural network 100 or by different neural networks). In some other implementations, one or more of the network blocks 110a-l each determine a set of patches of block inputs to the network blocks 110a-l (which are intermediate representations of the input image 102) and process the determined patches to generate block outputs of the network blocks 110a-l. The operation of the network blocks 110a-l is described in more detail below.

[0023] In some implementations, the neural network 100 processes one or more other network inputs in addition to the input image 102. For example, the neural network 100 may process data identifying a class to which the input image 102 belongs. As another example, the neural network 100 may process a text input that corresponds to the input image 102, e.g., a text input representing text that characterizes the input image 102.

[0024] As another example, in an implementation in which, for example, the neural network 100 determines a sequence of patches (or their embeddings) of the input image 102 that are successively updated by the sequence of network blocks 110a-l, the neural network 100 may process a predefined token, sometimes referred to as a "class" token. That is, the predefined token is the same for each input image 102 received by the neural network 100, and may be, for example, a predefined tensor. For example, the predefined token may be machine-learned, for example, simultaneously with the training of the neural network 100. As a particular example, the predefined token may be appended to the sequence of patches processed by the neural network 100, such that each network block of the sequence of network blocks 110a-l generates an updated representation of the predefined token. In some implementations, the network output 108 is (or includes) the final representation of the predefined token generated by the last network block 110l, and the final representation may be considered as an embedding of the input image 102.

[0025] Some or all of the network blocks 110a-l include a respective collection of multiple different expert modules. As used herein, an expert module is a component of a network block that is configured to process a true subset of the elements of the block input to the network block, i.e., a portion of the elements. The expert module may be configured to process the true subset of the elements of the block input using one or more neural network layers to generate an updated representation of the true subset of the elements. As used herein, a network block that includes multiple expert modules is also referred to as an expert network block.

[0026] For example, each network block 110a~l of the sequence can be an expert network block. As another example, every other network block 110a~l of the sequence can be an expert network block. As another example, every other network block 110a~l of the sequence up to a specific network block 110i where i < l can be an expert network block.

[0027] For example, as shown in FIG. 1, the j-th network block 110j of the sequence of network blocks 110a~l is an expert network block including e expert modules 130a~e, where e > 1. The following description particularly refers to the j-th network block 110j, but it should be understood that the same technique can be implemented for each expert network block of the neural network 100.

[0028] The j-th network block can segment its block input 112 into a plurality of patches, and each patch includes a different subset of the elements of the block input 112. For example, if the block input 112 is a three-dimensional representation of the input image 102, the j-th network block 110j can determine that each patch has dimensions L×W×C, where C represents the number of channels of the representation of the input image 102 (for example, if the block input 112 is an RGB image, C = 3, and if the block input 112 is an intermediate representation of an RGB image, it is a larger number).

[0029] In some implementations, each patch of the block input 112 is of the same size, that is, includes the same number of elements and has the same dimensions. In some other implementations, different patches can be of different sizes, for example, can include different numbers of elements.

[0030] In some implementations, each element of the block input 112 is in exactly one patch. In some other implementations, some or all of the elements of the block input 112 can be in multiple different patches, i.e., the patches of the block input 112 can overlap.

[0031] In some implementations, each network block 110a-l of the sequence of network blocks processes the same set of patches of the input image 102. That is, each block input and block output may have the same dimensions, such that the number of patches, patch sizes, and spatial relationships of the patches determined by each network block 110a-l are the same. For example, the neural network 100 may process the input image 102 to determine an input sequence of patches of the input image 102, and then use each network block 110a-l of the sequence of network blocks to process a respective updated representation of the determined sequence of patches.

[0032] The patches of the block input 112 may be represented in any suitable manner. For example, each patch may be represented as a two-dimensional tensor that includes the elements of the block input 112, e.g., a tensor that preserves the spatial relationships of the elements of the block input 112. As another example, each patch may be represented as a one-dimensional sequence of the elements of the patch. As a particular example, if a patch is a two-dimensional region of the block input 112, the patch may be, for example, a flattened version of the two-dimensional region, where the elements of the block input 112 are ordered in a one-dimensional tensor according to their respective positions within the block input 112. As another example, each initial image patch 112a-n may be represented as an unordered set of pixels of the initial image patch 112a-n.

[0033] An exemplary technique for determining patches of an image or an intermediate representation of an image is discussed below with reference to FIG.

[0034] After determining the patches of the block inputs 112, the jth network block 110j can assign each patch to one or more of the expert modules 130a-e of the jth network block 110j using the routing module 120. The jth network block can then process each image patch using the corresponding assigned expert modules 130a-e to generate a respective expert module output, and combine the expert module outputs of the expert modules 130a-e using the combination module 140 to generate the block output 114 of the jth network block 110j.

[0035] In other words, the routing module 120 is configured to determine, for each expert module 130a-e, a corresponding set 122a-e of patches of the block input 112 that should be routed to the expert module 130a-e for processing. Each expert module 130a-e then processes its assigned set of patches 122a-e to generate a corresponding set 132a-e of expert module outputs, i.e., a respective expert module output for its respective assigned patch. The expert module outputs may also be referred to simply as "module outputs."

[0036] To assign patches to respective expert modules 130a-e, the routing module 120 may process each patch using a routing neural network including one or more neural network layers. In response to processing a particular patch, the routing neural network may generate an output including a respective routing score of each expert module 130a-e of the j-th network block 110j. The routing module 120 may determine to assign the patch to the one or more expert modules 130a-e corresponding to the highest routing score of the patch.

[0037] In some implementations, each patch is assigned to the same number of expert modules 130a-e, i.e., each patch may be assigned to the k expert modules 130a-e that have the highest routing scores for the patch, where k is the same for all patches.

[0038] In some other implementations, some patches may be assigned to a different number of expert modules. For example, a patch may be assigned to an expert module 130a-e if the routing score of the expert module and the patch meets a predetermined threshold. As another example, each expert module 130a-e may be guaranteed to process exactly p patches by assigning the p patches with the highest routing score corresponding to the expert module among all the routing scores corresponding to the expert module and the respective patch, where p≧1.

[0039] In some implementations, not every expert module 130a-e of the jth network block 110j processes a patch of the block input 112, i.e., for some input images 102, some expert modules 130a-e of the jth network block 110j are idle for block inputs 112 generated from the input images 102. Allowing some expert modules 130a-e to be idle during some executions of the neural network 100 can improve the efficiency of the neural network 100 by reducing the number of calculations required to generate the network output 108, since not every parameter of the neural network 100 is used to process each input image 102.

[0040] The routing neural network of the routing module 120 may be configured, through training, to assign each patch to an expert module 130a-e that extracts the most information from the patch. That is, in some implementations, at least some of the expert modules 130a-e may be "specialized" for particular types of patches, e.g., patches that depict particular semantic objects or categories of semantic objects. Exemplary techniques for training the neural network 100 are discussed in more detail below.

[0041] In some implementations, the jth network block 110j limits the number of patches that each expert module 130a-e can process, so that no one expert module 130a-e is assigned a disproportionate amount of work (i.e., computation) relative to the other expert modules and thus becomes a bottleneck in the execution of the neural network 100. For example, the network block 110j limits the number of patches that can be assigned to any one expert module 130a-e.

[0042]

number

[0043] It can be restricted to the following:

[0044] where B is the number of images in the batch of input images 102 being processed by neural network 100 (e.g., the batch of images 102 provided to neural network 100 during inference and processed in parallel by neural network 100), K is the number of expert modules 130a-e to which each patch is assigned, P is the number of patches in block input 112, and E is the number of expert modules 130a-e in network block 110j. C may be a hyperparameter of neural network 100 representing the capacity of each expert module 130a-e, where C=1 indicates the capacity of each expert module 130a-e to assign a number of patches to each expert module 130a-e during a given run of neural network 100.

[0045]

number

[0046] indicates that only

[0047] In some implementations, the neural network 100 sets, for example, C<<1, e.g., C=0.5, 0.1, 0.05, or 0.01.

[0048] In some implementations, if more image patches are allocated to an expert module 130a-e than allowed, the patches that exceed the limit (e.g., the patches allocated to the expert module 130a-e with the lowest routing score corresponding to the expert module 130a-e) are reallocated, for example, to the expert module 130a-e with the next highest routing score for each patch. In some other implementations, the patches that exceed the limit are not reallocated, but rather are "dropped" so that they are processed by one less expert module 130a-e than the other patches, which is equivalent to processing the dropped patch using an expert module that is simply an identity function.

[0049] In some implementations, the neural network 100 operates with different values ​​of C during different runs of the neural network 100, i.e., the value of C may be changed after the neural network 100 is trained. For example, during training of the neural network 100, the value of C may be relatively high, so that relatively few patches are dropped, and the training system may use the module outputs of the expert modules 130a-e to determine updates to the network parameters of the neural network layers of the expert modules 130a-e. Then, during inference, the value of C may be lowered to improve the time and computational efficiency of the neural network 100. As a particular example, the neural network 100 may be deployed in an inference environment with scarce resources, for example, on a device such as a mobile phone or tablet that has relatively few computational resources, or in an environment where the neural network 100 must continually receive new images 102 and generate corresponding network outputs 108 using a fixed time budget.

[0050] In some implementations, for each patch of the block input 112, the routing module 120 randomly samples (e.g., from a normal distribution) a noise value for each expert module 130a-e and adds the sampled noise value to the corresponding routing score before determining the k highest routing scores for the patch. For example, the routing module 120 can add noise values ​​during training of the neural network 100 to promote diversity and thus encourage each expert module 130a-e to receive parameter updates. The routing module 120 can be configured not to add noise values ​​during inference since it was trained to allocate high-quality routing scores.

[0051] In some implementations, the network block may apply a nonlinear activation function, such as a softmax, Tanh, or ReLU function, to the routing scores before determining the k highest routing scores.

[0052] The routing neural network may include any suitable configuration of neural network layers.

[0053] For example, the routing neural network may include one or more feedforward neural network layers. As a particular example, if a patch has dimensions L×W×C, the routing module 120 may reshape the patch to have dimensions 1×(L·W·C) and process the reshaped patch using a feedforward neural network layer.

[0054] Alternatively or additionally, the routing neural network may include one or more convolutional neural network layers. As a particular example, if a patch has dimensions L1×W1×C, the routing neural network may perform a two-dimensional convolution on the patch using a convolution kernel having dimensions L2×W2×C, where L1>L2 and W1>W2.

[0055] Alternatively or additionally, the routing neural network may include one or more self-attention layers, which are discussed in more detail below with reference to FIG.

[0056] In some implementations, the operations of the e expert modules 130a-e are performed in parallel, thus improving the time efficiency of the neural network 100.

[0057] In some implementations, each expert module 130a-e executes on the same processing device, e.g., parallel processing hardware configured to execute at least a portion of the expert modules 130a-e on respective different threads, e.g., a graphics processing unit (GPU) or a tensor processing unit (TPU). In some other implementations, each expert module 130a-e may be executed by a different processing device, e.g., a distributed system of multiple different parallel processing devices. Executing at least a portion of the expert modules 130a-e in parallel can significantly reduce the latency of execution of the neural network 100, e.g., by a factor of approximately e if all expert modules 130a-e were executed on different devices. Thus, for a given number of network parameters, a neural network that includes at least one expert network block of an expert module can enjoy reduced latency compared to a neural network with only a dense neural network layer, because (i) each element of the input image is processed by a true subset of the network parameters of the expert module, and (ii) in some implementations, the execution of the expert module can be parallelized, while the dense neural network layer cannot. Similarly, for the same reason, given a particular time budget (i.e., for a particular latency threshold), a neural network that includes at least one expert network block of an expert module can include significantly more network parameters, and thus increased network capacity (i.e., ability to encode more information), compared to a neural network with only a dense neural network layer.

[0058] In some implementations where the expert modules 130a-e are distributed across multiple devices, the processing device executing the routing module 120 is configured to provide only the set of patches 122a-e allocated to a particular expert module 130a-e to the processing device executing the particular expert module 130a-e. In some other such implementations, the processing device executing the routing module 120 is configured to provide the entire block input 112 to the processing device executing a particular expert module 130a-e along with data specifying which patches of the block input 112 should be processed by the particular expert module 130a-e.

[0059] After each patch has been processed by the corresponding one or more expert modules to generate a respective module output, the combination module 140 of the network block 110j can combine the module outputs corresponding to each patch to generate the block output 114 of the network block 110j. In particular, the combination module 140 can, for each patch of the block input 112, combine the corresponding module outputs generated by the respective expert modules 130a-e in response to processing the patch to generate a patch output for the patch. The combination module 140 can then reconfigure the patch outputs to generate the block output 114 such that the patch outputs have the same spatial relationship to one another at the block output 114 as the corresponding patches had at the block input 112.

[0060] For example, in some implementations, the patch outputs have the same dimensions as the corresponding patches. Thus, the combination module can configure the patch outputs to generate block outputs 114 having the same dimensions L×W×C as the block inputs 112, with each patch output occupying the same position in the block output 114 as its corresponding patch occupied in the block input 112.

[0061] In some implementations, for each patch of the block input, the combination module 140 combines the corresponding module outputs by calculating a sum of the module outputs of the patch. In some such implementations, to generate a patch output for each patch, the combination module 140 can weight each module output using a routing score corresponding to the patch and the expert module 130a-e that generated the module output. In some other implementations, the combination module 140 combines the module outputs of the expert modules 130a-e by processing the module outputs using one or more neural network layers, e.g., one or more self-attention layers or one or more recurrent neural network layers.

[0062] In some implementations, before allocating patches of the block inputs 112 to the respective expert modules 130a-e, the network block 110j first processes the block inputs using one or more neural network layers to generate updated representations of the block inputs 102, and then allocates the patches of the updated representations of the block inputs 102 to the respective expert modules 130a-e. For example, the network block 110j can first process the block inputs using one or more self-attention neural network layers and / or one or more element-wise forward propagation neural network layers configured to process each patch of the block inputs 102. In other words, the inputs to the expert modules can be a true subset of elements of the updated representations of the block inputs 102, rather than of the block inputs 102 themselves. Generally, in this specification, we refer to patches of the block inputs 102, but it should be understood that the same procedure can be applied to patches of the updated representations of the block inputs 102. Similarly, one or more neural network layers preceding the routing module 120 may be considered to be part of different network blocks in the sequence of network blocks 110a-l.

[0063] The block output 114 of the j-th network block 110j may be provided as a block input to a subsequent network block in the sequence of network blocks 110a-l.

[0064] As mentioned above, in some implementations, the block output of the last network block 110l in the sequence is the network output 108. For example, if the block output of the last network block 110l includes a respective representation of each of multiple patches of the input image 102, the network output 108 may include a representation of each patch. As another example, the network output 108 may include a combined representation generated from the representations of each patch, e.g., by determining an average of the patch representations or by applying a pooling operation (e.g., average pooling, max pooling, or global pooling).

[0065] As another example, as mentioned above, in some implementations, the neural network 100 processes the given class token together with the input image 102 such that each network block in the sequence of network blocks generates an updated representation of the given class token in addition to an updated representation of the patch. In these implementations, the network output 108 can be the final representation of the given class token generated by the last network block 110l.

[0066] Neural network 100 may be trained using any suitable technique. For example, a training system may process training images using neural network 100 to generate respective training outputs and determine an error in the training outputs using, for example, ground-truth outputs corresponding to the training images. The training system may then backpropagate the error through neural network 100 and use gradient descent to determine updates to parameters of neural network 100.

[0067] Alternatively or additionally, the training system can train the neural network 100 using an "importance loss" that encourages, for each expert network block in the sequence of network blocks 110a-l, the routing scores corresponding to each expert module 130a-e in the expert network block to have the same mean. For example, for each expert module i, the training system can

[0068]

number

[0069] can be calculated.

[0070] In the formula, f i is a function defined by the routing neural network of the routing module 120 of the expert network block for calculating the routing score of the expert module i or the intermediate routing score of the expert module i. For example, f i (x) may be the routing score of expert module i and patch x before random noise is added and before the nonlinear activation function is applied. Then, for each expert module i, the training system performs imp i Importance loss L to encourage (X) to have the same meanimp (X). As a specific example, the training system can compute

[0071]

number

[0072] can be calculated.

[0073] where E is the number of expert modules.

[0074] Alternatively or additionally, the training system can train the neural network 100 using a "load loss" that encourages, for each expert network block and for each patch, each expert module 130a-e of the expert network block to have the same mean likelihood of being assigned a patch. For example, for each expert module i and for each image x, the training system can

[0075]

number

[0076] can be calculated.

[0077] In the formula, g i (x) is a function defined by the routing neural network of the routing module 120 of the expert network block for calculating the routing score of the expert module i or the intermediate routing score of the expert module i. For example, g i (x) may be the routing score for expert module i and patch x after random noise has been added but before the nonlinear activation function is applied. Then, for each training image x and for each expert module i, the training system calculates p i (x) = P(gi (x)≧threshold k (x) can be calculated.

[0078] where P(e) is the probability that event e occurs. Given the distribution of random noise, p i (x) can be directly computed. Then, for each expert module i, the training system

[0079]

number

[0080] Calculate

[0081] For each expert module i, load i (x) to encourage the load loss L to have the same average load (X). As a specific example, the training system can compute

[0082]

number

[0083] can be calculated.

[0084] where E is the number of expert modules.

[0085] In some implementations, the neural network 100 may be pre-trained using a first machine learning task to generate initial values ​​for the network parameters of the neural network 100. The initial values ​​may then be fine-tuned (i.e., updated) using a second machine learning task, e.g., simultaneously with one or more head neural networks each configured to process the network output 108 to generate a final prediction for the respective second machine learning task. For example, the training system may pre-train the neural network 100 in an unsupervised or self-supervised manner using unlabeled images, e.g., to train the neural network 100 to segment unlabeled images into different classes based on their similarity, or to train the neural network 100 to perform unsupervised semantic segmentation of unlabeled images. The training system can then fine-tune the network parameters of neural network 100 in a supervised manner (e.g., together with a new untrained head neural network) using a task-specific dataset of labeled images to perform a particular task, such as object recognition, pose detection, face recognition, optical flow tracking, or any other suitable machine learning task.

[0086] The image 102 can be any suitable type of image. For example, the image 102 can be a two-dimensional image, e.g., a two-dimensional image having multiple channels (e.g., an RGB image). As another example, the image 102 can be a hyperspectral image that represents a continuous spectrum of wavelengths, e.g., by determining a spectral distribution for each pixel of the image 102. As another example, the image 102 can be a point cloud that includes a plurality of points, each point having, e.g., a respective coordinate in a three-dimensional or higher-dimensional coordinate space, and as a particular example, the image 102 can be a point cloud generated by a LIDAR sensor. As another example, the image 102 can be a medical image generated by a medical imaging device, and as a particular example, the image 102 can be a computed tomography (CT) image, a magnetic resonance imaging (MRI) image, an ultrasound image, an X-ray image, a mammogram image, a fluoroscopy image, or a positron emission tomography (PET) image.

[0087] Although the description below refers to generating image patches of image 102, each including a respective "pixel" of image 102, it should be understood that neural network system 100 may generate image patches including any suitable type of components of image 102. For example, if image 102 is a point cloud, each image patch of image 102 may include a subset of the points of the point cloud. As another example, if image 102 is an MRI image including multiple voxels in a three-dimensional voxel grid, each image patch of image 102 may include a subset of the voxels in the voxel grid.

[0088] The neural network may be configured to process images to perform any suitable image processing task, for example a classification task, a regression task, or a combination thereof.

[0089] As a particular example, the neural network 100 may be configured to generate a classification output that includes a respective score corresponding to each of a plurality of categories. The score for a category indicates the likelihood that the image belongs to that category. In some cases, the categories may be classes 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 categories 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.

[0090] As another particular example, the neural network 100 may be configured to generate an element-level classification output (e.g., a pixel-level classification output for an RGB image or a point-level classification output for a LIDAR image) that includes, for each element in an image, a respective score corresponding to each of a plurality of categories. For a given element (e.g., a given pixel or point), the score of a category indicates the likelihood that the element belongs to that category. In some cases, the categories may be classes of objects, and an element may belong to a category if it is part of an object included in the object class that corresponds to that category. That is, the element-level classification output may be a semantic segmentation output.

[0091] As another particular example, the neural network 100 may be configured to generate a regression output that estimates one or more continuous variables (i.e., which may assume infinitely many possible numerical values) that characterize an image. In a particular example, the regression output may estimate the coordinates of a bounding box that encloses each object depicted in the image. The coordinates of the bounding box may be defined by the (x, y) coordinates of the vertices of the bounding box.

[0092] In some implementations, the neural network 100 may be configured to process multiple images, e.g., multiple frames of a video. For example, the neural network can receive multiple images that are video frames of a video and can process each video frame as described above to generate an output that characterizes the video frame, e.g., by characterizing whether the video frame depicts a person performing a particular action.

[0093] In some such implementations, the neural network 100 processes each video frame at each different time to generate a respective network output for each video frame that characterizes a prediction for the video frame. For example, the neural network can generate a network output that predicts a classification of the video frame. In some such implementations, the neural network combines multiple network outputs corresponding to each video frame to generate a final network output that characterizes the video. For example, the neural network can process each network output using a downstream neural network, e.g., a recurrent neural network.

[0094] In some other implementations, the neural network 100 processes each video frame in parallel to generate a single network output that characterizes the video. As a particular example, the system can generate one or more respective input elements in the input sequence for each video frame.

[0095] 1B is a diagram of an example neural network block 150 that includes multiple expert modules 180. Neural network block 150 is an example of a system, implemented as a computer program on one or more computers at one or more locations, in which the systems, components, and techniques described below may be implemented.

[0096] Network block 150 may be one of a sequence of L network blocks of a neural network configured to process an input image to generate a network output that characterizes the input image, e.g., represents an embedding of the input image or represents a prediction about the input image, where L ≧ 1. For example, network block 150 may be one of the network blocks of the sequence of network blocks 110a-l described above with reference to FIG. 1A.

[0097] Network block 150 is configured to process a set of embedded patches 152 of the input image, i.e., updated representations of the patches determined from the input image generated by one or more preceding network blocks in the sequence of network blocks, to generate block output 156. If network block 150 is the first network block in the sequence, network block 150 may be configured to process the original patches determined from the input image.

[0098] The network block 150 includes one or more self-attention neural network layers. As shown in FIG. 1B, the network block 150 includes a single self-attention neural network layer 160. In some implementations, the self-attention neural network layer 160 is configured to take the embedded patch 152 and apply an attention mechanism across the embedded patch 152. In some other implementations, the self-attention neural network layer 160 is configured to take a processed version of each of the embedded patches 152 and apply an attention mechanism to the processed embedded patch 152. For example, as shown in FIG. 1B, the network block 150 can first apply the layer normalization layer 154 to the embedded patch 152 before providing an output of the layer normalization layer 154 to the self-attention neural network layer 160. Alternatively or in addition, the network block 150 can apply one or more other neural network layers, e.g., one or more element-wise forward propagation neural network layers, to the embedded patch 152 before the self-attention neural network layer 160.

[0099] For example, the network block 150 may generate, for each embedded patch 152, a one-dimensional tensor that includes elements of the embedded patch 152. If each embedded patch 152 has dimensions L1×W1×C, the network block 150 may generate a one-dimensional tensor with dimensions 1×(L·W·C). The network block 150 may then generate an input sequence from the one-dimensional tensor, where at least some of the input positions of the input sequence are one-dimensional tensors, and the network block 150 may process the input sequence using the self-attention neural network layer 160.

[0100] In particular, for each embedded patch 152 (or a processed version thereof) corresponding to each particular input position, the self-attention neural network layer 160 is configured to apply an attention mechanism to the embedded patch 152 of the input position (i.e., other embedded patches 152, and in some implementations, itself) using one or more queries derived from the embedded patch 152 of the particular input position to generate a respective output for the particular position. The output of the self-attention neural network layer 160 is an output sequence including a respective output element corresponding to each input position.

[0101] In some implementations, some or all of the self-attention neural network layers of network block 150 (e.g., self-attention neural network layer 160 shown in FIG. 1B ) are multi-head self-attention neural network layers that apply h different attention mechanisms in parallel to generate respective sequences of output elements, and then combine the multiple sequences of output elements to generate a final sequence of output elements.

[0102] In some implementations, some or all of the self-attention neural network layers of the network block 150 (e.g., the self-attention neural network layer 160 shown in FIG. 1B ) incorporate positional information of each embedded patch 152 of the input sequence into the attention mechanism. For example, when applying attention with respect to a particular embedded patch 152 (i.e., when generating each output element corresponding to the particular embedded patch 152), the self-attention neural network layer 160 can identify an attention positional embedding that represents the position of the embedded patch 152 corresponding to the image (or within the block input).

[0103] As used herein, an embedding is an ordered collection of numbers that represent the input in a particular embedding space. For example, an embedding can be a vector of floating-point or other numbers with a fixed dimension.

[0104] In some implementations, the network block 150 includes a residual connection layer 162 that combines the output of the self-attention neural network layer 160 with the input to the self-attention neural network layer 160. Alternatively or in addition, the network block 150 may include a layer normalization layer 164 that applies layer normalization to the output of the self-attention neural network layer 160 (or the residual connection layer 162).

[0105] An output of the self-attention neural network layer 160 (or the residual connection layer 162 or the layer normalization layer 164) includes updated representations of the embedded patches 152. The updated representations of the embedded patches 152 may be provided to a routing module 170 configured to allocate, for each updated representation of each embedded patch 152, the updated representation to one or more expert modules 180 of the collection of multiple expert modules of the network block 150. The routing module 170 may be configured similarly to the routing module 120 described above with reference to FIG. 1A.

[0106] The network block can provide, for each expert module 180, an allocated updated representation of each embedded patch 152 for the process of generating a respective expert module output for the respective allocated updated representation. The expert modules 180 can be configured similarly to the expert modules 130a-e described above with reference to FIG. 1A.

[0107] 1B, each expert module 180 is a multi-layer perceptron (MLP) that includes one or more feedforward neural network layers. For example, an expert module 180 may include two or more learned linear transforms that are each separated by an activation function, e.g., a nonlinear element-wise activation function, e.g., a ReLU activation function.

[0108] The expert module outputs may then be combined by a combination module 190 of the network block 150 to generate the block output 156 of the network block. The combination module 190 may be configured similarly to the combination module 140 described above with reference to FIG. 1A.

[0109] In some implementations, the network block 150 includes a residual connection layer 192 that combines the output of the combination module 190 with the input to the routing module 170. That is, the block output 156 can be the output of the residual connection layer 192 rather than the combination module 190.

[0110] The block output 156 may be provided to one or more subsequent network blocks in the sequence of network blocks for further processing to generate a network output of the neural network.

[0111] After generating the network output, the neural network can provide the network output 108 to one or more downstream systems. For example, the neural network can provide the network output to one or more head neural networks to generate predictions for the respective machine learning tasks, as described above with reference to FIG. 1A. As another example, the neural network can provide the network output, which may represent embeddings of the input image, to a database or one or more downstream machine learning models for further processing.

[0112] FIG. 2 shows example images 210, 220, 230, 240, and 250 that have been segmented into image patches.

[0113] The images 210-250 may be provided as inputs to a neural network system, such as the neural network 100 described above with reference to FIG. 1A, configured to process the images 210-250 to generate predictions for the images 210-250. The neural network system may be configured to segment the images 210-250 into a number of image patches and process the image patches using a sequence of network blocks including one or more expert network blocks.

[0114] Although the following description refers to determining patches of an image, it should be understood that the same techniques may be used to determine patches of an intermediate representation of an image, e.g., patches of block inputs to a network block, such as block input 112 described above with reference to FIG. 1A.

[0115] Images 210, 220, 230, 240, and 250 illustrate different possibilities for segmenting an image into image patches. In particular, in FIG. 2, each image 210-250 is depicted as being segmented into a collection of image patches that are visually different, i.e., use different shading or hatching. In general, the neural network system is configured to segment all received images according to the same schema. That is, the same neural network system would not necessarily segment each of images 210, 220, 230, 240, and 250 as depicted, since the images are segmented according to different schemas.

[0116] As shown in the first image 210, in some implementations, the neural network system can generate image patches each having the same size and shape, e.g., each image patch can be rectangular. Furthermore, in some implementations, the neural network system can segment the first image 210 such that every pixel is a component of exactly one image patch. As a particular example, the image patches can represent rectangular grids of the same size, as shown in FIG. 2. As another particular example, the image patches can represent hexagonal grids of the same size.

[0117] As shown in the second image 220, in some implementations, the neural network system can generate image patches having different sizes.

[0118] As shown in the third image 230, in some implementations, the neural network system can segment the third image 230 such that some pixels are components of multiple different image patches. For example, a first image patch can include a true superset of pixels of a second image patch, i.e., the second image patch can be completely contained within the first image patch. As a particular example, for each of a plurality of first image patches determined from the image (e.g., for each first image patch in a grid of first image patches), the neural network system can generate a respective second image patch having the same center pixel but different sizes and / or aspect ratios.

[0119] As shown in the fourth image 240, in some implementations, the neural network system can segment the fourth image 240 such that some pixels are not components of any image patch. For example, the neural network system can process the fourth image 240 using a machine learning model to identify one or more regions of interest, and the neural network system can generate a respective patch for each identified region of interest. For example, the machine learning model can be configured to identify one or more pixels, and the neural network system can generate a respective patch centered on each identified pixel.

[0120] As shown in the fifth image 250, in some implementations, the neural network system can generate image patches of any shape. That is, the image patches need not be rectangular. For example, the neural network system can process the fifth image 250 using a machine learning model configured to segment the fifth image 250, e.g., by assigning a respective class to each pixel of the fifth image 250. The neural network system can then generate a respective patch for each successive set of pixels that are assigned the same class by the machine learning model.

[0121] Although images 210-250 are shown in FIG. 2 as two-dimensional images (or two-dimensional images with multiple channels, e.g., RGB images), in general, the neural network system may be configured to generate predictions for any type of image, as described above with reference to FIG. 1A.

[0122] 3 is a flow diagram of an exemplary process for executing a network block including multiple expert modules. For convenience, the process 300 is described as being executed by one or more computer systems located at one or more locations. For example, a neural network system appropriately programmed in accordance with this specification, such as the neural network 100 described above with reference to FIG. 1A, may execute the process 300.

[0123] The network block may be one network block of a sequence of network blocks of a neural network configured to process an input image and generate a network output related to the input image.

[0124] The system obtains a block input that represents an intermediate representation of an input image (step 302).

[0125] The system determines a set of multiple patches of the block input (or an updated representation of the block input, e.g., an updated representation generated by applying the self-attention mechanism described above with reference to FIG. 1B), each patch including a different subset of elements of the block input (or the updated representation of the block input) (step 304).

[0126] The system assigns each patch to one or more respective expert modules of the network block (step 306).

[0127] For each patch of the plurality of patches, the system processes the patch using one or more corresponding expert modules to generate one or more respective module outputs (step 308).

[0128] The system generates a block output for the network block by combining the module outputs generated by each expert module (step 310).

[0129] This specification uses the term "configured" in reference to systems and computer program components. One or more computer systems being configured to perform a particular operation or action means that the system has installed thereon software, firmware, hardware, or a combination thereof that, during operation, causes the system to perform the operation or action. One or more computer programs being configured to perform a particular operation or action means that one or more programs contain instructions that, when executed by a data processing device, cause the device to perform the operation or action.

[0130] The embodiments and functional operations of the subject matter described herein may be implemented in digital electronic circuitry, tangibly embodied computer software or firmware, computer hardware, including the structures disclosed herein and their structural equivalents, or a combination of one or more of them. The embodiments of the subject matter described herein may 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 for controlling the operation of a data processing device. The computer storage medium may be a machine-readable storage device, a machine-readable storage substrate, a random or serial access memory device, or a combination of one or more of them. Alternatively or additionally, the program instructions may be encoded on an artificially generated propagated signal, e.g., a machine-generated electrical, optical, or electromagnetic signal, generated to encode information for transmission to a receiver device suitable for execution by a data processing device.

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

[0132] 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, or 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 as part of a file that holds other programs or data, for example one or more scripts stored in a markup language document, a single file dedicated to the program in question, or in multiple organized files, for example files that store one or more modules, subprograms, or code portions. A computer program may be deployed to be executed on one computer or on multiple computers that are located at one location or distributed across multiple locations and connected together by a data communication network.

[0133] As used herein, the term "database" is used broadly to refer to any collection of data, which 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.

[0134] 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 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.

[0135] 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 performing operations on input data and generating output. The processes and logic flows may also be performed by special purpose logic circuitry, e.g., FPGAs or ASICs, or a combination of special purpose logic circuitry and one or more programmed computers.

[0136] A computer suitable for executing a computer program can be based on a general-purpose or special-purpose microprocessor or both, or any other type of central processing unit. In general, the central processing unit receives instructions and data from a read-only memory or a random access memory or both. The essential elements of a computer are a central processing unit for executing or 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 in special-purpose logic circuitry. In general, the 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 coupled to receive data from or transfer data to the mass storage devices, or both. However, the computer may not have such devices. Additionally, a computer may be incorporated into another device, 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 just a few.

[0137] Computer-readable media suitable for storing computer program instructions and data include, by way of example, all forms of non-volatile memory, media, and memory devices including 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 disks and DVD-ROM disks.

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

[0139] A data processing apparatus for implementing machine learning models may also include dedicated hardware accelerator units, for example for handling the often computationally intensive parts of the machine learning training or generation, i.e., inference, workload.

[0140] The machine learning model may 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.

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

[0142] 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 having a client-server relationship to each other. In some embodiments, a server sends data, e.g., HTML pages, to a user device for the purpose of, e.g., displaying data to a user interacting with the device acting as a client and receiving user input from such a user. Data generated at the user device, e.g., results of user interaction, may be received from the device at the server.

[0143] In addition to the above-mentioned embodiments, the following embodiments are also innovative.

[0144] Embodiment 1 is a system including a neural network configured to process an input image and generate a network output related to the input image, the neural network including a sequence of one or more network blocks, the one or more network blocks comprising: obtaining a block input representing an intermediate representation of an input image; determining a plurality of patches of the block input or of the updated representation of the block input, each patch including a different subset of elements of the block input or of the updated representation of the block input; allocating each patch to one or more respective expert modules of a plurality of expert modules of the network block; For each patch of the plurality of patches, processing the patch using one or more corresponding expert modules to generate one or more respective module outputs; and generating a block output by combining module outputs generated by the respective expert modules.

[0145] In a second embodiment, allocating each patch to one or more respective expert modules of the network block includes: processing the patch using the routing sub-network to generate, for each expert module of the plurality of expert modules, a respective routing score; determining one or more highest routing scores of the plurality of routing scores; and allocating the patch to an expert module corresponding to each of the one or more highest routing scores.

[0146] In a third embodiment, for each expert module of the plurality of expert modules, processing the patch using the routing sub-network to generate a respective routing score is performed. processing the patch using one or more feedforward neural network layers to generate a first intermediate routing score for each respective expert module; for each expert module, sampling a noise value and adding the noise value to the first intermediate routing score of the expert module to generate a second intermediate routing score of the expert module; 3. The system of embodiment 2, comprising, for each expert module, applying a non-linear activation function to the second intermediate routing score of the expert module to generate a routing score for the expert module.

[0147] Embodiment 4 further comprises, for each patch of the plurality of patches, processing the patch using one or more corresponding expert modules to generate one or more respective module outputs; For each of one or more expert modules corresponding to the patch, processing the patch using the expert module to generate a respective initial module output; and for each of one or more expert modules corresponding to the patch, scaling a corresponding initial module output according to a routing score corresponding to the expert module and the patch.

[0148] In a fifth embodiment, the neural network comprises: Importance loss, which encourages the routing scores corresponding to each expert module to have the same measure of central tendency, or The system of any one of embodiments 2 to 4, wherein each expert module is trained using one or more of the load losses that encourage each probability of being allocated a patch to have the same measure of central tendency.

[0149] The sixth embodiment calculates an importance loss L for a set of training images X. imp Determining (X) For each expert module i,

[0150]

number

[0151] Calculating f i is a function defined by the routing sub-network for calculating the routing score of an expert module i or the intermediate routing score of an expert module i; For each expert module i, imp i Importance loss L to encourage (X) to have the same measure of central tendency imp and calculating (X).

[0152] In the seventh embodiment, the importance loss L imp (X) is

[0153]

number

[0154] 7. The system of embodiment 6, wherein E is the number of expert modules.

[0155] In embodiment 8, Each patch is assigned to k different expert modules, k ≥ 1; Load loss L on a set of training images X load Determining (X) For each training image x in the set of training images X,

[0156]

number

[0157] Calculating the formula: i (x) is a function defined by the routing sub-network for calculating the routing score of an expert module i or the intermediate routing score of an expert module i; For each training image x and for each expert module i, p i (x) = P(g i (x) ≧ threshold k (x) where P(e) is the probability of event e occurring; and For each expert module i,

[0158]

number

[0159] Calculate the load for each expert module i. i (x) to encourage the load loss L to have the same measure of central tendency load and calculating (X).

[0160] In the ninth embodiment, the load loss L load (X) is

[0161]

number

[0162] 9. The system of embodiment 8, wherein E is the number of expert modules.

[0163] Example 10 is the system of any one of Examples 1 to 9, wherein, for each patch of the plurality of patches, processing the patch using a corresponding expert module includes processing the patch in parallel across the expert modules.

[0164] An eleventh embodiment is the system of any one of the first to tenth embodiments, wherein the number of the plurality of expert modules is greater than the number of the plurality of patches.

[0165] A twelfth embodiment is a method including any one of the operations of the first to eleventh embodiments.

[0166] Embodiment 13 is one or more computer storage media storing instructions that, when executed by one or more computers, cause the one or more computers to perform the operations of any one of embodiments 1 to 11.

[0167] Although this specification contains many specific implementation details, these should not be considered as limitations on the scope of any invention or what may be claimed, but rather as descriptions of features that may be specific to certain embodiments of a particular invention. Certain features described herein in the context of separate embodiments may also be implemented in combination in a single embodiment. Conversely, various features described in the context of a single embodiment may also be implemented in multiple embodiments separately or in any suitable subcombination. Furthermore, although features may be described above as working in a particular combination, and may even be initially claimed as such, one or more features of 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.

[0168] Similarly, although operations are illustrated in the figures and claimed in a particular order, this should not be understood as requiring that such operations be performed in the particular order illustrated, or in sequential order, or that all of the illustrated operations be performed to achieve a desired result. In certain situations, multitasking and parallel processing may be advantageous. Furthermore, the division of various system modules and components in the above-described embodiments should not be understood as requiring such division in all embodiments, and it should be understood that the described program components and systems may generally be integrated together in a single software product or packaged in multiple software products.

[0169] 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. As an example, the processes depicted in the accompanying figures do not necessarily require the particular order depicted, or sequential order, to achieve desirable results. In some cases, multitasking and parallel processing may be advantageous. [Explanation of symbols]

[0170] 100 Neural Networks, Neural Network Systems 102 Input images 104 1st block output 108 Network Output 110a~l network block 112 Block Input 114 Block Output 120 Routing Module 122a~e Collection of patches 130a~e Expert Module 132a-e Corresponding sets of expert module outputs 140 Combination Module 150 Neural Network Blocks, Network Blocks 152 Embedded Patches 154-layer normalization layer 156 Block Output 160 Self-Attention Neural Network Layer 162 Residual Connection Layer 164-layer normalization layer 170 Routing Module 180 Expert Module 190 Combination Module 192 Residual Connection Layer 210 1st image 220 Second Image 230 Third Image 240 4th image 250 5th image 300 processes

Claims

1. 1. A system including a neural network configured to process an input image and generate a network output related to the input image, the neural network including a sequence of one or more network blocks, the one or more network blocks comprising: obtaining a block input representing an intermediate representation of the input image; determining a plurality of patches of the block input or of the updated representation of the block input, each patch including a different subset of elements of the block input or of the updated representation of the block input; allocating each patch to one or more respective expert modules of a plurality of expert modules of the network block; For each patch of the plurality of patches, processing the patch using one or more corresponding expert modules to generate one or more respective module outputs; generating a block output by combining the module outputs generated by the respective expert modules; and processing the plurality of patches using one or more self-attention neural network layers to apply a self-attention mechanism across the plurality of patches to update each patch prior to allocating each patch to one or more respective expert modules.

2. allocating each patch to one or more respective expert modules of said network blocks; processing the patch using a routing sub-network to generate, for each expert module of the plurality of expert modules, a respective routing score; determining one or more highest routing scores of the plurality of routing scores; and allocating the patch to the expert module corresponding to each of the one or more highest routing scores.

3. A system including a neural network configured to process an input image and generate a network output related to the input image, the neural network including a sequence of one or more network blocks, the one or more network blocks comprising: obtaining a block input representing an intermediate representation of the input image; determining a plurality of patches of the block input or of the updated representation of the block input, each patch including a different subset of elements of the block input or of the updated representation of the block input; allocating each patch to one or more respective expert modules of a plurality of expert modules of the network block; For each patch of the plurality of patches, processing the patch using one or more corresponding expert modules to generate one or more respective module outputs; generating a block output by combining the module outputs generated by the respective expert modules, processing the patch using a routing sub-network to generate a respective routing score for each expert module of the plurality of expert modules; processing the patches using one or more feedforward neural network layers to generate a first intermediate routing score for each respective expert module; for each expert module, sampling a noise value and adding the noise value to the first intermediate routing score of the expert module to generate a second intermediate routing score for the expert module; and for each expert module, applying a non-linear activation function to the second intermediate routing score of the expert module to generate the routing score for the expert module.

4. A system including a neural network configured to process an input image and generate a network output related to the input image, the neural network including a sequence of one or more network blocks, the one or more network blocks comprising: obtaining a block input representing an intermediate representation of the input image; determining a plurality of patches of the block input or of the updated representation of the block input, each patch including a different subset of elements of the block input or of the updated representation of the block input; allocating each patch to one or more respective expert modules of a plurality of expert modules of the network block; For each patch of the plurality of patches, processing the patch using one or more corresponding expert modules to generate one or more respective module outputs; generating a block output by combining the module outputs generated by the respective expert modules, for each patch of the plurality of patches, processing the patch using the one or more corresponding expert modules to generate one or more respective module outputs; for each of the one or more expert modules corresponding to the patch, processing the patch using the expert module to generate a respective initial module output; and for each of the one or more expert modules corresponding to the patch, scaling a corresponding initial module output according to a routing score corresponding to the expert module and the patch.

5. A system including a neural network configured to process an input image and generate a network output related to the input image, the neural network including a sequence of one or more network blocks, the one or more network blocks comprising: obtaining a block input representing an intermediate representation of the input image; determining a plurality of patches of the block input or of the updated representation of the block input, each patch including a different subset of elements of the block input or of the updated representation of the block input; allocating each patch to one or more respective expert modules of a plurality of expert modules of the network block; For each patch of the plurality of patches, processing the patch using one or more corresponding expert modules to generate one or more respective module outputs; generating a block output by combining the module outputs generated by the respective expert modules, The neural network comprises: Importance loss, which encourages the routing scores corresponding to each expert module to have the same measure of central tendency, or The system, in which each expert module is trained using one or more of the loading losses that encourage each probability that is assigned a patch to have the same measure of central tendency.

6. Determining the importance loss Limp(X) for a set of training images X, For each expert module i, [0010] where fi is a function defined by a routing sub-network for calculating the routing score of the expert module i or an intermediate routing score of the expert module i; and for each expert module i, computing the importance loss Limp(X) to encourage impi(X) to have the same measure of central tendency.

7. The importance loss Limp(X) is [0025] 7. The system of claim 6, wherein E is a number of the plurality of expert modules.

8. Each patch is assigned to k different expert modules, k>1; Determining the load loss Lload(X) for a set of training images X, For each training image x in the set of training images X, [0030] where gi(x) is a function defined by a routing sub-network for calculating the routing score of the expert module i or an intermediate routing score of the expert module i; For each training image x and for each expert module i, pi(x)=P(gi(x)≧thresholdk(x)) where P(e) is the probability of event e occurring; and For each expert module i, [0045] and for each expert module i, calculating the load loss Lload(X) to encourage loadi(x) to have the same measure of central tendency.

9. The load loss Lload(X) is [0050] 9. The system of claim 8, wherein E is a number of the plurality of expert modules.

10. The system of claim 1 , wherein, for each patch of the plurality of patches, processing the patch using the corresponding expert module comprises processing the patch in parallel across the expert modules.

11. The system of claim 1 , wherein a number of the plurality of expert modules is greater than a number of the plurality of patches.

12. 2. The system of claim 1, wherein the block output of a last network block in the sequence of network blocks is the network output or the block output of the last network block in the sequence of network blocks is further processed using one or more neural network layers to generate the network output.

13. the network output is a classification output including a respective score corresponding to each of a plurality of categories; the network output is an element-level classification output including, for each of one or more elements in the input image, a respective score corresponding to each of a plurality of categories; the network output is a regression output that estimates one or more continuous variables that characterize the input image; or 10. The system of claim 1, wherein the neural network is configured to process a plurality of images that are video frames to generate an output that characterizes the video frames.

14. A method comprising operations performed by a network block of a system according to any one of claims 1 to 13.

15. One or more computer-readable storage media storing 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