Generating neural network output by cross attention of query embedding over set of latent embeddings

Abstract

To provide a method of processing data using a machine learning model, a system, and a non-transitory computer storage medium.SOLUTION: A method comprises the steps of: obtaining a network input to a neural network, and a set of query embeddings; and processing the network input using the neural network to generate a network output that comprises a respective dimension corresponding to each query embedding in the set of query embeddings. The step of processing the network input includes: processing the network input using an encoder block of the neural network to generate a representation of the network input as a set of latent embeddings; and processing the set of latent embeddings, and the set of query embeddings, using a cross-attention block that generates each dimension of the network output by cross-attention of a corresponding query embedding over the set of latent embeddings.SELECTED DRAWING: Figure 1

Description

Technical Field

[0001] Cross - reference to Related Applications This application claims the benefit of the filing date of U.S. Provisional Patent Application No. 63 / 194,874, filed on May 28, 2021, entitled "GENERATING NEURAL NETWORK OUTPUTS BY CROSS ATTENTION OF QUERY EMBEDDINGS OVER A SET OF LATENT EMBEDDINGS", the content of which is hereby incorporated by reference in its entirety.

[0002] This specification relates to the processing of data using a machine learning model.

Background Art

[0003] A machine learning model receives an input and generates an output, such as a predictive output, based on the received input. Some machine learning models are parametric models and generate an output based on the received input and the values of the model's parameters.

[0004] Some machine learning models are deep models that use multiple model layers to generate an output for the received input. For example, a deep neural network is a deep machine learning model that includes an output layer and one or more hidden layers, each of which applies a non - linear transformation to the input received to generate the output.

Summary of the Invention

Means for Solving the Problems

[0005] This specification generally describes a system implemented as a computer program on one or more computers at one or more positions in a neural network to process network inputs and generate network outputs.

[0006] In one aspect, a method is described that is performed by one or more data processing devices. The method comprises the steps of obtaining (i) a network input to a neural network and (ii) a set of query embeddings that collectively define a prediction task performed by the neural network by processing the network input. The method processes the network input using a neural network to generate a network output having respective dimensions corresponding to each query embedding in the set of query embeddings. The method processes the network input using an encoder block of the neural network to generate a representation of the network input as a set of potential embeddings. The method processes (i) the set of potential embeddings and (ii) the set of query embeddings using a cross-attention block that generates each dimension of the network output by cross-attention of the corresponding query embeddings to the set of potential embeddings.

[0007] Throughout this specification, an embedding refers to an ordered set of numerical values, such as a numerical vector or matrix.

[0008] Furthermore, throughout this specification, a neural network "block" refers to a group of one or more neural network layers within a neural network.

[0009] According to one aspect, a method is provided that is executed by one or more data processing devices. The method includes: obtaining (i) a network input to a neural network and (ii) a set of query embeddings that collectively define a prediction task to be performed by the neural network by processing the network input; processing the network input using the neural network to generate a network output having respective dimensions corresponding to each query embedding in the set of query embeddings, including processing the network input using an encoder block of the neural network to generate a representation of the network input as a set of potential embeddings; and processing (i) the set of potential embeddings and (ii) the set of query embeddings using a cross-attention block that generates each dimension of the network output by cross-attention of the corresponding query embedding to the set of potential embeddings.

[0010] In some implementations, for each one or more of the query embeddings, the query embedding defines a respective spatial position within the network input, and the dimension of the network output corresponding to the query embedding defines a respective prediction related to the spatial position within the network input.

[0011] In some implementations, for each one or more of the query embeddings, the query embedding comprises one or more input features from a respective spatial position within the network input, and the dimension of the network output corresponding to the query embedding defines a respective prediction related to the spatial position within the network input.

[0012] In some implementations, for each one or more of the query embeddings, the query embedding defines the modality of the network input.

[0013] In some implementations, the set of query embeddings collectively defines a plurality of prediction tasks, and for each one or more of the query embeddings, the query embedding specifies a respective prediction task from the plurality of prediction tasks, and the dimension of the network output corresponding to the query embedding defines the prediction output of the prediction task specified by the query embedding.

[0014] In some implementations, the plurality of prediction tasks comprises one or more of a classification task, a regression task, a segmentation task, or an autoencoding task.

[0015] In some implementations, the step of generating the dimension of the network output by cross-attention of the corresponding query embeddings with respect to the set of latent embeddings comprises: (i) generating respective attention weights for each latent embedding based on the query embedding and (ii) the latent embedding, and generating the dimension of the network output corresponding to the query embedding based on the attention weights for the latent embedding.

[0016] In some implementations, the step of generating the dimension of the network output corresponding to the query embedding based on the attention weights for the latent embedding comprises: processing each latent embedding to generate a value embedding of the latent embedding, combining the value embeddings using the attention weights, and generating the dimension of the network output corresponding to the query embedding based at least in part on the result of combining the value embeddings using the attention weights.

[0017] In some implementations, the cross-attention is query-key-value attention.

[0018] In some implementations, the number of query embeddings in the set of query embeddings is greater than the number of latent embeddings in the set of latent embeddings.

[0019] In some implementations, the number of query embeddings in the set of query embeddings is at least twice as large as the number of potential embeddings in the set of potential embeddings.

[0020] In some implementations, the number of potential embeddings in the set of potential embeddings is predefined, the number of query embeddings in the set of query embeddings is variable, and is independent of the number of potential embeddings in the set of potential embeddings.

[0021] In some implementations, the network input comprises multimodal data.

[0022] In some implementations, the method further comprises the step of initializing a set of potential embeddings, the network input comprises a set of data element embeddings, and the encoder block of the neural network comprises one or more cross-attention blocks each performing an operation of updating each potential embedding in the set of potential embeddings using attention for some or all of the data element embeddings in the set of data element embeddings, and the encoder block of the neural network comprises one or more self-attention blocks each performing an operation of updating each potential embedding in the set of potential embeddings using attention for the set of potential embeddings.

[0023] In some implementations, the number of potential embeddings in the set of potential embeddings is less than the number of data element embeddings in the set of data element embeddings.

[0024] In some implementations, each data element embedding corresponds to a respective spatial position in the network input and comprises a feature embedding based on the features of the network input at that spatial position.

[0025] In some implementations, each data element embedding corresponds to a respective spatial position in the network input and comprises a position embedding characterizing the spatial position.

[0026] In some implementations, each data element embedding comprises a modality embedding that defines the modality corresponding to the data element embedding.

[0027] According to another aspect, a system is provided that includes one or more computers and one or more storage devices communicatively coupled to the one or more computers, the one or more storage devices storing instructions that, when executed by the one or more computers, cause the one or more computers to perform operations of methods described herein.

[0028] According to another aspect, one or more non-transitory computer storage media are provided that store instructions that, when executed by one or more computers, cause the one or more computers to perform operations of the methods described herein.

[0029] The subject matter described herein can be implemented in particular embodiments to realize one or more of the following advantages.

[0030] The neural network systems described herein may process network inputs to generate representations of the network inputs as a set of potential embeddings, and then generate each dimension of the network output by cross-attention of the corresponding query embedding to the set of potential embeddings. Generating each dimension of the network output by cross-attention of the query embedding to the set of potential embeddings allows the system to flexibly modify the dimensionality of the network output, for example, by modifying the number of query embeddings, without having to reconstruct hidden layers of the neural network.

[0031] Furthermore, the prediction tasks executed by the neural network are defined by query embeddings and can be modified using query embeddings. Thus, the neural network system provides a flexible and general neural network architecture that can be implemented to execute various prediction tasks for network inputs of various modalities with minimal changes. Thereby, the neural network system can use resources more efficiently by reducing or eliminating the need to design, implement, and train a new neural network for each new domain or prediction task.

[0032] In contrast to conventional neural networks that generate fixed-size outputs, such as fixed-size image segmentation outputs, the neural network system described herein can use query embeddings to generate sparse or partial network outputs. For example, instead of generating a complete segmentation of an image, the neural network system can generate a segmentation of only a specified appropriate subset of the image, for example, by executing queries on a set of potential embeddings using only the relevant query embeddings. Thus, the neural network system can reduce the consumption of computational resources (e.g., memory and computing power) by providing the option to generate sparse or partial network outputs when appropriate. For example, the neural network system can be used to segment only the region of interest of a larger image without segmenting the remaining parts of the image.

[0033] More generally, since the dimension of the network output corresponds to the number of query embeddings within the set of query embeddings, the size of the network output can be defined regardless of the size of the network input. Further, using cross-attention avoids the quadratic dependence of the calculations performed on the sizes of the network input and output. Instead, the dependencies are linear, which facilitates, for example, the processing of video, audio, and multimodal data. Similarly, for example, since the network output is generated by performing queries on potential embeddings using cross-attention, the implementation of the described techniques can use a set of potential embeddings that can be independent of the spatial or temporal structure of the network input and output. This facilitates the processing of multimodal data, i.e., network inputs having different structures or numbers of dimensions. Further, in some implementations of the system, a simpler architecture can be used to perform tasks, such as an optical flow determination task, more efficiently than some other techniques.

[0034] Details of one or more embodiments of the subject matter described in this specification are set forth in the accompanying drawings and the description below. Other features, aspects, and advantages of the subject matter will become apparent from the description, the drawings, and the claims.

Brief Description of the Drawings

[0035]

Figure 1

Figure 2

Figure 3

Figure 4

Figure 5

Figure 6

Figure 7A

Figure 7B

[0036] Like reference numerals and designations in the various drawings indicate like elements.

[0037] FIG. 1 is a block diagram of an exemplary neural network system 100. The neural network system 100 is an example of a system implemented as a computer program on one or more computers at one or more locations where the systems, components, and techniques described below are implemented.

[0038] The neural network system 100 can be configured to obtain a neural network input 102 and process the input 102 using a neural network 150 to generate a neural network output 108.

[0039] The network input 102 can include any suitable type of data. That is, the neural network 150 can be configured to process inputs of multiple modalities, such as images, videos, audio, text, or other suitable modalities, or any combination thereof. The neural network 150 can have a single set of parameters and is substantially flexible so as to be able to process any one or combination of multiple modalities.

[0040] In some implementations, the network input 102 can characterize an entity. The entity can include an image, an audio waveform, a point cloud (e.g., generated by a lidar or radar sensor), a protein, a sequence of words (e.g., forming one or more sentences or paragraphs), a video (e.g., representing a sequence of video frames), or other suitable types of data, or combinations thereof. The images used herein, i.e., still images or videos, can include point clouds.

[0041] As a specific example, the entity can include a plurality of units arranged in a spatial or temporal structure. For example, as shown in FIG. 4, the entity can be an image, and each unit can be a pixel within the image. Each unit within the entity can have an associated data element embedding that can characterize, for example, the position of the unit in the spatial or temporal structure and / or the features associated with the unit in the spatial structure. In another specific example, the entity can be text. For example, the entity can be a sequence of words, phrases, characters, or word fragments in one or more languages. In yet another specific example, the entity can include a combination of different modalities. For example, the entity can include a combination of audio, video, and label data. As another example, the entity can comprise sensor data from one or more sensors configured to recognize one or more characteristics of a real-world environment, such as image data, audio data, or so-called "fine touch" sensor data (which enables localization). Although some examples of the network input 102 were described above, generally, the network input 102 can have any suitable number of dimensions and structure and can include any suitable type of data.

[0042] The neural network system 100 can process the network input 102 using the neural network 150 to generate a network output 108. The network output 108 can be, for example, a classification output, a regression output, a sequence output (i.e., including a sequence of output elements), a segmentation output, an autoencoding output, or other suitable network output, or a combination thereof.

[0043] The neural network system 100 can flexibly modify the dimensionality of the network output 108, for example, to generate a sparse or partial network output 108. For example, if the network input 102 is an image, the network output 108 can be a segmentation of only a specified (appropriate) subset of the image. In general, the network output 108 can be represented as a set of output vectors, and the number of output vectors can be referred to as the "dimensionality" of the network output 108. As will be described in more detail below, the system 100 can control the dimensionality of the network output 108 using a set of query embeddings 106.

[0044] In addition to obtaining the network input 102, the system 100 can further obtain a set of query embeddings 106. The network input 102 can be provided, for example, by a user of the system 100 through an application programming interface (API) made available by the system 100.

[0045] The set of query embeddings 106 may be provided in the same way by a user of system 100, for example, or may be determined in advance, i.e., the system may be configured to perform a particular task. The query embeddings 106 may be used by system 100 to generate a network output 108 having a desired number of dimensions. If required for a particular task, each query embedding within the set of query embeddings 106 can define relevant information for a corresponding output dimension within the network output 108. In other words, each query embedding can define, for example, the position of the corresponding output dimension and / or modality-specific features, etc., such that the network output comprises an output for a particular 1D, 2D, or 3D position, or a particular data modality. Such query embeddings can be constructed by combining, for example concatenating or adding, a set of vectors to the query embedding to provide the information required for the corresponding dimension of the network output.

[0046] The query embeddings 106 can be defined in any suitable way, for example, manually engineered, learned, or defined as a function of the network input 102. In general, the neural network system 100 does not depend on the use of any particular type of query embedding and can produce good results with simple embeddings. During training, the system is provided with a set of network inputs and query embeddings to generate the network output. The system is trained to generate a network output that performs a prediction task. This will be further explained later.

[0047] Collectively, the set of query embeddings 106 can define a prediction task to be performed by neural network 150 by processing network input 102. The prediction task can be, for example, a classification task, a regression task, a segmentation task, an autoencoding task, or any other suitable task, or a combination thereof. In some implementations, each query embedding within the set of query embeddings 106 can be generated by combining (e.g., concatenating or summing) appropriate other features related to the spatial, modality, task-specific, and / or corresponding output dimensions of network output 108.

[0048] As a specific example, network input 102 can include representations of a first image and a second image, each representing the same scene, and the prediction task can be to estimate a two-dimensional displacement of points (e.g., optical flow), e.g., pixels or a suitable subset of pixels (i.e., not all pixels) within the first image. In some implementations, such points can comprise an image patch comprising pixels or multiple pixels. In such cases, each query embedding within the set of query embeddings 106 can, for example, specify the x and y coordinates of the corresponding points within the first image, and the optical flow prediction task can include one or more additional embeddings for the representations of the first and second images of the points. That is, the set of query embeddings 106 can define a set of points within the first image to generate a prediction, and the prediction is to estimate the optical flow. For example, the magnitude and direction of the optical flow at a point can be defined by the dimensions of the network output. Examples of query embeddings and prediction tasks are described in more detail below with reference to FIGS. 4 and 5.

[0049] After obtaining the set of network input 102 and query embeddings 106, system 100 can use neural network 150 to process the input 102 and generate a network output 108 having a desired number of dimensions, for example, an output including respective dimensions corresponding to each query embedding within the set of query embeddings 106. This process will be described in more detail below.

[0050] System 100 can use the encoder block 110 of neural network 150 to process network input 102 to generate a representation of network input 102 as a set of potential embeddings 104. "Potential embedding" can refer to an embedding within a potential space. Generally, the potential embeddings 104 can be included in a space of any suitable number of dimensions. Generally, the number of dimensions of the set of embeddings can be defined by the number of embeddings within the set of embeddings (e.g., N). In some implementations, the encoder block processes the network input and an initial set of potential embeddings to generate a set of potential embeddings.

[0051] In some implementations, the number of dimensions of network input 102 can be much larger than the number of dimensions of potential embeddings 104 and / or query embeddings 106. For example, if network input 102 includes an image having dimensions of 224×224 pixels and the number of data element embeddings (e.g., corresponding to each pixel within the image) is M = 50176, the number of potential embeddings 104 can be, for example, N = 512, and thus N << M. The encoder block 110 can process network input 102 to generate a representation of the input as a set of potential embeddings 104 in any suitable manner. A specific example of this process will be described in more detail below with reference to FIG. 2.

[0052] After generating a set of potential embeddings 104 using the encoder block 110, the neural network system 100 can process the set of potential embeddings 104 and the set of query embeddings 106 using the output block 120 of the neural network 150 to generate a network output 108. In particular, the output block 120 can be configured to perform an "attention" operation. The attention operation can include updating each embedding within the first set of embeddings using attention to the second set of embeddings. In some implementations, after performing the attention operation (e.g., after updating each embedding), the output block 120 can process each updated embedding using one or more neural network layers, such as a fully connected neural network layer. There are many different attention mechanisms that can be used, and exemplary attention operations are described in more detail below with reference to FIG. 2.

[0053] Specifically, the output block 120 can update each query embedding within the set of query embeddings 106 with respect to the set of potential embeddings 104. Continuing with the above example of optical flow, each query embedding within the set of query embeddings 106 can specify, for example, the x and y coordinates of the corresponding pixel in the first image, as well as an embedding that specifies the optical flow prediction task. The output block 120 can generate a particular dimension of the network output 108 through cross-attention of the corresponding query embeddings with respect to the set of potential embeddings 106. Returning to the previous example by way of illustration, the dimension of the network output 108 can be, for example, a single flow vector representing the optical flow of the pixel in the first image defined by the query embedding. The output block 120 can repeat this process for each query embedding within the set of query embeddings 106, thereby generating each respective dimension of the network output 108. In some implementations, the output corresponding to each query embedding within the set of query embeddings 104 can be generated in parallel. This process will be described in more detail below with reference to FIG. 2.

[0054] The encoder block 110 and the output block 120 can have any suitable neural network architecture that enables them to perform their predetermined functions. For example, the encoder block 110 and the output block 120 can have any suitable number (e.g., 2 layers, 5 layers, 10 layers) of any suitable neural network layers (e.g., convolutional layers, fully connected layers, recurrent layers, attention layers, etc.) and can be connected in any suitable configuration (e.g., as a linear sequence of layers). The neural network system 100 can also further include any number of neural network blocks configured to perform any suitable operations. Specific examples of the encoder block 110 and the output block 120 will be described in more detail below with reference to FIG. 2.

[0055] In this way, the system 100 can flexibly modify the dimensionality of the network output 108, for example, by modifying the number of query embeddings 106 without reconstructing the hidden layer of the neural network 150. Further, as described above, the prediction task executed by the neural network 150 is defined by the query embedding 106, and in some implementations, it can be modified using the query embedding 106. Therefore, the neural network system 100 can provide a flexible and general neural network architecture that can be implemented to execute various prediction tasks for network inputs of various modalities with minimal changes.

[0056] The neural network system 100 can further include a training engine that can train the neural network 150 with respect to a set of training data over a plurality of training iterations. The training data can include a set of training examples, and each training example specifies (i) a training input and (ii) a target output to be generated by the neural network 160 by processing the training input.

[0057] In each training iteration, the training engine can sample a batch of training examples from the training data and use a sequence of neural network blocks included in neural network 150 (e.g., encoder block 110 and output block 120) to process the training input specified by the training example in order to generate the corresponding network output. In particular, for each training input, neural network 150 processes the training input using the current model parameter values of encoder block 110 to generate a set of potential embeddings 104. Neural network 150 processes the set of potential embeddings 104 and the set of query embeddings 106 using the current model parameter values of output block 120 to generate network output 108 corresponding to the training input.

[0058] The training engine can adjust the model parameter values of encoder block 110 and output block 120 to optimize an objective function that measures the similarity between (i) the network output generated by neural network 150 and (ii) the target network output specified by the training example. The objective function can be, for example, a cross-entropy objective function, a mean squared error objective function, or any other suitable objective function. In some implementations, during training, the network output generated by neural network 150 can include only a subset of the total set of possible outputs. In such cases, the objective function can be evaluated with reference to only the generated outputs, rather than the total set of outputs. This accelerates the training of the neural network and improves training efficiency.

[0059] The training engine can determine the gradient of the objective function, for example, using backpropagation techniques. The training engine can use the gradient to update the model parameter values of the encoder block 110 and the output block 120, for example, using any suitable gradient descent optimization algorithm, such as Adam. The training engine can determine the performance metric of the neural network 150 based on a set of validation data that is not used during the training of the neural network 150. After training, the neural network system 100 can be used to process network inputs and generate network outputs, for example, to perform a machine learning task.

[0060] In some implementations, the training engine can train the neural network 150 using reinforcement learning techniques. For example, the training engine can train the neural network 150 by repeatedly backpropagating the gradient of the reinforcement learning objective function through the neural network 150 (e.g., an encoder block, a cross-attention block, or both) to iteratively adjust the model parameter values of the neural network 150. The reinforcement learning function can be any suitable reinforcement learning objective function.

[0061] The neural network system 100 can be configured to perform any suitable machine learning task. Some examples are shown below.

[0062] In some implementations, system 100 uses neural network 150 to perform image or audio segmentation tasks. For example, the neural network can process a network input that includes an image representing an audio waveform or a series of audio samples to generate a network output that defines a respective score distribution for a set of possible classes for each pixel in the input image or for each audio sample. The class scores can define the likelihood that the corresponding pixel or sample is included in the class. For example, in the case of an image, the possible classes can include, for example, water, buildings, vehicles, pedestrians, etc. In this example, a query embedding can specify the position of the corresponding pixel in the input image or the position of the corresponding sample in the audio, and the dimension of the network output corresponding to the query embedding can define the score distribution of the corresponding pixel in the input image or the corresponding sample in the audio. The image can be a video, i.e., a video. In a similar manner, system 100 can use neural network 150 to perform an image depth prediction task in which the network output defines a predicted depth value for the corresponding pixel to obtain a (spatial 3D) depth map of the image.

[0063] In some implementations, system 100 uses neural network 150 to perform protein modeling tasks. For example, the network input characterizes a protein (e.g., a multiple sequence alignment of proteins), and the network output characterizes the predicted structure of the protein (e.g., the respective three-dimensional (3-D) spatial positions and orientations of each amino acid in the protein structure). In this example, a query embedding can specify the index of an amino acid in the amino acid sequence of the protein, and optionally, the type of the amino acid (e.g., alanine, arginine, asparagine, etc.) at the specified index. The dimension of the network output corresponding to the query embedding can define the predicted spatial position and orientation of the amino acid specified by the query embedding.

[0064] In some implementations, system 100 uses neural network 150 to execute agent control tasks, where the network input represents one or more observations or other data sequences characterizing the state of the environment, and the output defines an action to be executed by the agent in response to the most recent data in the sequence. The environment can be a real environment or a simulated environment, and the agent can be, for example, a robot, an autonomous land, sea, or air vehicle, or a control system for an industrial facility. Each query embedding can specify a respective task to be executed by the agent in the environment. For example, the task is for the agent to move to a target position in the environment, and the query embedding can specify the target position.

[0065] In some implementations, system 100 uses neural network 150 to process network inputs representing audio samples in an audio waveform to perform speech recognition, for example, to generate a network output characterizing a sequence of phonemes, graphemes, characters, or words corresponding to the audio waveform. The audio samples define a sequence of samples, and optionally, the network input can also include an encoding of the position of the samples in the sequence. Each dimension of the network output can correspond to a respective time interval in the audio waveform, for example, and can define a respective score distribution over a set of possible phonemes, graphemes, characters, or words. In this example, the neural network can generate the network output based on, for example, (i) each time interval in the audio waveform, such as the position in the sequence, and optionally (ii) a query embedding defining a prediction task for the time intervals in the audio waveform, such as decoding a phoneme, grapheme, character, or word corresponding to the time interval.

[0066] In some implementations, system 100 uses neural network 150 to perform a reconstruction task, such as an autoencoding task, by processing network input to generate a network output that defines a predicted reconstruction of the network input. The network input can be, for example, an image or video, an audio waveform, a point cloud, or a sequence of text. In this example, each query embedding can include data that specifies a corresponding position within the network input to be reconstructed, such as a spatial position within an image, a spatial and / or temporal position within a video, or a time window within an audio waveform. Optionally, each query embedding can also define the modality of the network input to be reconstructed, for example, define whether the network input is an image, an audio waveform, or a point cloud. Each dimension of the network output can correspond to a reconstructed point of the 1D, 2D, or 3D network input specified by the query embedding. More or fewer points than the original network input may be reconstructed. The points can be reconstructed sequentially and / or in parallel. For example, the set of query embeddings may comprise a query embedding for each point within the network input, such as a pixel of an image or a sample of an audio signal, or may comprise a query embedding for just one point, and the network output may be reconstructed sequentially, one point at a time, by specifying each desired point, or the set of query embeddings may reconstruct only a subset of the points within the network input.

[0067] In some implementations, system 100 uses neural network 150 for neural machine translation tasks to generate a network output that translates a sequence of text into another language, i.e., a sequence of text in another language that is a translation of the input sequence of text. For example, to process a network input that represents a sequence of text, such as a sequence of words, phrases, characters, or word pieces, in one language. In this example, each query embedding can include data that identifies, for example, the natural language (e.g., English, French, German, etc.) into which the input text sequence is to be translated. Additionally, each query embedding can specify the corresponding position within the output sequence of text.

[0068] In some implementations, system 100 uses neural network 150 to perform audio or audiovisual processing tasks. For example, if the network input represents an oral utterance, the network output generated by the neural network can be scores for each of a set of texts, where each score represents the estimated likelihood that the text is the correct transcription of the utterance. In this example, each query embedding can correspond to each text representation. As another example, if the network input represents an oral utterance, the output generated by the neural network can indicate whether a particular word or phrase ("hotword") was spoken within the utterance. In this example, each query embedding can be, for example, a one-hot embedding corresponding to each word or phrase, or can correspond to it.

[0069] In some implementations, system 100 uses neural network 150 to perform a text reading task, the network input represents text in natural language or features of text in natural language, and the network output is a spectrogram, a waveform, or other data that defines the audio of the text spoken in natural language. In this example, each query embedding can be for identifying each voice used when speaking the text (e.g., if each voice corresponds to a person's voice). Further, each query embedding can specify a corresponding time window within the network output.

[0070] In some implementations, system 100 uses neural network 150 to perform a health prediction task, the network input represents data obtained from a patient's electronic health record data, and the output is a prediction related to the patient's future health, such as a predicted treatment to be prescribed to the patient, the likelihood of a health - adverse event occurring to the patient, or the patient's predicted diagnosis, etc. In this case, each query embedding can specify the corresponding point or period in the patient's life. In some implementations, system 100 uses neural network 150 to perform a text generation task, the network input represents a sequence of text, and the output is another sequence of text, such as the completion of the input sequence of text, a response to a question raised in the input sequence, or a sequence of text related to the topic specified by the first text sequence. As another example, the network input can represent data other than text, such as an image, and the output sequence can be text that describes the data represented by the network input. In this example, each query embedding can specify the corresponding position within the text output sequence.

[0071] In some implementations, system 100 uses neural network 150 to perform an image generation task, where the network input represents an adjustment input and the output is a sequence of pixel intensity values of an image.

[0072] In some implementations, system 100 uses neural network 150 to perform a genomics task, where the network input represents a DNA sequence or a fragment of another molecular sequence, and the network output is, for example, either an embedding of the fragment for use in a downstream task by utilizing an unsupervised learning technique for a dataset of DNA sequence fragments, or an output for a downstream task. Examples of downstream tasks include prediction of promoter sites, methylation analysis, prediction of the functional effects of non-coding variants, and the like. In this example, each query embedding can specify a corresponding position within the DNA sequence.

[0073] In some implementations, system 100 uses neural network 150 to perform image, video, or audio classification tasks. The network input can then represent, as described above, the pixels of an image or video, or the samples of an audio waveform. Optionally, the network input can include positional encodings of the pixels of an image or video, or the samples of an audio, and in the case of video, the positional encoding can be a spatial position and / or a temporal position. The network output defines the classification of the image, video, or audio. The classification can include respective scores for each object category within a set of possible object categories (e.g., image, vehicle, pedestrian, bicycle, etc.). The score of an object category can define the likelihood that the network input comprises an object belonging to that object category, e.g., that an image or video depicts an object. The classification of a video can include, for example, gesture recognition and the classification of actions depicted within the video. The set of query embeddings can comprise a single (learnable) query embedding. In a similar manner, system 100 can use neural network 150 to perform image or audio bounding box tasks, and the network output comprises a classification and a vector defining the coordinates of the bounding box of the object in 1, 2, or 3 dimensions (counting time as a dimension here).

[0074] As described above, generally in the above examples, the reference to an image or video includes a point cloud. For example, in some implementations, system 100 uses neural network 150 to perform point cloud processing tasks, e.g., the network input represents a point cloud (e.g., generated by a lidar or radar sensor), and the network output characterizes, e.g., the type of object represented by the point cloud. In this example, each query embedding can specify the corresponding spatial position of one or more points within the point cloud.

[0075] In some implementations, system 100 uses neural network 150 to perform image or video captioning tasks. For example, the neural network can process a network input that includes an image or a video comprising a series of images to generate a network output that defines a set of tokens that describe the image or video for one or more of the images or series of images, where the tokens can represent words, parts of words, or sentences. The tokens may be deterministically defined, for example, by sampling from one or more defined distributions, or may be probabilistically defined. In the case of video, the temporal position of each of the series of images can be specified (but not necessarily) by query embedding. As used herein, temporal position includes temporal order.

[0076] In some implementations, the network input is multimodal input, and system 100 uses neural network 150 to perform multimodal tasks. Generally, such multimodal input is a combination of data of two or more different types or modalities, and different types of data may represent the same or overlapping objects within the network input using different modalities, but this is not necessarily the case. Such multimodal data may comprise audiovisual data comprising a combination of pixels of an image or video and audio data representing samples of an audio waveform. As another example, such multimodal data may comprise a combination of text data representing text in natural language, such as tokens, and pixels of an image or video, or samples of an audio waveform. A multimodal task is any of the tasks described above, but may include additional mode data in the network input. As some examples, multimodal tasks may comprise classification tasks, segmentation tasks, speech recognition tasks, reconstruction tasks, recognition tasks, captioning tasks, and agent control tasks. To construct the network input, modality-specific embeddings may be provided for each modality type, and then the multimodal data may be serialized into a 2D input array 202. Generally, the query embedding may be as described above, but in some implementations, modality-specific embeddings are added to generate an output for each modality. The embeddings may be learned embeddings. If a spatial or temporal position is incorporated into the query embedding, such as in the case of video or audio data, this may be as described above, as it allows the system to select each modality by means of the modality-specific embeddings. The network output may correspond to those described above, but what is applied is provided to the multimodal input. For example, a multimodal classification task may generate a network output that classifies the multimodal input into one or more of a plurality of categories, for example by defining category scores for each of the possible categories for the input.Similarly, the caption task generates a network output of tokens that describe multimodal input, and the corresponding query may optionally include, for example, a spatial or temporal position for each query embedding. As another example, when the multimodal data comprises a combination of text data and image, video, or audio data, the task may comprise, for example, processing that combination to provide a network output that defines whether an image, video, or audio waveform is described by a particular caption, for example by defining a score for the text or caption. The query may be as described above with respect to classification and may include position data for generating respective scores for each position represented by one of the query embeddings within the set of query embeddings.

[0077] Generally, in the above examples, the network input may comprise feature embeddings that define raw data, such as raw pixel data or audio samples, or spatial, temporal, or spatio-temporal features obtained by preprocessing, for example by a feature encoder neural network.

[0078] An exemplary architecture of the neural network system 100 will be described in more detail below with reference to FIG. 2.

[0079] FIG. 2 is a more detailed block diagram of an exemplary architecture of the neural network system 100. The neural network system 100 is an example of a system implemented as a computer program on one or more computers at one or more locations where the systems, components, and techniques described below are implemented.

[0080] The neural network system 100 can be configured to obtain a network input 202 that can include a set of data element embeddings (e.g., an "input array") that characterize an entity. The entity can include, for example, a plurality of units arranged in a spatial and / or temporal structure.

[0081] In some implementations, each data element embedding can correspond to a respective spatial position within the network input and can include a feature embedding based on the features of the network input at that spatial position. For example, if the entity is an image and the units are pixels within the image, the system can select a patch of the image around that pixel and concatenate the pixels within the patch into a vector to obtain the feature embedding. As used herein, references to spatial positions should generally be understood to include positions within a sequence such as a sequence of words or a sequence of audio waveform samples.

[0082] In some implementations, each data element embedding can include a position embedding that characterizes the spatial position corresponding to the data element embedding. For example, if the entity is a sequence of words and each unit within the entity is a word, the system can generate a position embedding based on the index of the word within the sequence of words. In some implementations, each data element embedding can include a modality embedding that defines the modality corresponding to the data element embedding. For example, if the network input is an audio waveform, the embedding of the data element can include a modality embedding that specifies that the network input is an audio waveform.

[0083] Generally, the network input can have dimensions of M×C, where M is the number of data element embeddings and C is the number of channels of each data element embedding.

[0084] The neural network system 100 can further obtain a set of query embeddings 206 (e.g., an "output query array") that collectively define prediction tasks executed by the neural network 250 by processing the network input 202. The query embeddings 206 can have any suitable number of dimensions. For example, the query embeddings can include O query vectors.

[0085] In some implementations, the neural network system 100 can further obtain an initial set of latent embeddings 204 (e.g., a "latent array"). The latent embeddings 204 can be pre-defined and / or randomly initialized, or they can be learned, i.e., they can comprise a set of learned parameters each defining an element of an array such as learned weights. That is, in some implementations, the encoder block processes the network input and the initial set of latent embeddings to generate a representation of the network input as a set of latent embeddings. The set of latent embeddings 204 can have dimensions N×D, where N is the number of latent embeddings and D is the number of channels of each latent embedding, both of which can be hyperparameters of the neural network system 100. In some implementations, the number O of query embeddings can be, for example, at least twice as large as the number N of latent embeddings. Further, the number O of query embeddings can be variable and independent of the number N of latent embeddings.

[0086] The neural network system 100 can process the network input 202 using the neural network 250 to generate a network output 208 (e.g., an "output array") that includes respective dimensions corresponding to each query embedding within the set of query embeddings 206. For example, the network output 208 can be represented by an O×E array, where O is the number of output vectors and E is the number of output elements (hyperparameters) within each output vector. The neural network system 100 can generate respective dimensions (e.g., respective output vectors) corresponding to each query embedding within the set of query embeddings 206. The number of output vectors can be the same as the number of query vectors.

[0087] The neural network system 100 can generate the network output 208 by using the encoder block 210 and the output block 250. The encoder block 210 can include a sequence of one or more neural network blocks, e.g., (i) one or more cross-attention blocks 230, and (ii) one or more self-attention blocks 240. The encoder block 210 can be configured to process the network input 202 to generate a representation of the network input 202 as a set of potential embeddings 204. The output block 250 can be configured to process the set of potential embeddings 204 and the set of query embeddings 206 generated by the encoder block 210 to generate the network output 208. The attention blocks included in the neural network 250 will be described in more detail below.

[0088] The cross-attention block 230, the self-attention block 240, and the output block 250 can each be configured to perform an attention operation, e.g., using attention for a second set of embeddings, to update each embedding within the first set of embeddings.

[0089] For example, for each target embedding in the first set of embeddings, each attention block can generate respective attention weights for each embedding in the second set of embeddings and generate a combined embedding based on the second set of embeddings and the corresponding attention weights. As a specific example, each attention block can generate a combined embedding as a weighted sum of the second set of embeddings, for example, by multiplying each embedding in the second set of embeddings by a corresponding weight and summing the weighted embeddings. Then, each attention block can use the combined embedding to update the target embedding in the first set of embeddings, for example, by replacing the target embedding with the combined embedding, adding the combined embedding to the target embedding, or by other suitable methods.

[0090] In some implementations, the attention block can perform a query-key-value (QKV) attention operation and, for example, use attention on the second set of embeddings, using query (Q), key (K), and value (V) embeddings, to update each embedding in the first set of embeddings. In particular, each attention block can include (i) a query subnetwork, (ii) a key subnetwork, and (iii) a value subnetwork. For each target embedding in the first set of embeddings, the query subnetwork can be configured to process the target embedding in the first set of embeddings to generate a respective query embedding (Q) for the target embedding. The key subnetwork can be configured to process each embedding in the second set of embeddings to generate a respective key embedding (K) for each embedding in the second set of embeddings. Similarly, the value subnetwork can be configured to process each embedding in the second set of embeddings to generate a respective value embedding (V) for each embedding in the second set of embeddings.

[0091] Next, each attention block can use query embeddings (Q), key embeddings (K), and value embeddings (V) to update each target embedding within the first set of embeddings using attention for the second set of embeddings. Specifically, each attention block can generate attention weights for each embedding within the second set of embeddings, for example, as the inner product (e.g., dot) of each of the query embeddings (Q) and the key embeddings (K). Based on the second set of embeddings and the attention weights, each attention block can generate a combined embedding, for example, as a linear combination of the value embeddings (V) weighted by their respective attention weights. Finally, each attention block can update the target embeddings within the first set of embeddings by using the combined embedding to replace, for example, the target embeddings within the first set of embeddings with a weighted sum of the value embeddings (V).

[0092] In some implementations, the first set of embeddings and the second set of embeddings can be the same set of embeddings. In such cases, the attention operation (e.g., the QKV attention operation) may be referred to as a "self-attention" operation. The self-attention operation can be performed, for example, by a self-attention block 240. For example, the first set of embeddings can be a set of latent embeddings 204, the second set of embeddings can also be a set of latent embeddings 204, and the self-attention block 240 can use self-attention for the set of latent embeddings 204 to update each latent embedding within the set of latent embeddings 204. In some implementations, the self-attention block 240 can repeatedly update each latent embedding within the set of latent embeddings 204 using self-attention for the set of latent embeddings 204.

[0093] In some implementations, the first set of embeddings and the second set of embeddings can be different sets of embeddings. In such cases, the attention operation (e.g., the QKV attention operation) may be referred to as a "cross-attention" operation. The cross-attention operation can be performed, for example, by the cross-attention block 230 and the output block 250 included in the encoder 210.

[0094] As a specific example, in the case of the cross-attention block 230 included in the encoder 210, the first set of embeddings can be a set of latent embeddings 204 (e.g., initialized randomly or initialized to previously learned values), and the second set of embeddings can be the data element embeddings 202 provided as input to the neural network 250. The cross-attention block 230 can update each latent embedding in the set of latent embeddings 204 using cross-attention for some or all of the data element embeddings within the set of data element embeddings 202.

[0095] As another specific example, in the case of the output block 250, the first set of embeddings can be a set of query embeddings 206, and the second set of embeddings can be a set of latent embeddings 204 generated by the self-attention block 240. The output block 250 can update each query embedding in the set of query embeddings 206 using cross-attention for the set of latent embeddings 204. In other words, the output block 250 can generate each dimension of the network output 208 by cross-attention of the corresponding query embedding in the set of query embeddings 206 with respect to the set of latent embeddings 204. After cross-attending each query embedding in the set of query embeddings 206 with respect to the set of latent embeddings 204, the output block 250 can generate a complete neural network output 208 including each respective dimension of the query embeddings in the set of query embeddings 206.

[0096] When neural network system 100 obtains a set of second different query embeddings 206 (e.g., specifying a different number of query vectors, different spatial positions, and / or different modality-specific features), neural network system 100 can process network input 202 and the set of second query embeddings 206 as described above to generate a network output 208 corresponding to each query embedding within the set of second different query embeddings 206. Thus, by modifying query embeddings 206, neural network system 100 can flexibly modify the dimensionality of network output 208 and / or the prediction task performed by neural network 250.

[0097] Next, an exemplary process of using neural network system 100 to generate neural network output 208 will be described in more detail.

[0098] FIG. 3 is a flowchart of an exemplary process 300 for generating neural network output using a neural network system. For convenience, process 300 will be described as being executed by a system of one or more computers located at one or more positions. For example, a neural network system, such as neural network system 100 of FIG. 1 or neural network system 100 of FIG. 2, can execute process 300 when appropriately programmed according to this specification.

[0099] The system obtains (302) a set of query embeddings that collectively define (i) a network input to a neural network and (ii) a prediction task to be performed by the neural network by processing the network input. In some implementations, the network input can include, for example, multimodal data such as video data, audio data, and label data. Generally, as shown in FIG. 6A, the network input can include any suitable type of data or combinations thereof.

[0100] The system processes the network input using a neural network to generate a network output having respective dimensions corresponding to each query embedding in the set of query embeddings (304).

[0101] For example, the system can process the network input using an encoder block of the neural network to generate a representation of the network input as a set of potential embeddings.

[0102] Next, the system can process (i) a set of latent embeddings and (ii) a set of query embeddings using a cross-attention block, such as the output block 250 described above, that generates each dimension of the network output by way of corresponding query embedding cross-attention on the set of latent embeddings (e.g., query-key-value attention, etc.). For example, as described above with reference to FIG. 2, the system can generate respective attention weights for each latent embedding based on (i) the query embedding and (ii) the latent embedding, and generate the dimension of the network output corresponding to the query embedding based on the attention weights of the latent embeddings. In particular, the system can process each latent embedding to generate a value embedding of the latent embedding. The system can then combine the value embeddings using the attention weights. Finally, the system can generate the dimension of the network output corresponding to the query embedding based at least in part on the result of combining the value embeddings using the attention weights.

[0103] In some implementations, the query embedding can define each spatial position within the network input. In such cases, each prediction related to the spatial position of the network input can be defined by the dimension of the network output corresponding to the query embedding. For example, as shown in FIG. 4, the network input can be an image having a size of 100×100 pixels, and each query embedding can define the respective spatial position (e.g., x coordinate and y coordinate) of the corresponding pixel within the image. The prediction task can be, for example, reconstructing the original image. The set of query embeddings can include 10,000 query vectors, such as vectors for each pixel within the image. The dimension of the network output corresponding to the query embedding can define each prediction related to the spatial position of the network input, such as the reconstruction of the pixel within the image defined by the query embedding.

[0104] In some implementations, the query embedding can include one or more input features from each spatial position within the network input. In such cases, each prediction related to the spatial position within the network input can be defined by the dimension of the network output corresponding to the query embedding. For example, as shown in FIG. 5, in the case of an optical flow prediction task, each query embedding can specify one or more input features at each spatial position within the network input, such as the spatial position of a pixel within an image (e.g., x - coordinate and y - coordinate) and one or more input features associated with that pixel. For example, in the network input, two images can be concatenated such that the corresponding pixels are indexed by the same (M) dimension of the input array (concatenated in the channel dimension C), or in another approach, the two images are not concatenated, and then the network input, and optionally also the query embedding, may be provided with temporal encoding features. In one example, the input features of a pixel within an image can be the patch of the image surrounding that pixel concatenated into a vector. The dimension of the network output corresponding to the query embedding that defines that pixel within the image can specify, for example, the predicted flow vector of that pixel within the image.

[0105] In some implementations, the query embedding can define the modality of the network input. In the case of multi - modal input, such as input including video, audio, and labels, all three modalities can be specified by each query embedding within a set of query embeddings. Generally, the query embedding can be represented in a variety of different ways. As a specific example, the query embedding can be represented as a one - hot embedding.

[0106] In some implementations, a set of query embeddings can collectively define multiple prediction tasks. For example, each query embedding can define a respective prediction task. In such a case, the prediction output of the prediction task specified by the query embedding can be defined by the dimension of the network output corresponding to the query embedding. The prediction tasks can include, for example, classification tasks, regression tasks, segmentation tasks, autoencoding tasks, or other suitable tasks. For example, as shown in FIG. 5, the embedded "task_id" can specify a particular prediction task. In the case of eight prediction tasks, the set of query embeddings can include eight query embeddings each specifying a particular prediction task.

[0107] Exemplary query embeddings and prediction tasks are described in more detail below with reference to FIGS. 4 and 5.

[0108] FIG. 4 shows an exemplary query embedding 400 that can be used by a neural network system (e.g., system 100 of FIG. 1 or system 100 of FIG. 2) to generate a neural network output 408.

[0109] As shown in FIG. 4, the network input 402 can be an image. The prediction task executed by the neural network included in the neural network system can be an autoencoding task, for example, the reconstruction of the original image 402. In this case, each query embedding in the set of query embeddings can define the respective spatial position (e.g., x coordinate and y coordinate) of the corresponding pixel in the image 402. The image 402 can have a size of 100×100 pixels, and the set of query embeddings can include 10,000 query vectors, for example, vectors for each pixel in the image 402.

[0110] As described above with reference to FIGS. 1 and 2, the neural network can include an encoder 410 and an output block 420. The encoder 410 can process the image 402 to generate a representation of the image 402 as a set of potential embeddings 404. The output block 420 can process the set of potential embeddings 404 and the set of query embeddings 404 to generate the network output 408 by cross-attention of the corresponding query embeddings with respect to the set of potential embeddings 404. The network output 408 can be, for example, a reconstruction of the image 402, where each pixel is represented by the corresponding query embedding within the set of query embeddings 400.

[0111] FIG. 5 shows another example of query embeddings 500. In general, query embeddings can be constructed using output-specific features to generate network outputs having various semantics.

[0112] For example, if the network output is a sequence of words, the query embeddings can include positional embeddings, where each positional embedding specifies the position of each word within the sequence of words. As a specific example, the positional embeddings can be based on the indices of the words within the sequence of words. In another example, if the network output is a two-dimensional array of pixels, the positional embeddings can be based on the x-y coordinates of the pixels within the array of pixels. In yet another example, if the network output is a point cloud, the positional embeddings can be based on the x-y-z coordinates of the points within the point cloud. In some implementations, the positional embeddings can be, for example, Fourier feature position encodings having frequency bands logarithmically spaced over a predefined target frequency range. In some implementations, both "raw" coordinates, such as x, y, or z, and Fourier features can be included.

[0113] In some implementations, the input features of the target output can also be used, either alone or together with the location features, to execute a query. For example, if the network output is an image, the input features of the pixels in the image can be, for example, the patches of the image around that pixel concatenated into a vector. In another example, if the output is an audio waveform, the input features at a particular point in the audio waveform can be the amplitude of the waveform at that point.

[0114] In the case of multi-task or multi-modal settings, embeddings can be used for each task or modality respectively. In a classification task, a single learned embedding can be used. In particular, as described above with reference to FIG. 1, the training engine can train a neural network. As part of the training, the training engine can train the embedding simultaneously with the neural network, for example, by backpropagating the gradient of the loss through the neural network into the embedding. In the case of tasks with heterogeneous outputs such as multi-modal autoencoding, features specific to some queries (for example, spatial locations represented by x and y coordinates) can be combined with the modality embedding. In the case of multi-modal autoencoding, each query embedding can specify the modality that it is desirable to reconstruct.

[0115] FIG. 6 shows an exemplary modality 600 of an input that can be processed by a neural network system (e.g., system 100 of FIG. 1, or system 100 of FIG. 2). This system can be used in domains with various input-output spaces, including multi-task language understanding, high-density visual tasks such as optical flow, and hybrid high-density / sparse multi-modal tasks such as video + audio + class auto-encoding. In some implementations, the input can be pre-processed and / or post-processed, for example, to reduce the size of the input and / or output, respectively, which can be very large. The last two columns show the number of dimensions of the input and the respective output.

[0116] FIGS. 7A and 7B show exemplary multi-modal results 700, such as auto-encoding of audio-video-labels. The input is shown on the left side and the reconstruction is shown on the right side. The neural network system can jointly represent modalities with very different characteristics and achieve substantially high prediction accuracy. In FIG. 7B, "PSNR" refers to the peak signal-to-noise ratio.

[0117] As used herein, the term "configured" is used in relation to systems and computer program components. A system consisting of one or more computers being configured to perform a particular operation or action means that the system has software, firmware, hardware, or a combination thereof installed thereon that causes the system to perform that operation or action during operation. A one or more computer programs being configured to perform a particular operation or action means that the one or more programs include instructions that, when executed by a data processing apparatus, cause the apparatus to perform the operation or action.

[0118] Embodiments and functional operations of the subject matter described in this specification can be implemented in digital electronic circuitry, tangibly embodied computer software or firmware, computer hardware including the structures disclosed in this specification and their structural equivalents, or one or more combinations of them. Embodiments of the subject matter described in this specification can be implemented as one or more computer programs, i.e., as 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, a data processing apparatus. The computer storage medium can be a machine-readable storage device, a machine-readable storage substrate, a random access memory device, or a serial access memory device, or one or more combinations of them. Alternatively or additionally, 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 an appropriate receiver device for execution by a data processing apparatus.

[0119] The term "data processing apparatus" refers to data processing hardware and includes any kind of apparatus, device, and machine for processing data, including by way of example a programmable processor, a computer, or multiple processors or computers. The apparatus can be, or further include, a special purpose logic circuit, e.g., an FPGA (field programmable gate array) or an ASIC (application specific integrated circuit). The apparatus can optionally include code that creates an execution environment for computer programs, e.g., processor firmware, a protocol stack, a database management system, an operating system, or one or more combinations of them.

[0120] A computer program, which may also be referred to or described as a program, software, software application, app, module, software module, script, or code, can be written in any form of programming language, including compiler languages or interpreter languages, or declarative languages or procedural languages, and can be deployed in any form, such as a stand-alone program or as a module, component, subroutine, or other unit suitable for use in a computing environment. The program may or may not correspond to a file in a file system. The program can be stored as part of a file that holds other programs or data, such as one or more scripts stored in a markup language document, a single file dedicated to the program in question, or multiple coordinated files, such as files that hold one or more modules, subprograms, or portions of code. The computer program can be deployed to be executed on one computer or on multiple computers located at one site or distributed across multiple sites and interconnected by a data communication network.

[0121] As used herein, the term "engine" is widely used to refer to a software-based system, subsystem, or process programmed to perform one or more specific functions. Generally, an engine is implemented as one or more software modules or components and installed on one or more computers at one or more locations. In some cases, one or more computers may be dedicated to a particular engine, and in some cases, multiple engines can also be installed and executed on one or more of the computers in question.

[0122] The processes and logical flows described herein can be performed by one or more programmable computers executing one or more computer programs to manipulate input data and generate output in order to perform functions. The processes and logical flows can also be performed by, for example, a special-purpose logic circuit such as an FPGA or ASIC, or by a combination of special-purpose logic circuits and one or more programmed computers.

[0123] Computers suitable for the execution of a computer program can be based on general or special purpose microprocessors or both, or on other types of central processing units. Generally, a central processing unit receives instructions and data from a read only memory, a random access memory, or both. Essential elements of a computer are a central processing unit for executing or performing 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 circuits. Generally, a computer also includes, or is operatively coupled to, one or more mass storage devices for storing data, such as, by way of example only, magnetic, magneto-optical disks, or optical disks, or receives data from them, or transfers data to them, or both. However, a computer need not have such devices. Further, a computer can be embedded in another device, such as a mobile telephone, a personal digital assistant (PDA), a mobile audio or video player, a game console, a global positioning system (GPS) receiver, a portable storage device, such as a universal serial bus (USB) flash drive, to name just a few.

[0124] Computer-readable media suitable for storing computer program instructions and data include, by way of example, semiconductor memory devices such as, for example, EPROM, EEPROM, flash memory devices, magnetic disks such as internal hard disks or removable disks, magneto-optical disks, and any form of non-volatile memory, media, and memory devices such as CD-ROM and DVD-ROM disks.

[0125] To provide for interaction with a user, embodiments of the subject matter described herein can be implemented on a computer having a display device for displaying information to the user such as, for example, a CRT (cathode ray tube) or LCD (liquid crystal display) monitor, as well as a keyboard and a pointing device such as a mouse or 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 such as, for example, visual feedback, auditory feedback, or tactile feedback, and input received from the user can be in any form including, but not limited to, acoustic, speech, or tactile input. Further, the computer can interact with the user by sending and receiving documents between the devices used by the user, for example, by sending a web page to a web browser on the user's device in response to a request received from the web browser. Also, the computer can interact with the user by sending a text message or other form of message to a personal device such as a smartphone running a messaging application and receiving, in return, a response message from the user.

[0126] A data processing apparatus for implementing a machine learning model can also include, for example, a dedicated hardware accelerator unit for processing the general and computationally intensive parts of machine learning, i.e., inference, workloads, for example, during training or production.

[0127] The machine learning model can be implemented and deployed using a machine learning framework such as the TensorFlow framework.

[0128] Embodiments of the subject matter described in this specification can be implemented in a computing system that includes, for example, a back-end component as a data server, or a computing system that includes a middleware component such as an application server, or a front-end component such as a graphical user interface, a web browser, or a client computer equipped with an app that allows a user to interact with an implementation of the subject matter described in this specification, or a computing system that includes 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, such as a communication network. Examples of communication networks include local area networks (LANs) and wide area networks (WANs), such as the Internet.

[0129] The computing system can include a client and a server. The client and the server are typically remote from each other and typically interact through a communication network. The relationship between the client and the server is generated by computer programs that are executed on respective computers and have a client-server relationship with each other. In some embodiments, the server sends data, such as an HTML page, to a user device for the purpose of displaying the data to a user who interacts with a device that functions as a client, for example, and receiving user input from the user. Data generated at the user device, such as the result of a user interaction, can be received at the server from the device.

[0130] Although this specification contains many details, these should not be construed as limiting the scope of the invention or the scope of the claims that may be claimed, but rather as descriptions of features that may be specific to particular embodiments of a particular invention. Specific features described in the context of separate embodiments herein may also be implemented in combination in a single embodiment. Conversely, various features described in the context of a single embodiment may also be implemented separately in multiple embodiments or in any suitable sub-combination. Furthermore, features have been described above as functioning in a particular combination and may initially be claimed as such in the claims, but in some cases, one or more features from the claimed combination may be removed from the claimed combination, and the claimed combination may be directed to a sub-combination or a variation of a sub-combination.

[0131] Similarly, operations are shown in the drawings and described in the claims in a particular order, but this should not be understood as requiring that such operations be performed in the particular order or sequence shown, or that all of the operations shown be performed, in order to achieve a desired result. In certain circumstances, multitasking and parallel processing may be advantageous. Furthermore, the separation of the 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 program components and systems described may generally be integrated into a single software product or packaged into multiple software products.

[0132] Specific 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 the desired result. As one example, the processes shown in the accompanying drawings need not be in the particular order or sequence shown to achieve the desired result. In some cases, multitasking and parallel processing may be advantageous.

Description of the Reference Numerals

[0133] 100 Neural network system 102 Neural network input 104 Latent embedding 106 Query embedding 108 Network output 110 Encoder block 120 Output block 150 Neural network 160 Neural network 202 Network input 202 2D input array 202 Data element embedding 204 Latent embedding 206 Query embedding 208 Network output 210 Encoder block 230 Cross-attention block 240 Self-attention block 250 Neural network 250 Output block 300 Process 400 Query embedding 402 Network input 402 Image 404 Latent embedding 408 Neural network output 410 Encoder 420 Output block 500 Query Embedding 600 Modality 700 Exemplary Multimodal Results

Claims

1. A method executed by one or more data processing devices, comprising: obtaining (i) a network input to a neural network and (ii) a set of query embeddings that collectively define a prediction task performed by the neural network by processing the network input; processing the network input using the neural network to generate a network output having respective dimensions corresponding to each query embedding in the set of query embeddings, processing the network input using an encoder block of the neural network to generate a representation of the network input as a set of potential embeddings, and processing (i) the set of potential embeddings and (ii) the set of query embeddings using a cross-attention block that generates each dimension of the network output by cross-attention of the corresponding query embedding to the set of potential embeddings. A method comprising the steps above.

2. For each one or more of the query embeddings, wherein the query embedding defines a respective spatial position within the network input, and wherein the dimension of the network output corresponding to the query embedding defines a respective prediction related to the spatial position within the network input. The method according to claim 1.

3. For each one or more of the query embeddings, wherein the query embedding comprises one or more input features from a respective spatial position within the network input, and wherein the dimension of the network output corresponding to the query embedding defines a respective prediction related to the spatial position within the network input. The method according to claim 1 or 2.

4. For each one or more query embeddings, wherein the query embedding defines a modality of the network input. The method according to any one of claims 1 to 3.

5. wherein the set of query embeddings collectively defines a plurality of prediction tasks, and for each one or more of the query embeddings, wherein the query embedding specifies a respective prediction task from the plurality of prediction tasks. The method according to any one of claims 1 to 4, wherein the dimension of the network output corresponding to the query embedding defines the prediction output of the prediction task specified by the query embedding.

6. The method according to claim 5, wherein the plurality of prediction tasks comprises one or more of a classification task, a regression task, a segmentation task, or an autoencoding task.

7. The step of generating the dimension of the network output by cross-attention of the corresponding query embedding with respect to the set of potential embeddings comprises: (i) generating respective attention weights for each potential embedding based on (i) the query embedding and (ii) the potential embedding; generating the dimension of the network output corresponding to the query embedding based on the attention weights for the potential embedding. The method according to any one of claims 1 to 6.

8. The step of generating the dimension of the network output corresponding to the query embedding based on the attention weights for the potential embedding comprises: processing each potential embedding to generate a value embedding of the potential embedding; combining the value embeddings using the attention weights; generating the dimension of the network output corresponding to the query embedding based at least in part on a result of combining the value embeddings using the attention weights. The method according to claim 7.

9. The method according to claim 8, wherein the cross-attention is query-key-value attention.

10. The method according to any one of claims 1 to 9, wherein the number of query embeddings in the set of query embeddings is greater than the number of potential embeddings in the set of potential embeddings.

11. The method according to claim 10, wherein the number of query embeddings in the set of query embeddings is at least twice greater than the number of potential embeddings in the set of potential embeddings.

12. The method according to any one of claims 1 to 11, wherein the number of potential embeddings in the set of potential embeddings is predefined, the number of query embeddings in the set of query embeddings is variable, and is independent of the number of potential embeddings in the set of potential embeddings.

13. The method according to any one of claims 1 to 12, wherein the network input comprises multimodal data.

14. Further comprising the step of initializing the set of potential embeddings, wherein the network input comprises a set of data element embeddings, wherein the encoder block of the neural network, uses attention for some or all of the data element embeddings within the set of data element embeddings to update each potential embedding within the set of potential embeddings, and each performs an operation comprising the step of updating each potential embedding within the set of potential embeddings, and comprises one or more cross-attention blocks, wherein the encoder block of the neural network, uses attention for the set of potential embeddings to update each potential embedding within the set of potential embeddings, and each performs an operation comprising the step of updating each potential embedding within the set of potential embeddings, and comprises one or more self-attention blocks; the method according to any one of claims 1 to 13.

15. The method according to claim 14, wherein the number of potential embeddings within the set of potential embeddings is less than the number of data element embeddings within the set of data element embeddings.

16. The method according to claim 14 or 15, wherein each data element embedding corresponds to a respective spatial position within the network input and comprises a feature embedding based on the features of the network input at the spatial position.

17. The method according to any one of claims 14 to 16, wherein each data element embedding corresponds to a respective spatial position within the network input and comprises a position embedding characterizing the spatial position.

18. The method according to any one of claims 14 to 17, wherein each data element embedding comprises a modality embedding defining the modality corresponding to the data element embedding.

19. One or more computers, and one or more storage devices communicatively coupled to the one or more computers, wherein the one or more storage devices store instructions that, when executed by the one or more computers, cause the one or more computers to perform the operations of the respective methods according to any one of claims 1 to 18; a system.

20. One or more non-transitory computer storage media that, when executed by one or more computers, store instructions that cause the one or more computers to perform the operations of the respective methods of any one of claims 1 to 18.

Images