Processing context inputs using compression tokens in neural networks

Abstract

Methods, systems, and apparatus, including computer programs encoded on a computer storage medium, for generating a response to a query using a context input. One of the methods includes: receiving a context input comprising a sequence of context tokens; generating a modified sequence of context tokens, comprising modifying the sequence of context tokens to insert a plurality of compression tokens into the sequence of context tokens, wherein the compression tokens of the plurality of compression tokens are interspersed throughout the sequence of context tokens; processing the modified sequence of context tokens using a neural network to, for each of a plurality of layers of the neural network, generate a respective layer output embedding of each of the compression tokens; and generating a response to a query that relates to the context input using the respective layer output embeddings of the compression tokens for the plurality of layers.

Description

[0001] Attorney Docket No.: 45288-0626WO1

[0002] PROCESSING CONTEXT INPUTS USING COMPRESSION TOKENS IN NEURAE NETWORKS CROSS-REFERENCE TO RELATED APPLICATIONS

[0003] [1] This application claims the benefit under 35 U.S.C. § 119(e) of the filing date of U.S. Patent Application No. 63 / 751,829, which was filed on January 30. 2025, and which is incorporated here by reference.

[0004] BACKGROUND

[0005] [2] This specification relates to processing inputs using neural networks.

[0006] [3] 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 another layer in the network, e.g., 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 weights.

[0007] SUMMARY

[0008] [4] This specification describes a method for generating a response to a query' that is related to a context input by compressing information represented in the context input into compression tokens.

[0009] [5] A system can receive a context input including a sequence of context tokens. The system generates a modified sequence of context tokens by modifying the sequence of context tokens to insert a plurality of compression tokens into the sequence of context tokens. The compression tokens represent the information encoded in the context tokens using a smaller number of tokens than the context tokens. In this way, the compression tokens “compress” the information encoded in the context tokens. The compression tokens or the plurality of compression tokens can be interspersed throughout the sequence of context tokens.

[0010] [6] The system processes the modified sequence of context tokens using a neural network to, for each of a plurality of layers of the neural network, generate a respective layer output embedding of each of the compression tokens. The system generates a response to a query that relates to the context input using the respective layer output embeddings of the compression tokens for the plurality of layers.Attorney Docket No.: 45288-0626WO1

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

[0012] [8] This specification describes a computer-implemented method of using a neural network to process a context input that includes a sequence of context tokens in order to generate a response to a query' that relates to the context input. The techniques described herein include making use of one or more of a number of methods by which the processing of long contexts can be improved, thus allowing for the processing of long contexts with increased effectiveness and reduced consumption of computational resources.

[0013] [9] Existing techniques for processing context inputs using neural networks struggle to process long contexts. For example, for existing generative neural network architectures, processing long contexts requires excessive memory and processing costs. For example, performing self-attention operations using existing generative neural network architectures has quadratic complexity with respect to sequence length. Modifying the architecture of a neural network can be computationally difficult. Such modifications can also make it difficult to integrate neural networks with existing systems. Additionally, there is a high computational cost associated with processing long contexts. Thus, computer systems on which neural networks are implemented can lack the resources to be able to handle processing long contexts.

[0014]

[0010] The techniques described in this specification allow for more efficient and effective processing of long contexts by a neural network. The techniques include, in response to receiving a context input including a sequence of context tokens, generating a modified sequence of context tokens by modifying the sequence of context tokens to insert a plurality of compression tokens into the sequence of context tokens. The compression tokens compress relevant information represented by the context tokens in the sequence of context tokens into a smaller number of tokens. Instead of processing the entire sequence of context tokens, e.g., to generate a response to a query relating to the context, a neural network can process only the compression tokens. Thus, the use of compression tokens can result in more efficient processing of long contexts because a smaller number of tokens can be processed, e.g., to generate responses to queries. For example, compressing the context into compression tokens can reduce the computational overhead and memory’ footprint of a key -value (KV) cache used to perform attention operations to generate a response to a query. For example, this reduction of the computational overhead and memory' footprint in turn can enable a system employing the disclosed techniques to process significantly longer contexts than traditional models while maintaining high throughput.Attorney Docket No.: 45288-0626WO1

[0015]

[0011] Implementing the use of compression tokens in a neural network also requires minimal modifications to the architecture of the neural network, therefore reducing the computational difficulty of modifying neural network architectures, while also allowing for seamless integration of the modified neural network into existing systems.

[0016]

[0012] The techniques described in this specification avoid discrepancies between the layer, in a sequence of layers of a neural network, at which a model expects to receive information and the layer in the sequence at which it receives the information. For example, some existing techniques that employ compression tokens struggle to process long contexts, in spite of using compression tokens, due to such discrepancies. The techniques described herein include auto-regressively generating an output sequence of tokens using a neural network by processing, for each token in the output sequence and for a self-attention layer of the neural network, respective layer output embeddings of compression tokens that are generated by an earlier layer that is before the immediately preceding layer to the self-attention layer. This avoids the discrepancy between the layer at which a model expects to receive information and the layer at which it receives the information because, according to these techniques, the model receives information from the earlier layer. Avoiding such discrepancies can allow for more efficient processing of long contexts by neural networks, reducing the use of computational resources in processing long contexts. Avoiding such discrepancies can improve responses to queries generated by a system employing the disclosed techniques, e.g., by increasing a likelihood that the system maintains correct information flow between layers of the neural network used to generate the responses.

[0017]

[0013] The techniques described in this specification include using separate sets of parameters to generate layer output embeddings for compression tokens and for context tokens. This allows the neural network to effectively generate layer output embeddings of the compression tokens, while still maintaining high performance of downstream tasks by the model by allowing parameters corresponding to the layer output embeddings of the context tokens to be modified separately from parameters corresponding to the layer output embeddings of the compression tokens.

[0018]

[0014] The techniques described in this specification include using neural networks with at least one self-attention layer that applies self-attention in which each compression token attends to a different subset of context tokens from at least one of the other compression tokens. In some implementations of the techniques described herein, this can be accomplished by interspersing compression tokens through a sequence of context tokens. For example, at least one compression token is not allowed to attend to all of the context tokens. Limiting the contextAttorney Docket No.: 45288-0626WO1

[0019] tokens that can be attended to by the compression tokens improves the effectiveness with which the neural network processes long contexts using compression tokens. For example, limiting the context tokens that can be attended to by the compression tokens can reduce the amount of information from the context tokens that is lost via compression into the smaller number of compression tokens, since each compression token only encodes the limited amount of information stored in the limited number of context tokens to which it attends. Limiting the context tokens that can be attended to by the compression tokens can increase a likelihood that the information encoded in the compression tokens is relevant, accurately represents the information encoded in the context tokens, or both. For example, in performing self-attention operations for each compression token, a system employing the disclosed techniques can be more likely to identify relevant information in a limited number of context tokens, e.g., than if self-attention is performed over a larger number of context tokens. With more effective processing of long contexts by neural networks, the use of computational resources in processing long contexts can be reduced.

[0020]

[0015] A system employing the disclosed techniques can effectively perform tasks that require processing a large amount of high-dimensional data. Examples of such tasks are described below.

[0021]

[0016] 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.

[0022] BRIEF DESCRIPTION OF THE DRAWINGS

[0023]

[0017] FIG. 1 A is a block diagram of an example response generation system.

[0024]

[0018] FIG. IB is a block diagram of an example implementation of a response generation system that generates modified sequences of context tokens.

[0025]

[0019] FIG. 1C is a block diagram of an example implementation of a response generation system that uses average pooling.

[0026]

[0020] FIG. 2 is a diagram illustrating attention masks of one or more self-attention layers of a neural network used to generate responses to queries.

[0027]

[0021] FIG. 3 is a diagram illustrating layer-wise information flow7across a plurality of neural netw ork layers of a neural network in a response generation system.

[0028]

[0022] FIG. 4 is a flowchart of an example process for generating a response to a query using compression tokens.Attorney Docket No.: 45288-0626WO1

[0029]

[0023] FIG. 5 is a flowchart of an example process for generating a response to a query' using context compression via average pooling.

[0030]

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

[0031] DETAILED DESCRIPTION

[0032]

[0025] Methods and systems are provided for processing context inputs using compression tokens in neural networks to facilitate efficient long-context processing. Large language models and multimodal models are often required to process substantial context lengths, such as extensive document histories, web pages, or lengthy audio and video data. Standard attention mechanisms in decoder-only transformer architectures exhibit quadratic computational complexity relative to sequence length, which can lead to high latency and excessive resource consumption. Furthermore, model effectiveness can decline as context length increases. The technology described herein addresses these challenges by compressing contextual information into a reduced number of activations, thereby facilitating more efficient response generation without requiring significant architectural modifications to the underlying model.

[0033]

[0026] A system employing the disclosed techniques can effectively perform tasks that require processing a large amount of high-dimensional data. Examples of such tasks are described below.

[0034]

[0027] For example, a system employing the techniques described herein can process context input that includes context tokens representing a video. The system can receive an input query that, for example, includes a question about the video. The system can process the input video tokens using the techniques described herein to generate a response to the question related to the video.

[0035]

[0028] As another example, a system employing the techniques described herein can process context input that includes context tokens representing an audio signal. The system can receive an input query that, for example, includes a question about the audio signal. The system can process the input audio tokens using the techniques described herein to generate a response to the question related to the audio signal.

[0036]

[0029] A system employing the techniques described herein can perform a variety of machine learning tasks. Some examples of these machine learning tasks follow.Attorney Docket No.: 45288-0626WO1

[0037]

[0030] In any of the implementations described herein, the neural network may be deployed as part of a chat bot, dialogue agent, or other software tool that receives inputs from users and provides outputs in response to the received input, e.g., as part of a conversation or dialogue. In these implementations, the input sequences received by the neural network are (generated from) user inputs and the output sequences generated by the neural network can be used to generate responses to the user inputs.

[0038]

[0031] In implementations the neural network may be configured as, or include, a generative (large) language model or a multi-modal model, e.g., a visual and language model, to perform these example machine learning tasks.

[0039]

[0032] In some cases, the neural network is a neural network that is configured to perform an image processing task, i.e., receive an input image and to process the input image to generate a network output for the input image. For example the input sequence may comprise tokens representing pixel values for pixels in regions or patches of the image. For example, the task may be image classification and the output generated by the neural network for a given image may be 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. As another example, the task can be image embedding generation and the output generated by the neural network can be a numeric embedding of the input image. As yet another example, the task can be object detection and the output generated by the neural network can identify locations in the input image at which particular types of objects are depicted. As yet another example, the task can be image segmentation and the output generated by the neural network can assign each pixel of the input image to a category from a set of categories. In some other cases, the neural network is a neural network that is configured to perform an image generation task, where the input is a conditioning input and the output is a sequence of intensitv value inputs for the pixels of an image.

[0040]

[0033] As one example, the task may be a neural machine translation task. For example, if the input to the neural network is a sequence of text, e.g., a sequence of words, phrases, characters, or word pieces, in one language, the output generated by the neural network may be a translation of the sequence of text into another language, i.e., a sequence of text in the other language that is a translation of the input sequence of text. The vocabulary for the input tokens may be words, wordpieces or characters of the first language, and the vocabulary for the output tokens may be words, wordpieces or characters of the other language. As a particular example, the task may be a multi-lingual machine translation task, where a single neural network is configured to translate between multiple different source language - target language pairs. InAttorney Docket No.: 45288-0626WO1

[0041] this example, the source language text may be augmented with an identifier that indicates the target language into which the neural network should translate the source language text.

[0042]

[0034] Some implementations may be used for automatic code generation. For example the input tokens may represent words, wordpieces or characters in a first natural language and the output tokens may represent instructions in a computer programming or markup language, or instructions for controlling an application program to perform a task, e.g., build a data item such as an image or web page.

[0043]

[0035] As another example, the task may be an audio processing task. For example, if the input to the neural network is a sequence representing a spoken utterance, the output generated by the neural network may be a score for each of a set of pieces of text, each score representing an estimated likelihood that the piece of text is the correct transcript for the utterance, e.g. a speech to text task. As another example, if the input to the neural network is a sequence representing a spoken utterance, the output generated by the neural network can indicate whether a particular word or phrase (“hotword”) was spoken in the utterance. As another example, if the input to the neural network is a sequence representing a spoken utterance, the output generated by the neural network can be a classification of the spoken utterance into one of a plurality of categories, for example an identity of the natural language in which the utterance was spoken.

[0044]

[0036] As another example, the task can be a natural language processing or understanding task, e.g., an entailment task, a paraphrase task, a textual similarity task, a sentiment task, a sentence completion task, a grammaticality task, and so on, that operates on a sequence of text in some natural language.

[0045]

[0037] As another example, the task can be a text to speech task, where the input is text in a natural language or features of text in a natural language and the network output is a spectrogram, a waveform, or other data defining audio of the text being spoken in the natural language.

[0046]

[0038] As another example, the task can be a health prediction task, where the input is a sequence derived from electronic health record data for a patient and the output is a prediction that is relevant to the future health of the patient, e.g., a predicted treatment that should be prescribed to the patient, the likelihood that an adverse health event will occur to the patient, or a predicted diagnosis for the patient. Such electronic health data may, for example, comprise one or more sequences of physiological data taken from a patient, with the output being a corresponding prediction that relates to those sequences of data. Examples of physiological data and a corresponding prediction include: blood glucose measurements, with the predictionAttorney Docket No.: 45288-0626WO1

[0047] being a predicted future blood glucose measurement or the prediction of a hyper- or hypoglycemic event; a heart rate, with the prediction being the presence or absence of a heart condition, or a future cardiac event; blood pressure measurements, with the prediction being the risk of a future heart condition; or the like.

[0048]

[0039] As another example, the task can be a text generation task, where the input is a sequence of text, and the output is another sequence of text, e.g., a completion 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 first sequence of text. As another example, the input to the text generation task can be an input other than text, e.g., an image, and the output sequence can be text that describes the input.

[0049]

[0040] In some implementations the context input represents data to be compressed, e.g., image data, text data, audio data, or any other type of data; and the output sequence a compressed version of the data. The input and output tokens may each comprise any representation of the data to be compressed / compressed data, e.g., symbols or embeddings generated / decoded by a respective neural network. In some complementary implementations the input sequence represents compressed data and the output sequence represents a decompressed version of the data, e.g., image data, text data, audio data, or any other type of data.

[0050]

[0041] As another example, the task can be an agent control task, where the input is a sequence of observations or other data characterizing states of an environment and the output defines an action to be performed by the agent in response to the most recent data in the sequence. The agent can be, e.g., a real-world or simulated robot, a control system for an industrial facility, or a control system that controls a different kind of agent. The observations may comprise sensor data captured by sensors associated with (e.g.. part of) the agent, for example visual data, LIDAR data, sonar data, agent configuration data (e.g.. joint angles), agent orientation data, or the like.

[0051]

[0042] In some implementations, the environment is a real-world environment, the agent is a mechanical (or electro-mechanical) agent interacting with the real-world environment, e.g., a robot or an autonomous or semi-autonomous land, air, or sea vehicle operating in or navigating through the environment, and the actions are actions taken by the mechanical agent in the real-world environment to perform the task. For example, the agent may be a robot interacting with the environment to accomplish a specific task, e.g., to locate or manipulate an object of interest in the environment or to move an object of interest to a specified location in the environment or to navigate to a specified destination in the environment.Attorney Docket No.: 45288-0626WO1

[0052]

[0043] In these implementations, the observations may include, e.g., one or more of: images, object position data, and sensor data to capture observations as the agent interacts with the environment, for example sensor data from an image, distance, or position sensor or from an actuator. For example in the case of a robot, the observations may include data characterizing the current state of the robot, e.g., one or more of: joint position, joint velocity, joint force, torque or acceleration, e.g., gravity-compensated torque feedback, and global or relative pose of an item held by the robot. In the case of a robot or other mechanical agent or vehicle the observations may similarly include one or more of the position, linear or angular velocity, force, torque or acceleration, and global or relative pose of one or more parts of the agent. The observations may be defined in 1, 2 or 3 dimensions, and may be absolute and / or relative observations. The observations may also include, for example, sensed electronic signals such as motor current or a temperature signal; and / or image or video data for example captured by a camera or a LIDAR sensor, e.g., data from sensors of the agent or data from sensors that are located separately from the agent in the environment.

[0053]

[0044] In these implementations, the actions may be control signals to control the robot or other mechanical agent, e.g., torques for the joints of the robot or higher-level control commands, or the autonomous or semi-autonomous land, air, sea vehicle, e.g., torques to the control surface or other control elements, e.g., steering control elements of the vehicle, or higher-level control commands. The control signals can include for example, position, velocity, or force / torque / accel eration data for one or more joints of a robot or parts of another mechanical agent. The control signals may also or instead include electronic control data such as motor control data, or more generally data for controlling one or more electronic devices within the environment the control of which has an effect on the observed state of the environment. For example in the case of an autonomous or semi-autonomous land or air or sea vehicle the control signals may define actions to control navigation, e.g., steering, and movement e.g., braking and / or acceleration of the vehicle.

[0054]

[0045] In some implementations the environment is a simulation of the above-described real-world environment, and the agent is implemented as one or more computers interacting with the simulated environment. For example, a system implementing the neural network may be used to select actions in the simulated environment during training or evaluation of the system and, after training, or evaluation, or both, are complete, the action selection policy may be deployed for controlling a real-world agent in the particular real- world environment that was the subject of the simulation. This can avoid unnecessary wear and tear on and damage to the real-world environment or real-world agent and can allow the control neural network to beAttorney Docket No.: 45288-0626WO1

[0055] trained and evaluated on situations that occur rarely or are difficult or unsafe to re-create in the real-world environment. For example the system may be partly trained using a simulation of a mechanical agent in a simulation of a particular real-world environment, and afterwards deployed to control the real mechanical agent in the particular real-world environment. Thus in such cases the observations of the simulated environment relate to the real-world environment, and the selected actions in the simulated environment relate to actions to be performed by the mechanical agent in the real-world environment.

[0056]

[0046] In some implementations, as described above, the agent may not include a human being (e.g., it is a robot). Conversely, in some implementations the agent comprises a human user of a digital assistant such as a smart speaker, smart display, or other device. Then the information defining the task can be obtained from the digital assistant, and the digital assistant can be used to instruct the user based on the task.

[0057]

[0047] For example, a system implementing the neural network may output to the human user, via the digital assistant, instructions for actions for the user to perform at each of a plurality' of time steps. The instructions may for example be generated in the form of natural language (transmitted as sound and / or text on a screen) based on actions chosen by the system. The system chooses the actions such that they contribute to performing a task. A monitoring system (e.g., a video camera system) may be provided for monitoring the action (if any) which the user actually performs at each time step, in case (e.g., due to human error) it is different from the action which the system instructed the user to perform. Using the monitoring system the system can determine whether the task has been completed. The system may identify actions which the user performs incorrectly with more than a certain probability. If so, when the system instructs the user to perform such an identified action, the system may warn the user to be careful. Alternatively or additionally, the system may leam not to instruct the user to perform the identified actions, i.e., ones which the user is likely to perform incorrectly.

[0058]

[0048] More generally, the digital assistant instructing the user may comprise receiving, at the digital assistant, a request from the user for assistance and determining, in response to the request, a series of tasks forthe user to perform, e.g.. steps or sub-tasks of an overall task. Then for one or more tasks of the series of tasks, e.g., for each task, e.g., until a final task of the senes the digital assistant can be used to output to the user an indication of the task, e.g., step or subtask, to be performed. This may be done using natural language, e.g., on a display and / or using a speech synthesis subsystem of the digital assistant. Visual, e.g., video, and / or audio observations of the user performing the task may be captured, e.g.. using the digital assistant. A system as described above may then be used to determine whether the user has successfullyAttorney Docket No.: 45288-0626WO1

[0059] achieved the task, e.g., step or sub-task, i.e., from the answer as previously described. If there are further tasks to be completed the digital assistant may then, in response, progress to the next task (if any) of the series of tasks, e.g., by outputting an indication of the next task to be performed. In this way the user may be led step-by-step through a series of tasks to perform an overall task. During the training of the neural network, training rew ards may be generated, e.g., from video data representing examples of the overall task (if corpuses of such data are available) or from a simulation of the overall task.

[0060]

[0049] In a further aspect there is provided a digital assistant device including a system as described above. The digital assistant can also include a user interface to enable a user to request assistance and to output information. In implementations this is a natural language user interface and may comprise a keyboard, voice input-output subsystem, and / or a display. The digital assistant can further include an assistance subsystem configured to determine, in response to the request, a series of tasks for the user to perform. In implementations this may comprise a generative (large) language model, in particular for dialog, e.g., a conversation agent such as Sparrow (Glaese et al. arXiv:2209.14375) or Chinchilla (Hoffmann et al. arXiv:2203.15556). The digital assistant can have an observation capture subsystem to capture visual and / or audio observations of the user performing a task; and an interface for the abovedescribed language model neural network (which may be implemented locally or remotely). The digital assistant can also have an assistance control subsystem configured to assist the user. The assistance control subsystem can be configured to perform the steps described above, for one or more tasks, e g., of a series of tasks, e.g., until a final task of the series. More particularly the assistance control subsystem and output to the user an indication of the task to be performed, capture, using the observation capture subsystem, visual or audio obser ations of the user performing the task, determine from the above-described answer whether the user has successfully achieved the task. In response the digital assistant can progress to a next task of the series of tasks and / or control the digital assistant, e.g., to stop capturing observations.

[0061]

[0050] As another example, the task can be a genomics task, where the input is a sequence representing a fragment of a DNA sequence or other molecule sequence and the output is either an embedding of the fragment for use in a downstream task, e.g., by making use of an unsupervised learning technique on a data set of DNA sequence fragments, or an output for the downstream task. Examples of downstream tasks include promoter site prediction, methylation analysis, predicting functional effects of non-coding variants, and so on.

[0062]

[0051] In some cases, the machine learning task is a combination of multiple individual machine learning tasks, i.e., the system is configured to perform multiple different individualAttorney Docket No.: 45288-0626WO1

[0063] machine learning tasks, e.g., two or more of the machine learning tasks mentioned above. For example, the system can be configured to perform multiple individual natural language understanding tasks, with the network input including an identifier for the individual natural language understanding task to be performed on the network input.

[0064]

[0052] In some cases, the machine learning task is a multi-modal processing task that requires processing multi-modal data. In general, multi-modal data is a combination of two or more different types of data, e.g., two or more of audio data, image data, text data, or graph data. As one example the multi-modal data may comprise audio-visual data, comprising a combination of pixels of an image or of video and audio data representing values of a digitized audio waveform. As another example the multi-modal data may comprise a combination of i) text data representing text in a natural language and ii) pixels of an image or of video or audio data representing values of an audio waveform. Optionally, but not necessarily, the different types of data may represent the same or overlapping objects using the different modalities (types), and when processing multi-modal data the data may be mapped into a common embedding space.

[0065]

[0053] As a particular example, the task is a multi-modal processing task that requires processing both text and image inputs, so that the neural network includes both a computer vision neural network and a text processing neural network. That is, the target output to be generated by the computer vision neural network for a given image depends on one or more outputs generated by the text processing neural network for one or more corresponding text inputs (and vice versa). Examples of such tasks include open-vocabulary image classification, open-vocabulary object detection, image captioning, text-based image search, image-based retrieval, and so on.

[0066]

[0054] As some further examples a multi-modal processing task can involve processing a text input comprising a sequence of text or audio data representing values of an audio waveform, e.g., instantaneous amplitude data or time-frequency domain data, or an image or video (or encoded versions of these inputs) to generate the network output. The network output may comprise any form of output appropriate to the task performed. For example the network output may comprise text in a natural or computer language that defines a result of the task, e.g., for tasks such as image captioning, video or audio question answering (answering a natural language question about a visual or audio input), or object detection or instance segmentation. For example in a video or audio question answering task the question can define an information content extraction task, to extract information from the content of the video or audio, or the question can define a reasoning task such as a predictive reasoning task (e.g. '‘what wouldAttorney Docket No.: 45288-0626WO1

[0067] happen next?”), a counterfactual reasoning task (e.g. “what would happen if..?”), or a causal reasoning task (e.g. “why did X happen?”). The network output can provide an answer in any convenient form, e.g. tokens representing natural language. An input to the system may be obtained from a sensor sensing the real world, e.g. a condition or characteristic of the real world. For example the video or audio may be captured from the real-world. The network output can then provide an answer, e.g. in natural language, to a question asked about the real-world input.

[0068]

[0055] Also or instead the response to the query may comprise data defining an image, video or audio object, e.g., as specified by the input (e.g. by a natural language description of one or more characteristics of the object), e.g., in a generative task. As a further alternative the network output may comprise non-textual action selection data for selecting an action to be performed by an agent controlled by the network output, e g. as described above, e.g. in response to an input that includes a natural language description of a physical or other task to be performed by the agent. As another example the network output may also or instead define an intermediate step to be performed during the task, e.g., a call to a software API for a software tool that is used when performing the task; the input may then receive an output from the software tool that is used to generate a final network output that performs the task.

[0069]

[0056] More generally, the multi-modal processing task may correspond to any of the tasks previously described for any of the types of data making up the multi-modal combination. For example, an accuracy of the previously described tasks may be increased when the task is applied to multi-modal data combining the data for which the task has been previously described and another type of data. For example detection or classification of an object or event may be improved when data of multiple different types (modalities) is processed.

[0070]

[0057] FIG. 1 A is a block diagram of an example response generation system 100. The response generation 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.

[0071]

[0058] The system 100 is configured to process a context input including a sequence of context tokens 102 and a query 101 to generate a response 110 to the query 101 that relates to the context input. The context input can include a sequence of context tokens 102.

[0072]

[0059] The system 100 includes a compression engine 104 and a response generation engine 108. Each of the compression engine 104 and the response generation engine 108 can include a neural network, e.g., the same neural network. The neural network can be a generative neural network, such as a decoder-only transformer. The neural network can include aAtorney Docket No.: 45288-0626WO1

[0073] sequence of self-atention layers. The neural network can be configured to process a sequence of tokens, which can include the sequence of context tokens 102, query tokens for the query 101, tokens representing the response 110, compression tokens which are defined below, or any combination of these. For example, each self-atention layer of the neural network can process the sequence of tokens to update embeddings for the tokens in the sequence. Each self-atention layer can thus generate respective layer output embeddings representing the tokens in the sequence, where the respective layer output embeddings are the embeddings for the tokens that have been updated by the layer. The layer output embedding for a token generated by a layer can be an output generated by the layer for an embedding corresponding to, e.g.. occupying the same position as, the token in the sequence. In some implementations, one or more layers of the neural network can apply average pooling to tokens in the sequence to generate compressed layer output embeddings.

[0074]

[0060] The system 100 receives a query 101 and the sequence of context tokens 102 as an input. The sequence of context tokens 102 can include a sequence of vectors representing various data modalities, such as text, audio, or video data. Each context token in the sequence 102 can be a vector or other ordered collection of numerical values, e.g., a vector of floating point or other numerical values, that represents data related to a context for the query 101. For example, a context token can represent a sub-word unit in a natural language string that describes data related to the context, a patch of an image that depicts data related to the context, or a segment of an audio waveform that includes audio related to the context.

[0075]

[0061] The query 101 can be a query that relates to the data represented by the sequence of context tokens 102. For example, the sequence 102 can represent an image, and the query 101 can include a question about the image. For example, the sequence 102 can represent an audio signal, and the query 101 can include a question about the audio signal. In some implementations, the query 101 includes a plurality of query tokens.

[0076]

[0062] During a compression phase, the system processes the query 101 and the sequence of context tokens 102 using the neural network to generate a compressed representation of the query 101 and the sequence of context tokens 102. During a prediction phase that follows the compression phase, the system uses the neural network to, conditioned on the compressed representation, generate the response 110 to the uery 101.

[0077]

[0063] The compressed representation is compressed, e.g., as compared to an original representation of the sequence of context tokens and query. In some implementations, the compressed representation includes, for each of the self-atention layers of the neural network, layer output embeddings generated by the self-atention layer for the compressionAttorney Docket No.: 45288-0626WO1

[0078] tokens. Each of the compression tokens or compressed layer output embeddings corresponds to a plurality of context tokens and encodes information related to the context that is represented by the corresponding context tokens. The compression tokens or compressed layer output embeddings collectively can represent the data represented by the sequence of context tokens using a smaller number of tokens or embeddings than the number of context tokens in the sequence of context tokens. In order to generate the response to the query, the system can process only the layer output embeddings for the compression tokens or compressed layer output embeddings, e.g., and not the context tokens. In this way, the system can generate the response to the query' by processing fewer tokens.

[0079]

[0064] In some implementations, the system 100 uses a Key-Value (KV) cache included in the system 100. Various example implementations described in this specification comprise compressing the KV cache, such that the KV cache comprises a compressed KV cache. A KV cache is a specialized memory' buffer or data structure configured to store intermediate activation tensors — specifically the key (k) and value (v) vectors — generated during the inference of a transformer-based neural network. For example, as each self-attention layer of the neural network processes the sequence of tokens, an attention head of the self-attention layer generates key vectors ("‘keys”) and value vectors (“values”) for each of the tokens in the sequence. The attention head generates the keys and values for each of the tokens from the layer output embedding for the token. In these implementations, the compressed representation can include, for each attention head of each attention layer, the keys and values for each of the compression tokens that are generated by the attention head from the layer output embeddings of the compression tokens. In these implementations, the compressed representation can include, for each attention head of each attention layer, the keys and values for each of the compressed layer output embeddings generated by the attention head from the compressed layer output embeddings.

[0080]

[0065] In more detail, during the compression phase, the compression engine 104 is configured to process the sequence of context tokens 102 to generate compression outputs 106. The compression outputs 106 can include layer output embeddings generated by the self-attention layers of the neural network included in the compression engine 104. In some implementations, during the compression phase, the layer output embeddings can be generated by the neural network in response to processing a modified sequence of context tokens that is based on the sequence of context tokens 102, but additionally includes compression tokens interspersed throughout the sequence. The layer output embeddings can be embeddings for the compression tokens that have been updated by one or more of the selfAtorney Docket No.: 45288-0626WO1

[0081] atention layers. Such implementations are described in further detail below with reference to FIG. IB. In some implementations, during the compression phase, the layer output embeddings can be generated by the neural network in response to processing a segmented sequence of the context tokens. The segmented sequence of context tokens can be generated by partitioning the sequence of context tokens 102 into a plurality of segments. One or more of the layers of the neural network can apply average pooling to each segment of the sequence of context tokens to generate a respective compressed layer output embedding for the segment.. Such implementations are described in further detail below with reference to FIG. 1C.

[0082]

[0066] In some implementations, the compression outputs 106 can include keys and values for the layer output embeddings from the KV cache. In some implementations, the compression outputs 106 can include keys and values for the compressed layer output embeddings from the KV cache. For example, the compression outputs 106 can include keys and values rather than the layer output embeddings or the compressed layer output embeddings.

[0083]

[0067] The response generation engine 108 processes the compression outputs 106 and the query 101 using the neural network to generate a response 110 to the uery 101. For example, the neural network can process the compression outputs 106 and the query 101 to generate the response 110 by processing a sequence of tokens including query tokens included in the query 101, the compression tokens, and a sequence of output tokens. Processing the sequence of tokens can mean performing atention operations over embeddings corresponding to the tokens. For example, the embeddings corresponding to the compression tokens can be the compression outputs 106. In some implementations, processing the sequence of tokens can mean, for one or more tokens in the sequence, accessing keys and values from the KV cache that were computed for the one or more tokens by preceding layers of neural network, and updating embeddings for other tokens in the sequence using the accessed keys and values. For example, the neural network can generate the response 110 using techniques such as those described below with reference to FIG. 5. The neural network can apply self- atention using an atention mask, such as one of the atention masks described below with reference to FIG. 2.

[0084]

[0068] By processing the compression outputs 106 instead of the sequence of context tokens 102, the response generation engine 108 can generate the response 110 with reduced computational overhead, lower memory usage for key-value (KV) caches, improved throughput, or any combination of these. For example, the compression outputs can includeAttorney Docket No.: 45288-0626WO1

[0085] embeddings for tokens that compress relevant information represented by the context tokens, thereby enabling the response generation engine 108 to process fewer tokens to generate the response 110. The response 110 can include a sequence of tokens representing a natural language answer to a query, a predicted action for an agent, or any other output relevant to the context input.

[0086]

[0069] FIG. IB is a block diagram of an example implementation of a response generation system 120 that generates modified sequences of context tokens. The system 120 can represent an example implementation of the response generation system 100 of FIG. 1 A. For example, the response generation engine 108 included in the system 120 can represent the response generation engine 108 of FIG. 1A. The modified sequence generation engine 122 can be included in or represent one or more parts of the compression engine 104 of FIG. 1A. The neural network 130 can be the neural network included in the compression engine 104 and the response generation engine 108 of FIG. 1A.

[0087]

[0070] The modified sequence generation engine 122 is configured to process the sequence of context tokens 102 that is received by the system 120. The modified sequence generation engine 122 can process the sequence 102 using a compression token generation engine 124. The compression token generation engine 124 is configured to process the sequence of context tokens 102 to generate a plurality7of compression tokens 126, and to insert the plurality of compression tokens 126 into the sequence 102.

[0088]

[0071] Each of the compression tokens 126 corresponds to a plurality7of context tokens in the sequence of context tokens 102. In some implementations, each of the compression tokens 126 can correspond to a plurality7of context tokens that come before the compression token in the sequence of context tokens.

[0089]

[0072] For example, each of the compression tokens 126 can be an initial embedding. By processing each compression token, the neural network 130 can modify the initial embedding to generate a respective layer output embedding for the compression token that represents the information related to the context that is represented by the corresponding context tokens. The respective layer output embedding of each of the compression tokens 126 can represent information of more than one other context token, thereby “compressing’7the information of the more than one other context token into a single embedding. In this way, the number of compression tokens 126 is less than the number of context tokens included in the sequence of context tokens 102.

[0090]

[0073] The compression token generation engine 124 can insert each of the compression tokens 126 in any of a variety of locations within the sequence of context tokens 102. In someAttorney Docket No.: 45288-0626WO1

[0091] implementations, the compression token generation engine 124 can intersperse the compression tokens 126 throughout the sequence 102. For example, the compression token generation engine 124 can intersperse the compression tokens 126 throughout the sequence of context tokens 102 such that not all of the compression tokens 126 are adjacent to one another within the sequence of context tokens 102. For example, at least one compression token can be separated from at least one other compression token by at least one context token of the sequence of context tokens. In some examples, the compression tokens can be interspersed uniformly throughout the sequence of context tokens, such that the same number of context tokens separate each compression token. As a result of processing the sequence of context tokens 102 using the compression token generation engine 124, the modified sequence generation engine 122 generates a modified sequence of context tokens 128 that includes both the original context tokens and the inserted compression tokens 126.

[0092]

[0074] The neural network 130 is configured to process the modified sequence of context tokens 128 to generate layer output embeddings 132. As described above with reference to FIG. 1A, the neural network 130 can include a plurality of layers, such as self-attention layers and feed-forward layers. In some implementations, the neural network 130 includes a sequence of self-attention layers.

[0093]

[0075] The layer output embeddings 132 include a respective layer output embedding for each compression token 126 and for each layer of the neural network 130. The respective layer output embedding for a given compression token that is generated by a layer can be the respective layer output embedding generated by the layer in response to processing the layer output embedding for the given compression token that was generated by the preceding layer. Processing the layer output embedding for the given compression token that was generated by the preceding layer can include updating the layer output embedding by performing attention operations over a sequence of tokens that includes the modified sequence of context tokens 128. The respective layer output embedding generated by the layer can then be the updated layer output embedding that results from performing the attention operations.

[0094] Performing attention operations over a sequence of tokens can include performing attention operations over embeddings that represent the tokens. The neural network 130 can perform the self-attention operations according to an attention mask. Examples of attention masks that can be used by the neural network 130 are described with reference to FIG. 2 below.

[0095]

[0076] For example, a first layer of the neural network 130 can process each compression token 126 to generate a respective layer output embedding for the compression token. A second layer of the neural network 130 can then process the respective layer outputAttorney Docket No.: 45288-0626WO1

[0096] embedding for each compression token that was generated by the first layer to generate a respective layer output embedding for the compression token for the second layer. A third layer of the neural network 130 can then process the respective layer output embedding for each compression token that was generated by the second layer to generate a respective layer output embedding for the compression token for the third layer. In this way, each subsequent layer of the neural network 130, e.g., following the first layer, generates a respective layer output embedding for each compression token in response to processing the respective layer output embedding for the compression token that was generated by the preceding layer.

[0097]

[0077] Each layer output embedding can be a vector or other ordered collection of numerical values e.g., a vector of floating point or other numerical values. Each layer output embedding represents information encoded in the context tokens corresponding to the compression token from which the layer output embedding was generated. Because each compression token corresponds to more than one context token as described above, each layer output embedding represents information encoded in more than one context token. Thus, collectively the layer output embeddings 132 can represent the information encoded in the sequence of context tokens 102 using a smaller number of embeddings than the number of context tokens included in the sequence of context tokens 102. In this way, the layer output embeddings 132 can be considered a compressed representation of the sequence of context tokens 102.

[0098]

[0078] In some implementations, the compressed representation of the sequence of context tokens 102 includes keys and values from a KV cache that were computed for the layer output embeddings 132 by one or more self-attention layers of the neural network 130. In these implementations, the keys and values corresponding to the layer output embeddings 132 can include a smaller number of keys and values than the keys and values that were computed for the sequence of context tokens 102. In this way, the keys and values for the layer output embeddings 132 can be considered a compressed representation of the keys and values for the sequence of context tokens 102.

[0099]

[0079] As described below, the system 120 can process the compressed representation of the sequence of context tokens 102 to generate the response 110 to the query 101. For example, the system 120 can generate the response 110 from the layer output embeddings 132. In some implementations, generating the response 110 from the layer output embeddings 132 can involve directly including the layer output embeddings in the compressed representation of the sequence of context tokens 102, and generating the response 110 using the compressed representation. In some implementations, generating the response 110 from the layer output embeddings 132 can involve including keys and values generated from the layer outputAtorney Docket No.: 45288-0626WO1

[0100] embeddings by one or more self-atention layer and stored in a KV cache in the compressed representation of the sequence of context tokens 102, and generating the response 110 using the compressed representation.

[0101]

[0080] By processing the layer output embeddings 132, the system 120 can generate the response 110 by processing a smaller number of embeddings or a smaller number of keys and values, e.g., as compared to the number of tokens or the number of keys and values that would be processed if the system 120 were to process the entire sequence of context tokens 102 or the keys and values for the entire sequence of context tokens 102 to generate the response 110.

[0102]

[0081] In some implementations, the layer output embeddings 132 include a respective layer output embedding for each context token in the sequence 102 and for each layer of the neural network 130. For example, a first layer can generate a respective layer output embedding for each context token in response to processing the context token. Each layer other than the first layer can generate a respective layer output embedding for each context token by processing a respective layer output embedding for the context token generated by the preceding layer, as described above in reference to the compression tokens. The neural network 130 can generate the layer output embeddings for the context tokens by performing self-atention operations, as described above. In some implementations, the layer output embeddings for the context tokens can be keys and values that were computed for the context tokens.

[0103]

[0082] In some implementations, the neural network 130 uses a first set of parameters to process the context tokens and generate the respective layer output embeddings for the context tokens. The neural network 130 uses a second, different set of parameters to process the compression tokens and generate the respective layer output embeddings for the compression tokens. The use of different respective parameters for each of the context tokens and the compression tokens can allow the neural network to achieve high levels of the compression functionality without degrading the performance of its base language modeling capabilities.

[0104]

[0083] Training the neural network 130 can include updating the parameters of the neural network 130. In some implementations, the rates at which each of the first set of parameters and the second set of parameters are updated can be different. For example, the second set of parameters can be updated during training using a higher learning rate, such as a rate 100 times higher than that used for the first set. Additional description of how the neural network 130 can be trained is provided below with reference to FIG. 4.Atorney Docket No.: 45288-0626WO1

[0105]

[0084] The response generation engine 108 processes the query 101 and the layer output embeddings 132 to generate the response 110 to the query’ 101.

[0106]

[0085] In some implementations, the response generation engine 108 generates the response 110 using the neural network 130. For example, the response generation engine 108 can use the neural network 130 to autoregressively generate an output sequence of tokens, as described in further detail below with reference to FIG. 4. For example, the response generation engine 108 can generate the output sequence of tokens by processing a sequence of tokens including the output sequence of tokens. Processing the sequence of tokens can include using self-attention layers of the neural network 130 to perform self-attention operations over the sequence. As described above, performing self-attention operations over the sequence can include performing self-attention operations over embeddings corresponding to tokens in the sequence. In some implementations, performing self-attention operations over the sequence can include accessing keys and values for one or more tokens in the sequence from the KV cache and using the accessed keys and values to update embeddings for other tokens in the sequence. The response generation engine 108 can perform the self-attention operations according to an attention mask, such as one of the attention masks described below with reference to FIG. 2.

[0107]

[0086] In some implementations, autoregressively generating the output sequence of tokens includes, for each token in the output sequence, shifting or "‘offsetting’' the compression tokens or layer output embeddings that are processed by’ a particular self-attention layer of the neural network 130. For example, a particular self-attention layer can process query tokens and the output sequence of tokens from an immediately preceding layer such that the query' tokens attend to compression tokens from an earlier layer that is before the immediately- preceding layer. For example, a particular self-attention layer can update embeddings for tokens in the output sequence of tokens using keys and values for query' tokens generated by an immediately preceding layer and keys and values for compression tokens generated by an earlier layer that is before the immediately preceding layer. This offsetting process can help to improve the response generated by the response generation engine 108 by increasing a likelihood that the neural network 130 maintains correct information flow. For example, the offsetting process can enable a given self-attention layer of the neural network to receive the compressed context information from the layer expected by its weights, avoiding delays in information availability' that might otherwise occur, e.g., if the given self-attention layer were to receive the information from the immediately preceding layer. The offsetting process is described in further detail below with reference to the implementation (c) of FIG. 3.Attorney Docket No.: 45288-0626WO1

[0108]

[0087] FIG. 1 C is a block diagram of an example implementation of a response generation system 150 that uses average pooling. The system 150 can represent an example implementation of the response generation system 100 of FIG. 1A. For example, the response generation engine 108 included in the system 150 can represent the response generation engine 108 of FIG. 1A. The segment identification engine 152 included in the system 150 can be included in or represent one or more parts of the compression engine 104 of FIG. 1A. The neural network 130 can be the neural network included in the compression engine 104 and the response generation engine 108 of FIG. 1A.

[0109]

[0088] The segment identification engine 152 is configured to process the sequence of context tokens 102 that is received by the system 150. In response to processing the sequence of context tokens 102, the segment identification engine 152 identifies a plurality of segments of the sequence of context tokens that each include a different proper subset of the sequence. As used in this specification, a proper subset of the sequence of context tokens is a set of tokens included in the sequence of context tokens such that at least one token in the sequence of context tokens is not included in the set. The proper subsets identified by the segment identification engine 152 can be non-overlapping. That the proper subsets are nonoverlapping can mean that, for any pair of the identified proper subsets, each token included in a first proper subset of the pair is not included in a second proper subset of the pair, and each token included in the second proper subset is not included in the first proper subset.

[0110]

[0089] In some implementations, the segment identification engine 152 can identify the segments of the sequence 102 based on a compression rate. The compression rate can be a pre-determined rate that indicates a number of tokens to be included in the identified segments, which can in turn determine the number of tokens to be ‘‘compressed’' into a representative compression token, as described below. For example, if the sequence 102 has 1000 tokens and the compression rate is 10, the segment identification engine 152 can identify' 100 segments of 10 tokens each. In response to identifying the segments of the sequence 102, the segment identification engine 152 can generate a segmented sequence of context tokens 154. The segmented sequence of context tokens 154 includes the original sequence of context tokens 102 and data identifying the plurality of segments of the sequence 102.

[0111]

[0090] The neural network 130 of the system 150 is configured to process the segmented sequence of context tokens 154 to generate a respective layer output embedding 132 for each of the context tokens included in the sequence and for each layer of the plurality of layers of the neural network 130. As described above with reference to FIG. 1A, the respective layerAttorney Docket No.: 45288-0626WO1

[0112] output embedding for a given context token that is generated by a layer can be the respective layer output embedding generated by the layer in response to processing the layer output embedding for the given context token that was generated by the preceding layer. The neural network 130 can include or represent the neural network 130 of FIG. IB, e.g., including a plurality7of layers.

[0113]

[0091] In some implementations, for one or more layers of the neural network 130, generating the layer output embeddings for the context tokens includes performing selfattention operations over the sequence of context tokens, as described above. The neural network 130 can perform the self-attention operations according to an attention mask.

[0114] Examples of attention masks that can be used by the neural network 130 are described with reference to FIG. 2 below.

[0115]

[0092] In the system 150, the neural network 130 is further configured to, for each of the plurality of layers and for each of the plurality of segments of the sequence of context tokens, generate a compressed layer output embedding for the segment by applying average pooling on the respective layer output embeddings for the layer and for the context tokens in the segment. Applying average pooling to the respective layer output embeddings can include generating the compressed layer output embedding such that the compressed layer output embedding represents information that includes information represented by all of the context tokens in the segment. For example, if the layer output embeddings are vectors, applying average pooling to the respective layer output embeddings can include determining an average over the vectors of each of the constituent elements of the vectors. The compressed layer output embedding can be a vector of which the elements are the determined averages of the constituent elements of the layer output embeddings.

[0116]

[0093] In this way, for each layer of the neural network and for each segment, the neural network 130 generates a plurality of compressed layer output embeddings 134. The compressed layer output embeddings 134 include a compressed layer output embedding corresponding to each of the segments of the segmented sequence of context tokens 154. Each compressed layer output embedding can be a vector that represents the average of the layer output embeddings of the context tokens in the corresponding segment.

[0117]

[0094] The response generation engine 108 processes the compressed layer output embeddings 134 using the neural network 130 to generate the response 110. For example, the response generation engine 108 can use the neural network 130 to autoregressively generate an output sequence of tokens, as described in further detail below with reference to FIG. 4.Attorney Docket No.: 45288-0626WO1

[0118] For example, the response generation engine 108 can generate the output sequence of tokens using self-attention layers of the neural network 130 to perform self-attention operations over a sequence including the output sequence of tokens. The response generation engine 108 can perform the self-attention operations according to an attention mask, such as one of the attention masks described below with reference to FIG. 2.

[0119]

[0095] In some implementations, the system 120 of FIG. IB can implement an average pooling approach. For example, the system 120 can generate the layer output embeddings 132 by using the neural network 130 with a modified attention mask. For instance, the modified attention mask can be a pooling mask where each compression token is restricted to attending only to context tokens within a respective segment of context tokens in the sequence. This restriction can prevent the dispersion of attention scores over very long sequences and allow the neural network to more effectively focus on and summarize specific context subsequences.

[0120]

[0096] FIG. 2 is a diagram 200 illustrating attention masks of one or more self-attention layers of a neural network used to generate responses to queries. The neural network can be a neural network configured to process a sequence of tokens to generate a response to the query. The neural network can be included in a response generation system, e.g., the response generation system 100 of FIG. 1, configured to process a context input and a query to generate a response to the query that relates to the context input. For example, the neural network can represent or include the neural network 130 of FIG. 1.

[0121]

[0097] The sequence of tokens can include one or more of: query tokens representing the query , such as the query tokens 206a-g; context tokens representing a context input relevant to responding to the query, such as the context tokens 202a and 202c-g; compression tokens that represent information encoded in a plurality of context tokens, such as the compression tokens 204c-g; or an output sequence of tokens representing the response to the query. Each output sequence of tokens includes a sequence of answer tokens that represent the response to the query, such as the answer tokens 208a-g.

[0122]

[0098] As illustrated, the query tokens 206 can represent a natural language string that represents a question. In the example of FIG. 2, the query tokens represent the natural language string "Who found the dragon?", including a question (‘’Who found the dragon?”). The context tokens can represent a natural language string such as "The prince found the dragon. The END." The natural language string can represent context that is relevant for responding to the query.Attorney Docket No.: 45288-0626WO1

[0123]

[0099] The neural network can include one or more self-attention layers that apply selfattention to generate the response. For example, in order to generate the response, the neural network can apply self-attention over the sequence of tokens to update the output sequence of tokens representing the response. These self-attention operations can be used to generate respective layer output embeddings for each of the compression tokens in the sequence for each of the one or more self-attention layers and to update the output sequence of tokens at each of the one or more self-attention layers using the respective layer output embeddings.

[0124]

[0100] In some implementations, the response generation system can include a Key-Value (KV) cache, as described above with reference to FIG. 1 A. In order to apply self-attention over the sequence of tokens, a given self-attention layer can access keys and values for one or more of the tokens in the sequence of tokens that were computed during the self-attention operations for a preceding self-attention layer. For example, the given self-attention layer can update the output sequence of tokens based on the keys and values accessed from the KV cache. For example, the given self-attention layer can update the output sequence of tokens without using layer output embeddings from preceding layers, e.g., by instead using keys and values for tokens that were computed by preceding layers.

[0125]

[0101] The neural network can apply self-attention using an attention mask. In each of the implementations (a)-(g) illustrated in FIG. 2, the one or more self-attention layers of the neural network can apply self-attention using an attention mask corresponding to the implementation. Each attention mask defines, for each token in the sequence, the other tokens in the sequence to which the token can attend during the self-attention operations. In implementations in which the system includes a KV cache, the attention mask can define, for each token in the sequence, the other tokens in the sequence for which the keys and values from the KV cache can be used to update the token at the self-attention layer.

[0126]

[0102] The diagram 200 compares seven different attention mechanisms — labeled (a) through (g) — corresponding to seven different implementations of the neural network. In each of its seven implementations, the neural network uses the respective attention mechanism for generating the response to the query.

[0127]

[0103] The first implementation (a) illustrates an attention mechanism corresponding to a full context baseline. In this implementation, the attention mask 210a is configured as a causal mask where every token in the sequence - that is, each of the query' tokens 206a, context tokens 202a, and answer tokens 208a - is permitted to attend to all preceding tokens. For example, for every token in the sequence, the token can be updated based on keys and values from the KV cache for all preceding tokens in the sequence.Attorney Docket No.: 45288-0626WO1

[0128]

[0104] The second implementation (b) illustrates an attention mechanism corresponding to a no context baseline. In this implementation, the sequence of tokens does not include context tokens. The attention mask 210b is configured such that query tokens 206b and answer tokens 208b only attend to preceding tokens, and do not attend to context tokens. For example, for each of the query tokens 206b and answer tokens 208b, the token can be updated based only on keys and values from the KV cache for preceding tokens in the sequence, which do not include context tokens.

[0129]

[0105] The third implementation (c) illustrates an attention mechanism corresponding to including compression tokens 204c in the sequence of tokens. For example, the system can generate a modified sequence of context tokens by inserting a plurality of compression tokens 204c into the sequence of context tokens 202c. The compression tokens 204c can be interspersed throughout the sequence of context tokens 202c, e.g., such that not all of the compression tokens are adjacent to one another within the sequence of context tokens. Thus, the modified sequence of context tokens can include the context tokens 202c and the compression tokens 204c. The system can generate the modified sequence of context tokens using techniques described below with reference to FIG. 4.

[0130]

[0106] In this implementation, the attention mask 210c is configured such that: the compression tokens 204c attend to the context tokens 202c; the query tokens 206c and the answer tokens 208c do not attend to the context tokens 202c; and the query tokens 206c and the answer tokens 208c attend to the compression tokens 204c. For example, for each of the compression tokens 204c, the compression token can be updated based only on keys and values from the KV cache for context tokens. For each of the query' tokens 206c and the answer tokens 204c, the token can be updated based only on keys and values from the KV cache for compression tokens, and not for context tokens.

[0131]

[0107] The fourth implementation (d) illustrates an attention mechanism corresponding to applying average pooling to the context tokens. For example, in this implementation, the compression tokens 204d are generated by applying average pooling to the context tokens 202d, e.g., using the techniques described below with reference to FIG. 5. For example, each of the compression tokens 204d can have been generated by averaging a respective set of context tokens of the context tokens 202d. As illustrated in FIG. 2, each of the compression tokens 204d can have been generated by averaging the context tokens 202d that immediately precede the compression token in the sequence. For example, the system can insert compression tokens 204d into the sequence of context tokens 202d such that eachAttorney Docket No.: 45288-0626WO1

[0132] compression token is located immediately following a set of context tokens of which it represents the average.

[0133]

[0108] In this implementation, the attention mask 210d can be configured such that the query tokens 206d and the answer tokens 208d attend to the compression tokens 204d. The query' tokens 206d and the answer tokens 208d do not attend to the context tokens 202d. For example, for each of the query tokens 206d and the answer tokens 208d, the token can be updated based only on keys and values from the KV cache for compression tokens, and not for context tokens.

[0134]

[0109] The fifth implementation (e) illustrates an attention mechanism corresponding to applying self-attention over the sequence of tokens by offsetting the compression tokens that are processed by a particular self-attention layer. The attention mask 210e used in this implementation can be the same as the attention mask 210c used in the third implementation (c) described above. In this implementation, the sequence of tokens over which a given selfattention layer performs self-attention can include query' tokens from an immediately- preceding layer and compression tokens from an earlier layer that is before the immediately preceding layer. As explained below, offsetting the compression tokens that are processed in this way improves the self-attention process by increasing a likelihood that information encoded in compression tokens that are processed by a given layer is relevant to the output of the given layer.

[0135] [HO]

[0136] [Hl] For example, for each of the compression tokens 204c, the compression token can be updated based on keys and values from the KV cache for context tokens computed for the preceding layer. For each of the query tokens 206c, the query token can be updated based on keys and values from the KV cache computed for compression tokens for the preceding layer. For each of the answer tokens 204c, the answer token can be updated based on keys and values from the KV cache computed for compression tokens for a layer before the preceding layer. The sixth implementation (f) illustrates an attention mechanism corresponding to using different sets of parameters to process the context tokens and the compression tokens. The attention mask 21 Of used in this implementation can be the same as the attention mask 210c used in the third implementation (c) described above. As described in further detail below, in some implementations of the disclosed techniques, the system can process context tokens using a first set of parameters and compression tokens using a second, different set of parameters.Attorney Docket No.: 45288-0626WO1

[0137]

[0112] The seventh implementation (g) illustrates an attention mechanism corresponding to combining features of implementations (d), (e), and (f). The attention mask 210g is an example attention mask that can be used in implementations combining these features.

[0138] [H3] For example, in this implementation, each of the compression tokens 204g is generated by combining a respective subset of the context tokens 202g. In the illustration of FIG. 2, each of the compression tokens 204g is generated by combining the context tokens 202g that immediately precede the compression token in the sequence.

[0139]

[0114] In this implementation, a given self-attention layer applies self-attention over the sequence of tokens by offsetting the compression tokens 204g that are processed by a particular self-attention layer, as described above. In this implementation, a second, different set of parameters is used to process the compression tokens 204g. e.g., different from a first set used to process the context tokens 204g.

[0140]

[0115] FIG. 3 is a diagram 300 illustrating layer-wise information flow across a plurality of neural network layers of a neural network in a response generation system. The neural network can be a neural network configured to process a sequence of tokens to generate a response to a query. The sequence of tokens can include one or more of: query tokens representing the query, such as the query tokens 306a-c; context tokens representing a context input relevant to responding to the query', such as the context tokens 302a-c; or compression tokens that represent information encoded in a plurality of context tokens, such as the compression tokens 304b-c. The neural network can process the sequence of tokens by performing self-attention operations over the sequence of tokens using one or more selfattention layers of the neural network. As described above, performing self-attention operations over the sequence of tokens can include, at each layer of the neural network, updating layer output embeddings for each of one or more tokens in the sequence based on keys and values computed at preceding layers for other tokens in the sequence that are stored in a KV cache.

[0141]

[0116] The diagram compares the internal data movement and processing mechanics of three different implementations of the neural network: a base model implementation (a), a compression token model implementation (b), and an activation-offset compression token model implementation (c). The implementations (a) - (c) are alternative implementations for processing context tokens and generating summaries or responses.

[0142]

[0117] In the base model implementation (a), the neural network can have a standard autoregressive transformer architecture where layer output embeddings ("‘Activations”) 303a for the context tokens 302a from a first layer ('‘Layer i”) are directly consumed by a secondAttorney Docket No.: 45288-0626WO1

[0143] layer (“Layer i+1”). For example, query tokens 305a at Layer i+1 attend directly to the layer output embeddings 303a of context tokens 302a generated by Layer i. For example, layer output embeddings for each of the query tokens 305a at Layer i+1 can be updated based on keys and values for the layer output embeddings 303a computed at Layer i from the KV cache.

[0144]

[0118] In the compression token implementation (b), the neural network is configured to perform self-attention operations over a sequence of tokens that includes context tokens 302b, compression tokens 304b, and query tokens 306b. The compression tokens 304b are utilized to summarize context information encoded in the context tokens 302b, as described above. In this implementation, a first layer (“Layer i”) generates activations, e.g., layer output embeddings, 303b for the context tokens. At the subsequent layer (“Layer i+1”), the neural network processes these activations to generate a "Summary of Layer i”, including layer output embeddings 305b for compression tokens that encode a compressed version of the information encoded in the activations 303b. Because the summary 305b itself is an output of the self-attention operation at Layer i+1. the compression tokens of the summary 305b are not available to the query tokens 307b until a third layer (“Layer i+2”). For example, rather than attending to layer output embeddings that represent information corresponding to the immediately preceding layer, the query tokens 307b at Layer i+2 attend to the layer output embeddings 305b at Layer i+1 that represent information encoded in the layer output embeddings 303b generated at Layer i. This implementation introduces a one-layer information delay, where the query tokens "consume the summary of Layer i" at a greater depth, e.g., within the layers of the neural network, than the original information would have been processed in the base model. This delay can cause a mismatch between the state of the context information and the layer weights trained to interpret it. potentially degrading the accuracy of the response generation system.

[0145]

[0119] In the activation-offset compression token model implementation (c), the neural network is configured to perform self-attention in a way that can help to mitigate the information delay observed in the compression token implementation (b) described above. In this implementation, the system performs an offsetting process where layer output embeddings 305c for the compression tokens 304c (the "Summary of Layer i") generated by a given layer are shifted or "offset" down by at least one layer. Specifically, the summary ■ generated by processing context activations, e.g., layer output embeddings, 303c from Layer i is made available for consumption by the query tokens 307c at Layer i+1 instead of Layer i+2. For example, layer output embeddings for each of the query tokens 307c at Layer i+1Attorney Docket No.: 45288-0626WO1

[0146] can be updated based on keys and values for the layer output embeddings 305c computed at Layer i+1 from the KV cache.

[0147]

[0120] For example, for a particular self-attention layer i in a sequence of layers, the response generation system can generate a response to the query by processing query token embeddings from an immediately preceding layer and compression token embeddings from an earlier layer before the immediately preceding layer. According to the techniques disclosed in this specification, the neural network of the response generation system can generate responses to queries by performing attention operations in which, for each of the self-attention layers of the neural network and for each of an output sequence of tokens representing the response, a respective layer output embedding of the token attends to query token embeddings from an immediately preceding layer and compression token embeddings from an earlier layer before the immediately preceding layer. For example, for each of the self-attention layers of the neural network and for each of an output sequence of tokens representing the response, a respective layer output embedding of the token can be updated based on keys and values from the KV cache for query tokens computed at an immediately preceding layer, and keys and values from the KV cache for compression tokens computed at an earlier layer before the immediately preceding layer.

[0148]

[0121] Offsetting the layer from which compression tokens are processed in this way can help to increase a likelihood that the compressed context information encoded in the compression tokens is aligned with the layer at which self-attention parameters of the neural network expect to receive it. Realigning the information flow in this way can improve the performance of the response generation system, e.g., by allowing the system to handle long-context reasoning tasks without requiring significant architectural modifications to the underlying neural network.

[0149]

[0122] FIG. 4 is a flowchart of an example process 400 for generating a response to a query using compression tokens. The process 400 can be performed by a system of one or more computers, such as the response generation system 100 of FIG. 1A or the system 120 of FIG. IB.

[0150]

[0123] The system receives context input including a sequence of context tokens (402). In some implementations, the system can also receive a query relating to the context input. The context tokens can include one or more of video tokens, text tokens, or audio tokens. The context tokens can represent various data modalities, such as text tokens from a document, image patches from a frame of video, or segments of an audio waveform. Each context token represents data related to a context for the query, as described above with reference to FIG.Attorney Docket No.: 45288-0626WO1

[0151] 1 A. For example, the sequence of context tokens included in the context input can include a sequence of vectors representing a story, an instruction for an agent to complete a task, or a set of facts.

[0152]

[0124] In some implementations, the context input is received from a user. In some implementations, the query related to the context input is received from a user.

[0153]

[0125] The system generates a modified sequence of context tokens including compression tokens interspersed throughout the sequence (404). The system can generate the modified sequence of context tokens by modifying the sequence of context tokens to insert a plurality of compression tokens into the sequence of context tokens. In some implementations, the plurality of compression tokens are inserted into the sequence of context tokens in such a way that the compression tokens are interspersed throughout the sequence of context tokens. For the plurality of compression tokens to be interspersed throughout the sequence of context tokens can mean that at least one compression token is separated from at least one other compression token by at least one context token. For example, the plurality of compression tokens are interspersed throughout the sequence of context tokens if not all the compression tokens of the plurality of compression tokens are adjacent to one another within the sequence of context tokens.

[0154]

[0126] In some implementations, the system generates the modified sequence of context tokens by first determining a number of compression tokens to insert into the sequence of context tokens. The system can determine the number of compression tokens based on a total number of context tokens in the sequence. The system can determine the number of compression tokens based on a target compression rate. The target compression rate can be a predetermined compression rate that indicates a ratio of context tokens in the sequence to compression tokens inserted into the sequence. For example, if the sequence of context tokens includes 1000 tokens and the compression rate is 10, the system can determine that the ratio of context tokens to compression tokens in the modified sequence is to be 10. As a result, the system can insert 100 compression tokens into the sequence to generate the modified sequence (e.g., because the ratio of 1000 to 100 is equal to 10).

[0155]

[0127] Each compression token can be an initial embedding. The initial embedding can be a representation of data in an embedding space. By processing each compression token as described below, the system can modify each initial embedding for each compression token to generate respective layer output embeddings for the compression tokens for each layer of a neural network. Each layer output embedding can be a representation of data in the sameAttorney Docket No.: 45288-0626WO1

[0156] embedding space corresponding to the initial embeddings. Rather than being random, the data represented by the layer output embeddings can be related to the context for the query.

[0157]

[0128] In some implementations, each compression token corresponds to a plurality of context tokens of the sequence of context tokens. The respective layer embeddings for each compression token can represent data related to the data for the context for the query that is represented by the corresponding context tokens. For example, each compression token can correspond to a respective set of a plurality of context tokens that immediately precede the compression token in the sequence (e g., and follow a preceding compression token in the sequence, such that they are in between the preceding compression token and the compression token). The respective layer output embedding for each compression token can then represent data related to the data for the context that is represented by the context tokens included in the corresponding set.

[0158]

[0129] The system can insert the compression tokens into the sequence in any of a variety of locations such that they are interspersed throughout the sequence. To be interspersed throughout the sequence can mean that at least one compression token can be separated from at least one other compression token by at least one context token of the sequence of context tokens. In some implementations, the system can intersperse the compression tokens uniformly throughout the sequence, for example, by inserting one compression token after every fixed number of context tokens. The fixed number can be any integer. For example, the fixed number can be determined by the compression rate, e.g., equal to the quotient of the number of context tokens in the sequence divided by the compression rate. By interspersing the compression tokens rather than appending them only at the end of the sequence, the system enables each compression token to attend to a local subset of context tokens.

[0159]

[0130] The system uses a neural network to generate a respective layer output embedding of each compression token (406). The neural network includes a plurality of layers, which can be self-attention layers that apply self-attention. In some implementations, the neural network includes at least one self-attention layer. In some implementations, at least one self-attention layer of the neural network applies self-attention that is causally masked. In some implementations, the neural network applies self-attention according to one of the attention masks described above with reference to FIG. 2. The neural netw ork can process the modified sequence of context tokens to, for each of the plurality of layers, generate a respective layer output embedding of each of the compression tokens.

[0160]

[0131] Processing the modified sequence of context tokens to generate the respective layer output embeddings can include performing self-attention operations over a sequence ofAttorney Docket No.: 45288-0626WO1

[0161] tokens at each of the layers of the neural network. Performing self-attention operations over the sequence of tokens can include performing self-attention operations over embeddings representing the tokens in the sequence. In some implementations, the neural network can perform the self-attention operations using a KV cache. For example, performing the selfattention operations at a given layer of the neural network can include accessing keys and values stored in the KV cache for tokens that were computed at a preceding layer of the neural network. Performing the self-attention operations at a given layer can include updating each layer output embedding in the sequence using keys and values stored in the KV cache for a specific set of tokens that are indicated by an attention mask, as described with reference to FIG. 2.

[0162]

[0132] For example, in some implementations, for at least one self-attention layer of the neural network, the self-attention layer performs the attention operations by applying a modified attention mask where each compression token or layer output embedding of a compression token is restricted to attending only to a specific subset of the context tokens of the sequence. The compression token or layer output embedding of the compression token can be updated using only keys and values from the KV cache computed for the specific subset of the context tokens. For example, each compression token or layer output embedding of the compression token can attend to a different subset of the context tokens from at least one of the other compression tokens or layer output embeddings of the compression token. For example, each compression token or layer output embedding of the compression token can be restricted to attending only to the context tokens to which the compression token corresponds.

[0163]

[0133] In some implementations, the modified attention mask can restrict the tokens to which the context tokens or layer output embeddings of the context tokens attend, e.g., such that the context tokens or layer output embeddings do not attend to the compression tokens. This restricted attention mechanism can provide a pooling inductive bias that overcomes the technical challenge of attention dispersion in long sequences, where a standard softmax operation might otherwise fail to focus on relevant context subsets.

[0164]

[0134]

[0165]

[0135] Each layer output embedding can be a representation of data in an embedding space. For example, the embedding space can be the same embedding space of initial embeddings corresponding to the compression tokens. The represented data can relate to the context for the query represented by the context tokens from the received context input. In some implementations, each layer output embedding for a given compression token can representAttorney Docket No.: 45288-0626WO1

[0166] data related to the data for the context for the query that is represented by a plurality of context tokens corresponding to the given compression token.

[0167]

[0136] In some implementations, the system processes the modified sequence of context tokens using the neural network to, for each of the plurality of layers of the neural network, generate a respective layer output embedding of each of the context tokens of the sequence of context tokens. For example, the system can generate respective layer output embeddings for each of the context tokens using techniques similar to those used to generate the respective layer output embeddings for the compression tokens.

[0168]

[0137] In some implementations, the neural netw ork uses a first set of parameters to generate the respective layer output embeddings for the context tokens and a second, different set of parameters to generate the respective layer output embeddings for the compression tokens. As described above with reference to FIG. 1 A, this separation of parameters allows the system to optimize the compression functionality without degrading the base language modeling capabilities. While the neural network is being trained, the second set of parameters can be updated with a second learning rate that is higher than a first learning rate used to update the first set of parameters. The training of the neural network is described in further detail below.

[0169]

[0138] The system generates a response to a query' that relates to the context input (408). The system generates the response using the respective layer output embeddings of the compression tokens for the plurality of layers. The response can include one or more of a natural language answer, a predicted action for a robotic agent, or a summary of the context. In implementations in which the context input includes an instruction for an agent to complete a task, the response can include the completion of the task. In some implementations, the query’ includes a request for information that is included in information represented by the context input. The response can then include the information requested by the query'. The response can include response tokens representing a response to the query', e.g., response tokens representing one or more of a natural language answer, a predicted action for a robotic agent, or a summary of the context. In implementations in which the context input is received from a user, the system can provide the generated response to the query' to the user on a user device.

[0170]

[0139] The query' can be a query that is received by the system that relates to the context represented by the context input received at the operation 402. The query can include a plurality of query’ tokens that represent the query. To generate the response, the system can process the layer output embeddings of the compression tokens and the query'. In someAttorney Docket No.: 45288-0626WO1

[0171] implementations, the system can process the layer output embeddings of the compression tokens and the query without using any embeddings of the context tokens. For example, the system can generate the response using only information encoded in the compression tokens. In some implementations, the system can process the layer output embeddings of the compression tokens and the query, and additionally process a proper subset of the context tokens. A proper subset of the context tokens can be a set of the context tokens that such that at least one of the context tokens is not included in the set. In some implementations, the proper subset includes less than 25% of the context tokens.

[0172]

[0140] In some implementations, the system can use a neural network to generate the response. The neural network can be the same neural network used to generate the layer output embeddings.

[0173]

[0141] In some implementations, the neural network can include a sequence of self-attention layers. The neural network can generate the response by auto-regressively generating an output sequence of tokens. The system can initialize an output sequence of tokens, e.g., such that each token in the sequence is an initial embedding that is a representation of data in an embedding space. The neural network can process the output sequence of tokens by updating each of the output sequence of tokens at each of a plurality of self-attention layers. The generated output sequence of tokens can include the output sequence tokens after it has been updated by a final self-attention layer.

[0174]

[0142] For example, the neural network can generate each token in the output sequence by processing, using each of the self-attention layers of the sequence of self-attention layers, a sequence that includes the tokens in the output sequence, the query tokens included in the query , and the compression tokens. In some implementations, the sequence also includes context tokens. The tokens in the output sequence used to generate each token can include the token and tokens preceding the token in the output sequence. As described above, processing the sequence of tokens can include performing attention operations over embeddings representing the tokens in the sequence. Performing self-attention operations over the embeddings for the sequence at a self-attention layer can include updating the embeddings for the tokens in the output sequence based on keys and values from the KV cache that were computed at a preceding layer for a specific set of tokens that are indicated by an attention mask, such as those described with reference to FIG. 2. Each token in the output sequence that is autoregressively generated by the neural network can be the embedding for the token in the output sequence after it has been updated by a final self-attention layer of the neural network.Attorney Docket No.: 45288-0626WO1

[0175]

[0143] For example, the neural network can first generate a first token in the output sequence by using each of the self-attention layers to process the initial embedding for the first token, the layer output embeddings for the query tokens, and the layer output embeddings for the compression tokens. The neural network can next generate a second token in the output sequence by using each of the self-attention layers to process the initial embedding for the second token, the layer output embeddings for the generated first token, the layer output embeddings for the query tokens, and the layer output embeddings for the compression tokens.

[0176]

[0144] In this way, for each subsequent token in the output sequence, the neural network can generate the subsequent token by using each of the self-attention layers to process the initial embedding for the subsequent token, layer output embeddings for the generated preceding tokens in the output sequence, the layer output embeddings for the query tokens, and the layer output embeddings for the compression tokens. In some implementations, to generate a particular token of the output sequence, each of the self-attention layers processes specific embeddings for each of the particular token, the generated preceding tokens, the query tokens and the compression tokens. For example, the specific embeddings can be embeddings that were generated by particular other self-attention layers in particular positions with respect to the self-attention layer processing the specific embeddings.

[0177]

[0145] In some examples, the neural network auto-regressively generates the output sequence of tokens by. for each particular token in the output sequence and for a particular self-attention layer in the sequence of self-attention layers: processing (i) respective layer output embeddings of the query tokens generated by an immediately preceding self-attention layer in the sequence, (ii) a respective layer output embedding of the particular token generated by the immediately preceding layer in the sequence, (iii) respective layer output embeddings of any preceding tokens in the output sequence generated by the immediately preceding layer in the sequence, and (iv) the respective layer output embeddings of the compression tokens generated by an earlier layer that is before the immediately preceding layer in the sequence using the particular self-attention layer to generate a layer output embedding of the particular token.

[0178]

[0146] In implementations in which the system includes a KV cache, the neural network can auto-regressively generate the output sequence of tokens by, for each particular token in the output sequence and for a particular self-attention layer in the sequence of self-attention layers, updating the layer output embedding for the particular token using keys and values from the KV cache that were computed for: (i) respective layer output embeddings of theAttorney Docket No.: 45288-0626WO1

[0179] query tokens at an immediately preceding self- attend on layer in the sequence, (ii) respective layer output embeddings of any preceding tokens in the output sequence at the immediately preceding layer in the sequence, and (iii) the respective layer output embeddings of the compression tokens at an earlier layer that is before the immediately preceding layer in the sequence using the particular self-attention layer to generate a layer output embedding of the particular token.

[0180]

[0147] For example, for each particular token in the output sequence and for a particular self-attention layer (e.g., layer i), the system can process query token embeddings from an immediately preceding layer (e.g., layer i-1) and compression token embeddings from an earlier layer (e.g., layer i-2). Offsetting the layer corresponding to the compression tokens that are processed by each layer to generate each token in the output sequence can increase a likelihood of correct information flow in the neural network. For example, this offsetting process can help to align the layer at which the neural network expects to receive compressed context information from a particular other layer, e.g., represented by the embeddings for the compression tokens generated by the particular other layer, with the layer of the neural network that ends up processing the compressed context information. This offsetting process is further illustrated and described with respect to the implementation (c) of FIG. 3.

[0181]

[0148] The neural network used to generate the respective layer output embeddings at the operation 406 can be trained in any of a variety of ways. For example, the neural network can be trained on a set of training examples that each include: i) a training network input including a training modified sequence of context tokens into which compression tokens have been inserted and a query including query tokens; and ii) a target network output that includes a target response to the query. The neural network can process, for each training example, the training network input to generate a respective predicted network output that includes predicted response to the query. As described above, processing each training example can include generating layer output embeddings for compression tokens in the training modified sequence, and optionally for context tokens in the training modified sequence.

[0182]

[0149] For example, the training of the neural network can include, for each training example, determining respective gradients of an objective function with respect to current values of a set of neural network parameters of the neural network. The objective function measures an error between: (i) the respective predicted network output generated by the neural network by processing the training network input of the training example; and (ii) theAttorney Docket No.: 45288-0626WO1

[0183] target network output of the training example. The training of neural network can include, for each training example, updating the current values of the set of neural network parameters of the neural network using the respective gradients determined for the training example. In some implementations, the current values of the set of neural network parameters can be updated using the respective gradients and a learning rate that defines how much the current values are changed during the update.

[0184]

[0150] As described above, in some implementations, the neural network can generate the predicted layer output embeddings for the context tokens using a first set of parameters, and generate the predicted layer output embeddings for the compression tokens using a second, different set of parameters. In such implementations, current values of the first set of parameters can be updated using the respective gradients for the first set of parameters and a first learning rate. Cunent values of the second set of parameters can be updated using the respective gradients for the second set of parameters and a second, different learning rate. For example, the second learning rate can be higher than the first learning rate, e.g., the second set of parameters can be changed by a larger proportional amount during the update than the first set of parameters. In some implementations, the second learning rate is 100 times higher than the first learning rate.

[0185]

[0151] The gradients used for training the neural network can be computed using backpropagation. The gradients can be used to update the parameter values of the property prediction neural network using an update rule of any appropriate gradient descent optimization algorithm, e g., RMSprop or Adam. The objective function used for training the neural network can be any appropriate objective function.

[0186]

[0152] In some implementations, the neural network is first trained on a first set of training examples of which the training network input includes a training sequence of context tokens that does not include compression tokens. For example, each training example in the first set of training examples can include: i) a training network input including a training sequence of context tokens that does not include compression tokens and a query7including query tokens; and ii) a target network output that includes target response to the query7. The neural network can be trained on the first set of training examples as described above. In such implementations, the neural network is next trained on a second set of training examples of which the training network input includes a training modified sequence of context tokens that includes compression tokens, e.g., as described above.Attorney Docket No.: 45288-0626WO1

[0187]

[0153] FIG. 5 is a flowchart of an example process 500 for generating a response to a query using context compression via average pooling. The process 500 can be performed by a system of one or more computers, such as the response generation system 100 of FIG. 1A or the system 150 of FIG. 1C.

[0188]

[0154] The system receives context input including a sequence of context tokens (502). For example, the received context input can be similar to the context input described above with reference to the operation 402 of FIG. 4.

[0189]

[0155] The system identifies segments of the sequence that each include a proper subset of the sequence (504). In some implementations, the identified segments can be nonoverlapping. In some implementations, the system can identify the segments based on a predetermined compression rate, as described above with reference to FIG. 1C. For example, if the compression rate is N, where N is any positive integer, the system can identify segments such that the segments are non-overlapping and each include N tokens.

[0190]

[0156] Each identified segment can include a different proper subset of the original tokens. Segmenting the sequence of context tokens can enable the system to organize the sequence into manageable units for local summarization, increasing a likelihood that each part of the context is represented in the final compressed output.

[0191]

[0157] The system uses a neural network to generate a respective layer output embedding of each context token (506). The neural network can include a plurality of layers, such as selfattention layers or feed-forward layers. The neural network can be similar to the neural network used to perform the operation 406 of FIG. 4, as described above. The neural network can have been trained according to the training techniques described with reference to FIG. 4.

[0192]

[0158] The neural network processes the sequence of context tokens to generate, for each of the plurality of layers, a respective layer output embedding of each context token. For example, the system can generate a layer output embedding for every context token in the original sequence for each layer of the network. Each layer output embedding for each context token can be a vector or ordered collection of numerical values representing the state of the context token at a particular layer, e.g., depth, of the neural network. The system can generate the layer output embeddings using techniques such as those described with reference to the operation 406 of FIG. 4. For example, the techniques can include at least one selfattention layer of the neural network performing attention operations over a sequence that includes the sequence of context tokens and query tokens for the query. Performing the selfattention operations can include updating layer output embeddings for the context tokensAttorney Docket No.: 45288-0626WO1

[0193] using keys and values for other context tokens and for query tokens computed at previous self-attention layers that are stored in the KV cache.

[0194]

[0159] For each of the plurality of layers and for each of the plurality of segments of the sequence of context tokens, the system generates a compressed layer output embedding for the segment by applying average pooling on the respective layer output embeddings for the layer and for the context tokens in the segment (508). For example, for each layer and for each segment, the system can apply average pooling by first calculating a mean vector for the layer and for the segment. The mean vector can be an average of the layer output embeddings generated by the layer across the context tokens in the segment. For example, the system can calculate the mean vector by determining, for each constituent element of the mean vector, an average across the context tokens in the segment of the corresponding constituent elements of the layer output embeddings for the context token. Each constituent element of the mean vector can be the determined average across the context tokens in the segment. The compressed layer output embedding for the segment can include the mean vector.

[0195]

[0160] The system generates a response to a query that relates to the context input using the compressed layer output embeddings for the segments for the layers (510). The system can generate the response by processing the compressed layer output embeddings and the query. For example, the query can include one or more query tokens representing the query. The system can use a generative neural network included in the system to process the query tokens and the compressed layer output embeddings to generate the response. In some implementations, the system can process the query tokens and the compressed layer output embeddings without processing any of the context tokens. In some implementations, the system can process the query tokens, the compressed layer output embeddings, and a proper subset of the context tokens, e.g., as described with reference to FIG. 4.

[0196]

[0161] In some implementations, the generative neural network can generate the response by performing self-attention operations over a sequence of tokens, e.g., including an output sequence of tokens, the query7tokens, the compressed layer output embeddings, and optionally a proper subset of context tokens. The generative neural network can perform the self-attention operations over the sequence of tokens using techniques described above with reference to FIG. 4. For example, the generative neural network can generate the response by updating, at each layer, tokens in the output sequence of tokens. The generative neural network can update the output sequence of tokens using keys and values from a KV cache that were computed for other tokens in the sequence by preceding layers of the generativeAttomey Docket No.: 45288-0626WO1

[0197] neural network. The response can include the output sequence of tokens after the output sequence of tokens has been updated by a final layer of the generative neural network.

[0198]

[0162] The response can represent a natural language answer to the query, a predicted action for an agent that is responsive to the query, or a classification output that is responsive to the query. For example, the response can include any of a variety of elements, as described with reference to FIG. 4.

[0199]

[0163] In some implementations, the system can generate the response by autoregressively generating an output sequence of tokens using a neural network with a plurality of selfattention layers included in the system, e.g., as described above with reference to the operation 408 of FIG. 4. For example, in some implementations, the system can autoregressively generate the output sequence by processing, for each layer of the neural network, compression tokens generated by a layer before the immediately preceding, e.g., by offsetting the compression tokens that are processed to generate the output sequence of tokens.

[0200]

[0164] The process 500 can enable the system to generate the response to the query by processing fewer tokens (the compressed embeddings instead of the full sequence of context tokens). By processing fewer tokens, the system can increase throughput and reduce memory consumption, e.g., as compared to processing the full sequence.

[0201]

[0165] 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.

[0202]

[0166] 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-Attorney Docket No.: 45288-0626WO1

[0203] 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.

[0204]

[0167] The term “data processing apparatus7’ 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.

[0205]

[0168] 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.

[0206]

[0169] 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.

[0207]

[0170] 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 moreAttorney Docket No.: 45288-0626WO1

[0208] 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.

[0209]

[0171] 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.

[0210]

[0172] 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.

[0211]

[0173] 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.

[0212]

[0174] 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 theAttorney Docket No.: 45288-0626WO1

[0213] 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.

[0214]

[0175] 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.

[0215]

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

[0216]

[0177] 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.

[0217]

[0178] 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.

[0218]

[0179] In addition to the embodiments described above, the following embodiments are also innovative:Attorney Docket No.: 45288-0626WO1

[0219]

[0180] Embodiment 1 is a method performed by one or more computers, the method comprising: receiving a context input comprising a sequence of context tokens and a query; generating a modified sequence of context tokens, comprising modifying the sequence of context tokens to insert a plurality of compression tokens into the sequence of context tokens, wherein the compression tokens of the plurality of compression tokens are interspersed throughout the sequence of context tokens; processing the modified sequence of context tokens using a neural network to, for each of a plurality of layers of the neural network, generate a respective layer output embedding of each of the compression tokens; and generating a response to the query that relates to the context input using the respective layer output embeddings of the compression tokens for the plurality of layers.

[0220]

[0181] Embodiment 2 is the method of embodiment 1, wherein generating the response to the query using the respective layer output embeddings of the compression tokens comprises: generating the response to the uery using the respective layer output embeddings of the compression tokens without using any embeddings of the context tokens.

[0221]

[0182] Embodiment 3 is the method of embodiment 1, wherein generating the response to the query using the respective layer output embeddings of the compression tokens comprises: generating the response to the uery using the respective layer output embeddings and respective layer output embeddings for a proper subset of the context tokens.

[0222]

[0183] Embodiment 4 is the method of embodiment 3, wherein the proper subset includes less than 25% of the context tokens.

[0223]

[0184] Embodiment 5 is the method of any of embodiments 1-4, wherein processing the modified sequence of context tokens using a neural network to, for each of a plurality of layers of the neural network, generate a respective layer output embedding of each of the compression tokens comprises processing the modified sequence of context tokens using the neural network to, for each of the plurality of layers of the neural network, generate a respective layer output embedding of each of the context tokens of the sequence of context tokens.

[0224]

[0185] Embodiment 6 is the method of embodiment 5, processing the modified sequence of context tokens using the neural network to, for each of the plurality of layers of the neural network, generate a respective layer output embedding of each of the context tokens of the sequence of context tokens comprises generating a respective layer output embedding of each of the context tokens using a first set of parameters; and processing the modified sequence of context tokens using a neural network to, for each of a plurality of layers of the neural network, generate a respective layer output embedding of each of the compression tokensAttorney Docket No.: 45288-0626WO1

[0225] comprises generating a respective layer output embedding of each of the compression tokens using a second set of parameters that is different from the first set of parameters.

[0226]

[0186] Embodiment 7 is the method of embodiment 6. wherein, during training of the neural network, the second set of parameters have been updated with a second learning rate that is higher than a first learning rate used to update the first set of parameters.

[0227]

[0187] Embodiment 8 is the method of embodiment 7, wherein the second learning rate is 10 times higher than the first learning rate.

[0228]

[0188] Embodiment 9 is the method of any of embodiments 1-8, wherein the compression tokens of the plurality of compression tokens are interspersed uniformly throughout the sequence of context tokens.

[0229]

[0189] Embodiment 10 is the method of any of embodiments 1-9, wherein the plurality of layers of the neural network include one or more self-attention layers that apply selfattention.

[0230]

[0190] Embodiment 11 is the method of embodiment 10, wherein the self-attention is causally masked.

[0231]

[0191] Embodiment 12 is the method of embodiment 10 or embodiment 11, wherein, for each the one or more self-attention layers, the respective context tokens do not attend to the compression tokens.

[0232]

[0192] Embodiment 13 is the method of any of embodiments 1-12, wherein the neural network has been trained on input sequences of tokens that do not include compression tokens and further trained on input sequences of tokens that include compression tokens.

[0233]

[0193] Embodiment 14 is the method of any of embodiments 1-13, further comprising: determining a number of compression tokens to include in the modified context sequence based on a total number of context tokens in the sequence of context tokens.

[0234]

[0194] Embodiment 15 is a method performed by one or more computers, the method comprising: receiving a context input comprising a sequence of context tokens; identifying a plurality of segments of the sequence of context tokens that each include a different proper subset of the sequence of context tokens; processing the sequence of context tokens using a neural network to, for each of a plurality of layers of the neural network, generate a respective layer output embedding of each of the context tokens; for each of the plurality' of layers and for each of the plurality of segments of the sequence of context tokens, generating a compressed layer output embedding for the segment by applying average pooling on the respective layer output embeddings for the layer and for the context tokens in theAttorney Docket No.: 45288-0626WO1

[0235] segment; and generating a response to a query that relates to the context input using the compressed layer output embeddings for the segments for the layers.

[0236]

[0195] Embodiment 16 is a method performed by one or more computers, the method comprising: receiving a context input comprising a sequence of context tokens; generating a modified sequence of context tokens, comprising modifying the sequence of context tokens to insert a plurality of compression tokens into the sequence of context tokens; processing the modified sequence of context tokens using a neural network to, for each of a plurality of layers of the neural network, generate a respective layer output embedding of each of the compression tokens, wherein the neural network comprises a self-attention layer that applies self-attention and wherein, during the self-attention (i) the context tokens do not attend to the compression tokens and (ii) each compression token attends to a different subset of the context tokens from at least one of the other compression tokens; and generating a response to a query that relates to the context input using the respective layer output embeddings of the compression tokens for the plurality of lay ers.

[0237]

[0196] Embodiment 17 is a method performed by one or more computers, the method comprising: receiving a context input comprising a sequence of context tokens; generating a modified sequence of context tokens, comprising modifying the sequence of context tokens to insert a plurality of compression tokens into the sequence of context tokens; processing the modified sequence of context tokens using a neural network to, for each of a sequence of self-attention layers of the neural network, generate a respective layer output embedding of each of the compression tokens; and generating a response to a query that comprises a plurality of query tokens using the respective layer output embeddings of the compression tokens by auto-regressively generating an output sequence of tokens using the neural network, comprising, for each particular token in the output sequence and for a particular self-attention layer in the sequence of self-attention layers: processing (i) respective layer output embeddings of the query tokens generated by an immediately preceding self-attention layer in the sequence, (ii) a respective layer output embedding of the particular token generated by the immediately preceding layer in the sequence, (iii) respective layer output embeddings of any preceding tokens in the output sequence generated by the immediately preceding layer in the sequence, and (iv) the respective layer output embeddings of the compression tokens generated by an earlier layer that is before the immediately preceding layer in the sequence using the particular self-attention layer to generate a layer output embedding of the particular token.Attorney Docket No.: 45288-0626WO1

[0238]

[0197] Embodiment 18 is the method of embodiment 17, wherein the particular self-attention layer has index i in the sequence of self-attention layers and the earlier layer has index i-2.

[0239]

[0198] Embodiment 19 is the method of embodiment 17 or embodiment 18, wherein processing the modified sequence of context tokens using a neural network to, for each of a sequence of self-attention layers of the neural network, generate a respective layer output embedding of each of the compression tokens comprises processing the modified sequence of context tokens using the neural network to, for each of the sequence of self-attention layers of the neural network, generate a respective layer output embedding of each of the context tokens of the sequence of context tokens.

[0240]

[0199] Embodiment 20 is the method of embodiment 19, wherein: processing the modified sequence of context tokens using the neural network to. for each of the sequence of selfattention layers of the neural network, generate a respective layer output embedding of each of the context tokens of the sequence of context tokens comprises generating a respective layer output embedding of each of the context tokens using a first set of parameters; and processing the modified sequence of context tokens using a neural network to, for each of a sequence of self-attention layers of the neural network, generate a respective layer output embedding of each of the compression tokens comprises generating a respective layer output embedding of each of the compression tokens using a second set of parameters that is different from the first set of parameters.

[0241]

[0200] Embodiment 21 is the method of embodiment 16, wherein the compression tokens of the plurality of compression tokens are interspersed throughout the sequence of context tokens.

[0242]

[0201] Embodiment 22 is the method of embodiment 17, wherein the compression tokens of the plurality of compression tokens are interspersed throughout the sequence of context tokens.

[0243]

[0202] Embodiment 23 is the method of embodiment 16, wherein: the neural network comprises a sequence of self-attention layers; the response to the query comprises a plurality of query tokens; and generating the response to the query that relates to the context input using the respective layer output embeddings of the compression tokens for the plurality of layers comprises auto-regressively generating an output sequence of tokens using the neural network, comprising, for each particular token in the output sequence and for a particular self-attention layer in the sequence of self-attention layers: processing (i) respective layer output embeddings of the query tokens generated by an immediately preceding self-attention layer in the sequence, (ii) a respective layer output embedding of the particular tokenAttorney Docket No.: 45288-0626WO1

[0244] generated by the immediately preceding layer in the sequence, (iii) respective layer output embeddings of any preceding tokens in the output sequence generated by the immediately preceding layer in the sequence, and (iv) the respective layer output embeddings of the compression tokens generated by an earlier layer that is before the immediately preceding layer in the sequence using the particular self-attention layer to generate a layer output the embedding of the particular token.

[0245]

[0203] Embodiment 24 is the method of embodiment 23, wherein processing the modified sequence of context tokens using a neural network to, for each of a plurality of layers of the neural network, generate a respective layer output embedding of each of the compression tokens comprises processing the modified sequence of context tokens using the neural network to. for each of the plurality of layers of the neural network, generate a respective layer output embedding of each of the context tokens of the sequence of context tokens.

[0246]

[0204] Embodiment 25 is the method of embodiment 24, wherein: processing the modified sequence of context tokens using the neural network to, for each of the plurality of layers of the neural network, generate a respective layer output embedding of each of the context tokens of the sequence of context tokens comprises generating a respective layer output embedding of each of the context tokens using a first set of parameters; and processing the modified sequence of context tokens using a neural network to, for each of a plurality' of layers of the neural network, generate a respective layer output embedding of each of the compression tokens comprises generating a respective layer output embedding of each of the compression tokens using a second set of parameters that is different from the first set of parameters.

[0247]

[0205] Embodiment 26 is the method of embodiment 23, wherein each self-attention layer of the neural network comprises an attention head configured to generate one or more key vectors and one or more value vectors for each of the query tokens, the compression tokens, and the tokens in the output sequence of tokens.

[0248]

[0206] Embodiment 27 is the method of embodiment 26, wherein the one or more key and value vectors are stored in a Key -Value cache.

[0249]

[0207] Embodiment 28 is the method of embodiment 27, wherein the Key- Value cache is a specialized memory buffer or data structure.

[0250]

[0208] Embodiment 29 is the method of any of embodiments 1-28, wherein the context tokens comprise one or more of video tokens, text tokens, or audio tokens.Attorney Docket No.: 45288-0626WO1

[0251]

[0209] Embodiment 30 is the method of any of embodiments 1-29, wherein the context input indicates an instruction for an agent to complete a task and the response to the query comprises the completion of the task.

[0252]

[0210] Embodiment 31 is the method of any of embodiments 1-30, wherein the query comprises a request for information that is included in information represented by the context input.

[0253]

[0211] Embodiment 32 is the method of any of embodiments 1-31, wherein the context input is received from a user and the method further comprises providing the response to the query to the user on a user device.

[0254]

[0212] Embodiment 33 is the method of any of embodiments 1-32, wherein that modified sequence of context tokens includes at least one compression token that is separated from at least one other compression token by at least one context token.

[0255]

[0213] 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. 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.

[0256]

[0214] 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.

[0257]

[0215] Particular embodiments of the subject matter have been described. Other embodiments are within the scope of the following claims. For example, the actions recitedAttorney Docket No.: 45288-0626WO1

[0258] 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.

[0259]

[0216] What is claimed is:

Claims

Attorney Docket No.: 45288-0626WO1CLAIMS1. A method performed by one or more computers, the method comprising:receiving a context input comprising a sequence of context tokens;receiving a query relating to the context input;generating a modified sequence of context tokens, comprising modifying the sequence of context tokens to insert a plurality of compression tokens into the sequence of context tokens, wherein the compression tokens of the plurality of compression tokens are interspersed throughout the sequence of context tokens;processing the modified sequence of context tokens using a neural network to, for each of a plurality of layers of the neural network, generate a respective layer output embedding of each of the compression tokens; andgenerating a response to the query that relates to the context input using the respective layer output embeddings of the compression tokens for the plurality' of layers.

2. The method of claim 1, wherein generating the response to the query using the respective layer output embeddings of the compression tokens comprises:generating the response to the query using the respective layer output embeddings of the compression tokens without using any embeddings of the context tokens.

3. The method of claim 1, wherein generating the response to the query using the respective layer output embeddings of the compression tokens comprises:generating the response to the query using the respective layer output embeddings and respective layer output embeddings for a proper subset of the context tokens.

4. The method of claim 3, wherein the proper subset includes less than 25% of the context tokens.

5. The method of any preceding claim, wherein processing the modified sequence of context tokens using a neural network to, for each of a plurality of layers of the neural network, generate a respective layer output embedding of each of the compression tokens comprises processing the modified sequence of context tokens using the neural network to, for each of the plurality of layers of the neural network, generate a respective layer outputAttorney Docket No.: 45288-0626WO1embedding of each of the context tokens of the sequence of context tokens.

6. The method of claim 5, wherein:processing the modified sequence of context tokens using the neural network to, for each of the plurality7of layers of the neural network, generate a respective layer output embedding of each of the context tokens of the sequence of context tokens comprises generating a respective layer output embedding of each of the context tokens using a first set of parameters; andprocessing the modified sequence of context tokens using a neural network to, for each of a plurality of layers of the neural network, generate a respective layer output embedding of each of the compression tokens comprises generating a respective layer output embedding of each of the compression tokens using a second set of parameters that is different from the first set of parameters.

7. The method of claim 6, wherein, during training of the neural network, the second set of parameters have been updated with a second learning rate that is higher than a first learning rate used to update the first set of parameters.

8. The method of claim 7, wherein the second learning rate is 10 times higher than the first learning rate.

9. The method of any preceding claim, wherein the compression tokens of the plurality of compression tokens are interspersed uniformly throughout the sequence of context tokens.

10. The method of any preceding claim, wherein the plurality of layers of the neural network include one or more self-attention layers that apply self-attention.

11. The method of claim 10, wherein the self-attention is causally masked.

12. The method of claim 10 or claim 11, wherein, for each the one or more self-attention layers, the respective context tokens do not attend to the compression tokens.

13. The method of any preceding claim, wherein the neural network has been trained onAttorney Docket No.: 45288-0626WO1input sequences of tokens that do not include compression tokens and further trained on input sequences of tokens that include compression tokens.

14. The method of any preceding claim, further comprising:determining a number of compression tokens to include in the modified context sequence based on a total number of context tokens in the sequence of context tokens.

15. A method performed by one or more computers, the method comprising:receiving a context input comprising a sequence of context tokens;identifying a plurality of segments of the sequence of context tokens that each include a different proper subset of the sequence of context tokens;processing the sequence of context tokens using a neural network to, for each of a plurality of layers of the neural network, generate a respective layer output embedding of each of the context tokens;for each of the plurality of layers and for each of the plurality of segments of the sequence of context tokens, generating a compressed layer output embedding for the segment by applying average pooling on the respective layer output embeddings for the layer and for the context tokens in the segment; andgenerating a response to a query that relates to the context input using the compressed layer output embeddings for the segments for the layers.

16. A method performed by one or more computers, the method comprising:receiving a context input comprising a sequence of context tokens;generating a modified sequence of context tokens, comprising modifying the sequence of context tokens to insert a plurality of compression tokens into the sequence of context tokens;processing the modified sequence of context tokens using a neural network to, for each of a plurality of layers of the neural network, generate a respective layer output embedding of each of the compression tokens, wherein the neural network comprises a selfattention layer that applies self-attention and wherein, during the self-attention (i) the context tokens do not attend to the compression tokens and (ii) each compression token attends to a different subset of the context tokens from at least one of the other compression tokens; and generating a response to a query that relates to the context input using the respective layer output embeddings of the compression tokens for the plurality of layers.Attorney Docket No.: 45288-0626WO117. A method performed by one or more computers, the method comprising:receiving a context input comprising a sequence of context tokens;generating a modified sequence of context tokens, comprising modifying the sequence of context tokens to insert a plurality of compression tokens into the sequence of context tokens;processing the modified sequence of context tokens using a neural network to, for each of a sequence of self-attention layers of the neural network, generate a respective layer output embedding of each of the compression tokens; andgenerating a response to a query that comprises a plurality of query tokens using the respective layer output embeddings of the compression tokens by auto-regressively generating an output sequence of tokens using the neural network, comprising, for each particular token in the output sequence and for a particular self-attention layer in the sequence of self-attention layers:processing (i) respective layer output embeddings of the query tokens generated by an immediately preceding self-attention layer in the sequence, (ii) a respective layer output embedding of the particular token generated by the immediately preceding layer in the sequence, (iii) respective layer output embeddings of any preceding tokens in the output sequence generated by the immediately preceding layer in the sequence, and (iv) the respective layer output embeddings of the compression tokens generated by an earlier layer that is before the immediately preceding layer in the sequence using the particular selfattention layer to generate a layer output embedding of the particular token.

18. The method of claim 17, wherein the particular self-attention layer has index i in the sequence of self-attention layers and the earlier layer has index i-2.

19. The method of claim 17 or claim 18, wherein processing the modified sequence of context tokens using a neural network to, for each of a sequence of self-attention layers of the neural network, generate a respective layer output embedding of each of the compression tokens comprises processing the modified sequence of context tokens using the neural network to, for each of the sequence of self-attention layers of the neural network, generate a respective layer output embedding of each of the context tokens of the sequence of context tokens.Attorney Docket No.: 45288-0626WO120. The method of claim 19, wherein:processing the modified sequence of context tokens using the neural network to, for each of the sequence of self-attention layers of the neural network, generate a respective layer output embedding of each of the context tokens of the sequence of context tokens comprises generating a respective layer output embedding of each of the context tokens using a first set of parameters; andprocessing the modified sequence of context tokens using a neural network to, for each of a sequence of self-attention layers of the neural network, generate a respective layer output embedding of each of the compression tokens comprises generating a respective layer output embedding of each of the compression tokens using a second set of parameters that is different from the first set of parameters.

21. The method of claim 16, wherein the compression tokens of the plurality of compression tokens are interspersed throughout the sequence of context tokens.

22. The method of claim 17, wherein the compression tokens of the plurality of compression tokens are interspersed throughout the sequence of context tokens.

23. The method of claim 16, wherein:the neural network comprises a sequence of self-attention layers;the response to the query comprises a plurality of query tokens; andgenerating the response to the query that relates to the context input using the respective layer output embeddings of the compression tokens for the plurality of layers comprises auto-regressively generating an output sequence of tokens using the neural network, comprising, for each particular token in the output sequence and for a particular self-attention layer in the sequence of self-attention layers:processing (i) respective layer output embeddings of the query' tokens generated by an immediately preceding self-attention layer in the sequence, (ii) a respective layer output embedding of the particular token generated by the immediately preceding layer in the sequence, (iii) respective layer output embeddings of any preceding tokens in the output sequence generated by the immediately preceding layer in the sequence, and (iv) the respective layer output embeddings of the compression tokens generated by an earlier layer that is before the immediately preceding layer in the sequence using the particular selfattention layer to generate a layer output the embedding of the particular token.Attorney Docket No.: 45288-0626WO124. The method of claim 23, wherein processing the modified sequence of context tokens using a neural network to, for each of a pl urality of layers of the neural network, generate a respective layer output embedding of each of the compression tokens comprises processing the modified sequence of context tokens using the neural network to, for each of the plurality' of layers of the neural network, generate a respective layer output embedding of each of the context tokens of the sequence of context tokens.

25. The method of claim 24, wherein:processing the modified sequence of context tokens using the neural network to, for each of the plurality of layers of the neural network, generate a respective layer output embedding of each of the context tokens of the sequence of context tokens comprises generating a respective layer output embedding of each of the context tokens using a first set of parameters; andprocessing the modified sequence of context tokens using a neural network to, for each of a plurality of layers of the neural network, generate a respective layer output embedding of each of the compression tokens comprises generating a respective layer output embedding of each of the compression tokens using a second set of parameters that is different from the first set of parameters.

26. The method of claim 23, wherein each self-attention layer of the neural network comprises an attention head configured to generate one or more key vectors and one or more value vectors for each of the query tokens, the compression tokens, and the tokens in the output sequence of tokens.

27. The method of claim 26, wherein the one or more key and value vectors are stored in a Key -Value cache.

28. The method of claim 27, wherein the Key-Value cache is a specialized memory' buffer or data structure.

29. The method of any preceding claim, wherein the context tokens comprise one or more of video tokens, text tokens, or audio tokens.Attorney Docket No.: 45288-0626WO130. The method of any preceding claim, wherein the context input indicates an instruction for an agent to complete a task and the response to the query comprises the completion of the task.

31. The method of any preceding claim, wherein the query comprises a request for information that is included in information represented by the context input.

32. The method of any preceding claim, wherein the context input is received from a user and the method further comprises providing the response to the query' to the user on a user device.

33. The method of any preceding claim, wherein that modified sequence of context tokens includes at least one compression token that is separated from at least one other compression token by at least one context token.

Images