Full attention with sparse computation cost
The present disclosure is directed to machine learning model architectures which provide full attention capability in each attention head while maintaining low computation and memory complexity. Specifically, according to one aspect of the present disclosure, example attention models provided herein can treat the self-attention mechanism as a conditional expectation over embeddings at each location and approximate the conditional distribution with a structured factorization. Each location can attend to all other locations, either via direct attention, or through indirect attention to group representations, which are again conditional expectations of embeddings from corresponding local regions.
1 . A computing system for performing an attention mechanism with reduced computational requirements, the computing system comprising:
one or more processors; and
one or more non-transitory computer-readable media that collectively store a machine- learned attention model configured to receive and process a model input to generate a model output, wherein the machine-learned attention model comprises one or more attention layers, wherein at least one of the attention layers comprises one or more attention heads, and wherein at least one of the attention heads is configured to:
receive a sequence of input data elements; and
apply a structured attention pattern to the sequence of input data elements to generate a sequence of output data elements;
wherein, for each input data element in the sequence of input data elements, the structured attention pattern specifies one or more locations of direct expectation and one or more groups of locations of local expectation; and
wherein, for each of the one or more groups of locations with local expectation, the at least one of the attention heads is configured to:
determine a single group probability for the group of locations; and
determine an individual local expectation for each location in the group of locations.
2 . The computing system of claim 1 , wherein the structured attention pattern comprises a full attention pattern that has a support that covers an entirety of the sequence of input data elements.
3 . The computing system of claim 1 , wherein the sequence of input data elements comprises a sequence of input embeddings.
4 . The computing system of claim 1 , wherein the structured attention pattern specifies a plurality of groups of locations of local expectation.
5 . The computing system of claim 1 , wherein the machine-learned attention model comprises a plurality of attention layers, wherein each of the plurality of attention layers comprises a plurality of attention heads, and wherein each of the plurality of attention heads is configured to apply the structured attention pattern.
6 . The computing system of claim 1 , wherein the structured attention pattern comprises a partition tree having two or more hierarchical partition levels.
7 . The computing system of claim 1 , wherein the at least one of the attention heads is configured to:
for each of the one or more groups of locations with local expectation, normalize the individual local expectations for the group of locations; and
normalize the one or more locations of direct expectation and the single group probabilities for the one or more groups of locations.
8 . The computing system of claim 1 , wherein the structured attention pattern comprises a combiner-fixed attention pattern.
9 . The computing system of claim 1 , wherein the structured attention pattern comprises a combiner-logsparse attention pattern.
10 . The computing system of claim 1 , wherein the structured attention pattern comprises a combiner-axial attention pattern.
11 . The computing system of claim 1 , wherein the structured attention pattern comprises a machine-learned factorization plan that specifies the one or more locations of direct expectation and the one or more groups of locations of local expectation.
12 . The computing system of claim 1 , wherein the model input comprises natural language data.
13 . The computing system of claim 1 , wherein the model input comprises image data, audio data, protein data, or computer-readable code data.
14 . A computer-implemented method for performing an attention mechanism with reduced computational requirements, the method comprising:
receiving a sequence of input data elements; and
applying a structured attention pattern to each of the sequence of input data elements to generate a sequence of output data elements, wherein applying the structured attention pattern to each input data element comprises:
determining one or more locations of direct expectation and one or more groups of locations of local expectation; and
for each of the one or more locations of direct expectation, determining a direct expectation; and
for each of the one or more groups of locations with local expectation:
determining a single group probability for the group of locations; and
determining an individual local expectation for each location in the group of locations.
15 . The computer-implemented method of claim 14 , wherein the structured attention pattern has a support that covers an entirety of the sequence of input data elements.
16 . The computer-implemented method of claim 14 , wherein, for at least one of the one or more groups of locations with local expectation, determining the individual local expectation for each location in the group of locations comprises re-using the individual local expectation for each location in the group of locations that was previously computed for a different input data element in the sequence of input data elements.
17 . The computer-implemented method of claim 14 , wherein the sequence of input data elements comprises a sequence of input embeddings.
18 . The computer-implemented method of claim 14 , wherein the structured attention pattern specifies a plurality of groups of locations of local expectation.
19 . One or more non-transitory computer-readable media for performing an attention mechanism with reduced computational requirements, wherein the one or more non-transitory computer-readable media collectively store:
a machine-learned attention model configured to receive and process a model input to generate a model output, wherein the machine-learned attention model comprises one or more attention layers, wherein at least one of the attention layers comprises one or more attention heads, and wherein at least one of the attention heads is configured to:
receive a sequence of input data elements; and
apply a structured attention pattern to the sequence of input data elements to generate a sequence of output data elements;
wherein, for each input data element in the sequence of input data elements, the structured attention pattern specifies one or more locations of direct expectation and one or more groups of locations of local expectation; and
wherein, for each of the one or more groups of locations with local expectation, the at least one of the attention heads is configured to:
determine a single group probability for the group of locations; and
determine an individual local expectation for each location in the group of locations.