Processing an image using expert blending

Abstract

Methods, systems, and devices for generating predictions about images, including computer programs encoded on computer storage media. One of the systems includes a neural network comprising a sequence of one or more network blocks each configured to perform operations comprising: obtaining a block input representing an intermediate representation of an input image; determining a plurality of tile blocks of the block input or an updated representation of the block input, wherein each tile block comprises a different subset of elements of the block input or the updated representation of the block input; assigning each tile block to one or more respective expert modules of a plurality of expert modules of the network block; for each tile block of the plurality of tile blocks, processing the tile block using the corresponding expert module to generate a respective module output; and generating a block output by combining the module outputs.

Description

[0001] Cross-references to related applications

[0002] This application claims priority to U.S. Provisional Application Serial No. 63 / 194,926, filed May 28, 2021, pursuant to Section 119 of Title 35 of the United States Code, the entire contents of which are incorporated herein by reference. Background Technology

[0003] This manual relates to image processing using neural networks.

[0004] A neural network is a machine learning model that uses layers of one or more non-linear units to predict the output of a received input. Some neural networks include one or more hidden layers in addition to the output layer. The output of each hidden layer serves as the 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 the received input based on the current values ​​of the corresponding set of parameters. Summary of the Invention

[0005] This specification describes a system implemented as a computer program on one or more computers at one or more locations, the system being configured to process an image using a neural network and generate a network output representing the image. The neural network comprises a sequence of one or more network blocks, each network block being configured to process a block input including an image or an intermediate representation of an image and generate a block output.

[0006] At least one of the network blocks is an expert network block, which comprises a collection of multiple distinct expert modules, each configured to process a proper subset of the elements of the block input. Specifically, the expert network block is configured to identify multiple "patches" of the elements of the block input, and to process each patch using a routing module to identify a proper subset of the expert modules to which the patch is routed. Then, for each patch routed to an expert module, each expert module processes the patch to generate a corresponding expert module output for the routed patch. The expert network block can then combine the corresponding expert module outputs generated by the corresponding expert module in response to processing the patch for each patch to generate a corresponding patch output. The expert network block can then assemble the patch outputs corresponding to each patch of the block input to generate a block output, wherein the patch outputs have the same spatial relationships as the corresponding patches in the block input.

[0007] The subject matter described in this specification can be implemented in certain embodiments to achieve one or more of the following advantages.

[0008] Some existing systems use “dense” feedforward neural networks (i.e., for each layer of a feedforward neural network, each network parameter of that layer is used to process each element of the layer input to that layer) to process images. For large models configured to process high-resolution images, such dense neural networks can require significant time and computational resources to process a single image. Using the techniques described in this specification, systems can process images using feedforward neural networks by selectively activating a subset of neural network parameters based on the network input, thereby significantly improving the time and computational efficiency of image processing. Furthermore, introducing this sparsity allows feedforward neural networks to include more network parameters than previously feasible, since only a subset of the parameters are used to process any given input.

[0009] Using the techniques described in this specification, a system can implement a neural network comprising expert network blocks having multiple expert modules—each expert module having been trained and configured to process different types of images (or different types of image patches), thereby allowing the modules to be "specialized" and further improving the efficiency and performance of the neural network.

[0010] As described in this specification, a self-attention-based neural network configured to process images can achieve the same performance as a state-of-the-art convolutional neural network with less computation. In other words, for a fixed computational budget, a self-attention-based neural network performs better than a convolutional neural network configured to perform the same machine learning task on an image (measured, for example, by testing accuracy, precision, recall, etc.). This is because applying self-attention is generally more computationally efficient than convolutional kernels across the entire image, as the self-attention mechanism can focus on different regions of the image with less computation than convolution.

[0011] Details of one or more embodiments of the subject matter of this specification are set forth in the following drawings and description. Other features, aspects, and advantages of this subject matter will become apparent from the description, drawings, and claims. Attached Figure Description

[0012] Figure 1A This is a schematic diagram of an example neural network configured to generate predictions about an image.

[0013] Figure 1B This is a schematic diagram of an example neural network block that includes multiple expert modules.

[0014] Figure 2 This shows an example image that has been segmented into image patches.

[0015] Figure 3 This is a flowchart of an example process for executing a network block that includes multiple expert modules.

[0016] Similar reference numerals and names in the various figures indicate similar elements. Detailed Implementation

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

[0018] Figure 1A This is a schematic diagram of an example neural network 100. The neural network 100 is an example of a system implemented as a computer program on one or more computers at one or more locations, wherein the systems, components and techniques described below can be implemented.

[0019] Neural network 100 is configured to process input image 102 based on a set of trained network parameters to generate network output 108 representing an embedding of input image 102 or a prediction about input image 102. Neural network 100 can be configured to perform any appropriate machine learning task using input image 102. Example machine learning tasks are discussed below. In this specification, processing an image refers to processing the intensity values ​​of the pixels in the image.

[0020] The neural network 100 comprises a sequence of L network blocks 110a-l, where L ≥ 1. Each network block 110a-l can be configured to process block inputs including images or intermediate representations of images and generate block outputs.

[0021] like Figure 1A As shown, the 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, which is an intermediate representation of the image 102. Then, each subsequent network block 110b-l can process the block output of the previous network block in the sequence. That is, the block input of the (i+1)th network block 110a-n is the block output of the ith network block 110a-l.

[0022] In some other implementations, the first network block 110a is configured to process the embedding of the input image 102, rather than the input image 102 itself. For example, although the following description generally refers to a neural network that processes the image 102, in some implementations, the neural network 100 may be configured to receive the embedding of the image 102 as input, for example, in response to processing the embedding of the image 102 generated by a different neural network.

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

[0024] In some other implementations, one or more neural network layers are used to further process the block outputs of the final network blocks 110a-l in the sequence to generate the 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 use one or more neural network layers, including a final logistic layer or a softmax layer, to process the block outputs of the final network blocks 110a-l in the sequence to generate the network output 108, which identifies a value representing the probability that the input image 102 belongs to each of the candidate categories.

[0025] As a specific example, neural network 100 may include one or more "head" subnetworks, each configured to use one or more neural network layers to process the block outputs of the final network blocks 110a-l in the sequence to generate predictions about the input image 102 for corresponding different machine learning tasks. That is, after neural network 100 is trained, the outputs of the final network blocks 110a-l in the sequence (which can be viewed as embeddings of the input image 102) can encode information useful for multiple different machine learning tasks.

[0026] In some implementations, neural network 100 first processes the input image 102 to determine a set of patches of the input image, for example, by determining an input sequence comprising multiple elements representing determined patches of the input image 102. Then, neural network 100 can process the determined patches using a sequence of network blocks 110a-l (or embeddings of determined patches generated by one or more initial neural network layers of neural network 100 or by different neural networks). In some other implementations, one or more of network blocks 110a-l each determine a set of patches (which are intermediate representations of the input image 102) of the input block to network block 110a-l, and process the determined patches to generate a block output for network block 110a-l. The operation of network blocks 110a-l is described in more detail below.

[0027] In some implementations, in addition to the input image 102, the neural network 100 processes one or more other network inputs. For example, the neural network 100 can process data that identifies the class to which the input image 102 belongs. As another example, the neural network 100 can process a text input corresponding to the input image 102, for example, a text input representing text that characterizes the input image 102.

[0028] As another example, in implementations where the neural network 100 determines a sequence of patches (or an embedding thereof) of the input image 102 to be sequentially updated through the sequence of network blocks 110a-l, the neural network 100 can process a predetermined token, sometimes referred to as a "class" token. That is, the predetermined token can be the same, for example a predetermined tensor, for each input image 102 received by the neural network 100. For example, the predetermined token can be machine-learned, for example, concurrently with the training of the neural network 100. As a specific example, the predetermined token can be appended to the sequence of patches processed by the neural network 100, such that each network block in the sequence of network blocks 110a-l generates an updated representation for the predetermined token. In some implementations, the network output 108 is (or includes) the final representation of the predetermined token generated by the final network block 110l, wherein the final representation can be regarded as the embedding of the input image 102.

[0029] Some or all of the network blocks 110a-l include a respective set of a plurality of different expert modules. In this specification, an expert module is a component of a network block that is configured to process a proper subset, that is, fewer than all, of the elements of an input block to the network block. An expert module can be configured to process the proper subset of elements of a block input using one or more neural network layers to generate an updated representation of the proper subset of elements. In this specification, a network block including a plurality of expert modules is also referred to as an expert network block.

[0030] For example, each network block 110a-l in the sequence can be an expert network block. As another example, every other one of the network blocks 110a-l in the sequence can be an expert network block. As another example, every other one of the network blocks 110a-l in the sequence up to a specific network block 110i (i<l) can be an expert network block.

[0031] For example, as Figure 1A shown, the j-th network block 110j in the sequence of network blocks 110a-l is an expert network block including e expert modules 130a-e (where e>1). The following description specifically references the j-th network block 110j, but it should be understood that the same technique can be implemented for each expert network block in the neural network 100.

[0032] The j-th network block can partition its block input 112 into multiple patches, where each patch includes a distinct 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, then the j-th network block 110j can determine that each patch has dimensions L×W×C, where C represents the number of channels in the representation of the input image 102 (e.g., C = 3 if the block input 112 is an RGB image; or C is a larger number if the block input 112 is an intermediate representation of an RGB image).

[0033] In some implementations, each piece of block input 112 has the same size, meaning it includes the same number of elements and has the same dimensions. In other implementations, different pieces can have different sizes, for example, they can include different numbers of elements.

[0034] In some implementations, each element of block input 112 is in exactly one slice. In other implementations, some or all of the elements of block input 112 may be in multiple different slices; that is, the slices of block input 112 may overlap.

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

[0036] The patches of block input 112 can be represented in any suitable way. For example, each patch can be represented as a two-dimensional tensor that includes the elements of block input 112, such as a tensor that maintains the spatial relationships between the elements in block input 112. As another example, each patch can be represented as a one-dimensional sequence of the elements of the patch. As a specific example, if the patch is a two-dimensional region of block input 112, the patch can be a flattened version of the two-dimensional region, such as where the elements of block input 112 are ordered in a one-dimensional tensor according to their respective positions in block input 112. As another example, each initial image patch 112a-n can be represented as an unordered set of pixels of the initial image patch 112a-n.

[0037] The following is for reference. Figure 2 This paper discusses example techniques for determining patches or intermediate representations of an image.

[0038] After determining the segments of block input 112, network block 110j can use routing module 120 to assign each segment to one or more of the expert modules 130a-e of network block 110j. Then, network block 110j can use the corresponding assigned expert modules 130a-e to process each image segment to generate the corresponding expert module output, and use combination module 140 to combine the expert module outputs of expert modules 130a-e to generate block output 114 of network block 110j.

[0039] In other words, routing module 120 is configured to determine, for each expert module 130a-e, a corresponding set 122a-e of the blocks to be routed to the block input 112 processed by expert modules 130a-e. Then, each expert module 130a-e processes the assigned set 122a-e of blocks to generate a corresponding set 132a-e of expert module outputs, i.e., the corresponding expert module output for each assigned block. The expert module output can also be simply referred to as "module output".

[0040] To assign a fragment to the appropriate expert module 130a-e, routing module 120 can use a routing neural network comprising one or more neural network layers to process each fragment. In response to processing a specific fragment, the routing neural network can generate an output including the corresponding routing score for each expert module 130a-e of network block j 110j. Routing module 120 can determine which expert module 130a-e to assign the fragment to, and which corresponds to, the highest routing score for that fragment.

[0041] In some implementations, each segment is assigned to the same number of expert modules 130a-e. That is, each segment can be assigned to the k expert modules 130a-e that have the highest routing score for that segment, where k is the same for all segments.

[0042] In some other implementations, fragments can be assigned to different numbers of expert modules. For example, if the routing scores for expert modules 130a-e and the fragments meet a predetermined threshold, the fragments can be assigned to expert modules 130a-e. As another example, by assigning p (p≥1) fragments with the highest routing score corresponding to an expert module among all routing scores corresponding to both the expert module and the corresponding fragment, it is guaranteed that each expert module 130a-e processes exactly p fragments.

[0043] In some implementations, not every expert module 130a-e of the j-th network block 110j processes pieces of the block input 112; that is, for some input images 102, some expert modules 130a-e of the j-th network block 110j are idle for the block input 112 generated from the input images 102. Allowing some expert modules 130a-e to be idle during some execution of the neural network 100 can improve the efficiency of the neural network 100 by reducing the amount of computation required to generate the network output 108, since not every parameter of the neural network 100 is used to process every input image 102.

[0044] The routing neural network of routing module 120 can be trained to assign each piece to an expert module 130a-e that will extract the most information from the piece. That is, in some implementations, at least some expert modules 130a-e can be "specialized" for certain types of pieces—e.g., pieces describing a specific semantic object or a category of semantic objects. Example techniques for training neural network 100 are discussed in more detail below.

[0045] In some implementations, the j-th network block 110j limits the number of pieces that each expert module 130a-e can process, so that no expert module 130a-e can be assigned a disproportionate amount of work (i.e., computation) relative to other expert modules, and thus becomes a bottleneck for the execution of the neural network 100. For example, network block 110j can limit the number of pieces that can be assigned to any expert module 130a-e to less than or equal to:

[0046]

[0047] Where B is the number of images in the batch of input images 102 processed by neural network 100 (e.g., the batch of images 102 provided to neural network 100 and processed in parallel by neural network 100 during inference time), K refers to the number of expert modules 130a-e assigned to each piece, P is the number of pieces in block input 112, and E is the number of expert modules 130a-e in network block 110j. C can be a hyperparameter of neural network 100, representing the capacity of each expert module 130a-e, where C=1 indicates that each expert module 130a-e is only allowed to process a portion of the pieces during a given execution of neural network 100.

[0048] In some implementations, the neural network 100 is set to C << 1, for example, C = 0.5, 0.1, 0.05, or 0.01.

[0049] In some implementations, if more image patches are assigned to expert modules 130a-e than allowed, then for each patch with the second-highest routing score, patches exceeding the limit (e.g., patches assigned to expert modules 130a-e that have the lowest routing score corresponding to expert modules 130a-e) will be reassigned to expert modules 130a-e. In some other implementations, patches exceeding the limit are not reassigned but are "dropped," and thus processed by one fewer expert module 130a-e than the other patches; this is equivalent to using an expert module that performs only an identity function to process the dropped patch.

[0050] In some implementations, the neural network 100 operates using different values ​​of C during its different execution periods; that is, the value of C can be changed after the neural network 100 is trained. For example, during the training of the neural network 100, the value of C can be relatively high, thus discarding relatively high pieces, and the training system can use the module outputs of expert modules 130a-e to determine the updates to the network parameters of the neural network layers of expert modules 130a-e. Then, during inference time, the value of C can be reduced to improve the time and computational efficiency of the neural network 100. As a specific example, the neural network 100 can be deployed in resource-scarce inference environments, such as on devices with relatively limited computing resources, such as mobile phones or tablets, or in environments where the neural network 100 continuously receives new images 102 and must generate corresponding network outputs 108 using a fixed time budget.

[0051] In some implementations, for each piece of block input 112, routing module 120 randomly samples noise values ​​(e.g., according to a normal distribution) for each expert module 130a-e and adds the sampled noise values ​​to the corresponding routing scores before determining the k highest routing scores for that piece. For example, routing module 120 may add noise values ​​during the training of neural network 100 to encourage variation and thus encourage each expert module 130a-e to receive parameter updates. At inference time, routing module 120 may be configured not to add noise values ​​because routing module 120 has already been trained to assign high-quality routing scores.

[0052] In some implementations, the network block can apply a non-linear activation function, such as softmax, Tanh, or ReLU, to the routing scores before determining the k highest routing scores.

[0053] Routing neural networks can include any appropriate neural network layer configuration.

[0054] For example, the routing neural network may include one or more feedforward neural network layers. As a specific 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 use a feedforward neural network layer to process the reshaped patch.

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

[0056] Alternatively or additionally, a routing neural network may include one or more self-attention layers. See below for reference. Figure 1B Let's discuss self-attention in more detail.

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

[0058] In some other embodiments, each expert module 130a-e is executed on the same processing device, for example, on parallel processing hardware configured to execute at least a portion of expert modules 130a-e on corresponding different threads—such as graphics processing units (GPUs) or tensor processing units (TPUs). In some other implementations, the respective expert modules 130a-e can be executed by different processing devices, for example, by a distributed system composed of multiple different parallel processing devices. Parallel execution of at least some of the expert modules 130a-e can significantly reduce the latency of the execution of the neural network 100, for example, by approximately a factor of e if all expert modules 130a-e are executed on different devices. Therefore, for a given number of network parameters, a neural network including at least one expert network block of expert modules can enjoy shorter latency compared to a neural network with only dense neural network layers, for two reasons: (i) each element of the input image is processed by a strict subset of the network parameters of the expert modules, and (ii) the execution of the expert modules can be parallelized, whereas in some implementations, dense neural network layers cannot be parallelized. Similarly, for the same reason, given a specific time budget (i.e., for a specific delay threshold), a neural network that includes at least one expert network block with an expert module can include significantly more network parameters, and thus increases network capacity (i.e., the ability to encode more information) compared to a neural network with only dense neural network layers.

[0059] In some implementations where expert modules 130a-e are distributed across multiple devices, the processing unit of execution routing module 120 is configured to provide only the set of fragments 122a-e already assigned to the specific expert modules 130a-e to the processing unit executing the specific expert modules 130a-e. In some other such implementations, the processing unit of execution routing module 120 is configured to provide the entire block input 112 along with data identifying which fragments of the block input 112 will be processed by the specific expert modules 130a-e.

[0060] After processing each piece through one or more corresponding expert modules to generate the corresponding module output, the combination module 140 of network block 110j can combine the module outputs corresponding to each piece to generate block output 114 of network block 110j. Specifically, the combination module 140 can combine the corresponding module outputs generated by the corresponding expert modules 130a-e in response to processing the piece to generate the piece output for each piece of block input 112. Then, the combination module 140 can reassemble the block outputs to generate block output 114 such that the block outputs in block output 114 have the same spatial relationships as the corresponding pieces in block input 112.

[0061] For example, in some implementations, the piece outputs have the same dimensions as the corresponding pieces. Therefore, the composition module can assemble the piece outputs to generate a block output 114 with the same dimensions L×W×C as the block input 112, wherein each piece output occupies the same position in the block output 114 as the corresponding piece occupies in the block input 112.

[0062] In some implementations, the combination module 140 combines the corresponding module outputs for each piece of block input by calculating the sum of the module outputs of the piece. In some such implementations, to generate the piece output for each piece, the combination module 140 can weight each module output using a routing score corresponding to the piece and the expert modules 130a-e that generate the module outputs. 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).

[0063] In some implementations, before assigning pieces of block input 102 to the corresponding expert modules 130a-e, network block 110j first processes the block input using one or more neural network layers to generate an updated representation of block input 102, and then assigns pieces of the updated representation of block input 102 to the corresponding expert modules 130a-e. For example, network block 110j may first process the block input using one or more self-attention neural network layers and / or one or more element-wise feedforward neural network layers configured to process each piece of block input 102. In other words, the input to the expert modules may be a strict subset of the elements of the updated representation of block input 102, rather than a strict subset of the elements of block input 102 itself. Generally, in this specification, we refer to pieces of block input 102, but it should be understood that the same process can be applied to pieces of the updated representation of block input 102. Similarly, one or more neural network layers preceding routing module 120 can be considered as part of different network blocks in the sequence of network blocks 110a-l.

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

[0065] As described above, in some implementations, the block output of the final network block 110l in the sequence is network output 108. For example, if the block output of the final network block 110l includes a corresponding representation of each of a plurality of patches of the input image 102, then network output 108 may include a representation of each patch. As another example, network output 108 may include a combined representation generated based on the representation of each patch, for example, by determining the average of the patch representations or by applying a pooling operation (e.g., average pooling, max pooling, or global pooling).

[0066] As another example, as described above, in some implementations, the neural network 100 processes a predetermined type token along with the input image 102, such that each network block in the sequence of network blocks generates an updated representation of the predetermined type token in addition to the updated representation of the fragments. In these implementations, the network output 108 may be the final representation of the predetermined type token generated by the final network block 110l.

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

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

[0069]

[0070] Among them, f i This is a function defined by the routing neural network of the routing module 120 of the expert network block, used to calculate the routing score of expert module i or the intermediate routing score of expert module i. For example, f i (x) can be the routing scores of expert module i and piece x before adding random noise and applying a nonlinear activation function. The training system can then compute the importance loss. To encourage imp for each expert module i (X) have the same mean. As a specific example, the training system can calculate:

[0071]

[0072] Where E is the number of expert modules.

[0073] Alternatively or additionally, the training system can use "load loss" to train the neural network 100. For each expert network block and for each piece, "load loss" encourages that each expert module 130a-e of the expert network block has the same average probability of being assigned to a piece. For example, for each expert module i and for each image x, the training system can calculate:

[0074]

[0075] Among them, 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 expert module i or the intermediate routing score of expert module i. For example, g i (x) can be the routing scores of expert module i and piece x after adding random noise but before applying a non-linear activation function. Then, for each training image x and for each expert module i, the training system can compute:

[0076] p i (x)=P(g i (x)≥threshold k (x))

[0077] Where P(e) represents the probability of event e occurring. Given the distribution of random noise, p i (x) is directly computable. Then, for each expert module i, the training system can compute:

[0078]

[0079] And calculate load loss To encourage load for each expert module i i (x) have the same mean. As a specific example, the training system can calculate:

[0080]

[0081] Where E is the number of expert modules.

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

[0083] Image 102 can be any suitable type of image. For example, image 102 can be a two-dimensional image, such as a two-dimensional image with multiple channels (e.g., an RGB image). As another example, image 102 can be a hyperspectral image representing a continuous spectrum of wavelengths, for example, by identifying the distribution on the spectrum for each pixel in image 102. As another example, image 102 can be a point cloud comprising multiple points, where each point has corresponding coordinates in a three-dimensional or higher-dimensional coordinate space, for example; as a particular example, image 102 can be a point cloud generated by a LIDAR sensor. As another example, image 102 can be a medical image generated by a medical imaging device; as a particular example, image 102 can be a computed tomography (CT) image, magnetic resonance imaging (MRI) image, ultrasound image, X-ray image, mammogram image, fluorescence fluoroscopy image, or positron emission tomography (PET) image.

[0084] While the following description refers to image patches that generate image 102, each patch comprising a corresponding “pixel” of image 102, it should be understood that the neural network system 101 can generate image patches of image 102 having any suitable type of components. For example, if image 102 is a point cloud, each image patch of image 102 may include a subset of points in the point cloud. As another example, if image 102 is an MRI image comprising multiple voxels in a three-dimensional voxel grid, each image patch of image 102 may include a subset of voxels in the voxel grid.

[0085] Neural networks can be configured to process images to perform any appropriate image processing task, such as classification, regression, or a combination thereof.

[0086] As a specific example, neural network 100 can be configured to generate classification output including a score corresponding to each of a plurality of categories. The category score indicates the likelihood that an image belongs to that category. In some cases, a category can be a kind of object (e.g., dog, cat, person, etc.), and if an image depicts an object included in the object kind corresponding to the category, the image may belong to that category. In other cases, a category can represent a global image attribute (e.g., whether an image depicts a daytime or nighttime scene, or whether an image depicts a summer or winter scene), and if an image has a global attribute corresponding to the category, the image may belong to that category.

[0087] As another specific example, the neural network 100 can be configured to generate element-level classification output (e.g., pixel-level classification output for an RGB image or point-level classification output for a LiDAR image), which, for each element in the image, includes a corresponding score for each of a plurality of categories. For a given element (e.g., for a given pixel or point), the category score indicates the probability that the element belongs to that category. In some cases, the category can be a kind of object, and if an element is part of an object included in the object kind corresponding to the category, it is likely to belong to that category. That is, the element-level classification output can be semantic segmentation output.

[0088] As another specific example, the neural network 100 can be configured to generate a regression output that can estimate one or more continuous variables characterizing an image (i.e., an infinite number of possible values ​​can be assumed). In this specific example, the regression output can estimate the coordinates of the bounding box surrounding the corresponding object depicted in the image. The coordinates of the bounding box can be defined by the (x, y) coordinates of the vertices of the bounding box.

[0089] In some implementations, the neural network 100 can be configured to process multiple images, such as multiple frames of a video. For example, the neural network can receive multiple images as video frames and can process each video frame as described above to generate an output characterizing the video frame, for example, by characterizing whether the video frame depicts a person performing a specific action.

[0090] In some such implementations, the neural network 100 processes each video frame at corresponding different time points to generate a corresponding network output for each video frame, which represents the prediction of the video frame. For example, the neural network can generate a network output predicting the classification of a video frame. In some such implementations, the neural network combines multiple network outputs corresponding to a given video frame to generate a final network output representing the video. For example, the neural network can use a downstream neural network (e.g., a recurrent neural network) to process the corresponding network outputs.

[0091] In some other implementations, the neural network 100 processes each video frame in parallel to generate a single network output representing the video. As a specific example, the system can generate one or more corresponding input elements in the input sequence for each video frame.

[0092] Figure 1B This is a schematic diagram of an example neural network block 150 including multiple expert modules 180. The neural network 150 is an example of a system implemented as a computer program on one or more computers at one or more locations, wherein the systems, components and techniques described below can be implemented.

[0093] Network block 150 may be one of a sequence of L network blocks (L≥1) of a neural network configured to process the input image to generate a network output representing the input image, for example, representing an embedding of the input image or a network output representing a prediction about the input image. For example, network block 150 may be one of the network blocks mentioned above. Figure 1A One of the network blocks in the sequence of network blocks 110a-l.

[0094] Network block 150 is configured to process a set of embedded fragments 152 of the input image, i.e., updated representations of fragments determined based on the input image that have been generated by one or more previous 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 can be configured to process the original fragments determined based on the input image.

[0095] Network block 150 includes one or more self-attention neural network layers. For example... Figure 1B As shown, 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 obtain the embedding piece 152 and apply an attention mechanism to the embedding piece 152. In some other implementations, the self-attention neural network layer 160 is configured to obtain a corresponding processed version of the embedding piece 152 and apply an attention mechanism to the processed embedding piece 152. For example, as... Figure 1B As shown, network block 150 may first apply layer normalization layer 154 to embedding patch 152 before providing the output of layer normalization layer 154 to self-attention neural network layer 160. Alternatively or additionally, network block 150 may apply one or more other neural network layers to embedding patch 152 before self-attention neural network layer 160 (e.g., one or more element-wise feedforward neural network layers).

[0096] For example, network block 150 can generate a one-dimensional tensor including the elements of each embedding piece 152 for each embedding piece 152. If each embedding piece 152 has dimensions L1×W1×C, then network block 150 can generate a one-dimensional tensor with dimensions 1×(L·W·X). Network block 150 can then generate an input sequence based on the one-dimensional tensor, wherein at least a portion of the input positions of the input sequence is a one-dimensional tensor, and process the input sequence using a self-attention neural network layer 160.

[0097] Specifically, for each specific input location, a self-attention neural network layer 160 is configured to apply an attention mechanism to the embedding piece 152 (or its processed version) at the input location (i.e., other embedding pieces 152, and in some implementations, itself) using one or more queries derived from the embedding piece 152 at that specific input location, to generate a corresponding output for that location. The output of the self-attention neural network layer 160 is an output sequence that includes the corresponding output element for each input location.

[0098] In some implementations, some or all of the self-attention neural network layers in network block 150 (e.g. Figure 1B The self-attention neural network layer 160 shown is a multi-head self-attention neural network layer. The multi-head self-attention neural network layer applies h different attention mechanisms in parallel to generate the corresponding sequences of output elements, and then combines the multiple sequences of output elements to generate the final sequence of output elements.

[0099] In some implementations, some or all of the self-attention neural network layers in network block 150 (e.g. Figure 1B The self-attention neural network layer 160 shown incorporates positional information of the corresponding embedded patch 152 in the input sequence into the attention mechanism. For example, when attention is applied with respect to a specific embedded patch 152 (i.e., when generating a corresponding output element corresponding to a specific embedded patch 152), the self-attention neural network layer 160 can recognize attentional position embeddings representing the position of the embedded patch 152 corresponding to the image (or within the block input).

[0100] In this specification, an embedding is an ordered set of numerical values ​​representing inputs in a specific embedding space. For example, an embedding can be a floating-point vector or other numerical values ​​with a fixed dimension.

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

[0102] The output of the self-attention neural network layer 160 (or residual connection layer 162 or layer normalization layer 164) includes an updated representation of the embedded piece 152. This updated representation of the embedded piece 152 can be provided to a routing module 170, which is configured to assign each updated representation of the corresponding embedded piece 152 to one or more expert modules 180 from a set of multiple expert modules of network block 150. This can be similar to the above reference. Figure 1B The routing module 120 is used to configure the routing module 170.

[0103] The network block can process the assigned update representation of the corresponding embedded fragment 152 for each expert module 180 to generate the corresponding expert module output for each assigned update representation. This can be similar to the above reference. Figure 1A The expert modules 130a-e are used to configure the expert module 180.

[0104] like Figure 1B As shown, each expert module 180 is a multilayer perceptron (MLP) comprising one or more feedforward neural network layers. For example, expert module 180 may include two or more learned linear transformations, each transformation being separated by an activation function (e.g., a non-linear element activation function, such as the ReLU activation function).

[0105] Then, the expert module output can be combined using the combination module 190 of network block 150 to generate block output 156 for the network block. This can be similar to the above reference. Figure 1A The combination module 140 is used to configure the combination module 190.

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

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

[0108] 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 a corresponding machine learning task, as referenced above. Figure 1A As another example, a neural network can provide its output (which may represent the embedding of the input image) to a database or one or more downstream machine learning models for further processing.

[0109] Figure 2 Example images 210, 220, 230, 240, and 250 are shown, segmented into image patches.

[0110] Images of 210 to 250 can be provided as input to the neural network system, as shown in the reference above. Figure 1AThe neural network 101 is configured to process images 210 to 250 to generate predictions about images 210 to 250. The neural network system can be configured to segment images 210 to 250 into multiple image patches and process the image patches using a sequence of network blocks including one or more expert network blocks.

[0111] While the following description refers to determining patches of an image, it should be understood that the same techniques can be used to determine patches of intermediate representations of an image, for example, by feeding block inputs to network blocks (e.g., the above references). Figure 1A The block input 112) is a piece.

[0112] Images 210, 220, 230, 240, and 250 illustrate the different possibilities for segmenting an image into image patches. Specifically, in Figure 2 In this context, each image 210 to 250 is described as a set of image patches segmented into which each image patch is visually distinct, i.e., using different shading or shadow lines. Generally, the neural network system will be configured to segment all received images according to the same pattern. That is, the same neural network system does not necessarily segment each image 210, 220, 230, 240, and 250 as shown, because these images have already been segmented according to different patterns.

[0113] As shown in the first image 210, in some implementations, the neural network system can generate image patches that are each of the same size and shape; for example, each image patch can be rectangular. Furthermore, in some implementations, the neural network system can segment the first image 210 such that each pixel is exactly a member of an image patch. As a specific example, such as Figure 2 As shown, image patches can represent a grid of rectangles of the same size. As another specific example, image patches can represent a grid of hexagons of the same size.

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

[0115] 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 members of multiple distinct image patches. For example, a first image patch may include a strict superset of pixels from a second image patch; that is, the second image patch may be completely enclosed within a first image patch. As a particular example, for each of the multiple first image patches determined based on the image (e.g., for each first image patch in a grid of first image patches), the neural network system can generate a corresponding second image patch with the same center pixel but different sizes and / or aspect ratios.

[0116] 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 members of any image patch. For example, the neural network system can use a machine learning model to process the fourth image 240 to identify one or more regions of interest, and the neural network system can generate corresponding patches 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 corresponding patches centered on each identified pixel.

[0117] As shown in the fifth image 250, in some implementations, the neural network system can generate image patches of arbitrary shapes. That is, the image patches are not required to be rectangular. For example, the neural network system can use a machine learning model to process the fifth image 250, which is configured to segment the fifth image 250, for example, by assigning a corresponding class to each pixel in the fifth image 250. The neural network system can then generate corresponding patches for each consecutive set of pixels that have been assigned the same class by the machine learning model.

[0118] Although Figure 2 Images 210 to 250 are shown as two-dimensional images (or two-dimensional images with multiple channels, such as RGB images), but neural network systems can generally be configured to generate predictions for any type of image, as referenced above. Figure 1A As stated above.

[0119] Figure 3 This is a flowchart of an example process for executing a network block comprising multiple expert modules. For convenience, process 300 is described as being executed via a system of one or more computers at one or more locations. For example, a neural network system appropriately programmed according to this specification (e.g., the above reference). Figure 1A The neural network 100 can execute process 300.

[0120] A network block can be one of a sequence of network blocks in a neural network that is configured to process an input image and generate network outputs for the input image.

[0121] The system obtains a block input representing an intermediate representation of the input image (step 302).

[0122] The system determines the block input (or the updated representation of the block input, for example, by applying the above reference). Figure 1B The set of multiple pieces of the updated representation generated by the self-attention mechanism, wherein each piece includes a different subset of the elements of the block input (or the updated representation of the block input) (step 304).

[0123] The system assigns each piece to one or more corresponding expert modules among the multiple expert modules of the network block (step 306).

[0124] For each of the multiple slices, the system uses one or more corresponding expert modules to process the slice to generate one or more corresponding module outputs (step 308).

[0125] The system generates block outputs for network blocks by combining the module outputs generated by the corresponding expert modules (step 310).

[0126] This specification uses the term "configured to" in conjunction with system and computer program components. For a system of one or more computers, being configured to perform a particular operation or action means that software, firmware, hardware, or a combination thereof has been installed on the system that, in operation, causes the system to perform the operation or action. For one or more computer programs, being configured to perform a particular operation or action means that one or more programs include instructions that, when executed by a data processing device, cause the device to perform an operation or action.

[0127] Embodiments of the subject matter and functional operation described in this specification may be implemented in digital electronic circuits, in computer software or firmware embodied in a tangible manner, in computer hardware (including the structures disclosed herein and their structural equivalents), or in one or more combinations thereof. Embodiments of the subject matter described in this specification 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 executing or 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 one or more combinations thereof. Alternatively or additionally, program instructions may be encoded on artificially generated propagation signals, such as machine-generated electrical, optical, or electromagnetic signals, which are generated as coded information for transmission to a suitable receiver device for execution by the data processing device.

[0128] The term "data processing device" refers to data processing hardware and encompasses all types of devices, apparatuses, and machines used for processing data, including programmable processors, computers, or multiple processors or computers. Furthermore, the device may be or may further include special-purpose logic circuitry, such as FPGAs (Field-Programmable Gate Arrays) or ASICs (Application-Specific Integrated Circuits). In addition to hardware, the device may optionally include code that creates an execution environment for computer programs, such as code constituting processor firmware, protocol stacks, database management systems, operating systems, or combinations thereof.

[0129] A computer program, also referred to or described as a program, software, software application, application, module, software module, script, or code, can be written in any form of programming language, including compiled or interpreted languages, or declarative or procedural languages; it can be deployed in any form, including as a standalone program or as a module, component, subroutine, or other unit suitable for use in a computing environment. A program may (but does not need to) correspond to a file in a file system. A program may be stored as part of a file that holds other programs or data (e.g., one or more scripts stored in a markup language document), in a single file dedicated to the program in question, or in multiple co-located files (e.g., a file storing one or more modules, subroutines, or portions of code). A computer program can be deployed to execute on a single computer or on multiple computers located at a site or distributed across multiple sites and interconnected via a data communication network.

[0130] 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 at all, and it can be stored on storage devices in one or more locations. Thus, for example, an indexed database may include multiple collections of data, each of which can be organized and accessed in different ways.

[0131] Similarly, in this specification, the term "engine" is used broadly to refer to a software-based system, subsystem, or process programmed to perform one or more specific functions. Typically, an engine will be implemented as one or more software modules or components, installed on one or more computers at one or more locations. In some cases, one or more computers will be dedicated to a particular engine; in others, multiple engines may be installed and run on the same computer or on multiple computers.

[0132] The processes and logic flows described in this specification can be executed by one or more programmable computers that execute one or more computer programs to perform functions by manipulating input data and generating outputs. These processes and logic flows can also be executed by special-purpose logic circuitry (such as FPGAs or ASICs), or by a combination of special-purpose logic circuitry and one or more programmable computers.

[0133] A computer suitable for executing computer programs can be based on a general-purpose or special-purpose microprocessor, or both, or on any other type of central processing unit (CPU). Typically, the CPU receives instructions and data from read-only memory or random access memory, or both. The basic components of a computer are the CPU for executing instructions and one or more memory devices for storing instructions and data. The CPU and memory can be complemented or integrated therein by dedicated logic circuitry. Typically, a computer will also include, or be operatively coupled to, one or more mass storage devices (e.g., disks, magneto-optical disks, or optical disks) for storing data, transfer data to them, or both. However, a computer does not require such a device. Furthermore, a computer can be embedded in another device, such as a mobile phone, personal digital assistant (PDA), mobile audio or video player, game console, GPS receiver, or portable storage device (e.g., a Universal Serial Bus (USB) flash drive), to name just a few.

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

[0135] To provide interaction with the user, embodiments of the subject matter described in this specification can be implemented on a computer having a display device for displaying information to the user (e.g., a CRT (cathode ray tube) or LCD (liquid crystal display) monitor) and a keyboard and clicking device (e.g., a mouse or trackball) from which the user can provide input to the computer. Other types of devices can also be used to provide interaction with the user; for example, feedback provided to the user can be any form of sensory feedback, such as visual, auditory, or tactile feedback; input from the user can be received in any form, including sound, speech, or tactile input. Furthermore, the computer can interact with the user by sending documents to and receiving documents from the device used by the user; for example, by sending a webpage to a web browser on the user's device in response to a request received from a web browser. Additionally, the computer can interact with the user by sending text messages or other forms of messages to a personal device (e.g., a smartphone running a messaging application) and receiving response messages in return from the user.

[0136] Data processing devices used to implement machine learning models may also include, for example, dedicated hardware accelerator units for processing the general and computationally intensive portions of the machine learning training or output, i.e., inference, workload.

[0137] Machine learning frameworks (such as TensorFlow, Microsoft Cognitive Toolkit, Apache Singa, or Apache MXNet) can be used to implement and deploy machine learning models.

[0138] Embodiments of the subject matter described in this specification can be implemented in computing systems that include backend components (e.g., as data servers), or middleware components (e.g., application servers), or frontend components (e.g., client computers with graphical user interfaces, web browsers, or applications from which users can interact with implementations of the subject matter described in this specification), or any combination of one or more such backend, middleware, or frontend components. The components of the system can be interconnected via any form or medium of digital data communication (e.g., communication networks). Examples of communication networks include local area networks (LANs) and wide area networks (WANs), such as the Internet.

[0139] A computing system may include clients and servers. Clients and servers are typically geographically separated and generally interact via a communication network. The client-server relationship arises from computer programs running on corresponding computers and having a client-server relationship with each other. In some embodiments, the server transmits data (e.g., HTML pages) to a user device, for example, to display data to a user interacting with the device (acting as a client) and to receive user input from the user. Data generated on the user device (e.g., the result of user interaction) can be received from the device at the server.

[0140] In addition to the above embodiments, the following embodiments are also innovative:

[0141] Example 1 is a system including a neural network configured to process an input image and generate a network output for the input image. The neural network includes a sequence of one or more network blocks, each of which is configured to perform the following operations:

[0142] Obtain a block input representing an intermediate representation of the input image;

[0143] Determine multiple segments of the block input or the updated representation of the block input, wherein each segment includes a distinct subset of the elements of the block input or the updated representation of the block input;

[0144] Each piece is assigned to one or more corresponding expert modules from among the multiple expert modules of the network block;

[0145] For each of the plurality of fragments, the fragment is processed using one or more corresponding expert modules to generate one or more corresponding module outputs; and

[0146] Block outputs are generated by combining the module outputs generated by the corresponding expert modules.

[0147] Example 2 is the system according to Example 1, wherein assigning each fragment to one or more corresponding expert modules of the network block includes:

[0148] The segments are processed using a routing subnet to generate a corresponding routing score for each of the plurality of expert modules;

[0149] Determine one or more of the highest routing scores among the plurality of routing scores; and

[0150] The fragment is assigned to the expert module corresponding to each of the one or more highest routing scores.

[0151] Example 3 is a system according to Example 2, wherein processing the segments using a routing subnetwork to generate a corresponding routing score for each of the plurality of expert modules includes:

[0152] The pieces are processed using one or more feedforward neural network layers to generate a corresponding first intermediate routing score for each expert module;

[0153] For each expert module, a noise value is sampled and added to the first intermediate route score of the expert module to generate a second intermediate route score for the expert module;

[0154] For each expert module, a non-linear activation function is applied to the second intermediate routing score of the expert module to generate the routing score for the expert module.

[0155] Example 4 is a system according to any one of Examples 2 or 3, wherein, for each of the plurality of chips, processing the chip using one or more corresponding expert modules to generate one or more corresponding module outputs includes:

[0156] For each of one or more expert modules corresponding to the said fragment, the expert module is used to process the fragment to generate a corresponding initial module output; and

[0157] For each of the one or more expert modules corresponding to the segment, the corresponding initial module output is scaled according to the routing score corresponding to the expert module and the segment.

[0158] Example 5 is a system according to any one of Examples 2 to 4, wherein the neural network has been trained using one or more of the following:

[0159] Importance loss, which encourages a measure of central tendency that has the same value as the routing score corresponding to each expert module, or

[0160] Load loss, which encourages each expert module to have the same central tendency measure of the corresponding probability of being assigned a piece.

[0161] Example 6 is the system according to Example 5, wherein the importance loss is determined for batches of training images X. include:

[0162] For each expert module i, calculate:

[0163]

[0164] Among them, f i It is a function defined by the routing subnetwork for calculating the routing score for expert module i or for calculating the intermediate routing score for expert module i; and

[0165] Calculate the importance loss To encourage imp for each expert module i (X) have the same measure of central tendency.

[0166] Example 7 is the system according to Example 6, wherein the importance loss for:

[0167]

[0168] Where E is the number of the plurality of expert modules.

[0169] Example 8 is a system according to any one of Examples 5 to 7, wherein:

[0170] Each piece is assigned to k different expert modules, where k ≥ 1, and

[0171] The load loss is determined for a batch of training images X. include:

[0172] For each training image x in a batch of training images X, calculate:

[0173]

[0174] Among them, g i(x) is a function defined by the routing subnetwork for calculating the routing score for expert module i or for calculating the intermediate routing score for expert module i;

[0175] For each training image x and for each expert module i, compute:

[0176] p i (x)=P(g i (x)≥threshold k (x))

[0177] Where P(e) represents the probability of event e occurring; and

[0178] For each expert module i, calculate:

[0179]

[0180] And calculate the load loss. To encourage load for each expert module i i (x) have the same measure of central tendency.

[0181] Example 9 is the system according to Example 8, wherein the load loss for:

[0182]

[0183] Where E is the number of the plurality of expert modules.

[0184] Example 10 is a system according to any one of Examples 1 to 9, wherein processing the chip using the corresponding expert module for each of the plurality of chips includes parallel processing of the chip across the expert modules.

[0185] Example 11 is a system according to any one of Examples 1 to 10, wherein the number of the plurality of expert modules is greater than the number of the plurality of chips.

[0186] Example 12 is a method that includes the operation according to any one of Examples 1 to 11.

[0187] Example 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 operation according to any one of Examples 1 to 11.

[0188] While this specification contains numerous specific implementation details, these details should not be construed as limiting the scope of any invention or the scope that may be claimed, but rather as descriptions of features specific to particular embodiments of the invention. Certain features described herein in the context of individual 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 individually in multiple embodiments or in any suitable sub-combination. Furthermore, while features may be described above as functioning in certain combinations, and even claimed in this way initially, in some cases one or more features may be removed from the claimed combination, and the claimed combination may refer to a sub-combination or a variation of a sub-combination.

[0189] Similarly, although operations are described in the accompanying drawings and referenced in the claims in a specific order, this should not be construed as requiring these operations to be performed in the specific order shown or sequentially, or to perform all the shown operations to achieve the desired result. In some cases, multitasking and parallel processing may be advantageous. Furthermore, the separation of various system modules and components in the above embodiments should not be construed as requiring such separation in all embodiments, and it should be understood that the program components and systems can generally be integrated into a single software product or encapsulated within multiple software products.

[0190] Specific embodiments of this subject matter have been described. Other embodiments are within the scope of the following claims. For example, the actions referenced in the claims may be performed in a different order and still achieve the desired result. As an example, the processes described in the figures do not necessarily require the specific order or sequential sequence shown to achieve the desired result. In some cases, multitasking and parallel processing may be advantageous.

Claims

1. A system including a neural network configured to process an input image and generate a network output for the input image, the neural network comprising a sequence of one or more network blocks, each of the one or more network blocks being configured to perform an operation, the operation comprising: Obtain block input, wherein the block input represents an intermediate representation of the input image; Determine multiple fragments of the block input or multiple fragments of the updated representation of the block input, wherein each fragment includes a distinct subset of the elements of the block input or a distinct subset of the elements of the updated representation of the block input; Each piece is assigned to one or more corresponding expert modules from among the multiple expert modules of the network block; For each of the plurality of fragments, the fragment is processed using one or more corresponding expert modules to generate one or more corresponding module outputs; and A block output with the same dimensions as the block input and which is an updated representation of the input image is generated by combining the module outputs generated by the corresponding expert modules.

2. The system according to claim 1, wherein, Assigning each fragment to one or more corresponding expert modules of the network block includes: The segments are processed using a routing subnetwork to generate a corresponding routing score for each of the plurality of expert modules; Determine one or more of the highest routing scores among the plurality of routing scores; and The fragment is assigned to an expert module corresponding to each of the one or more highest routing scores.

3. The system according to claim 2, wherein, Processing the segments using a routing subnet to generate a corresponding routing score for each of the plurality of expert modules includes: The pieces are processed using one or more feedforward neural network layers to generate a corresponding first intermediate routing score for each expert module; For each expert module, a noise value is sampled and added to the first intermediate route score of the expert module to generate a second intermediate route score for the expert module; For each expert module, a non-linear activation function is applied to the second intermediate routing score of the expert module to generate the routing score for the expert module.

4. The system according to claim 2, wherein, For each of the plurality of fragments, the fragment is processed using one or more corresponding expert modules to generate one or more corresponding module outputs, including: For each of one or more expert modules corresponding to the said fragment, the expert module processes the fragment to generate a corresponding initial module output; and For each of the one or more expert modules corresponding to the segment, the corresponding initial module output is scaled according to the routing score corresponding to the expert module and the segment.

5. The system according to claim 2, wherein, The neural network has been trained using one or more of the following: Importance loss, which encourages a measure of central tendency that corresponds to the routing scores of each expert module, or Load loss, which encourages each expert module to have the same central tendency measure of the corresponding probability of being assigned a piece.

6. The system according to claim 5, wherein, Batch importance loss for training image X include: For each expert module i ,calculate: in, f i Defined by the routing subnetwork, for use against the expert module i Calculate routing scores or for the expert module i A function to calculate intermediate route scores; and Calculate the importance loss To target each expert module i encourage Measures that have the same central tendency.

7. The system according to claim 6, wherein, The importance of the loss for: in, E This refers to the number of the multiple expert modules.

8. The system according to claim 5, wherein: Each piece is assigned to k A different expert module, k ≥1, and Determine the load loss for each batch of training images X. include: For each training image x in the batch of the training image X, calculate: in, g i ( x ) is defined by the routing subnetwork and is used for the expert module. i Calculate routing scores or for the expert module i A function to calculate intermediate route scores; For each training image x and for each expert module i ,calculate: Wherein, P( e ) for event e The probability of it happening; and For each expert module i ,calculate: And calculate the load loss. To target each expert module i Encourage load i ( x ) have the same measure of central tendency.

9. The system according to claim 8, wherein, The load loss for: in, E This refers to the number of the multiple expert modules.

10. The system according to any one of claims 1 to 9, wherein, For each of the plurality of slices, processing the slice using the corresponding expert module includes parallel processing of the slice across the expert modules.

11. The system according to any one of claims 1 to 9, wherein, The number of the multiple expert modules is greater than the number of the multiple pieces.

12. The system according to any one of claims 1 to 9, wherein, The block output of the final network block in the sequence of network blocks is the network output, or the block output of the final 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 system according to any one of claims 1 to 9, wherein: The network output is a classification output, which includes a corresponding score for each of a plurality of categories. The network output is an element-level classification output. For each of one or more elements in the image, the classification output includes a corresponding score for each of a plurality of categories. The network output is a regression output, which estimates one or more continuous variables characterizing the image, or The neural network is configured to process multiple images as video frames to generate an output representing the video frames.

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

15. 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 claimed in any one of claims 1 to 13.

Images