Neural networks with per layer embedding injection

Abstract

Methods, systems, and apparatus, including computer programs encoded on computer storage media, for processing input sequences using a neural network that includes one or more additional embedding layers.

Description

[0001] NEURAL NETWORKS WITH PER LAYER EMBEDDING INJECTION

[0002] CROSS-REFERENCE TO RELATED APPLICATION

[0003] This application claims priority to U.S. Application No. 63 / 688,252, filed August 28, 2024. The disclosure of the foregoing application is hereby incorporated by reference in its entirety.

[0004] BACKGROUND

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

[0006] Neural networks are machine learning models that employ one or more layers of nonlinear units to predict an output for a received input. Some neural networks include one or more hidden layers in addition to an output layer. The output of each hidden layer is used as input to the next layer in the network, i.e., the next hidden layer or the output layer. Each layer of the network generates an output from a received input in accordance with current value inputs of a respective set of parameters.

[0007] SUMMARY

[0008] This specification describes a system implemented as computer programs on one or more computers in one or more locations that processes sequences of tokens to perform one or more machine learning tasks.

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

[0010] This specification describes a system that processes inputs using a neural network, e.g., an attention neural network that includes attention layer blocks, that makes use of perlayer embedding injection when processing a given input token.

[0011] Generally, when processing an input token, a conventional neural network would embed the input token and then process the embedding of the input token through a sequence of layers of the neural network to generate a network output for the input token.

[0012] By contrast, by making use of per-layer embedding injection, the described techniques generate multiple different embeddings of the given input token, i.e., generate multiple different representations of the given input token, at least one of which is then provided as an additional input to a corresponding layer block within the neural network.

[0013] For example, in some cases, each attention layer block of an attention neural network can have a corresponding embedding layer (also referred to as a “lookup” layer), with the embedding generated by the corresponding embedding layer being provided as an additional input to the corresponding attention layer block.

[0014] By providing these additional inputs to the layer block(s), the described techniques inject additional knowledge into the layer blocks of the neural network, resulting in improved performance of the neural network on the task(s) for which the neural network has been trained, i.e., in higher-quali ty network outputs being generated by the additional neural network.

[0015] Moreover, this improved performance that results from introducing these additional inputs can be achieved with minimal additional computational overhead, e.g., relative to using only the input embedding layer.

[0016] For example, while making use of the additional embedding layers requires additional parameters to be added to the neural network (in the form of the additional embeddings for input tokens maintained by the additional embeddings), incorporating the intermediate embeddings can be done using computationally efficient operations at each layer block, resulting in minimal additional latency being incurred. That is, because these additional parameters are simply looked up from memory rather than computed over (e.g., like the weight matrix of one of the layers of the neural network), they are more computationally efficient to provide, i.e., to incorporate into the operations required to perform inference using the neural network. That is, relative to adding parameters to the neural network by adding additional layers that perform matrix multiplication operations on their inputs, adding additional embedding layers is significantly more computationally efficient because the operations of the additional embedding layers can be implemented as a look-up operation to retrieve the corresponding embeddings from memory, i.e., because a look-up operation is much computationally “cheaper” than a matrix multiplication or other similar operation. Thus, even when a large number of parameters are added to the neural network, the computational cost of adding these parameters remains very low.

[0017] Additionally, many inference computing environments, e.g., those that use hardware accelerators with limited on-chip memory or those that run on an edge device with limited device memory, have significant limitations on the amount of random access memory (RAM) available to perform neural network inference. However, the additional embeddings can be stored on the other memory that is relatively slower, more abundant, and cheaper to access, e.g., flash memory or other types of non-volatile memory or, more generally, secondarystorage, rather than relatively faster, more limited, and expensive to access random access memory (RAM), as only one contiguous tensor needs to be read for each token being processed, resulting in only a single read from the other memory being required per token, which can be accomplished without impacting inference latency. As a result, any impact on the limited amount of RAM available to the neural network by needing to store the additional embeddings is mitigated, i.e., because the additional embeddings can be stored in the other memory described above rather than needing to be stored in RAM like at least some of the other parameters of the neural network. That is. because each new token requires reading only one contiguous tensor that includes all of the embeddings for the token, the additional embeddings can be stored in the other memory without impacting the latency of performing neural network inference, thereby preventing the additional embeddings from having to be stored in RAM.

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

[0019] BRIEF DESCRIPTION OF THE DRAWINGS

[0020] FIG. 1 is a diagram of an example neural network system.

[0021] FIG. 2 is a flow diagram of an example process for processing an input token using the neural network.

[0022] FIG. 3 is a diagram showing an example of the operations performed by the neural network.

[0023] FIG. 4 is a flow diagram of an example process for processing a continuous input token using a neural network.

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

[0025] DETAILED DESCRIPTION

[0026] This specification generally describes techniques for performing machine learning tasks on input sequences using a neural network.

[0027] As will be described below, the neural network can be used to perform any of a variety of tasks that require processing an input sequence that includes a respective input token at each of a plurality of input positions to generate network output.

[0028] FIG. 1 is a diagram of an example neural network system 100. The neural network 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.

[0029] The neural network system 100 is a system that processes an input sequence 102 to perform one or more machine learning tasks using an attention neural network 110. Generally, the input sequence 102 includes a respective input token at each of a plurality of input positions. That is, the system 100 receives an input sequence 102 and processes the input sequence 102 using the neural network 110 to generate a network output 112.

[0030] The neural network 110 can be configured through training to perform any kind of machine learning task, i.e., can be configured to receive any kind of input sequence and to generate any kind of score, classification, or regression output based on the input sequence.

[0031] In some situations, the neural network 110 can be referred to as an auto-regressive neural network, i.e., because the neural network auto-regressively generates an output sequence of tokens. More specifically, the auto-regressively generated output sequence is created by generating each particular token in the output sequence conditioned on a current input sequence that includes at least some of the tokens that precede the particular token in the output sequence, i.e., the tokens that have already been generated for any previous positions in the output sequence that precede the particular position of the particular token. In some cases, the current input sequence also includes one or more tokens representing a conditioning or context input for the output sequence.

[0032] For example, the neural network 1 10 can be an auto-regressive attention neural network that includes (i) a plurality7of attention blocks that each apply a self-attention operation and (ii) an output subnetwork that processes an output of the last attention block to generate the score distribution, e.g., a score distribution used for selecting an output token, e.g., by sampling from the score distribution or selecting a most likely token according to the score distribution. For example, the plurality of attention blocks can include a sequence of attention blocks in which the first attention block of the sequence applies an attention mechanism to an input that includes an “‘input embedding” of an input token, and each attention block, except the first layer block of the sequence, applies an attention mechanism to a corresponding input which includes an output of the preceding attention block of the sequence.

[0033] In this example, the neural network can have any of a variety of Transformer-based neural network architectures. Examples of such architectures include those described in J. Hoffmann, S. Borgeaud, A. Mensch, E. Buchatskaya, T. Cai, E. Rutherford, D. d. L. Casas, L. A. Hendricks, J. Welbl, A. Clark, et al. Training compute-optimal large language models, arXiv preprint arXiv:2203.15556, 2022; J.W. Rae, S. Borgeaud, T. Cai, K. Millican, J. Hoffmann, H. F. Song, J. Aslanides, S. Henderson, R. Ring, S. Young, E. Rutherford, T. Hennigan, J. Menick, A. Cassirer, R. Powell, G. van den Driessche, L. A. Hendricks, M. Rauh, P. Huang, A. Glaese, J. Welbl, S. Dathathri, S. Huang, J. Uesato, J. Mellor, I. Higgins, A. Creswell, N. McAleese, A.Wu, E. Eisen, S. M. Jayakumar, E. Buchatskaya, D. Budden, E. Sutherland. K. Simonyan, M. Paganini, L. Sifre, L. Martens, X. L. Li, A. Kuncoro. A. Nematzadeh, E. Gribovskaya, D. Donato, A. Lazaridou, A. Mensch, J. Lespiau, M. Tsimpoukelli, N. Grigorev, D. Fritz, T. Sottiaux, M. Pajarskas, T. Pohlen, Z. Gong, D. Toyama, C. de Masson d’Autume. Y. Li, T. Terzi, V. Mikulik, I. Babuschkin, A. Clark, D. de Las Casas, A. Guy, C. Jones, J. Bradbury, M. Johnson, B. A. Hechtman, L. Weidinger. I. Gabriel, W. S. Isaac, E. Lockhart, S. Osindero, L. Rimell, C. Dyer, O. Vinyals, K. Ayoub, J. Stanway, L. Bennett, D. Hassabis, K. Kavukcuoglu, and G. Irving. Scaling language models: Methods, analysis & insights from training gopher. CoRR, abs / 2112. 11446, 2021; Colin Raffel, Noam Shazeer. Adam Roberts, Katherine Lee. Sharan Narang, Michael Matena, Yanqi Zhou. Wei Li, and Peter J Liu. Exploring the limits of transfer learning with a unified text-to-text transformer. arXiv preprint arXiv: 1910.10683, 2019; Daniel Adiwardana, Minh- Thang Luong, David R. So, Jamie Hall, Noah Fiedel, Romal Thoppilan, Zi Yang, Apoorv Kulshreshtha, Gaurav Nemade, Yifeng Lu, and Quoc V. Le. Towards a human-like opendomain chatbot. CoRR, abs / 2001.09977, 2020; and Tom B Brown. Benjamin Mann. Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sas try. Amanda Askell, et al. Language models are few-shot learners. arXiv preprint arXiv:2005. 14165, 2020.

[0034] More specifically, the neural network includes plurality of layer blocks 120. The neural network can also include an output subnetwork as described above.

[0035] A layer block, as used in this specification, is a collection of one or more neural network layers.

[0036] In the example of FIG. 1, the plurality of layer blocks 120 include one or more attention layer blocks. For example, at least one layer block 120. and optionally multiple ones or all of the layer blocks 120, may include one (or more) of the attention blocks (e g. as a first sub-layer of the layer block 120) as described above.

[0037] During training (and during inference when inference is not auto-regressive), the neural network 110 receives an input sequence that includes a respective input token at each of a plurality of input positions. The neural network 110 processes the input sequence to generate a network output.

[0038] During auto-regressive sequence generation after training, the system 100 auto- regressively generates an output sequence that includes a respective output token at each of a plurality of output positions.

[0039] To do this, for each output position (after the very first output position in the sequence), the system 100 receives a preceding output token (“new input token”) generated for the preceding output position and processes the preceding output token using the neural network to generate a network output 112 that specifies the output token at the output position.

[0040] During this processing, a given layer block 120 receives an embedding for the new input token and generates as output an output embedding for the new input token, i.e., processes the embedding for the new input token to update the embedding for the new input token.

[0041] An embedding of a given input is an ordered collection of numerical values, e.g., a vector of floating point or other numerical values.

[0042] When the layer block 120 is an attention layer block, the layer block 120 updates the embeddings by, at least in part, applying an attention mechanism, e.g., using an attention sublayer.

[0043] A “sub-layer” as used in this specification is a component within a layer block that performs a corresponding subset of the operations of the layer block.

[0044] Generally, each attention mechanism uses one or more attention heads.

[0045] Each attention head generates a set of queries, a set of keys, and a set of values (for example, as the respective products of the input to the attention head with a query matrix, a key matrix and a value matrix associated with the attention head), and then applies any of a variety of variants of query -key -value (QKV) attention, e.g., a dot product attention function or a scaled dot product attention function, using the queries, keys, and values to generate an output. Each query, key, value can be a vector that includes one or more vector elements. “Self-attention” means that all of the queries, keys and values are derived from the same input sequence.

[0046] When there are multiple attention heads, the attention sub-layer then combines the outputs of the multiple attention heads, e.g., by concatenating the outputs and, optionally, processing the concatenated outputs through a linear layer. The attention mechanisms can be local or global (or some attention mechanisms can be local while others for other layer blocks are global). For local attention mechanisms, for each position, the positions (in the input to the attention mechanism) that are used to generate the queries, keys, and values for the position are defined by a local window size for the local attention mechanism, i.e., non-zero attention weights for a given position are computed only for positions that are within the local w indow of the given position, w here the local window is composed of all the input positions which are no more than the local window7size before the current input position.

[0047] In some cases, because the attention applied by the attention layers is causal, the system 100 can store, in memory and for any given attention mechanism and when generating the output for any given input position, the embeddings or the keys and values already computed for earlier input positions (i.e., for the “context” tokens that precede the current token in the sequence) rather than re-computing the embeddings (or the keys and values) for earlier input positions. Storing the keys and values in a memory is also referred to as maintaining a KV cache.

[0048] Some or all of the layer blocks in the neural network can also include other types of sub-layers, e.g., normalization layers, residual connection layers, feedforward sub-layers, and so on.

[0049] In some cases, some or all of the feedforward sub-layers within the layer blocks in the neural network are implemented as sparse mixture of experts (MoE) layers while in other cases all the feedforward layers are dense multi-layer perceptrons (MLPs).

[0050] In some cases, the neural network 110 uses an input embedding layer 130 to process the input token, i.e., to generate an input embedding of the token. The neural network 110 then provides the input embedding generated by the input embedding layer 130 as the embedding for the new7token that is processed as input by the first layer block in the sequence. In some other cases, e.g., for continuous tokens as will be described below7, the input token is a continuous or “soft” token, and the neural network 110 projects the input token, e.g., using a learned projection layer, to generate the input embedding of the token.

[0051] Unlike conventional neural netw orks, how ever, the neural netw ork 110 also includes or, equivalently, receives outputs generated by one or more additional embedding layer(s) 140. Each additional embedding layer 140 corresponds to one of the layer blocks 120, i.e., to one of the layer blocks 120 in a set of one or more of the layer blocks 120 within the neural network 110.

[0052] In some implementations, the set of one or more layer blocks 120 includes all of the layer blocks 120. i.e., so that each layer block has a corresponding additional embedding layer 140. In some other implementations, the set of one or more layer blocks 120 is a proper subset of the layer blocks 120. i.e., so that not all of the layer blocks 120 have a corresponding additional embedding layer 140. For example, the set can include every other layer block 120 or every third layer block 120 within the neural network 110.

[0053] Generally, the set of one or more layer blocks 120 includes one or more of the subsequent layer blocks 120, i.e., the layer blocks 120 after the sequence of layer blocks 120, and, optionally, the first layer block 120 in the sequence.

[0054] As part of processing the token, the neural network 110 processes the input token using each of the additional embedding layer(s) 140 to generate a respective intermediate embedding of the token for each additional embedding layer 140. As will be described below, processing the input token using a given embedding layer can be implemented as a lookup operation that looks up the intermediate embedding of the token. Moreover, in some cases, the intermediate embeddings for all of the additional embedding layer(s) can be looked up in a single operation.

[0055] Then, as part of processing the token, each layer block 120 in the set of one or more layer blocks processes both (i) the respective input embedding of the input token, i.e., the respective input embedding after being updated by the preceding layer of the neural network 110, and (ii) the respective intermediate embedding for the input token for the layer block 120 generated by the additional embedding layer for the layer block 120 to generate the output embedding of the input token. That is. the neural network 110 “injects” an additional embedding as an additional input to the layer block 120, thereby injecting additional information into the layer block 120.

[0056] Some examples of machine learning tasks that a neural network when implemented using one of the architectures described above or other known architectures can be configured to perform follow.

[0057] In any of the implementations below, the neural network may be deployed as part of a chat hot, 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.

[0058] 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. 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 identitylocations 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 intensity value inputs for the pixels of an image.

[0059] 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 w ords, w ordpieces or characters of the first language, and the vocabulary for the output tokens may be w ords, w-ordpieces 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 languages - target language pairs. In this example, the source language text may be augmented with an identifier that indicates the target language into which the neural netw ork should translate the source language text.

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

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

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

[0063] 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 (e g. comprising at least one parameter of the patient measured numerically, such as a temperature, blood pressure measurement, and / or at least one medical image captured from the patient, such as an X-ray, ultrasound or MRI image) 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 (e.g. at a sequence of corresponding times), 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 prediction being a predicted future blood glucose measurement or the prediction of a hyper- or hypo-glycemic 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. 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.

[0064] In some implementations the input sequence represents data to be compressed, e.g., image data, text data, audio data, or any other t pe 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 generat ed / 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.

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

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

[0067] 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, j oint 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.

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

[0069] 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 be 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.

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

[0071] 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 learn not to instruct the user to perform the identified actions, i.e., ones which the user is likely to perform incorrectly.

[0072] 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 for the 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 series the digital assistant can be used to output to the user an indication of the task, e.g., step or sub-task, 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 successfully 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 rewards 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.

[0073] 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 above-described 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 observations of the user performing the task, determine from the abovedescribed 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.

[0074] 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 dow nstream task. Examples of dow nstream tasks include promoter site prediction, methylation analysis, predicting functional effects of non-coding variants, and so on.

[0075] 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 individual 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.

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

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

[0078] 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 would 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.

[0079] Also or instead the network output 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.

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

[0081] More generally, the task to be performed by the neural network can be specified by the input sequence. As a particular example, the input sequence can include a prompt or an instruction that specifies the task that is to be performed by the neural network. Optionally, in this example, the input sequence also includes context for performing the task.

[0082] In general in implementations of the described techniques the input data, e.g., text, audio, and / or an image or video, may be encoded into a sequence of input tokens in any convenient manner; and output tokens may be similarly decoded into text, audio, and / or image or video data according to the particular task or tasks to be performed.

[0083] FIG. 2 is a flow diagram of an example process 200 for processing an input token using a neural network. For convenience, the process 200 will be described as being performed by a system of one or more computers located in one or more locations. For example, a neural network system, e.g., the neural network system 100 depicted in FIG. 1, appropriately programmed in accordance with this specification, can perform the process 200.

[0084] As described above, the neural network includes a sequence of layer blocks and a plurality7of embedding layers. For example, some or all of the layer blocks in the sequence can be attention layer blocks that apply an attention mechanism over corresponding embeddings of the input token and a respective set of context tokens as part of generating the output of the layer block. For self-attention layer blocks, the context tokens are tokens that precede the input token in the current input sequence. For cross-attention layer blocks, the context tokens are tokens from a context input that is being processed with the input token.

[0085] The plurality of embedding layers includes a respective embedding layer for each layer block in a set of one or more of the layer blocks. As described above, the plurality of embedding layers can also include an input embedding layer that generates an input embedding of any given input token to the neural network, and which is processed by the first layer block of the sequence.

[0086] The system receives an input that includes an input token for processing using the neural network (step 202). During training or when the system is performing non-auto- regressive inference, the input can also include one or more other tokens from the input sequence. During auto-regressive inference, the system can also obtain the earlier tokens in the current input sequence or can maintain, e.g.. in a KV cache or other memory, keys and values corresponding to the earlier tokens in the current input sequence.

[0087] The system processes the input using the neural network to generate a network output (step 204). For example, during auto-regressive inference, the network output can indicate the next token to follow the input token in the current input sequence. During training or during non-autoregressive inference, the network output can be a different type of network output.

[0088] As part of the processing, for each of the one or more layer blocks in the set, the system processes the input token using the respective embedding layer for the layer block to generate a respective intermediate embedding for the input token (step 206).

[0089] Generally, each embedding layer corresponds to an embedding table (also referred to as an embedding “matrix”) that maps each token in at least a portion of vocabulary7of tokens to a respective embedding.

[0090] In some implementations, rather than maintain separate embedding tables, the system can maintain a single mapping that maps each input token in the (portion of) the vocabulary to a tensor that includes all of the embeddings for the input token for all of the layer blocks in the set. That is, the mapping maps a respective identifier for each token to the tensor for the token that includes, e.g., is a concatenation of, the respective embeddings of the token for all of the layer blocks in the set. As will be described in more detail below, these embedding can be learned as part of the training of the neural network.

[0091] The system can then use the single mapping to read out the single contiguous tensor out.

[0092] As a result, the system can effectively store the single mapping in secondary storage or other relatively slower memory, e.g., Flash or other non-volatile memory, rather than needing to store the single mapping in random access memory' (RAM). That is, the system can perform a single look-up operation to read out all of the intermediate embeddings for the input token.

[0093] The system can also process the input token using the input embedding layer to generate an input embedding for the input token and then provide the input embedding as the respective layer input embedding of the input token for the first layer block of the sequence. In some implementations, the system looks up the input embedding as part of the single lookup operation performed to read out the intermediate embeddings for the input tokens. In some other implementations, the system maintains a separate embedding table for the input embedding layer.

[0094] For each of the sequence of layer blocks, the system processes a respective layer input embedding of the input token to generate a respective layer output embedding of the input token (step 208). That is, for the first layer block of the sequence, the respective layer input embedding of the input token is the input embedding generated by the input embedding layer. For each subsequent layer block, i.e., for each layer block that is after the first layer block of the sequence, the respective layer input embedding is the respective layer output embedding of the input token generated by the preceding layer block in the sequence.

[0095] More specifically , for each layer block in the set of one or more layer blocks, processing the respective layer input embedding of the input token to generate the respective layer output embedding of the input token includes processing (i) the respective layer input embedding of the input token and (ii) the respective intermediate embedding for the input token for the layer block to generate the respective layer output embedding.

[0096] A given layer block can perform this processing in any of a variety of ways that incorporate the intermediate embedding as additional information. For example, the layer block can process the respective layer input embedding of the input token using one or more sub-layers to generate a respective intermediate output embedding of the input token and then process the respective intermediate output embedding and the respective intermediate embedding for the input token for the layer block to generate the respective layer output embedding. For example, for a self-attention layer block, the one or more sub-layers can include a self-attention sub-layer, optionally followed by a feedforward sub-layer.

[0097] The layer block can combine the respective intermediate output embedding and the respective intermediate embedding for the input token for the layer block in any of a variety of ways.

[0098] For example, in some cases, to limit the additional number of parameters added to the neural network, the additional embeddings can have a smaller dimensionality than the intermediate output embedding.

[0099] In these cases, the system can generate a projected embedding from the respective intermediate embedding for the input token, e g., by applying a projection weight matrix to the respective intermediate embedding for the input token and then combine the projected embedding with the respective intermediate output embedding to generate the respective layer output embedding, e.g., by summing the projected embedding and the respective intermediate output embedding.

[0100] Another example of combining the respective intermediate output embedding and the respective intermediate embedding for the input token for the layer block is described below with reference to FIG. 3.

[0101] The system then processes at least the respective layer output embedding of the input token generated using the last layer block in the sequence to generate the network output (step 210). For example, the system can process the respective layer output embedding of the input token generated using the last layer block in the sequence using an output layer block of the neural netw ork to generate a respective score for each token in a vocabulary' of tokens. The system can then select a token from the vocabulary of tokens using the respective scores.

[0102] Generally, the plurality of embedding layers and the layer blocks in the sequence of layer blocks have been trained jointly during the training of the neural network. For example, the plurality of embedding layers and the layer blocks in the sequence of layer blocks can have been learned during pre-training of the neural netw ork and then fine-tuned during finetuning of the neural network. As another example, the neural network can have been pre- trained with only the input embedding layer, and then the additional embedding layers can be trained during fine-tuning of the neural network.

[0103] In more detail, the neural network, i.e., the plurality of embedding layers and the layer blocks and any other learned components, can have been trained across one or more training stages.

[0104] For example, the one or more training stages can include a pre-training stage. During the pre-training stage, the neural network can have been trained by the system or a separate system on a next token prediction task, e.g., a task that requires predicting, given a current sequence of tokens, the next token that follows the current sequence in the training data.

[0105] As a particular example, the neural network can have been trained on a maximumlikelihood objective on a large dataset of text in one or more natural languages, e.g., text that is publicly available from the Internet or another text corpus, a large dataset of computer code in one or more programming languages, e.g., Python, C++, C#, Java, Ruby, PHP, and so on, e.g., computer code that is publicly available from the Internet or another code repository, a large dataset of audio samples, e.g., audio recordings or waveforms that represent the audio recordings, a large dataset of images where each image includes an array of pixels, a large dataset of videos where each video includes a temporal sequence of frames, or a large multimodal dataset that includes a combination of two or more of these datasets.

[0106] As another example, the one or more previous training stages can include one or more additional training stages (“fine-tuning” or “post-training” stages), e.g.. that occur after the pre-training stage. For example, the one or more previous training stages can include any one or more of: a supervised fine-tuning stage, a reinforcement learning stage, a preference learning stage, an instruction tuning stage, and so on.

[0107] FIG. 3 is a diagram that shows an example 300 of the operations performed by the neural network 110. In the example 300, the neural network includes K layer blocks 120. Additionally, each layer block operates on a batch of T tokens that have dimensionality D.

[0108] The input embedding to the first layer block 120 in the sequence is generated by processing the new token using the input embedding layer 130 (that uses an embedding matrix E) and then applying a positional embedding 302 to the output of the input embedding layer 130. This application of positional embedding 302 can also be considered to be part of the operations performed by the input embedding layer 130.

[0109] In particular, in the example 300, each of the K layer blocks 120 is an attention layer block that includes an attention sub-layer 310 and a feedforward sub-layer 320. Moreover, the neural network 110 includes a respective additional embedding layer 140 for each of the layer blocks 120. The additional embedding layer 140 maintains a respective X dimensional embedding for each of the V tokens in the vocabulary. Thus, introducing an additional embedding layer 140 for each of the K blocks results in K * V * X additional model parameters. As shown in the example 300, the respective additional embedding layers 140 can be implemented as a single mapping that maps an identifier for each of the V tokens to all T A dimensional embeddings of the token.

[0110] As shown in the example 300, each layer block processes the respective layer input embedding of the input token 112 using one or more sub-layers to generate a respective intermediate output embedding of the input token and then processes the respective intermediate output embedding and the respective intermediate embedding for the input token for the layer block to generate the respective layer output embedding.

[0111] More specifically, as indicated above, the one or more sub-layers include a selfattention sub-layer 310 and a feed-forward sub-layer 320. In the example 300, the intermediate output embedding of the input token is the output of a residual connection 322 applied to the output of the feed-forward sub-layer 320 and the output of the self-attention sub-layer 310.

[0112] Additionally, in the example 300, the respective intermediate embeddings have a different dimensionality X than the respective layer input and respective layer output embeddings. For example, as described above, to improve the computational efficiency of the neural network, i.e., to mitigate the impact of adding the additional embedding layers, the respective intermediate embeddings can have a smaller dimensionality than the respective layer input and respective layer output embeddings.

[0113] To process the respective intermediate output embedding and the respective intermediate embedding for the input token for the layer block, the layer block 120 generates a gating embedding from the respective intermediate output embedding. As part of this, the layer block 120 applies a gating weight matrix G 330 to the respective intermediate output embedding. Optionally, the layer block 120 can also apply an activation function, e.g., a rectified linear unit (ReLU) activation function 370. to the output of applying the gating w eight matrix to the respective intermediate output embedding.

[0114] The layer block 120 then combines the gating embedding and the respective intermediate embedding to generate a gated embedding. As shown in the example 300, the layer block 120 "gates” the intermediate embedding with the gating embedding by element- wise multiplying 360 the gating embedding and the respective intermediate embedding to generate the gated embedding.

[0115] Once the gated embedding has been generated, the layer block 120 generates a projected embedding from the gated embedding, e.g., by applying a projection weight matrix P 340 to the gated embedding.

[0116] The layer block 120 then combines the projected embedding with the respective intermediate output embedding to generate the respective layer output embedding, e.g. by summing 350 the projected embedding with the respective intermediate output embedding to generate the respective layer output embedding.

[0117] Descriptions of input tokens that can be processed by the neural network now follows.

[0118] For example, each input token can be a vector of a specified dimensionality and can be generated by processing a corresponding input item representing a portion of the input text, audio, image, video, and so on, e.g., using a tokenizer or a separate encoder neural network.

[0119] For example, when the input sequence includes tokens that represent a visual input, e.g., a video or an image, the system can receive the visual input and process the visual input using a visual encoder neural network to generate the tokens representing the visual input. For example, the visual encoder neural network can be a variant of a vision Transformer or of a convolutional neural and can generate either continuous or discrete tokens. Examples of visual encoders include VQ-VAE. ViT-VQGAN, NaViT. or CoCa. The visual encoder neural network can be pre-trained prior to the training of the attention neural network and then held fixed during the training of the attention neural network, can be pre-trained prior to the training of the attention neural network and then fine-tuned during the training of the attention neural network, or can be trained from scratch jointly with the attention neural network (here ‘'jointly” means that that updates to the visual encoder neural network are interleaved with, or simultaneous with, updates to the attention neural network). As another example, the system can generate the tokens representing the video input by applying a specified set of transformations to patches of the visual input, e.g., splitting, dimensionality’ reduction, quantization and so on.

[0120] For example, when the input sequence includes tokens that represent an audio input, the system can receive the audio input and process the audio input using an audio encoder neural network to generate the tokens representing the audio input. For example, the audio encoder neural network can include any one of BeST-RQ, aUSM encoder, a SoundStream encoder, or an AudioLM encoder. The audio encoder neural network can be pre-trained prior to the training of the attention neural network and then held fixed during the training of the attention neural network, can be pre-trained prior to the training of the attention neural network and then fine-tuned during the training of the attention neural network, or can be trained from scratch jointly with the attention neural network. As another example, the system can generate the tokens representing the audio input by applying a specified set of transformations to segments of the audio input, e.g.. splitting, dimensionality reduction, quantization and so on.

[0121] When the input sequence is a multi-modal input sequence that includes tokens representing multiple different modalities, the tokens can be arranged within the input sequence in any appropriate way. In some cases, the sequence can include tokens of one modality followed by tokens of another modality. For example, when the sequence represents a text question about an input of another modality, e.g., audio, images, or video, the tokens representing the input of the other modality' can be followed by the tokens representing the text question in the input sequence. In some cases, the sequence can include tokens of multiple modalities interleaved with one another. For example, the sequence can represent a multi-modal query that uses one modality to make multiple references to instances of data of another modality, e.g., “In this image [image 1], the cat is in the suitcase. Where is it in this image? [image 2].” In this example, the tokens for the different modalities can be interleaved to provide context as to which tokens of one modality refer to which tokens of the other modality.

[0122] By being trained on multi-modal sequences arranged as described above, the neural network can effectively process different input sequences with different arrangements of input tokens from different modalities after training.

[0123] When the network outputs generated by the system are predictions of tokens, this can refer to a prediction of the input item represented by the token, e.g., so that the system can generate outputs that include text, images, audio, or videos or multi-modal outputs.

[0124] In some cases, e.g., when the tokens are text tokens or when the encoder neural network generates quantized outputs, the tokens that are input to the neural network are discrete tokens. In these cases, the tokens can be processed using embedding layers as described above.

[0125] In some other cases, the tokens representing one or more other modalities other than text can be continuous tokens, i.e., are vectors that are not constrained to belong to a fixed vocabulary of tokens. In some of these cases, the system only maintains additional embedding lay ers for discrete tokens and does not incorporate additional embeddings of continuous tokens when performing inference.

[0126] In some of these cases, for a given type of continuous token, i.e., for a given encoder neural network, the system maintains a respective projection neural network, e.g., a single projection layer or a multi-layer perceptron (MLP), for each of the one or more layer blocks that have corresponding additional embedding layers. The system can then use these projection neural networks to generate the intermediate embeddings. This is described in more detail below with reference to FIG. 4.

[0127] FIG. 4 is a flow diagram of an example process 400 for processing a continuous input token using a neural network. For convenience, the process 400 will be described as being performed by a system of one or more computers located in one or more locations. For example, a neural network system, e.g., the neural network system 100 depicted in FIG. 1, appropriately programmed in accordance with this specification, can perform the process 400.

[0128] As described above, the neural network includes a sequence of layer blocks and a plurality of embedding layers. For example, some or all of the layer blocks in the sequence can be attention layer blocks that apply an attention mechanism as part of generating the output of the layer block. The neural network also includes one or more projection neural networks, one for each layer block that has a corresponding embedding layer as described above. In some implementations, the system maintains multiple different projection neural networks for each of these layer blocks, one for each different type of continuous token that can be processed by the neural network.

[0129] The system receives an input that includes an input token for processing using the neural network (step 402). Generally, the input token is a continuous token, e.g., a token that represents a portion of an image, video, or audio signal. During training or when the system is performing non-auto-regressive inference, the input can also include one or more other tokens from the input sequence. During auto-regressive inference, the system can also obtain the earlier tokens in the current input sequence or can maintain, e.g., in a K.V cache or other memory, keys and values corresponding to the earlier tokens in the current input sequence.

[0130] The system processes the input using the neural network to generate a network output (step 404). For example, during auto-regressive inference, the network output can indicate the next token to follow the input token in the current input sequence. During training or during non-autoregressive inference, the network output can be a different type of network output.

[0131] As part of the processing, for each of the one or more layer blocks in the set, the system processes the input token using the respective projection neural network for the layer block to generate a respective intermediate embedding for the input token (step 406). As described above, the projection neural network can be a single projection layer or can be a multi-layer perceptron (MLP).

[0132] For each of the sequence of layer blocks, the system processes a respective layer input embedding of the input token to generate a respective layer output embedding of the input token (step 408). That is, for the first layer block of the sequence, the respective layer input embedding of the input token is the input embedding generated by the input embedding layer. For each subsequent layer block, i.e., for each layer block that is after the first layer block of the sequence, the respective layer input embedding is the respective layer output embedding of the input token generated by the preceding layer block in the sequence.

[0133] More specifically, for each layer block in the set of one or more layer blocks, processing the respective layer input embedding of the input token to generate the respective layer output embedding of the input token includes processing (i) the respective layer input embedding of the input token and (ii) the respective intermediate embedding for the input token for the layer block, i.e., generated by the projection neural network for the layer block, to generate the respective layer output embedding.

[0134] A given layer block can perform this processing in any of a variety of ways that incorporate the intermediate embedding as additional information. For example, the layer block can perform the processing described above with reference to FIGS. 2 or 3.

[0135] The system then processes at least the respective layer output embedding of the input token generated using the last layer block in the sequence to generate the network output (step 410). For example, the system can process the respective layer output embedding of the input token generated using the last layer block in the sequence using an output layer block of the neural network to generate a respective score for each token in a vocabulary of tokens. The system can then select a token from the vocabulary of tokens using the respective scores.

[0136] Generally, the projection neural network(s) can be trained jointly with the plurality of embedding layers and the layer blocks in the sequence of layer blocks during the training of the neural network, e.g., using one of the training schemes described above.

[0137] This specification uses the term "‘configured7’ 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.

[0138] 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, e.g., one or more modules of computer program instructions encoded on a tangible non transitory storage medium for execution by, or to control the operation of, data processing apparatus. The computer storage medium can be a machine- readable storage device, a machine-readable storage substrate, a random or serial access memory device, or a combination of one or more of them. Alternatively or in addition, the program instructions can be encoded on an artificially generated propagated signal, e.g., a machine-generated electrical, optical, or electromagnetic signal, that is generated to encode information for transmission to suitable receiver apparatus for execution by a data processing apparatus.

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

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

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

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

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

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

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

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

[0147] 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, e.g., inference, workloads.

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

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

[0150] The computing system can include clients and servers. A client and server are generally remote from each other and ty pically 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.

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

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

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

[0154] Aspects of the present disclosure may be as set out in the following clauses:

[0155] Clause 1. A method performed by one or more computers, the method comprising: receiving an input comprising an input token for processing using a neural network, wherein: the neural network comprises a sequence of layer blocks and a plurality of embedding layers, and the plurality of embedding layers comprises a respective embedding layer for each layer block in a set of one or more of the layer blocks; and processing the input using the neural network to generate a network output, wherein the processing comprises: for each of the one or more layer blocks in the set, processing the input token using the respective embedding layer for the layer block to generate a respective intermediate embedding for the input token; for each of the sequence of layer blocks, processing a respective layer input embedding of the input token to generate a respective layer output embedding of the input token, and wherein, for each layer block in the set of one or more layer blocks, processing the respective layer input embedding of the input token to generate the respective layer output embedding of the input token comprises: processing the respective layer input embedding of the input token and the respective intermediate embedding for the input token for the layer block to generate the respective layer output embedding; and processing at least the respective layer output embedding of the input token generated using a last layer block in the sequence to generate the network output.

[0156] Clause 2. The method of clause 1, wherein: the input token is a last input token in a current input sequence; and the network output identifies a token that is predicted to follow the input token in the current input sequence. Clause 3. The method of clause 2, wherein the neural network further comprises an output layer block and wherein processing at least the respective layer output embedding of the input token generated using a last layer block in the sequence to generate the network output comprises: processing the respective layer output embedding of the input token generated using the last layer block in the sequence to generate a respective score for each token in a vocabulary of tokens; and selecting a token from the vocabulary of tokens using the respective scores.

[0157] Clause 4. The method of any preceding clause, wherein the plurality of embedding layers further comprises an input embedding layer, wherein processing the input using the neural network to generate a network output comprises processing the input token using the input embedding layer to generate an input embedding for the input token, and wherein the respective layer input embedding of the input token for a first layer block of the sequence is the input embedding for the input token.

[0158] Clause 5. The method of clause 4, wherein the respective layer input embedding of the input token for each layer block that is after the first layer block of the sequence is the respective layer output embedding of the input token generated by the preceding layer block in the sequence.

[0159] Clause 6. The method of any preceding clause, wherein the sequence of layer blocks comprise one or more attention layer blocks that each apply an attention mechanism over the input token and a respective set of context tokens.

[0160] Clause 7. The method of clause 5, when dependent on clause 2, wherein for one or more of the attention layer blocks, the respective set of context tokens comprises one or more tokens from the current input sequence.

[0161] Clause 8. The method of any preceding clause, wherein processing the respective layer input embedding of the input token and the respective intermediate embedding for the input token for the layer block to generate the respective layer output embedding comprises: processing the respective layer input embedding of the input token using one or more sub-layers to generate a respective intermediate output embedding of the input token; and processing the respective intermediate output embedding and the respective intermediate embedding for the input token for the layer block to generate the respective layer output embedding.

[0162] Clause 9. The method of clause 8, wherein the one or more sub-layers comprise a selfattention sub-layer.

[0163] Clause 10. The method of clause 9, wherein the one or more sub-layers comprise a feedforward sub-layer.

[0164] Clause 11. The method of any one of clauses 8-10, wherein processing the respective intermediate output embedding and the respective intermediate embedding for the input token for the layer block to generate the respective layer output embedding comprises: generating a gating embedding from the respective intermediate output embedding, comprising applying a gating weight matrix to the respective intermediate output embedding; and combining the gating embedding and the respective intermediate embedding to generate a gated embedding.

[0165] Clause 12. The method of clause 11, wherein processing the respective intermediate output embedding and the respective intermediate embedding for the input token for the layer block to generate the respective layer output embedding further comprises: generating a projected embedding from the gated embedding, comprising applying a projection weight matrix to the gated embedding; and combining the projected embedding with the respective intermediate output embedding to generate the respective layer output embedding. Clause 13. The method of clause 12, wherein combining the projected embedding with the respective intermediate output embedding to generate the respective layer output embedding comprises: summing the projected embedding with the respective intermediate output embedding to generate the respective layer output embedding.

[0166] Clause 14. The method of any one of clauses 11-13, wherein combining the gating embedding and the respective intermediate embedding to generate a gated embedding comprises: element-wise multiplying the gating embedding and the respective intermediate embedding to generate the gated embedding.

[0167] Clause 15. The method of any one of clauses 11-14, wherein generating a gating embedding from the respective intermediate output embedding further comprises: applying an activation function to an output of applying a gating weight matrix to the respective intermediate output embedding.

[0168] Clause 16. The method of clause 15. wherein the activation function is a rectified linear unit (ReLU) activation function.

[0169] Clause 17. The method of any one of clauses 8-10, wherein processing the respective intermediate output embedding and the respective intermediate embedding for the input token for the layer block to generate the respective layer output embedding comprises: generating a projected embedding from respective intermediate embedding for the input token, comprising applying a projection weight matrix to the respective intermediate embedding for the input token; and combining the projected embedding with the respective intermediate output embedding to generate the respective layer output embedding.

[0170] Clause 18. The method of clause 17, wherein combining the projected embedding with the respective intermediate output embedding to generate the respective layer output embedding comprises: summing the projected embedding with the respective intermediate output embedding to generate the respective layer output embedding. Clause 19. The method of any preceding clause, wherein the respective intermediate embeddings have a different dimensionality than the respective layer input and respective layer output embeddings.

[0171] Clause 20. The method of any preceding clause, wherein the plurality of embedding layers and the layer blocks in the sequence of layer blocks have been trained jointly.

[0172] Clause 21. The method of any preceding clause, wherein the set of one or more layer blocks comprises all of the layer blocks in the sequence.

[0173] Clause 22. The method of any preceding clause, wherein the set of one or more layer blocks comprises a plurality of layer blocks, and wherein the respective embedding layers for each of the layer blocks are implemented as a single mapping that maps the input token to a tensor that includes the respective intermediate embedding for the input token for each of the plurality of layer blocks in the set.

[0174] Clause 23. The method of clause 22, wherein for each of the one or more layer blocks in the set, processing the input token using the respective embedding layer for the layer block to generate a respective intermediate embedding for the input token comprises: reading, using the single mapping, the tensor that includes the respective intermediate embedding for the input token for each of the plurality of layer blocks in the set.

[0175] Clause 24. The method of clause 23. wherein the single mapping from which the tensor is read is stored in non-volatile memory.

[0176] Clause 25. The method of clause 24. wherein the non-volatile memory is Flash memory.

[0177] Clause 26. A system comprising one or more computers and one or more storage devices storing instructions that when executed by the one or more computers cause the one or more computers to perform the operations of the respective method of any one of clauses 1-25.

[0178] Clause 27. One or more computer storage media storing instructions that when executed by one or more computers cause the one or more computers to perform the operations of the respective method of any one of clauses 1-25.

[0179] What is claimed is:

Claims

1. CLAIMS1. A method performed by one or more computers, the method comprising: receiving an input comprising an input token for processing using a neural network, wherein: the neural network comprises a sequence of layer blocks and a plurality of embedding layers, and the plurality of embedding layers comprises a respective embedding layer for each layer block in a set of one or more of the layer blocks; and processing the input using the neural network to generate a network output, wherein the processing comprises: for each of the one or more layer blocks in the set, processing the input token using the respective embedding layer for the layer block to generate a respective intermediate embedding for the input token; for each of the sequence of layer blocks, processing a respective layer input embedding of the input token to generate a respective layer output embedding of the input token, and wherein, for each layer block in the set of one or more layer blocks, processing the respective layer input embedding of the input token to generate the respective layer output embedding of the input token comprises: processing the respective layer input embedding of the input token and the respective intermediate embedding for the input token for the layer block to generate the respective layer output embedding; and processing at least the respective layer output embedding of the input token generated using a last layer block in the sequence to generate the network output.

2. The method of claim 1, wherein: the input token is a last input token in a current input sequence; and the network output identifies a token that is predicted to follow the input token in the current input sequence.

3. The method of claim 2, wherein the neural network further comprises an output layer block and wherein processing at least the respective layer output embedding of the input token generated using a last layer block in the sequence to generate the network output comprises: processing the respective layer output embedding of the input token generated using the last layer block in the sequence using the output layer block to generate a respective score for each token in a vocabulary of tokens; and selecting a token from the vocabulary of tokens using the respective scores.

4. The method of any preceding claim, wherein the plurality of embedding layers further comprises an input embedding layer, wherein processing the input using the neural network to generate a network output comprises processing the input token using the input embedding layer to generate an input embedding for the input token, and wherein the respective layer input embedding of the input token for a first layer block of the sequence is the input embedding for the input token.

5. The method of claim 4, wherein the respective layer input embedding of the input token for each layer block that is after the first layer block of the sequence is the respective layer output embedding of the input token generated by the preceding layer block in the sequence.

6. The method of any preceding claim, wherein the sequence of layer blocks comprise one or more attention layer blocks that each apply an attention mechanism over corresponding embeddings of the input token and a respective set of context tokens.

7. The method of claim 5, when dependent on claim 2, wherein for one or more of the attention layer blocks, the respective set of context tokens comprises one or more tokens from the current input sequence.

8. The method of any preceding claim, wherein processing the respective layer input embedding of the input token and the respective intermediate embedding for the input token for the layer block to generate the respective layer output embedding comprises: processing the respective layer input embedding of the input token using one or more sub-layers of the layer block to generate a respective intermediate output embedding of the input token; and processing the respective intermediate output embedding and the respective intermediate embedding for the input token for the layer block to generate the respective layer output embedding.

9. The method of claim 8, wherein the one or more sub-layers of the layer block comprise a self-attention sub-layer.

10. The method of claim 9, wherein the one or more sub-layers of the layer block comprise a feed-forward sub-layer.

11. The method of any one of claims 8-10, wherein processing the respective intermediate output embedding and the respective intermediate embedding for the input token for the layer block to generate the respective layer output embedding comprises: generating a gating embedding from the respective intermediate output embedding, comprising applying a gating weight matrix to the respective intermediate output embedding; and combining the gating embedding and the respective intermediate embedding to generate a gated embedding.

12. The method of claim 11, wherein processing the respective intermediate output embedding and the respective intermediate embedding for the input token for the layer block to generate the respective layer output embedding further comprises: generating a projected embedding from the gated embedding, comprising applying a projection weight matrix to the gated embedding; and combining the projected embedding with the respective intermediate output embedding to generate the respective layer output embedding.

13. The method of claim 12, wherein combining the projected embedding with the respective intermediate output embedding to generate the respective layer output embedding comprises: summing the projected embedding with the respective intermediate output embedding to generate the respective layer output embedding.

14. The method of any one of claims 1 1-13, wherein combining the gating embedding and the respective intermediate embedding to generate a gated embedding comprises: element-wise multiplying the gating embedding and the respective intermediate embedding to generate the gated embedding.

15. The method of any one of claims 11-14, wherein generating a gating embedding from the respective intermediate output embedding further comprises: applying an activation function to an output of applying a gating weight matrix to the respective intermediate output embedding.

16. The method of claim 15, wherein the activation function is a rectified linear unit (ReLU) activation function.

17. The method of any one of claims 8-10, wherein processing the respective intermediate output embedding and the respective intermediate embedding for the input token for the layer block to generate the respective layer output embedding comprises: generating a projected embedding from respective intermediate embedding for the input token, comprising applying a projection weight matrix to the respective intermediate embedding for the input token; and combining the projected embedding with the respective intermediate output embedding to generate the respective layer output embedding.

18. The method of claim 17, wherein combining the projected embedding with the respective intermediate output embedding to generate the respective layer output embedding comprises: summing the projected embedding with the respective intermediate output embedding to generate the respective layer output embedding.

19. The method of any preceding claim, wherein the respective intermediate embeddings have a different dimensionality than the respective layer input and respective layer output embeddings.

20. The method of any preceding claim, wherein the plurality of embedding layers and the layer blocks in the sequence of layer blocks have been trained jointly.

21. The method of any preceding claim, wherein the set of one or more layer blocks comprises all of the layer blocks in the sequence.

22. The method of any preceding claim, wherein the set of one or more layer blocks comprises a plurality of layer blocks, and wherein the respective embedding layers for each of the layer blocks are implemented as a single mapping that maps the input token to a tensor that includes the respective intermediate embedding for the input token for each of the plurality of layer blocks in the set.

23. The method of claim 22, wherein for each of the one or more layer blocks in the set, processing the input token using the respective embedding layer for the layer block to generate a respective intermediate embedding for the input token comprises: reading, using the single mapping, the tensor that includes the respective intermediate embedding for the input token for each of the plurality of layer blocks in the set.

24. The method of claim 23. wherein the single mapping from which the tensor is read is stored in non-volatile memory.

25. The method of claim 24, wherein the non-volatile memory' is Flash memory'26. The method of any preceding claim, wherein the neural network further comprises a respective projection neural network for each layer block in the set of one or more of the layer blocks, and the method further comprising: receiving an additional input comprising an input continuous token; processing the additional input using the neural network to generate an additional network output, wherein the processing comprises: for each of the one or more layer blocks in the set. processing the input token using the respective projection neural network for the layer block to generate a respective intermediate embedding for the input continuous token; for each of the sequence of layer blocks, processing a respective layer input embedding of the input continuous token to generate a respective layer output embedding of the input continuous token, and wherein, for each layer block in the set of one or more layer blocks, processing the respective layer input embedding of the input continuous token to generate the respective layer output embedding of the input continuous token comprises: processing the respective layer input embedding of the input continuous token and the respective intermediate embedding for the input continuous token for the layer block to generate the respective layer output embedding; and processing at least the respective layer output embedding of the input continuous token generated using the last layer block in the sequence to generate the additional network output.

27. A system comprising one or more computers and one or more storage devices storing instructions that when executed by the one or more computers cause the one or more computers to perform the operations of the respective method of any one of claims 1-26.

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

Images