Neural networks with area attention

Methods, systems, and apparatus, including computer programs encoded on computer storage media, for implementing an area attention layer in a neural network system. The area attention layer area implements a way for a neural network model to attend to areas in the memory, where each area contains a group of items that are structurally adjacent.

BACKGROUND

This specification relates to processing inputs using neural networks.

Neural networks are machine learning models that employ one or more layers of nonlinear 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 another layer in the network, i.e., a hidden layer or an output layer. Each layer of the network generates an output from a received input in accordance with current values of a respective set of parameters.

Attentional mechanisms have been used to boost the accuracy on a variety of deep learning tasks. They allow a model to focus selectively on specific pieces of information; a piece of information can be, for example, a word in a sentence for neural machine translation or a region of pixels in image captioning.

An attentional mechanism typically follows a memory-query paradigm, where the memory M contains a collection of items of information from a source modality, e.g., the embeddings of an image or the hidden states of encoding an input sentence, and the query q comes from a target modality, e.g., the hidden state of a decoder model. In architectures such as Transformer, self-attention involves queries and memory from the same modality for either encoder or decoder. Each item in the memory has a key-value pair, (ki, vi), where the key is used to compute the probability airegarding how well the query matches the item, as expressed in equation 1.

Typical choices for fattinclude dot products qkiand a multilayer perceptron. The output OqMfrom querying the memory M with q is then calculated as the sum of all the values in the memory weighted by their probabilities, as expressed in equation 2, which can be fed to other parts of the model for further calculation.

During training, the model learns to attend to specific pieces of information given a query. For example, it can associate a word in the target sentence with a word in the source sentence for translation tasks.

Attention mechanisms are typically designed to focus on a predetermined granularity of individual items in the entire memory, where each item defines the granularity of what the model can attend to. For example, it can be a character for a character-level translation model, a word for a word-level model, a grid cell for an image-based model, or a hidden state in a latent space.

SUMMARY

This specification describes a system implemented as computer programs on one or more computers in one or more locations that generates a machine learning model output from a machine learning model input.

In particular, this specification describes a system that uses a self-attention neural network to generate model output from a model input, with at least one of the self-attention layers in the neural network employing an area attention mechanism.

Generally, at least one of the input and output is a sequence. For example, the output can be a target sequence that includes a respective output at each of multiple positions in an output order, the input can be an input sequence that includes a respective input at each of multiple positions in an input order, or both the input and output can be a sequence, i.e., the system transduces an input sequence into an target sequence.

For example, the system may be a neural machine translation system. That is, if the input sequence is a sequence of words or characters in an original language, e.g., a sentence or phrase, the target sequence may be a translation of the input sequence into a target language, i.e., a sequence of words or characters in the target language that represents the sequence of words in the original language.

As another example, the system may be a speech recognition system. That is, if the input sequence is a sequence of audio data representing a spoken utterance, the target sequence may be a sequence of graphemes, characters, or words that represents the utterance, i.e., is a transcription of the input sequence.

As another example, the system may be a natural language processing system. For example, if the input sequence is a sequence of words in an original language, e.g., a sentence or phrase, the target sequence may be a summary of the input sequence in the original language, i.e., a sequence that has fewer words than the input sequence but that retains the essential meaning of the input sequence. As another example, if the input sequence is a sequence of words that form a question, the target sequence can be a sequence of words that form an answer to the question. As yet another example, if the input sequence is a sequence of words, the target sequence may be a sequence that defines a parse tree, e.g., a dependency parse or a constituency parse, of the input sequence. As yet another example, if the input sequence is a sequence of words, the output of the neural network can be a natural language understanding output, e.g., an output for an entailment task, a paraphrase task, a textual similarity task, a sentiment task, a grammaticality task, and so on.

As another example, the system may be part of a computer-assisted medical diagnosis system. For example, the input sequence can be a sequence of data from an electronic medical record and the output can be a sequence of predicted treatments or a single predicted treatment to be performed after the last input in the input sequence.

As another example, the system may be an image generation system that generates images conditioned on a particular type of input, e.g., a smaller image, an object category, or a natural language text sequence. In these examples, the system may receive the input and then generate the output image as a sequence of color values, i.e., of color channel values for the pixels of the output image, or as a two-dimensional structure of color values.

As another example, the system may be part of an image captioning system. That is, the system can receive an image, e.g., as a sequence of the color values of the image or a 3D tensor, or an embedding of an image, and then generate a sequence of text that captions the image, i.e., that is a natural language description of the image.

As another example, the system may be a system that receives as input one or more video frames. For example, the system can receive as input a sequence of video frames and then predict the next frame, e.g. as a sequence of color values or as a two-dimensional structure of color values. As another example, the system can receive as input a sequence of video frames from a video and then generate a natural language text sequence that describes the video.

As another example, the system may be a system that processes sequences of biological data, e.g., genome sequencing reads.

In particular, the neural network includes one or more area attention layers.

Each area attention layer is configured to, during the processing of each neural network input by the neural network, receive data specifying a memory including a plurality of items and, for each item, a respective key and a respective value.

The area attention layer then determines a plurality of areas within the memory, wherein (i) each area includes one or more items in the memory, and (ii) one or more of the areas include multiple adjacent items in the memory. The area attention layer determines, for each of the plurality of areas, a respective area key and a respective area value from at least the keys and values of the items in the area.

The area attention layer then receives an attention query and applies an attention mechanism between the attention query and the area keys for each area to generate a respective attention weight for each area.

The area attention layer then generates an area attention layer output by combining the area values for each area in accordance with the attention weights.

An area attention layer can be used in a neural network in place of some or all of the conventional attention layers in the neural network.

For example, some neural networks include an encoder neural network, a decoder neural network, and an attention layer. The encoder neural network generates a respective encoded representation for each of multiple portions of the neural network input and provides the encoded representation to the attention layer. For example, when the network input is a sequence, the encoder neural network can generate a respective encoded representation of each element in the sequence. When the network input is an image, the encoder neural network can generate a respective encoded representation of multiple spatial regions within the image or can receive an embedding of the image and generate the respective encoded representation for each of the spatial regions from the embedding. The decoder neural network generates the neural network output by querying the attention layer, i.e., by providing hidden states of certain components to the decoder neural network layer and using received attention layer outputs to generate the network output.

In this example, the area attention neural network layer can be used in place of the attention layer, with the items in the memory corresponding to the multiple portions, the keys and values both being the encoded representations, and the attention query being the query from the decoder.

As another example, some neural networks employ an attention-based architecture that are made up of multiple layers that employ multi-head attention. An example of such a neural network architecture is the Transformer architecture described in Vaswani, et al., Attention is all you need, http://arxiv.org/abs/1706.03762.

In these cases the area attention layer can be employed in place of some or all of the attention heads in the multi-head attention layers. When the area attention layer is an attention head in a self-attention layer, the queries, keys, and values are all derived from the same input, i.e., by the multi-head attention layer before being provided to the area attention layer. When the area attention layer is an attention head in an encoder-decoder attention layer, the queries are derived from the decoder input to the encoder-decoder attention layer, and the keys and values are derived from the encoded representations of the neural network input.

Conventional attention layers within a neural network generate attention weights for only individual items within a memory. The described area attention layer, on the other hand, also generates attention weights for areas of items, i.e., groups of multiple items within a memory. This allows the area attention layer to capture rich alignment distributions by focusing on multiple items when beneficial for processing the current input instead of being limited to focusing on a single item.

By using an area attention layer within a neural network to replace some or all of the convention attention layers in the neural network, the described systems can achieve high quality and even state-of-the-art results on a variety of tasks, e.g., machine translation and image captioning

With an encoder and decoder that are area attention-based, a sequence transduction neural network can transduce sequences more accurately than existing networks that are based on convolutional layers or recurrent layers. Additionally, the use of an attention-based architecture allows the sequence transduction neural network to transduce sequences quicker, to be trained faster, or both, e.g., because the operation of the network can be more easily parallelized.

DETAILED DESCRIPTION

This specification describes an attention mechanism that will be referred to as area attention, which is a general mechanism for a model to attend to a group of items in the memory that are structurally adjacent. In general, in this context, a memory contains a collection of individual items with a predefined, fixed granularity, e.g., a word token or an image grid. In area attention, each unit for attention calculation is an area that can contain or represent one or more than one item of the original memory. Because each of the areas can aggregate a varying number of items, in area attention, the granularity of attention is learned from training data rather than being predetermined. Area attention can be applied to both single or multi-head attention mechanisms. In multi-head attention, with each head using area attention, a model can attend to multiple areas in the memory.

Area attention can be applied by any or all of the attention layers in neural network. That is, the area attention mechanism can replace a conventional attention mechanism employed by any or all of the attention layers in a neural network.

FIG.1illustrates example areas of attention in relation to a one-dimensional memory, i.e., a memory that contains a one-dimensional structure of items. An area is a group of structurally adjacent items in the memory, e.g., spatially for a two-dimensional memory arrangement, e.g., for images, or temporally for a one-dimensional memory arrangement, e.g., for natural language sentences. When the memory consists of a sequence of items, i.e., has a one-dimensional structure, an area is a range of items that are sequentially or temporally adjacent; and the number of items in the area can be one or more than one. Many language-related tasks, e.g., machine translation or sequence prediction tasks, fall in the one-dimensional case.FIG.1shows an example in which the original memory100is a 4-item sequence of items11,12,13,14. By combining the adjacent items in the sequence, area memory110is formed where each item21-29in the area memory is a combination of multiple adjacent items in the original memory. Items21-24are combinations of one item; items25-27are combinations of two; and items28-29are combinations of three. The maximum area size a system will consider for a task can be limited. In the example ofFIG.1, the maximum area size is three items.

FIG.2illustrates example areas of attention in relation to a two-dimensional memory200, i.e., a memory that contains a two-dimensional structure of items31-39. When the memory contains a grid of items, a two-dimensional structure, an area can be any rectangular region in the grid. The maximum area size the system will consider for a task can be limited in the two-dimensional case as well. For a two-dimensional area, this can be done by setting the maximum height and width for each area. In the example ofFIG.2, the original memory is a 3×3 grid of items and the maximum height and width allowed for each area is two items. The area memory210thus contains areas of dimension 1×1 with area items41-49; areas of dimension 1×2 with area items51-56, each of dimension 1×2; areas of dimension 2×1 with area items61-66, each of dimension 2×1; and areas of dimensions 2×2 with area items71-74, each of dimension 2×2.

As particular examples, area item47is a 1×1 area and therefore includes only item37, while area item56is a 1×2 area and includes items33and36, area item66is a 2×1 area and includes items38and39, and area item71is a 2×2 area and includes items32,33,36, and36.

Because the output of area attention includes attended outputs over information with varying granularity, the remaining components of the neural network can learn which levels of granularity are important for a particular task during training of the neural network system.

For a model be able to attend to each area as a unit of attention, a key and value are defined for each area that contains or represents one or multiple items in the original memory. The system can then receive an attention query and apply an attention mechanism using the area keys and values.

FIG.3illustrates a system300of one or more computers310implementing a neural network system320with an area attention layer330. The neural network system is configured to receive a neural network input322and to generate a neural network output324.

The area attention layer is configured to perform operations to implement area attention during the processing of the neural network input. The attention layer is configured to receive data specifying a memory332that contains multiple items and, for each item, a respective key and a respective value, and to determine multiple areas within the memory324. Each area includes one or more items in the memory, and one or more of the areas include multiple adjacent items in the memory. The attention layer is further configured to determine, for each of the areas, a respective area key and a respective area value326from at least the keys and values of the items in the area; receive an attention query328; apply an attention mechanism between the attention query and the area keys for each area to generate a respective attention weight for each area330; and generate an area attention layer output332by combining the area values for each area in accordance with the attention weights. The area attention layer is further configured to provide the area attention layer output334to another component of the neural network system.

In some cases, the memory is arranged as a sequence of items, and determining multiple areas includes identifying, as different areas, each combination of adjacent items that includes no more than a maximum number of items. In other cases, the memory is arranged as a two-dimensional grid of items, and determining multiple areas includes identifying, as different areas, each rectangular region of items within the two-dimensional grid that has no more than a maximum height and no more than a maximum width. The maximums can be predetermined as hyperparameter values or determined using a hyperparameter search technique.

The area attention layer is configured to determine, for each of the areas, a respective area key and a respective area value. In some implementations, the layer does so by determining the area value to be a sum of the values of the items in the area or a mean of the keys of the items in the area, or by determining a plurality of features of the items in the area and combining the features to generate the area key of the area. This combining of features can be done by summing or concatenating the features to generate a combined feature and applying one or more learned non-linear transformations to the combined feature to generate the area key. The features that can be combined include one or more of an embedding corresponding to a number of items in the area, a mean of the keys of the items in the area, or a variance of the keys of the items in the area.

For areas that include only respective a single item from the memory, the area key and area value for the area are the key and value for the single item.

In some cases, each item corresponds to a portion of the neural network input, and the value for each memory item is an encoded representation of the neural network input.

In some implementations, other components of the neural network system provide as input to the attention layer the values and keys for the items in the memory and the attention query.

In some cases, the key is the same as the value the item in the memory. In other cases, the key and the value are different for each item in the memory.

In some implementations, the key of an area, μi, is defined simply as the mean vector of the key of each item in the area, as expressed in equation 3.

μi=1ri⁢∑j=1ri⁢ki,j(3)
Here |ri| is the size of the area ri.

In these implementations, the value of an area is defined as the sum of all value vectors of the items in the area, as expressed in equation 4.

With the keys and values defined, a system can use the standard way for calculating attention as described above in reference to equations 1 and 2. This basic form of area attention, expressed in equations 3 and 4, is parameter-free; that is, it does not introduce any parameters to be learned. Essentially, equation 3 and 4 use average and sum pooling over an area of vectors. It is possible to use other pooling methods, e.g., max pooling, to compute the key and value vector for each area.

Richer representations of each area can also be derived by using features other than the mean of the key vectors of the area. For example, a system can consider the standard deviation of the key vectors within each area, shown in equation 5, below.

The system can also consider other features, e.g., the height and width of each area.

When considering these other features, the system can process the features, i.e., the mean, the standard deviation and, if used, features identifying the height and width of the area, using a multi-layer perceptron neural network that is trained jointly with the neural network system to generate the key for the area.

To efficiently compute μi, σi, and νiri, a system can implement an optimization technique known as a summed area table. The summed area table is based on an integral image, I, which can be efficiently computed in a single pass of the memory. With the integral image, the system can calculate the key and value of each area in constant time.

FIG.4andFIG.5show pseudocode for performing area attention calculations. In particular, they perform the calculations of equations 3, 4 and 5 for all areas that have less than the maximum size given an input memory grid as well as determine the shape size of each area. The operations in the pseudocode can be used to efficiently compute the area keys and area values for all of the areas for a given two-dimensional memory grid. The operations in the pseudocode are based on the tensor operations of the TensorFlow open-source software library.