Local self-attention computer vision neural networks

One or more non-transitory computer-readable storage media storing instructions that when executed by 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 an 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.

A method performed by one or more computers, the method comprising:

• receiving an input image;
• and
• processing the input image using a computer vision neural network comprising a backbone neural network having one or more local self-attention layers and configured to process the input image to generate a feature representation of the input image and an output neural network configured to process the input image to generate an output for a computer vision task from the feature representation of the input image, wherein each local self-attention layer is configured to apply a local self-attention mechanism to a layer input representing the input 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 input 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 input 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.