Processing image using self-attention based neural network

Abstract

To provide a method, a system and an apparatus, including computer programs, for processing images using self-attention based neural networks.SOLUTION: A method includes the steps of: obtaining one or more images comprising a plurality of pixels; determining, for each image, a plurality of image patches each comprising a different subset of the pixels of the image; processing, for each image, the image patches to generate an input sequence comprising a respective input element at each of a plurality of input positions, the input elements corresponding to respective different image patches; and processing the input sequences using a neural network to generate a network output that characterizes the one or more images, the neural network comprising one or more self-attention neural network layers.SELECTED DRAWING: Figure 5

Description

[Technical Field]

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

[0002] 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 contain one or more hidden layers in addition to an output layer. The output of each hidden layer is used as the input to the next layer in the network, i.e., the next hidden layer or the output layer. Each layer of the network generates an output from the received input according to the current values ​​of its respective set of parameters. Summary of the Invention [Means for solving the problem]

[0003] This specification describes a system implemented as a computer program on one or more computers at one or more locations that executes a self-attention-based neural network configured through training to process one or more images and generate a network output that characterizes the one or more images.

[0004] The self-attention-based neural network can be configured to process an input sequence representing an image and generate an output sequence by applying a self-attention mechanism across elements of the input sequence. At least some of the elements of the input sequence can correspond to respective patches of the input image. That is, the system can segment the image into patches and process the pixels of each patch to generate respective elements of the input sequence. By applying the self-attention mechanism to these elements, the self-attention-based neural network can focus on the entire image and leverage both local and global information to generate the output sequence.

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

[0006] Some existing systems use self-attention-based neural networks in natural language processing (NLP) use cases to process text sequences and generate predictions about the text sequences. An advantage of self-attention-based neural networks in the NLP domain is scalability; generally, the performance of self-attention-based neural networks improves as the size of the neural network increases. However, the same cannot be said for existing systems that apply self-attention-based neural networks to images; generally, self-attention-based neural networks are unable to scale to larger architectures and therefore do not perform as well as other computer vision systems, such as convolutional neural networks. For example, some such existing systems do not apply self-attention across the entire input image, but instead to local neighborhoods of the input image. Therefore, a first local neighborhood of an image cannot attend to a second local neighborhood of the image.

[0007] The techniques described herein enable systems to process images directly using self-attention-based neural networks and enjoy high performance even as the size of the neural network increases. Specifically, the techniques described herein leverage the parallelization possible with self-attention-based neural networks to enable large-scale training, resulting in improved accuracy for image processing tasks. As a specific example, the systems described herein may be trained on datasets containing 14 million to 300 million images. Furthermore, exemplary implementations described herein apply global self-attention to full-size images. That is, the self-attention-based neural network applies self-attention across the entire input image, such that any region of the image can attend to other regions of the image.

[0008] As described herein, a self-attention-based neural network configured to process images can require much less computation to achieve the same performance as a state-of-the-art convolutional neural network. That is, for a fixed computational budget, a self-attention-based neural network performs better than a convolutional neural network. 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 focus on different regions of the image with fewer computations than a convolution. As a particular example, the self-attention-based neural network described herein can achieve comparable or better performance than a large convolutional neural network while requiring half, five, ten, one hundred, or one thousand times less computation.

[0009] 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 explanation of the drawings]

[0010] [Figure 1] FIG. 1 is a diagram of an exemplary neural network system. [Figure 2] FIG. 1 is a diagram of an exemplary self-attention-based neural network. [Figure 3] FIG. 1 illustrates an example image segmented into image patches. [Figure 4] FIG. 1 is a diagram of an exemplary training system. [Figure 5] 1 is a flow diagram of an example process for generating predictions for one or more images using a self-attention-based neural network. DETAILED DESCRIPTION OF THE INVENTION

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

[0012] This specification describes a system implemented as a computer program on one or more computers at one or more locations configured to execute a self-attention-based neural network configured to process one or more images, i.e., process intensity values ​​of pixels of the one or more images, to generate a network output that characterizes the one or more images.

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

[0014] Neural network system 100 is configured to process image 102 and generate network output 152 that represents a prediction about the image. Neural network system 100 may be configured to perform any suitable machine learning task using image 102. Exemplary machine learning tasks are described below.

[0015] The image may be any suitable type of image. For example, the image may be a two-dimensional image, such as a two-dimensional image having multiple channels (e.g., an RGB image). As another example, the image 102 may be a hyperspectral image that represents a continuous spectrum of wavelengths, for example, by identifying a spectral distribution for each pixel in the image 102. As another example, the image 102 may be a point cloud including multiple points, each point having, for example, a respective coordinate in a three-dimensional or higher-dimensional coordinate space; as a specific example, the image 102 may be a point cloud generated by a LIDAR sensor. As another example, the image 102 may be a medical image generated by a medical imaging device; as a specific example, the image 102 may 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.

[0016] While the following description 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 that include 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 points in the point cloud. As another example, if image 102 is an MRI image that includes multiple voxels in a three-dimensional voxel grid, each image patch of image 102 may include a subset of voxels in the voxel grid.

[0017] Neural network system 100 includes an image patch generation system 110, an image patch embedding system 120, and a neural network 130. As described in more detail below, neural network 130 is a self-attention-based neural network that includes a self-attention-based sub-network 140.

[0018] A self-attention-based neural network is a neural network that includes one or more self-attention neural network layers. The self-attention neural network layers are configured to receive a sequence of layer input elements as input and apply an attention mechanism to the sequence of layer input elements to generate a sequence of layer output elements. Specifically, for each layer input element, the self-attention neural network layer applies an attention mechanism to the layer input element in the sequence of layer input elements using one or more queries derived from the layer input element to generate a respective output element.

[0019] 1, neural network 130 is configured to use self-attention-based sub-network 140 to process an input sequence including input elements representing respective patches of image 102. Thus, neural network 130 can apply an attention mechanism to the input sequence to focus on different patches in different locations in image 102. It will be appreciated that the patches of image 102 can be processed by self-attention-based sub-network 140 using parallel processing, i.e., at least a portion of the processing can be performed in parallel.

[0020] The image patch generation system 110 is configured to process the image 102 and generate n distinct patches 112a-n of the image 102. As used herein, an image patch of an image is a strict subset of the image's pixels. Generally, each image patch 112a-n includes multiple contiguous pixels of the image 102. That is, for each particular image patch 112a-n, and for any pair of pixels in a particular image patch 112a-n, there exists a path from the first pixel of the pair to the second pixel of the pair, and the path includes only pixels in the particular image patch 112a-n.

[0021] In some implementations, each pixel in image 102 is contained in exactly one of the image patches 112a-n. In some other implementations, one or more image patches 112a-n may contain the same pixel from image 102, i.e., two or more of the images may overlap. Alternatively or additionally, one or more pixels from image 102 may be excluded from each of the image patches 112a-n, i.e., one or more pixels are not contained in any of the image patches.

[0022] The image patches 112a-n may be represented in any suitable manner, for example, each image patch 112a-n may be represented as a two-dimensional image that includes the pixels of the image patch 112a-n, e.g., as an image that preserves the spatial relationships of the pixels in the image patch 112a-n.

[0023] As another example, each image patch 112a-n may be represented as a one-dimensional sequence of the pixels of the image patch 112a-n. As a particular example, if the image patch 112a-n is a two-dimensional region of the image 102, the image patch 112a-n may be a flattened version of the two-dimensional region, as described in more detail below. As another particular example, if the image patch 112a-n includes only pixels that share the same column or row of the image 102 (i.e., the image patch 112a-n is a one-dimensional region of the image 102), the image patch 112a-n may be represented as a one-dimensional sequence that preserves the relative positions of the pixels.

[0024] As another example, each image patch 112a-n can be represented as an unordered set of the pixels of the image patch 112a-n.

[0025] An exemplary image patch is described in more detail below with reference to FIG.

[0026] The image patch embedding system 120 is configured to take n image patches 112a-n of an image 10 and generate a respective embedding 122a-n for each of the n image patches 112a-n. Each image patch embedding 122a-n represents pixels of the corresponding image patch 112a-n and may be generated by processing the pixels of the corresponding image patch 112a-n. As used herein, an embedding is an ordered set of numeric values ​​that represent an input in a particular embedding space. For example, an embedding may be a vector of floating-point or other numeric values ​​having a fixed dimensionality.

[0027] In some implementations in which each image patch 112a-n is represented as a two-dimensional sub-image of image 102, each image patch embedding 122a-n is a reshaped version of the corresponding image patch 112a-n. For example, the image patch embedding system 120 may "flatten" each image patch 112a-n to generate an image patch embedding 122a-n that is a one-dimensional tensor containing each pixel in the image patch 112a-n. As a particular example, if each image patch 112a-n has dimensionality L×W×C, where C represents the number of channels in the image (e.g., C=3 for an RGB image), the image patch embedding 122a-n may generate an image patch embedding 122a-n with dimensionality 1×(L·W·C).

[0028] In some other implementations, the image patch embedding system 120 can process a one-dimensional tensor containing the pixels of the image patches 112a-n (e.g., a flattened version of the image patches 112a-n) to generate corresponding image patch embeddings 122a-n. As described in more detail below, the image patch embeddings 122a-n are to be processed by a neural network 130, which, through training, is configured to accept inputs having a particular format, e.g., a particular size and shape. Thus, the image patch embedding system 120 can project each image patch 112a-n into a coordinate space having the number of dimensions required by the neural network 130.

[0029] For example, the image patch embedding system 120 may process each image patch 112a-n using a linear projection. z i =x i E i +b i however,

[0030]

number

[0031] is the i-th image patch embedding 122a-n, D is the number of input dimensions required by the neural network 130,

[0032]

number

[0033] is a one-dimensional tensor containing the i-th image patch 112a-n, N is the number of pixels in the i-th image patch 112a-n, and E i ∈R N×D is the projection matrix,

[0034]

number

[0035] is a linear bias term.

[0036] In some implementations, the image patch embedding system 120 may include a projection matrix E i In some implementations, the image patch embedding system 120 uses the same projection matrix E to generate each image patch embedding 122a-n. Similarly, in some implementations, the image patch embedding system 120 uses each different bias term b i, and in some other implementations, the image patch embedding system 120 uses the same bias term b i is used to generate each image patch embedding 122a-n.

[0037] In some implementations, the linear projection is machine-learned. For example, during training of the neural network 130, the training system determines the parameters of the linear projection (e.g., the projection matrix E i and the bias term b i ) can be simultaneously updated. As a particular example, the training system can update the parameters of the linear projection by backpropagating the training error of the neural network 130 through the neural network 130 to the image patch embedding system 120 and using stochastic gradient descent on the backpropagated error to determine the updates. An exemplary technique for training the neural network 130 is described in more detail below with reference to FIG. 4.

[0038] Instead of, or in addition to, processing one-dimensional tensors corresponding to image patches 112a-n using linear projections, image patch embedding system 120 can process one-dimensional tensors using an embedded neural network. For example, embedding system 120 can be considered a component of neural network 130. That is, embedding system 120 can be an embedded sub-network of neural network 130 that includes one or more neural network layers configured to process one-dimensional tensors and generate image patch embeddings 122a-n.

[0039] For example, the embedded neural network may include one or more feed-forward neural network layers configured to process one-dimensional tensors corresponding to the image patches 112a-n.

[0040] As another example, the embedded neural network may include one or more self-attention neural network layers configured to simultaneously process each one-dimensional tensor corresponding to each image patch 112a-n using a self-attention mechanism, which is described in more detail below.

[0041] As another example, the embedding neural network may include one or more convolutional neural network layers configured to process the image patches 112a-n using convolutional filters. As a particular example, if the image patches 112a-n are represented as two-dimensional images, the image patch embedding system 120 may process each (unflattened) image patch 112a-n using one or more convolutional neural network layers to generate a feature map for the image patch 112a-n. The image patch embedding system 120 may then flatten the feature maps and process the flattened feature maps using linear projection as described above to generate corresponding image patch embeddings 122a-n.

[0042] As another particular example, the image patch embedding system 120 may process the entire image 102 using one or more convolutional neural network layers to generate a feature map for the image 102. The feature map may be two-dimensional (or may be two-dimensional, such as the image 102, with each element having multiple channels). The neural network system 100 may then determine n patches of the feature map for the image 102, each patch including one or more elements of the feature map. That is, instead of segmenting the image 102 itself into image patches 112a-n, the image patch generation system 110 may segment the feature map for the image 102 generated by the embedded neural network of the image patch embedding system 120. As a particular example, each patch may include a single element of the feature map. The image patch embedding system 120 may then generate image patch embeddings 122a-n from the n patches of the feature map, e.g., by applying a linear projection to the patches of the feature map as described above.

[0043] After the image patch embedding system 120 generates the image patch embeddings 122a-n, the neural network system 100 may generate an input sequence from the image patch embeddings 122a-n to be provided as input to the neural network 130. Generally, the input sequence includes one or more input elements corresponding to each image patch embedding 122a-n. For example, the input sequence may include a respective input element corresponding to each of the n image patch embeddings 122a-n. As a particular example, the input elements corresponding to the n image patch embeddings 122a-n may be sorted into the input sequence in the raster order of the corresponding image patches 112a-n.

[0044] In some implementations, the input element of the input sequence corresponding to the image patch embedding 122a-n is equal to the image patch embedding 122a-n itself.

[0045] In some other implementations, to generate input elements of the input sequence corresponding to the image patch embeddings 122a-n, the neural network system 100 can combine (i) the image patch embeddings 122a-n and (ii) positional embeddings that represent the locations within the image 102 of the image patches 112a-n that correspond to the image patch embeddings 122a-n. For example, the neural network system 100 can append the positional embeddings to the image patch embeddings 122a-n. By incorporating the positional embeddings, the neural network system 100 can encode spatial information, such as the relative location of each image patch in the image, that can be leveraged by the neural network 130 to generate the network output 152.

[0046] In some implementations, the positional padding corresponding to each image patch 112a-n of image 102 is an integer. For example, the first image patch in the top left of image 102 may have a positional padding of "1," the second image patch immediately to the right of the first image patch may have a positional padding of "2," etc.

[0047] In some implementations, the location embeddings are machine-learned. For example, during training of the neural network 130, the training system can concurrently learn the location embeddings by backpropagating the training error of the neural network 130 through the neural network 130 to the location embeddings. In some such implementations, the training system can generate a different location embedding for each image patch (e.g., assuming that all images 102 received by the neural network system 100 are segmented into the same number of patches).

[0048] In some other implementations, the training system can incorporate two-dimensional information into the positional embedding by learning, for both dimensions of the image 102, a respective positional embedding for each coordinate along the dimension. For example, if the image 102 is segmented into a two-dimensional grid of image patches 112a-n, the training system can generate two sets of positional embeddings: a first set including a respective positional embedding for each index along the vertical axis of the grid, and a second set including a respective embedding for each index along the horizontal axis of the grid. To generate the positional embedding for a particular image patch 112a-n, the neural network system can combine, e.g., by concatenation, (i) the positional embedding corresponding to the particular image patch 112a-n's index along the vertical axis and (ii) the positional embedding corresponding to the particular image patch 112a-n's index along the horizontal axis.

[0049] In some implementations, one or more of the input elements in the input sequence do not correspond to any of the image patches 112a-n of the images 102. For example, the input sequence may include class embeddings 124 that are the same for all received images 102. For example, the class embeddings 124 may be tensors having the same dimensionality as the image patch embeddings 122a-n. As a particular example, the class embeddings 124 may be tensors of all "0's" or all "1's."

[0050] The class embedding 124 can be inserted anywhere in the input sequence; for example, the class embedding 124 can be the first input element of the input sequence or the last input element of the input sequence.

[0051] In some implementations, the class embeddings 124 are machine-learned. For example, during training of the neural network 130, the training system can simultaneously learn the parameters for the class embeddings 124 by backpropagating the training error of the neural network 130 through the neural network 130 to the class embeddings 124.

[0052] In implementations in which the input elements corresponding to each image patch 112a-n include a positional embedding corresponding to the image patch 112a-n, the neural network system 100 can similarly append the positional embedding to the class embedding 124, for example, appending a machine-learned positional embedding or a predetermined positional embedding (e.g., an all "0" or all "1" positional embedding).

[0053] After generating the input sequence, the neural network system 130 may provide the input sequence as an input to the neural network 130. The neural network 130 may process the input sequence to generate a network output 152.

[0054] In particular, the neural network 130 can process the input sequence using a self-attention-based sub-network 140 to generate an output sequence. In some implementations, the neural network 130 generates an output sequence that is the same length as the input sequence, i.e., that includes a respective output element for each input element in the input sequence. In particular, the output sequence can include class outputs 144 generated from the class embeddings 124 and respective image patch outputs 142a-n corresponding to each image patch embedding 122a-n in the input sequence.

[0055] The self-attention-based sub-network 140 may include one or more self-attention neural network layers, each configured to receive a layer input sequence and apply a self-attention mechanism to the layer input sequence to generate a layer output sequence. In some such implementations, the self-attention-based sub-network 140 may include a sequence of network blocks, each configured to receive a respective block input sequence including a respective element corresponding to each input element in the input sequence, process the block input sequence, and generate a respective block output sequence including a respective element for each input element in the input sequence. Each network block may include one or more self-attention neural network layers. An exemplary self-attention-based neural network is described in more detail below with reference to FIG. 2.

[0056] After the self-attention-based sub-network 140 generates the output sequence, the neural network 130 may provide one or more elements of the output sequence to the head sub-network 150 .

[0057] For example, head sub-network 150 may be configured to process n image patch outputs 142a-n. As a particular example, head sub-network 150 may combine the n image patch outputs 142a-n (e.g., using global average pooling) to generate a combined patch output, and then process the combined patch output to generate network output 152. For example, head sub-network 150 may process the combined patch output using one or more feed-forward neural network layers and / or linear classifiers.

[0058] As another example, head sub-network 150 may be configured to process only class output 144 to generate network output 152. That is, class output 144 may represent a final representation of image 102, and head sub-network 150 may process class output 144 to generate network output 152 that represents a prediction for image 102. For example, head sub-network 150 may include a multi-layer perceptron with one or more feed-forward neural network layers.

[0059] In some implementations, the self-attention-based subnetwork 140 and the head subnetwork 150 are trained end-to-end simultaneously on a single machine learning task. For example, the training system can perform a supervised training process using a training dataset including multiple training examples, each including a training input sequence (representing a respective training image) and a corresponding ground truth network output, i.e., an output representing the network output 152 that the neural network 130 should generate in response to processing the training input sequence. The training system can process the training input sequences using the neural network 130 to generate respective predicted network outputs and determine parameter updates to the head subnetwork 150 and the self-attention-based subnetwork 140 according to the error between (i) the predicted network output and (ii) the corresponding ground truth network output. For example, the training system can determine the parameter updates by backpropagating the error through both the head subnetwork 150 and the self-attention-based subnetwork 140 and performing stochastic gradient descent.

[0060] In some other implementations, the self-attention-based subnetwork 140 has been trained using transfer learning with one or more other head subnetworks different from the head subnetwork 150, e.g., configured to perform different respective machine learning tasks than the head subnetwork 150. For example, a training system may simultaneously train the self-attention-based subnetwork 140 and one or more other head subnetworks, then remove the one or more other head subnetworks and replace them with the head subnetwork 150 to generate the neural network 130. The training system may then fine-tune the neural network 130 to generate trained parameters for the head subnetwork 150. An exemplary technique for training the neural network 130 using transfer learning is described in more detail below with reference to FIG. 4.

[0061] In some implementations, the neural network includes one or more additional sub-networks, for example, one or more sub-networks immediately before the self-attention-based sub-network 140 (e.g., a sub-network including one or more recurrent neural network layers configured to process an input sequence) or immediately after the self-attention-based sub-network 140 (e.g., a sub-network including one or more recurrent neural network layers configured to process an input sequence).

[0062] In some implementations, neural network 130 does not include head sub-network 150. For example, neural network system 100 may be configured to generate an embedding of image 102, where the embedding includes (or is generated from) one or more of image patch outputs 142a-n and / or class output 144. Neural network system 100 can then provide the embedding of image 102 to a downstream system for storage or further processing, for example, by one or more other neural networks.

[0063] For example, the neural network system may be configured to receive images 102 from an external system and provide embeddings of the images 102 to the original external system, e.g., by providing image patch outputs 142a-n and / or class output 144 to the original external system. For each image, the external system may be configured to process the image embeddings using a neural network to generate predictions for the image; e.g., the external system may include head sub-network 150. As a particular example, the neural network system 100 may be configured to receive images from an edge device, e.g., a mobile phone, a tablet computer, or an autonomous vehicle. The edge device may then execute head sub-network 150 to generate predictions for the images.

[0064] As described in more detail below with reference to FIG. 4, in some implementations, the self-attention-based subnetwork 140 includes many more parameters than the head subnetwork 150 and may therefore be more computationally expensive to execute. Thus, an edge device may not have the computational resources to execute the self-attention-based subnetwork 140. Thus, the neural network system 100 may be configured to execute the self-attention-based subnetwork 140 (e.g., using one or more parallel processing devices such as a GPU or TPU), allowing the edge device to perform the relatively computationally inexpensive task of executing the head subnetwork 150. For example, the neural network system 100 may be deployed in the cloud and communicatively connected to multiple different edge devices.

[0065] The neural network system 100 may be configured to perform any suitable machine learning task on the images 102, such as a classification task, a regression task, or a combination thereof.

[0066] As a particular example, the neural network system 100 may be configured to generate a classification output including a respective 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 may be a class of objects (e.g., dog, cat, person, etc.), and an image may belong to a category if it depicts an object included in the object class corresponding to the category. In some cases, a category may represent a global image property (e.g., whether the image depicts a daytime scene or a nighttime scene, or whether the image depicts a summer scene or a winter scene), and an image may belong to a category if it has a global property corresponding to the category.

[0067] As another specific example, the neural network system 100 may be configured to generate a pixel-level classification output for each pixel of an image, the pixel-level classification output including a respective score corresponding to each of a plurality of categories. For a given pixel, the category score indicates the likelihood that the pixel belongs to that category. In some cases, the categories may be classes of objects, and a pixel may belong to a category if the pixel is part of an object included in the object class corresponding to the category. That is, the pixel-level classification output may be a semantic segmentation output.

[0068] As another particular example, the neural network system 100 may be configured to generate a regression output that estimates one or more continuous variables (i.e., capable of assuming an infinite number of possible values) that characterize an image. In a particular example, the regression output may estimate the coordinates of a bounding box that surrounds each object shown in the image. The coordinates of the bounding box may be defined by the (x, y) coordinates of the vertices of the bounding box. For example, the system may output two (x, y) coordinates of the bounding box coordinates, or it may output the coordinates of the center of the bounding box, as well as the height and width of the bounding box.

[0069] In some implementations, the neural network system 100 may be configured to perform a video analysis task. For example, the neural network system 100 may receive a plurality of images 102 that are video frames of a video and may process each video frame as described above to generate an output that characterizes the video frame, for example, by characterizing whether the video frame shows a person performing a particular action.

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

[0071] In some other implementations, the neural network system 100 processes each video frame simultaneously to generate a single network output 152 that characterizes the video. That is, the neural network system 100 may be configured to process multiple images 102 simultaneously. For example, the neural network system 100 may generate a respective input sequence for the neural network 130 corresponding to each image 102 as described above. The neural network system 100 may then combine the multiple input sequences into a single combined input sequence, e.g., by concatenating the input sequences, and then process the combined input sequence using the neural network 130.

[0072] Figure 2 is a diagram of an exemplary self-attention-based neural network 200. The self-attention-based neural network 200 may be a component of a neural network system configured to process images and generate predictions about the images, such as the neural network system 100 described above with reference to Figure 1. In particular, the self-attention-based neural network may be a sub-network of a neural network configured to generate network outputs representing predictions about the images, such as the neural network 130 described above with reference to Figure 1.

[0073] The self-attention-based neural network 200 is configured to process an input sequence 202 that represents an image and includes a respective input element at each of a plurality of input locations. For example, the input sequence 202 may include a respective input element that represents each of a plurality of image patches of the image, as described above with reference to FIG. 1. The self-attention-based neural network 200 is configured to process the input sequence 202 to generate an output sequence 204 that has the same length as the input sequence 202, i.e., has as many output elements as there are input elements in the input sequence 202.

[0074] The self-attention-based neural network 200 includes M network blocks 210a-m, where M≧1. Each network block 210a-m is configured to receive a block input sequence including a respective block input element for each input position in the input sequence 202, i.e., each block input element corresponds to a respective input element in the input sequence 202. Each network block 210a-m is configured to process the block input sequence and generate a block output sequence including a respective block output element for each of a plurality of input positions in the input sequence. That is, each block input sequence 212 preserves the number of elements in the input sequence 202 as the input sequence is processed by the neural network 200.

[0075] The first network block 210a in the sequence may receive input sequence 202. Each subsequent network block 210a-m in the sequence may receive as its block input sequence the respective block output sequence generated by the previous network block 210a-m in the sequence. The block output sequence of the Mth and final network block 210m may be output sequence 204.

[0076] Each network block 210a-m includes one or more self-attention neural network layers. Referring to the kth network block 210k, the network block 210k includes a single self-attention neural network layer 220. In some implementations, the self-attention neural network layer 220 is configured to obtain each block input element in the block input sequence 212 and apply an attention mechanism to the block input element. In some other implementations, the self-attention neural network layer 220 is configured to obtain a processed version of each block input element in the block input sequence 212 and apply an attention mechanism to the processed block input element. For example, as shown in FIG. 2, the network block 210k may first apply a layer normalization layer to the block input sequence 212 before providing the output of the layer normalization layer to the self-attention neural network layer 220. Alternatively or additionally, the network block 210k may apply one or more other neural network layers to the block input sequence 212 before the self-attention neural network layer 220, for example, an element-wise feedforward neural network layer.

[0077] In particular, for each block input element (or a processed version thereof) corresponding to each particular input position, the self-attention neural network layer 220 is configured to apply an attention mechanism to the block input element at the input position (i.e., other block input positions and, in some implementations, itself) using one or more queries derived from the block input element at the particular input position to generate a respective output for the particular position. The output of the self-attention neural network layer 220 is a layer output sequence that includes a respective layer output element corresponding to each input position.

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

[0079] In some implementations, some or all of the self-attention neural network layers of the self-attention-based neural network 200 (e.g., the self-attention neural network layer 220 shown in FIG. 2) incorporate position information of each block input element in the block input sequence into the attention mechanism.

[0080] For example, when applying attention with respect to a particular block input element (i.e., when generating each layer output element corresponding to the particular block input element), the self-attention neural network layer can identify an attention position embedding that represents the position of the image patch corresponding to the particular block input element within the image. For example, the attention position embedding corresponding to each image patch can be the same as the position embedding embedded in the input sequence 202.

[0081] Upon generating the respective layer output elements corresponding to particular block input elements, the self-attention neural network layer may perform, e.g., sequentially or in parallel, the following two different attention computations: (i) a first attention computation (i.e., the attention mechanism described above) in which queries generated from the particular block input elements focus on the set of keys generated from the respective block input elements, and (ii) a second attention computation in which queries generated from the attention location embeddings of the particular block input elements focus on the set of keys generated from the attention location embeddings of the respective block input elements. The self-attention neural network layer may then combine the outputs of the two attention computations to generate a final layer output element for the particular block input element, e.g., by determining the sum of the outputs of the two attention computations.

[0082] As another example, applying attention with respect to a particular block input element can cause the self-attention neural network layer to determine the respective offsets between (i) the particular block input element and (ii) each other block input element. For example, adjacent block input elements in a block input sequence may have an offset such as "1."

[0083] The self-attention neural network layer can identify a respective offset embedding corresponding to each offset, e.g., the offset embedding corresponding to each offset can be machine-learned during training of the neural network 200, e.g., as described above with respect to the position embeddings embedded in the input sequence 202.

[0084] The self-attention neural network layer can then perform two attention computations as described above, except that during the second attention computation, the offset embedding is used instead of the attention location embedding.

[0085] In some implementations, one or more of the network blocks 210a-m include a residual connection layer that combines the outputs of the self-attention neural network layer with the inputs of the self-attention neural network layer. Alternatively or additionally, one or more of the network blocks 210a-m may include a layer normalization layer that applies layer normalization to the inputs and / or outputs of the self-attention neural network layer. These layers are referred to as "Norm" in FIG. 2.

[0086] In some implementations, one or more network blocks 210a-m include one or more position-wise feedforward neural network layers. For example, the kth network block 210k includes a feedforward neural network layer 230. The feedforward layer 230 is configured to, for each input position of the input sequence 202, receive an input element for that position and apply a sequence of transformations to the input element for that position to generate an output element for that position. For example, the sequence of transformations may include two or more trained linear transformations, each separated by an activation function, e.g., a nonlinear element-wise activation function, e.g., a ReLU activation function. As a particular example, the feedforward neural network may be a multilayer perceptron including one, two, or more feedforward neural network layers. The input element received by the position-wise feedforward layer 230 may be the output of a layer normalization layer following the self-attention neural network layer 220, or the input element received by the position-wise feedforward layer 230 may be the output of the self-attention neural network layer 220 itself when there is no layer normalization layer.

[0087] In some implementations, one or more of the network blocks 210a-m includes a residual connection layer that combines the output of a feedforward neural network layer for each position with the input to the feedforward neural network layer for each position.

[0088] As described above, in some implementations, each input element of the input sequence 202 includes or is generated from a respective positional embedding. Instead of or in addition to incorporating a positional embedding into the input elements of the input sequence 202, the self-attention-based neural network 200 can incorporate a positional embedding into the block input elements of one or more of the respective block input sequences 212 of the network blocks 210 a-m. For example, before processing each of the block input sequences 212 of each of the network blocks 210 a-m, the self-attention-based neural network can append a respective positional embedding, e.g., a machine-learned positional embedding, to each block input element of the block input sequence 212. In some implementations, each of the network blocks 210 a-m uses a different set of trained positional embeddings. In some other implementations, each of the network blocks 210 a-m uses the same set of trained positional embeddings.

[0089] After generating output sequence 204, self-attention-based neural network 200 can provide output sequence 204 to one or more downstream systems. For example, self-attention-based neural network 200 can provide output sequence 204 to one or more head neural networks to generate predictions for respective machine learning tasks, as described above with reference to FIG. 1. As another example, self-attention-based neural network 200 can provide output sequence 204, which can represent embeddings of images corresponding to input sequence 202, to a database or to one or more downstream machine learning models for further processing.

[0090] FIG. 3 shows example images 310, 320, 330, 340, 350, and 360 segmented into image patches.

[0091] The images 310-360 may be provided as inputs to a neural network system, such as the neural network system 100 described above with reference to Figure 1, configured to process the images 310-360 and generate predictions for the images 310-360. The neural network system may include an image patch generation system, such as the image patch generation system 110 described above with reference to Figure 1, that segments the images 310-360 into a plurality of image patches. The image patches, or network inputs generated from the image patches, may then be processed by a self-attention-based neural network to generate predictions for the images.

[0092] Images 310, 320, 330, 340, 350, and 360 illustrate different possibilities for segmenting an image into image patches. Specifically, in FIG. 3, each image 310 is shown segmented into a set of image patches that are each visually distinct, i.e., using different shading or shading. Generally, an image generation system is configured to segment all received images according to the same scheme. That is, the image generation system does not necessarily segment each of the illustrated images 310, 320, 330, 340, 350, and 360 because they are segmented according to different schemes.

[0093] As shown in the first image 310, in some implementations, the image patch generation system can generate image patches each having the same size and shape, e.g., each image patch can be rectangular. Further, in some implementations, the image patch generation system can segment the first image 310 such that every pixel is a member of exactly one image patch. As a particular example, the image patches can represent rectangular grids of the same size, as shown in FIG. 3. As another particular example, the image patches can represent hexagonal grids of the same size.

[0094] As shown in the second image 320, in some implementations, the image patch generation system can generate image patches having different sizes.

[0095] As shown in the third image 330, in some implementations, the image patch generation system can segment the third image 330 such that some pixels are members of multiple different image patches.

[0096] As shown in fourth image 340, in some implementations, the image patch generation system can segment fourth image 340 such that some pixels are not members of any image patch. For example, the image generation system can process fourth image 340 using a machine learning model to identify one or more regions of interest, and the image patch generation 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 image patch generation system can generate a respective patch centered on each identified pixel.

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

[0098] As shown in sixth image 360, in some implementations, the image patch generation system may generate a one-dimensional space-filling curve that includes each pixel of the image. The image patch generation system may then segment the one-dimensional space-filling curve to generate a set of one-dimensional image patches. As a particular example, the image patch generation system may segment each column or row of the image into n subsequences of pixels in that column or row, with each subsequence representing an image patch.

[0099] Although images 310-360 are shown in FIG. 3 as two-dimensional images (or images having multiple channels and that are two-dimensional, e.g., RGB images), in general, a neural network system may be configured to generate predictions for any type of image, as described above with reference to FIG. 1.

[0100] 4 is a diagram of an example training system 400. Training system 400 is an example of a system in which the systems, components, and techniques described below may be implemented, implemented as a computer program on one or more computers at one or more locations.

[0101] The training system 400 is configured to train a base neural network 420 to perform a first machine learning task and to train one or more task neural networks 450 to perform respective second machine learning tasks that are different from the first machine learning task. In particular, the training system 400 can use the trained parameters from the base neural network 420 to generate parameters for the one or more task neural networks 450.

[0102] The base neural network 420 and the one or more task neural networks 450 are each configured to process an input sequence representing an image, the input sequence including one or more elements corresponding to respective image patches of the corresponding image, as described above. The first machine learning task and the one or more second machine learning tasks can each be any suitable machine learning task. For example, the first machine learning task and the one or more second machine learning tasks can include one or more of the tasks described above with reference to FIG. 1.

[0103] In some implementations, the first machine learning task (i.e., the machine learning task on which the self-attention-based subnetwork 430 was pre-trained) is a self-supervised machine learning task. That is, the training system 400 can train the base neural network 420 using a training dataset that does not include ground truth labels, but instead uses a portion of the training dataset, the ground truth labels for the remainder of the training dataset. As a particular example, the first machine learning task can be a masked image prediction task, in which the base neural network 420 processes an input sequence 412 representing a portion of a respective image (i.e., an image in which one or more pixels are “masked”) and generates a base network output 422 representing a prediction of the content of the masked portion of the image. For example, the training system 400 can generate an initial input sequence including a respective element for each image patch in the complete, unmasked image, and then remove one or more of the elements or replace one or more of the elements using the same “mask” token that identifies the corresponding image patch as having been masked from the image.

[0104] Each of the neural networks includes a respective self-attention-based subnetwork and a respective head subnetwork. Specifically, base neural network 420 includes self-attention-based subnetwork 430 and base head subnetwork 440, and task neural networks 450 include a respective self-attention-based subnetwork 460 and a respective task head subnetwork 470. Each self-attention-based subnetwork and each head subnetwork 440 may be configured as described above with reference to self-attention-based subnetwork 140 and head subnetwork 150 shown in FIG. 1, respectively.

[0105] Each of the self-attention-based subnetworks 430 and 460 can be configured similarly to one another, e.g., they can have the same network architecture with the same number and size of neural network layers. Each of the base head subnetworks 440 and 470, however, can be configured specifically for a corresponding machine learning task. That is, the base head subnetwork 440 can be configured specifically for a first machine learning task, and each task head subnetwork 470 can be configured specifically for a corresponding second machine learning task. For example, each head subnetwork 440 and 470 can be configured to generate a respective network output having a format required for the corresponding machine learning task. Thus, different head subnetworks can have different network architectures. As a specific example, one or more of the head subnetworks can be a multilayer perceptron including one, two, or more feedforward neural network layers.

[0106] The training system 400 includes a training data store configured to maintain respective training datasets for a first machine learning task and one or more second machine learning tasks. In some implementations, the training dataset for the second machine learning task is smaller than the training dataset for the first machine learning task, and thus the training system 400 may be configured to leverage trained parameters of the base neural network 420, which was trained using a larger training dataset, to compensate for the training of the task neural network 450, which may be hindered by a relative lack of training data.

[0107] Each training dataset can be generated from a respective set of training images, i.e., each training example in a particular training dataset can be generated from a respective training image from the set of training images corresponding to the training dataset. In some implementations, each training dataset is generated from the same set of images, while in some other implementations, different training datasets can be generated from different sets of images.

[0108] The training system can provide an input sequence 412 from a training dataset corresponding to a first machine learning task to a base neural network 420. The training system 400 can process the input sequence using a self-attention-based sub-network 430 and a base head sub-network 440 to generate a base network output 422.

[0109] The training system 400 includes a training engine 480 configured to obtain a base network output 422, determine an error in the base network output 422, and generate parameter updates 482 for the base neural network 420 according to the error. The training engine 480 may generate the parameter updates 482 using any suitable training technique. For example, the training engine 480 may use one or more of supervised learning, unsupervised learning, semi-supervised learning, self-supervised learning, distillation learning (in which the base neural network 420 is trained to generate base network outputs 422 that match the outputs of a “teacher” neural network), or adversarial learning (in which the base neural network 420 is trained by a discriminator neural network to generate base network outputs 422 that are predicted not to have been generated by the base neural network 420).

[0110] As a particular example, in an implementation in which training system 400 performs supervised learning using a training dataset that includes ground truth network outputs, training engine 480 can determine the difference between base network output 422 and the ground truth network output corresponding to input sequence 412. Training engine 480 can generate parameter updates 482 by backpropagating errors through base neural network 420 and performing stochastic gradient descent. Parameter updates 482 can include respective updates for parameters of both self-attention-based sub-network 430 and base head sub-network 440.

[0111] After completing training of the base neural network 420, the training system 100 can generate one or more task neural networks 450 using the trained parameters of the base neural network 420.

[0112] In particular, training system 400 can obtain trained parameters 432 of the self-attention-based sub-network and apply the trained parameters to self-attention-based sub-network 460, i.e., set the parameters of self-attention-based sub-network 460 to be the same as trained parameters 432. As explained above, each of self-attention-based sub-networks 430 and 460 can be configured similarly to one another, and thus training of sub-network 430 can be transferred to each of sub-networks 460.

[0113] During training of the base neural network 420, the self-attention-based sub-network 430 can learn to generate representations of input sequences that encode information about the corresponding input images, which are useful for performing multiple different machine learning tasks, including a first machine learning task and one or more second machine learning tasks. That is, the information encoded into the representations by the self-attention-based sub-network 430 can be useful for a second machine learning task even if the self-attention-based sub-network 430 was not trained using the second machine learning task.

[0114] However, instead of taking the trained parameters of the base head subnetwork 440 and applying them to the task head subnetwork 470, the training system 100 can generate initialized parameters 472 for the task subnetwork 470, for example, by randomly sampling the initialized parameters. Because each of the head subnetworks 440 and 470 is configured specifically for its respective machine learning task, the training of the base head subnetwork 440 cannot be transferred to the task head subnetwork 470.

[0115] In other words, to generate one of the task neural networks 450, the training system 400 can discard the base head subnetwork 440 from the base neural network 420 and replace the base head subnetwork 440 with a newly initialized task head subnetwork 470.

[0116] In some implementations, the self-attention-based subnetworks 430 and 460 are larger (e.g., have larger network parameters) than the head subnetworks 440 and 470. As a particular example, the self-attention-based subnetworks may include millions, hundreds of millions, billions, or hundreds of billions of parameters, and the head subnetworks may include hundreds, thousands, or hundreds of thousands of parameters. Thus, by pre-training the self-attention-based subnetwork 430 and using the pre-trained subnetwork 430 to determine parameters for one or more task head subnetworks 470, the training system 400 can train the task neural networks 450 much more efficiently than if each task neural network 450 were trained from scratch. Most of the time and computational cost of training the task neural network 450 can be completed “upfront” during the pre-training of the self-attention-based subnetwork 430. That is, the cost of training the self-attention-based subnetwork 430 can be amortized over multiple task neural networks 450.

[0117] The training system can train the task neural network to generate trained parameters for the new task head sub-network 470 .

[0118] In particular, for each task neural network 450, training system 100 may use task neural network 450 to process an input sequence 414 from a training dataset that corresponds to a second machine learning task of task neural network 450 and generate task network output 452. In some implementations, input sequence 414 has a different format (e.g., a different number or size of elements) than input sequence 412. In these implementations, task neural network 450 may include one or more input neural network layers before self-attention-based sub-network 460 configured to project input sequence 414 into the number of dimensions that self-attention-based sub-network 460 is configured to handle.

[0119] In some implementations in which machine-learned position embeddings are incorporated into elements of the input sequence 414, the training system 400 fine-tunes the position embeddings during training of the task neural network 450. In some other such implementations, the training system 400 does not fine-tune the position embeddings.

[0120] In some implementations, the set of training images used to generate the training data set for base neural network 420 includes images having a different resolution than the set of training images used to generate the training data set for task neural network 450. That is, input sequence 414 can represent images having a different resolution than the images represented by input sequence 412.

[0121] As a particular example, input sequence 414 may represent images of higher resolution than input sequence 412, i.e., the self-attention-based sub-network may be fine-tuned on larger images than during its original training. In some implementations where each image patch is the same size and each pixel of each image is contained in exactly one image patch, input sequence 412 and input sequence 414 may have the same number of elements, i.e., each element of input sequence 414 represents a larger image patch than each element of input sequence 412.

[0122] In some other such implementations, elements of input sequence 412 and input sequence 414 may represent image patches of the same size, i.e., input sequence 414 may be longer than input sequence 412. This may create problems in implementations that incorporate machine-learned positional embeddings for elements of input sequences 412 and 414 (e.g., positional embeddings learned during training of base neural network 420, as described above), as the additional elements of input sequence 414 do not have learned positional embeddings.

[0123] Thus, in some implementations, training system 400 can determine positional embeddings for additional elements using positional embeddings learned for the input sequences of base neural network 420. For example, training system 400 can initialize positional embeddings for additional elements (e.g., by initializing each positional embedding to be zero or by randomly initializing the positional embeddings) and train the positional embeddings during training of task neural network 450.

[0124] As another example, the training system 400 may perform two-dimensional interpolation on the learned positional embeddings for the base neural network 420 according to the locations of the additional image patches in the training images. To name just a few examples, the training system 400 may use two-dimensional linear interpolation, two-dimensional bicubic interpolation, or two-dimensional Lanczos interpolation.

[0125] The training engine 480 can take the task network output 452, determine an error in the task network output 452, and generate parameter updates 484 for the task neural network 450 according to the error.

[0126] In some implementations, the parameter updates 484 include respective updates to the parameters of both the self-attention-based sub-network 460 and the base head sub-network 470. That is, the training system 400 can further fine-tune the parameters of the self-attention-based sub-network 460 even though the parameters have already been trained during training of the base neural network 420.

[0127] In some other implementations, the parameter updates 484 include only updates to the parameters of the task head subnetwork 470. That is, the training system 100 can "freeze" the parameters of the self-attention-based subnetwork 460 while training the task neural network 450.

[0128] In some implementations, instead of pre-training the base neural network 420, the training system 400 can obtain the trained parameters 432 of the self-attention-based sub-network 430 from an external system.

[0129] In some implementations, after one or more task neural networks 450 are trained, as described above with reference to FIG. 1, the self-attention-based subnetwork 460, each task head subnetwork 470, may be separately deployed, e.g., on separate, communicatively connected computing devices. For example, the self-attention-based subnetwork 460 may be deployed at a data center, and the task head subnetwork 470 may be deployed on an edge device with limited computational resources. The edge device may then provide the image to the data center, which may process the image using the self-attention-based subnetwork 460 to generate an embedding for the image. The data center may then provide the embedding back to the edge device, which may process the embedding using the respective task head subnetwork 470 to generate a prediction for the image.

[0130] 5 is a flow diagram of an exemplary process 500 for generating predictions for one or more images using a self-attention-based neural network. For convenience, process 500 is described as being performed by one or more computer systems at one or more locations. For example, a neural network system suitably programmed in accordance with this specification, such as neural network system 100 shown in FIG. 1, can perform process 500.

[0131] The self-attention-based neural network may include one or more self-attention neural network layers, for example, the self-attention-based neural network may be the neural network 130 described above with reference to FIG.

[0132] The system acquires one or more images (step 502), each of which includes multiple pixels.

[0133] For each image of the one or more images, the system determines a set of image patches of the image (step 504), each image patch comprising a different subset of the pixels of the image.

[0134] For each image of the one or more images, the system processes a corresponding set of image patches to generate an input sequence (step 506). The input sequence may include a respective element at each of a plurality of input locations, one or more of the input elements corresponding to a respective patch of the image.

[0135] In some implementations, for each image patch, the system can generate respective one-dimensional initial input elements that include pixels of the image patch. For example, the initial input elements can be a flattened version of the image patch. The system can then use the initial input elements to generate input elements that correspond to the image patch. For example, the system can generate the input elements by processing the initial input elements using a second neural network. The second neural network can be an embedded neural network that is a component of an image patch embedding system, e.g., image patch embedding system 120 described above with reference to FIG. 1. As a particular example, the embedded neural network can include one or more fully-connected neural network layers.

[0136] In some implementations, the system can process the image patch to generate respective intermediate input elements. For example, the intermediate input elements can be flattened versions of the image patch or processed versions thereof (e.g., as generated by an embedded neural network as described above). The system can then combine each intermediate input element with a respective positional embedding that represents the position of the corresponding image patch in the image. For example, each positional embedding can be an integer. As another example, each positional embedding can be machine learned.

[0137] In some implementations, an input sequence corresponding to a particular image includes one or more input elements in addition to the input elements corresponding to the image patches of the image. For example, the input sequence may include machine-learned tensors, such as the class embeddings 124 described above with reference to FIG.

[0138] The system processes the input sequence using a self-attention-based neural network to generate a network output that characterizes one or more images (step 508).

[0139] For example, the system may process an input sequence using a self-attention-based sub-network of a self-attention-based neural network (e.g., the self-attention-based sub-network 140 described above with reference to FIG. 1 ) to generate a respective output element for each input element of the input sequence.

[0140] The system can then process one or more output elements using a third neural network element to generate a network output. For example, the third neural network can be another sub-network of the self-attention-based neural network, such as a head sub-network configured similarly to head sub-network 150 described above with reference to FIG. 1. As a particular example, the head sub-network can be configured to process only output elements that correspond to machine-learned tensors (e.g., class embeddings 124) to generate the network output.

[0141] In some such implementations, the head sub-network is configured to generate a first type of network output (e.g., corresponding to a first machine learning task), and the self-attention-based sub-network is trained simultaneously with a fourth neural network to generate a second type of network output (e.g., a second machine learning task) that is different from the first type. The fourth neural network may be a different head sub-network, e.g., the base head sub-network 440 described above with reference to FIG. 4.

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

[0143] Embodiments of the subject matter and functional operations described herein may be implemented in digital electronic circuitry, in tangibly embodied computer software or firmware, in computer hardware containing the structures disclosed herein and their structural equivalents, or in one or more combinations thereof. 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 controlling the operation of a data processing apparatus. The computer storage medium may be a machine-readable storage device, a machine-readable storage substrate, a random or serial access memory device, or a combination of one or more of these. 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 suitable receiver apparatus for execution by the data processing apparatus.

[0144] 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 can also be or include special-purpose logic circuitry, e.g., an FPGA (field-programmable gate array) or an ASIC (application-specific integrated circuit). In addition to hardware, an apparatus can sometimes include code that creates an execution environment for a computer program, e.g., code that constitutes processor firmware, a protocol stack, a database management system, an operating system, or a combination of one or more of these.

[0145] A computer program, sometimes 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, or declarative or procedural, 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, but need not, correspond to a file in a file system. A program can be stored as part of a file with other programs or data, e.g., one or more scripts stored in a markup language document, in a single file dedicated to the program, or in multiple cooperating files, e.g., files storing one or more modules, subprograms, or portions of code. A computer program can be deployed to be executed on one computer or on multiple computers located at one site or distributed across multiple sites and interconnected by a data communications network.

[0146] 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 which 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 differently.

[0147] 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 at one or more locations. In some cases, one or more computers are dedicated to a particular engine, and in other cases, multiple engines can be installed and executed on the same one or more computers.

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

[0149] A computer suitable for running a computer program can be based on a general-purpose microprocessor, a special-purpose microprocessor, or both, or other types of central processing unit. Typically, the central processing unit will receive instructions and data from a read-only memory, a random-access memory, or both. The essential elements of a computer are a central processing unit for performing 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. Generally, a computer also includes one or more mass storage devices, e.g., magnetic, magneto-optical, or optical disks, for storing data, or is operatively coupled to the mass storage devices to receive data from, transfer data to, or both. However, a computer need not have such devices. Furthermore, a computer can be incorporated in another device, e.g., 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.

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

[0151] To provide for user interaction, embodiments of the subject matter described herein may be implemented on a computer having a display device, such as a CRT (cathode ray tube) or LCD (liquid crystal display) monitor, for displaying information to the user, as well as a keyboard and a pointing device, such as a mouse or trackball, through which the user can provide input to the computer. Other types of devices may also be used to provide for user interaction. For example, 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 input, voice input, or tactile input. Additionally, a computer may interact with a user by sending documents to and receiving documents from a device used by the user, e.g., by sending a web page to a web browser on the user's device in response to a request received from that 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 reply message from the user in return.

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

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

[0154] Embodiments of the subject matter described herein can be implemented in a computing system that includes back-end components, e.g., as a data server, or includes middleware components such as an application server, or includes front-end components, e.g., a client computer having a graphical user interface, a web browser, or an 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 can 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.

[0155] 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, e.g., to display data to and receive user input from a user interacting with the device acting as a client. Data generated at the user device, e.g., a result of a user interaction, can be received at the server from the device.

[0156] In addition to the embodiments described above, the following embodiments are also innovative.

[0157] Embodiment 1 is acquiring one or more images comprising a plurality of pixels; for each image of the one or more images, determining a plurality of image patches of the image, each image patch comprising a different subset of pixels of the image; for each image of the one or more images, processing a corresponding plurality of image patches to generate an input sequence including a respective input element at each of a plurality of input locations, the plurality of input elements corresponding to a respective different image patch; processing an input sequence using a neural network to generate a network output that characterizes one or more images, the neural network including one or more self-attention neural network layers; The method includes:

[0158] In a second embodiment, the step of processing a plurality of image patches corresponding to the image to generate the input sequence comprises the steps of: generating respective one-dimensional initial input elements comprising pixels of the image patch; generating a respective input element using a respective initial input element; 2. The method of embodiment 1, comprising:

[0159] Embodiment 3 is the method of embodiment 2, wherein each image patch has dimensionality L×W×C, where C represents the number of channels in the image, and each initial input element has dimensionality 1×(L·W·C).

[0160] Example 4 is the method of example 2 or 3, wherein generating each input element using each initial input element includes processing the initial input element using a second neural network.

[0161] Example 5 is the method of Example 4, wherein the second neural network includes one or more fully connected neural network layers.

[0162] In a sixth embodiment, the step of processing a plurality of image patches corresponding to an image to generate an input sequence comprises: processing a plurality of image patches to generate respective intermediate input elements; combining the intermediate input element with a position embedding representing the position of the corresponding image patch in the image for each intermediate input element to generate a respective input element; 6. The method of any one of embodiments 1 to 5, comprising:

[0163] Embodiment 7 is the method of embodiment 6, wherein each position embedding is an integer.

[0164] Example 8 is the method of Example 6, in which each position embedding is machine-learned.

[0165] Example 9 is the method of any one of Examples 1 to 8, wherein certain input elements in the input sequence are machine-learned tensors.

[0166] Embodiment 10 is a method for processing an input sequence using a neural network to generate a network output that characterizes an image, the method comprising: processing an input sequence using a neural network to generate a respective output element for each input element in the input sequence; processing one or more of the output elements using a third neural network to generate a network output; 10. The method of any one of embodiments 1 to 9, comprising:

[0167] Embodiment 11 is a third neural network configured to generate a first type of network output; 11. The method of embodiment 10, wherein the neural network is trained simultaneously with a fourth neural network to generate a second type of network output that is different from the first type.

[0168] Example 12 is the method of Example 11, wherein a plurality of network parameters of the neural network are updated during training of the third neural network.

[0169] Embodiment 13 is the method of any one of embodiments 10 to 12, wherein the third neural network is a multilayer perceptron.

[0170] In a fourteenth embodiment, for each input sequence, A particular input element in the input sequence is a machine-learned tensor, and processing one or more output elements using a third neural network includes processing output elements corresponding to particular input elements using the third neural network to generate a prediction of the image. 14. The method of any one of embodiments 10 to 13.

[0171] Embodiment 15 is the method of any one of embodiments 1 to 14, wherein one or more of the self-attention neural network layers is a multi-head self-attention neural network layer.

[0172] Embodiment 16 is the method of any one of embodiments 1 to 15, wherein the neural network includes a sequence of one or more sub-networks, each sub-network configured to receive a respective sub-network input for each of a plurality of input positions and to generate a respective sub-network output for each of the plurality of input positions, and each sub-network includes a self-attention neural network layer and a feed-forward neural network layer for each position.

[0173] In a seventeenth embodiment, each subnetwork is a first normalization layer that applies a layer normalization to the sub-network input for each of a plurality of input locations; a first residual connection layer that combines the subnetwork inputs and the output of the self-attention neural network layer for each of a plurality of input positions; a second layer normalization layer that applies layer normalization to the output of the first residual connection layer, or A second residual connection layer that combines the output of the first residual connection layer with the output of a position-wise feedforward neural network layer 17. The method of embodiment 16, further comprising one or more of:

[0174] Embodiment 18 is the network output includes a classification output including a respective score corresponding to each of a plurality of categories, the category score indicating the likelihood that the image belongs to that category; the network output includes a pixel-level classification output including, for each pixel in the image, a respective score corresponding to each of a plurality of categories, the category score indicating the likelihood that the pixel belongs to that category; or the network output includes the coordinates of one or more bounding boxes surrounding each object shown in the image, or The neural network receives a plurality of images that are video frames of a video, and the network output includes an output that characterizes the video frames. 18. The method of any one of embodiments 1 to 17.

[0175] Example 19 is the method of example 18, wherein the output characterizing the video frames includes an output characterizing whether the video frames show a person performing a particular action.

[0176] Embodiment 20 is a system comprising one or more computers and one or more storage devices storing instructions operable, when executed by the one or more computers, to cause the one or more computers to perform any one of the methods of embodiments 1 to 19.

[0177] Embodiment 21 is a computer storage medium encoded with a computer program comprising instructions operable, when executed by a data processing device, to cause the data processing device to perform the method of any one of embodiments 1 to 19.

[0178] While the specification contains many specific implementation details, these should not be construed as limitations on the scope of any invention or the scope that may be claimed, but rather as descriptions of features that may be specific to particular embodiments of a particular invention. Some features that are described in this specification in the context of separate embodiments can also be implemented in combination in a single embodiment. Conversely, various features that are described in the context of a single embodiment can also be implemented in multiple embodiments separately or in any suitable subcombination. Furthermore, while features are described above as acting in some combinations and are initially claimed as such, in some cases, one or more features from a claimed combination can be deleted from the combination, and the claimed combination may be directed to a subcombination or a variation of a subcombination.

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

[0180] Specific embodiments of the subject matter have been described. Other implementations are within the scope of the following claims. For example, the actions recited in the claims can 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 shown or sequential order to achieve desirable results. In some cases, multitasking and parallel processing may be advantageous. [Explanation of symbols]

[0181] 100 Neural Network Systems 102 images 110 Image Patch Generation System 112 image patches 120 Image Patch Embedding System 122 Image Patch Embedding 124 Class Embedding 130 Neural Networks 140 Self-Attention Based Subnetworks 142 image patch output 144 class output 150 head subnetwork 152 Network Output 200 Self-Attention-Based Neural Networks 202 Input Sequence 204 Output Sequence 210 network blocks 212 Block Input Sequence 220 Self-Attention Neural Network Layer 230 Feedforward Neural Network Layers 310 images 320 images 330 images 340 images 350 images 360 images 400 Training System 412 Input Sequence 420-based neural network 422 base network output 430 Self-Attention Based Subnetworks 432 trained parameters 440 Base Head Subnetwork 450 Task Neural Network 452 Task Network Output 460 Self-Attention Based Subnetworks 470 Task Head Subnetwork 472 Initialized Parameters 480 Training Engine 482 Parameter Update 484 Parameter Update

Claims

1. acquiring one or more images comprising a plurality of pixels; for each image of the one or more images, determining a plurality of image patches of the image, each image patch comprising a different subset of the plurality of pixels of the image; for each image of the one or more images, processing a corresponding plurality of image patches to generate an input sequence including a respective input element at each of a plurality of input locations, a plurality of the input elements corresponding to respective different image patches; processing the input sequence using a neural network to generate a network output that characterizes the one or more images, the neural network including one or more self-attention neural network layers; A method comprising:

2. The step of processing the plurality of image patches corresponding to an image to generate an input sequence comprises, for each image patch: generating respective one-dimensional initial input elements comprising the plurality of pixels of the image patch; generating a respective input element using each of the one-dimensional initial input elements; 2. The method of claim 1, comprising:

3. Each image patch has dimensions L×W×C, C represents the number of channels of the image; The method of claim 2 , wherein each initial input element has a dimensionality of 1×(L·W·C).

4. 4. The method of claim 2, wherein generating each input element using the respective one-dimensional initial input elements comprises processing each one-dimensional initial input element using a second neural network.

5. The method of claim 4 , wherein the second neural network includes one or more fully connected neural network layers.

6. processing the plurality of image patches corresponding to images to generate an input sequence, processing the plurality of image patches to generate respective intermediate input elements; combining, for each intermediate input element, a position embedding representing the position of the corresponding image patch in the image to generate a respective input element; 6. The method of any one of claims 1 to 5, comprising:

7. The method of claim 6 , wherein each position embedding is an integer.

8. The method of claim 6 , wherein each positional embedding is machine-learned.

9. 9. The method of claim 1, wherein a particular input element in the input sequence is a machine-learned tensor.

10. processing an input sequence using the neural network to generate a network output that characterizes the image; processing the input sequence using the neural network to generate a respective output element for each input element in the input sequence; processing one or more of the output elements using a third neural network to generate the network output; 10. The method of any one of claims 1 to 9, comprising:

11. the third neural network is configured to generate a first type of network output; the neural network is trained simultaneously with a fourth neural network to generate a second type of network output different from the first type; The method of claim 10.

12. 12. The method of claim 11, wherein a plurality of network parameters of the neural network are updated during training of the third neural network.

13. 13. The method of claim 10, wherein the third neural network is a multi-layer perceptron.

14. For each input sequence, a particular input element in the input sequence is a machine-learned tensor; wherein processing one or more output elements using the third neural network includes processing the output element corresponding to the particular input element using the third neural network to generate a prediction of the image.

14. The method according to any one of claims 10 to 13.

15. 15. The method of claim 1, wherein one or more of the self-attention neural network layers is a multi-head self-attention neural network layer.

16. the neural network comprises a sequence of one or more sub-networks; Each subnetwork is receiving a respective sub-network input for each of the plurality of input locations; generating a respective sub-network output for each of said plurality of input locations; configured to:

16. The method of claim 1, wherein each sub-network comprises a self-attention neural network layer and a position-wise feedforward neural network layer.

17. Each subnetwork is a first normalization layer that applies layer normalization to the sub-network input for each of the plurality of input locations; a first residual connection layer that combines the sub-network inputs and the outputs of the self-attention neural network layer for each of the plurality of input positions; a second layer normalization layer that applies layer normalization to the output of the first residual connection layer; or a second residual connection layer that combines the output of the first residual connection layer with the output of the position-wise feedforward neural network layer; 17. The method of claim 16, further comprising one or more of:

18. the network output includes a classification output including a respective score corresponding to each of a plurality of categories, the category score indicating the likelihood that the image belongs to that category; or the network output comprises a pixel-level classification output including, for each pixel in the image, a respective score corresponding to each of a plurality of categories, the category score indicating the likelihood that the pixel belongs to that category; or the network output includes the coordinates of one or more bounding boxes surrounding each object shown in the image; or the neural network receives a plurality of images that are video frames of a video, and the network output includes an output that characterizes the video frames; 18. The method of any one of claims 1 to 17.

19. The method of claim 18 , wherein the output characterizing the video frame includes an output characterizing whether the video frame shows a person performing a particular action.

20. one or more computers; one or more storage devices storing instructions that, when executed by said one or more computers, cause said one or more computers to perform the method of any one of claims 1 to 19; A system comprising:

21. One or more computer storage media having stored thereon instructions that, when executed by one or more computers, cause the one or more computers to perform the method of any one of claims 1 to 19.

Images