Neural network inference using a quantized key-value cache

Abstract

Methods, systems, and apparatus, including computer programs encoded on computer storage media, for using a Transformer-based neural network to generate output sequences. To generate the output sequences, the Transformer-based neural network is configured to perform quantized inference.

Description

[0001] NEURAL NETWORK INFERENCE USING A QUANTIZED KEY- VALUE CACHE

[0002] CROSS-REFERENCE TO RELATED APPLICATION

[0003] This application claims priority to U.S. Provisional Application No. 63 / 540,889, filed on September 27, 2023. The disclosure of the prior application is considered part of and is incorporated by reference in the disclosure of this application.

[0004] BACKGROUND

[0005] This specification relates to neural networks.

[0006] 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 the next layer in the network, i.e., the next hidden layer or the 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.

[0007] SUMMARY

[0008] This specification describes a quantized inference system implemented as computer programs on one or more computers in one or more locations that implements and uses a Transformer-based neural network to generate output sequences. To generate the output sequences, the Transformer-based neural network is configured to perform quantized inference.

[0009] According to an aspect, there is provided a computer-implemented method of generating an output sequence that comprises an output token at each of a plurality of output positions using a neural network, wherein the neural network comprises a plurality of attention layers, wherein each attention layer of one or more of the plurality of attention layers comprises one or more attention heads, and wherein the method comprises, for each attention head of the one or more attention heads and for a particular output position in the plurality of output positions: maintaining, in one or more memory devices, a plurality of quantized keys that were generated prior to receiving an embedded representation of an output token at the particular output position; receiving the embedded representation of the output token at the particular output position; applying a key linear transformation to the embedded representation to generate an original key for the particular output position; quantizing the original key to generate a quantized original key that is represented using a lower precision format than the original key; and generating an updated embedded representation of the output token at the particular output position based on applying an attention mechanism based on the quantized original key for the particular output position and the plurality of quantized keys in the one or more memory devices.

[0010] Applying the attention mechanism based on the quantized original key for the particular output position and the plurality of quantized keys in the one or more memory devices may comprise: retrieving the plurality of quantized keys from the one or more memory devices; dequantizing the plurality of quantized keys to generate a plurality of dequantized keys; dequantizing the quantized original key for the particular output position to generate a dequantized original key; and computing an output of the attention mechanism based on the dequantized original key for the particular output position and the plurality of dequantized keys.

[0011] Applying the attention mechanism based on the quantized key for the particular output position and the plurality of quantized keys in the one or more memory devices may comprise: retrieving the plurality of quantized keys from the one or more memory devices; and computing an output of the attention mechanism based on the quantized original key for the particular output position and the plurality of quantized keys.

[0012] The original key may be a vector comprising one or more elements, and wherein quantizing the original key to generate the quantized original key may comprise, for each element in the original key: identifying, from among a set of anchor points, a nearest anchor point to the element as the identified anchor point that corresponds to the element in the original key; and setting the element of the original key to the identified nearest anchor point.

[0013] The original key may be a vector comprising one or more elements, and wherein quantizing the original key to generate the quantized original key may comprise: determining a quantization range based on the one or more elements in the original key; and mapping each element of the original key to a corresponding, mapped element within the quantization range based on applying a scale factor.

[0014] Mapping the element in the original key to the mapped element within the quantization range may comprise: determining an upper limit of the quantization range based on a maximal absolute value of the one or more elements in the original key.

[0015] Mapping the element in the original key to the mapped element within the quantization range may comprise: determining that a particular element of the original key is either above or below the quantization range; and in response, clipping the mapped element of the particular element to be equal respectively to an upper limit or a lower limit of the quantization range.

[0016] The method may further comprise: generating, based at least on the updated embedded representation of the output token at the particular output position that has been generated by the attention head, a time step output for a current generation time step that defines a score distribution over a set of possible tokens; and selecting, using the time step output, a selected token from the set of possible tokens as the output token at the particular output position.

[0017] The method may be performed by a Transformer neural network. The respective quantization range for the attention head may be predetermined prior to generating the output sequence by using the Transformer neural network, and the method may further comprise: determining an update to parameter values of the Transformer neural network based at least on the selected token.

[0018] The one or more memory devices may comprise a plurality of memory devices that each correspond to a different level of precision and that each store a corresponding subset of the quantized keys that have been quantized to have the corresponding level of precision.

[0019] Applying the attention mechanism may comprise applying the attention mechanism also using (i) an original query that is generated from the embedded representation of the output token at the particular output position and (ii) a plurality of quantized values that were generated prior to receiving the embedded representation of the output token at the particular output position.

[0020] The plurality of quantized values may be quantized to have a same precision as the plurality of quantized keys.

[0021] The plurality of quantized values may be quantized to have a different precision from the plurality of quantized keys.

[0022] According to another aspect, there is provided a computer-implemented method of generating an output sequence that comprises an output token at each of a plurality of output positions, based on an input sequence of input tokens at each of a plurality of input positions corresponding to the output positions, the method being performed by using a plurality of attention layers arranged in a series with the output of each attention layer except a first of the attention layers being an input to the next attention layer of the series, each of the attention layers being arranged to receive, for each of the input positions, a corresponding attention layer input vector based on a corresponding input token, and to generate from the attention layer input vector a corresponding attention layer output vector for the input position, the output token for the corresponding output position being based on the attention layer output vector for a last of the attention layers: for each attention layer, the generating of the attention layer output vector being performed by, for each of one or more attention heads of the attention layer, generating one or more transformed vectors for the input position by applying respective linear transformations to the attention layer input vector, the one or more transformed vectors including a key vector for the input position; quantizing the key vector for the input position to generate a quantized key vector for the input position, the quantized key vector being represented using a lower precision format than the key vector, and storing the quantized key vector in memory; and for each input position except the first, obtaining the attention layer output vector for the input position using (i) the key vector for the input position or the quantized key vector for the input position, and (ii) a corresponding said quantized key vector for one or more previous input positions retrieved from the memory.

[0023] The one or more transformed vectors may further include a query vector for the input position, and a value vector for the input position, and the attention layer output vector for the input position may be obtained by: calculating a weight for the input position based on the product of the query vector for the input position and the key vector for the input position, calculating weights for the previous input positions based on the product of respective products of (i) the query vector for the input position and (i) the key vector for previous input positions, forming the attention layer output vector based on a sum, weighted by the corresponding weights, of the value vector for the input position and the value vectors of the previous input positions.

[0024] The one or more transformed vectors may further include a query vector for the input position, and a value vector for the input position, and the attention layer output vector for the input position may be obtained by: calculating a weight for the input position based on the product of the query vector for the input position and the key vector for the input position, calculating weights for the previous input positions based on the product of respective products of (i) the query vector for the input position and (i) the key vector for previous input positions, quantizing the value vector for the input position to generate a quantized value vector for the input position, the quantized value vector being represented using a lower precision format than the value vector for the position, and storing the quantized value vector in memory, forming the attention layer output vector based on a sum, weighted by the corresponding weights, of the value vector or quantized value vector for the input position and the quantized value vectors of the previous input positions. The plurality of quantized value vectors may be quantized to have a same precision as the plurality of quantized key vectors.

[0025] The plurality of quantized value vectors may be quantized to have a different precision from the plurality of quantized key vectors.

[0026] Quantizing the key vector for the input position may comprise: determining a quantization range based on the plurality of elements in the key vector; and mapping each of the plurality of elements in the key vector to a corresponding, mapped element within the quantization range based on applying a scale factor.

[0027] Quantizing the key vector for the input position may comprise: determining an upper limit of the quantization range based on a maximal absolute value of the one or more elements in the key vector for the input position.

[0028] Quantizing the key vector for the input position may comprise: determining that a particular element of the key vector for the position is either above or below a predetermined range that corresponds the quantization range; and in response, setting a corresponding element of the quantized key vector as respectively an upper limit or a lower limit of the predetermined range.

[0029] The method may further comprise: at the last of the attention layers: generating, based at least on the attention layer output vectors generated by the corresponding one or more attention heads, a time step output that defines a score distribution over a set of possible tokens; and selecting, using the score distribution, a token from the set of possible tokens as the output token at the particular output position.

[0030] The one or more memories may comprise a plurality of memories that each correspond to a different level of precision and that each store a corresponding subset of the quantized key vectors, the method may further comprise moving at least one of the quantized key vectors from a memory corresponding to a higher level of precision to a lower level of precision and correspondingly reducing the precision of the at least one of the quantized key vectors.

[0031] The method may be performed following a training phase of generating parameter values defining the linear transformations based on a training database of tuplets composed of training input sequences and corresponding training output sequences, the training process comprising, for successive tuplets or groups of tuplets, providing the training input sequences as inputs to the first of the attention layers; generating provisional output sequences of output tokens based on corresponding outputs of the last of the attention layers; and modifying the parameter values defining at least one of the linear transformations to increase a measure of the similarity of the provisional output sequences and the corresponding training output sequences; wherein the generating of the provisional output sequences comprises quantizing key vectors generated based on the training input sequences.

[0032] Modifying the parameter values may comprise converting the quantized key vectors generated based on the training input sequences into a higher precision format than the corresponding quantized key vectors.

[0033] The input sequence may represent an input image and the output sequence may be a network output which is the result of performing a task on the input sequence, wherein the task may be image classification and the network output is scores for each of a set of object categories, with each score representing an estimated likelihood that the image contains an image of an object belonging to the category; the task may be image embedding generation and the network output is a numeric embedding of the input image; the task may be object detection and the network output identifies locations in the input image at which particular types of objects are depicted; or the task may be image segmentation and the network output assigns each pixel of the input image to a category from a set of categories.

[0034] According to another aspect, there is provided one or more computer-readable storage media encoded with instructions that, when executed by one or more computers, cause the one or more computers to perform the operations of the above method aspects.

[0035] According to a further aspect, there is provided a system comprising one or more computers and one or more storage devices storing instructions that when executed by one or more computers cause the one or more computers to perform the respective operations of the above method aspects.

[0036] It will be appreciated that features described in the context of one aspect may be combined with features described in the context of another aspect.

[0037] The subject matter described in this specification can be implemented in particular embodiments so as to realize one or more of the following advantages.

[0038] A Transformer-based neural network is a neural network that includes one or more attention layers. Each attention layer can include any variations of a query-key-value attention mechanism. When executing the operations of the Transformer-based neural network to generate an output over multiple time steps, an inference system running on one or more data processing apparatuses stores the keys and values generated at each attention layer in a local memory of the data processing apparatuses. For example, when running on specialized hardware (e.g., special purpose neural network hardware accelerators), the inference system can store the keys and values in one or more memory devices local to the specialized hardware, e.g., in an on-device high-bandwidth memory (HBM).

[0039] While scaling Transformer-based neural networks to a larger model size (e.g., 100 billion, 500 billion, or more parameters), to having a larger context window size (e.g., 1 million tokens, 10 million tokens, or more), or both has led to state-of-the-art performance on a range of machine learning tasks and thus expanded the practical utility of Transformerbased neural networks, the total number of keys and values that need to be stored has also increased, thereby increasing the memory footprint of a Transformer-based neural network at inference time. Moreover, the large memory footprint gives rise to a large amount of memory traffic to load the stored keys and values from high-bandwidth memory (HBM) into the compute cores at each time step. Thus, both the memory and communication overhead can be high to execute large scale Transformer-based neural networks.

[0040] The quantization techniques described in this specification quantize original keys and values into quantized keys and quantized values, which can be stored in the memory using lower precision formats, in a manner that results in negligible accuracy loss when generating an output based on the quantized keys and quantized values. The lower precision format can be any format that uses fewer number of bits to represent each value compared to the higher precision format that was used during the training of the Transformer-based neural network. The quantization techniques can be applied either during training (“quantization-aware” training) or post-training. In particular, in the latter case, no modification to the training pipeline or re-training of the neural networks is necessary.

[0041] Accordingly, the described quantization techniques allow for a quantized inference system to effectively perform machine learning tasks while having a reduced memory footprint (e.g., occupying less space in the memory) and a reduced communication burden compared with conventional inference systems (i.e., inference systems that store the original, unquantized keys and values). Therefore the quantized inference system enables more efficient use of computational resources (in particular, memory and network resources). As a result, the quantized inference system can more readily be deployed to resource-constrained environments than conventional inference systems.

[0042] The details of one or more embodiments of the subject matter of this specification are set forth in the accompanying drawings and the description below. Other features, aspects, and advantages of the subject matter will become apparent from the description, the drawings, and the claims. BRIEF DESCRIPTION OF THE DRAWINGS

[0043] FIG. 1 shows an example quantized inference system.

[0044] FIG. 2 is a flow diagram of an example process for generating an output sequence based on an input sequence.

[0045] FIGS. 3A-B are flow diagrams of sub-steps of one of the steps of the process of FIG. 2.

[0046] FIG. 4 is a flow diagram of sub-steps of an implementation of one of the steps of the process of FIG. 2.

[0047] FIG. 5 is a flow diagram of sub-steps of another implementation of one of the steps of the process of FIG. 2.

[0048] FIG. 6 is a flow diagram of an example process for training an attention neural network.

[0049] Like reference numbers and designations in the various drawings indicate like elements.

[0050] DETAILED DESCRIPTION

[0051] FIG. 1 shows an example quantized inference system 100. The quantized inference system 100 is an example of a system implemented as computer programs on one or more computers in one or more locations, in which the systems, components, and techniques described below can be implemented.

[0052] The quantized inference system 100 is a system that can receive an input sequence 102 and generate an output sequence 132 based on the input sequence 102. The input sequence 102 includes a respective input token at each of multiple input positions in an input order. The output sequence 132 includes a respective output token at each of multiple output positions in an output order. For example, the input sequence 102 provides context for the output sequence 132, and the output sequence 132 is a response to the input sequence 102.

[0053] The input tokens and the output tokens can be selected from a vocabulary of tokens. For example, the vocabulary of tokens can include one or more of characters, sub-words, words, punctuation marks, numbers, or other symbols that appear in a corpus of natural language text and / or computer code.

[0054] Additionally, or alternatively, the vocabulary of tokens can include tokens that can represent data other than text. For example, the vocabulary of tokens can include image tokens that represent a discrete set of image patch embeddings of an image that can be generated by an image encoder neural network based on processing the image patches of the image. As another example, the vocabulary of tokens can include audio tokens that represent code vectors in a codebook of a quantizer, e.g., a residual vector quantizer.

[0055] The output sequences 132 generated by the quantized inference system 100 can be used in any of a variety of ways. For example, the system can provide the output sequences 132 for presentation to the user on a client device. As another example, the system can provide the output sequences 132 to another component in the system for further processing. As yet another example, the system can store the output sequences 132 in a memory device for some future purpose.

[0056] As an example, the quantized inference system 100 can generate text sequences, i.e., each output sequence 132 generated by the system is a sequence of tokens from a vocabulary of tokens that includes, e.g., one or more of characters, sub-words, words, punctuation marks, numbers, or other symbols that appear in natural language text. For example, the quantized inference system 100 can generate text sequences in response to input sequences 102 and provide the text sequences for presentation to users.

[0057] As another example, the quantized inference system 100 can generate images or videos that each have multiple frames (where each frame is an image) by generating images as sequences of pixels. For example, each output sequence 132 generated by the system is a sequence of color values for pixels in an image arranged according to a specified order. As a similar example, the quantized inference system 100 can be a video generation system that generates video as a temporal sequence of video frames, where each frame is an output image. As another example, each output sequence 132 generated by the system is a sequence of tokens that represent image patch embeddings of an image which can then be processed by a decoder neural network to generate the image (pixel values).

[0058] As another example, the quantized inference system 100 can generate multi-modal output sequences, i.e., each output sequence 132 generated by the system is a sequence of tokens that represents data in multiple data modalities, e.g., an output sequence of tokens that represents both audio and video data, an output sequence of tokens that represents both text and image data, or an output sequence of tokens that represents both text and audio data.

[0059] As another example, the input sequence 102 can be an input sequence of text and the output sequence is another sequence of text, e.g., a completion of the input sequence of text, a paraphrase of the input sequence of text, a response to a question posed in the input sequence, or a sequence of text that is about a topic specified by the input sequence of text. As another example, the input sequence 102 can be an input other than text, e.g., an image, a video, an audio, or some combination of these, and the output sequence 132 can be text that describes the input.

[0060] As another example, the input sequence 102 represents data to be compressed, e.g., image data, text data, audio data, or any other type of data; and the output sequence 132 is a compressed version of the data. The tokens included in the output sequence can include any representation of compressed data, e.g., symbols or embeddings to be decoded by a respective neural network.

[0061] As a particular example, the quantized inference system 100 can be part of a dialog system and the input sequence 102 can include audio or text from the most recent conversational turn submitted by a user of the dialog system during the dialog while the output sequence 132 is the next turn in the conversation, e.g., either text or audio that is a response to the most recent conversational turn. Optionally, the input sequence 102 can also include one or more historical conversational turns that occurred earlier in the conversation.

[0062] As another particular example, the quantized inference system 100 can be part of a machine translation system and the input sequence 102 can include text in a source language while the output sequence 132 is text in a target language that is a translation of the source text into the target language.

[0063] As another particular example, the quantized inference system 100 can be part of a natural language processing system. For example, if the input sequence 102 is a sequence of words in an original language, e.g., a sentence or phrase, the output sequence 132 can 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 102 is a sequence of words that form a question, the output sequence 132 can be a sequence of words that form an answer to the question.

[0064] As another particular example, the quantized inference system 100 can be part of a computer-assisted medical diagnosis system. For example, the input sequence 102 can be a sequence of data from an electronic medical record and the output sequence 132 can be a sequence of predicted treatments.

[0065] As another particular example, the quantized inference system 100 can be part of a computer code generation system and the input sequence 102 can be a text description of a desired piece of code or a snippet of computer code in a programming language and the output sequence 132 can be computer code, e.g., a snippet of code that is described by the input sequence 102 or a snippet of code that follows the input sequence 102 in a computer program. As another particular example, the quantized inference system 100 can be part of a multi-modal system that processes, e.g., both text and image input sequences, or both text and audio input sequences, and generates the output sequences 132 that are either in a single data modality or in multiple data modalities, e.g., text and image output sequences, or text and audio output sequences. Examples of such multi-modal systems include an image captioning system, a text-based image search system, an image-based question answering system, and so on.

[0066] As another particular example, the quantized inference system 100 can be part of or associated with a search system that facilitate searching of resources on the Internet. A resource can be any data that can be provided over the Internet. A resource can be identified by a resource address that is associated with the resource. Resources include web pages, word processing documents, portable document format (PDF) documents, images, video, and news feed sources, to name a few.

[0067] In this particular example, the search system can receive search queries submitted by client devices and, in response, identify resources that are relevant to the search query in the form of search results and returns the search results to the user devices in search results pages. A search result page can include search result data generated by the search system that identifies a resource responsive to a search query, and includes a link to the resource. The search result page can additionally include a result in the form of an output sequence 132 that is generated by the quantized inference system 100 based on an input sequence 102 derived from the search query.

[0068] To generate the output sequences 132, the quantized inference system 100 includes an attention neural network 105. The attention neural network 105 operates auto-regressively to generate the output tokens in the output sequence 132 over a plurality of time steps. Each time step corresponds to an output position in the output sequences 132. That is, the attention neural network 105 generates an output token at each time step to add to the output sequences 132.

[0069] The attention neural network 105 auto-regressively generates an output sequence 132 by generating each given output token in the output sequence 132 conditioned on a current input sequence 106 that includes (i) the input tokens included in the input sequence 102 followed by (ii) any output tokens that precede the given output token in the output sequence 132, i.e., the output tokens that have already been generated for any previous output positions in the output sequence that precede the given output position of the given output token. To generate a given output token at a given output position within an output sequence, the attention neural network 105 can process the current input sequence to generate a score distribution, e.g., a probability distribution, that assigns a respective score, e.g., a respective probability, to each token in a vocabulary of tokens.

[0070] The attention neural network 105 can then select, as the given output token, a token from the vocabulary using the score distribution. For example, the attention neural network 105 can greedily select the highest-scoring token or can sample, e.g., using nucleus sampling or another sampling technique, a token from the score distribution.

[0071] The attention neural network 105 can have any of a variety of Transformer-based neural network architectures, e.g., encoder-only Transformer architectures, encoder-decoder Transformer architectures, decoder-only Transformer architectures, diffusion Transformer architectures, other attention-based architectures, and so on.

[0072] Examples of such Transformer-based neural network architectures include those described in Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J Liu. Exploring the limits of transfer learning with a unified text-to-text transformer. arXiv preprint arXiv: 1910.10683, 2019; Daniel Adiwardana, Minh-Thang Luong, David R. So, Jamie Hall, Noah Fiedel, Romal Thoppilan, Zi Yang, Apoorv Kulshreshtha, Gaurav Nemade, Yifeng Lu, and Quoc V. Le. Towards a human-like open-domain chatbot. CoRR, abs / 2001.09977, 2020; Tom B Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are fewshot learners. arXiv preprint arXiv:2005.14165, 2020; Aakanksha Chowdhery, et al. PaLM: Scaling Language Modeling with Pathways, arXiv preprint arXiv: 2204.02311; Rohan Anil, et al. Palm 2 technical report. arXiv preprint arXiv:2305.10403, 2023; and Gemini Team, et al. Gemini: a family of highly capable multimodal models. arXiv preprint arXiv:2312.11805 (2023).

[0073] As a particular example, the attention neural network 105 can be one of the language model neural networks described in Sparrow (Glaese et al. arXiv:2209.14375), Chinchilla (Hoffmann et al. arXiv:2203.15556), and PaLM 2 (Anil, et al. arXiv:2305: 10403). As another particular example, the attention neural network 105 can be one of the vision language model neural networks described in Flamingo (Alayrac et al. arXiv: 2204.14198), PaLI (Chen et al. arXiv: 2209.06794), and PaLI-X (Chen et al. arXiv: 2305.18565). As yet another particular example, the attention neural network 105 can be one of the foundation model neural networks described in Imagen (Saharia et al. arXiv: 2205.11487) and Parti (Yu et al. arXiv: 2206.10789). A foundation model is a large-scale machine learning model trained on a broad data set that can be adapted and fine-tuned for a wide variety of applications and downstream tasks.

[0074] Generally, however, the attention neural network 105 includes (i) a sequence of multiple attention layers 110 that each apply an attention mechanism (ii) and an output subnetwork that processes an output of the last attention layer to generate the score distribution. The sequence of multiple attention layers 110 can be arranged in a series, such that each attention layer (except the last) passes its output to the next.

[0075] Thus, although one attention layer is depicted in FIG. 1 for convenience, the attention neural network 105 generally includes many other layers, including, for example, embedding layers, output layers, and other attention layers.

[0076] During the processing of a current input sequence 106, each attention layer 110 in the sequence receives an input vector for at least the last token in the current input sequence 106. The input vector will also be referred to as “attention layer input vector” or “embedded representation” in this specification.

[0077] The attention layer 110 then updates the input vector for at least the last token in the current input sequence 106 at least in part by applying a (learned) function (an “attention mechanism”) to generate an output vector for at least the last token in the current input sequence 106. The input vector and the output vector typically have the same dimension, i.e., include the same number of vector elements. The output vector will also be referred to as “attention layer output vector” or “updated embedded representation” in this specification. In this specification the term “learned” refers to a function or value that has been adjusted during the training of the system.

[0078] The input vector for the first attention layer is an embedding of the last token in the current input sequence 106 that is generated by an embedding layer of the attention neural network 105, and the input vector for each subsequent attention layer is the output vector generated by the preceding attention layer.

[0079] The output subnetwork processes the output vector generated by the last attention layer in the sequence for the last token in the current input sequence 106 to generate the probability distribution. For example, the output subnetwork can include a linear layer followed by a softmax layer. As another example, the output subnetwork can include more layers, e.g., can be a multi-layer perceptron (MLP).

[0080] To apply the attention mechanism, each attention layer 110 uses one or more attention heads 120. Each attention head 120 generates a set of queries, a set of keys, and a set of values, and then applies any of a variety of variants of query-key-value (QKV) attention, e.g., a dot product attention function or a scaled dot product attention function, using the queries, keys, and values to generate an output. Each query, key, value can be a vector that includes one or more vector elements (and hence a query, a key, a value can also be referred to as a “query vector,” a “key vector,” and a “value vector,” respectively). When there are multiple attention heads, the attention layer 110 then combines the outputs of the multiple attention heads, e.g., by concatenating the outputs and, optionally, processing the concatenated outputs through a linear layer.

[0081] A detailed explanation of the attention mechanism is provided below.

[0082] The attention mechanism performed by each attention layer 110 generates, for a given input vector x that corresponds to a given token in the current input sequence 106, e.g., the last token in the current input sequence 106, a respective query vector q, a respective key vector k, and a respective value vector v.

[0083] The given output vector for the given input vector x is computed as a weighted sum of up to N+l value vectors (where N is an integer): the value vector for the given input vector x, and the value vectors for the N preceding input vectors (that correspond respectively to the preceding token in the current input sequence).

[0084] The N+l weights (i.e. the weight for value vector for the given input vector x, and the N weights for the value vectors for the N preceding input vectors) may each be computed by a compatibility function e.g. a dot product or scaled dot product, of the query vector for the given input vector x with the (N+l) key vectors corresponding to (N+l) value vectors.

[0085] The attention mechanism is configured to apply a self-attention mechanism over the given input vector x; this may be followed by one or more feed-forward neural network layers to generate the given output vector. In general an attention mechanism determines a relationship between two sequences; a self-attention mechanism is configured to relate different positions in the same sequence to determine a transformed version of the sequence as an output. For example the attention layer input may comprise an input vector for each token in the current input sequence. These input vectors provide an input to the self-attention mechanism and are used by the self-attention mechanism to determine a new representation of the same sequence for the attention layer output, which may comprise an output vector for each token in the current input sequence. An output of the self-attention mechanism may be used as the output vector, or it may be processed by one or more feed-forward layers to provide the output vector. In some implementations the attention mechanism is a self-attention mechanism configured to apply each of a query transformation e.g. defined by a query matrix IV<?, a key transformation e.g. defined by a key matrix WK, and a value transformation e.g. defined by a value matrix Wv, to the input vector x that corresponds to each token in the current input sequence (e.g. a matrix in which each row is one of the input vectors x that correspond respectively to the tokens in the current input sequence) to derive a respective matrix of query vectors Q=XWQ(e.g. a matrix in which each row is one of the query vectors that correspond respectively to the tokens in the current input sequence), matrix of key vectors K=XWK(e.g. a matrix in which each row is one of the key vectors that correspond respectively to the tokens in the current input sequence), and matrix of value vectors M^e.g. a matrix in which each row is one of the value vectors that correspond respectively to the tokens in the current input sequence), which are used determine an attended sequence for the output. (Note that in an option, the calculation of Q, K and V may also include adding a constant bias, but that is neglected in this this explanation; i.e. the bias is assumed to be zero.) For example the attention mechanism may be a dot product attention mechanism applied by applying each query vector to each key vector to determine respective weights for each value vector, then combining the value vectors using the respective weights to determine the output vector that corresponds to each token in current input sequence. The output vector may be scaled by a scaling factor e.g. by the square root of the dimensions of the queries and keys, to implement scaled dot product attention. Thus, for example, an output of the attention mechanism may be determined as softmax V where d is a dimension of each key vector and each query vector (and typically each value vector). A summation over the value vectors is assumed here, weighted by the values softmax In another implementation the attention mechanism be comprise an “additive attention” mechanism that computes the compatibility function using a feed-forward network with a hidden layer. As previously mentioned, the output of the attention mechanism may be further processed, within the attention layer, by one or more fully-connected, feed forward neural network layers.

[0086] The above detailed explanation of the attention mechanism assumes that there is only a single “attention head” in the attention mechanism for the attention layer 110. However, alternatively, as noted above there may be multiple attention heads in each attention layer, each using a different matrix WQ, and, optionally, a different respective set of matrices WKand Wv(the attention mechanism may be referred to as a “multi-query attention mechanism” when the same set of matrices WKand Wvare used by the multiple attention heads). In this case, the attention layer output vector is a concatenation of a respective vector produced by each of the attention heads.

[0087] In practice, to apply the multi-query attention mechanism, the attention layer 110 need only apply the set of matrices WKand Wvonce to generate the matrix of key vectors and the matrix of value vectors, and the matrix of key vectors and the matrix of value vectors are shared across the multiple attention heads within the attention layer 110.

[0088] When the attention neural network 105 operates auto-regressively to generate the output tokens in the output sequence 132 over a plurality of time steps, the attention layer 110 produces output vectors one-by-one (i.e. one output vector at each time step), based on the last input vector (i.e. the input vector corresponding to the last token in the current input sequence) and the N immediately preceding input vectors (i.e. the input vectors corresponding to the TV immediately preceding tokens in the current input sequence).

[0089] As illustrated in FIG. 1, the quantized inference system 100 includes or accesses one or more memory devices 130. The one or more memory devices 130 can either be or include logical memory devices, or can alternatively be or include physical memory devices. For example, when the quantized inference system 100 is implemented on a plurality of hardware accelerators, e.g., graphics processing units (“GPUs”), field-programmable gate arrays (“FGPAs”), or application-specific integrated circuits (“ASICs”), including tensor processing units (“TPUs”), the one or more memory devices 130 can include high bandwidth memory (HBM) devices disposed within the hardware accelerators, such that the one or more memory devices 130 are local to or co-located with computing resources of the hardware accelerators.

[0090] The one or more memory devices 130 stores cached intermediate computation results of the attention mechanisms performed by the attention layers 110, and possibly the intermediate states of other layers, in the attention neural network 105. The cached intermediate computation results are used to accelerate both the speed and computational resource efficiency of the inference of the attention neural network 105.

[0091] Such intermediate computation results — which include the respective key vectors k, and optionally, the respective value vectors v, which are generated from the input vectors x that correspond respectively to the tokens in the current input sequence — are saved (cached) in the one or more memory devices 130 so that the inference can be executed faster and more computational resource efficiently because re-generating all of these intermediate computation results at each time step can be avoided. In particular, as noted above, the quantized inference system 100 performs quantized inference. That is, the quantized inference system 100 stores and calculates numbers with reduced precision, i.e., in a compact format, while performing inference. For example, the quantized inference system 100 may perform inference using special-purpose hardware that is configured to perform computation using fixed point arithmetic in order to allow the hardware to compute inferences with reduced latency. As another example, the quantized inference system 100 may be implemented on a mobile device and may perform quantized inference to limit the amount of computational resources consumed by computing inferences.

[0092] Some conventional inference systems executing a Transformer-based neural network to generate an output at inference time represent and store the intermediate computation results of the attention mechanisms performed by the attention layers 110 in a first format. The first format can for example be the same format as what would be used by a conventional training system if it were to train the attention neural network 105.

[0093] The quantized inference system 100 described in this specification, however, quantizes at least some of the intermediate computation results, e.g., the respective key vectors k and, optionally, the respective value vectors v from the first format to a second format and then stores the respective quantized key vectors k (“quantized keys 132”) and, optionally, the respective quantized value vectors v (“quantized values 134”) in the one or more memory devices 130 accessible by the quantized inference system 100 to facilitate the computation of the outputs of the attention mechanisms.

[0094] The second format can be any format that has a lower precision, lower range, or both than the first format. For example, the first format can be a floating-point format, while the second format can be an integer format. More generally, the second format can be any format that uses fewer number of bits to represent a value than the first format.

[0095] In machine learning, values are often represented in either an integer format or a floating-point format. A value occupies a number of bits in computer memory. Depending on the ranges of numbers used in machine learning, different number of bits may be allocated. For example, an integer may take 8 bits, which has the range of -128 to 127 if it is a signed integer (-27to 27-l), and the range of 0 to 255 if it is an unsigned integer (0 to 28-l). The 8- bit integer format is often shorthanded as int8, e.g., int8 for the signed integers and unit8 for the unsigned integers.

[0096] A floating-point format uses formulaic representation of real numbers as an approximation to support a trade-off between range and precision. A value may be represented approximately to a fixed number of significant digits (the significand, which may also be referred to as mantissa) and scaled using an exponent in some fixed base (note that the exponent may also referred to as a scale, or exponential scale); the base for the scaling may be two, ten, or sixteen (or another integer). The value may be represented in the following form: significand*baseexponent, where the significand may be an integer or a floating-point value, base is an integer greater than or equal to two, and exponent is an integer. For example, a bfloatl6 floating-point format uses 8 exponent bits and 8 significant bits to represent numbers in the range of ~le-38to ~3e38; and thus each value occupies 16 bits in computer memory when represented using bfloatl6 format.

[0097] Converting a value from the first format, e.g., a floating-point format (e.g., a BF16, FP32, or FP64 format) to a number in the second format, e.g., an integer format (e.g., int4, int8, or inti 6 format) therefore means that the value can be stored using less memory, because fewer bits are used to represent the value. As a few examples for illustration, a value in a bfloatl6 format can be converted a quantized value in an int8 format (thus reducing memory usage by a factor of 2), or to another quantized value in an int4 format (thus reducing memory usage by a factor of 4).

[0098] Denoting the number of vector elements of each input vector by n (which is typically equal to t / ), a given input vector may be denoted by the vector x E Rn. The query vector for the given input vector is q=xIV<?, the key vector for the given input vector is k=xIF<, and the value vector for the given input vector is v=xWv. The data which is attended to (in the case of a self-attention mechanism, the N previous input vectors) may be denoted C E RNxn(i.e. the X used by the self-attention mechanism is a concatenation of x and C).

[0099] In an implementation of the present concept, for a given token in the current input sequence 106, e.g., the last token in the current input sequence 106, the key vectors k for the N preceding input vectors (i.e. CWK). and the N value vectors v for the N preceding input vectors (i.e. CIV7), have already been calculated (at a time when the input vector x that corresponds to the given token is used to produce an output vector), and have been stored in the one or more memory devices 130, with at least the N key vectors k and, optionally, the N value vectors v having been stored in the quantized form in the one or more memory devices 130, as described above.

[0100] Thus, the value of the output vector, softmax V can be obtained from the respective query vector q, the respective key vector k, the respective value vector v, and the stored CWKand CIV7(since KTis a concatenation of k and CWK, and V is a concatenation of v and CIV7). Optionally, quantized forms of the respective key vector k and (optionally) the respective value vector v may be used in this calculation, instead of the respective query vector k and the respective value vector v in their original form (i.e. each quantized key vector k and quantized value vector v, produced from the input vector, may be used to obtain the output vector, not just to obtain the subsequent TV output vectors).

[0101] FIG. 2 is a flow diagram of an example process 200 for generating an output sequence based on an input sequence. The input sequence includes an input token at each of a plurality of input positions. The output sequence includes an output token at each of a plurality of output positions. For convenience, the process 200 will be described as being performed by a system of one or more computers located in one or more locations. For example, a quantized inference system, e.g., the quantized inference system 100 of FIG. 1, appropriately programmed in accordance with this specification can perform the process 200.

[0102] The system uses an attention neural network that includes a sequence of attention layers. Each attention layer includes one or more attention heads. To generate the output sequence, the system can, for each particular output position of the plurality of output positions of the output sequence, repeatedly perform an iteration of the process 200 at each attention head of each attention layer in the sequence to generate an output token to occupy the particular output position of the output sequence. Each attention head of each attention layer in the attention neural network can perform an iteration of the process 200 in parallel with each other.

[0103] The system has access to one or more memory devices. The one or more memory devices stores a plurality of quantized keys. The plurality of quantized keys have been generated and stored in the one or more memory devices prior to the time step that corresponds to the particular output position, i.e., prior to the attention neural network receives a current input sequence that includes the most recently generated (selected) output token, i.e., the output token that occupies an immediately previous position that immediately precedes the particular output position.

[0104] In some implementations, the one or more memory devices include a plurality of memory devices, where each memory device corresponds to a different attention layer, or corresponds to a different attention head of a different attention layer.

[0105] In some implementations, the one or more memory devices include a plurality of memory devices, where each memory device corresponds to a different format that is a different level of precision (or range) and stores a corresponding subset of the quantized keys that have been quantized to be represented in the different format to have the different level of precision (or range). In some of these implementations, while repeatedly performing iterations of the process 200, the system can move at least one of the quantized key vectors from a memory corresponding to a higher level of precision to a lower level of precision (or vice versa). Correspondingly, the precision of the at least one of the quantized key vectors is reduced (or increased).

[0106] The system receives, at the attention head of the attention layer, an embedded representation of an output token at the particular output position (step 202). The output token can be the most recently generated output token in the output sequence. That is, the output token can be an output token that occupies an immediately previous output position in the plurality of output positions of the output sequence that immediately precede the particular output position.

[0107] For the first attention layer in the sequence of attention layers the embedded representation can be generated by the system from the last token in a current input sequence by processing the most recently generated output token using an embedding layer.

[0108] The last token in a current input sequence can either be the most recently generated output token or, when the particular output position is the first output position and hence no output token has been generated yet, the last input token included in the input sequence.

[0109] For each subsequent attention layer in the sequence of attention layers the embedded representation can be the updated embedded representation of the output token that is generated by the preceding layer, e.g., a preceding attention layer or a different type of neural network layer.

[0110] The system applies, at the attention head of the attention layer, a (learned) key linear transformation to the embedded representation to generate an original key for the particular output position (step 204). The original key represents an intermediate computation result of the attention mechanisms to be quantized.

[0111] In some implementations, the key linear transformation that is used to generate the original key is specific to the attention head. Thus different attention heads within the attention layer apply different key linear transformations to generate different original keys.

[0112] In some other implementations, e.g., where a multi-query attention mechanism is used, the key linear transformation that is used to generate the original key is shared across the multiple attention heads within the attention layer. Thus the key linear transformation need only be applied once to generate an original key, and the original key are shared across the different attention heads within the attention layer. The system quantizes, at the attention head of the attention layer, the original key for the particular output position to generate a quantized original key for the particular output position (step 206). The quantized original key is represented using a second format that has a lower precision, lower range, or both than a first format used to represent the original key. The quantized original key can then be stored in the one or more memory devices.

[0113] As noted above, the original key is a vector that includes one or more elements. Accordingly, the system can apply any of a variety of known vector quantization techniques to quantize the original key to generate the quantized original key. A few example ways of applying vector quantization techniques are described in FIG. 3 A and FIG. 3B.

[0114] FIG. 3 A a flow diagram of sub-steps of a first implementation of step 206 of the process of FIG. 2.

[0115] The system determines a quantization range for the original key (step 302). In some implementations, the vector quantization technique can be a dynamic vector quantization technique. In these implementations, the quantization range can be a scalar range that is determined based on a distribution of values of the one or more elements in the original key. Thus the quantization ranges may vary from one original key to another original key.

[0116] In general, the quantization range can be determined based on the range between the maximum value and the minimum value of the one or more elements in the original key. For example, an upper limit of the quantization range can be determined based on a maximal absolute value of the one or more elements in the original key, while a lower limit of the quantization range can be determined based on a minimal absolute value of the one or more elements in the original key.

[0117] In some other implementations, the vector quantization technique can be a static vector quantization technique, where the quantization range remains fixed from one original key to another original key. For example, the quantization range can be a predetermined scalar range.

[0118] The system maps each of the one or more elements in the original key to a corresponding mapped (quantized) element within the quantization range based on applying a scale factor to the element of the original key (step 304). That is, the system replaces each element in the original key by the corresponding, mapped (quantized) element.

[0119] The scale factor is how a larger set of values to be quantized (the values of the one or more elements in the original key) is mapped to a corresponding smaller set of values (the values of the one or more mapped elements in the quantized original key), by scaling the values in the larger set of values to correspond to a mapped value in the smaller set of values. Here “small” and “large” refer to the cardinality of the set of values: a first set is “smaller” or “larger” than a second set if it has fewer or more elements respectively. Quantization can be considered as transforming an element of the larger set of values to a corresponding element of the smaller set of values.

[0120] Calculating the scale factor may vary depending on the distribution of values of the one or more elements in the original key and the size of the quantization range and, in some implementations, on the quantization technique used (e.g., symmetric quantization vs. asymmetric quantization).

[0121] In some implementations where a static vector quantization technique is applied, and the quantization range is a predetermined scalar range, the system can account for any values of the corresponding, mapped elements of the one or more elements in the original key that fall out of the predetermined scalar range by using value clipping.

[0122] Specifically, the system can determine that a particular element in the original key falls out of (is either above or below) a predetermined maximum range that corresponds to the quantization range, and in response, set the corresponding, mapped element of the particular element to be equal to either the upper limit or the lower limit of the quantization range.

[0123] FIG. 3B a flow diagram of sub-steps of an alternative implementation of step 206 of the process of FIG. 2.

[0124] For each of the one or more elements in the original key, the system identifies, from among a set of anchor points, a nearest anchor point to the element of the original key as the identified anchor point that corresponds to the element in the original key (step 312). Each of the anchor points is a scalar value.

[0125] The system can generate the set of anchor points based on the distribution of values of the one or more elements in the original key. The system can generate the set of anchor points so that the values of most or all (e.g., no less than a predetermined portion) of the one or more elements in the original key are reasonably close (e.g., within a predetermined threshold distance) to at least one anchor point.

[0126] For each of the one or more elements in the original key, the system generates the corresponding, mapped (quantized) element by setting the element in the original key to the identified nearest anchor point (step 314). That is, the system replaces the value of each element in the original key by the value of the corresponding anchor point. In general, the system can determine the anchor point corresponding to a value of an element to be the anchor point with the value that is closest to the value of the element. The system generates, at the attention head of the attention layer, an updated embedded representation of the output token at the particular output position based on applying an attention mechanism (step 208). To apply the attention mechanism, each attention head uses (i) the quantized original key for the particular output position and (ii) the plurality of quantized keys in the one or more memories. The original key may be discarded (e.g. deleted from memory, so it is no longer available to the computer system) following step 206, e.g. before step 208. In a variation, the original key may be used in step 208, in which case it may be discarded after step 208, e.g. before starting the next performance of method 200 to generate the next output token.

[0127] In some implementations, the attention neural network is configured to compute the outputs of the attention mechanism using values represented in the first format in which the original key is represented, as discussed further below with reference to FIG. 4.

[0128] FIG. 4 is a flow diagram of sub-steps 402-408 of an implementation of step 208 of the process of FIG. 2. Each attention head of each attention layer in the attention neural network can perform an iteration of steps 402-408 in parallel with each other to compute an output of the attention mechanism.

[0129] The system retrieves, for the attention head of each attention layer, the plurality of quantized keys from the one or more memory devices (step 402). The plurality of quantized keys have been generated and stored in the one or more memory devices prior to the time step that corresponds to the particular output position, i.e., prior to the attention neural network receives a current input sequence that includes the most recently generated (selected) output token, i.e., the output token that occupies an immediately previous position that immediately precedes the particular output position.

[0130] The system dequantizes the plurality of quantized keys to generate a plurality of dequantized keys (step 404). That is, the system converts a plurality of quantized keys that are each represented in the second format back to a plurality of dequantized keys that are each represented in the first format.

[0131] The way the system maps the plurality of quantized keys to the plurality of dequantized keys depends on the quantization technique used initially to quantize the original keys, e.g., by applying an inverse of a scale factor (used in quantization) to each of the one or more elements in the quantized key to a corresponding mapped (dequantized) element, or by mapping each of the one or more elements in the quantized key to a corresponding mapped (dequantized) element in accordance with a mapping between the set of anchor points and the one or more elements in the quantized key. The dequantized element may in fact the corresponding anchor point but represented in a numerical format having a larger number of bits than the quantized key.

[0132] Optionally, the system dequantizes the quantized original key for the particular output position to generate a dequantized original key (step 406). That is, the system converts the quantized original key that is represented in the second format back to a dequantized original key that is represented in the first format. Step 406 is optional because, in some implementations, the original keys may still be available to the system, and hence dequantizing the quantized original key to generate the dequantized original key is not needed.

[0133] The system computes an output of the attention mechanism based on the dequantized original key (or the original key in case it is still available) for the particular output position and the plurality of dequantized keys (step 408). The output of the attention mechanism can be used to generate the updated embedded representation of the output token at the particular output position.

[0134] The attention head can apply any of a variety of variants of query-key -value (QKV) attention (as described above in FIG. 1) using the dequantized original key for the particular output position, together with the plurality of dequantized keys which are generated from the plurality of quantized keys obtained from the one or more memory devices, to generate the output of the attention mechanism.

[0135] In some other implementations, each attention head of the attention neural network is configured to compute the outputs of the attention mechanism using the second format in which the quantized original key is represented, as discussed further below with reference to FIG. 5.

[0136] FIG. 5 is a flow diagram of sub-steps 502-504 of another implementation of step 208 of the process of FIG. 2. Each attention head of each attention layer in the attention neural network can perform an iteration of steps 502-504 in parallel with each other to compute an output of the attention mechanism.

[0137] The system retrieves, for the attention head of each attention layer, the plurality of quantized keys from the one or more memory devices (step 502). The plurality of quantized keys have been generated and stored in the one or more memory devices prior to the time step that corresponds to the particular output position, i.e., prior to the attention neural network receives a current input sequence that includes the most recently generated (selected) output token, i.e., the output token that occupies an immediately previous position that immediately precedes the particular output position. The system computes an output of the attention mechanism based on the quantized original key for the particular output position and the plurality of quantized keys (step 504). The output of the attention mechanism can be used to generate the updated embedded representation of the output token at the particular output position.

[0138] The attention head can apply any of a variety of variants of query-key -value (QKV) attention (as described above in FIG. 1) using the quantized original key for the particular output position, together with the plurality of quantized keys retrieved from the one or more memory devices, to generate the output of the attention mechanism.

[0139] In either implementation, when there is a single attention head, the output of the attention head is used as the updated embedded representation of the attention layer. When there are multiple attention heads, the attention layer then combines the outputs of the multiple attention heads, e.g., by concatenating the outputs and, optionally, processing the concatenated outputs through a linear layer, to generate the updated embedded representation.

[0140] When the attention layer is not the last layer in the attention neural network, the system can then provide the updated embedded representation as the embedded representation to the subsequent neural network layer in the attention neural network, which can for example be a feed-forward layer, a layer normalization layer configured to apply layer normalization to the updated embedded representation, or another attention layer.

[0141] By repeatedly performing the process 200 for all of the attention layers in the attention neural network and then by processing at least the updated embedded representation generated by the last attention layer in the sequence of multiple attention layers included in the attention neural network using one or more output layer(s), the system can generate the score distribution for the time step, which can be used to select the output token at the particular output position in the plurality of output positions in the output sequence.

[0142] While the description of FIG. 2 discusses quantized key vectors, it should be noted that the discussion is equally applicable to quantized value vectors.

[0143] That is, in some implementations, the one or more memory devices also stores a plurality of quantized values. The plurality of quantized values have been generated and stored in the one or more memory devices prior to the time step that corresponds to the particular output position, i.e., prior to the attention neural network receives a current input sequence that includes the most recently selected output token, i.e., the output token that occupies an immediately previous position that precedes the particular output position.

[0144] In some of these implementations, the plurality of quantized values are quantized to have a same format (and thus have a same precision and range as) the plurality of quantized keys, whereas, in others of these implementations, the plurality of quantized values are quantized to have a different format (and thus have a different precision and / or range) from the plurality of quantized keys.

[0145] In these implementations, at each attention head of each attention layer, the system can generate an original value for the particular output position, quantize the original value to generate a quantized original value, and then apply the attention mechanism using (i) the quantized original key for the particular output position, (ii) the plurality of quantized keys obtained from the one or more memories, (iii) the quantized original value for the particular output position, and (iv) the plurality of quantized values obtained from the one or more memories. Optionally, if the original key and / or value is still available (i.e. has / have not yet been discarded) it / they may be used in place of the quantized key and / or value.

[0146] In some implementations, the system does not use and store quantized queries in the one or more memory devices, and the attention mechanism applied at each attention head of each attention layer uses an original query that is generated from the embedded representation of the output token at the particular output position.

[0147] The process 200 can be performed as part of predicting an output sequence based on an input sequence for which the desired output sequence, i.e., the output sequence that should be generated by the system for the input sequence, is not known.

[0148] The process 200 can also be performed as part of processing training input sequences derived from a set of training data, i.e., input sequences derived from a set of training input sequences for which the output sequences that should be generated by the system is known, in order to train the attention neural network to determine trained values for the parameters of the attention neural network.

[0149] The system can repeatedly perform the process 200 on input sequences selected from a set of training data as part of a language model training technique to train the attention neural network, e.g., a gradient descent with backpropagation training technique that uses a conventional optimizer, e.g., stochastic gradient descent, RMSprop, or Adam optimizer, to determine updates to the values of the parameters of the attention neural network based on optimizing an objective function that is appropriate for the task that the attention neural network is configured to perform.

[0150] For example, the task can be a language modeling task, e.g., a task that requires predicting, given a current sequence of tokens, the next token that follows the current sequence in the training data, and the objective function can be an unsupervised or selfsupervised objective function, e.g., a maximum-likelihood objective function, that evaluates a difference between (i) the score distribution generated by the attention neural network for the time step that can be used to select the output token and (ii) a ground truth score distribution that defines the next token (e.g., a score distribution that assigns one to the next token, and zeros to remaining tokens in the vocabulary of tokens).

[0151] In some implementations, as part of the training, the system trains the attention neural network to learn to generate, at each attention head of each attention layer, key vectors having one or more elements that each fall within a predetermined maximum range that corresponds to a quantization range.

[0152] The predetermined maximum range can be a tunable hyperparameter which is set prior to the training of the attention neural network, while the training of the attention neural network can use the same objective function as mentioned above. Thus, in these implementations, once the attention neural network is trained, the system can more practically use the static vector quantization technique as mentioned above during inference without significant performance degradation.

[0153] In some implementations, as part of the training, the system trains the attention neural network to learn to generate, from a given input sequence, and based on applying an attention mechanism using quantized keys and / or quantized values, an output sequence that is not too much different from another output sequence (i.e. the output sequence and the other sequence meet a similarity criterion) that would be generated from the same given input sequence but based on applying an attention mechanism without quantizing the keys or values.

[0154] FIG. 6 is a flow diagram of an example process 600 for training an attention neural network on a training database of tuplets. Each tuplet includes a training input sequence and a corresponding training output sequence. For convenience, the process 600 will be described as being performed by a system of one or more computers located in one or more locations. For example, a system, e.g., the quantized inference system 100 of FIG. 1 or another training system, appropriately programmed in accordance with this specification can perform the process 600.

[0155] The system can repeatedly perform iterations of the process 600 on successive tuplets or groups of tuplets in the training database of tuplets. By repeatedly performing iterations of the process 600, the system can determine the trained values of the parameters of the attention neural network, e.g., including the trained values of the parameters of each attention head of each attention layer in the attention neural network that define the key, query, and value linear transformations. The system provides the training input sequences as inputs to the first of the attention layers in the attention neural network (step 602).

[0156] The system generates provisional output sequences of output tokens based on corresponding outputs of the last of the attention layers (step 604). As part of generating the provisional output sequence, the system quantizes the original keys generated by the attention layers based on the training input sequences to quantized original keys.

[0157] The system updates, i.e., modifies, the values of the parameters defining at least one of the key, query, or value linear transformations of each attention head of each attention layer in the attention neural network to increase a measure of the similarity of the provisional output sequences and the corresponding training output sequences (step 606). Updating the values of those parameters can involve converting the quantized original keys generated based on the training input sequences into a higher precision format than the original key.

[0158] For example, the system can do this based on optimizing an objective function that includes a term that measures, for each training input sequence, a distance between a provisional output sequence and a corresponding training output sequence. The distance can be computed in any appropriate distance measure, e.g., an edit distance in a text space, or a Euclidean distance or cosine similarity in an embedding space.

[0159] This specification uses the term “configured” in connection with systems and computer program components. For a system of one or more computers to be configured to perform particular operations or actions means that the system has installed on it software, firmware, hardware, or a combination of them that in operation cause the system to perform the operations or actions. For one or more computer programs to be configured to perform particular operations or actions means that the one or more programs include instructions that, when executed by data processing apparatus, cause the apparatus to perform the operations or actions.

[0160] Embodiments of the subject matter and the functional operations described in this specification can be implemented in digital electronic circuitry, in tangibly-embodied computer software or firmware, in computer hardware, including the structures disclosed in this specification and their structural equivalents, or in combinations of one or more of them. Embodiments of the subject matter described in this specification can be implemented as one or more computer programs, i.e., one or more modules of computer program instructions encoded on a tangible non transitory storage medium for execution by, or to control the operation of, data processing apparatus. The computer storage medium can be a machine- readable storage device, a machine-readable storage substrate, a random or serial access memory device, or a combination of one or more of them. Alternatively or in addition, the program instructions can be encoded on an artificially generated propagated signal, e.g., a machine-generated electrical, optical, or electromagnetic signal, that is generated to encode information for transmission to suitable receiver apparatus for execution by a data processing apparatus.

[0161] The term “data processing apparatus” refers to data processing hardware and encompasses all kinds of apparatus, devices, and machines for processing data, including by way of example a programmable processor, a computer, or multiple processors or computers. The apparatus can also be, or further include, special purpose logic circuitry, e.g., an FPGA (field programmable gate array) or an ASIC (application specific integrated circuit). The apparatus can optionally include, in addition to hardware, code that creates an execution environment for computer programs, e.g., code that constitutes processor firmware, a protocol stack, a database management system, an operating system, or a combination of one or more of them.

[0162] A computer program, which may also be referred to or described as a program, software, a software application, an app, a module, a software module, a script, or code, can be written in any form of programming language, including compiled or interpreted languages, or declarative or procedural languages; and it can be deployed in any form, including as a stand alone program or as a module, component, subroutine, or other unit suitable for use in a computing environment. A program may, but need not, correspond to a file in a file system. A program can be stored in a portion of a file that holds other programs or data, e.g., one or more scripts stored in a markup language document, in a single file dedicated to the program in question, or in multiple coordinated files, e.g., files that store one or more modules, sub programs, or portions of code. A computer program can be deployed to be executed on one computer or on multiple computers that are located at one site or distributed across multiple sites and interconnected by a data communication network.

[0163] In this specification, the term “database” is used broadly to refer to any collection of data: the data does not need to be structured in any particular way, or structured at all, and it can be stored on storage devices in one or more locations. Thus, for example, the index database can include multiple collections of data, each of which may be organized and accessed differently.

[0164] Similarly, in this specification the term “engine” is used broadly to refer to a software-based system, subsystem, or process that is programmed to perform one or more specific functions. Generally, an engine will be implemented as one or more software modules or components, installed on one or more computers in one or more locations. In some cases, one or more computers will be dedicated to a particular engine; in other cases, multiple engines can be installed and running on the same computer or computers.

[0165] The processes and logic flows described in this specification can be performed by one or more programmable computers executing one or more computer programs to perform functions by operating on input data and generating output. The processes and logic flows can also be performed by special purpose logic circuitry, e.g., an FPGA or an ASIC, or by a combination of special purpose logic circuitry and one or more programmed computers.

[0166] Computers suitable for the execution of a computer program can be based on general or special purpose microprocessors or both, or any other kind of central processing unit. Generally, a central processing unit will receive instructions and data from a read only memory or a random access memory or both. The essential elements of a computer are a central processing unit for performing or executing instructions and one or more memory devices for storing instructions and data. The central processing unit and the memory can be supplemented by, or incorporated in, special purpose logic circuitry. Generally, a computer will also include, or be operatively coupled to receive data from or transfer data to, or both, one or more mass storage devices for storing data, e.g., magnetic, magneto optical disks, or optical disks. However, a computer need not have such devices. Moreover, a computer can be embedded in another device, e.g., a mobile telephone, a personal digital assistant (PDA), a mobile audio or video player, a game console, a Global Positioning System (GPS) receiver, or a portable storage device, e.g., a universal serial bus (USB) flash drive, to name just a few.

[0167] Computer readable media suitable for storing computer program instructions and data include all forms of non volatile memory, media and memory devices, including by way of example semiconductor memory devices, e.g., EPROM, EEPROM, and flash memory devices; magnetic disks, e.g., internal hard disks or removable disks; magneto optical disks; and CD ROM and DVD-ROM disks.

[0168] To provide for interaction with a user, embodiments of the subject matter described in this specification can be implemented on a computer having a display device, e.g., a CRT (cathode ray tube) or LCD (liquid crystal display) monitor, for displaying information to the user and a keyboard and a pointing device, e.g., a mouse or a trackball, by which the user can provide input to the computer. Other kinds of devices can be used to provide for interaction with a user as well; for example, feedback provided to the user can be any form of sensory feedback, e.g., visual feedback, auditory feedback, or tactile feedback; and input from the user can be received in any form, including acoustic, speech, or tactile input. In addition, a computer can interact with a user by sending documents to and receiving documents from a device that is used by the user; for example, by sending web pages to a web browser on a user’s device in response to requests received from the web browser. Also, a computer can interact with a user by sending text messages or other forms of message to a personal device, e.g., a smartphone that is running a messaging application, and receiving responsive messages from the user in return.

[0169] Data processing apparatus for implementing machine learning models can also include, for example, special-purpose hardware accelerator units for processing common and compute-intensive parts of machine learning training or production, i.e., inference, workloads.

[0170] Machine learning models can be implemented and deployed using a machine learning framework, e.g., a TensorFlow framework or a JAX framework.

[0171] Embodiments of the subject matter described in this specification can be implemented in a computing system that includes a back end component, e.g., as a data server, or that includes a middleware component, e.g., an application server, or that includes a front end component, e.g., a client computer having a graphical user interface, a web browser, or an app through which a user can interact with an implementation of the subject matter described in this specification, or any combination of one or more such back end, middleware, or front end components. The components of the system can be interconnected by any form or medium of digital data communication, e.g., a communication network. Examples of communication networks include a local area network (LAN) and a wide area network (WAN), e.g., the Internet.

[0172] The computing system can include clients and servers. A client and server are generally remote from each other and typically interact through a communication network. The relationship of client and server arises by virtue of computer programs running on the respective computers and having a client-server relationship to each other. In some embodiments, a server transmits data, e.g., an HTML page, to a user device, e.g., for purposes of displaying data to and receiving user input from a user interacting with the device, which acts as a client. Data generated at the user device, e.g., a result of the user interaction, can be received at the server from the device.

[0173] While this specification contains many specific implementation details, these should not be construed as limitations on the scope of any invention or on the scope of what may be claimed, but rather as descriptions of features that may be specific to particular embodiments of particular inventions. Certain features that are described in this specification in the context of separate embodiments can also be implemented in combination in a single embodiment. Conversely, various features that are described in the context of a single embodiment can also be implemented in multiple embodiments separately or in any suitable subcombination.

[0174] Moreover, although features may be described above as acting in certain combinations and even initially be claimed as such, one or more features from a claimed combination can in some cases be excised from the combination, and the claimed combination may be directed to a subcombination or variation of a subcombination.

[0175] Similarly, while operations are depicted in the drawings and recited in the claims in a particular order, this should not be understood as requiring that such operations be performed in the particular order shown or in sequential order, or that all illustrated operations be performed, to achieve desirable results. In certain circumstances, multitasking and parallel processing may be advantageous. Moreover, the separation of various system modules and components in the embodiments described above should not be understood as requiring such separation in all embodiments, and it should be understood that the described program components and systems can generally be integrated together in a single software product or packaged into multiple software products.

[0176] Particular embodiments of the subject matter have been described. Other embodiments are within the scope of the following claims. For example, the actions recited in the claims can be performed in a different order and still achieve desirable results. As one example, the processes depicted in the accompanying figures do not necessarily require the particular order shown, or sequential order, to achieve desirable results. In some cases, multitasking and parallel processing may be advantageous.

Claims

CLAIMS1. A computer-implemented method of generating an output sequence that comprises an output token at each of a plurality of output positions using a neural network, wherein the neural network comprises a plurality of attention layers, wherein each attention layer of one or more of the plurality of attention layers comprises one or more attention heads, and wherein the method comprises, for each attention head of the one or more attention heads and for a particular output position in the plurality of output positions: maintaining, in one or more memory devices, a plurality of quantized keys that were generated prior to receiving an embedded representation of an output token at the particular output position; receiving the embedded representation of the output token at the particular output position; applying a key linear transformation to the embedded representation to generate an original key for the particular output position; quantizing the original key to generate a quantized original key that is represented using a lower precision format than the original key; and generating an updated embedded representation of the output token at the particular output position based on applying an attention mechanism based on the quantized original key for the particular output position and the plurality of quantized keys in the one or more memory devices.

2. The method of claim 1, wherein applying the attention mechanism based on the quantized original key for the particular output position and the plurality of quantized keys in the one or more memory devices comprises: retrieving the plurality of quantized keys from the one or more memory devices; dequantizing the plurality of quantized keys to generate a plurality of dequantized keys; dequantizing the quantized original key for the particular output position to generate a dequantized original key; and computing an output of the attention mechanism based on the dequantized original key for the particular output position and the plurality of dequantized keys.

3. The method of claim 1, wherein applying the attention mechanism based on the quantized key for the particular output position and the plurality of quantized keys in the one or more memory devices comprises: retrieving the plurality of quantized keys from the one or more memory devices; computing an output of the attention mechanism based on the quantized original key for the particular output position and the plurality of quantized keys.

4. The method of any one of claims 1-3, wherein the original key is a vector comprising one or more elements, and wherein quantizing the original key to generate the quantized original key comprises, for each element in the original key: identifying, from among a set of anchor points, a nearest anchor point to the element as the identified anchor point that corresponds to the element in the original key; and setting the element of the original key to the identified nearest anchor point.

5. The method of any one of claims 1-3, wherein the original key is a vector comprising one or more elements, and wherein quantizing the original key to generate the quantized original key comprises: determining a quantization range based on the one or more elements in the original key; and mapping each element of the original key to a corresponding, mapped element within the quantization range based on applying a scale factor.

6. The method of claim 5, wherein mapping the element in the original key to the mapped element within the quantization range comprises: determining an upper limit of the quantization range based on a maximal absolute value of the one or more elements in the original key.

7. The method of any one of claims 5-6, wherein mapping the element in the original key to the mapped element within the quantization range comprises: determining that a particular element of the original key is either above or below the quantization range; and in response, clipping the mapped element of the particular element to be equal respectively to an upper limit or a lower limit of the quantization range.

8. The method of any one of claims 1-7, further comprising: generating, based at least on the updated embedded representation of the output token at the particular output position that has been generated by the attention head, a time step output for a current generation time step that defines a score distribution over a set of possible tokens; and selecting, using the time step output, a selected token from the set of possible tokens as the output token at the particular output position.

9. The method of claim 8, when dependent on claim 5, which is performed by a Transformer neural network, and wherein the respective quantization range for the attention head is predetermined prior to generating the output sequence by using the Transformer neural network, and wherein the method further comprises: determining an update to parameter values of the Transformer neural network based at least on the selected token.

10. The method of any one of claims 1-9, wherein the one or more memory devices comprise a plurality of memory devices that each correspond to a different level of precision and that each store a corresponding subset of the quantized keys that have been quantized to have the corresponding level of precision.

11. The method of any one of claims 1-10, wherein applying the attention mechanism comprises applying the attention mechanism also using (i) an original query that is generated from the embedded representation of the output token at the particular output position and (ii) a plurality of quantized values that were generated prior to receiving the embedded representation of the output token at the particular output position.

12. The method of claim 11, wherein the plurality of quantized values are quantized to have a same precision as the plurality of quantized keys.

13. The method of claim 11, wherein the plurality of quantized values are quantized to have a different precision from the plurality of quantized keys.

14. A computer-implemented method of generating an output sequence that comprises an output token at each of a plurality of output positions, based on an input sequence of input tokens at each of a plurality of input positions corresponding to the output positions, the method being performed by using a plurality of attention layers arranged in a series with the output of each attention layer except a first of the attention layers being an input to the next attention layer of the series, each of the attention layers being arranged to receive, for each of the input positions, a corresponding attention layer input vector based on a corresponding input token, and to generate from the attention layer input vector a corresponding attention layer output vector for the input position, the output token for the corresponding output position being based on the attention layer output vector for a last of the attention layers: for each attention layer, the generating of the attention layer output vector being performed by, for each of one or more attention heads of the attention layer, generating one or more transformed vectors for the input position by applying respective linear transformations to the attention layer input vector, the one or more transformed vectors including a key vector for the input position; quantizing the key vector for the input position to generate a quantized key vector for the input position, the quantized key vector being represented using a lower precision format than the key vector, and storing the quantized key vector in memory; and for each input position except the first, obtaining the attention layer output vector for the input position using (i) the key vector for the input position or the quantized key vector for the input position, and (ii) a corresponding said quantized key vector for one or more previous input positions retrieved from the memory.

15. The method of claim 14 in which the one or more transformed vectors further include a query vector for the input position, and a value vector for the input position, and the attention layer output vector for the input position is obtained by: calculating a weight for the input position based on the product of the query vector for the input position and the key vector for the input position, calculating weights for the previous input positions based on the product of respective products of (i) the query vector for the input position and (i) the key vector for previous input positions,forming the attention layer output vector based on a sum, weighted by the corresponding weights, of the value vector for the input position and the value vectors of the previous input positions.

16. The method of claim 14 in which the one or more transformed vectors further include a query vector for the input position, and a value vector for the input position, and the attention layer output vector for the input position is obtained by: calculating a weight for the input position based on the product of the query vector for the input position and the key vector for the input position, calculating weights for the previous input positions based on the product of respective products of (i) the query vector for the input position and (i) the key vector for previous input positions, quantizing the value vector for the input position to generate a quantized value vector for the input position, the quantized value vector being represented using a lower precision format than the value vector for the position, and storing the quantized value vector in memory, forming the attention layer output vector based on a sum, weighted by the corresponding weights, of the value vector or quantized value vector for the input position and the quantized value vectors of the previous input positions.

17. The method of claim 16, wherein the plurality of quantized value vectors are quantized to have a same precision as the plurality of quantized key vectors.

18. The method of claim 16, wherein the plurality of quantized value vectors are quantized to have a different precision from the plurality of quantized key vectors.

19. The method of any one of claims 14-18, wherein quantizing the key vector for the input position comprises: determining a quantization range based on the plurality of elements in the key vector; and mapping each of the plurality of elements in the key vector to a corresponding, mapped element within the quantization range based on applying a scale factor.

20. The method of claim 19, wherein quantizing the key vector for the input position comprises: determining an upper limit of the quantization range based on a maximal absolute value of the one or more elements in the key vector for the input position.

21. The method of any one of claims 19-20, wherein quantizing the key vector for the input position comprises: determining that a particular element of the key vector for the position is either above or below a predetermined range that corresponds the quantization range; and in response, setting a corresponding element of the quantized key vector as respectively an upper limit or a lower limit of the predetermined range.

22. The method of any one of claims 14-21, further comprising the last of the attention layers: generating, based at least on the attention layer output vectors generated by the corresponding one or more attention heads, a time step output that defines a score distribution over a set of possible tokens; and selecting, using the score distribution, a token from the set of possible tokens as the output token at the particular output position.

23. The method of any one of claims 14-22, wherein the one or more memories comprise a plurality of memories that each correspond to a different level of precision and that each store a corresponding subset of the quantized key vectors, the method further comprising moving at least one of the quantized key vectors from a memory corresponding to a higher level of precision to a lower level of precision and correspondingly reducing the precision of the at least one of the quantized key vectors.

24. The method of any of claims 14-23 which is performed following a training phase of generating parameter values defining the linear transformations based on a training database of tuplets composed of training input sequences and corresponding training output sequences, the training process comprising, for successive tuplets or groups of tuplets, providing the training input sequences as inputs to the first of the attention layers; generating provisional output sequences of output tokens based on corresponding outputs of the last of the attention layers; andmodifying the parameter values defining at least one of the linear transformations to increase a measure of the similarity of the provisional output sequences and the corresponding training output sequences; wherein the generating of the provisional output sequences comprises quantizing key vectors generated based on the training input sequences.

25. The method of claim 24 in which the step of modifying the parameter values comprises converting the quantized key vectors generated based on the training input sequences into a higher precision format than the corresponding quantized key vectors.

26. The method of any preceding claim in which the input sequence represents an input image and the output sequence is a network output which is the result of performing a task on the input sequence, wherein the task is image classification and the network output is scores for each of a set of object categories, with each score representing an estimated likelihood that the image contains an image of an object belonging to the category; the task is image embedding generation and the network output is a numeric embedding of the input image; the task is object detection and the network output identifies locations in the input image at which particular types of objects are depicted; or the task is image segmentation and the network output assigns each pixel of the input image to a category from a set of categories.

27. A 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 perform the operations of the method of any one of claims 1-26.

28. One or more computer storage media storing instructions that when executed by one or more computers cause the one or more computers to perform the operations of the method of any one of claims 1-26.

Images