Google patents image processing method using local attention blocks

1. A system for processing an image with parallel processing, comprising: one or more computers; and one or more storage devices storing instructions that when executed by the one or more computers cause the one or more computers to implement a computer vision neural network comprising a backbone neural network having one or more local self-attention layers and configured to process the image to generate a feature representation of the image and an output neural network configured to process the image to generate an output for a computer vision task from the feature representation of the image, wherein each local self-attention layer is configured to apply a local self-attention mechanism to a layer input representing the image and to generate a self-attention layer output, wherein applying the local self-attention mechanism to generate the self-attention layer output comprises: determining, by the local self-attention layer, a plurality of query blocks, wherein each query block is a respective non-overlapping subset of elements of the layer input that represents the image and comprises a respective plurality of neighboring elements of the layer input; determining, by the local self-attention layer, a plurality of context blocks, each context block corresponding to a respective query block, wherein each context block comprises elements in a corresponding query block and a plurality of elements of the layer input in a local window surrounding the corresponding query block within the layer input that represents the image, and wherein the local self-attention layer pads boundaries of the layer input such that the plurality of context blocks are of a same size, the size of each of the plurality of context blocks being determined by a halo value of the local self-attention layer; and generating, by the local self-attention layer, in parallel across the plurality of query blocks that are each a respective non-overlapping subset of the elements of the layer input, a respective block attention output for each query block, comprising: determining a respective query for each element in the query block based on applying a first learned linear transformation to the element, determining a respective key for each element of the corresponding context block based on applying a second learned linear transformation to each element of the corresponding context block, determining a respective value for each element of the corresponding context block based on applying a third learned linear transformation to each element of the corresponding context block, and using the determined query, keys, and values to generate a respective attention output for each element of the query block.