generating a neural network output by querying cross-attention embedded on a set of latent embeddings

Abstract

Methods, systems, and apparatus, including computer programs encoded on computer storage media, for generating a network output using a neural network. In one aspect, a method includes obtaining: (i) a network input to a neural network, and (ii) a set of query embeddings; processing the network input using the neural network to generate a network output comprising a respective dimension 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 latent embeddings; and, processing: (i) the set of latent 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 over the set of latent embeddings.

Description

[0001] Cross-reference to related applications

[0002] This application claims priority to U.S. Provisional Patent Application Serial No. 63 / 194,874, filed May 28, 2021, entitled “Generating Neural Network Outputs by Cross Attention of Qury Embeddings Over a Set of Latentembeddings”, the entire contents of which are incorporated herein by reference. Technical Field

[0003] This manual relates to processing data using machine learning models. Background Technology

[0004] Machine learning models receive input and generate outputs, such as predicted outputs, based on the received input. Some machine learning models are parametric models, generating outputs based on the received input and the values ​​of the model parameters.

[0005] Some machine learning models are deep models, which employ multiple layers to generate outputs from received inputs. For example, a deep neural network is a type of deep machine learning model that includes an output layer and one or more hidden layers, each of which applies a nonlinear transformation to the received input to generate an output. Summary of the Invention

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

[0007] In one aspect, a method performed by one or more data processing devices is described. The method includes obtaining: (i) network inputs for a neural network, and (ii) a set of query embeddings, which collectively define a prediction task to be performed by the neural network by processing the network inputs. The method uses the neural network to process the network inputs to generate a network output comprising a corresponding dimension corresponding to each query embedding in the query embedding set. The method uses an encoder block of the neural network to process the network inputs to generate a representation of the network inputs as a latent embedding set. The method uses a cross-attention block to process (i) the latent embedding set and (ii) the query embedding set, the cross-attention block generating each dimension of the network output through cross-attention of the corresponding query embeddings on the latent embedding set.

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

[0009] Furthermore, throughout this specification, a neural network “block” refers to a set of one or more neural network layers.

[0010] According to one aspect, a method is provided performed by one or more data processing devices, the method comprising obtaining: (i) network inputs to a neural network and (ii) a set of query embeddings, which 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 including a corresponding dimension corresponding to each query embedding in the set of query embeddings, comprising: processing the network inputs using an encoder block of the neural network to generate a representation of the network inputs as a latent embedding set; and processing (i) the latent embedding set and (ii) the query embedding set using a cross-attention block, the cross-attention block generating each dimension of the network output by cross-attention of the corresponding query embeddings on the latent embedding set.

[0011] In some implementations, for each of one or more query embeddings: the query embedding defines the corresponding spatial location in the network input; the dimension of the network output corresponding to the query embedding defines the corresponding prediction related to the spatial location in the network input.

[0012] In some implementations, for each of one or more query embeddings: the query embedding includes one or more input features from the corresponding spatial location in the network input; the dimension of the network output corresponding to the query embedding defines the corresponding prediction related to the spatial location in the network input.

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

[0014] In some implementations, the set of query embeddings collectively defines multiple prediction tasks, and wherein, for each of one or more query embeddings: the query embedding specifies the corresponding prediction task from the multiple 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.

[0015] In some implementations, multiple prediction tasks include one or more of the following: classification tasks, regression tasks, segmentation tasks, or autoencoding tasks.

[0016] In some implementations, generating the dimensions of the network output through cross-attention of the corresponding query embedding on the set of potential embeddings includes: generating the corresponding attention weights for each potential embedding based on: (i) the query embedding and (ii) the potential embedding; and generating the dimensions of the network output corresponding to the query embedding based on the attention weights of the potential embeddings.

[0017] In some implementations, generating the dimension of the network output corresponding to the query embedding based on the attention weights of the latent embedding includes: processing each latent embedding to generate a value embedding of the latent embedding; combining the value embeddings using 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 attention weights.

[0018] In some implementations, this cross-attention is query-key-value attention.

[0019] In some implementations, the number of query embeddings in the query embedding set is greater than the number of potential embeddings in the potential embedding set.

[0020] In some implementations, the number of query embeddings in the query embedding set is at least a factor of 2 greater than the number of potential embeddings in the potential embedding set.

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

[0022] In some implementations, the network input includes multimodal data.

[0023] In some implementations, the method further includes initializing the latent embedding set, wherein the network input includes a set of data element embeddings, wherein the encoder block of the neural network includes one or more cross-attention blocks, each cross-attention block performing an operation including updating each latent embedding in the latent embedding set using attention to some or all of the data element embeddings in the data element embedding set; wherein the encoder block of the neural network includes one or more self-attention blocks, each self-attention block performing an operation including updating each latent embedding in the latent embedding set using attention to the latent embedding set.

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

[0025] In some implementations, each data element embedding corresponds to a specific spatial location in the network input and includes a feature embedding based on features of the network input at that spatial location.

[0026] In some implementations, each data element embedding corresponds to a specific spatial location in the network input and includes a location embedding that represents that spatial location.

[0027] In some implementations, each data element embedding includes a modal embedding that defines the modality corresponding to the data element embedding.

[0028] According to another aspect, a system is provided, comprising: 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 methods described herein.

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

[0030] The subject matter described in this specification can be implemented in specific embodiments to achieve one or more of the following advantages.

[0031] The neural network system described in this specification can process network inputs to generate representations of the network inputs as a set of latent embeddings, and then generate each dimension of the network output by cross-attention corresponding to query embeddings on that set of latent embeddings. Generating each dimension of the network output by cross-attention corresponding to query embeddings on the set of latent embeddings allows the system to flexibly modify the dimensions of the network output, for example, by modifying the number of query embeddings, without reconstructing the hidden layers of the neural network.

[0032] Furthermore, the prediction task performed by a neural network is defined by the query embedding and can be modified using the query embedding. Therefore, neural network systems offer a flexible and general neural network architecture that can be implemented to perform various prediction tasks with minimal modifications to network inputs from various modalities. Consequently, neural network systems can achieve more efficient resource utilization by reducing or eliminating the need to design, implement, and train new neural networks for each new domain or prediction task.

[0033] Compared to conventional neural networks that generate fixed-size outputs (such as fixed-size image segmentation outputs), the neural network systems described in this specification can use query embeddings to generate sparse or partial network outputs. For example, a neural network system can generate segmentation of only a specified proper subset of an image, rather than generating a complete segmentation of the image, by querying a latent embedding set using only relevant query embeddings. Therefore, neural network systems can reduce the consumption of computational resources (e.g., memory and computing power) by providing the option to generate sparse or partial network outputs where appropriate. For example, a neural network system can be used to segment only the region of interest of a large image while avoiding segmentation of the rest of the image.

[0034] More generally, the size of the network output can be defined without considering the size of the network input, since the dimension of the network output corresponds to the number of query embeddings in the query embedding set. Furthermore, the use of cross-attention avoids a secondary dependency of the performed computation on the sizes of the network input and output; instead, the dependency is linear, which is advantageous for processing data such as video, audio, and multimodal data. Similarly, the described techniques are implemented using a set of latent embeddings whose spatial or temporal structure can be independent of the network input and output, for example, since the network output is generated by querying latent embeddings using cross-attention. This facilitates processing multimodal data, i.e., network inputs with different structures or dimensions. Moreover, some implementations of the system can perform tasks more efficiently and with simpler architectures compared to other techniques, such as optical flow determination tasks.

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

[0036] Figure 1 This is a block diagram of an example neural network system that can generate neural network outputs.

[0037] Figure 2 This is a more detailed block diagram of the example neural network system.

[0038] Figure 3 This is a flowchart illustrating an example process of generating neural network output using a neural network system.

[0039] Figure 4 The illustration shows an example query embedding.

[0040] Figure 5 Another example of query embedding is illustrated.

[0041] Figure 6 The illustration shows an example modality of input that can be processed by a neural network system.

[0042] Figure 7A The illustration shows an example of multimodal results.

[0043] Figure 7B The illustration shows an example of multimodal results.

[0044] The same reference numerals and names in the various figures indicate the same elements. Detailed Implementation

[0045] Figure 1This is a block diagram of an example 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, in which the systems, components, and techniques described below are implemented.

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

[0047] Network input 102 can include any suitable type of data. That is, neural network 150 can be configured to process inputs of any of a variety of modalities, such as images, videos, audio, text, or any other suitable modality, or any combination thereof. Neural network 150 can have a single set of parameters while being substantially flexible, enabling it to process any of a variety of modalities or combinations thereof.

[0048] In some implementations, network input 102 can characterize an entity. This entity can include images, audio waveforms, point clouds (e.g., generated by LiDAR or radar sensors), proteins, word sequences (e.g., forming one or more sentences or paragraphs), videos (e.g., representing a sequence of video frames), or any other suitable type of data or combinations thereof. As used herein, images, i.e., still or moving images (videos), can include point clouds.

[0049] As a specific example, an entity may include multiple units arranged in a spatial or temporal structure, such as... Figure 4 As shown, an entity can be an image, and each unit can be a pixel in the image. Each unit in an entity can have an associated data element embedding, which can characterize, for example, the location of the unit in a spatial or temporal structure and / or features associated with the unit in the spatial structure. In another specific example, an entity can be text; for example, an entity can be a sequence of words, phrases, characters, or word fragments in one or more languages. In yet another specific example, an entity can include a combination of different modalities; for example, an entity can include a combination of audio, video, and tag data. As another example, an entity can include sensor data from one or more sensors configured to perceive one or more characteristics of the real-world environment, such as image data, audio data, or so-called “fine touch” sensor data (which allows for localization). Although several examples of network input 102 have been described above, in general, network input 102 can have any suitable dimension and structure and can include any suitable type of data.

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

[0051] The neural network system 100 can flexibly modify the dimension of the network output 108, for example, to generate sparse or partial network outputs 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 “dimension” of the network output 108. As described in more detail below, the system 100 can control the determination of the dimension of the network output 108 using a set of query embeddings 106.

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

[0053] The set of query embeddings 106 can be provided, for example, by a user of system 100 in the same manner, or the set of query embeddings 106 can be predetermined, i.e., the system can be configured to perform a specific task. System 100 can use the query embeddings 106 to generate a network output 108 with desired dimensions. Where necessary for a specific task, each query embedding in the set of query embeddings 106 can define relevant information about the corresponding output dimension in the network output 108. In other words, each query embedding can define, for example, the location of the corresponding output dimension and / or modality-specific features, such as where the network output includes outputs with specific 1D, 2D, or 3D locations or specific data modalities. Such query embeddings can be constructed by combining (e.g., concatenating or adding) sets of vectors to provide the information required for the corresponding dimension of the network output.

[0054] Query embeddings 106 can be defined in any suitable manner; for example, they can be hand-designed, learned, or defined as functions of the network input 102. Generally, the neural network system 100 does not depend on using any particular type of query embedding, and even 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 prediction tasks. This will be described further later.

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

[0056] As a specific example, network input 102 may include representations of a first image and a second image, each depicting the same scene, and the prediction task may be to estimate the two-dimensional displacement (e.g., optical flow) of points in the first image (e.g., pixels or an appropriate subset of pixels (i.e., not all pixels)). In some implementations, such points may include pixels or image patches comprising multiple pixels. In this case, each query embedding in the set of query embeddings 106 may specify, for example, the x and y coordinates of a corresponding point in the first image, and may include one or more additional embeddings for the optical flow prediction task, such as representations of the first and second images of that point. That is, the set of query embeddings 106 may define a set of points in the first image against which a prediction is generated, wherein the prediction is an estimate of optical flow. For example, the dimensions of the network output may define the magnitude and direction of the optical flow at the point. See below for reference. Figure 4 and Figure 5 The example query embedding and prediction task are described in more detail.

[0057] After obtaining the set of network input 102 and query embeddings 106, system 100 can use neural network 150 to process input 102 and generate network output 108 with desired dimensions, for example, including the output corresponding to the respective dimension of each query embedding in the set of query embeddings 106. The process will be described in more detail below.

[0058] System 100 can use encoder block 110 of neural network 150 to process network input 102 to generate a set of latent embeddings 104 from the representation of network input 102. "Latent embedding" can refer to embeddings in a latent space. Generally, latent embeddings 104 can be included in a space of any appropriate dimension. Generally, the dimension of the embedding set can be determined by the number of embeddings in that set (e.g., ...). N The term is defined as follows: In some implementations, the encoder block processes the network input and an initial set of latent embeddings to generate that set.

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

[0060] After generating a set of latent embeddings 104 using encoder block 110, neural network system 100 can use output block 120 of neural network 150 to process the set of latent embeddings 104 and the set of query embeddings 106 to generate network output 108. Specifically, output block 120 can be configured to perform an "attention" operation. The attention operation can include updating each embedding in the first embedding set using attention to a second set of embeddings. In some implementations, after performing the attention operation (e.g., after updating each embedding), output block 120 can use one or more neural network layers (e.g., fully connected neural network layers) to process each updated embedding. Many different attention mechanisms are available; see below. Figure 2 A more detailed description of the example attention operation.

[0061] Specifically, output block 120 can update each query embedding in the set of query embeddings 106 on the set of latent 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 the corresponding pixel in the first image, and an embedding specifying the optical flow prediction task. Output block 120 can generate a specific dimension of network output 108 through cross-attention of the corresponding query embedding on the set of latent embeddings 106. Returning to the previous example as an illustration, the dimension of network output 108 can be, for example, a single flow vector representing the optical flow of a pixel in the first image as defined by the query embedding. Output block 120 can repeat this process for each query embedding in the set of query embeddings 106 to generate each corresponding dimension of network output 108. In some implementations, the output corresponding to each query embedding in the set of query embeddings 104 can be generated in parallel. See below for reference. Figure 2 Describe the process in more detail.

[0062] Encoder block 110 and output block 120 can have any suitable neural network architecture that enables them to perform their specified functions. For example, encoder block 110 and output block 120 can have any suitable number (e.g., 2, 5, or 10 layers) of suitable neural network layers (e.g., convolutional layers, fully connected layers, recurrent layers, attention layers, etc.) connected in any suitable configuration (e.g., as a linear sequence of layers). Neural network system 100 may also additionally include any number of neural network blocks configured to perform any suitable operation. Reference is made below. Figure 2 A specific example of encoder block 110 and output block 120 is described in more detail.

[0063] In this way, system 100 can flexibly modify the dimension of network output 108, for example, by changing the number of query embeddings 106, without refactoring the hidden layers of neural network 150. Furthermore, as described above, the prediction task performed by neural network 150 is defined by query embeddings 106 and can be modified using query embeddings 106 in some implementations. Therefore, neural network system 100 can provide a flexible and general neural network architecture that can be implemented to perform various prediction tasks on network inputs of various modalities with minimal modification. Neural network system 100 may also include a training engine that can train 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 that should be generated by neural network 160 by processing the training input.

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

[0065] The training engine can adjust the model parameter values ​​of encoder block 110 and output block 120 to optimize an objective function that measures (i) the network output generated by neural network 150 and (ii) the similarity between targets specified in 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 output generated by neural network 150 may include only a portion of the total set of possible outputs. In this case, the objective function can be evaluated by referring only to the generated outputs (e.g., rather than the total set of outputs). This can accelerate the training of the neural network and improve training efficiency.

[0066] The training engine can, for example, use backpropagation to determine the gradient of the objective function. The training engine can use the gradient, for example, any suitable gradient descent optimization algorithm (e.g., Adam), to update the model parameter values ​​of encoder block 110 and output block 120. The training engine can determine performance measures of neural network 150 on a set of validation data that was not used during the training of neural network 150. After training, neural network system 100 can be used to perform machine learning tasks, such as processing network inputs and generating network outputs.

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

[0068] The neural network system 100 can be configured to perform any appropriate machine learning task. Below are some examples.

[0069] In some implementations, system 100 uses neural network 150 to perform image or audio segmentation tasks. For example, the neural network can process network input including an image or a series of audio samples representing an audio waveform to generate a network output that defines a corresponding score distribution across a set of possible categories for each pixel in the input image or each audio sample. The score of a category defines the likelihood that a corresponding pixel or sample is included in that category. For example, for an image, possible categories could include, for example, water, buildings, vehicles, pedestrians, etc. In this example, a query embedding can specify the location of a corresponding pixel in the input image or the location of a 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 audio sample. The image can be a moving image, i.e., a video. Similarly, system 100 can use neural network 150 to perform an image depth prediction task, where the network output defines predicted depth values ​​for corresponding pixels to obtain a (spatial 3D) depth map of the image.

[0070] In some implementations, system 100 uses neural network 150 to perform protein modeling tasks, where the network input represents a protein (e.g., multiple sequence alignment of a protein), and the network output represents a predicted structure of the protein (e.g., the corresponding 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 amino acid index in the amino acid sequence of the protein, and optionally specify the type of amino acid at the specified index (e.g., alanine, arginine, asparagine, etc.). The dimension 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.

[0071] In some implementations, system 100 uses neural network 150 to perform agent control tasks, wherein the network input represents a series of one or more observations or other data characterizing the state of the environment, and the output defines the action to be performed by the agent in response to the latest data in the sequence. The environment can be a real-world or simulated environment, and the agent can be, for example, a control system of a robot, autonomous land, sea, or air vehicle, or an industrial facility. Each query embedding can specify a corresponding task to be performed by the agent in the environment; for example, the task could be that the agent navigates to a target location in the environment, and the query embedding can specify the target location.

[0072] In some implementations, system 100 uses neural network 150 to process network input representing audio samples in an audio waveform to perform speech recognition, for example, to generate network output representing a sequence of phonemes, epigraphs, characters, or words corresponding to the audio waveform. The audio samples define a sequence of samples, and optionally, the network input may also include encodings of the positions of samples within the sequence. For example, each dimension of the network output may correspond to a corresponding time interval in the audio waveform, and a corresponding score distribution may be defined on a set of possible phonemes, epigraphs, characters, or words. In this example, the neural network may generate the network output, for example, by defining a query embedding that includes: (i) the corresponding time interval in the audio waveform, such as a position in the sequence, and optionally (ii) a prediction task of the time interval in the audio waveform, such as decoding the phonemes, epigraphs, characters, or words corresponding to the time interval.

[0073] In some implementations, system 100 uses neural network 150 to perform a reconstruction task, such as an autoencoding task, by processing the network input to generate a predicted reconstruction of the network output that defines the network input. The network input can be, for example, an image or video, an audio waveform, a point cloud, or a text sequence. In this example, each query embedding can include data specifying the corresponding location in the network input to be reconstructed (e.g., a spatial location in an image, a spatial and / or 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 being reconstructed, for example, defining whether the network input is an image, an audio waveform, or a point cloud. For example, 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 can be reconstructed. Points can be reconstructed serially and / or in parallel. For example, the query embedding set can include the query embedding for each point in the network input, such as pixels of an image or samples of an audio signal, or it can include the query embedding for only one point and can serially reconstruct the network output one point at a time by specifying each desired point; or the query embedding set can reconstruct some but not all points in the network input.

[0074] In some implementations, system 100 uses neural network 150 to perform neural machine translation tasks, such as processing network input representing a sequence of texts (e.g., sequences of words, phrases, characters, or word fragments) in one language to generate a network output that is a translation of the text sequence into another language, i.e., a sequence of texts in another language that serves as the translation of the input text sequence. 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. Furthermore, each query embedding may specify a corresponding position in the output text sequence.

[0075] In some implementations, system 100 uses neural network 150 to perform audio or audiovisual processing tasks. For example, if the network input represents spoken utterance, the network output generated by the neural network could be a score for each of a set of text segments, each score representing an estimated probability that the text segment is a correct transcription of the utterance. In this example, each query embedding could be a representation of or correspond to the corresponding text segment. As another example, if the network input represents spoken utterance, the output generated by the neural network could indicate whether a specific word or phrase (“hot word”) was spoken in the utterance. In this example, each query embedding could be or correspond to, for example, a one-hot embedding corresponding to the corresponding word or phrase.

[0076] In some implementations, system 100 uses neural network 150 to perform text-to-speech tasks, wherein the network input represents natural language text or features of natural language text, and the network output is a spectrogram, waveform, or other data defining the audio of the text spoken in natural language. In this example, each query embedding can be used to identify the corresponding spoken voice to be used when uttering the text (e.g., where each spoken voice corresponds to the voice of the corresponding person). Furthermore, each query embedding can specify a corresponding time window in the network output.

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

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

[0079] In some implementations, system 100 uses neural network 150 to perform genomics tasks, where, for example, unsupervised learning techniques are used on a dataset of DNA sequence fragments or the output of a downstream task. The network input represents fragments of DNA sequences or other molecular sequences, and the network output is an embedding of the fragment for the downstream task. Examples of downstream tasks include promoter site prediction, methylation analysis, prediction of functional effects of non-coding variants, etc. In this example, each query embedding can specify a corresponding position in the DNA sequence.

[0080] In some implementations, system 100 uses neural network 150 to perform image, video, or audio classification tasks. The network input can then represent samples of pixels from an image or video, or samples of audio waveforms, as previously described. Optionally, the network input can include location encodings of the pixels from an image or video, or samples of audio; for video, the location encoding can be spatial and / or temporal location. The network output defines the classification of the image, video, or audio. The classification can include a corresponding score for each of a set of possible object categories (e.g., for images, vehicles, pedestrians, cyclists, etc.). The score of the object category can define the probability that the network input includes objects belonging to that object category (e.g., an image or video depicting that object). Video classification can include, for example, a classification of actions depicted in the video for gesture recognition. The query embedding set can include a single query embedding (which can be learned). In a similar manner, system 100 can use neural network 150 to perform image or audio bounding box tasks, where the network output includes a vector of 1, 2, or 3-dimensional coordinates (with temporal count as the dimension) classifying and defining the bounding boxes of objects.

[0081] As previously mentioned, references to images or videos in the above examples typically include point clouds. For example, in some implementations, system 100 uses neural network 150 to perform point cloud processing tasks, where the network input represents a point cloud (e.g., generated by a LiDAR or radar sensor) and the network output characterizes, for example, 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.

[0082] In some implementations, system 100 uses neural network 150 to perform image or video captioning tasks. For example, the neural network can process network input including an image or a video comprising a series of images to generate a network output that is one or more defined terms describing the image or video, where terms can represent words, parts of words, or sentences. Terms can be defined deterministically or stochastically, for example, by sampling from one or more defined distributions. In the case of video, the query embedding can (but does not need to) specify the temporal position of each image series. As used herein, temporal position includes temporal order.

[0083] In some implementations, the network input is a multimodal input and system 100 uses neural network 150 to perform a multimodal task. Generally, such a multimodal input is a combination of two or more different types or modalities of data, wherein the different types of data may, but do not necessarily, represent the same or overlapping objects in the network input using different modalities. Such multimodal data can include audiovisual data, including combinations of pixels of an image or combinations of video and audio data representing samples of audio waveforms. As another example, such multimodal data can include combinations of text data, such as lexical units, text representing natural language, and samples of pixels or audio waveforms of an image or video. The multimodal task can be any of the above tasks, but with additional data modalities included in the network input. As some examples, multimodal tasks can include 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 can be provided with a modality-specific embedding; then, the multimodal data can be serialized into a 2D input array 202. Generally, query embeddings can be as described previously, but in some implementations, modality-specific embeddings are added to generate outputs for each modality. Embeddings can be learned embeddings. This can be done as described previously when spatial or temporal locations are incorporated into query embeddings, for example, for video or audio data, because modality-specific embeddings enable the system to select each modality. Network outputs can correspond to those previously described, but the ones applied are provided for multimodal inputs. For example, a multimodal classification task can generate a network output that classifies the multimodal input into one or more of several categories, for example, by defining scores for each of the multiple possible categories of the input. Similarly, a captioning task generates a network output for terms describing multimodal inputs; the corresponding query may, but does not necessarily, include the spatial or temporal location of each query embedding; and so on. As another example, where the multimodal data includes a combination of text data and image or video or audio data, the task may include processing the combination to provide a network output that defines whether an image or video or audio waveform is described by text (e.g., by a specific caption), for example, by defining scores for the text or caption. The query can be as described previously for the classification; or it can include location data to generate a corresponding score for each location represented by one of the query embeddings in the query embedding set.

[0084] Generally, in the above example, the network input may include: raw data, such as raw pixel data or audio samples; or feature embeddings, such as defining spatial, temporal, or spatiotemporal features obtained through preprocessing by a feature encoder neural network.

[0085] The following is for reference. Figure 2 A more detailed description of the example architecture of the neural network system 100.

[0086] Figure 2 This is a block diagram of a more detailed example architecture of a 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, in which the following systems, components, and techniques are implemented.

[0087] The neural network system 100 can be configured to receive network input 202, which may include a set of data element embeddings representing an entity (e.g., an "input array"). The entity may, for example, include multiple units arranged in a spatial and / or temporal structure.

[0088] In some implementations, each data element embedding may correspond to a specific spatial location in the network input and may include a feature embedding based on 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 can obtain the feature embedding by selecting small patches of the image surrounding the pixel and concatenating the pixels within those patches into a vector. As used herein, references to spatial locations should be understood to typically include locations within sequences such as word sequences or audio waveform sample sequences.

[0089] In some implementations, each data element embedding may include a positional embedding that represents the 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 word sequence. In some implementations, each data element embedding may include a modal embedding that defines the modality corresponding to the data element embedding. For example, if the network input is an audio waveform, the data element embedding may include a modal embedding that specifies the network input as an audio waveform.

[0090] Generally speaking, the dimensions of network input can be MxC ,in, M It is the number of embedded data elements. C It is the number of channels embedded in each data element.

[0091] The neural network system 100 can also obtain a set of query embeddings 206 (e.g., an "output query array") by processing the network input 202, which collectively define the prediction task to be performed by the neural network 250. The query embeddings 206 can have any suitable dimension; for example, the query embeddings may include... O Query vectors.

[0092] In some implementations, the neural network system 100 may also obtain an initial set of latent embeddings 204 (e.g., a "latent array"). The latent embeddings 204 may be predefined and / or randomly initialized, or they may be learned; that is, they may include a set of learned parameters, each of which defines an element of the learned array, such as weights. In other words, 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 the latent embedding set. The set of latent embeddings 204 may have dimensions. NxD ,in, N It is the number of potential embeddings, and D This refers to the number of channels per potential embedding; both of these can be hyperparameters of the neural network system. In some implementations, the number of query embeddings is... O It can be greater than the number of potential embeddings. N For example, at least twice as large. Furthermore, the number of query embeddings... O It can be variable and independent of the number of potential embeddings. N .

[0093] The neural network system 100 can use the neural network 250 to process the network input 202 to generate a network output 208 (e.g., an "output array"), which includes a corresponding dimension corresponding to each query embedding in the set of query embeddings 206. For example, the network output 208 can be generated by... OxE Represented by an array, where... O It is the number of output vectors and E This is the number of output elements in each output vector (hyperparameter). The neural network system 100 can generate a corresponding dimension (e.g., a corresponding output vector) for 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.

[0094] The neural network system 100 can generate a network output 208 by using an encoder block 210 and an output block 250. The encoder block 210 may include a series of one or more neural network blocks, such as (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 set of latent embeddings 204 from the representation of the network input 202. The output block 250 can be configured to process the set of latent embeddings 204 generated by the encoder block 210 and a set of query embeddings 206 to generate the network output 208. The attention blocks included in the neural network 250 are described in more detail below.

[0095] Cross-attention block 230, self-attention block 240, and output block 250 can all be configured to perform attention operations, such as using attention to a second set of embeddings to update each embedding in a first set of embeddings.

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

[0097] In some implementations, attention blocks can perform query-key-value pairs. QKV Attention operations, for example, using queries ( Q ),key( K ) and value ( V The embedding uses attention to a second set of embeddings to update each embedding in the first set of embeddings. Specifically, each attention block may 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 may be configured to process the target embedding in the first set of embeddings to generate a corresponding query embedding for that target embedding. Q The key network can be configured to process each embedding in the second embedding set to generate a corresponding key embedding for each embedding in the second embedding set. K Similarly, the value subnetwork can be configured to process each embedding in the second embedding set to generate a corresponding value embedding for each embedding in the second embedding set. V ).

[0098] Then, each attention block can use query embeddings ( Q ), key embedding ( K ) and value embedding ( V The attention of each unit on the second embedding set is used to update each target embedding in the first embedding set. Specifically, each attention block can generate attention weights for each embedding in the second embedding set, for example, as a query embedding ( Q ) and each key embedding ( K The inner (e.g., dot) product of the attention blocks. Based on the second set of embeddings and attention weights, each attention block can generate a combined embedding, for example, as a value embedding weighted by its corresponding attention weights. VA linear combination of ). Finally, each attention block can be used with combination embeddings, for example, by using value embeddings ( V The target embedding in the first embedding set is updated by weighting and replacing the target embedding in the first embedding set.

[0099] In some implementations, the first and second embedding sets can be the same set of embeddings. In this case, the attention operation (e.g., a QKV attention operation) can be referred to as a "self-attention" operation. The self-attention operation can be performed by, for example, a self-attention block 240. For example, the first embedding set can be a set of latent embeddings 204, the second embedding set can also be a set of latent embeddings 204, and the self-attention block 240 can use self-attention on the set of latent embeddings 204 to update each latent embedding in the set of latent embeddings 204. In some implementations, the self-attention block 240 can use self-attention on the set of latent embeddings 204 to repeatedly update each latent embedding in the set of latent embeddings 204.

[0100] In some implementations, the first and second embedding sets can be different embedding sets. In this case, attention operations (e.g., QKV Attention operations can be referred to as "cross-attention" operations. Cross-attention operations can be performed by, for example, cross-attention blocks 230 and output blocks 250 included in encoder 210.

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

[0102] As another specific example, in the case of output block 250, the first set of embeddings may be a set of query embeddings 206, and the second set of embeddings may be a set of latent embeddings 204 generated by self-attention block 240. Output block 250 can use cross-attention on the set of latent embeddings 204 to update each query embedding in the set of query embeddings 206. In other words, output block 250 can generate each dimension of network output 208 by cross-attention of the corresponding query embeddings in the set of query embeddings 206 on the set of latent embeddings 204. After performing the cross-attention operation on each query embedding in the set of query embeddings 206 on the set of latent embeddings 204, output block 250 can generate the complete neural network output 208, which includes the corresponding dimension of each query embedding in the set of query embeddings 206.

[0103] If the neural network system 100 obtains a second set of distinct 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 input 202 and the set of second query embeddings 206 as described above to generate a network output 208 corresponding to each query embedding in the set of second distinct query embeddings 206. Therefore, by modifying the query embeddings 206, the neural network system 100 can flexibly modify the dimension of the network output 208 and / or the prediction task performed by the neural network 250.

[0104] The following describes in more detail an example process of using neural network system 100 to generate neural network output 208.

[0105] Figure 3 This is a flowchart of an example 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 in one or more locations. For example, a neural network system appropriately programmed according to this specification (e.g., Figure 1 The neural network system 100 or Figure 2 The neural network system 100 in the middle can execute process 300.

[0106] The system obtains: (i) the network input to the neural network, and (ii) a query embedding set, which together define the prediction task to be performed by the neural network by processing the network input (302). In some implementations, the network input may include multimodal data, such as video data, audio data, and labeled data. Generally, as Figure 6 As shown, network input can include any appropriate type of data or a combination thereof.

[0107] The system uses a neural network to process the network input to generate a network output that includes the corresponding dimension (304) for each query embedding in the query embedding set.

[0108] For example, a system can use encoder blocks of a neural network to process network inputs to generate a set of latent embeddings from the representations of the network inputs.

[0109] The system can then use cross-attention blocks (e.g., output block 250 above) to process: (i) the latent embedding set and (ii) the query embedding set. The cross-attention blocks (e.g., output block 250 above) generate each dimension of the network output through cross-attention (e.g., query-key-value attention) on the corresponding query embeddings within the latent embedding set. For example, as referenced above... Figure 2As described, the system can generate corresponding 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 embedding. Specifically, the system can process each latent embedding to generate a value embedding of the latent embedding. Then, the system can 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.

[0110] In some implementations, the query embedding can define the corresponding spatial location in the network input. In this case, the dimension of the network output corresponding to the query embedding can define the corresponding prediction related to the spatial location in the network input. For example, as... Figure 4 As shown, the network input can be of size 1000. 100×100 An image of pixels, and each query embedding can define the corresponding spatial location of the corresponding pixel in the image (e.g., x and y (Coordinates). For example, a prediction task could be reconstructing the original image. Query embeddings can include... 10,000 A query vector, such as the vector for each pixel in an image. The dimension of the network output corresponding to the query embedding can define the corresponding prediction related to the spatial location in the network input, for example, the reconstruction of pixels in an image defined by the query embedding.

[0111] In some implementations, the query embedding can include one or more input features from the corresponding spatial location in the network input. In this case, the dimension of the network output corresponding to the query embedding can define the corresponding prediction related to the spatial location in the network input. For example, as... Figure 5 As shown, for the optical flow prediction task, each query embedding can specify one or more input features at a corresponding spatial location in the network input, such as the spatial location of pixels in an image (e.g., ...). x and y (Coordinates) and one or more input features associated with that pixel. For example, in network input, two images can be concatenated so that the same (coordinates) of the input array are used to generate the image. M Index the corresponding pixel in the channel dimension. C (In one approach, the two images are concatenated); or in another approach, the two images are not concatenated, and the network input and optional query embeddings can also be provided with temporally encoded features. In one example, the input features of a pixel in an image can be a vector concatenated from small patches of the image surrounding that pixel. The dimension of the network output corresponding to the query embedding defining that pixel in the image can specify, for example, the prediction flow vector of that pixel in the image.

[0112] In some implementations, query embeddings can define the modality of the network input. For multimodal inputs, such as inputs including video, audio, and tags, all three modalities can be specified by each query embedding in a query embedding set. Generally, query embeddings can be represented in several different ways. As a specific example, a query embedding can be represented as a one-hot embedding.

[0113] In some implementations, a set of query embeddings can collectively define multiple prediction tasks; for example, each query embedding can define a corresponding prediction task. In this case, the dimension of the network output corresponding to the query embedding can define the prediction output of the prediction task specified by the query embedding. Prediction tasks can include, for example, classification tasks, regression tasks, segmentation tasks, autoencoding tasks, or any other suitable task. For example, such as... Figure 5 As shown, the "task_id" embedding can specify a specific prediction task. In the case of eight prediction tasks, the query embedding set can include eight query embeddings, each specifying a specific prediction task.

[0114] The following is for reference. Figure 4 and Figure 5 The example query embedding and prediction task are described in more detail.

[0115] Figure 4 The diagram illustrates what can be achieved by a neural network system (e.g., Figure 1 System 100 or Figure 2 The example query embedding 400 in system 100 is used to generate the neural network output 408.

[0116] like Figure 4 As shown, the network input 402 can be an image. The prediction task to be performed by the neural network included in the neural network system can be an autoencoding task, such as the reconstruction of the original image 402. In this case, each query embedding in the query embedding set can define the corresponding spatial location of the corresponding pixel in the image 402 (e.g., x and y (Coordinates). Image 402 can have 100×100 The size of the pixels, and the query embedding set can include 10,000 A query vector, for example, the vector for each pixel in image 402.

[0117] As referenced above Figure 1 and Figure 2As described, the neural network may include an encoder 410 and an output block 420. Encoder 410 may process image 402 to generate a representation of image 402 as a set of latent embeddings 404. Output block 420 may process the set of latent embeddings 404 and the set of query embeddings 404 to generate a network output 408 through cross-attention of the corresponding query embeddings on the set of latent embeddings 404. Network output 408 may be, for example, a reconstruction of image 402, wherein each pixel is represented by a corresponding query embedding in the set of query embeddings 400.

[0118] Figure 5 The illustration shows another example of query embedding 500. Generally, query embeddings can be constructed using output-specific features to produce network outputs with different semantics.

[0119] For example, if the network output is a sequence of words, the query embedding can include positional embeddings, where each positional embedding specifies the position of the corresponding word within the word sequence. As a specific example, the positional embedding can be based on the index of the word within the word sequence. In another example, if the network output is a two-dimensional pixel array, the positional embedding can be based on the xy coordinates of the pixels in the pixel array. In yet another example, if the network output is a point cloud, the positional embedding can be based on the xyz coordinates of points in the point cloud. In some implementations, the positional embedding can be, for example, a Fourier feature positional encoding with frequency bands that are log-linearly 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.

[0120] In some implementations, the input features of the target output can also be used for querying, either alone or in conjunction with location features. For example, if the network output is an image, the input features of a pixel in the image could be, for instance, small patches of the image surrounding that pixel, concatenated into a vector. In another example, if the output is an audio waveform, the input features of a specific time point in the audio waveform could be the amplitude of the waveform at that time point.

[0121] For multi-task or multi-modal settings, embeddings for each task or modality can be used separately. For classification tasks, a single learned embedding can be used. Specifically, see the reference above. Figure 1 As described, the training engine can train a neural network. As part of the training, the training engine can, for example, train the embedding simultaneously and concurrently using the backpropagation gradient of the loss through the neural network and into the embedding. For tasks with heterogeneous outputs, such as multimodal autoencoding, features specific to certain queries (e.g., by...) x and ySpatial location represented by coordinates can be combined with modality embeddings. In the case of multimodal autocoding, each query embedding can specify the modality that needs to be reconstructed.

[0122] Figure 6 The diagram illustrates what can be achieved by a neural network system (e.g., Figure 1 System 100 or Figure 2 The system 100 in the example shows an input modality 600 processed by the system. This system can be used in domains with multiple input and output spaces, including multi-task language understanding, dense vision tasks such as optical flow, and hybrid dense / sparse multimodal tasks such as video + audio + class autocoding. In some implementations, the input can be preprocessed and / or post-processed, for example, to reduce the size of very large inputs and / or outputs, respectively. The last two columns illustrate the dimensions of the input and the corresponding output.

[0123] Figure 7A and Figure 7B The illustration shows example multimodal results 700, such as audio-video-label autoencoding. The input is shown on the left, and the reconstruction is shown on the right. Neural network systems are able to jointly represent modalities with distinctly different properties and achieve remarkably high prediction accuracy. Figure 7B In this context, "PSNR" refers to Peak Signal-to-Noise Ratio.

[0124] This specification uses the term "configured" in conjunction with system and computer program components. For a system of one or more computers configured to perform a specific operation or action, this means that software, firmware, hardware, or a combination thereof have been installed on the system, which, in operation, causes the system to perform the operation or action. For one or more computer programs configured to perform a specific operation or action, this means that the one or more programs include instructions that, when executed by a data processing device, cause the device to perform the operation or action.

[0125] Embodiments of the subject matter and functional operation described in this specification can be implemented in digital electronic circuit systems, in tangibly implemented computer software or firmware, in computer hardware (including the structures disclosed in this specification and their equivalents), or in one or more combinations thereof. Embodiments of the subject matter described in this specification can be implemented as one or more computer programs, i.e., one or more modules of computer program instructions encoded on a tangible, non-transitory storage medium, for execution by a data processing apparatus or for controlling 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 or serial access memory device, or one or more combinations thereof. Alternatively or additionally, program instructions can be encoded on artificially generated propagation signals, such as machine-generated electrical, optical, or electromagnetic signals, which are generated to encode information for transmission to a suitable receiver device for execution by the data processing apparatus.

[0126] The term "data processing device" refers to data processing hardware and encompasses all kinds of devices, apparatuses, and machines used for processing data, including, for example, programmable processors, computers, or multiple processors or computers. The device may also be or include dedicated logic circuit systems, such as FPGAs (Field-Programmable Gate Arrays) or ASICs (Application-Specific Integrated Circuits). In addition to hardware, the device may optionally include code that creates an execution environment for computer programs, such as code constituting processor firmware, protocol stacks, database management systems, operating systems, or combinations thereof.

[0127] A computer program (also known as a program, software, software application, app, module, software module, script, or code) can be written in any form of programming language, including compiled or interpreted languages ​​or declarative or procedural languages, and can be deployed in any form, including as a standalone program or as a module, component, subroutine, or other unit suitable for use in a computing environment. A program may, but does not need to, correspond to a file in a file system. A program may be stored as a part of a file that holds other programs or data (e.g., one or more scripts stored in a markup language document), in a single file dedicated to the program in question, or in multiple co-located files (e.g., a file storing one or more modules, subroutines, or code portions). A computer program can be deployed to execute on a single computer or on multiple computers located in one location or distributed across multiple locations and interconnected by a communication network.

[0128] Similarly, in this specification, 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 will be implemented as one or more software modules or components installed on one or more computers in one or more locations. In some cases, one or more computers will be dedicated to a particular engine; in other cases, multiple engines may be installed and run on the same computer or multiple computers.

[0129] The processes and logic flows described in this specification can be executed by one or more programmable computers, which execute one or more computer programs to perform functions by manipulating input data and generating output. The processes and logic flows can also be executed by a dedicated logic circuit system (such as an FPGA or ASIC) or a combination of a dedicated logic circuit system and one or more programmable computers.

[0130] A computer suitable for executing computer programs can be based on a general-purpose or special-purpose microprocessor, or both, or any other type of central processing unit. Typically, the central processing unit receives instructions and data from read-only memory or random access memory, or both. The basic components of a computer are the central processing unit for 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 into a special-purpose logic circuitry system. Typically, a computer will also include, or be operatively coupled to, one or more mass storage devices for receiving or transferring data to, such as magnetic disks, magneto-optical disks, or optical disks. However, a computer does not require such devices. Furthermore, a computer can 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, or a portable storage device (e.g., a universal serial bus (USB) flash drive), to name just a few.

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

[0132] To provide interaction with the user, embodiments of the subject matter described in this specification can be implemented on a computer having a display device (e.g., a CRT (cathode ray tube) or LCD (liquid crystal display) monitor) and a keyboard and pointing device (e.g., a mouse or trackball) for displaying information to the user, through which the user can provide input to the computer. Other types of devices can also be used to provide interaction with the user; for example, feedback provided to the user can be any form of sensory feedback, such as visual feedback, auditory feedback, or tactile feedback; and input from the user can be received in any form, including acoustic, voice, or tactile input. Additionally, the computer can interact with the user by sending documents to and receiving documents from the device used by the user; for example, sending a webpage to a web browser on the user's device in response to a request received from a web browser. Furthermore, the computer can 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 response messages from the user in return.

[0133] The data processing apparatus used to implement machine learning models may also include, for example, dedicated hardware accelerator units for processing the common and computationally intensive parts of machine learning training or production, namely inference and workloads.

[0134] Machine learning frameworks (such as the TensorFlow framework) can be used to implement and deploy machine learning models.

[0135] Embodiments of the subject matter described in this specification can be implemented in a computing system that includes back-end components such as data servers, middleware components such as application servers, front-end components such as client computers with graphical user interfaces, web browsers, or apps, or any combination of one or more such back-end, middleware, or front-end components through which a user can interact with the implementation of the subject matter described in this specification via the graphical user interface, web browser, or app. The components of the system can be interconnected via digital data communication (e.g., a communication network) of any form or medium. Examples of communication networks include local area networks (LANs) and wide area networks (WANs) (e.g., the Internet).

[0136] A computing system may include clients and servers. Clients and servers are typically geographically separated and usually interact via a communication network. The client-server relationship is established by means of computer programs running on separate computers that have a client-server relationship with each other. In some embodiments, the server sends data (e.g., HTML pages) to a user device, for purposes such as displaying data to a user interacting with the device as a client and receiving user input from the user. Data generated at the user device, such as the result of user interaction, may be received at the server from the device.

[0137] While this specification contains numerous specific implementation details, these should not be construed as limiting the scope of any invention or the scope that may be claimed, but rather as descriptions of features specific to particular embodiments of a particular invention. Certain features described herein in the context of individual embodiments 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 individually or in any suitable sub-combination in multiple embodiments. Furthermore, although features may be described above as functioning in certain combinations and even initially claimed in this way, in some cases one or more features from the claimed combination may be removed from the combination, and the claimed combination may be for sub-combinations or variations thereof.

[0138] Similarly, although operations are depicted in the accompanying drawings and recited in the claims in a specific order, this should not be construed as requiring these operations to be performed in the specific order shown or sequentially, or to perform all of the shown operations to achieve the desired result. In some cases, multitasking and parallel processing may be advantageous. Furthermore, the separation of various system modules and components in the above embodiments should not be construed as requiring such separation in all embodiments, and it should be understood that the described program components and systems can generally be integrated into a single software product or packaged into multiple software products.

[0139] Specific embodiments of the subject matter have been described. Other embodiments are within the scope of the appended claims. For example, the actions recited in the claims can be performed in a different order and still achieve the desired result. As an example, the processes depicted in the drawings do not necessarily require the specific or sequential order shown to achieve the desired result. In some cases, multitasking and parallel processing may be advantageous.

Claims

1. A method performed by one or more data processing devices, the method comprising: Obtain: (i) the network input of the neural network, wherein the network input includes a set of data element embeddings, and (ii) a set of query embeddings that collectively define the prediction task performed by the neural network by processing the network input; The neural network is used to process the network input to generate a network output, the network output including a corresponding dimension for each query embedding in the query embedding set, including: The network input is processed using the encoder block of the neural network to generate a representation of the network input as a latent embedding set, including initializing the latent embedding set. The encoder block of the neural network includes one or more cross-attention blocks, each of which performs operations including the following: Each potential embedding in the potential embedding set is updated using attention to some or all of the data element embeddings in the data element embedding set. The encoder block of the neural network includes one or more self-attention blocks, each self-attention block performing operations including the following: The attention of the latent embedding set is used to update each latent embedding in the latent embedding set; and Cross-attention blocks are used to process (i) the latent embedding set and (ii) the query embedding set, wherein the cross-attention blocks generate each dimension of the network output by cross-attention of the corresponding query embeddings on the latent embedding set, and The network inputs include text sequences, images, video, or audio waveforms.

2. The method according to claim 1, wherein, For each of one or more of the query embeddings: The query embedding defines the corresponding spatial location in the network input; and The dimension of the network output corresponding to the query embedding defines the corresponding prediction related to the spatial location in the network input.

3. The method according to claim 1, wherein, For each of one or more of the query embeddings: The query embedding includes one or more input features from the corresponding spatial location in the network input; as well as The dimension of the network output corresponding to the query embedding defines the corresponding prediction related to the spatial location in the network input.

4. The method according to claim 1, wherein, For each of one or more query embeddings: The query embedding defines the modality of the network input.

5. The method according to claim 1, wherein, The query embedding set collectively defines multiple prediction tasks, and wherein, for each of one or more of the query embeddings: The query embedding specifies the corresponding prediction task among 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.

6. The method according to claim 5, wherein, The multiple prediction tasks include one or more of the following: classification tasks, regression tasks, segmentation tasks, and autoencoding tasks.

7. The method according to claim 1, wherein, The dimensions used to generate the network output through cross-attention embedded in the latent embedding set corresponding to the query include: Based on: (i) the query embedding and (ii) the latent embedding, a corresponding attention weight is generated for each latent embedding; and The attention weights of the latent embedding are used to generate the dimensions of the network output corresponding to the query embedding.

8. The method according to claim 7, wherein, The dimensions for generating the network output corresponding to the query embedding based on the attention weights of the latent embedding include: Process each latent embedding to generate the value embedding of the latent embedding; The value embedding is combined using the attention weights; and The dimensions of the network output corresponding to the query embedding are generated at least in part based on the result of combining the value embeddings using the attention weights.

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

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

11. The method according to claim 10, wherein, The number of query embeddings in the query embedding set is greater than the number of potential embeddings in the potential embedding set by at least a factor of 2.

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

13. The method according to claim 1, wherein, The network input includes multimodal data.

14. The method according to any one of claims 1-13, wherein, The number of potential embeddings in the potential embedding set is less than the number of data element embeddings in the data element embedding set.

15. The method according to any one of claims 1-13, wherein, Each data element embedding corresponds to a specific spatial location in the network input and includes a feature embedding based on features of the network input at that spatial location.

16. The method according to any one of claims 1-13, wherein, Each data element embedding corresponds to a specific spatial location in the network input and includes a location embedding that characterizes the spatial location.

17. The method according to any one of claims 1-13, wherein, Each data element embedding includes a modality embedding that defines the modality corresponding to the data element embedding.

18. A system comprising: One or more computers; as well as 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 operation of the method according to any one of claims 1-17.

19. A non-transitory computer storage medium for storing one or more instructions, said instructions, when executed by one or more computers, causing said one or more computers to perform the operation of the method according to any one of claims 1-17.

Images