Attention neural networks with sparse attention mechanisms
AssigneeGoogle Inc.
InventorsJoshua Timothy Ainslie, Santiago Ontañón, Philip Pham, Manzil Zaheer, Guru Guruganesh, Kumar Avinava Dubey, Amr Ahmed
Life of this patent
Estimating the term…
Abstract
Methods, systems, and apparatus, including computer programs encoded on computer storage media, for processing network inputs using an attention neural network that has one or more sparse attention sub-layers. Each sparse attention sub-layer is configured to apply a sparse attention mechanism that attends differently for input positions that are in a first proper subset of the input positions in the input to the sub-layer than for positions that are not in the first proper subset.
Claims20
A system for performing a machine learning task on a network input to generate a network output, the system 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:
- an attention neural network configured to perform the machine learning task, the attention neural network comprising one or more sparse attention layers, each sparse attention layer comprising one or more sparse attention sub-layers, each sparse attention sub-layer configured to:
- receive a sequence of queries derived from an input sequence to the sparse attention layer, the sequence of queries having a respective query at each of a plurality of input positions;
- receive a sequence of keys derived from the input sequence to the sparse attention layer, the sequence of keys having a respective key at each of the plurality of input positions;
- receive a sequence of value inputs derived from the input sequence to the sparse attention layer, the sequence of value inputs having a respective value input at each of the plurality of input positions;
- and
- generate an attended input sequence comprising a respective attended input at each of the plurality of input positions, comprising:
- for each input position in a first proper subset of the input positions, generating the attended input at the input position by:
- using the query at the input position to attend over all of the keys in the sequence of keys to generate a respective weight for all of the input positions and computing a weighted sum of the value inputs at all of the input positions in accordance with the respective weights;
- and
- for each input position in a second proper subset of the input positions, generating the attended input at the input position by:
- using the query at the input position to attend over only the keys at a corresponding proper subset of the input positions to generate a respective weight for each of the input positions in the corresponding proper subset and computing a weighted sum of the value inputs at the corresponding proper subset of the input positions in accordance with the respective weights for the corresponding proper subset of the input positions, the corresponding proper subset of input positions for each input position in the second proper subset including:
- the first proper subset of the input positions;
- and
- one or more input positions outside of the first proper subset of the input positions.
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 an attention neural network configured to perform a machine learning task,
- the attention neural network comprising one or more sparse attention layers, each sparse attention layer comprising one or more sparse attention sub-layers, each sparse attention sub-layer configured to:
- receive a sequence of queries derived from an input sequence to the sparse attention layer, the sequence of queries having a respective query at each of a plurality of input positions;
- receive a sequence of keys derived from the input sequence to the sparse attention layer, the sequence of keys having a respective key at each of the plurality of input positions;
- receive a sequence of value inputs derived from the input sequence to the sparse attention layer, the sequence of value inputs having a respective value input at each of the plurality of input positions;
- and
- generate an attended input sequence comprising a respective attended input at each of the plurality of input positions, comprising:
- for each input position in a first proper subset of the input positions, generating the attended input at the input position by:
- using the query at the input position to attend over all of the keys in the sequence of keys to generate a respective weight for all of the input positions and computing a weighted sum of the value inputs at all of the input positions in accordance with the respective weights;
- and
- for each input position in a second proper subset of the input positions, generating the attended input at the input position by:
- using the query at the input position to attend over only the keys at a corresponding proper subset of the input positions to generate a respective weight for each of the input positions in the corresponding proper subset and computing a weighted sum of the value inputs at the corresponding proper subset of the input positions in accordance with the respective weights for the corresponding proper subset of the input positions, the corresponding proper subset of input positions for each input position in the second proper subset including:
- the first proper subset of the input positions;
- and
- one or more input positions outside of the first proper subset of the input positions.
A method comprising:
- receiving a network input;
- and
- processing the network input using an attention neural network to generate a network output for a machine learning task for the network input, the attention neural network comprising one or more sparse attention layers, each sparse attention layer comprising one or more sparse attention sub-layers, each sparse attention sub-layer configured to:
- receive a sequence of queries derived from an input sequence to the sparse attention layer, the sequence of queries having a respective query at each of a plurality of input positions;
- receive a sequence of keys derived from the input sequence to the sparse attention layer, the sequence of keys having a respective key at each of the plurality of input positions;
- receive a sequence of value inputs derived from the input sequence to the sparse attention layer, the sequence of value inputs having a respective value input at each of the plurality of input positions;
- and
- generate an attended input sequence comprising a respective attended input at each of the plurality of input positions, comprising:
- for each input position in a first proper subset of the input positions, generating the attended input at the input position by:
- using the query at the input position to attend over all of the keys in the sequence of keys to generate a respective weight for all of the input positions and computing a weighted sum of the value inputs at all of the input positions in accordance with the respective weights;
- and
- for each input position in a second proper subset of the input positions, generating the attended input at the input position by:
- using the query at the input position to attend over only the keys at a corresponding proper subset of the input positions to generate a respective weight for each of the input positions in the corresponding proper subset and computing a weighted sum of the value inputs at the corresponding proper subset of the input positions in accordance with the respective weights for the corresponding proper subset of the input positions, the corresponding proper subset of input positions for each input position in the second proper subset including:
- the first proper subset of the input positions;
- and
- one or more input positions outside of the first proper subset of the input positions.
Dependencies are inferred from the claim text.