Discrete lexical processing using diffusion model

Abstract

Methods, systems, and apparatus, including computer programs encoded on a computer storage medium, for generating an output sequence of discrete lexical elements using a diffusion model. In one aspect, a method includes generating a final potential representation of a sequence of discrete lexical elements by using a diffusion model, the final potential representation including a determined value for each of a plurality of potential variables; applying a de-embedding matrix to a final potential representation of the output sequence of discrete lexical elements to generate a de-embedded final potential representation, the de-embedded final potential representation comprising, for each of the plurality of potential variables, a respective digital score for each discrete lexical element in the word list of the plurality of discrete lexical elements; for each of the plurality of potential variables, selecting a discrete lexical element with the highest number score from a plurality of discrete lexical elements in the word list; and generating an output sequence of discrete lexical elements comprising the selected discrete lexical elements.

Description

Cross - Reference to Related Applications

[0001] This application claims the priority of U.S. Provisional Application No. 63 / 411,045, filed on September 28, 2022. The disclosure of the prior application is considered part of the disclosure of this application and is incorporated herein by reference. Background Art

[0002] This specification relates to processing data using neural networks.

[0003] A neural network is a machine - learning model that uses one or more layers of non - linear units to predict an output for a received input. In addition to the output layer, some neural networks also include one or more hidden layers. The output of each hidden layer is used as an 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 based on the current values of a corresponding set of parameters. Summary of the Invention

[0004] According to a first aspect, there is provided a computer - implemented method of using a diffusion model to generate an output sequence of discrete tokens. The method includes generating a final latent representation of a sequence of discrete tokens by using a diffusion model, the final latent representation including a determined value for each of a plurality of latent variables. Generating the final latent representation includes, at each of a plurality of reverse diffusion time steps: processing a diffusion - model input that includes an intermediate latent representation of the sequence of discrete tokens for the reverse diffusion time step to generate an estimate of the sequence of discrete tokens up to the reverse diffusion time step; using the estimate to define a distribution over a continuous space of possible values for each of the plurality of latent variables; and generating an updated intermediate latent representation of the sequence of discrete tokens for the reverse diffusion time step by sampling from the distribution. The method further includes applying a de - embedding matrix with learned values to the final latent representation of the output sequence of discrete tokens to generate a de - embedded final latent representation, the de - embedded final latent representation including, for each of the plurality of latent variables, a corresponding numerical score for each of a plurality of discrete tokens in a token vocabulary. The method further includes: for each of the plurality of latent variables, selecting the discrete token having the highest numerical score from the plurality of discrete tokens in the token vocabulary; and generating an output sequence of discrete tokens including the selected discrete tokens.

[0005] For the first reverse diffusion time step, the intermediate latent representation may be an initial latent representation. For subsequent reverse diffusion time steps, the intermediate latent representation may be the updated intermediate latent representation generated in the immediately preceding reverse diffusion time step.

[0006] In an implementation, an estimate of a sequence of discrete tokens up to a reverse diffusion time step is an estimate of a final latent representation of an output sequence of discrete tokens up to the reverse diffusion time step.

[0007] In an implementation, discrete tokens include text, symbols, or signals.

[0008] In an implementation, the diffusion model input further includes an estimate of a sequence of discrete tokens generated up to the previous reverse diffusion time step. In an implementation, an estimate of a sequence of discrete tokens up to the previous reverse diffusion time step is an estimate of a final latent representation of an output sequence of discrete tokens up to the previous reverse diffusion time step.

[0009] In an implementation, using a diffusion model to generate an output sequence of discrete tokens includes generating unconditional discrete tokens.

[0010] In an implementation, using a diffusion model to generate an output sequence of discrete tokens includes generating discrete tokens conditioned on an input sequence of discrete tokens, and wherein the method includes:

[0011] Receiving an input sequence of discrete tokens;

[0012] Converting each discrete token in the input sequence into a one-hot vector; and

[0013] Applying an embedding matrix with pre-trained values to each one-hot vector to embed the one-hot vector into a continuous vector.

[0014] In an implementation, using a diffusion model to generate an output sequence of discrete tokens includes generating discrete tokens conditioned on an input sequence of discrete tokens, and wherein the method includes:

[0015] Receiving an input sequence of discrete tokens;

[0016] Converting each discrete token in the input sequence into a one-hot vector; and

[0017] Applying an embedding matrix to each one-hot vector to embed the one-hot vector into an embedding.

[0018] In an implementation, the output sequence of discrete tokens further includes the input sequence of discrete tokens received by the diffusion model.

[0019] In an implementation, the input sequence of discrete tokens includes discrete tokens representing an audio data input including spoken words, e.g., a waveform characterizing the audio in the time domain or time-frequency domain.

[0020] In an implementation, the output token represents audio data including spoken words, e.g., a waveform characterizing the audio in the time domain or time-frequency domain.

[0021] In an implementation, the input sequence is generated from an audio signal that includes spoken utterances. In an implementation, the input can be received in the form of an audio (speech) signal, which is converted by a speech-to-text converter to form the input sequence.

[0022] In an implementation, the output sequence of discrete tokens represents text. In an implementation, the output sequence of discrete tokens is converted by a text-to-speech converter to form an audio signal.

[0023] In an implementation, the input sequence of discrete tokens represents a sequence of actions to be performed by an agent, such as a mechanical agent in a real-world environment that implements actions to perform mechanical tasks.

[0024] In an implementation, the output sequence of discrete tokens represents a sequence of actions to be performed by an agent, such as a mechanical agent in a real-world environment that implements actions to perform mechanical tasks.

[0025] In an implementation, the method further includes:

[0026] Applying a linear projection to a continuous vector to generate a projected continuous vector; and

[0027] Using a diffusion model to process the projected continuous vector.

[0028] In an implementation, generating discrete tokens conditioned on an input sequence of discrete tokens includes using classifier-free guidance techniques.

[0029] In an implementation, the de-embedding matrix has been learned during the training of the diffusion model, while the embedding matrix is fixed during the training of the diffusion model.

[0030] In an implementation, the method further includes training the diffusion model on unlabeled discrete token data that includes discrete token inputs to minimize the mean squared error between each discrete token input and an estimate of the discrete token input generated by the diffusion model as of the sampled reverse diffusion time steps.

[0031] In an implementation, the training also minimizes the cross-entropy loss evaluated on the final latent representation of the sequence of discrete tokens generated by the diffusion model from the discrete token inputs.

[0032] In an implementation, the training includes:

[0033] Applying a binary mask to the sequence of discrete token inputs, the binary mask including one or more first masks that define conditional tokens in the sequence and one or more second masks that define padding tokens in the sequence; and

[0034] A diffusion model is trained on a sequence of masked discrete token inputs to generate an estimate of a sequence of discrete token inputs, the sequence of discrete token inputs including padding tokens in the sequence.

[0035] In an implementation, the training includes learning the values of a de-embedding matrix while keeping the pre-trained values of the embedding matrix fixed.

[0036] According to another aspect, a system is provided that includes: one or more computers; and one or more storage devices that store instructions that, when executed by the one or more computers, are operable to cause the one or more computers to perform the operations of any of the above methods.

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

[0038] This specification describes a system implemented as a computer program on one or more computers in one or more locations that uses a diffusion model to generate an output sequence of discrete tokens, the diffusion model performing a reverse diffusion process on continuous embeddings in an embedding space.

[0039] In some cases, the generation process can be unconditional, where the output sequence can be generated by the diffusion model from random noise. In other cases, the generation process can be conditional, where, for example, the output text, computer program code, symbol, or signal sequence generated by the diffusion model is a completion or extension of the input text, computer program code, symbol, or signal sequence.

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

[0041] The diffusion model described in this specification is a model that performs a diffusion process in a continuous latent space but is capable of handling data that is inherently discrete, such as text data, i.e., receiving discrete input data and / or generating discrete output data. By performing a diffusion process on continuous embeddings in the latent space and then converting the final latent representation back to the discrete space using continuous turn steps that use a learnable de-embedding matrix, the described diffusion model is both flexible - for example, can be easily configured for both conditional and unconditional discrete data generation - and scalable - for example, can be extended to arbitrarily large models with 100 million, 500 million, or more parameters - and can thus better meet the needs of various generation tasks. The described diffusion model achieves or even exceeds state-of-the-art performance on a wide range of conditional and unconditional text generation tasks implemented by many existing autoregressive models.

[0042] Advantageously, unlike most existing autoregressive models that predict output tokens one by one in a way that earlier tokens do not provide information to later tokens, the described diffusion model can predict all tokens in a sequence at once. This allows for bidirectional (e.g., non-causal) attention; thus, the selection of a token later in the sequence can influence the selection of earlier tokens to achieve higher-quality data generation. This also makes the inference process significantly more parallelizable, i.e., it is possible to generate any amount of data within a fixed time budget by executing multiple diffusion processes in parallel with each other, since the generation of all tokens can occur concurrently with the diffusion model, rather than sequentially as required by autoregressive models.

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

[0044] Figure 1 is a block diagram of an example data generation system.

[0045] Figure 2 is a flowchart of an example process for generating a discrete output sequence.

[0046] Figure 3 is Figure 2 a flowchart of an example process for a sub-step of one of the steps of the process of

[0047] Like reference numerals and names in the various figures indicate like elements. DETAILED DESCRIPTION

[0048] Figure 1 is a diagram of an example data generation system 100. The data generation system 100 is an example of a system implemented as a computer program on one or more computers in one or more locations, where the systems, components, and techniques described below may be implemented.

[0049] The data generation system 100 is a system that generates a discrete output sequence 162 in response to a received request. The data generation system 100 may store the discrete output sequence 162 in an output data repository or provide the discrete output sequence 162 for some other immediate purpose. For example, the data generation system 100 may then output the discrete output sequence 162 for presentation, e.g., on a client device that submitted the data generation request.

[0050] The discrete output sequence 162 includes a plurality of discrete tokens. Each discrete token can be an individual, 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 numerical format used by a computer system. In contrast, discrete data has additional constraints that further limit the possible values that the data can take on within the specified range so as not to exceed those possible values in the numerical format used by the computer system.

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

[0052] For example, the discrete output sequence 162 can include, at each of a plurality of positions, a corresponding token from a vocabulary of discrete tokens. The vocabulary of tokens can include any of a variety of discrete tokens representing text, symbols, or signals. For example, the vocabulary of discrete tokens can include one or more of characters, sub-words, words, punctuation marks, numbers, or other symbols that occur in a corpus of text. For example, the text can be natural language text or computer program code. In some examples, the output tokens represent audio data including spoken words, e.g., a waveform that characterizes audio in the time domain or the time-frequency domain. In some examples, the discrete token output sequence represents text. In one example, the discrete token output sequence is converted by a text-to-speech converter to form an audio signal. In some examples, the discrete token output sequence represents a sequence of actions to be performed by an agent, such as a robotic agent in a real-world environment that implements actions to perform mechanical tasks.

[0053] In some cases, the data generation system 100 can be configured as an unconditional output sequence generation system that generates an unconditional (i.e., not conditioned on any conditional input) discrete output sequence 162. In these implementations, the discrete output sequence 162 generated by the system 100 is a sample of the distribution of the training input sequences used during the training of the system 100.

[0054] In other cases, the data generation system 100 can be configured as a conditional output sequence generation system that generates a discrete output sequence 162 conditioned on the input sequence 102. In some of these cases, the input sequence 102 includes a discrete input sequence. Like the discrete output sequence 162, the discrete input sequence can include, at each of a plurality of positions, a corresponding token from a vocabulary of tokens. The vocabulary of tokens can include any of a variety of discrete tokens representing text, symbols, or signals. In other of these cases, the input sequence 102 includes a continuous input sequence, e.g., the continuous input sequence represents a sequence of pixels of an image, or represents a waveform or spectrogram of audio.

[0055] For example, when configured as a conditional output sequence generation system, the data generation system 100 may receive an input sequence 102 that includes text data, additional data other than text (e.g., embeddings of data of different data types (e.g., images, videos, speech, etc.)), or both, and generate a discrete output sequence 162 that is a text sequence, e.g., a completion of the input sequence 102, an extension of the input sequence 102, a response to a question posed in the input sequence 102, a text sequence about a topic specified by the input sequence 102, a text description of data of different data types, etc.

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

[0057] In both cases, it should be understood that text (both natural language text and computer program code), symbols, and signals are only understood as examples of discrete data for illustration purposes, and the discrete tokens may be discrete data items in many other formats or modalities. For example, the discrete output sequence 162 may be or include a biological data sequence (e.g., a gene expression sequence), an electronic health record data sequence (e.g., a sequence of health events), a clinical procedure data sequence (e.g., a sequence of medical orders, clinical documents, notes, diagnostic codes, medications, etc.).

[0058] In one example, the input may be received in the form of an audio (speech) signal captured by a microphone, which is converted by a speech-to-text converter to form the input sequence 102. In some examples, the discrete token input sequence includes tokens representing audio data inputs that include spoken words, e.g., waveforms that characterize the audio in the time domain or time-frequency domain. In some examples, the discrete token input sequence represents a sequence of actions to be performed by an agent, such as a mechanical agent in a real-world environment that implements actions to perform mechanical tasks.

[0059] To generate the discrete output sequence 162, the data generation system 100 uses an initialization engine 120 to initialize a latent representation of the discrete output sequence 162, i.e., generate an initial latent representation 122, and then uses a diffusion model neural network 130 (or simply "diffusion model 130") and an update engine 140 to generate a final latent representation 136 of the discrete output sequence 162 based on performing a reverse diffusion process to update the initial latent representation 122 over multiple reverse diffusion time steps. Thus, the final latent representation 136 is generated in the last reverse diffusion time step of the reverse diffusion process.

[0060] The initial latent representation 122 has the same dimension as the final latent representation 136, but has different values. That is, the initial latent representation 122 includes a plurality of latent variables, and the final latent representation 136 includes the same number of latent variables, but the values of these latent variables will generally be different between the two representations 122 and 136.

[0061] When configured as an unconditional output sequence generation system, the initialization engine 120 generates the initial latent representation 122 by sampling an initial value for each of the plurality of latent variables included in the initial latent representation 122 from a corresponding noise distribution (e.g., a Gaussian distribution or another predetermined distribution). Thus, the initial latent representation 122 includes a plurality of latent variables, where the initial value of each latent variable is sampled from a corresponding noise distribution.

[0062] When configured as a conditional output sequence generation system, the preprocessing engine 110 (which is an optional component of the data generation system 100) processes the received input sequence 102 to generate a conditional embedding 116 that includes a numerical sequence. Then, based on using the numerical values included in the conditional embedding 116 as the initial values of some of the latent variables included in the initial latent representation 122 and based on sampling an initial value for each of the remaining latent variables from a corresponding noise distribution (e.g., a Gaussian distribution or another predetermined distribution), the initialization engine 120 generates the initial latent representation 122. Thus, the initial latent representation 122 includes a plurality of latent variables, where the initial value of each of some of the latent variables is determined from the input sequence 102 and the initial value of each of the other latent variables is sampled from a corresponding noise distribution.

[0063] More specifically, the preprocessing engine 110 includes an embedding matrix 112 that defines a mapping from a discrete space to an embedding space. To define this mapping, the embedding matrix 112 can be a matrix having dimensions i.e., can have predefined values included as entries of the matrix, where D is the embedding size (the size of each consecutive vector in the embedding space) and V is the vocabulary size (the total number of discrete tokens included in the token vocabulary). Thus, in one example, each token ω in the vocabulary has an associated embedding e ω ϵ ℝ D with a fixed norm Match the norm of the random Gaussian samples in dimension D used for noise cleaning of the data. The embedding values can be generated during a pre-training phase. In one example, the embedding matrix is the matrix of all embeddings of the vocabulary. The embedding of an entity (e.g., a word) can refer to the representation of that entity as an ordered set of numerical values, such as a numerical vector. Embeddings of entities can be generated, for example, as the output of a neural network processing data representing the entity, or as the result of some other encoding process.

[0064] The preprocessing engine 110 converts each token included in the input sequence 102 into a discrete one-hot vector and applies the embedding matrix 112 to each discrete one-hot vector, for example, based on an element-wise multiplication operation, to map the discrete one-hot vector into a continuous vector having numerical values for the token in the embedding space. Thus, in some implementations, the conditional embedding 116 consists of the numerical values included in the continuous vectors that have been generated as a result of applying the embedding matrix 112 to the input sequence 102. In other implementations, the preprocessing engine 110 additionally applies a linear projection to each continuous vector to generate a projected continuous vector and then generates the conditional embedding 116 from the numerical values included in the projected continuous vector. The conversion to tokens can be performed using a SentencePiece tokenizer containing 32,000 words, as described in Kudo et. al., “Sentencepiece: A simple and language independent subword tokenizer and detokenizer for neural text processing (Sentencepiece: A simple and language independent subword tokenizer and detokenizer for neural text processing)”, In EMNLP (Demonstration), pp. 66–71, Association for Computational Linguistics, 2018, the entire content of which is incorporated herein by reference.

[0065] The diffusion model 130 can have any suitable architecture that allows the diffusion model to process, at any given reverse diffusion time step, the diffusion model input for that reverse diffusion time step (the diffusion model input includes an intermediate latent representation (as of that reverse diffusion time step) having the same dimension as the discrete output sequence 162) to generate the diffusion model output 132 for the given reverse diffusion time step. The diffusion model output 132 includes an estimate of the discrete output sequence 162 (as of the reverse diffusion time step), which also has the same dimension as the discrete output sequence 162.

[0066] For example, the diffusion model 130 can be a convolutional neural network including one or more convolutional residual blocks, such as, for example, a U-Net or other architectures. As another example, the diffusion model 130 can be an attention neural network including one or more attention blocks (e.g., self-attention blocks, gated attention blocks, or cross-attention blocks). As yet another example, the diffusion model 130 can include both convolutional residual blocks and attention blocks.

[0067] At a given reverse diffusion time step, the data generation system 100 then uses the update engine 140 to update the intermediate latent representation based on the estimate of the discrete output sequence 162 for the given reverse diffusion time step. The intermediate latent representation updated by the update engine 140 at a given reverse diffusion time step will be referred to in this specification as the updated intermediate latent representation 134; the intermediate latent representation updated at the last reverse diffusion time step in the reverse diffusion process will be referred to in this specification as the final intermediate latent representation 136.

[0068] Thus, at any given reverse diffusion time step, the intermediate latent representation provided as part of the diffusion model input to the diffusion model 130 will be the updated intermediate latent representation 134 generated at the immediately preceding reverse diffusion time step. For the first reverse diffusion time step, the intermediate latent representation is the initial latent representation 122 generated by the initialization engine 120.

[0069] More specifically, at any given reverse diffusion time step, the update engine 140 at least uses the estimate of the discrete output sequence 162 for the given reverse diffusion time step to define a corresponding distribution over a continuous space of possible values for each of the plurality of latent variables included in the initial latent representation 122. The update engine 140 then generates the updated intermediate latent representation 134 of the discrete output sequence 162 for the reverse diffusion time step by selecting a value for each latent variable from the corresponding distribution (e.g., by selecting the value with the highest probability or by sampling from the distribution). Thus, the updated latent representation 134 includes a plurality of latent variables, where the value of each latent variable is sampled from the corresponding distribution defined by the diffusion model output 132 of the diffusion model 130.

[0070] To define such a distribution for each latent variable, the update engine 140 can calculate the mean of the distribution, and optionally the variance of the distribution, based on the estimate of the discrete output sequence 162 and based on the intermediate latent representation before being updated at the given reverse diffusion time step. The update step will be further described below with reference to Figure 3 The update step is further described.

[0071] Some implementations of the data generation system 100 utilize self - conditional inputs 138 when performing the reverse diffusion process. When the data generation system 100 is configured as a conditional output sequence generation system, using self - conditioning can improve the quality of the discrete output sequence 162.

[0072] At any given reverse diffusion time step, the self - conditional input 138 can be an estimate of the discrete output sequence 162 included in the diffusion model output generated at the immediately preceding reverse diffusion time step. Thus, for any given reverse diffusion time step, the diffusion model 130 can process a diffusion model input that includes: (i) an intermediate latent representation (up to the reverse diffusion time step); and (ii) an estimate of the discrete output sequence 162 generated at the immediately preceding reverse diffusion time step, to generate a diffusion model output 132 for the given reverse diffusion time step.

[0073] Some implementations of the data generation system 100 utilize guidance when performing the reverse diffusion process. That is, the reverse diffusion process is sometimes a guided reverse diffusion process. To alleviate the need for a separately trained guidance neural network, the guidance can be classifier - free guidance.

[0074] Using classifier - free guidance can involve the diffusion model 130 processing a conditional token sequence that is part of the diffusion model input at any given reverse diffusion time step to generate a diffusion model output. The conditional tokens can be defined by a binary conditional mask that is set to one for positions in the discrete output sequence that each have a predetermined discrete token (“conditional positions”), and set to zero for positions in the discrete output sequence that each have an unknown discrete token (“padding positions”) (i.e., positions where the system needs to generate a discrete token).

[0075] After the last reverse diffusion time step, the data generation system 100 provides a final latent representation 136 that is in the same embedding space as the initial latent representation 122 to the post - processing engine 150. The final latent representation 136 includes a plurality of latent variables. Each latent variable has a determined value that is determined as a result of the reverse diffusion process and may be different from its initial value. The post - processing engine 150 processes the final latent representation 136 to generate a discrete output sequence 162 that is in the same discrete space as the input sequence 102.

[0076] To this end, the post - processing engine 150 includes a de - embedding matrix 152 that defines a mapping from the embedding space to the discrete space. To define this mapping, the de - embedding matrix 152 can be a matrix with dimensions i.e., can have a learned value, where V is the vocabulary size (the total number of discrete tokens included in the token vocabulary) and D is the embedding size (the size of each continuous vector in the embedding space).

[0077] The post-processing engine 150 applies the de-embedding matrix 152 to each latent variable included in the final latent representation 136, for example, based on an element-wise multiplication operation, to generate a de-embedded final latent representation. The de-embedded final latent representation includes, for each latent variable included in the final latent representation 136, a distribution corresponding to that latent variable over the token vocabulary. For each latent variable, the corresponding distribution includes a respective numerical score for each token included in the vocabulary.

[0078] Then, the post-processing engine 150 generates a discrete output sequence 162 by selecting a token from the vocabulary for each latent variable according to the corresponding distribution (e.g., by selecting the token with the highest numerical score or by sampling from the corresponding distribution).

[0079] Because the computations involved in the generation of each discrete output sequence are parallelizable over the entire length of the discrete output sequence, the data generation system 100 better utilizes the computational resources of a hardware accelerator on which the generation system 100 may be implemented.

[0080] A hardware accelerator is a computing device with specialized hardware configured to perform specialized computations, including, for example, parallel computations. Examples of accelerators include graphics processing units (“GPUs”), field-programmable gate arrays (“FPGAs”), and application-specific integrated circuits (“ASICs”), including tensor processing units (“TPUs”).

[0081] The training of the diffusion model 130 and other trainable components of the data generation system 100 to determine the trained parameter values of these components will be further described below.

[0082] Figure 2 is a flowchart of an example process 200 for generating a discrete output sequence. For convenience, process 200 will be described as being performed by a system of one or more computers located in one or more locations. For example, a data generation system (e.g., Figure 1 the data generation system 100 depicted in

[0083] The system performs a reverse diffusion process on an initial latent representation of the discrete output sequence to generate a final latent representation of the discrete output sequence (step 202). The discrete output sequence includes a respective token from the token vocabulary at each of a plurality of positions. The token vocabulary may include any of a variety of discrete tokens representing text, symbols, or signals.

[0084] The initial latent representation has the same dimension as the final latent representation, but has different values. That is, the initial latent representation includes a plurality of latent variables, and the final latent representation includes the same number of latent variables, but the initial values of these latent variables included in the initial latent representation and the determined values of these latent variables included in the final latent representation will generally be different from each other.

[0085] When configured as an unconditional output sequence generation system, the system can generate an initial latent representation by sampling initial values for each of the plurality of latent variables from a corresponding noise distribution (e.g., a Gaussian distribution or another predetermined distribution).

[0086] When configured as a conditional output sequence generation system, the system can receive an input sequence and then map the received input sequence into a conditional embedding including a numerical sequence. The mapping can be performed by applying an embedding matrix (where D is the embedding size and V is the vocabulary size) to one-hot vectors corresponding to each token included in the input sequence. Thus, based on using the numerical values included in the conditional embedding as the initial values of some of the latent variables included in the initial latent representation and based on sampling initial values for each of the remaining latent variables from a corresponding noise distribution, the system can generate an initial latent representation.

[0087] Then, the system generates a final latent representation by updating the initial latent representation over a plurality of reverse diffusion time steps using a diffusion model. In other words, the final latent representation is the updated intermediate latent representation generated in the last reverse diffusion time step.

[0088] The following explains in reference to Figure 3 more details the updating of the initial latent representation. The figure shows a flowchart of an example process 300 of sub-steps 302 to 306 of step 202. Generally, the system can repeatedly execute process 300 at each of the plurality of reverse diffusion time steps in the reverse diffusion process.

[0089] The system processes the diffusion model input for the reverse diffusion time step to generate a diffusion model output (step 302). The diffusion model output can include an estimate of the discrete output sequence up to the reverse diffusion time step. On the other hand, the diffusion model input can include: (i) an intermediate latent representation of the discrete output sequence (up to the reverse diffusion time step) ; and (ii) a time index t indicating the current reverse diffusion time step. In one example, the diffusion model is a trained neural network that takes as input an intermediate latent representation of the discrete output sequence (up to the reverse diffusion time step) takes as input the time index t indicating the current reverse diffusion time step and outputs an estimate of the final latent representation of the discrete output sequence up to that reverse diffusion time step .

[0090] For the first reverse diffusion time step, the intermediate latent representation is the initial latent representation. For any subsequent reverse diffusion time step, the intermediate latent representation is the updated intermediate latent representation generated in the immediately preceding reverse diffusion time step.

[0091] For example, the estimate of the discrete output sequence can thus be defined as:

[0092] ,

[0093] where θ represents the parameters of the diffusion model.

[0094] Optionally, the diffusion model input further includes a self-conditioning input. At any given reverse diffusion time step, the self-conditioning input can be the estimate of the discrete output sequence included in the diffusion model output generated at the immediately preceding reverse diffusion time step. Thus, for any given reverse diffusion time step, the diffusion model input can additionally include the estimate of the discrete output sequence included in the diffusion model output generated at the immediately preceding reverse diffusion time step t+1 .

[0095] As another example, the estimate of the discrete output sequence can thus alternatively be defined as:

[0096] ,

[0097] where θ similarly represents the parameters of the diffusion model.

[0098] Further optionally, the diffusion model input further includes a guidance, e.g., classifier-free guidance. When using classifier-free guidance, for any given reverse diffusion time step, the diffusion model input can additionally include a fixed sequence of conditional tokens.

[0099] As yet another example, the estimate of the discrete output sequence can thus alternatively be defined as:

[0100]

[0101] where is a guidance scale, and c represents a conditional token used in classifier-free guidance. The conditional token c can be defined by a binary conditional mask which is set to one for positions in the discrete output sequence that each have a predetermined discrete token, and is set to zero for positions in the discrete output sequence that each have an unknown discrete token (i.e., positions where the system needs to generate a discrete token).

[0102] The system uses the estimated value of the discrete output sequence to define a distribution over a continuous space of possible values for each of a plurality of latent variables (step 304). To define such a distribution for each latent variable, the system may compute the mean of the distribution based on (i) the estimated value of the discrete output sequence and (ii) the intermediate latent representation included in the diffusion model input and optionally compute the variance of the distribution.

[0103] For example, the system may compute the mean of the distribution as:

[0104]

[0105] and compute the variance of the distribution as:

[0106] ,

[0107] where , and defines the variance for the current reverse diffusion time step according to a predetermined variance schedule In one example, the system uses a fixed schedule where the value of is predetermined and fixed.

[0108] The system generates an updated intermediate latent representation of the discrete output sequence for the reverse diffusion time step (step 306). The system may do this by selecting a value for each latent variable from the corresponding distribution (e.g., by sampling a value from the continuous space of possible values according to the distribution):

[0109] ,

[0110] where I is the identity matrix.

[0111] Thus, the updated latent representation includes a plurality of latent variables where the value of each latent variable is sampled from the corresponding distribution defined by the diffusion model output of the diffusion model for the current reverse diffusion time step.

[0112] Return toFigure 2 The system applies the de-embedding matrix to the final latent representation of the discrete output sequence to generate a de-embedded final latent representation (step 204). For each latent variable, the de-embedded final latent representation includes a distribution over the vocabulary of tokens corresponding to that latent variable. For each latent variable, the corresponding distribution includes respective numerical scores for each token in the vocabulary of tokens.

[0113] For example, the system may generate the de-embedded final latent representation by computing:

[0114]

[0115] where is a de-embedding matrix with vocabulary size V and embedding size D, is a one-hot representation of a token in the vocabulary, and represents the softmax probability of token k in the vocabulary of tokens with logit .

[0116] For each of the plurality of latent variables, the system selects a token from the tokens in the vocabulary according to the corresponding distribution (step 206). For each latent variable, the system may, for example, select the token with the highest numerical score from all the tokens in the vocabulary. The system generates a discrete output sequence (step 208). The discrete output sequence includes the tokens that have been selected from the vocabulary at step 206. When configured as a conditional output sequence generation system, the discrete output sequence may also optionally include the input sequence that the system receives as input.

[0117] By repeatedly performing process 200, the system may generate different discrete output sequences. That is, process 200 may be performed as part of predicting a discrete output sequence from an input sequence for which the desired output (i.e., the discrete output sequence that the system should generate from the input sequence) is unknown.

[0118] Some steps of process 200 (e.g., sub-step 302 of step 202) can also be performed as part of processing an input sequence derived from a training data set (i.e., an input derived from an input set for which the discrete output sequence to be generated by the system is known) to train the trainable components of the system to determine trained values of the parameters of these components. In one example, a diffusion model can be trained on the C4 data set as 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 content of which is incorporated herein by reference. In one example, a SentencePiece tokenizer with 32,000 words can be used to convert the training data into tokens as described in Kudo et. al., "Sentencepiece: A simple and language independent subword tokenizer and detokenizer for neural text processing", In EMNLP (Demonstration), pp. 66–71, Association for Computational Linguistics, 2018, the entire content of which is incorporated herein by reference. In one example, the diffusion model can be trained with a sequence length of 256. In one example, 10% padding tokens can be inserted uniformly (i.e., not necessarily at the end of the sequence) in the training set to allow the model to generate samples of different sizes and provide more flexibility.

[0119] Specifically, the system can repeatedly perform sub-step 302 of step 202 on input sequences selected from an unlabeled discrete training data set as part of the diffusion model training process to train the trainable components of the system to optimize an objective function suitable for the discrete data generation task that the diffusion model is configured to perform.

[0120] For example, the objective function can include a diffusion loss term that trains the diffusion model on input sequences selected from unlabeled discrete training data to map each input sequence Minimize the mean squared error between the estimate of the input sequence generated by the sampled reverse diffusion time steps within the reverse diffusion process by the diffusion model and:

[0121]

[0122] The objective function may additionally include a reconstruction loss term that trains the de - embedding matrix R to minimize the cross - entropy loss evaluated on the final latent representation of the discrete output sequence generated from the input sequence by the diffusion model:

[0123] ,

[0124] where where is a constant scaling factor of an order of magnitude similar to During training, the system may incorporate any number of techniques to improve the speed, effectiveness, or both, of the training process.

[0125] For example, the system may use span masking techniques to train the diffusion model on multiple filling tasks, such as middle filling and span filling tasks. In this example, for each input sequence selected from the unlabeled discrete training data, the system may apply a binary mask to the discrete tokens included in the input sequence. The binary mask includes one or more first masks that define the conditional tokens in the sequence and one or more second masks that define the padding tokens in the sequence. Then the system may train the diffusion model on the masked input sequence to generate estimates of the padding tokens included in the masked input sequence, i.e., to generate estimates of the original discrete tokens masked by the second mask in the input sequence.

[0126]

[0127] As another example, the system may initialize the de - embedding matrix R to the transpose of the embedding matrix E before training. The values of the embedding matrix E may be pre - trained. In the pre - training phase, a BERT model of a fixed size (e.g., 150m parameters) and a feature dimension of 896 may be trained to generate word embeddings.

[0128] As yet another example, to stabilize training and avoid unigram entropy decline, the system may specifically train the de - embedding matrix R and keep the embedding matrix E fixed during training. Thus, while the de - embedding matrix R is being learned during the training of the diffusion model, the embedding matrix E is not being learned during training.

[0129] ​In this example, "learning" means that during the training of the diffusion model, one or more values included as entries of the de-embedding matrix R will be adjusted. Conversely, "unlearning" means that predefined values included as entries of the embedding matrix E (which were determined before the training of the diffusion model (e.g., determined as a result of the training of another neural network that is learned together with the embedding matrix E)) will remain fixed throughout the training of the diffusion model.

[0130] As previously mentioned, in one example where self-conditioning is used, the x 0 estimate obtained at the previous sampling step is passed as input to the diffusion model, and this estimate is gradually refined. To approximate the inference behavior during training while maintaining computational efficiency, the first estimate is calculated with self-conditioning set to zero:

[0131] .

[0132] Then a second forward pass is performed using stop gradient on to obtain:

[0133]

[0134] Then the outputs from these two forward passes are used to optimize the diffusion model so as to accurately estimate x 0 .

[0135] This specification uses the term "configured" in connection with systems and computer program components. For a system of one or more computers that is configured to perform particular operations or actions, it means that the system has installed thereon software, firmware, hardware, or a combination thereof that in operation causes the system to perform those operations or actions. For one or more computer programs that are configured to perform particular operations or actions, it means that the one or more programs include instructions that, when executed by a data processing device, cause the device to perform the operations or actions.

[0136] Embodiments of the subject matter and the functional operations described in this specification can be implemented in digital electronic circuitry, in tangibly embodied computer software or firmware, in computer hardware, including the structures disclosed in this specification and their structural equivalents, or in combinations of one or more of them. Embodiments of the subject matter described in this specification can be implemented as one or more computer programs, i.e., one or more modules of computer program instructions encoded on a tangible non-transitory storage medium for execution by, or to control the operation of, data processing apparatus. A computer storage medium may be a machine-readable storage device, a machine-readable storage substrate, a random or serial access memory device, or a combination of one or more of them. Alternatively or additionally, the program instructions may be encoded on an artificially generated propagated signal, e.g., a machine-generated electrical, optical, or electromagnetic signal, that is generated to encode information for transmission to the appropriate receiver apparatus for execution by the data processing apparatus.

[0137] The term “data processing apparatus” refers to data processing hardware and includes all kinds of devices, apparatus, and machines for processing data, e.g., including programmable processors, computers, or multiple processors or computers. The apparatus may also be, or further include, special purpose logic circuitry, e.g., an FPGA (Field Programmable Gate Array) or an ASIC (Application Specific Integrated Circuit). In addition to hardware, the apparatus may optionally include code that creates an execution environment for the computer program, e.g., code that constitutes processor firmware, a protocol stack, a database management system, an operating system, or a combination of one or more of them.

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

[0139] In this specification, the term "database" is used broadly to refer to any collection of data: the data need not be structured in any particular way, or at all, and can be stored on a storage device in one or more locations. Thus, for example, an index database can include multiple collections of data, each of which can be organized and accessed differently.

[0140] 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 can be installed and run on the same one or more computers.

[0141] The processes and logical flows described in this specification can be performed by one or more programmable computers that execute one or more computer programs to perform functions by operating on input data and generating output. The processes and logical flows can also be performed by, for example, special logic circuitry such as an FPGA or ASIC, or by a combination of special logic circuitry and one or more programmed computers.

[0142] Computers suitable for executing a computer program can be based on a general or special purpose microprocessor or both, or any other kind of central processing unit. Typically, the central processing unit will receive instructions and data from a read only memory or a random access memory or both. The basic elements of a computer are a central processing unit for performing or executing instructions and one or more memory devices for storing the instructions and data. The central processing unit and the memory can be supplemented by, or incorporated in, special logic circuitry. Typically, a computer will also include one or more mass storage devices for storing data, such as magnetic disks, magneto-optical disks, or optical disks, or operatively coupled to receive data from one or more mass storage devices or to transfer data to one or more mass storage devices or both. However, a computer need not have such devices. In addition, 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, such as a universal serial bus (USB) flash drive, to name but a few.

[0143] Computer-readable media suitable for storing computer program instructions and data include all forms of non-volatile memory, media, and memory devices, by way of example including semiconductor memory devices (e.g., EPROM, EEPROM, and flash memory devices); magnetic disks (e.g., internal hard disks or removable disks); magneto-optical disks; and CD-ROM and DVD-ROM disks.

[0144] To provide for interaction with a user, embodiments of the subject matter described in this specification can be implemented on a computer having a display device (e.g., a CRT (cathode ray tube) or LCD (liquid crystal display) monitor) for displaying information to the user and a keyboard and a pointing device (e.g., a mouse or a trackball) by which the user can provide input to the computer. Other kinds of devices can also be used to provide for 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, speech, 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, by sending web pages to a web browser on the user device in response to a request received from the web browser. Further, the computer can interact with the user by transmitting text messages or other forms of messages to a personal device (e.g., a smart phone running a messaging application) and receiving responsive messages in return from the user.

[0145] The data processing device for implementing a machine learning model may also include, for example, a dedicated hardware accelerator unit for processing common and computationally intensive portions of machine learning training or production (i.e., inference, workload).

[0146] A machine learning framework (e.g., TensorFlow framework or JAX framework) can be used to implement and deploy a machine learning model.

[0147] Embodiments of the subject matter described in this specification can be implemented in a computing system that includes a backend component (e.g., as a data server), or includes a middleware component (e.g., an application server), or includes a frontend component (e.g., a client computer having a graphical user interface, a web browser, or an app through which a user can interact with an implementation of the subject matter described in this specification), or any combination of one or more such backend, middleware, or frontend components. The components of the system can be interconnected by any form or medium of digital data communication (e.g., a communication network). Examples of communication networks include a local area network (LAN) and a wide area network (WAN), such as the Internet.

[0148] A computing system may include a client and a server. The client and the server are typically far apart from each other and typically interact via a communication network. The relationship between the client and the server is created by computer programs that run on respective computers and have a client-server relationship with each other. In some embodiments, the server sends data (e.g., an HTML page) to a user device, e.g., for displaying data to and receiving user input from a user interacting with the device acting as the client. Data generated at the user device may be received at the server, e.g., due to user interaction.

[0149] Although this specification contains many specific implementation details, these should not be construed as limitations on the scope of any invention or of what may be claimed, but rather as descriptions of features that may be specific to particular embodiments of a particular invention. Certain features that are described in this specification in the context of separate embodiments may also be implemented in combination in a single embodiment. Conversely, various features that are described in the context of a single embodiment may also be implemented separately or in any suitable sub-combination in multiple embodiments. In addition, although features may be described above as acting in certain combinations and even initially claimed as such, in some cases one or more features from a claimed combination may be excluded from the combination, and the claimed combination may refer to a sub-combination or a variant of a sub-combination.

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

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

Claims

1. A computer-implemented method for generating an output sequence of discrete tokens using a diffusion model, the method comprising: generating a final latent representation of the sequence of discrete tokens using the diffusion model, the final latent representation including a determined value for each of a plurality of latent variables, wherein generating the final latent representation includes, at each of a plurality of reverse diffusion time steps: processing a diffusion model input that includes an intermediate latent representation of the sequence of discrete tokens for the reverse diffusion time step to generate an estimate of the sequence of discrete tokens up to the reverse diffusion time step; using the estimate to define a distribution over a continuous space of possible values for each of the plurality of latent variables; and generating an updated intermediate latent representation of the sequence of discrete tokens for the reverse diffusion time step by sampling from the distribution; applying a de-embedding matrix with learned values to the final latent representation of the output sequence of discrete tokens to generate a de-embedded final latent representation, the de-embedded final latent representation including, for each of the plurality of latent variables, a respective numerical score for each discrete token in a vocabulary of a plurality of discrete tokens; for each of the plurality of latent variables, selecting the discrete token having the highest numerical score from among the plurality of discrete tokens in the vocabulary; and generating the output sequence of discrete tokens including the selected discrete tokens.

2. The method of claim 1, wherein the discrete tokens include text, symbols, or signals.

3. The method of any one of claims 1 to 2, wherein the diffusion model input further includes an estimate of the sequence of discrete tokens generated up to the previous reverse diffusion time step.

4. The method of any one of claims 1 to 3, wherein generating the output sequence of discrete tokens using the diffusion model includes generating unconditional discrete tokens.

5. The method of any one of claims 1 to 3, wherein generating the output sequence of discrete tokens using the diffusion model includes generating discrete tokens conditioned on an input sequence of discrete tokens, and wherein the method comprises: receiving the input sequence of discrete tokens; converting each discrete token in the input sequence to a one-hot vector; and applying an embedding matrix with pre-trained values to each one-hot vector to embed the one-hot vector into a continuous vector.

6. The method of claim 5, wherein the output sequence of discrete tokens further includes the input sequence of discrete tokens received by the diffusion model.

7. The method of any one of claims 5 to 6, further comprising: applying a linear projection to the continuous vector to generate a projected continuous vector; and using the diffusion model to process the projected continuous vector.

8. The method of any one of claims 5 to 7, wherein generating the discrete tokens conditioned on the input sequence of discrete tokens includes using a classifier-free guidance technique.

9. The method according to any one of claims 1 to 7, wherein the de-embedding matrix has been learned during the training of the diffusion model, while the embedding matrix is fixed during the training of the diffusion model.

10. The method according to any one of claims 1 to 9, further comprising: training the diffusion model on unlabeled discrete token data comprising discrete token inputs to minimize the mean squared error between each discrete token input and an estimate of the discrete token input generated by the diffusion model up to the sampled reverse diffusion time steps.

11. The method according to claim 10, wherein the training also minimizes the cross-entropy loss evaluated on the final latent representation of the sequence of discrete tokens generated by the diffusion model from the discrete token inputs.

12. The method according to any one of claims 10 to 11, wherein the training comprises: applying a binary mask to a sequence of discrete token inputs, the binary mask comprising one or more first masks defining conditional tokens in the sequence and one or more second masks defining padding tokens in the sequence; and training the diffusion model on the masked sequence of discrete token inputs to generate an estimate of the sequence of discrete token inputs, the sequence of discrete token inputs including the padding tokens in the sequence.

13. The method according to any one of claims 10 to 12, wherein the training comprises learning the values of the de-embedding matrix while keeping the pre-trained values of the embedding matrix fixed.

14. A system, comprising: one or more computers; and one or more storage devices storing instructions that, when executed by the one or more computers, are operable to cause the one or more computers to perform the operations of the corresponding method according to any one of the preceding claims.

15. 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 corresponding method according to any one of the preceding claims.

Images