Context-aware bifurcated attention for efficient sampling
Techniques for improved attention calculations are described. In some examples, the attention calculations are for a Transformer-based model. In some examples, a trained model is applied to input data, wherein the application of the trained model at least includes determining at least one of logits or an output of an attention-block using incremental decoding, wherein the incremental decoding calculation is decomposed into at least an incremental calculation and a context calculation.
1 . A computer-implemented method comprising:
receiving a request to generate code using a trained language model, the request including input data;
processing the input data using the trained language model, wherein the processing includes generating logits using attention computations for an incremental decoding, wherein the incremental decoding is decomposed into at least a context calculation and an incremental calculation, wherein the context calculation includes deriving a value from a query and a key from a context, wherein the incremental calculation includes deriving another value from the query and a key from a previous incremental decoding and a new key from a current step, and wherein the incremental decoding further includes concatenating the context calculation with the incremental calculation; and
outputting generated code from the trained language model.
2 . The computer-implemented method of claim 1 , wherein the incremental decoding includes not duplicating at least one of a context or a weight during the attention computations.
3 . The computer-implemented method of claim 1 , further comprising broadcasting to multiple indices within a sampling batch without duplicating a memory access for context attentions.
4 . A computer-implemented method comprising:
receiving a request to generate code using a trained model, the request including input data;
processing the input data using the trained model, wherein the processing includes generating logits using attention computations for an incremental decoding, wherein the incremental decoding is decomposed into at least a context calculation and an incremental calculation, wherein the context calculation includes deriving a value from a query and a key from a context, wherein the incremental calculation includes deriving another value from the query and a key from a previous incremental decoding and a new key from a current step, and wherein the incremental decoding further includes concatenating the context calculation with the incremental calculation; and
outputting generated code from the trained model.
5 . The computer-implemented method of claim 4 , wherein the trained model utilizes multi-headed self-attention.
6 . The computer-implemented method of claim 4 , wherein the trained model utilizes multi-query self-attention.
7 . The computer-implemented method of claim 4 , wherein the trained model is Transformer-based.
8 . The computer-implemented method of claim 7 , wherein the Transformer-based model is a bidirectional encoder representations Transformer.
9 . The computer-implemented method of claim 7 , wherein the Transformer-based model is a generative pre-trained Transformer.
10 . The computer-implemented method of claim 7 , wherein an attention mechanism of the Transformer-based model utilizes quantization.
11 . The computer-implemented method of claim 4 , further comprising performing the incremental decoding without broadcasting the key from the context.
12 . The computer-implemented method of claim 11 , further comprising concatenating attention weights for the context calculation and attention weights for the incremental calculation.
13 . The computer-implemented method of claim 4 , wherein the request includes a batch size indicating a number of outputs to present.
14 . The computer-implemented method of claim 4 , wherein the incremental decoding includes not duplicating at least one of a context or a weight during the attention computations.
15 . A system comprising:
a first one or more electronic devices implementing a data storage service in a multi-tenant provider network; and
a second one or more electronic devices implementing a model hosting service in the multi-tenant provider network, the model hosting service hosting a trained model that includes instructions stored in non-transitory memory that upon execution by one or more processors cause the model hosting service to:
receive a request to generate code using the trained model, the request including input data stored in the data storage service;
process the input data using the trained model, wherein the processing includes generating logits using attention computations for an incremental decoding, wherein the incremental decoding is decomposed into at least a context calculation and an incremental calculation, wherein the context calculation includes deriving a value from a query and a key from a context, wherein the incremental calculation includes deriving another value from the query and a key from a previous incremental decoding and a new key from a current step, and wherein the incremental decoding further includes concatenating the context calculation with the incremental calculation; and
output generated code from the trained model.
16 . The system of claim 15 , wherein the instructions stored in the non-transitory memory comprise further instructions that, upon execution by the one or more processors, cause the model hosting service to broadcast to multiple indices within a sampling batch without duplicating a memory access for context attentions.
17 . The system of claim 15 , wherein the model hosting service is a part of a code generation service.
18 . The system of claim 15 , wherein the instructions stored in the non-transitory memory comprise further instructions that, upon execution by the one or more processors, cause the model hosting service to perform the incremental decoding without broadcasting the key from the context.
19 . The system of claim 18 , wherein the instructions stored in the non-transitory memory comprise further instructions that, upon execution by the one or more processors, cause the model hosting service to concatenate attention weights for the context calculation and attention weights for the incremental calculation.
20 . The system of claim 18 , wherein the incremental decoding includes not duplicating at least one of a context or a weight during the attention computations.