Score interpolated diffusion model

Abstract

In one aspect, the method includes: initializing an output sequence by assigning a respective embedding to each of a plurality of output positions; processing a diffusion model input including the current continuous representation using the diffusion model to generate a diffusion model output in each of a plurality of de-diffusion steps; processing each initial score using a softmax function to generate, for each of the plurality of output positions, a probability distribution over a plurality of embeddings in a vocabulary of embeddings; and updating the continuous representation of the output sequence using the probability distribution and the vocabulary of embeddings.

Description

[Technical Field]

[0001] CROSS-REFERENCE TO RELATED APPLICATIONS This application claims priority to U.S. Provisional Application No. 63 / 427,784, filed November 23, 2022. The disclosure of the prior application is considered part of the disclosure of this application and is incorporated by reference into the disclosure of this application.

[0002] This specification relates to the use of neural networks to perform machine learning tasks on discrete data.

[0003] A neural network is a machine learning model that uses one or more layers of nonlinear units to predict an output given a received input. Some neural networks include one or more hidden layers in addition to an output layer. The output of each hidden layer is used as the input to the next layer in the network, i.e., the next hidden layer or the output layer. Each layer of the network generates an output from the received input according to the current values ​​of its respective set of parameters.

[0004] In some cases, discrete data refers to data that can take on or have a finite or countably infinite number of values, while continuous data can take on any value within a continuous range of possible values. In a broader sense, continuous data refers to data that can take on any value within a specified range, constrained only by the precision of the numeric format used in the computer system. Discrete data has additional constraints beyond the possible values ​​of the numeric format used in the computer system that further limit the possible values ​​the data can take within the specified range. Summary of the Invention [Problem to be solved by the invention]

[0005] This specification describes a system, implemented as a computer program on one or more computers at one or more locations, that uses a diffusion model to perform a de-diffusion process on a continuous embedding to generate an output sequence of discrete tokens.

[0006] In some cases, the generative process is unconditional, where the output sequence can be generated by the diffusion model from random noise. In other cases, the generative process is conditional, where, for example, the output text, computer program code, symbol, or signal sequence generated by the diffusion model is a complement or extension of the input text, computer program code, symbol, or signal sequence, and the output sequence is generated based on a conditioning input.

[0007] According to a first aspect, there is provided a computer-implemented method for generating an output sequence at each of a plurality of output positions, the output sequence including a respective discrete token selected from a vocabulary of discrete tokens, the method including: initializing a continuous representation of the output sequence by assigning a respective embedding to each of one or more of the plurality of output positions, each embedding selected from a vocabulary of embeddings including a respective embedding corresponding to each discrete token in the vocabulary of discrete tokens; obtaining a current continuous representation of the output sequence in each of a plurality of de-diffusion steps; processing a diffusion model input including the current continuous representation using a diffusion model to generate a diffusion model output, the diffusion model output including, for each of the plurality of output positions, a respective initial score for each of a plurality of embeddings in the vocabulary of embeddings; processing the respective initial scores using a softmax function to generate, for each of the plurality of output positions, a probability distribution over a plurality of embeddings in the vocabulary of embeddings including a respective probability score for each embedding in the vocabulary of embeddings; and iteratively performing updating the current continuous representation of the output sequence using the probability distribution and the vocabulary of embeddings.

[0008] Updating the current continuous representation of the output sequence using the probability distributions and the vocabulary of embeddings may include determining score function estimates for the output sequence using the respective probability distributions and the vocabulary of embeddings, and updating the current continuous representation of the output sequence by applying an ordinary differential equation (ODE) solver to the score function estimates.

[0009] Updating the current continuous representation of the output sequence by applying an ODE solver to the score function estimates may include using the ODE solver to evaluate the score function estimates to determine, for at least one of the subset of output positions in the current continuous representation of the output sequence, a respective updated embedding.

[0010] Determining the score function estimate may include using a vocabulary of probability distributions and embeddings to determine a respective interpolated embedding for each of a subset of the plurality of output locations.

[0011] Determining the score function estimate for the output sequence may include calculating, at each of the plurality of output positions, a difference between a respective interpolated embedding of the output position and a current embedding of the output position in the current continuous representation of the output sequence.

[0012] Processing a diffusion model input that includes a current continuous representation using a diffusion model may include applying a normalization to a current embedding in the current continuous representation.

[0013] The method may further include receiving a conditioning input sequence including one or more conditioning tokens, and initializing the output sequence may include selecting embeddings at one or more given output positions of the output positions based on the one or more conditioning tokens from a vocabulary of embeddings, and generating initial embeddings at remaining output positions of the plurality of output positions with at least some degree of randomness.

[0014] The method may further include receiving a conditioned input and processing the conditioned input using an encoder neural network to generate an encoder network output, wherein processing the diffusion model input including the current continuous representation using the diffusion model to generate the diffusion model output includes processing the diffusion model input using the diffusion model while the diffusion model is conditioned to the encoder network output to generate the diffusion model output.

[0015] The diffusion model input may include (i) a current continuous representation, (ii) embeddings at one or more given output positions selected based on one or more conditioning tokens, and (iii) a sequence of Boolean conditioning masks that identify which positions among the multiple output positions are the one or more given output positions.

[0016] The diffusion model input may further include respective interpolation embeddings for a subset of the plurality of output locations.

[0017] The diffusion model input may further include data specifying the noise level.

[0018] The method may further include determining a noise level by selecting a time step for the despreading step, applying time warping to the selected time step by evaluating a monotonic piecewise linear function over the selected time step to determine a warped time step, and determining the noise level from the warped time step.

[0019] Obtaining a current continuous representation of the output sequence may include using the current continuous representation updated at the time of the immediately preceding despreading step as the current continuous representation of the output sequence in the despreading step.

[0020] The diffusion model can be an attention-based neural network that includes one or more attention neural network layers.

[0021] The method may further include, after performing the multiple de-diffusion steps, processing a final diffusion model input using the diffusion model, the final diffusion model input including a current continuous representation updated as of a last de-diffusion step in the multiple de-diffusion steps, to generate a final diffusion model output, the final diffusion model output including, for each of the multiple output positions, a respective final initial score for each of a plurality of embeddings in a vocabulary of embeddings; selecting, for each of the multiple output positions, an embedding from the vocabulary of embeddings using the respective final initial score; and generating an output sequence including discrete tokens corresponding to the selected embedding.

[0022] The method may further include, after performing the multiple despreading steps, selecting, for each of the multiple output positions, an embedding from a vocabulary of embeddings that has a closest distance to an embedding in the current continuous representation at the updated output position as of a last despreading step in the multiple despreading steps, and generating an output sequence including discrete tokens corresponding to the selected embedding.

[0023] The diffusion model may be trained by optimizing the cross-entropy loss between the probability distribution over multiple embeddings in the vocabulary of embeddings generated using the diffusion model output and the true probability distribution over multiple embeddings in the vocabulary of embeddings.

[0024] Embeddings within the vocabulary of embeddings can be learned jointly with training the diffusion model by backpropagating the gradients of the cross-entropy loss computed with respect to the parameters of the diffusion model to the embeddings.

[0025] Discrete tokens may include text characters, symbols, or signals.

[0026] The diffusion model may be configured to perform one or more of a language modeling task, including prompt completion and / or gap completion tasks, or a machine translation task, a speech transcription task, a video transcription task, an image captioning task, or an agent control task.

[0027] According to another aspect, there is provided a system comprising one or more computers and one or more storage devices storing instructions operable, when executed by the one or more computers, to cause the one or more computers to perform the operations of any of the aforementioned methods.

[0028] According to another aspect, a computer storage medium is provided having encoded thereon instructions that, when executed by one or more computers, cause the one or more computers to perform the operations of any of the methods described above.

[0029] Particular embodiments of the subject matter described herein can be implemented to achieve one or more of the following advantages: This specification describes a framework for generative modeling of data using any of a variety of diffusion models that are continuous in time space, input space, or both. The described framework provides effective continuous diffusion on data, such as text data, that is discrete in nature, to achieve competitive performance on a wide range of conditional and unconditional discrete data generation tasks.

[0030] The described framework provides flexibility in model architecture and adaptability in denoising procedures when generating discrete output data. Specifically, the described framework allows for a tradeoff between computational cost and sample quality during sampling by selecting an appropriate number of iterative refinement steps without requiring retraining or fine-tuning. This is advantageous over generating discrete output data using existing autoregressive models, where the number of steps is necessarily equal to the length of the sequence to be generated. Furthermore, under the described framework, the same diffusion model can be used on a variety of different hardware device types and / or for a number of different latency requirements (e.g., when a fast response is required or when computation is limited due to a sudden surge in server demand, using fewer iterations with a diffusion model reduces quality but still provides overall results, whereas an autoregressive model does not provide overall results).

[0031] Advantageously, by using score interpolation (e.g., rather than score matching), which configures the diffusion model to compute score function estimates used during the iterative denoising process by interpolating between embedding values ​​of a continuous representation of the discrete output sequence, the described framework enables end-to-end training of the diffusion model and latent embeddings using a single cross-entropy loss function.

[0032] Additionally, time warping, an active learning technique that automatically adapts the distribution of noise levels sampled during training, can be used to further improve the training efficiency of diffusion models. Therefore, implementing a framework for generative modeling of discrete data requires reduced consumption of computational resources, e.g., reduced processor cycles, reduced memory, and reduced power consumption, compared to existing approaches.

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

[0034] Like reference numbers and designations in the various drawings indicate like elements. [Brief explanation of the drawings]

[0035] [Figure 1A] FIG. 1 is a block diagram of an exemplary data generation system during inference. [Figure 1B] FIG. 1 is a block diagram of an exemplary data generation system during training. [Figure 2A] 1 illustrates exemplary operations performed by components of a data generation system. [Figure 2B] 1 illustrates exemplary operations performed by components of a data generation system. [Figure 2C] 1 illustrates exemplary operations performed by components of a data generation system. [Figure 2D] 1 illustrates exemplary operations performed by components of a data generation system. [Figure 3] FIG. 1 is a flow diagram of an exemplary process for generating an output sequence. DETAILED DESCRIPTION OF THE INVENTION

[0036] 1A is a block diagram of an exemplary data generation system 100 during inference. Data generation 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 herein below may be implemented.

[0037] Data generation system 100 is a system that generates a discrete output sequence 162 in response to a received request. Data generation system 100 may store discrete output sequence 162 in an output data repository or provide discrete output sequence 162 for some other immediate use. For example, data generation system 100 may subsequently output discrete output sequence 162 for display, such as on a client device that submitted the data generation request.

[0038] The discrete output sequence 162 includes a plurality of discrete tokens. Each discrete token may be a separate, discrete data item. In this specification, continuous data refers to data that can take on any value within a specified range, constrained only by the precision of the numeric format used by the computer system. In contrast, discrete data has additional constraints beyond the possible values ​​of the numeric format used by the computer system that further limit the possible values ​​the data can take on within the specified range.

[0039] In one example, the discrete output sequence 162 includes a number of tokens from a finite number of possible tokens.

[0040] For example, the discrete output sequence 162 may include, at each of a plurality of positions, a respective token from a vocabulary of discrete tokens. The vocabulary of tokens may include any of a variety of discrete tokens representing text, symbols, or signals. For example, the vocabulary of discrete tokens may include one or more letters, subwords, words, punctuation marks, numbers, or other symbols that appear in a text corpus. The text may be, for example, natural language text or computer program code.

[0041] In some examples, the discrete output sequence 162 represents audio data. The audio data may include, for example, spoken words characterizing audio waveforms in the time domain or the time-frequency domain. In some examples, the discrete output sequence 162 may be a representation of phonemes in the case of text-to-speech (TTS) or a MIDI file containing musical note / instrument definitions. In some examples, the discrete output sequence 162 represents text. In one example, the discrete output sequence 162 is converted by a speech synthesis converter to form an audio signal. In some examples, the discrete output sequence 162 represents a sequence of actions to be performed by an agent, such as a mechanical agent in a real-world environment, performing actions to perform a mechanical task.

[0042] In some cases, data generation system 100 may be configured as an unconditional output sequence generation system that generates discrete output sequences 162 unconditionally, i.e., without conditioning on any conditioning inputs. In these cases, the discrete output sequences 162 generated by system 100 approximate samples of the distribution of training input sequences used during training of data generation system 100.

[0043] In other cases, the data generation system 100 may be configured as a conditional output sequence generation system that generates a discrete output sequence 162 that is conditioned on the input sequence 102 .

[0044] In some of these cases, input sequence 102 comprises a discrete input sequence. Like discrete output sequence 162, the discrete input sequence may comprise one or more tokens selected from a vocabulary of tokens. The vocabulary of tokens may include any of a variety of discrete tokens representing text, symbols, or signals.

[0045] In these other cases, the input sequence 102 comprises a continuous input sequence (i.e., an input sequence of values ​​that are each a continuous value), for example, the continuous input sequence represents a sequence of pixels of an image (e.g., one or more color intensity values ​​for each pixel), which may be a still image or a video (e.g., the continuous input sequence may include portions of each of multiple images (image frames of a video), each portion being one or more color intensity values ​​for each pixel in succession), or may represent an audio waveform or an audio spectrogram.

[0046] When configured as a conditional output sequence generation system, for example, the data generation system 100 can receive an input sequence 102 that includes text data, additional non-text data, e.g., data representing images, video, audio, etc., or both, and generate a discrete output sequence 162 that is a sequence of text, such as a completion (e.g., gap fill) of the input sequence 102, an extension of the input sequence 102, a response to a question posed in the input sequence 102, a sequence of text on a topic specified by the input sequence 102, a text description of data of a different data type, e.g., a text caption of an image, an audio transcription, a transcription of the audio contained in one or more video frames, etc.

[0047] As another example, the data generation system 100 may receive an input sequence 102 including one or more code segments and generate a discrete output sequence 162 including code segments conditioned on the input sequence 102, such as one or more code segments that, when combined with the code segments included in the input sequence 102, constitute an executable application, program, object, or sequence of instructions.

[0048] In either case, text (both natural language text and computer program code), symbols, and signals will be understood as merely examples of discrete data for purposes of illustration, and it will be understood that the discrete tokens can be discrete data items in many other forms or formats. For example, the discrete output sequence 162 can be or include a sequence of biological data (e.g., a sequence of gene expression), a sequence of electronic health record data (e.g., a sequence of health events), a sequence of clinical procedure data (e.g., a sequence of physician orders, clinical documentation, notes, diagnostic codes, medications, etc.), etc.

[0049] In one example, the input may be received in the form of an audio (voice) signal captured by a microphone and converted by a speech recognition converter to form the input sequence 102 .

[0050] In some examples, the input sequence of discrete tokens 102 includes tokens that represent audio data input, including, for example, spoken words that characterize audio waveforms in the time domain or time-frequency domain.

[0051] In another example, the output sequence of discrete tokens 163 represents a sequence of actions to be performed by an agent. The agent may be, for example, a mechanical agent, a real-world or simulated robot, a control system for an industrial facility, or a control system controlling different types of agents. The actions may be discrete control inputs for controlling a robot, for example, a control command selected from a set of control commands for an agent or an autonomous or semi-autonomous land, air, or sea vehicle.

[0052] In another example, the input sequence 102 may be a query input that includes an input token string and at least one data item.

[0053] For example, the data item(s) may be image(s). Each data item may be a still image. For example, a data item may be a pixel value (e.g., red-green-blue (RGB) value) for each pixel of a pixel array. Alternatively, one or more of the data items may be a video image. For example, the data may be a pixel value for each pixel of a respective array for each of a plurality of frames. The image(s) may be captured by imaging the real world with a still or video camera. In another possibility, the data item(s) may be audio signal(s) (an audio signal is audio data representing values ​​of an audio waveform at each of a plurality of points in time, e.g., sound captured by a microphone over a certain period of time). In a further possibility, the data items may be video images with accompanying respective soundtracks.

[0054] In particular, to generate the discrete output sequence 162, the data generation system 100 uses the input embedding engine 110, the diffusion model neural network 120, the time embedding neural network 150, and the update engine 160 to perform a despreading process to generate the discrete output sequence 162 over multiple despreading steps. The discrete output sequence 162 is generated after the final despreading step of the despreading process.

[0055] Before the despreading process begins, the input embedding engine 110 initializes a continuous representation of the discrete output sequence 162, i.e., generates an initial continuous representation. The continuous representation may include multiple embeddings. An "embedding" may generally refer to an ordered collection of numbers, such as a vector, matrix, or other tensor of numbers. For example, each embedding may include multiple numbers for multiple embedding dimensions (e.g., embedding dimensions of 128, 256, or 512). Thus, the continuous representation provides a numerical encoding of the discrete input sequence 102 in a continuous embedding space.

[0056] The dimension of the initial continuous representation depends on the dimension of the discrete output sequence 162 (e.g., depends on the number of tokens included in the discrete output sequence 162). If the discrete output sequence 162 includes a discrete token at each of multiple output positions, the initial continuous representation of the discrete output sequence 162 may include the same number of output positions as the discrete output sequence 162 and may include one embedding at each output position.

[0057] When configured as an unconditional output sequence generation system, the input embedding engine 110 generates the initial continuous representation by sampling an initial value for each of multiple embeddings included in the initial continuous representation from a corresponding noise distribution, e.g., a Gaussian distribution or other predetermined distribution. Thus, the initial continuous representation includes multiple embeddings, and the initial value of each embedding is sampled from a corresponding noise distribution.

[0058] When configured as a conditional output sequence generation system, the input embedding engine 110 processes the received discrete input sequence 102 to generate one or more conditioned embeddings, each containing a learned numerical value. The input embedding engine 110 then generates an initial continuous representation based on using the conditioned embeddings as embeddings in some of the output positions of the initial continuous representation and sampling the initial numerical values ​​included in the embeddings in the remaining output positions of the initial continuous representation from a corresponding noise distribution, e.g., a Gaussian distribution or other predetermined distribution.

[0059] In particular, the input embedding engine 110 can generate conditioned embeddings from the received discrete input sequence 102 by selecting, from a vocabulary of embeddings, a respective embedding that maps to each discrete input token included in the discrete input sequence 102, and then assigning each selected embedding to some of a plurality of output positions of the initial continuous representation.

[0060] The vocabulary of embeddings maintained by the input embedding engine 110 may include a respective embedding corresponding to each discrete token in the vocabulary of discrete tokens, thereby defining a mapping between the vocabulary of discrete tokens and the vocabulary of embeddings. Each embedding in this vocabulary may have a numerical value that is learned during training of the data generation system 100.

[0061] After the input embedding engine 110 generates the initial continuous representation of the discrete output sequence 162, the data generation system 100 includes the initial continuous representation in the diffusion model input 119, which is provided as an input to the diffusion model neural network 120 for further processing.

[0062] The diffusion model input 119 also includes other data. In particular, for each de-diffusion step, the diffusion model input 119 also includes a corresponding time step embedding generated by the time embedding neural network 150 from data defining the corresponding time step t 142 (typically, t = t max From t=t min Time step t142 (performed until t142) identifies the current despreading step in the despreading process. Generally, different despreading steps throughout the despreading process correspond to different noise levels in the despreading process, and therefore time step t142 also identifies the noise level that corresponds to the current despreading step currently being performed.

[0063] Time embedding neural network 150 may have any suitable neural network architecture that enables the neural network to map data defining time step t 142 into a time embedding. Particular example architectures are described further below, but more generally, time embedding neural network 150 can include any suitable number (e.g., 1, or 5, or 10) of any suitable type of neural network layers (e.g., embedding layers, fully connected layers, convolutional layers, etc.), connected in any suitable configuration (e.g., as a directed graph of layers).

[0064] At each dediffusion step, the diffusion model neural network 120 processes the diffusion model input 119 to generate a diffusion model output 121, and then updates the continuous representation of the discrete output sequence 162, i.e., uses the diffusion model output 121 generated by the diffusion model neural network 120 to generate an updated continuous representation of the discrete output sequence 162 at the time of the dediffusion step.

[0065] The diffusion model neural network 120 may have any suitable architecture that enables the neural network to map a diffusion model input 119, which includes a current continuous representation of the discrete output sequence 162 (as of a given dediffusion step), to a diffusion model output 121.

[0066] Specific exemplary architectures of the diffusion model neural network 120 are described further below. More generally, the diffusion model neural network 120 can be and / or include a neural network having any suitable architecture. For example, the diffusion model neural network 120 can include one or more fully connected layers. As another example, the diffusion model neural network 120 can be an attention neural network, e.g., having a Transformer or other architecture including one or more attention layers. An attention layer is a layer that includes an attention mechanism, e.g., a scaled dot-product attention mechanism. For example, the diffusion model neural network 120 can include one or more self-attention layers (e.g., one or more multi-head self-attention layers), one or more cross-attention layers, or both. These examples are not exhaustive, and it can be readily understood that the diffusion model neural network 120 can include any suitable type of neural network layer, connected in any suitable number and in any suitable configuration.

[0067] In particular, the diffusion model neural network 120 is trained using a score interpolation objective on a set of training data to generate the diffusion model outputs 121. In other words, the data generation system 100 obtains score function estimates from the diffusion model outputs 121 through linear interpolation and then uses the score function estimates for sampling by an ordinary differential equation (ODE) solver 165 to update the current continuous representation of the discrete output sequence 162.

[0068] Using such a score interpolation objective during training, instead of, for example, a score matching objective, can be advantageous because at least score interpolation enables the use of cross-entropy loss during training, thereby enabling end-to-end training of the diffusion model neural network 120 contained in the input embedding engine 110 and the data generation system 100 including embeddings with a single loss function. This can improve the effectiveness of training by allowing these two system components to be jointly adaptive.

[0069] On the other hand, when a score matching objective is used, joint training of the diffusion model neural network 120 with the embeddings contained in the input embedding engine 110 can lead to a collapse of the embedding space (where multiple embeddings correspond to the same vector).

[0070] To facilitate score interpolation, at each de-diffusion step, the diffusion model output 121 includes multiple logit vectors for each of multiple output positions of the output sequence 162. Each logit vector includes a respective initial score for each of multiple embeddings in the vocabulary of embeddings. Thus, for a given output position and for a given embedding vector, each initial score included in the logit vector may represent the predicted probability of a discrete token mapped to a given embedding vector located at the given output position.

[0071] The data generation system 100 applies a softmax function 122 to the logit vector contained in the diffusion model output 121 to generate, for each of a plurality of output positions, a probability distribution over a plurality of embeddings in the vocabulary of embeddings in the de-diffusion step.

[0072] The data generation system 100 uses probability distributions to generate score function estimates 132 through linear interpolation, as described in more detail below, and then uses an update engine 160 to update the current continuous representation of the output sequence by applying an ODE solver 165 to the score function estimates 132.

[0073] One way that the update engine 160 can use the score function estimates to update the current continuous representation is to apply an ordinary differential equation (ODE) solver 165 to the score function estimates to update the current continuous representation.

[0074] In such an approach, update engine 160 may use any suitable ODE solver to update the current continuous representation. For example, ODE solver 165 may be an ODE solver using Euler's first-order method or Huygen's second-order method (e.g., one of the ODE solvers described in T. Karras, et al., "Elucidating the design space of diffusion-based generative models." arXiv preprint arXiv:2206.00364, 2022). As another example, ODE solver 165 may be an ODE solver using a stochastic or more advanced method, such as the ODE solver described in C. Lu, et al., "Dpmsolver: A fast ode solver for diffusion probabilistic model sampling in around 10 steps." arXiv preprint arXiv:2206.00927, 2022).

[0075] By repeatedly using the diffusion model neural network 120 to generate, for each of a plurality of dediffusion steps throughout the dediffusion process, a diffusion model output 121 that defines a corresponding score function estimate for the dediffusion step, and by repeatedly using the update engine 160 to update the current continuous representation with the score function estimate for each dediffusion step, the data generation system 100 may iteratively update the initial continuous representation to generate a final continuous representation. Thus, the final continuous representation of the discrete output sequence 162 is an updated continuous representation as of the last dediffusion step of the dediffusion process.

[0076] The initial continuous representation has the same dimensionality as the updated continuous representation (and similarly, the final continuous representation), but different values, i.e., the initial continuous representation includes multiple embeddings, and the updated continuous representation (and similarly, the final continuous representation) includes the same number of embeddings, but the numerical values ​​included in at least some of these embeddings are different.

[0077] After performing multiple despreading steps, data generation system 100 may generate discrete output sequence 162 from the final continuous representation of discrete output sequence 162. Discrete output sequence 162 includes, at each of multiple output positions, a respective discrete token selected from a vocabulary of discrete tokens.

[0078] In some cases, the data generation system 100 can do this by using the final continuous representation of the discrete output sequence 162 to select an embedding for each of multiple output positions from a vocabulary of embeddings.

[0079] Specifically, for each of a plurality of output positions of the discrete output sequence 162, the data generation system 100 selects an embedding from the vocabulary of embeddings that has the closest distance to the embedding in the final continuous representation, and then generates the discrete output sequence 162 including discrete tokens that each map to the selected embedding according to a mapping defined by the vocabulary of embeddings maintained by the input embedding engine 110.

[0080] For example, for a given output position, the data generation system 100 can use a distance measure such as Euclidean distance or cosine similarity to determine which embedding in the vocabulary of embeddings is the “closest” embedding to the embedding in the final continuous representation.

[0081] In some other cases, the data generation system 100 runs the diffusion model neural network 120 an additional time after the final dediffusion step to generate the discrete output sequence 162. Specifically, the data generation system 100 processes the final diffusion model input including the final continuous representation using the diffusion model neural network 120 to generate the final diffusion model output. The final diffusion model output includes, for each of the multiple output positions, a respective final initial score for each of the multiple embeddings in the vocabulary of embeddings.

[0082] The data generation system 100 then uses the respective final initial scores to select an embedding from the vocabulary of embeddings for each of the multiple output positions, and then generates a discrete output sequence 162 including discrete tokens that each map to the selected embedding according to the mapping defined by the vocabulary of embeddings maintained by the input embedding engine 110.

[0083] For example, for a given output position, the data generation system 100 may greedily select the highest scoring embedding from the vocabulary of embeddings, or may sample embeddings from the vocabulary of embeddings according to their respective final initial scores, e.g., using kernel sampling or other sampling techniques.

[0084] 1B is a block diagram of an exemplary data generation system 100 during training. In general, the data generation system 100 may be trained in an end-to-end manner using a single loss function, e.g., a cross-entropy loss function, a focal loss function, or a hinge loss function, to determine trained values ​​for embeddings contained in a vocabulary of embeddings maintained by the input embedding engine 110 and trained values ​​for the parameters of the diffusion model neural network 120 and the temporal embedding neural network 150 against a set of training data.

[0085] For example, the set of training data may include a plurality of discrete training sequences, each including at each of a plurality of positions a respective token from a vocabulary of tokens, which may include any of a variety of discrete tokens representing text, symbols, or signals.

[0086] In one example, the set of training data can include the MassiveText dataset described in JW Rae, et al. Scaling language models: Methods, analysis & insights from training gopher. arXiv preprint arXiv:2112.11446, 2021. In another example, the set of training data can include the C4 dataset described in Raffel et al., "Exploring the limits of transfer learning with a unified text-to-text transformer," J. Mach. Learn. Res., 21:140:1-140:67, 2020, the entire contents of which are incorporated herein by reference. In another example, the set of training data may include a set of training data items, each training data item including an audio file, an image file, and / or a video file (e.g., an audio file / image file and / or video file captured from the real-world environment by a microphone or camera) and an appropriate caption (e.g., if the training data item includes an audio file, or if the training data item includes an image or video file showing one or more elements in a location, a caption describing the element(s). In another example, the set of training data may include a set of training data items, each training data item including an audio file, an image file, and / or a video file (e.g., an audio file / image file and / or video file captured from the real-world environment by a microphone or camera) and an appropriate command for a real-world agent, such as a robot.In another example, the set of training data can include a set of training data items, each training data item including a text input and an audio file, an image file, and / or a video file (e.g., an audio file / image file and / or a video file captured from a real-world environment by a microphone or camera) that is an appropriate response to the text input (e.g., an image depicting skiing if the text input is "provide an image of skiing"). In another example, the set of training data can include a set of training data items, each training data item including a text input in a first natural language and a text output translating the text input into a second natural language.

[0087] The data generation system 100 uses discrete training sequences included in a set of training data to train the aforementioned components of the system 100. More specifically, for a given discrete training sequence 103 selected from the set of training data, the system uses an input embedding engine 110 to generate a continuous representation of the given discrete training sequence 103 based on the given discrete training sequence 103, and further adds time-varying noise to the continuous representation to generate a noisy continuous representation of the given discrete training sequence 103.

[0088] Next, the data generation system 100 performs a (partial) de-diffusion step on the noisy continuous representation of the given discrete training sequence 103. That is, the system processes the training diffusion model input 129, which includes the noisy continuous representation of the given discrete training sequence, using the diffusion model neural network 120 to generate the training diffusion model output 131, and then processes the training diffusion model output 131 using a softmax function to generate a probability distribution 133 across multiple embeddings in the vocabulary of embeddings for each of multiple output positions of the given discrete training sequence. Unlike during inference, the system does not need to apply the update engine 160 to use the probability distribution 133 to generate an updated continuous representation of the given discrete training sequence.

[0089] FIG. 1B shows that for each output position, the data generation system 100 computes a cross-entropy loss function that evaluates the cross-entropy loss between (i) a probability distribution 133 across multiple embeddings in the vocabulary of embeddings generated using the training diffusion model output 131 and (ii) the true probability distribution across multiple embeddings in the vocabulary of embeddings.

[0090] For each output position, the true probability distribution may be, for example, a probability distribution that assigns higher probability scores to true embeddings that map to a particular token located at the output position of the given discrete training sequence 103 and lower probability scores to the remaining embeddings that respectively map to other tokens in the vocabulary of tokens that are not located at the output position of the given discrete training sequence 103.

[0091] The data generation system 100 then updates the parameter values ​​of the diffusion model neural network 120 and the time embedding neural network 150 by using an appropriate machine learning training technique, e.g., gradient descent with a backpropagation training technique that uses a conventional optimizer, e.g., stochastic gradient descent, RMSprop, or Adam optimizer, to optimize the cross-entropy loss function.

[0092] Additionally, the data generation system 100 updates the values ​​of the embeddings contained in the embedding vocabulary maintained by the input embedding engine 110 by backpropagating the gradients of the cross-entropy loss computed with respect to the parameters of the diffusion model neural network 120 to the embeddings contained in the embedding vocabulary. Thus, the embedding vocabulary is learned jointly with the training of the diffusion model neural network 120.

[0093] During training, the data generation system 100 may incorporate any number of techniques to improve the efficiency, effectiveness, or both of the training process.

[0094] In particular examples, the data generation system 100 may employ an active learning strategy to improve efficiency by automatically adapting the distribution of noise levels sampled during training.

[0095] FIG. 1B illustrates that the data generation system 100 utilizes an optional time warping engine 140 during training. Without this time warping, there is a predetermined one-to-one relationship between the dediffusion steps performed by the diffusion model and the level of noise, which is a function of "time," denoted by t, σ(t), discretized as the number of time steps corresponding to the result of each dediffusion step. Time t corresponds to time in the diffusion process modeled (inversely) by the diffusion model. In the diffusion process, the continuous representation is gradually corrupted by increasing levels of noise. In the dediffusion process, each dediffusion step performed by the diffusion model corresponds to an update of the continuous representation, where time t is updated to a lower value of t, with a corresponding lower noise value σ(t). The update of the continuous representation is modeled as a differential equation and solved by the ODE solver 165, as described above. In this diffusion process, the time t = 0 (or more generally, t min ) may be a noise-free (ideal) continuous representation (e.g., corresponding to the desired discrete output sequence 162) and may be expressed as a time series of a signal at a later time T (or t max ), diffusion produces completely corrupted data. As mentioned above, at each time t, the noise level is σ(t), where σ t also denoted σ(t'), which may be an increasing function of t or may actually be equal to t. Time warping maps t to a value t', so that in the despreading step corresponding to time step t, the level of noise applied is σ(t') instead of σ(t).

[0096] In the example, for warping, for a given despreading step during training, the time warping engine 140 determines the range of possible time steps (t = t min From t=t maxThe time step for a given despreading step is selected based on uniformly sampling time steps t from (between t and t), and applying time warping to the sampled time steps t by evaluating a monotonic piecewise linear function F over the sampled time steps t to determine a warped time step t' 143. The warped time step t' 143 identifies an adapted noise level σ(t') for the given despreading step currently being performed during training.

[0097] Such monotonic piecewise linear functions can be parameterized by a set of learnable parameters, for example, by dividing both the input and output ranges of the function into a fixed number of bins and applying one of the parameterization techniques described in C. Durkan, A., et al., Neural spline flows. Advances in neural information processing systems, 32, 2019, and T. Mueller, et al., Neural importance sampling. ACM Transactions on Graphics (TOG), 38(5):1-19, 2019, the entire contents of which are incorporated herein by reference.

[0098] When such an active learning strategy is used, the time-varying noise added to the continuous representation by the input embedding engine 110 to generate a noisy continuous representation of a given discrete training sequence 103 is adapted noise determined according to a non-uniform warped time step t, i.e., instead of a uniformly sampled time step t. Furthermore, data identifying the warped time step t' is received and processed by the time embedding neural network 150 when generating the time step embeddings included in the training diffusion model input 129.

[0099] 2A illustrates exemplary operations performed by an input embedding engine 210. The input embedding engine 210 may be the same as or similar to the input embedding engine 110 included in the data generation system 100 of FIG.

[0100] At a high level, before the despreading process begins, the input embedding engine 210 can perform these operations to generate an initial continuous representation of a discrete output sequence including a plurality of output positions 202A-202F and to generate a sequence of Boolean conditioning masks, both of which are included as part of the diffusion model input for the first despreading step.

[0101] The input embedding engine 210 receives a discrete input sequence that includes a plurality of conditioning tokens. In the example of Figure 2A, the discrete input sequence includes three conditioning tokens: "the," "cat," and "mat."

[0102] To generate the initial continuous representation, the input embedding engine 210 selects a respective embedding (also called a "conditioning embedding") from the embedding vocabulary that maps to each conditioning token included in the discrete input sequence, and then assigns the selected conditioning embedding to some of a number of output positions (also called "clean" output positions) in the initial continuous representation. As described above, the selected conditioning embeddings have values ​​that were learned during training of the data generation system.

[0103] For example, FIG. 2A shows that the initial continuous representation includes conditioning embeddings at clean output positions 202A, 202B, and 202F that map to conditioning tokens of "the," "cat," and "mat," respectively.

[0104] The input embedding engine 210 generates an embedding at each remaining output position (also called a "noisy" output position) in the initial continuous representation by sampling the initial numerical value included in each such embedding from a corresponding noise distribution, e.g., a Gaussian distribution or other predetermined distribution.

[0105] For example, FIG. 2A shows that the initial continuous representation includes an embedding with randomly initialized numerical values ​​at noisy output locations 202C, 202D, and 202E.

[0106] The input embedding engine 210 also generates a sequence of Boolean conditioning masks. The sequence of Boolean conditioning masks identifies which of the multiple output locations are clean output locations and which are noisy output locations. For example, the sequence of Boolean conditioning masks can include a zero for each clean output location and a one for each noisy output location, or vice versa.

[0107] 2B illustrates exemplary operations performed by a time-embedding neural network 250. The time-embedding neural network 250 may be the same as or similar to the time-embedding neural network 150 included in the data generation system 100 of FIG.

[0108] The time embedding neural network 250 processes data defining time step t to generate a time-step embedding. The time step t identifies the noise level corresponding to the current despreading step in the despreading process. The time embedding neural network may have any architecture that allows the neural network to map the data defining time step t to a time-step embedding.

[0109] In the example of Figure 2B, the time embedding neural network 250 includes a random Fourier embedding layer, followed by a linear layer, followed by a GeLU activation layer, followed by another linear layer, followed by another GeLU activation layer. Each layer includes 128 units. The time embedding neural network 250 processes data defining a time step t through these layers to generate a time-step embedding.

[0110] However, in other examples, the temporal embedding neural network 250 may include a different number of layers, different types of layers, or both.

[0111] 2C illustrates exemplary operations performed by the diffusion model neural network 220. The diffusion model neural network 220 may be the same as or similar to the diffusion model neural network 120 included in the data generation system 100 of FIG.

[0112] In the example of Figure 2C, the diffusion model neural network 220 is configured as a transform neural network containing a total of N attention blocks arranged between an input linear layer with 1024 units and an output linear layer with 32000 units. Each attention block contains a self-attention layer with eight attention heads and a RoPE embedding layer, followed by two linear layers separated by a GeLU activation layer. Each attention block also contains two layer normalization layers.

[0113] However, in other examples, a diffusion model neural network may include different types of layers, different numbers of layers, or both. For example, a diffusion model neural network may include different numbers of attention blocks, each of which may include the same or different layers.

[0114] At each dediffusion step, the diffusion model neural network 220 receives diffusion model inputs, which may include (i) the current continuous representation of the output sequence, (ii) conditioning embeddings of clean output positions selected by the input embedding engine 210 based on the received discrete input sequence, and (iii) a sequence of Boolean conditioning masks.

[0115] In the very first despreading step, the current continuous representation is the initial continuous representation generated by the input embedding engine 210. In any subsequent despreading step, the current continuous representation is the updated current continuous representation generated in the immediately preceding despreading step. In either case, the current embeddings included in the current continuous embedding are not constrained to any embedding in the embedding vocabulary.

[0116] Optionally, the diffusion model input also includes a self-conditioning input, which includes an interpolated embedding for each of the noisy positions of the plurality of output positions. Generating the interpolated embeddings is described further below.

[0117] The diffusion model neural network 220 processes the diffusion model inputs to generate diffusion model outputs, which include, for each of a plurality of output positions, a respective initial score for each of a plurality of embeddings in the vocabulary of embeddings.

[0118] To that end, the output layer of the diffusion model neural network 220 may include a respective unit corresponding to each embedding. For example, as shown in FIG. 2C, in an implementation in which the input embedding engine 210 uses a vocabulary of 32,000 tokens associated with the SentencePiece tokenizer (described in Kudo et al., “Sentencepiece: A simple and language independent subword tokenizer and detokenizer for neural text processing,” EMNLP (Demonstration), pp. 66-71, Association for Computational Linguistics, 2018), the output layer may be a linear layer with 32,000 units, each configured to generate an initial score for each of the embeddings that map to one of the 32,000 tokens.

[0119] 2D illustrates exemplary operations performed by softmax layer 222 and update engine 260. Softmax layer 222 and update engine 260 may be the same as or similar to softmax layer 122 and update engine 160, respectively, included in data generation system 100 of FIG.

[0120] At each dediffusion step, the softmax layer 222 applies a softmax function to the diffusion model outputs generated by the diffusion model neural network 220 at the dediffusion step to generate, for each of the multiple output positions, a probability distribution p across multiple embeddings in the vocabulary of embeddings. The probability distribution p may include a respective probability score for each embedding in the vocabulary of embeddings.

[0121] The update engine 260 may then use the probability distribution p to update, i.e., denoise, the current continuous representation of the output sequence in the despreading step according to the score function estimate of the output sequence calculated from the probability distribution p.

[0122] For a given output position, the score function estimate is defined as:

number

[0123] The right side of this equation is the expected value

number

number

[0124] Thus, the update engine 260 may calculate the score function estimate as follows:

number

[0125] In other words, for each noisy output position in the current continuous representation of the output sequence, the update engine 260 first calculates the respective interpolated embedding

number

[0126] The score function estimate has the form of a numerator divided by a denominator, where the numerator depends on the difference between the interpolated embedding of the noisy output location and the current embedding of the noisy output location, and the denominator depends on the square of the time step t.

[0127] The sampling engine 260 applies an ordinary differential equation (ODE) solver to the score function estimates to update the current continuous representation of the output sequence. The ODE solver may be any suitable ODE solver, for example, using Euler's method or Hoyne's method. In particular, the sampling engine 260 uses the ODE solver to evaluate the score function estimates to determine a respective updated embedding for each noisy output location in the current continuous representation of the output sequence.

[0128] After performing multiple despreading steps, the update engine 260 (or other component of the data generation system) may generate a discrete output sequence. One way to do this is to use a final continuous representation of the discrete output sequence. The final continuous representation of the discrete output sequence is the updated continuous representation generated in the last despreading step of the despreading process.

[0129] In this approach, the update engine 260 selects, for each noisy output location, an embedding from the vocabulary of embeddings that has the closest distance to the embedding in the final continuous representation, and then generates a discrete output sequence containing discrete tokens at the noisy output locations that each map to the selected embedding. In the example of Figure 2D, three tokens, "ran," "off," and "the," are generated at noisy output locations 202C, 202D, and 202E, respectively, based on the selected embeddings.

[0130] 3 is a flow diagram of an exemplary process 300 for generating an output sequence. The output sequence includes, at each of a plurality of output locations, a respective discrete token selected from a vocabulary of discrete tokens. For convenience, process 300 will be described as being performed by one or more computer systems located at one or more locations. For example, a data generation system, such as data generation system 100 of FIG. 1, suitably programmed in accordance with this specification, can perform process 300.

[0131] In general, the system may repeatedly perform one iteration of steps 302-308 in each of multiple despreading steps (hereinafter referred to as the "current" despreading step) of the despreading process to generate an output sequence.

[0132] In some cases, the number of iterations is fixed. In other cases, the system or other system may adjust the number of iterations based on a latency requirement for generating the output sequence, i.e., the number of iterations may be selected such that the output sequence items are generated to meet the latency requirement. In still other cases, the system or other system may adjust the number of iterations based on a computational resource consumption requirement for generating the output sequence, i.e., the number of iterations may be selected such that the output sequence items are generated to meet the requirement. For example, the requirement may be the maximum number of floating-point operations (FLOPS) executed as part of generating the output sequence.

[0133] In some cases, the number of iterations can be any number between 10 and 500, inclusive. In some cases, the number of iterations can be any number between 10 and 50, inclusive. In some cases, the number of iterations can be less when a stochastic sampler or more advanced methods are used by the ODE solver.

[0134] Before performing the first iteration of steps 302-308, the system generates an initial continuous representation of the output sequence, which contains the same number of output positions as the discrete output sequence. The system can generally do this by assigning a respective embedding to each of the multiple output positions.

[0135] For example, when a discrete input sequence including multiple tokens is received, the system may select conditioned embeddings from the embedding vocabulary that correspond to each of the tokens included in the discrete input sequence. The embedding vocabulary may include a respective embedding that corresponds to each discrete token in the discrete token vocabulary, thereby defining a mapping between the discrete token vocabulary and the embedding vocabulary.

[0136] In this example, the system may generate an initial continuous representation based on using a conditioned embedding as the embedding at some of the output positions of the initial continuous representation (the "clean" output positions) and sampling the initial numerical values ​​included in the embedding at the remaining output positions of the initial continuous representation (the "noisy" output positions) from a corresponding noise distribution, e.g., a Gaussian distribution or other predetermined distribution.

[0137] The system obtains a current continuous representation of the output sequence for the current despreading step, step 302. The current continuous representation includes a current embedding for each of a plurality of output positions.

[0138] If the current despreading step is the very first despreading step, the current continuous representation is the initial continuous representation. If the current despreading step is a subsequent despreading step, the current continuous representation is the updated continuous representation generated in the immediately preceding despreading step.

[0139] In some implementations, the system applies normalization to the current embedding in the current continuous representation. For example, the normalization can be L2 normalization.

[0140] The system processes the diffusion model inputs using a diffusion model neural network to generate diffusion model outputs (step 304). The diffusion model inputs may include (i) a (normalized) current continuous representation of the output sequence, (ii) conditioning embeddings at clean output positions selected by the system based on the received discrete input sequence, (iii) a sequence of Boolean conditioning masks, and (iv) data defining the current de-diffusion step. The sequence of Boolean conditioning masks identifies which of multiple output positions in the current continuous representation are clean output positions and which are noisy output positions.

[0141] Optionally, in some implementations, the diffusion model inputs also include self-conditioning inputs, which may include respective interpolated embeddings for the noisy positions of the plurality of output positions, where, as described above, each interpolated embedding was generated by the system in a previous de-diffusion step to facilitate the calculation of a score function estimate for the output sequence.

[0142] Further optionally, in some implementations, the diffusion model input also includes conditioning data generated from additional data other than text, such as data representing images, video, audio, etc., received by the system as part of the conditioning input. In some of these implementations, the system may process the conditioning data using an encoder neural network to generate an encoder network output, and then include the encoder network output in the diffusion model input. Thus, the diffusion model neural network is conditioned on the encoder network output to generate the diffusion model output.

[0143] The diffusion model output includes, for each of a plurality of output positions, a respective initial score for each of a plurality of embeddings in the vocabulary of embeddings.

[0144] The system processes the respective initial scores using a softmax function to generate a probability distribution over multiple embeddings in the vocabulary of embeddings for each of the multiple output positions (step 306). For each output position, the probability distribution may include a respective probability score for each embedding in the vocabulary of embeddings.

[0145] The system updates the current continuous representation using the probability distributions and the embedding vocabulary to generate an updated continuous representation of the output sequence for the current despreading step (step 308).

[0146] Specifically, the system can do this by determining a score function estimate for the output sequence using the respective probability scores and a vocabulary of embeddings, and applying an ordinary differential equation (ODE) solver to the score function estimate. By applying the ODE solver, the system can generate a respective updated embedding for each noisy output position of a plurality of output positions in the current continuous representation of the output sequence.

[0147] The term "configured" is used herein in connection with systems and computer program components. When one or more computer systems are configured to perform a particular operation or action, it means that the system has installed thereon software, firmware, hardware, or a combination thereof that, during operation, causes the system to perform the operation or action. When one or more computer programs are configured to perform a particular operation or action, it means that the one or more programs contain instructions that, when executed by a data processing device, cause the device to perform the operation or action.

[0148] Embodiments of the subject matter and functional operations described herein may be implemented in digital electronic circuitry, tangibly embodied computer software or firmware, computer hardware, such as the structures disclosed herein and their structural equivalents, or one or more combinations thereof. Embodiments of the subject matter described herein may 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 controlling the operation of a data processing apparatus. The computer storage medium may be a machine-readable storage device, a machine-readable storage substrate, a random-access or serial-access memory device, or one or more combinations thereof. Alternatively or additionally, the program instructions may be encoded in an artificially generated propagated signal, such as a mechanically generated electrical, optical, or electromagnetic signal, generated to encode information for transmission to a receiving device suitable for execution by the data processing apparatus.

[0149] The term "data processing apparatus" refers to data processing hardware and encompasses all kinds of apparatus, devices, and machines for processing data, including, by way of example, a programmable processor, a computer, or multiple processors or multiple computers. An apparatus may also be or further include special-purpose logic circuitry, such as an FPGA (field-programmable gate array) or an ASIC (application-specific integrated circuit). An apparatus may optionally include, in addition to hardware, code that creates an execution environment for a computer program, such as code that constitutes processor firmware, a protocol stack, a database management system, an operating system, or one or more combinations thereof.

[0150] A computer program, which may be referred to or described as a program, software, software application, app, module, software module, script, or code, can be written in any form of programming language, including compiled or interpreted, or declarative or procedural, and can be deployed in any form, including as a stand-alone program or as a module, component, subroutine, or other unit suitable for use in a computing environment. A program may, but need not, correspond to a file in a file system. A program can be stored in a portion of a file that holds other programs or data, for example, one or more scripts stored in a markup language document, in a single file dedicated to the program, or in multiple cooperating files, for example, files that store one or more modules, subprograms, or portions of code. A computer program can be deployed to run on one computer or on multiple computers that are located at one site or distributed across multiple sites and interconnected by a data communications network.

[0151] The term "database" is used broadly herein to refer to any collection of data, which need not be structured in any particular way, or even at all, and which may be stored on storage devices in one or more locations. Thus, for example, an index database may contain multiple collections of data, each of which may be organized and accessed in a different way.

[0152] Similarly, the term "engine" is used broadly herein 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 installed on one or more computers at one or more locations. In some cases, one or more computers are dedicated to a particular engine, and in other cases, multiple engines can be installed and run on the same computer or computers.

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

[0154] A computer suitable for running a computer program may be based on a general-purpose or special-purpose microprocessor, or both, or on another type of central processing unit. Generally, the central processing unit receives instructions and data from a read-only memory or a random-access memory, or both. The basic components of a computer are a 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 in, special-purpose logic circuitry. Generally, a computer also includes one or more mass storage devices, such as magnetic, magneto-optical, or optical disks, for storing data, or is operatively coupled to receive data from or transmit data to them, or both. However, a computer need not have such devices. Furthermore, a computer may be incorporated in other devices, 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, such as a universal serial bus (USB) flash drive, to name just a few.

[0155] Computer-readable media suitable for storing computer program instructions and data include all forms of non-volatile memory, media, and memory devices, including, by way of example, semiconductor memory devices 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.

[0156] To provide for user interaction, embodiments of the subject matter described herein may be implemented in a computer having a display device, such as a CRT (cathode ray tube) or LCD (liquid crystal display) monitor, for displaying information to the user, as well as a keyboard and pointing device, such as a mouse or trackball, through which the user can provide input to the computer. Other types of devices may also be used to interact with the user. For example, feedback provided to the user may be any form of sensory feedback, such as visual feedback, auditory feedback, or tactile feedback, and input from the user may be received in any form, including acoustic, speech, or tactile input. Furthermore, a computer may interact with a user by sending and receiving documents to a device used by the user, for example, by sending a web page to a web browser on the user's device in response to a request received from the web browser. A computer may also interact with a user by sending text messages or other types of messages to a personal device, such as a smartphone running a messaging application, and then receiving a reply message from the user.

[0157] A data processing device for implementing machine learning models may also include dedicated hardware accelerator units, for example for handling the common and computationally intensive part of machine learning training or production, i.e., inference, workloads.

[0158] The machine learning model may be implemented and deployed using a machine learning framework, for example, the TensorFlow framework or the JAX framework.

[0159] Embodiments of the subject matter described herein may be implemented in a computing system that includes a back-end component, e.g., a data server, or includes a middleware component, e.g., an application server, or includes a front-end component, e.g., a client computer having a graphical user interface, a web browser, or an app through which a user can interact with an embodiment of the subject matter described herein, or includes any combination of one or more such back-end, middleware, or front-end components. The components of the system can be interconnected by any form or medium of digital data communication (e.g., a communication network). Examples of communication networks include local area networks (LANs) and wide area networks (WANs), e.g., the Internet.

[0160] A computing system may include clients and servers. Clients and servers are generally remote from each other and typically interact through a communication network. The relationship of client and server arises by virtue of computer programs running on the respective computers and having a client-server relationship to each other. In some embodiments, a server sends data, e.g., HTML pages, to a user device for the purpose of displaying the data to and receiving user input from a user interacting with the device acting as a client. Data generated at the user device, e.g., results of user interaction, can be received from the device at the server.

[0161] While the specification contains many specific implementation details, these should not be construed as limiting the scope of any invention or the scope of patentable subject matter, but rather as descriptions of features that may be unique to particular embodiments of a particular invention. Certain features described herein as separate embodiments may also be implemented in combination in a single embodiment. Conversely, various features described in a single embodiment may also be implemented in multiple embodiments, either separately or in any suitable subcombination. Furthermore, even if features may be described above as functioning in a particular combination and originally claimed as such, one or more features from a claimed combination may, in some cases, be deleted from the combination, and the claimed combination may be directed to subcombinations or variations of the subcombination.

[0162] Similarly, while operations are illustrated in the figures and described in the claims in a particular order, this should not be understood as requiring that such operations be performed in the particular order or sequential order shown, or that all of the operations shown 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.

[0163] 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 produce desirable results. As an example, the processes depicted in the accompanying figures do not necessarily require the particular order shown or sequential order to achieve desirable results. In some cases, multitasking and parallel processing may be advantageous.

Claims

1. 1. A computer-implemented method for generating, at each of a plurality of output locations, an output sequence including a respective discrete token selected from a vocabulary of discrete tokens, the method comprising: initializing a continuous representation of the output sequence by assigning a respective embedding to each of one or more of the plurality of output positions, each embedding selected from a vocabulary of embeddings including a respective embedding corresponding to each discrete token in the vocabulary of discrete tokens; In each of the plurality of despreading steps, obtaining a current continuous representation of the output sequence; processing a diffusion model input including the current continuous representation using a diffusion model to generate a diffusion model output, the diffusion model output including, for each of the plurality of output positions, a respective initial score for each of a plurality of embeddings in the vocabulary of embeddings; processing the respective initial scores using a softmax function to generate, for each of the plurality of output positions, a probability distribution over the plurality of embeddings in the vocabulary of embeddings, the probability distribution including a respective probability score for each embedding in the vocabulary of embeddings; updating the current continuous representation of the output sequence using the probability distributions and the vocabulary of the embeddings.

2. updating the current continuous representation of the output sequence using the probability distributions and the vocabulary of embeddings; determining a score function estimate for the output sequence using the respective probability scores and the vocabulary of the embeddings; and updating the current continuous representation of the output sequence by applying an ordinary differential equation (ODE) solver to the score function estimate.

3. updating the current continuous representation of the output sequence by applying the ODE solver to the score function estimate; 3. The method of claim 2, comprising using the ODE solver to evaluate the score function estimate to determine a respective updated embedding for at least one of the subset of output positions in the current continuous representation of the output sequence.

4. determining the score function estimate 4. The method of claim 2, comprising using the probability distributions and the vocabulary of embeddings to determine a respective interpolated embedding for each of the subset of output locations.

5. Determining the score function estimate for the output sequence at each of the plurality of output locations includes: The method of claim 4 , comprising calculating a difference between the respective interpolated embedding of the output position and a current embedding of the output position in the current continuous representation of the output sequence.

6. processing the diffusion model input including the current continuous representation using the diffusion model; The method of any one of claims 1 to 5, comprising applying a normalization to the current embedding in the current continuous representation.

7. The method further includes receiving a conditioning input sequence including one or more conditioning tokens, and initializing the output sequence includes: selecting an embedding at one or more given ones of the output positions from the embedding vocabulary based on the one or more conditioning tokens; and generating, with at least some degree of randomness, an initial embedding at the remaining output positions of the plurality of output positions.

8. The method comprises: receiving a conditioning input; processing the conditioned input using an encoder neural network to generate an encoder network output; and processing the diffusion model input, including the current continuous representation, using the diffusion model to generate the diffusion model output.

7. The method of claim 1, comprising processing the diffusion model input using the diffusion model while the diffusion model is conditioned to the encoder network output to produce the diffusion model output.

9. The diffusion model input is:

8. The method of claim 7, comprising: (i) the current continuous representation; (ii) the embedding at the one or more given output positions selected based on the one or more conditioning tokens; and (iii) a sequence of Boolean conditioning masks that identify which positions among the plurality of output positions are the one or more given output positions.

10. 10. The method of claim 9 when also dependent on claim 4, wherein the diffusion model input further comprises the respective interpolated embeddings for the subset of the plurality of output locations.

11. A method according to any one of claims 9 to 10, wherein the diffusion model input further comprises data specifying a noise level.

12. selecting a time step of said despreading step; applying time warping to the selected time steps by evaluating a monotonic piecewise linear function over the selected time steps to determine a warped time step; The method of claim 11 , further comprising determining the noise level by: determining the noise level from the warped time steps.

13. 13. The method according to claim 1, wherein obtaining the current continuous representation of the output sequence comprises using a current continuous representation updated at the time of a previous despreading step as the current continuous representation of the output sequence in the despreading step.

14. The method of any one of claims 1 to 13, wherein the diffusion model is an attention-based neural network comprising one or more attention neural network layers.

15. After performing the multiple despreading steps, processing a final diffusion model input, including the current continuous representation updated as of a last dediffusion step in the plurality of dediffusion steps, using the diffusion model to generate a final diffusion model output, the final diffusion model output including, for each of the plurality of output positions, a respective final initial score for each of the plurality of embeddings in the vocabulary of embeddings; for each of the plurality of output positions, selecting an embedding from the vocabulary of embeddings using the respective final initial scores; The method of any one of claims 1 to 14, further comprising: generating the output sequence comprising discrete tokens corresponding to the selected embedding.

16. After performing the multiple despreading steps, selecting, for each of the plurality of output positions, an embedding from the vocabulary of embeddings that has the closest distance to an embedding in the current continuous representation at the output position updated as of a last despreading step in the plurality of despreading steps; The method of any one of claims 1 to 14, further comprising: generating the output sequence comprising discrete tokens corresponding to the selected embedding.

17. 17. The method of any one of claims 1 to 16, wherein the diffusion model is trained by optimizing a cross-entropy loss between the probability distribution over the multiple embeddings in the vocabulary of embeddings generated using the diffusion model output and a true probability distribution over the multiple embeddings in the vocabulary of embeddings.

18. 18. The method of claim 17, wherein the embeddings in the embedding vocabulary are learned jointly with training the diffusion model by backpropagating gradients of the cross-entropy loss computed with respect to the parameters of the diffusion model to the embeddings.

19. The method of any preceding claim, wherein the discrete tokens comprise text characters, symbols, or signals.

20. 20. The method of any one of claims 1 to 19, wherein the diffusion model is configured to perform one or more of the following tasks: language modeling tasks, including prompt completion and / or gap completion tasks, or machine translation tasks, speech transcription tasks, video transcription tasks, image captioning tasks, or agent control tasks.

21. 21. A system comprising one or more computers and one or more storage devices storing instructions, the instructions being operable, when executed by the one or more computers, to cause the one or more computers to perform the operations of the respective methods of any one of claims 1 to 20.

22. A computer storage medium encoded with instructions that, when executed by one or more computers, cause the one or more computers to perform the operations of the respective methods recited in any one of claims 1 to 21.

Images