Discrete token processing using diffusion models

Abstract

In one aspect, the method includes: generating a final latent representation of the sequence of discrete tokens using the diffusion model, the final latent representation including determined values ​​for each of a plurality of latent variables; applying a de-embedding matrix to the final latent representation of the output sequence of discrete tokens to generate, for each of the plurality of latent variables, a de-embedding final latent representation including a respective numerical score for each discrete token in a vocabulary of the plurality of discrete tokens; selecting, for each of the plurality of latent variables, a discrete token having a highest numerical score from among the plurality of discrete tokens in the vocabulary; and generating an output sequence of discrete tokens including the selected discrete token.

Description

[Technical Field]

[0001] CROSS-REFERENCE TO RELATED APPLICATIONS This application claims priority to U.S. Provisional Patent Application No. 63 / 411,045, filed September 28, 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. [Background technology]

[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 nonlinear units to predict an output for a received input. Some neural networks contain one or more hidden layers in addition to an output layer. The output of each hidden layer is used as the input for 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. Summary of the Invention

[0004] According to a first aspect, a computer-implemented method for generating an output sequence of discrete tokens using a diffusion model is provided. The method includes using the diffusion model to generate a final latent representation of the sequence of discrete tokens, the final latent representation including determined values ​​for each of a plurality of latent variables. Generating the final latent representation includes, at each of a plurality of de-diffusion time steps, processing a diffusion model input including an intermediate latent representation of the sequence of discrete tokens for the de-diffusion time step to generate an estimate of the sequence of discrete tokens at the time of the de-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 de-diffusion time step through sampling from the distribution. The method further includes applying a de-embedding matrix having learned values ​​to the final latent representation of the output sequence of discrete tokens to generate, for each of the plurality of latent variables, a de-embedded final latent representation including a respective numerical score for each discrete token in a vocabulary of the plurality of discrete tokens. The method further includes selecting, for each of the plurality of latent variables, a discrete token from among the plurality of discrete tokens in the vocabulary having a highest numerical score, and generating an output sequence of discrete tokens that includes the selected discrete token.

[0005] In the first dediffusion time step, the intermediate latent representation may be the initial latent representation, and in subsequent dediffusion time steps, the intermediate latent representation may be the updated intermediate latent representation generated in the immediately preceding dediffusion time step.

[0006] In an embodiment, the estimate of the sequence of discrete tokens at the de-spreading time step is an estimate of the final latent representation of the output sequence of discrete tokens at the de-spreading time step.

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

[0008] In an embodiment, the diffusion model input further includes an estimate of the sequence of discrete tokens generated at the previous de-diffusion time step, which estimate of the sequence of discrete tokens at the previous de-diffusion time step is an estimate of the final latent representation of the output sequence of discrete tokens at the previous de-diffusion time step.

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

[0010] In an embodiment, generating an output sequence of discrete tokens using a diffusion model includes generating discrete tokens conditioned on an input sequence of discrete tokens, and the method further comprises: receiving an input sequence of discrete tokens; Transforming each discrete token in the input sequence into a one-hot vector; applying an embedding matrix with pre-trained values ​​to each one-hot vector to embed the one-hot vector into a continuous vector.

[0011] In an embodiment, generating an output sequence of discrete tokens using a diffusion model includes generating discrete tokens conditioned on an input sequence of discrete tokens, and the method further comprises: receiving an input sequence of discrete tokens; Transforming each discrete token in the input sequence into a one-hot vector; applying an embedding matrix to each one-hot vector to embed the one-hot vector into an embedding.

[0012] In an embodiment, the output sequence of discrete tokens also includes the input sequence of discrete tokens received by the diffusion model.

[0013] In an embodiment, the input sequence of discrete tokens comprises discrete tokens representing audio data input including, for example, spoken words that characterize audio waveforms in the time domain or time-frequency domain.

[0014] In an embodiment, the output tokens represent audio data including, for example, spoken words that characterize audio waveforms in the time domain or time-frequency domain.

[0015] In an embodiment, the input sequence is generated from an audio signal that includes speech. In an embodiment, the input may be received in the form of an audio (speech) signal that is converted by a speech-to-text converter to form the input sequence.

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

[0017] In an embodiment, an input sequence of discrete tokens represents a sequence of actions to be performed by an agent, for example a mechanical agent in a real-world environment that implements actions to perform a mechanical task.

[0018] In an embodiment, the output sequence of discrete tokens represents a sequence of actions to be performed by an agent, for example a mechanical agent in a real-world environment that implements actions to perform a mechanical task.

[0019] In an embodiment, the method further comprises: applying a linear projection to the continuous vector to generate a projected continuous vector; and processing the projected continuous vectors using a diffusion model.

[0020] In an embodiment, generating discrete tokens conditioned on the input sequence of discrete tokens includes using classifier-less guidance techniques.

[0021] In an embodiment, the embedding matrix is ​​fixed during training of the diffusion model, while the de-embedding matrix is ​​learned during training of the diffusion model.

[0022] In an embodiment, the method further includes training the diffusion model with unlabeled discrete token data including the 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 at the time of the sampled de-diffusion time step.

[0023] In an embodiment, the training also minimizes the cross-entropy loss evaluated with respect to the final latent representation of the sequence of discrete tokens generated by the diffusion model from the discrete token input.

[0024] In an embodiment, the training comprises: applying a binary mask to a sequence of discrete token inputs, the binary mask including one or more first masks defining conditioning tokens in the sequence and one or more second masks defining infilling tokens in the sequence; training a diffusion model with the sequence of masked discrete token inputs to generate estimates of the sequence of discrete token inputs including filler tokens in the sequence.

[0025] In an embodiment, training involves learning the values ​​of the de-embedding matrix while keeping the pre-trained values ​​of the embedding matrix fixed.

[0026] According to another aspect, a system is provided that includes one or more computers and one or more storage devices that store 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.

[0027] 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 methods described above.

[0028] This specification describes a system, implemented as a computer program on one or more computers at one or more locations, that generates an output sequence of discrete tokens using a diffusion model that performs a de-diffusion process on a continuous embedding in an embedding space.

[0029] In some cases, the generation process may be unconditional, where the output sequence can be generated by the diffusion model from random noise. In other cases, the generation process may be conditional, for example, where 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.

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

[0031] The diffusion models described herein perform a diffusion process in a continuous latent space, yet can process inherently discrete data, such as text data, i.e., receive discrete input data and / or generate discrete output data. By performing the diffusion process on a continuous embedding in the latent space and converting the final latent representation back to discrete space with a continuous-to-discrete transformation step using a learnable de-embedding matrix, the described diffusion models are both flexible (e.g., easily configurable for both conditional and unconditional discrete data generation) and scalable (e.g., to arbitrarily large models with 100 million, 500 million, or more parameters) to better meet the needs of a variety of generation tasks. The described diffusion models achieve or exceed state-of-the-art performance on a wide range of conditional and unconditional text generation tasks achieved by many existing autoregressive models.

[0032] Advantageously, unlike most existing autoregressive models, which predict output tokens one after the other in a way that earlier tokens are not informed by later tokens, the described diffusion model can predict all tokens in a sequence at once. This enables bidirectional (e.g., rather than causal) attention, so that token choices made later in the sequence influence those made earlier, resulting in higher-quality data generation. This significantly improves the parallelization of the inference process; because all tokens are generated simultaneously with a diffusion model rather than sequentially as required with an autoregressive model, multiple diffusion processes can be run in parallel with each other to generate any amount of data within a given time budget.

[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. [Brief explanation of the drawings]

[0034] [Figure 1]FIG. 1 is a block diagram of an exemplary data generation system. [Figure 2] FIG. 1 is a flow diagram of an exemplary process for generating a discrete output sequence. [Figure 3] FIG. 3 is a flow diagram of an exemplary process for one sub-step of the process step of FIG. 2. DETAILED DESCRIPTION OF THE INVENTION

[0035] Like reference symbols and designations in the various drawings refer to like elements.

[0036] 1 is a diagram of an exemplary data generation system 100. 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 may 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 sent 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 includes multiple tokens from a finite number of possible tokens.

[0040] For example, the discrete output sequence 162 may include a respective token from a vocabulary of discrete tokens at each of a plurality of positions. 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. In some examples, the output tokens represent audio data, including, for example, spoken words that characterize audio waveforms in the time domain or time-frequency domain. In some examples, the output sequence of discrete tokens represents text. In one example, the output sequence of discrete tokens is converted by a speech synthesis converter to form an audio signal. In some examples, the output sequence of discrete tokens represents a sequence of actions performed by an agent, such as, for example, a mechanical agent in a real-world environment that implements actions to perform a mechanical task.

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

[0042] In other cases, the data generation system 100 can be configured as a conditioned 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 a respective token from a vocabulary of tokens at each of a plurality of positions. The vocabulary of tokens can include any of a variety of discrete tokens representing text, symbols, or signals. In other cases, the input sequence 102 includes a continuous input sequence that represents, for example, a sequence of pixels in an image, or represents an audio waveform or audio spectrogram.

[0043] For example, when configured as a conditional output sequence generation system, the data generation system 100 can receive an input sequence 102 that includes textual data and / or additional non-textual data, such as, for example, images, video, speech, etc., and generate a discrete output sequence 162 that is a sequence of text, such as, for example, 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 sequence of text on a topic specified by the input sequence 102, a textual description of the data of the different data type, etc.

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

[0045] 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 many other forms or modalities of discrete data items. 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.

[0046] In one example, the input may be received in the form of an audio (speech) signal captured by a microphone and converted by a speech recognition converter to form the input sequence 102. In some examples, the input sequence of discrete tokens includes tokens that represent audio data input, including spoken words, e.g., characterizing audio waveforms in the time domain or time-frequency domain. In some examples, the input sequence of discrete tokens represents a sequence of actions to be performed by an agent, e.g., a mechanical agent in a real-world environment implementing actions to perform a mechanical task.

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

[0048] The initial latent representation 122 has the same dimensions as the final latent representation 136, but different values. That is, the initial latent representation 122 contains multiple latent variables and the final latent representation 136 contains the same number of latent variables, but the values ​​of these latent variables typically differ between these two representations 122 and 136.

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

[0050] When configured as a conditioned output sequence generation system, a preprocessing engine 110, an optional component of the data generation system 100, processes the received input sequence 102 to generate a conditioned embedding 116 that includes a sequence of numerical values. An initialization engine 120 then generates the initial latent representation 122 by using the numerical values ​​included in the conditioned embedding 116 as initial values ​​for some of the latent variables included in the initial latent representation 122 and sampling the initial values ​​for each of the remaining latent variables from a corresponding noise distribution, e.g., a Gaussian distribution or other predetermined distribution. Thus, the initial latent representation 122 includes multiple latent variables, some of which have initial values ​​determined from the input sequence 102 and other of which have initial values ​​sampled from a corresponding noise distribution.

[0051] More specifically, the preprocessing engine 110 includes an embedding matrix 112 that defines a mapping from the discrete space to the embedding space. To define this mapping, the embedding matrix 112 may be a matrix with dimensions D×V, i.e., a matrix that may have D×V predefined values ​​included as entries of the matrix, where D is the embedding size (the size of each continuous vector in the embedding space) and V is the vocabulary size (the total number of discrete tokens included in the vocabulary of tokens). Thus, in one example, each token ω in the vocabulary is assigned a norm ω fixed to match the norm of a random Gaussian sample of dimension D used to denoise the data.

number

number

[0052] The pre-processing engine 110 converts each token in the input sequence 102 into a discrete one-hot vector and applies an embedding matrix 112 to each discrete one-hot vector, mapping the discrete one-hot vector to a continuous vector in an embedding space that has the numerical values ​​of the token, e.g., based on an element-wise multiplication operation. Thus, in some implementations, the conditioned embedding 116 consists of the numerical values ​​contained in the continuous vector resulting from applying the embedding matrix 112 to the input sequence 102. In other implementations, the pre-processing engine 110 additionally applies a linear projection to each continuous vector to generate a projected continuous vector, and then generates the conditioned embedding 116 from the numerical values ​​contained in the projected continuous vectors. The conversion to tokens can be performed using the SentencePiece tokenizer 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 contents of which are incorporated herein by reference, and which consists of 32,000 words.

[0053] Diffusion model 130 may have any suitable architecture such that, at any specified de-diffusion time step, the diffusion model processes a diffusion model input for the de-diffusion time step that includes an intermediate latent representation (at the time of the de-diffusion time step) having the same dimensions as the discrete output sequence 162 to generate a diffusion model output 132 for the specified de-diffusion time step. The diffusion model output 132 includes an estimate of the discrete output sequence 162 (at the time of the de-diffusion time step) that also has the same dimensions as the discrete output sequence 162.

[0054] For example, the diffusion model 130 may be a convolutional neural network, such as a U-Net, or other architecture, including one or more convolutional residual blocks. As another example, the diffusion model 130 may be an attention neural network including one or more attention blocks, such as a self-attention block, a gated attention block, or a cross-attention block. As yet another example, the diffusion model 130 may include both a convolutional residual block and an attention block.

[0055] At the specified de-diffusion time step, data generation system 100 then uses update engine 140 to update the intermediate latent representations based on the estimates of discrete output sequence 162 for the specified de-diffusion time step. The intermediate latent representations after being updated by update engine 140 at the specified de-diffusion time step are referred to herein as updated intermediate latent representations 134, and the intermediate latent representations after being updated at the final de-diffusion time step in the de-diffusion process are referred to herein as final intermediate latent representations 136.

[0056] Thus, at any given de-diffusion time step, the intermediate latent representations provided as (part of) the diffusion model input to the diffusion model 130 will be the updated intermediate latent representations 134 generated at the immediately preceding de-diffusion time step. At the first de-diffusion time step, the intermediate latent representations are the initial latent representations 122 generated by the initialization engine 120.

[0057] More specifically, at any specified de-diffusion time step, the update engine 140 uses at least the estimates of the discrete output sequence 162 for that specified de-diffusion time step to define a corresponding distribution over a continuous space of possible values ​​for each of a plurality of latent variables included in the initial latent representation 122. The update engine 140 then generates an updated intermediate latent representation 134 of the discrete output sequence 162 for the de-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. The updated latent representation 134 thus includes a plurality of latent variables, each of which has a value sampled from a corresponding distribution defined by the diffusion model output 132 of the diffusion model 130.

[0058] To define such a distribution for each latent variable, the update engine 140 can calculate the mean, and optionally the variance, of the distribution based on estimates of the discrete output sequence 162 and the intermediate latent representations before updating at the specified de-diffusion time step. The update step is further described below with reference to FIG. 3.

[0059] Some implementations of data generation system 100 use self-conditioning input 138 when performing the despreading process. When data generation system 100 is configured as a conditioned output sequence generation system, the use of self-conditioning can improve the quality of discrete output sequence 162.

[0060] At any specified de-diffusion time step, the self-conditioning input 138 can be an estimate of the discrete output sequence 162 contained in the diffusion model output generated at the immediately preceding de-diffusion time step. Thus, for any specified de-diffusion time step, the diffusion model 130 can process the diffusion model input, including (i) the intermediate latent representation (as of the de-diffusion time step) and (ii) the estimate of the discrete output sequence 162 generated at the immediately preceding de-diffusion time step, to generate the diffusion model output 132 for the specified de-diffusion time step.

[0061] In some implementations of the data generation system 100, guidance is used when performing the dediffusion process. That is, the dediffusion process is optionally a guided dediffusion process. The guidance can be classifier-less guidance to alleviate the need for a separately trained guidance neural network.

[0062] Using classifier-less guidance may involve processing a sequence of conditioning tokens as part of the diffusion model input by the diffusion model 130 at any specified inverse diffusion time step to generate a diffusion model output. The conditioning tokens may be defined by a binary conditioning mask that is set to 1 for positions in the discrete output sequence that each have a predetermined discrete token ("conditioning positions") and 0 for positions in the discrete output sequence that each have an unknown discrete token ("filling positions"), i.e., positions where a discrete token should be generated by the system.

[0063] After the final de-diffusion time step, data generation system 100 provides post-processing engine 150 with a final latent representation 136 that is in the same embedding space as initial latent representation 122. Final latent representation 136 includes multiple latent variables. Each latent variable has a decision value determined as a result of the de-diffusion process, which may differ from its initial value. Post-processing engine 150 processes final latent representation 136 to generate a discrete output sequence 162 that is in the same discrete space as input sequence 102.

[0064] To that end, 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, de-embedding matrix 152 may be a matrix with dimensions V×D, i.e., a matrix that may have V×D pre-trained values ​​included as its entries, where V is the vocabulary size (the total number of discrete tokens in the vocabulary of tokens) and D is the embedding size (the size of each continuous vector in the embedding space).

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

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

[0067] Because the computations involved in generating each discrete output sequence can be parallelized over the entire length of the discrete output sequence, data generation system 100 makes better use of the computational resources of the hardware accelerator in which generation system 100 may be implemented.

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

[0069] Training the diffusion model 130 and other trainable components of the data generation system 100 and determining the trained parameter values ​​of these components is described further below.

[0070] 2 is a flow diagram of an exemplary process 200 for generating a discrete output sequence. For convenience, process 200 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 shown in FIG. 1, suitably programmed in accordance with this specification can perform process 200.

[0071] The system performs a de-diffusion process on the 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 a vocabulary of tokens at each of a plurality of positions. The vocabulary of tokens can include any of a variety of discrete tokens representing text, symbols, or signals.

[0072] The initial latent representation has the same dimensions as the final latent representation but different values, i.e., the initial latent representation contains multiple latent variables and the final latent representation contains the same number of latent variables, but the initial values ​​of these latent variables in the initial latent representation and the determined values ​​of these latent variables in the final latent representation are usually different from each other.

[0073] When configured as an unconditional output sequence generation system, the system can generate initial latent representations based on sampling initial values ​​for each of a plurality of latent variables from a corresponding noise distribution (e.g., a Gaussian distribution or other predetermined distribution).

[0074] When configured as a conditional output sequence generation system, the system can receive an input sequence and then map the received input sequence to a conditional embedding containing a sequence of numbers. The mapping is expressed as an embedding matrix D, where D is the embedding size and V is the vocabulary size.

number

[0075] The system then generates a final latent representation by updating the initial latent representation over multiple de-diffusion time steps using a diffusion model. In other words, the final latent representation is the updated intermediate latent representation generated in the last de-diffusion time step.

[0076] Updating the initial latent representation is described in further detail below with reference to Figure 3, which shows a flow diagram of an exemplary process 300 of substeps 302-306 of step 202. In general, the system may repeatedly perform process 300 at each of multiple despreading time steps T, T-1, ..., 2, 1.

[0077] The system processes the diffusion model inputs for the inverse diffusion time steps to generate diffusion model outputs (step 302). The diffusion model outputs can include estimates of the discrete output sequence at the inverse diffusion time steps, while the diffusion model inputs include (i) intermediate latent representations of the discrete output sequence (at the inverse diffusion time steps) x t , and (ii) a time index t indicating the current despreading time step. In one example, the diffusion model is a trained neural network that provides an intermediate latent representation of the discrete output sequence (at the despreading time step) x t and the time index t indicating the current despreading time step as input, and the final latent representation of the discrete output sequence at the despreading time step

number

[0078] In the first dediffusion time step, the intermediate latent representation x t is the initial latent representation. At any subsequent de-diffusion time step, the intermediate latent representation x t is the updated intermediate latent representation generated in the previous de-diffusion time step.

[0079] For example, estimates of the discrete output sequence

number

number

[0080] Optionally, the diffusion model inputs also include self-conditioning inputs. At any given de-diffusion time step, the self-conditioning inputs can be estimates of the discrete output sequence contained in the diffusion model output generated at the immediately preceding de-diffusion time step t+1. Thus, for any given de-diffusion time step, the diffusion model inputs include estimates of the discrete output sequence contained in the diffusion model output generated at the immediately preceding de-diffusion time step t+1.

number

[0081] Another example is the estimate of the discrete output sequence

number

number

[0082] Further optionally, the diffusion model input also includes guidance, such as, for example, classifier-less guidance. When using classifier-less guidance, for any specified de-diffusion time step, the diffusion model input may additionally include a fixed sequence of conditioning tokens.

[0083] As yet another example, the estimate of the discrete output sequence

number

number

[0084] The system estimates the discrete output sequence

number

number

[0085] For example, the system can calculate the mean μ of the distribution as follows:

number

number

number

[0086] The system generates an updated intermediate latent representation x of the discrete output sequence for the back-diffusion time step from the distribution t-1 (step 306). The system can do this by selecting values ​​for each latent variable from a corresponding distribution, for example, by sampling values ​​from a continuous space of possible values ​​according to the distribution:

number

[0087] Therefore, the updated latent representation x t-1 contains multiple latent variables, and the value of each latent variable is sampled from a corresponding distribution defined by the diffusion model output of the diffusion model for the current inverse diffusion time step.

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

[0089] For example, the system can generate the final unembedded latent representation by:

number

number

number

[0090] For each of a plurality of latent variables, the system selects a token from among the tokens in the vocabulary according to a corresponding distribution (step 206). For each latent variable, the system can, for example, select the token with the highest numerical score from among all tokens in the vocabulary.

[0091] The system generates a discrete output sequence (step 208), which includes the tokens selected from the vocabulary in step 206. If configured as a conditional output sequence generation system, the discrete output sequence may also optionally include an input sequence received as input by the system.

[0092] By repeatedly executing process 200, the system can generate different discrete output sequences. That is, process 200 can be performed as part of predicting a desired output, i.e., a discrete output sequence from an input sequence, where the discrete output sequence to be generated by the system from the input sequence is unknown.

[0093] Some steps of process 200, e.g., sub-step 302 of step 202, may also be performed as part of processing input sequences derived from a set of training data, i.e., inputs derived from a set of inputs for which the discrete output sequences to be produced by the system are known, to train trainable components of the system and determine training values ​​for the parameters of those components. As an example, a diffusion model can be trained on 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 one example, the training data may be converted into tokens using the SentencePiece tokenizer 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 contents of which are incorporated herein by reference, and consisting of 32,000 words. In one example, the diffusion model may be trained with a sequence length of 256. In one example, 10% of the padding tokens may be inserted uniformly in the training set (i.e., not necessarily at the end of the sequence) to allow the model to generate samples of various sizes and provide more flexibility.

[0094] Specifically, as part of the diffusion model training process, the system may repeatedly perform substep 302 of step 202 on input sequences selected from a set of unlabeled discrete training data to train a trainable component of the system and optimize an objective function appropriate to the discrete data generation task that the diffusion model is configured to perform.

[0095] For example, the objective function may include an input sequence x chosen from unlabeled discrete training data. 0で We train a diffusion model θ and, for each input sequence x0, estimates of the input sequence generated by the diffusion model at sampled dediffusion time steps within the dediffusion process.

number

number

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

number

number

[0097] During training, the system may incorporate any number of techniques to improve the speed, effectiveness, or both of the training process.

[0098] For example, the system can use span masking techniques to train a diffusion model on several filler tasks, such as fill-in-the-middle and spans-in-filling tasks. In this example, for each input sequence selected from unlabeled discrete training data, the system can apply a binary mask to the discrete tokens in the input sequence. The binary mask includes one or more first masks that define conditioning tokens in the sequence and one or more second masks that define filler tokens in the sequence. The system can then train a diffusion model on the masked input sequence to generate estimates of the filler tokens in the masked input sequence, i.e., estimates of the original discrete tokens in the input sequence masked by the second mask.

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

[0100] As yet another example, to stabilize training and avoid degradation of unigram entropy, the system can specifically train the de-embedding matrix R while keeping the embedding matrix E fixed during training. Thus, the de-embedding matrix R is learned during training of the diffusion model, while the embedding matrix E is not learned during training.

[0101] In this example, "learned" means that one or more values ​​included as entries in the de-embedding matrix R are adjusted during training of the diffusion model. In contrast, "unlearned" means that predefined values ​​included as entries in the embedding matrix E determined before training of the diffusion model (e.g., values ​​determined as a result of training another neural network with which the embedding matrix E was trained) remain fixed throughout training of the diffusion model.

[0102] As mentioned above, in one example where self-conditioning is used, the x0 estimate is gradually improved by passing the estimate obtained in the previous sampling step as input to the diffusion model. To approximate the inference behavior during training while maintaining computational efficiency, an initial estimate is calculated with self-conditioning set to zero:

number

[0103] next,

number

number

[0104] The diffusion model is then optimized using the outputs from the two forward propagations to accurately estimate x0 with or without self-conditioning.

[0105] The term "configured" is used herein in connection with systems and computer program components. One or more computer systems configured to perform particular operations or actions means that software, firmware, hardware, or a combination thereof is installed on the system that, when running, causes the system to perform the operation or action. One or more computer programs configured to perform particular operations or actions means that one or more programs contain instructions that, when executed by a data processing device, cause the device to perform the operation or action.

[0106] Embodiments of the subject matter and functional operations described herein may be implemented in digital electronic circuitry, in tangibly embodied computer software or firmware, in computer hardware, or one or more combinations thereof, including the structures disclosed herein and structural equivalents thereof. Embodiments of the subject matter described herein may be implemented as one or more modules of computer program instructions, i.e., as one or more computer programs 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 a data processing apparatus.

[0107] 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 computers. An apparatus can also be or further include special-purpose logic circuitry, such as an FPGA (field-programmable gate array) or an ASIC (application-specific integrated circuit). In addition to hardware, an apparatus can optionally include 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.

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

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

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

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

[0112] 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. Typically, the central processing unit receives instructions and data from a read-only memory, a random-access memory, or both. The essential 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. Typically, a computer also includes one or more mass storage devices, e.g., magnetic disks, magneto-optical disks, or optical disks, for storing data, or the computer is operatively connected to receive data from, transmit data to, or both, one or more mass storage devices, e.g., magnetic disks, magneto-optical disks, or optical disks for storing data. However, a computer need not have such devices. Additionally, computers may be incorporated into other devices, such as mobile phones, personal digital assistants (PDAs), mobile audio or video players, game consoles, global positioning system (GPS) receivers, or portable storage devices, such as universal serial bus (USB) flash drives, to name just a few.

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

[0114] 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 documents to and receiving documents from 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.

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

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

[0117] 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 a local area network (LAN), a wide area network (WAN), e.g., the Internet.

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

[0119] While this 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 multiple features are described above as functioning in a particular combination and are initially recited as such in the claims, one or more features may, in some cases, be deleted from the claimed combination, and the claimed combination may be directed to subcombinations or variations of subcombinations.

[0120] 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 embodiments should not be understood as requiring such separation in all embodiments, and it should be understood that the described program components and systems generally may be integrated into a single software product or packaged into multiple software products.

[0121] 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 an output sequence of discrete tokens using a diffusion model, comprising: generating a final latent representation of the sequence of discrete tokens using the diffusion model, the final latent representation including determined values ​​for each of a plurality of latent variables, wherein generating the final latent representation includes, at each of a plurality of de-diffusion time steps: processing a diffusion model input comprising an intermediate latent representation of the sequence of discrete tokens at the de-diffusion time step to generate an estimate of the sequence of discrete tokens at the time of the de-diffusion time step; using the estimates to define a distribution over a continuous space of possible values ​​for each of the plurality of latent variables; generating an updated intermediate latent representation of the sequence of discrete tokens for the de-diffusion time step through sampling from the distribution; applying a de-embedding matrix having learned values ​​to the final latent representation of the output sequence of discrete tokens to generate, for each of the plurality of latent variables, a final de-embedded latent representation including a respective numerical score for each discrete token in a vocabulary of a plurality of discrete tokens; selecting, for each of the plurality of latent variables, a discrete token from the plurality of discrete tokens in the vocabulary having a highest numerical score; generating an output sequence of the discrete tokens including the selected discrete tokens; method.

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

3. The method according to any one of claims 1 to 2, wherein the diffusion model input further comprises an estimate of the sequence of discrete tokens generated at a previous de-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 comprises generating unconditional discrete tokens.

5. Generating the output sequence of discrete tokens using the diffusion model includes generating discrete tokens conditioned on an input sequence of discrete tokens, the method comprising: receiving an input sequence of said discrete tokens; converting each discrete token in the input sequence into a one-hot vector; applying an embedding matrix having 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 also includes the input sequence of discrete tokens received by the diffusion model.

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

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

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

10. 10. The method of claim 1, further comprising training the diffusion model with unlabeled discrete token data comprising the discrete token inputs to minimize a mean squared error between each discrete token input and an estimate of the discrete token input generated by the diffusion model at a sampled de-diffusion time step.

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

12. The training includes: applying a binary mask to a sequence of discrete token inputs, the binary mask including one or more first masks defining conditioning tokens in the sequence and one or more second masks defining filler tokens in the sequence; and training the diffusion model with a masked sequence of discrete token inputs to generate an estimate of the sequence of discrete token inputs including the filler token in the sequence.

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

14. 14. 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 each of the methods of any one of claims 1 to 13.

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 each of the methods recited in any one of claims 1 to 13.

Images