Sequence packing for training image processing neural networks

Abstract

A method, system, and apparatus for training an image processing neural network, comprising a computer program encoded on a computer storage medium. For example, the method may include acquiring a batch of training examples, each training example comprising its respective training image; generating one or more input sequences, each input sequence being a sequence of tokens comprising tokens representing patches from two or more corresponding training images, each corresponding to two or more training images in the batch; processing each input sequence using an image processing neural network to generate its respective training output for each corresponding training image; and training the image processing neural network with the training outputs of the batch of training images using a first loss function for a first task.

Description

Technical Field

[0001] Cross - reference to Related Applications This application claims the benefit of priority of U.S. Provisional Patent Application No. 63 / 467,296, filed on May 17, 2023, the content of which is incorporated herein by reference.

Background Art

[0002] This specification relates to the training of neural networks.

[0003] A neural network is a machine - learning model that uses one or more layers of non - linear units to predict an output for a received input. Some neural networks include one or more hidden layers in addition to an output layer. The output of each hidden layer is used as 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 Problems

[0004] This specification describes a system implemented as a computer program on one or more computers to train an image - processing neural network at one or more locations.

[0005] An image - processing neural network is a neural network that receives an input image and processes the input image to generate an output of the image.

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

[0007] Vision transformers (ViTs) and other neural networks that process sequences of patches from images, such as Vision Transformer variants and MLP-Mixer neural networks, have demonstrated very strong performance in various computer vision tasks. However, these neural networks generally require that, during training, the input image be resized to a fixed resolution and fixed aspect ratio, such as a square aspect ratio, and then divided into a fixed number of patches before being processed by the neural network.

[0008] In contrast, this specification describes an alternative approach in which multiple patches from multiple different images are "packed" into a single input sequence during training. That is, a single input sequence during training may contain patches from multiple different input images, which may have different resolutions, different aspect ratios, or both.

[0009] Training neural networks in this way enables them to process images of variable resolution during inference, improving training efficiency and downstream inference performance, and allowing the incorporation of various other techniques that can further enhance training effectiveness. Examples of such techniques include randomly sampled token drop and resolution sampling.

[0010] In other words, by modifying the image processing neural network training scheme as described herein, the image processing neural network can be trained to improve performance in inference time while accurately processing images of different resolutions and aspect ratios. Furthermore, the neural network can achieve this performance improvement while reducing the computational resources consumed during training.

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

[0012] [Figure 1] An example of a neural network system is shown. [Figure 2] This is an exemplary process flowchart for performing training steps while training an image processing neural network. [Figure 3] This is an illustrative process flowchart for generating the input sequence when token drop is employed. [Figure 4] This is an exemplary process flowchart for obtaining a batch when resolution sampling is employed. [Figure 5] This example demonstrates how to generate an input sequence and process it using an image processing neural network. [Figure 6] Examples of the performance of the described technology are shown below. [Modes for carrying out the invention]

[0013] Similar reference numbers and names in various drawings refer to the same elements.

[0014] Figure 1 shows an example of a neural network system 100. The neural network system 100 is an example of a system implemented as a computer program on one or more computers in one or more locations, and may implement the systems, components, and technologies described below.

[0015] In detail, system 100 is a system that trains an image processing neural network 110 with training data 120.

[0016] The image processing neural network 110 is a neural network that receives an input image and processes the input image to generate an output of the image.

[0017] For example, the output of the image can be an image embedding of the input image in an embedding space. As used herein, an "embedding" is a numerical value having a predetermined dimensionality, for example, a vector of floating point values or other values. A space of possible vectors having a predetermined number of dimensions is referred to as an "embedding space". In other words, an image embedding is an encoding of an image.

[0018] As another example, the output of the image can be a classification output of a classification task of the input image. That is, the output includes respective scores for each of a set of object categories, and the score for an object category represents the likelihood that the image depicts an object belonging to that object category.

[0019] After being trained, the image processing neural network 110 can be adapted to one or more downstream tasks. In other words, the training that the system 100 performs on the training data 120 can be called a "pre-training" phase that is performed to improve whether the neural network 110 can successfully and easily adapt to one or more downstream tasks after being trained with the training data 120.

[0020] To adapt the image processing neural network 110, the system 100 can train a downstream neural network 130 that includes at least a portion of the layers of the image processing neural network 110 on training data for the downstream task.

[0021] For example, the downstream task can be an image classification task as described above.

[0022] As another example, the downstream task can be object detection, e.g., open vocabulary object detection, and the output for a given image identifies the positions of one or more bounding boxes in the image and, for each bounding box, the category to which the object depicted in each bounding box belongs. An example of an object detection task is an open vocabulary object detection task, and the set of possible categories may vary for different inputs and is specified by the embedding of category labels generated by a neural network of a text encoder. For example, in open vocabulary object detection, for each bounding box, the system can select the category label embedding that is most similar in terms of, e.g., cosine similarity or dot product, to the embedding of the bounding box generated by a downstream neural network.

[0023] As another example, the downstream task can be image segmentation. That is, the neural network can be configured to generate an element-level classification output (e.g., pixel-level classification output) that includes respective scores corresponding to each of a plurality of categories for each element of the network input. For a given element (e.g., a given pixel), the category score indicates the likelihood that the element belongs to the category. In some cases, the category may be an object class, and if the element is part of an object included in the object class corresponding to the category, it may belong to that category.

[0024] As another example, the downstream task can be image depth prediction. In a depth prediction task, the output generated by a neural network identifies, for each pixel of the image, the predicted depth of the scene at the pixel.

[0025] Another example is a downstream task, which could be a video comprehension task. That is, a neural network can be configured to process a sequence of video frames and produce an output that characterizes the video frames, for example, by characterizing whether a video frame depicts a person or other agent performing a particular action, by generating captions that describe the semantic content of the video, by classifying one or more objects in the video, and so on.

[0026] Other examples of downstream tasks include multimodal text and vision tasks, such as image captioning, where the input is an image and the output is a text caption describing the input image, or visual question answering, where the input is an image or video and a question about the image or video, and the output is a response to the question.

[0027] The image processing neural network 110 can have any suitable architecture that processes a sequence of tokens representing an image to produce an image output.

[0028] For example, the neural network 110 may have an architecture that includes multiple self-attention network blocks, each performing self-attention to update the tokens in the sequence. Examples of such architectures include the Visual Transformer (ViT) and other ViT variants. In a specific example, the neural network may have the architecture described in Dosovitskiy, et al, An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale, arXiv:2010.11929.

[0029] As another example, the neural network 110 may have an architecture that includes multiple network blocks that perform different types of operations to update the tokens in a sequence. An example of such an architecture is the MLP-Mixer architecture. As a specific example, the neural network may have the architecture described in Tolstikhin, et al, MLP-Mixer: An all-MLP Architecture for Vision, arXiv:2105.01601.

[0030] In any of the above examples, the neural network 110 may include a sequence of network blocks, each of which updates a token in the input sequence. The neural network 110 may also include one or more additional neural network layers that process the output of the last network block in the sequence to produce a training output.

[0031] When adapting the neural network 110 to a new downstream task, system 100 or other adaptive system may replace one or more additional neural network layers with a set of additional neural network layers specific to the downstream task. System 100 can then train the resulting downstream neural network on training data for the downstream task, either by fixing the network block and training only the new additional layers, or by training both the new additional layers and the network block.

[0032] More specifically, given an input image, the neural network 110 divides the input image into patches. Next, the neural network 110 generates an embedding for each patch and processes a sequence containing the embeddings for the patches to generate a training output for the input image.

[0033] The process of dividing an image into patches and generating patch embeddings is described in more detail below.

[0034] Generally, the training data 120 includes multiple different training examples, each of which includes at least a training image and optionally other data, such as classification labels or corresponding text sequences.

[0035] Training images in training examples generally have various resolutions and aspect ratios. However, conventional training schemes for training these neural networks generally require the training images to be resized to a fixed resolution and fixed aspect ratio, for example, a square aspect ratio, and then divided into a fixed number of patches before being processed by the neural network 110.

[0036] To improve upon these conventional training schemes, as will be described in more detail below, system 100 trains the neural network 110 by "packing" multiple images into an arbitrary given "packed" input sequence 140 that will be processed by the neural network 110 during training.

[0037] Training the neural network 110 in this way has many implications for the training of the neural network 110, and for its subsequent fine-tuning and use (at least partially) for performing inference.

[0038] For example, with a fixed computational budget, the described training scheme consistently outperforms conventional approaches. For instance, the described training scheme can yield a neural network 110 that performs comparably to a visual transducer (ViT) trained using a preceding state-of-the-art technique, but with a quarter of the computational effort. As one example, sequence packing allows for multiple images to be combined within the same input sequence, significantly increasing the number of training examples processed within the allocated computational budget, which contributes to improved performance. In another specific example, sequence packing combined with variable resolution input and variable token drop (described below) allows system 100 to process five times more images during training than conventional schemes. Furthermore, this efficiency improvement extends to the fine-tuning process, where a similar scheme can be applied.

[0039] Furthermore, by exposing the neural network 110 to multiple resolutions during both pre-training and fine-tuning, a single model demonstrates superior performance when evaluated at various resolutions, giving the neural network 110 a significant advantage in terms of inference cost compared to neural networks trained using conventional schemes.

[0040] Figure 2 is a flowchart of an exemplary process 200 for performing training steps while training an image processing neural network. For convenience, process 200 is described as being performed by one or more computer systems located at one or more locations. For example, a neural network system, e.g., the appropriately programmed neural network system 100 in Figure 1, can perform process 200.

[0041] Generally, the system can train an image processing neural network over multiple training steps. In each training step, the system takes a batch of training examples and trains the image processing neural network on the training examples in the batch, i.e., by performing 200 iterations of process 200.

[0042] To perform training, the system retrieves a batch of training examples, each containing its respective training image (step 202).

[0043] In some cases, the training examples may consist only of images. However, in other cases, the training examples may include other information, such as a corresponding text sequence describing the image, or data identifying the image's label.

[0044] Generally, a batch of training examples can include images with varying resolutions and aspect ratios.

[0045] In some embodiments, the system trains a neural network at the original resolution of the training images in a batch.

[0046] In some other embodiments, the system utilizes resolution sampling when generating batches of training images.

[0047] More specifically, traditional training schemes have a tension between improving throughput during training (training with smaller images) and improving performance (training with larger images to allow the use of higher resolution images during evaluation). As a result, models are often pre-trained at lower resolutions and fine-tuned at higher resolutions. However, the technique described becomes much more flexible by utilizing resolution sampling, because it enables mixed-resolution training by sampling from a distribution of image sizes while preserving the original aspect ratio of each image. This allows for both higher throughput and exposure to larger images, resulting in significantly improved performance (in terms of model size and training duration) compared to, for example, equivalent ViT.

[0048] Resolution sampling will be explained in more detail below, referring to Figure 4.

[0049] For each training image in the training example, the system divides the training image into multiple patches and generates a token to represent each patch (step 204).

[0050] To divide training images into patches, the system can divide each image into regions of a fixed size. Although the region size is fixed, different images can be divided into different numbers of patches because different images have different resolutions.

[0051] Generally, to generate tokens representing a patch, the system generates patch embeddings for the patch and location embeddings for the patch's position within the corresponding training image.

[0052] Next, the system combines the patch embedding and the location embedding of a patch to generate tokens that represent each patch. For example, the system can sum, average, or concatenate the patch embedding and location embedding.

[0053] To generate patch embeddings for any given patch, the system can use a patch embedding subnetwork to process the luminance values ​​of pixels within the patch to generate patch embeddings for the patch. For example, the patch embedding subnetwork may be a single linear projection layer, a multilayer subnetwork, such as an MLP, or have a different suitable architecture.

[0054] Generally, patch-embedded subnetworks are trained in conjunction with the rest of the image processing neural network, for example, using the gradient of a calculated loss function as described below.

[0055] The system can generate a location embedding for a given patch within the corresponding training image using one of several methods.

[0056] For example, the system could map each coordinate to a position on a flattened 1D representation of the image and use a one-dimensional (1D) positional embedding that is learned in collaboration with the rest of the neural network.

[0057] As another example, the system can use two-dimensional (2D) positional embeddings that are learned in conjunction with the rest of the neural network.

[0058] In other words, in this example, each embedding for each index of the [maxLen,maxLen] grid is learned, where the grid is indexed by the (x,y) coordinates of each patch, and maxLen is the pre-defined maximum length of any given dimension of the image. This allows for variable aspect ratios at resolutions up to R=P·maxLen, where P is the number of patches from any given image. However, all combinations of (x,y) coordinates must be checked during training. For example, the coordinates of a given patch for a given dimension can represent the index of a patch along a given dimension, with respect to the number of patches, for example, x=2 being the second patch along the x dimension.

[0059] However, while 2D embedding supports variable aspect ratios, none of the above encoding schemes readily extrapolate to an unknown resolution. To address this problem, in some embodiments, the system uses "factorized" position embeddings.

[0060] In this scheme, the system maintains a separate set of embeddings for each of several sets of x-coordinates, and a separate set of embeddings for each of several sets of y-coordinates.

[0061] In some cases, the embedding is an absolute embedding, so the embedding of a given coordinate is a function of the absolute index of the coordinate in the image.

[0062] In some cases, the embedding is a fractional embedding, and therefore the embedding of a given coordinate is a function of the ratio between the absolute index of the coordinate in the image and the length of the image edges along the corresponding dimension, i.e., the number of patches along the corresponding dimension of the given image.

[0063] When using factorized positional embeddings, in order to determine the positional embedding of a given patch, the system generates a first embedding of the x-coordinate of the patch's position in the corresponding training image, for example, by mapping the absolute index of the coordinate to a first embedding, or by mapping the ratio of the image's edge length along the x-dimension to the coordinate in the image to the first embedding.

[0064] The system generates a second embedding of the y-coordinate of the patch's position in the corresponding training image, for example, by mapping the absolute index of the coordinates to the second embedding, or by mapping the ratio of the image's edge length along the y-dimension to the coordinate in the image to the second embedding.

[0065] Next, the system combines the first embedding and the second embedding to generate a location embedding for the patch's location.

[0066] For example, the system may sum the first embedding and the second embedding to generate a location embedding for the patch location.

[0067] As another example, the system may stack, or concatenate, the first embedding and the second embedding to generate a location embedding for the patch's location.

[0068] As another example, the system may generate a positional embedding for the patch location by multiplying the first embedding and the second embedding element by element.

[0069] When using a factored scheme, the set of embeddings for the x-dimension and the set of embeddings for the y-dimension can be learned in conjunction with the rest of the neural network, for example, as standard pre-trained embeddings or trained Fourier position embeddings, or they can be determined before training, for example, as sinusoidal embeddings.

[0070] The system generates one or more input sequences (step 206). Generally, one or more input sequences collectively represent all images in a batch.

[0071] Each input sequence corresponds to two or more training images in a batch and is a sequence of tokens containing tokens that represent patches from two or more corresponding training images.

[0072] In other words, the system generates an input sequence using sequence packing by "packing" tokens from multiple images into the same input sequence.

[0073] In general, since images in a batch may have variable resolution and aspect ratio, a single sequence may contain tokens from two different images that have two different resolutions, two different aspect ratios, or both.

[0074] In detail, as mentioned above, the number of tokens representing an image can depend on the image resolution; therefore, images with different resolutions will be represented by different numbers of tokens. Consequently, any given input sequence can contain different numbers of tokens for two different images with two different resolutions.

[0075] In some embodiments, the system uses a token drop scheme when generating each input sequence. In the part that employs token drop, the system can determine for each image whether to remove any of the tokens representing the image from the input sequence, and then generate the input sequence using only the tokens that have not been removed from each of two or more input images.

[0076] This will be explained in more detail below, with reference to Figure 3.

[0077] The system can determine, in one of several ways, which images in a given batch are represented by which input sequence.

[0078] For example, in some embodiments, each input sequence must contain the same maximum number of tokens, that is, to optimize the performance of the hardware on which the neural network is being trained, such as a tensor processing unit (TPU) or other ASIC optimized for a fixed-size input.

[0079] In these embodiments, the system can generate input sequences using a greedy approach in which the system adds images to a first sequence that has enough remaining space. That is, the system can traverse a batch of images according to a random order of images and determine whether the tokens representing the images can "fit" into any of the already generated, partially complete input sequences. If so, the system adds the images to one of the partially complete input sequences. Otherwise, the system places the images in a new sequence. When no more images can fit into any of the generated sequences, the sequences are filled with padding ("masking") tokens to fix the length of the sequences.

[0080] As another example, when generating a given input sequence rather than using masking tokens, the system can dynamically select the resolution of the final image or the token drop rate in the sequence to precisely fit the remaining tokens in the given input sequence.

[0081] Next, the system uses an image processing neural network to process each input sequence and generate the respective training output for each corresponding training image (step 208).

[0082] For example, if an image processing neural network includes a sequence of self-attention layers, the system processes the input sequence through the sequence of self-attention layers to update the tokens in the input sequence, and then generates a training output for each training image from the updated tokens.

[0083] This example will be explained below with reference to Figure 5.

[0084] The system uses the training output of a batch of training images to train an image processing neural network with the task's loss function (step 210). For example, the loss function may be a control loss function or a classification loss function, such as a loss function with a cross-entropy term.

[0085] More specifically, in some cases, the system trains a neural network through comparative learning. In those cases, the training output generated by the image processing neural network for each image is the image embedding of the image in the embedding space.

[0086] Furthermore, the image processing neural network is trained in conjunction with a text encoder neural network having text encoder neural network parameters and configured to process text segments and generate text embeddings of the text segments in the embedding space. For example, the text encoder neural network may be a transformer neural network, a recurrent neural network (RNN), or any other suitable type of text processing neural network.

[0087] Therefore, in these cases, each training example also includes its respective training text segment, for example, a text segment determined to be semantically similar to the training image of the training example. For example, in a given training example, the text segment could be an image text annotation from a set of manually or automatically generated image annotations, or it could be alt text associated with an image from a set of alt text data. Alt text is text that appears on a web page in place of an image, for example, when the image cannot be rendered properly or loaded. For example, the system can obtain alt text data from data maintained by internet search engines or other software that automatically crawls web pages on the internet.

[0088] In these cases, to train the image processing neural network and the text encoder neural network, the system can process each training text segment using the text encoder neural network to generate training text embeddings for the training text segment.

[0089] Subsequently, the system can train image processing neural networks and text encoder neural networks on a contrast loss function based on the similarity between training text embeddings and training image embeddings.

[0090] In some other cases, the system trains an image processing neural network on an image classification task. In those cases, the training output produced by the image processing neural network can be a classification output, and each training example includes a ground truth classification of the training image of the training example. Thus, for each training example, the system trains the image processing neural network with a loss function that includes a classification loss for the ground truth classification of the training image, e.g., cross-entropy or another classification loss.

[0091] As described above, after repeatedly performing process 200 iterations to train the image processing neural network, the system can then adapt the image processing neural network to perform downstream computer vision tasks. Examples of such tasks include image classification, object detection, image segmentation, depth prediction, video understanding, and multimodal text and vision tasks.

[0092] Figure 3 is a flowchart of an exemplary process 300 for generating an input sequence when the system performs a token drop. For convenience, process 300 is described as being performed by a system of one or more computers located in one or more locations. For example, a neural network system, e.g., the appropriately programmed neural network system 100 in Figure 1, can perform process 300.

[0093] Token dropping (i.e., random omission of input patches during training) was developed to speed up training. However, typically the same proportion of tokens are dropped from all examples. By utilizing sequence packing, the system can enable sequential token dropping, thereby varying the token drop rate from image to image. This allows for the faster throughput benefits enabled by token dropping, while still allowing the model to process some complete images, reducing training / inference mismatch and improving inference performance. Furthermore, packing allows for variation in the distribution of drops throughout training to further improve training performance.

[0094] The system selects two or more training images from a batch of training examples (step 302). For example, the system can greedily select training images represented by the training sequence, as described above.

[0095] The system determines the token drop rate for each of two or more training images (step 304).

[0096] In some cases, the token drop rate for each image is the same across all training images in a batch. In these cases, the system can determine the token drop rate for a batch of training images using one of several methods.

[0097] For example, the system can randomly sample each token drop rate in a batch from a set of possible token drop rates. For instance, the system can sample token drop rates from a beta distribution or other distributions across a range of possible token drop rates.

[0098] As another example, the system can determine the token drop rate for each batch according to the index of the training step performed during training. In other words, the system can maintain the ability to map the index of the training step to the token drop rate.

[0099] In some other cases, different images in a batch may have different token drop rates.

[0100] For example, the system can determine the token drop rate for each given training image in a batch based on the index of the training image among the training images processed during training. In other words, the system can maintain the ability to map the index of a training image to a token drop rate.

[0101] As another example, the system can determine the token drop rate based on the resolution of the training images. For example, the system can sample the token drop rate of a given image from a distribution based on the resolution of the training images, for example, from a distribution having a mean value based on the resolution of the training images.

[0102] For each of the two or more training images, the system determines whether to select one of the tokens representing the patch from the two or more training images for removal, according to the token drop rate of the training images (step 306).

[0103] The system generates an input sequence containing only the tokens of two or more images that were not selected for removal (step 308).

[0104] Figure 4 is a flowchart of an exemplary process 400 for obtaining a batch of training examples when the system employs resolution sampling. For convenience, process 400 is described as being performed by one or more computer systems located in one or more locations. For example, a neural network system, e.g., a appropriately programmed neural network system 100 in Figure 1, can perform process 400.

[0105] The system samples the original batch of original training examples (step 402). As described above, each original training example contains the original training image, and each original training image has its own original resolution and its own original aspect ratio. For example, the system can sample an original batch from a larger set of training examples available to the system for use in training a neural network.

[0106] The system corrects the original image in the original training example (step 404).

[0107] More specifically, for each original training image, the system samples the resolution of the original training image, and then modifies the original training image to have the sampled resolution while maintaining the original aspect ratio of the original training image. In other words, the aspect ratio is preserved while the total number of pixels in the original image matches the sampled resolution.

[0108] For example, a system can sample resolutions from a fixed set of resolutions by, for instance, sampling resolutions from a uniform distribution across a fixed set of resolutions.

[0109] The system generates a batch of training examples containing the corrected images (step 406). That is, a training example corresponding to any given original training example contains the corrected training image generated from the original training image in the training example, and any other information from the original training example, such as the training text if the neural network was trained through controlled learning, or the ground truth label if the neural network was trained to perform classification.

[0110] Figure 5 shows an example 500 in which an input sequence is generated and processed between a given training step during the training steps.

[0111] More specifically, in Example 500, the image processing neural network is a neural network that includes one or more self-attention layers, each of which applies self-attention.

[0112] As shown in Example 500, the system selects three input images 502 (image 1, image 2, and image 3).

[0113] Next, the system “patches” each of the input images 502 504, dividing each of the images 502 into patches.

[0114] After dividing the image into patches, the system can perform a token drop 506 on each input image 502 to determine whether any of the tokens representing any of the patches of the image should be removed.

[0115] Next, the system generates an input sequence 510 ("pack sequence") that contains only the tokens that were not removed from the three input images. More specifically, input sequence 510 contains a subsequence for each of the three input images. More precisely, input sequence 510 contains a concatenation of multiple subsequences, each containing a token representing a patch from each of the three training images.

[0116] In Example 500, since the concatenation of multiple subsequences contains fewer than the maximum number of tokens, the system adds two masking tokens to the "padded" portion of input sequence 510 so that input sequence 510 has the maximum number of tokens.

[0117] Next, the system processes the input sequence 510 using a sequence of self-attention blocks 520. Each self-attention network block 520 updates at least one subset of the tokens in the input sequence by applying self-attention to the tokens in the input sequence 510. For example, masking tokens in the "pad" section of the input sequence 510 can be masked out so that they are not updated by the self-attention layer.

[0118] To ensure that tokens for a given image do not affect the final output for other images, for each self-attention network block 520, self-attention is masked such that the tokens for each corresponding image are updated only with the token for the corresponding image and not with tokens for any other corresponding image or padding.

[0119] In Example 500, as the neural network processes the sequence through block 520, it applies a pooling operation 530, such as global average pooling or other types of pooling, to the tokens to generate a pooled representation for each of the three images corresponding to the input sequence. Once the pooled representations are generated, the neural network can use the pooled representations of the images to generate a training output for each image.

[0120] For example, in a controlled learning task, the system can use pooling representations as embeddings for input images, or it can apply one or more transformations, such as learned projections or other learned behaviors, to the pooling representations to generate the embeddings.

[0121] As another example, in a classification task, the system can use an output subnetwork to process pooling representations and generate a classification output. For example, the output subnetwork may include one or more fully connected layers followed by a softmax layer.

[0122] To ensure that the tokens of a given image do not affect the final output of other images, pooling operation 530 is masked so that a pooled representation of each corresponding image is generated using only the tokens of the corresponding image and not the tokens of other images. That is, only the tokens of the corresponding image are pooled, and the tokens of other corresponding images or padding are not pooled when generating the pooled representation.

[0123] Figure 6 shows example 600 of the performance of the described techniques. Specifically, example 600 shows the performance of various architectures trained using the described techniques compared to the same architectures trained using the ViT training technique.

[0124] As can be seen from Figure 6, the described technique is superior to the ViT model across various dimensions, including model accuracy given pre-training time, fine-tuned accuracy at a given pre-training cost, and fine-tuned accuracy at a given inference cost.

[0125] This specification uses the term “configured” in relation to systems and computer program components. When one or more computer systems are configured to perform a particular operation or action, it means that, while in operation, software, firmware, hardware, or a combination thereof is installed on the system that causes the system to perform that operation or action. When one or more computer programs are configured to perform a particular operation or action, it means that one or more programs, when executed by a data processing device, contain instructions that cause the device to perform that operation or action.

[0126] The subject matter and functional embodiments described herein can be implemented in digital electronic circuits, tangibly embodied computer software or firmware, computer hardware, and include structures disclosed herein and their structural equivalents, or one or more combinations thereof. Embodiments of the subject matter described herein may be implemented as one or more modules of computer program instructions, i.e., as one or more computer programs encoded in a tangible non-temporary storage medium for execution by a data processing device or for controlling the operation of a data processing device. The computer storage medium may be a machine-readable storage device, a machine-readable memory board, a random-access memory device or a serial-access memory device, or one or more combinations thereof. Alternatively or additionally, program instructions may be encoded into artificially generated propagated signals, such as mechanically generated electrical signals, optical signals or electromagnetic signals, which are generated to encode information for transmission to a receiving device suitable for execution by a data processing device.

[0127] The term "data processing device" refers to data processing hardware and encompasses all types of devices, machines, and equipment for processing data, including, for example, programmable processors, computers, or multiple processors or computers. A device may also be, or further include, specialized logic circuits such as FPGAs (Field-Programmable Gate Arrays) or ASICs (Application-Specific Integrated Circuits). Optionally, in addition to hardware, a device may include code that constitutes the execution environment for computer programs, such as processor firmware, protocol stacks, database management systems, operating systems, or one or more combinations thereof.

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

[0129] In this specification, the term “database” is used broadly to refer to any collection of data. The data does not need to be structured in any particular way, or not structured at all, and can be stored on one or more storage devices. Therefore, for example, an index database may contain multiple collections of data, each of which may be organized and accessed in a different way.

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

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

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

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

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

[0135] Data processing equipment for implementing machine learning models may include, for example, dedicated hardware accelerator units for handling the general and computationally intensive parts of machine learning training or production, i.e., inference workloads.

[0136] Machine learning models can be implemented and deployed using machine learning frameworks, such as the TensorFlow framework or the Jax framework.

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

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

[0139] While this specification includes details of many specific embodiments, these should not be construed as limiting the scope of any invention or claimable content, but rather as descriptions of features that may be specific to a particular embodiment of a particular invention. Certain features described herein in the context of other embodiments may also be realized in combination in a single embodiment. Conversely, various features described in the context of a single embodiment may also be realized individually or in any suitable subcombination in multiple embodiments. Furthermore, features may be described above as functioning in a particular combination, and even if initially claimed as such, one or more features from the claimed combination may be removed from the combination, and the claimed combination may cover a subcombination or a variation of a subcombination.

[0140] Similarly, while operations are shown in the drawings and described in a specific order in the claims, this should not be understood as requiring that such operations be performed in a specific or sequential order shown, or that all shown operations be performed, in order to obtain the desired results. In certain situations, multitasking and parallel processing may be advantageous. Furthermore, the separation of various system modules and components in the embodiments described above should not be understood as requiring such separation in all embodiments, and it should be understood that the program components and systems described can generally be integrated into a single software product or packaged into multiple software products.

[0141] Specific embodiments of the subject matter have been described. Other embodiments are within the scope of the following claims. For example, the actions described in the claims may be performed in a different order, and this may still yield the desired results. As an example, the process shown in the accompanying diagram does not necessarily require to be performed in the specific order or sequence shown to achieve the desired results. In some cases, multitasking and parallel processing may be advantageous.

Claims

1. A method for training an image processing neural network having image processing neural network parameters to perform a first task, which is performed by one or more computers, wherein in each of a plurality of training steps, the method is This involves obtaining a batch of training examples, where each training example includes its respective training image. For each training image, Dividing the aforementioned training image into multiple patches, To generate a token representing each of the aforementioned patches, To generate one or more input sequences, each input sequence being a sequence of tokens containing tokens that correspond to two or more of the training images in the batch and represent patches from the two or more corresponding training images, To generate a training output for each corresponding training image, the image processing neural network is used to process each input sequence, and The image processing neural network is trained using the training output of the training images from the batch with the first loss function of the first task. Methods that include...

2. The training output generated by the image processing neural network for each image is an image embedding of the image in the embedding space. The image processing neural network is trained in conjunction with a text encoder neural network having text encoder neural network parameters and configured to process the text segments to generate text embeddings of the text segments in the embedding space. Each training example further includes its respective training text segment, and the method is as follows: For each training text segment, Processing the training text segment using the text encoder neural network to generate training text embeddings for the training text segment. It further includes, The method according to claim 1, wherein the first loss function is a control loss function based on the similarity between the training text embedding and the training image embedding.

3. The method according to claim 1, wherein the training output generated by the image processing neural network for each image is an image classification output of the image, and the loss function is the classification loss for ground truth classification of the image.

4. After training the image processing neural network with the first loss function for the first task, The method according to any one of claims 1 to 3, further comprising adapting the image processing neural network to perform a downstream computer vision task.

5. The aforementioned downstream computer vision task is Image classification, Object detection, Image segmentation, Depth prediction, Understanding the video, or The method according to claim 4, wherein one of the tasks is a multimodal text and a vision task.

6. The method according to claim 5, wherein the multimodal text and vision task is an image captioning or visual question answering task.

7. The method according to claim 5, wherein the computer vision task is object detection in an open vocabulary.

8. To generate one or more input sequences, each input sequence being a sequence of tokens containing tokens that correspond to two or more of the training images in the batch and represent patches from the two or more corresponding training images, Select two or more of the aforementioned training images. For each of the two or more training images mentioned above, determine the respective token drop rates. For each of the two or more training images, it is determined whether to select one of the tokens representing a patch from the two or more images for removal, according to the token drop rate of the training image, and To generate an input sequence that includes only the tokens for the two or more training images that were not selected for removal. The method according to any one of claims 1 to 7, including the method described in any one of claims 1 to 7.

9. The method according to claim 8, wherein each of the aforementioned token drop rates is the same for all of the training images in the batch.

10. The method according to claim 9, wherein determining the respective token drop rate for each of the two or more training images comprises randomly sampling the respective token drop rate for the batch from a set of possible token drop rates.

11. The method according to claim 9, wherein determining the respective token drop rate for each of the two or more training images includes determining the respective token drop rate for the batch according to the index of the training step during the training step performed during the training.

12. Generating the respective tokens representing each of the aforementioned patches is: To generate the patch embedding of the aforementioned patch, To generate a positional embedding of the patch location within the corresponding training image, and The patch embedding and position embedding of the patch are combined to generate the respective tokens representing the patch. The method according to any one of claims 1 to 11, including the method described in any one of claims 1 to 11.

13. Generating the patch embedding of the aforementioned patch is The method according to claim 12, comprising processing the luminance values ​​of the pixels of the patch using a patch embedding subnetwork to generate each of the patch embeddings.

14. Training the aforementioned image processing neural network is The method according to claim 13, comprising training the patch-embedded subnetwork using the gradient of the first loss function.

15. Generating a positional embedding of the patch location within the corresponding training image is: To generate a first embedding of the x-coordinate of the position of the patch within the corresponding training image, To generate a second embedding of the y-coordinate of the position of the patch within the corresponding training image, and To generate the positional embedding of the aforementioned location of the patch, the first embedding and the second embedding are combined, The method according to any one of claims 12 to 14, including the method described in any one of claims 12 to 14.

16. The method according to any one of claims 1 to 15, wherein each input sequence includes a concatenation of a plurality of subsequences, each containing a token representing a patch from one of the two or more corresponding training images.

17. Generating each input sequence is Determining whether the concatenation of the plurality of subsequences contains fewer tokens than the maximum number, and In response to the determination that the concatenation of the plurality of subsequences contains fewer than the maximum number of tokens, one or more masking tokens are added to the input sequence so that the input sequence has the maximum number of tokens. The method according to claim 16, including the method described in claim 16.

18. The method according to any one of claims 1 to 17, wherein the image processing neural network has a plurality of self-attention network blocks, each applying self-attention to the tokens of each input sequence in order to update at least one subset of the tokens of the input sequence, and for each of the self-attention network blocks, the self-attention is masked so that for each input sequence, the tokens of each corresponding image are updated only using the tokens of the corresponding image, and the tokens of any other corresponding image are not updated.

19. The method according to any one of claims 1 to 18, wherein for at least one of the input sequences, the two or more corresponding images include at least two images having different resolutions.

20. The method according to any one of claims 1 to 19, wherein for at least one of the input sequences, the two or more corresponding images include at least two images having different aspect ratios.

21. Obtaining a training example batch is Sampling the original batch of the original training examples, wherein each original training example includes the original training image having its original resolution and its original aspect ratio. For each original training image, Sampling the resolution of the original training image, Modifying the original training image to have the sampled resolution while maintaining the original aspect ratio of the original training image. The method according to any one of claims 1 to 20, including the method described in any one of claims 1 to 20.

22. The method according to claim 21, wherein two or more of the original training images have different original aspect ratios.

23. A system comprising one or more computers and one or more storage devices storing instructions, wherein when an instruction is executed by the one or more computers, the one or more computers cause the one or more computers to perform the operation of each of the methods described in any one of claims 1 to 22.

24. One or more computer storage media storing instructions, wherein, when executed by one or more computers, the instructions cause one or more computers to perform the operations of each of the methods described in any one of claims 1 to 22.

Images