Cross-attention of the query embedding against a set of latent embeddings to generate neural network outputs

Abstract

A method, system, and apparatus, including a computer program encoded on a computer storage medium, for generating a network output using a neural network. In one aspect, the method comprises obtaining (i) a network input to the neural network and (ii) a set of query embeddings, and processing the network input using the neural network to generate a network output having a respective dimension corresponding to each query embedding in the set of query embeddings, the method comprising: 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 (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 a corresponding query embedding on the set of potential embeddings.

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 May 28, 2021, for “GENERATING NEURAL NETWORK OUTPUTS BY CROSS ATTENTION OF QUERY EMBEDDINGS OVER A SET OF LATENT EMBEDDINGS,” the contents of which are incorporated by reference herein in their entirety.

[0002] This specification relates to processing data using machine learning models. [Background technology]

[0003] A machine learning model receives an input and generates an output, e.g., a predicted output, based on the received input. Some machine learning models are parametric models, which generate an output based on the received input and the values ​​of the parameters of the model.

[0004] Some machine learning models are deep models that use multiple model layers to generate an output for a 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 that each apply a nonlinear transformation to the received input to generate an output. Summary of the Invention [Means for solving the problem]

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

[0006] In one aspect, a method is described that is executed by one or more data processing devices. The method comprises 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 inputs. The method processes the network inputs using the neural network to generate a network output with a respective dimension corresponding to each query embedding in the set of query embeddings. The method processes the network inputs using an encoder block of the neural network to generate a representation of the network inputs as a set of latent 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 a corresponding query embedding on the set of potential embeddings.

[0007] Throughout this specification, an embedding refers to an ordered collection of numbers, for example a vector or matrix of numbers.

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

[0009] According to one aspect, there is provided a method, performed 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 to be performed by the neural network by processing the network inputs; processing the network inputs using the neural network to generate a network output having a respective dimension corresponding to each query embedding in the set of query embeddings, the processing comprising: processing the network inputs using an encoder block of the neural network to generate a representation of the network inputs 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 a corresponding query embedding on the set of potential embeddings.

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

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

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

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

[0014] In some implementations, the multiple prediction tasks comprise one or more of classification tasks, regression tasks, segmentation tasks, or auto-encoding tasks.

[0015] In some implementations, generating a dimensionality of the network output by cross-attention of a corresponding query embedding to a set of potential embeddings comprises generating respective attention weights for each potential embedding based on (i) the query embedding and (ii) the potential embeddings, and generating a dimensionality of the network output corresponding to the query embedding based on the attention weights for the potential embeddings.

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

[0017] In some implementations, the cross attention is a 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 potential embeddings in the set of potential embeddings.

[0019] In some implementations, the number of query embeddings in the set of query embeddings is at least two times greater than 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, and the number of query embeddings in the set of query embeddings is variable and independent of the number of potential embeddings in the set of potential embeddings.

[0021] In some implementations, the network input comprises multi-modal data.

[0022] In some implementations, the method further comprises initializing a set of potential embeddings, where the network input comprises a set of data element embeddings, the encoder block of the neural network comprises one or more cross-attention blocks each performing an operation comprising updating each potential embedding in the set of potential embeddings using attention to 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 comprising updating each potential embedding in the set of potential embeddings using attention to 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 location in the network input and comprises a feature embedding based on the features of the network input at that spatial location.

[0025] In some implementations, each data element embedding corresponds to a respective spatial location within the network input and comprises a positional embedding that characterizes the spatial location.

[0026] In some implementations, each data element embedding comprises a modality embedding that defines the modality that corresponds 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 the 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 system described herein may process the network inputs to generate a representation of the network inputs as a set of latent embeddings, and then generate each dimension of the network output by cross-attention of the corresponding query embedding to the set of latent embeddings. Generating each dimension of the network output by cross-attention of the query embedding to the set of latent 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 the hidden layers of the neural network.

[0031] Furthermore, the prediction task performed by the neural network is defined by and can be modified using the query embedding. Thus, the neural network system provides a flexible and general neural network architecture that can be implemented to perform a variety of prediction tasks for network inputs of different modalities with minimal modifications. This allows the neural network system to 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, e.g., fixed-size image segmentation outputs, the neural network systems described herein can use query embeddings to generate sparse or partial network outputs. For example, rather than generating a complete segmentation of an image, the neural network system can generate a segmentation of only a specified appropriate subset of the image, e.g., by querying a set of potential embeddings using only the relevant query embeddings. Thus, the neural network system can reduce consumption of computational resources (e.g., memory and computing power) by providing an option to generate sparse or partial network outputs, when appropriate. For example, the neural network system can be used to segment only regions of interest of a larger image, without segmenting the remainder of the image.

[0033] More generally, the size of the network output can be defined independent of the size of the network input, since the dimension of the network output corresponds to the number of query embeddings in the set of query embeddings. Furthermore, the use of cross-attention avoids the quadratic dependency of the computations performed on the size of the network inputs and outputs, and instead the dependency is linear, facilitating the processing of, for example, video, audio, and multimodal data. Similarly, for example, the network output is generated by querying potential embeddings using cross-attention, so implementations of the described technique use a set of potential embeddings that can be independent of the spatial or temporal structure of the network inputs and outputs. This facilitates the processing of multimodal data, i.e., network inputs with different structures or dimensionality. Furthermore, some implementations of the system can perform tasks, e.g., optical flow determination tasks, more efficiently and using simpler architectures than some other techniques.

[0034] The details of one or more embodiments of the subject matter herein 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, drawings, and claims. [Brief description of the drawings]

[0035] [Figure 1] FIG. 1 is a block diagram of an example neural network system capable of generating a neural network output. [Diagram 2] FIG. 2 is a more detailed block diagram of an example neural network system. [Diagram 3] 4 is a flow diagram of an example process for generating a neural network output using a neural network system. [Figure 4] FIG. 2 illustrates an exemplary query embedding. [Diagram 5] FIG. 13 is a diagram illustrating another example of query embedding. [Figure 6] FIG. 1 illustrates exemplary modalities of input that may be processed by a neural network system. [Figure 7A] FIG. 1 illustrates an exemplary multi-modal result. [Figure 7B] FIG. 1 illustrates an exemplary multi-modal result. DETAILED DESCRIPTION OF THE PREFERRED EMBODIMENTS

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

[0037] 1 is a block diagram of an exemplary neural network system 100. Neural network system 100 is one example of a system implemented as a computer program on one or more computers at one or more locations in which the systems, components, and techniques described below may be 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 may include any suitable type of data. That is, the neural network 150 may be configured to process inputs of any of multiple modalities, e.g., images, video, audio, text, or other suitable modalities, or any combination thereof. The neural network 150 may have a single set of parameters, while being substantially flexible such that it can process any one or combination of multiple modalities.

[0040] In some implementations, the network input 102 may characterize an entity. An entity may 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. As used herein, an image, i.e., a still image or a moving image (video), may include a point cloud.

[0041] As a particular example, an entity may include multiple units arranged in a spatial or temporal structure, e.g., as shown in FIG. 4, an entity may be an image and each unit may be a pixel in the image. Each unit in an entity may have an associated data element embedding that may characterize, e.g., the location of the unit in the spatial or temporal structure and / or features associated with the unit in the spatial structure. In another particular example, an entity may be text, e.g., an entity may be a sequence of words, phrases, characters, or word fragments in one or more languages. In yet another particular example, an entity may include a combination of different modalities, e.g., an entity may include a combination of audio, video, and label data. As another example, an entity may comprise sensor data from one or more sensors configured to recognize one or more characteristics of the real environment, such as image data, audio data, or so-called "fine touch" sensor data (which allows for localization). Although several examples of network inputs 102 are described above, in general, network inputs 102 may have any suitable dimensionality and structure and may include any suitable type of data.

[0042] The neural network system 100 may process the network input 102 using the neural network 150 to generate a network output 108. The network output 108 may be, for example, a classification output, a regression output, a sequence output (i.e., including a sequence of output elements), a segmentation output, an auto-encoding 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 (proper) subset of the image. In general, the network output 108 can be represented as a collection of output vectors, and the number of output vectors can be referred to as the "dimensionality" of the network output 108. As 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 network input 102, system 100 may further obtain a set of query embeddings 106. Network input 102 may be provided by a user of system 100, for example, through an application programming interface (API) made available by system 100.

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

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

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

[0048] As a particular example, the network input 102 may include representations of a first image and a second image each depicting the same scene, and the prediction task is to estimate the two-dimensional displacement (e.g., optical flow) of a point, e.g., a pixel or a proper subset of pixels (i.e., not all pixels) in the first image. In some implementations, such a point may comprise a pixel or an image patch comprising a plurality of pixels. In such a case, each query embedding in the set of query embeddings 106 may specify, e.g., the x and y coordinates of a corresponding point in the first image, and may include one or more further embeddings for the optical flow prediction task, e.g., the first and second image representations of a point. That is, the set of query embeddings 106 may define a set of points in the first image for generating a prediction, and the prediction is to estimate the optical flow. For example, the dimensionality of the network output may define the magnitude and direction of the optical flow at a point. Examples of query embeddings and prediction tasks are described in more detail below with reference to Figures 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 the network input 102 in order to generate a representation of the 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 the set of potential embeddings.

[0051] In some implementations, the number of dimensions of the network input 102 can be much larger than the number of dimensions of the potential embeddings 104 and / or the query embeddings 106. For example, if the 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 the network input 102 in any suitable way to generate a representation of the input as the set of potential embeddings 104. A specific example of this process will be described in more detail below with reference to FIG. 2.

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

[0053] In particular, the output block 120 can update each query embedding in the set of query embeddings 106 to the set of potential embeddings 104. Continuing with the optical flow example above, each query embedding in the set of query embeddings 106 can specify, for example, the x and y coordinates of a 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 embedding to the set of potential embeddings 106. Returning to the previous example as an illustration, the dimension of the network output 108 can be, for example, a single flow vector that represents 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 in the set of query embeddings 106, thereby generating the respective dimension of the network output 108. In some implementations, the outputs corresponding to each query embedding in the set of query embeddings 104 can be generated in parallel. This process is described in more detail below with reference to FIG. 2.

[0054] The encoder block 110 and the output block 120 may have any suitable neural network architecture that enables them to perform a given function. For example, the encoder block 110 and the output block 120 may 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 may be connected in any suitable configuration (e.g., as a linear sequence of layers). The neural network system 100 may also further include any number of neural network blocks configured to perform any suitable operation. Particular examples of the encoder block 110 and the output block 120 are described in more detail below with reference to FIG. 2.

[0055] In this manner, system 100 can flexibly modify the dimensionality of network output 108, for example, by modifying the number of query embeddings 106, without reconstructing the hidden layers of neural network 150. Furthermore, as described above, the prediction task performed by neural network 150 is defined by, and in some implementations can be modified using, query embeddings 106. Thus, neural network system 100 can provide a flexible and general neural network architecture that can be implemented to perform a variety of prediction tasks for network inputs of different modalities with minimal modifications.

[0056] The neural network system 100 may further include a training engine capable of training the neural network 150 on a set of training data over multiple training iterations. The training data may include a set of training examples, where 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 may sample a batch of training examples from the training data and process training inputs specified by the training examples using a sequence of neural network blocks included in neural network 150 (e.g., encoder block 110 and output block 120) to generate corresponding network outputs. In particular, for each training input, neural network 150 processes the training input using the current model parameter values ​​in 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 ​​in output block 120 to generate network outputs 108 corresponding to the training inputs.

[0058] The training engine can adjust model parameter values ​​of the encoder block 110 and the output block 120 to optimize an objective function that measures the similarity between (i) the network outputs generated by the neural network 150 and (ii) the target network outputs specified by the training examples. The objective function can be, for example, a cross-entropy objective function, a squared error objective function, or any other suitable objective function. In some implementations, during training, the network outputs generated by the neural network 150 can include only a portion of the full 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 a gradient of the objective function, for example, using backpropagation techniques. The training engine can use the gradient to update model parameter values ​​of the encoder block 110 and the output block 120, for example, using any suitable gradient descent optimization algorithm, for example, Adam. The training engine can determine a performance measure of the neural network 150 based on a set of validation data that is not used during training of the neural network 150. After training, the neural network system 100 can be used to perform machine learning tasks, for example, to process network inputs and generate network outputs.

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

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

[0062] In some implementations, the system 100 uses the neural network 150 to perform image or audio segmentation tasks. For example, the neural network can process a network input including an image representing an audio waveform or a series of audio samples to generate a network output that defines, for each pixel in the input image or for each audio sample, a respective score distribution for a set of possible classes. The class score can define the likelihood that the corresponding pixel or sample is included in the class. For example, for an image, the possible classes can include, for example, water, building, vehicle, pedestrian, etc. In this example, the query embedding can specify the location of the corresponding pixel in the input image or the location 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 moving image, i.e., a video. In a similar manner, the system 100 can use the neural network 150 to perform an image depth prediction task, where the network output defines a predicted depth value of the corresponding pixel to obtain a (spatial 3D) depth map of the image.

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

[0064] In some implementations, the system 100 uses a neural network 150 to perform an agent control task, where the network inputs represent a sequence of one or more observations or other data that characterize the state of the environment, and the outputs define actions to be performed by the agent in response to the most recent data in the sequence. The environment may be a real or simulated environment, and the agent may be, for example, a robot, an autonomous land, sea, or air vehicle, or a control system for an industrial facility. Each query embedding may specify a respective task to be performed by the agent in the environment, for example, the task may be for the agent to move to a goal location in the environment, and the query embedding may specify the goal location.

[0065] In some implementations, the system 100 performs speech recognition using the neural network 150 to process network inputs representing audio samples in an audio waveform to generate a network output that characterizes, for example, 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 may also include an encoding of the position of the samples in the sequence. Each dimension of the network output may, for example, correspond to a respective time interval in the audio waveform and define a respective score distribution over a set of possible phonemes, graphemes, characters, or words. In this example, the neural network may generate a network output by, for example, (i) a respective time interval in the audio waveform, for example, a position in a sequence, and optionally (ii) a query embedding that defines a prediction task for the time interval in the audio waveform, for example, decoding the phoneme, grapheme, character, or word that corresponds to the time interval.

[0066] In some implementations, the system 100 uses the neural network 150 to perform a reconstruction task, such as an auto-encoding task, by processing the network inputs to generate a network output that defines a predicted reconstruction of the network inputs. The network inputs 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 location in the network input to be reconstructed, for example a spatial location in an image, a spatial location and / or a temporal location in a video, or a time window in an audio waveform. Optionally, each query embedding can also define the modality of the network input to be reconstructed, for example, 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, for example, a 1D, 2D, or 3D network input specified by the query embedding. More or fewer points may be reconstructed than the original network input. The points may be reconstructed serially and / or in parallel. For example, the set of query embeddings may comprise a query embedding for each point in the network input, e.g., a pixel of an image or a sample of an audio signal, it may comprise a query embedding for just one point, the network output may be reconstructed successively one point at a time by specifying each desired point, or the set of query embeddings may reconstruct some but not all of the points in the network input.

[0067] In some implementations, the system 100 employs the neural network 150 for neural machine translation tasks to process a network input representing, for example, a sequence of text, e.g., a sequence of words, phrases, characters, or word fragments, in one language to generate a network output that translates 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. In this example, each query embedding may include data identifying, for example, the natural language (e.g., English, French, German, etc.) into which the input text sequence should be translated. Additionally, each query embedding may specify a corresponding position in the output sequence of text.

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

[0069] In some implementations, the system 100 uses a neural network 150 to perform a text-to-speech task, where the network input represents text in a natural language or features of text in a natural language, and the network output is a spectrogram, waveform, or other data defining the audio of the text spoken in the natural language. In this example, each query embedding may be for identifying a respective speaking voice used in uttering the text (e.g., where each speaking voice corresponds to a respective person's voice). Additionally, each query embedding may specify a corresponding time window in the network output.

[0070] In some implementations, the system 100 uses the neural network 150 to perform a health prediction task, where the network inputs represent data obtained from the patient's electronic health record data, and the outputs are predictions related to the patient's future health, such as predicted treatments to be prescribed to the patient, the likelihood of the patient experiencing an adverse health event, or the patient's predicted diagnosis. In this case, each query embedding can specify a corresponding time or period in the patient's life. In some implementations, the system 100 uses the neural network 150 to perform a text generation task, where the network inputs represent a sequence of text, and the outputs are another sequence of text, such as a completion of the input sequence of text, a response to a question posed in the input sequence, or a sequence of text related to a topic specified by the first text sequence. As another example, the network inputs can represent data other than text, such as images, and the output sequence can be text describing the data represented by the network inputs. In this example, each query embedding can specify a corresponding position in the output sequence of text.

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

[0072] In some implementations, the system 100 uses a neural network 150 to perform genomics tasks, where the network inputs represent fragments of DNA sequences or other molecular sequences, and the network outputs are either embeddings of the fragments for use in downstream tasks, e.g., by utilizing unsupervised learning techniques on a dataset of DNA sequence fragments, or outputs for downstream tasks. Examples of downstream tasks include promoter site prediction, methylation analysis, predicting functional effects of non-coding variants, etc. In this example, each query embedding may specify a corresponding position in the DNA sequence.

[0073] In some implementations, the system 100 uses the neural network 150 to perform image, video, or audio classification tasks. The network inputs may then represent image or video pixels, or audio waveform samples, as described above. Optionally, the network inputs may include positional encodings of the image or video pixels, or audio samples, where in the case of video, the positional encodings may be spatial and / or temporal locations. The network output defines a classification of the image, video, or audio. The classification may include a respective score for each object category (e.g., image, vehicle, pedestrian, bicycle, etc.) in a set of possible object categories. The score of an object category may 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 may comprise a classification of an action depicted in the video, e.g., for gesture recognition. The set of query embeddings may comprise a single query embedding (which may be learned). In a similar manner, system 100 can use neural network 150 to perform image or audio bounding box tasks, with the network output comprising a classification and a vector defining the coordinates of the object's bounding box in one, two, or three dimensions (time counts here as a dimension).

[0074] As mentioned above, in general in the above examples, the references to images or videos include point clouds. For example, in some implementations, the system 100 uses a 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 may specify the corresponding spatial location of one or more points in the point cloud.

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

[0076] In some implementations, the network input is a multimodal input, and the system 100 uses the neural network 150 to perform a multimodal task. In general, such a multimodal input is a combination of two or more different types or modalities of data, where the different types of data may, but need not, represent the same object or overlapping objects in the network input using different modalities. 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 textual data, e.g., tokens, representing text in a natural language, and pixels of an image or video, or samples of an audio waveform. A multimodal task may be any of the above tasks, but may include additional modes of 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, each modality type may be provided with a modality-specific embedding, and then the multimodal data may be serialized into the 2D input array 202. In general, the query embedding may be as described above, but in some implementations, modality-specific embeddings are added to generate outputs for each modality. The embeddings may be learned embeddings. If spatial or temporal locations are incorporated into the query embedding, e.g., for video or audio data, this may be as described above, since the modality-specific embeddings allow the system to select each modality. The network output may correspond to those described above, but adapted to be 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 multiple categories, e.g., by defining a per-category score for multiple possible categories for the input.Similarly, a captioning task generates a network output of tokens describing the multimodal input, and the corresponding queries may, but need not, include spatial or temporal locations for each query embedding, etc. As another example, if the multimodal data comprises a combination of text data and image, video, or audio data, the task may comprise processing the combination to provide a network output defining whether an image, video, or audio waveform is described by the text, e.g., by defining a score for the text or caption, by a particular caption. The queries may be as described above with respect to classification, and may include location data for generating a respective score for each location represented by one of the query embeddings in the set of query embeddings.

[0077] In general, in the above examples, the network input may comprise raw data, e.g. raw pixel data or audio samples, or feature embeddings that define spatial, temporal, or spatio-temporal features obtained by pre-processing, e.g. with a feature encoder neural network.

[0078] An exemplary architecture for neural network system 100 is described in more detail below with reference to FIG.

[0079] 2 is a more detailed block diagram of an example architecture of neural network system 100. Neural network system 100 is one 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, which can include a set of data element embeddings (e.g., an "input sequence") that characterize an entity. An entity can include, for example, multiple units arranged in a spatial and / or temporal structure.

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

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

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

[0084] The neural network system 100 may further obtain a set of query embeddings 206 (e.g., “output query arrays”) that collectively define a prediction task to be performed by the neural network 250 by processing the network inputs 202. The query embeddings 206 may have any suitable number of dimensions, for example, the query embedding may include O query vectors.

[0085] In some implementations, the neural network system 100 may further obtain an initial set of potential embeddings 204 (e.g., a "potential array"). The potential embeddings 204 may be predefined and / or randomly initialized, or they may be learned, i.e., they may comprise a set of learned parameters, e.g., weights, each of which defines an element of the array. That is, in some implementations, the encoder block processes the network input and the initial set of potential embeddings to generate a representation of the network input as a set of potential embeddings. The set of potential embeddings 204 may have dimensions N×D, where N is the number of potential embeddings and D is the number of channels of each potential embedding, both of which may be hyperparameters of the neural network system 100. In some implementations, the number O of query embeddings may be greater than the number N of potential embeddings, e.g., by at least a factor of two. Furthermore, the number O of query embeddings may be variable and independent of the number N of potential embeddings.

[0086] The neural network system 100 can process the network inputs 202 using a neural network 250 to generate a network output 208 (e.g., an "output array") that includes a respective dimension corresponding to each query embedding in 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) in each output vector. The neural network system 100 can generate a respective dimension (e.g., a respective output vector) corresponding to each query embedding in 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 an encoder block 210 and an output block 250. The encoder block 210 can include a sequence of one or more neural network blocks, for example, (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 are 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 attention operations, for example, to update each embedding in the first set of embeddings using attention to the second set of embeddings.

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

[0090] In some implementations, the attention block can perform a query key value (QKV) attention operation, e.g., updating each embedding in the first set of embeddings using attention to the second set of embeddings using a query (Q), key (K), and value (V) embedding. In particular, each attention block can include (i) a query sub-network, (ii) a key sub-network, and (iii) a value sub-network. For each target embedding in the first set of embeddings, the query sub-network 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 sub-network 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 sub-network 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] Each attention block can then use the query embedding (Q), key embedding (K), and value embedding (V) to update each target embedding in the first set of embeddings using attention to the set of second embeddings. Specifically, each attention block can generate an attention weight for each embedding in the second set of embeddings, e.g., as a dot product (e.g., dot) of the query embedding (Q) with each of the key embeddings (K). Based on the set of second embeddings and the attention weights, each attention block can generate a combined embedding, e.g., as a linear combination of the value embeddings (V) weighted by their respective attention weights. Finally, each attention block can use the combined embedding to update the target embedding in the first set of embeddings, e.g., by replacing the target embedding in 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., QKV attention operation) may be referred to as a "self-attention" operation. The self-attention operation may be performed, for example, by the self-attention block 240. For example, the first set of embeddings may be a set of potential embeddings 204, the second set of embeddings may also be a set of potential embeddings 204, and the self-attention block 240 may use self-attention on the set of potential embeddings 204 to update each potential embedding in the set of potential embeddings 204. In some implementations, the self-attention block 240 may iteratively update each potential embedding in the set of potential embeddings 204 using self-attention on the set of potential 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 may be performed, for example, by a cross-attention block 230 and an output block 250 included in the encoder 210.

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

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

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

[0097] An exemplary process for using the neural network system 100 to generate the neural network output 208 will now be described in more detail.

[0098] 3 is a flow diagram of an exemplary process 300 for generating a neural network output using a neural network system. For convenience, the process 300 is described as being performed by one or more computer systems located at one or more locations. For example, a neural network system, such as the neural network system 100 of FIG. 1 or the neural network system 100 of FIG. 2, when appropriately programmed in accordance with this specification, can perform the process 300.

[0099] The system obtains (302) (i) network inputs to the 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 inputs. In some implementations, the network inputs can include multi-modal data, such as video data, audio data, and label data, by way of example. In general, as shown in FIG. 6A, the network inputs can include any suitable type or combination of data.

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

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

[0102] The system can then process (i) the set of potential embeddings and (ii) the set of query embeddings using a cross-attention block, e.g., output block 250 described above, which generates each dimension of the network output by cross-attention (e.g., query-key-value attention, etc.) of the corresponding query embedding to the set of potential embeddings. For example, as described above with reference to FIG. 2, the system can generate respective attention weights for each potential embedding based on (i) the query embedding and (ii) the potential embedding, and generate a dimension of the network output corresponding to the query embedding based on the attention weights of the potential embedding. In particular, the system can process each potential embedding to generate a value embedding for the potential embedding. The system can then combine the value embeddings using the attention weights. Finally, the system can generate a 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 embeddings may define respective spatial locations in the network input. In such a case, the dimensions of the network output corresponding to the query embeddings may define respective predictions associated with the spatial locations of the network input. For example, as shown in FIG. 4, the network input may be an image having a size of 100×100 pixels, and each query embedding may define respective spatial locations (e.g., x and y coordinates) of corresponding pixels in the image. The prediction task may be, for example, reconstructing the original image. The set of query embeddings may include 10,000 query vectors, e.g., vectors for each pixel in the image. The dimensions of the network output corresponding to the query embeddings may define respective predictions associated with the spatial locations of the network input, e.g., reconstruction of pixels in the image defined by the query embeddings.

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

[0105] In some implementations, the query embedding can define the modality of the network input. For a multimodal input, e.g., an input that includes video, audio, and labels, all three modalities can be specified by each query embedding in the set of query embeddings. In general, 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, the set of query embeddings can collectively define multiple prediction tasks, e.g., each query embedding can define a respective prediction task. In such cases, the dimensionality of the network output corresponding to the query embedding can define the prediction output of the prediction task specified by the query embedding. The prediction task can include, for example, a classification task, a regression task, a segmentation task, an auto-encoding task, or other suitable task. For example, as shown in FIG. 5, the embedded "task_id" can specify a particular prediction task. For eight prediction tasks, the set of query embeddings can include eight query embeddings, each specifying a particular prediction task.

[0107] An exemplary query embedding and prediction task is described in more detail below with reference to FIGS.

[0108] FIG. 4 illustrates an example query embedding 400 that may be used by a neural network system (eg, system 100 of FIG. 1 or system 100 of FIG. 2) to generate a neural network output 408.

[0109] 4, the network input 402 may be an image. The prediction task performed by a neural network included in the neural network system may be an auto-encoding task, e.g., reconstruction of the original image 402. In this case, each query embedding in the set of query embeddings may define a respective spatial location (e.g., x and y coordinates) of a corresponding pixel in the image 402. The image 402 may have a size of 100×100 pixels, and the set of query embeddings may include 10,000 query vectors, e.g., a vector for each pixel in the image 402.

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

[0111] 5 shows another example of a query embedding 500. In general, the query embedding can be constructed using output-specific features to generate network outputs with different semantics.

[0112] For example, if the network output is a sequence of words, the query embedding may include positional embeddings, each positional embedding specifying the position of a respective word in the sequence of words. As a particular example, the positional embedding may be based on the index of the word in the sequence of words. In another example, if the network output is a two-dimensional array of pixels, the positional embedding may be based on the xy coordinates of a pixel in the array of pixels. In yet another example, if the network output is a point cloud, the positional embedding may be based on the xyz coordinates of a point in the point cloud. In some implementations, the positional embedding may be a Fourier feature positional encoding, e.g., with frequency bands spaced log-linearly across a predefined target frequency range. In some implementations, both "raw" coordinates, e.g., x, y, or z, and Fourier features may be included.

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

[0114] For multi-task or multi-modal settings, embeddings can be used for each task or modality, respectively. For classification tasks, 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 embeddings simultaneously with the neural network, for example, by backpropagating the gradient of the loss through the neural network to the embeddings. For tasks with heterogeneous outputs, such as multi-modal auto-encoding, features specific to some queries (e.g., spatial locations represented by x and y coordinates) can be combined with the modality embeddings. For multi-modal auto-encoding, each query embedding can specify the modality for which it is desired to reconstruct.

[0115] FIG. 6 illustrates an example modality 600 of inputs that may be processed by a neural network system (e.g., system 100 of FIG. 1 or system 100 of FIG. 2). The system can be used in domains with a variety of input / output spaces, including multitask language understanding, dense vision tasks such as optical flow, and hybrid dense / sparse multimodal tasks such as video+audio+class auto-encoding. In some implementations, the inputs can be pre-processed and / or post-processed, e.g., to reduce the size of very large inputs and / or outputs, respectively. The last two columns indicate the dimensionality of the inputs and the dimensionality of the respective outputs.

[0116] 7A and 7B show an example multimodal result 700, e.g., automatic encoding of audio-video-labels. The input is shown on the left and the reconstruction is shown on the right. 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] The term "configured" is used herein in connection with systems and computer program components. A system of one or more computers configured to perform a particular operation or action means that the system has installed thereon software, firmware, hardware, or a combination thereof that causes the system to perform the operation or action during operation. A computer program or programs configured to perform a particular operation or action means that the program or programs contain instructions that, when executed by a data processing device, cause the device to perform the operation or action.

[0118] The embodiments and functional operations of the subject matter described herein can be implemented in digital electronic circuitry, tangibly embodied computer software or firmware, computer hardware including the structures disclosed herein and their structural equivalents, or one or more combinations thereof. The embodiments of the subject matter described herein can be implemented as one or more computer programs, i.e., one or more modules of computer program instructions encoded on a tangible non-transitory storage medium for execution by or control of 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 thereof. 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 a suitable receiver device for execution by the data processing apparatus.

[0119] The term "data processing apparatus" refers to data processing hardware and encompasses 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 may be or may even include special purpose logic circuitry, such as, for example, an FPGA (field programmable gate array) or an ASIC (application specific integrated circuit). In addition to hardware, the apparatus may optionally include code that creates an environment for the execution of computer programs, such as code constituting a processor firmware, a protocol stack, a database management system, an operating system, or any combination of one or more of these.

[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, may be written in any form of programming language, including compiled or interpreted languages, or declarative or procedural languages, and may 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. A program may correspond to a file in a file system, but does not have to. A program may be stored in part of a file that holds other programs or data, e.g., one or more scripts stored in a markup language document, a single file dedicated to the program in question, or in multiple coordinated files, e.g., files that store one or more modules, subprograms, or portions of code. A computer program may be deployed to run on one computer, or on multiple computers located at one site or distributed across multiple sites and interconnected by a data communications network.

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

[0122] The processes and logic flows described herein may be performed by one or more programmable computers executing one or more computer programs by operating on input data to generate output to perform a function. The processes and logic flows may also be performed by special purpose logic circuitry, such as, for example, an FPGA or ASIC, or a combination of special purpose logic circuitry and one or more programmed computers.

[0123] A computer suitable for executing a computer program may be based on a general-purpose or dedicated microprocessor or both, or on other types of central processing units. Typically, the central processing unit receives instructions and data from a read-only memory, a random access memory, or both. The essential elements of a computer are a central processing unit for executing or executing instructions, and one or more memory devices for storing instructions and data. The central processing unit and memory may be supplemented by or incorporated in special-purpose logic circuitry. Typically, a computer also includes one or more mass storage devices for storing data, e.g., magnetic, magneto-optical, or optical disks, or is operatively coupled to receive data from or transfer data to them, or both. However, a computer need not be equipped with such devices. Furthermore, a computer may be embedded in another device, such as a mobile phone, a personal digital assistant (PDA), a mobile audio or video player, a game console, a global positioning system (GPS) receiver, a portable storage device, e.g., 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, all types of non-volatile memory, media, and memory devices, such as semiconductor memory devices, e.g., EPROMs, EEPROMs, flash memory devices, magnetic disks, e.g., internal hard disks or removable disks, magneto-optical disks, and CD-ROM and DVD-ROM disks.

[0125] To provide for user interaction, embodiments of the subject matter described herein 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, as well as a keyboard and a pointing device, such as a mouse or trackball, through which the user can provide input to the computer. Other types of devices can also be used to provide for user interaction, e.g., feedback provided to the user can be any form of sensory feedback, e.g., visual feedback, auditory feedback, tactile feedback, etc., and input from the user can be received in any form, such as acoustic, speech, or tactile input. Additionally, the computer can interact with the user by sending and receiving documents to and from a device used by the user, e.g., by sending a web page to a web browser on the user's device in response to a request received from the web browser. The computer can also interact with the user by sending text messages or other forms of messages to a personal device, e.g., a smartphone running a messaging application, and receiving a response message from the user in return.

[0126] A data processing apparatus for implementing machine learning models may also include dedicated hardware accelerator units, for example for handling typical, computationally intensive parts of machine learning training or production, i.e., inference, workloads.

[0127] Machine learning models can be implemented and deployed using machine learning frameworks, such as the TensorFlow framework.

[0128] An embodiment of the subject matter described herein can be implemented in a computing system that includes a back-end component, e.g., as a data server, or a middleware component, e.g., an application server, or a front-end component, e.g., a client computer with a graphical user interface, a web browser, or an app through which a user can interact with an implementation of the subject matter described herein, 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 local area networks (LANs) and wide area networks (WANs), e.g., the Internet.

[0129] A computing system may include clients and servers. Clients and servers are typically remote from each other and typically interact through a communications 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, the server sends data, e.g., HTML pages, to a user device, e.g., for the purpose of displaying the data to a user interacting with the device functioning as a client and receiving user input from the user. Data generated at the user device, e.g., results of user interaction, can be received at the server from the device.

[0130] Although the present specification contains many specific implementation details, these should not be construed as limiting the scope of the invention or the scope of the claims that may be claimed, but should be construed as descriptions of features that may be specific to certain embodiments of a particular invention. Certain 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 in multiple embodiments separately or in any suitable subcombination. Furthermore, although features may be described above as functioning in a particular combination and may initially be claimed as such, in some cases one or more features from the claimed combination may be deleted from that combination, and the claimed combination may be directed to a subcombination or a variation of the subcombination.

[0131] Similarly, although operations are illustrated in the figures and claimed in a particular order, this should not be understood as requiring that such operations be performed in the particular order or sequence illustrated, or that all of the illustrated operations be performed, to achieve desirable results. In certain situations, multitasking and parallel processing may be advantageous. Furthermore, the separation of various system modules and components in the above-described embodiments 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 may be performed in a different order and still achieve desirable results. By way of example, the processes depicted in the accompanying figures do not necessarily require the particular order or sequence shown to achieve desirable results. In some cases, multitasking and parallel processing may be advantageous. [Explanation of symbols]

[0133] 100 Neural Network System 102 Neural Network Input 104 Potential embedding 106 Query Embedding 108 Network Output 110 Encoder Block 120 Output Block 150 Neural Networks 160 Neural Networks 202 Network Input 202 2D Input Array 202 Data Element Embedding 204 Potential embedding 206 Query Embedding 208 Network Output 210 Encoder Block 230 Cross Attention Block 240 Self-Attention Block 250 Neural Networks 250 Output Blocks 300 processes 400 Query Embed 402 Network Input 402 images 404 Potential embedding 408 Neural Network Output 410 Encoder 420 Output Block 500 Query Embedding 600 Modalities 700 Exemplary Multimodal Results

Claims

1. A method performed by one or more data processing apparatus, comprising: obtaining (i) network inputs 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 inputs; processing the network inputs using the neural network to generate network outputs having respective dimensions corresponding to each query embedding in the set of query embeddings; processing the network inputs using an encoder block of the neural network to generate a representation of the network inputs as a set of latent embeddings; 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 a corresponding query embedding to the set of potential embeddings; The steps include: The method comprises: initializing the set of potential embeddings; the network input comprises a set of data element embeddings; The encoder block of the neural network one or more cross-attention blocks each performing an operation comprising: updating each potential embedding in the set of potential embeddings using attention to some or all of the data element embeddings in the set of data element embeddings; The encoder block of the neural network 11. A method comprising: one or more self-attention blocks, each performing an operation comprising updating each potential embedding in the set of potential embeddings using attention for the set of potential embeddings.

2. A method performed by one or more data processing devices, comprising: obtaining (i) network inputs 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 inputs; processing the network inputs using the neural network to generate network outputs having respective dimensions corresponding to each query embedding in the set of query embeddings; processing the network inputs using an encoder block of the neural network to generate a representation of the network inputs as a set of latent embeddings; 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 a corresponding query embedding to the set of potential embeddings; The steps include: The method comprises: For each of one or more of the query embeddings, the query embedding defines respective spatial locations within the network input; wherein the dimensions of the network output corresponding to the query embedding define respective predictions associated with the spatial locations within the network input.

3. A method performed by one or more data processing devices, comprising: obtaining (i) network inputs 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 inputs; processing the network inputs using the neural network to generate network outputs having respective dimensions corresponding to each query embedding in the set of query embeddings; processing the network inputs using an encoder block of the neural network to generate a representation of the network inputs as a set of latent embeddings; 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 a corresponding query embedding to the set of potential embeddings; The steps include: The method comprises: For each of one or more of the query embeddings, the query embedding comprises one or more input features from each spatial location in the network input; wherein the dimensions of the network output corresponding to the query embedding define respective predictions associated with the spatial locations within the network input.

4. A method performed by one or more data processing devices, comprising: obtaining (i) network inputs 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 inputs; processing the network inputs using the neural network to generate network outputs having respective dimensions corresponding to each query embedding in the set of query embeddings; processing the network inputs using an encoder block of the neural network to generate a representation of the network inputs as a set of latent embeddings; 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 a corresponding query embedding to the set of potential embeddings; The steps include: The method comprises: For each of the one or more query embeddings, The method, wherein the query embedding defines a modality of the network input.

5. A method performed by one or more data processing devices, comprising: obtaining (i) network inputs 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 inputs; processing the network inputs using the neural network to generate network outputs having respective dimensions corresponding to each query embedding in the set of query embeddings; processing the network inputs using an encoder block of the neural network to generate a representation of the network inputs as a set of latent embeddings; 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 a corresponding query embedding to the set of potential embeddings; The steps include: The method comprises: the set of query embeddings collectively defines a plurality of prediction tasks, and for each of one or more of the query embeddings: the query embedding specifies a respective prediction task from the plurality of prediction tasks; wherein the dimensions of the network output corresponding to the query embedding define a predicted output for the prediction task specified by the query embedding.

6. 6. The method of claim 5, wherein the plurality of prediction tasks comprises one or more of a classification task, a regression task, a segmentation task, or an auto-encoding task.

7. generating the dimensionality of the network output by cross-attention of corresponding query embeddings to the set of potential embeddings; generating respective attention weights for each potential embedding based on (i) the query embedding and (ii) the potential embeddings; generating the dimensionality of the network output corresponding to the query embedding based on the attention weights for the latent embeddings; The method of claim 1 , comprising:

8. generating the dimensionality of the network output corresponding to the query embedding based on the attention weights for the latent embeddings; processing each potential embedding to generate a value embedding for said potential embedding; combining the value embeddings using the attention weights; generating the dimensionality 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 of claim 7 comprising:

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

10. The method of claim 1 , wherein a number of query embeddings in the set of query embeddings is greater than a number of potential embeddings in the set of potential embeddings.

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

12. 2. The method of claim 1 , wherein a number of potential embeddings in the set of potential embeddings is predefined, and a number of query embeddings in the set of query embeddings is variable and independent of the number of potential embeddings in the set of potential embeddings.

13. A method performed by one or more data processing devices, comprising: obtaining (i) network inputs 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 inputs; processing the network inputs using the neural network to generate network outputs having respective dimensions corresponding to each query embedding in the set of query embeddings; processing the network inputs using an encoder block of the neural network to generate a representation of the network inputs as a set of latent embeddings; 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 a corresponding query embedding to the set of potential embeddings; The steps include: The method comprises: The method, wherein the network input comprises multi-modal data.

14. The method of claim 1 , wherein a number of potential embeddings in the set of potential embeddings is less than a number of data element embeddings in the set of data element embeddings.

15. 2. The method of claim 1 , wherein each data element embedding corresponds to a respective spatial location within the network input and comprises a feature embedding based on features of the network input at that spatial location.

16. 2. The method of claim 1, wherein each data element embedding corresponds to a respective spatial location within the network input and comprises a positional embedding that characterizes the spatial location.

17. The method of claim 1 , wherein each data element embedding comprises a modality embedding that defines a modality corresponding to the data element embedding.

18. One or more computers; one or more storage devices communicatively coupled to the one or more computers; 18. A system comprising: 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 method of any one of claims 1 to 17.

19. 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 method of any one of claims 1 to 17.

Images