Leveraging redundancy in attention with reuse transformers
Provided are systems and methods that improve the computational efficiency of Transformers or other attention-based neural networks or machine learning models by re-using a number of attention scores between layers and/or heads of the model. To reduce the computational cost of self-attention-based models while achieving comparable or even superior results, example aspects of the present disclosure propose a novel architecture that reuses attention scores computed in one layer in one or multiple subsequent layers.
1 . A computing system for performing attention with improved computational efficiency, the computing system comprising:
one or more processors; and
one or more non-transitory computer-readable media that collectively store:
a machine-learned model configured to perform attention operations on a model input to generate a model output, wherein the machine-learned model comprises a plurality of attention layers, and wherein each attention layer comprises a plurality of attention heads; and
instructions that, when executed by the one or more processors, cause the computing system to perform operations, the operations comprising:
obtaining the model input; and
processing the model input with the machine-learned model to generate the model output, wherein processing the model input with the machine-learned model comprises, for one or more reuse layers of the plurality of attention layers:
receiving a layer input;
accessing one or more attention scores computed for a previous layer of the plurality of attention layers; and
for one or more reuse heads of the plurality of attention heads of the one or more reuse layers, re-using the one or more attention scores computed for the previous layer to generate respective head outputs for the one or more reuse heads based on the layer input, wherein, for at least one of the one or more reuse layers, the one or more reuse heads comprise a proper subset of the plurality of attention heads of the one or more reuse layers; and
generating a layer output based on the respective head outputs of the one or more reuse heads of the plurality of attention heads of the reuse layer.
2 . The computing system of claim 1 , wherein:
the operations further comprise, for each of one or more non-reuse heads of the plurality of attention heads of the reuse layer:
computing one or more new attention scores based on the layer input; and
using the one or more new attention scores to generate a respective head output of the non-reuse head;
wherein generating the layer output comprises concatenating the respective head outputs of the one or more reuse heads and the respective head outputs of the one or more non-reuse heads.
3 . The computing system of claim 2 , wherein:
accessing the one or more attention scores computed for the previous layer comprises accessing a carry-forward tensor that contains the one or more attention scores computed for the previous layer; and
the operations further comprise inserting the one or more new attention scores into the carry-forward tensor.
4 . The computing system of claim 1 , wherein the one or more reuse layers comprise all the plurality of attention layers of the machine-learned model except an initial layer.
5 . The computing system of claim 1 , wherein the one or more reuse layers comprise a proper subset of the plurality of attention layers of the machine-learned model.
6 . The computing system of claim 1 , wherein re-using the one or more attention scores computed for the previous layer to generate the respective head output based on the layer input comprises directly using the one or more attention scores computed for the previous layer in the reuse head, wherein directly using the one or more attention scores comprises multiplying the one or more attention scores by a value matrix that is based on the layer input.
7 . The computing system of claim 1 , wherein re-using the one or more attention scores computed for the previous layer to generate the respective head output based on the layer input comprises using a weighted combination of the one or more attention scores computed for two or more previous layers in the reuse head.
8 . The computing system of claim 7 , wherein the weighted combination is generated using one or more learned weighting parameters.
9 . The computing system of claim 1 , wherein a number of the reuse layers and a number of the one or more reuse heads in each reuse layer comprise user-defined hyperparameters.
10 . The computing system of claim 1 , wherein a number of the one or more reuse heads in each reuse layer is constant.
11 . The computing system of claim 1 , wherein a number of the one or more reuse heads in each of the one or more reuse layers decreases in relation to layer depth within the machine-learned model.
12 . The computing system of claim 1 , wherein the one or more reuse heads and the one or more reuse layers were specified and effectuated during training of the machine-learned model.
13 . The computing system of claim 1 , wherein the one or more reuse heads and the one or more reuse layers were specified and effectuated subsequent to training of the machine-learned model.
14 . The computing system of claim 1 , wherein, for each of the one or more reuse layers, the previous layer comprises a directly sequentially previous layer.
15 . The computing system of claim 1 , wherein the model input comprises imagery and the model output comprises a prediction descriptive of a characteristic of the imagery.
16 . The computing system of claim 1 , wherein the machine- learned model comprises an encoder portion and a decoder portion, wherein at least one of the one or more reuse layers is contained within the decoder portion, and wherein the previous layer in contained within the encoder portion.
17 . The computing system of claim 1 , wherein the machine-learned model comprises a decoder-only model.
18 . The computing system of claim 1 , wherein the one or more attention scores computed for the previous layer comprise attention scores only for preceding tokens in a sequence.
19 . The computing system of claim 1 , wherein the machine-learned model comprises an encoder-only model.
20 . The computing system of claim 1 , wherein the one or more attention scores computed for the previous layer comprise attention scores for both preceding and subsequent tokens in a sequence.
21 . A computer-implemented method for performing attention with improved computational efficiency, the method comprising:
obtaining a model input; and
processing the model input with a machine-learned model to generate a model output;
wherein processing the model input with the machine-learned model comprises:
for at least an initial layer of the machine-learned model that comprises a first plurality of heads:
newly generating a first plurality of attention scores; and
computing a first layer output based on the plurality of attention scores; and
for each of one or more reuse layers of the machine-learned model that are subsequent to the initial layer:
determining a second plurality of attention scores, wherein at least some of the second plurality of attention scores are reused from one or more previous layers in one or more reuse attention heads of the reuse layer, the one or more reuse attention heads comprising a proper subset of a plurality of attention heads of the reuse layer; and
computing a second layer output based on the second plurality of attention scores.
22 . One or more non-transitory computer-readable media that collectively store:
instructions for implementing a second neural network layer in a neural network, wherein the neural network also comprises a first neural network layer prior to the second neural network layer, wherein the first neural network layer computes one or more attention scores based on first layer inputs to generate first layer outputs, and wherein the instructions cause one or more computers to perform operations comprising:
receiving second layer inputs to the second neural network layer;
obtaining the one or more attention scores computed by the first neural network layer; and
generating second layer outputs based on the second layer inputs and the one or more attention scores;
wherein generating the second layer outputs comprises re-using, by one or more reuse attention heads of the second neural network layer, the one or more attention scores at least in part in place of computation of new attention scores from the second layer inputs, the one or more reuse attention heads comprising a proper subset of a plurality of attention heads of the second neural network layer.
23 . The one or more non-transitory computer-readable media of claim 22 , wherein re-using the one or more attention scores at least in part in place of computation of new attention scores from the second layer inputs comprises:
computing a linear combination of input embeddings in the second layer inputs using the one or more attention scores.