Accelerating diffusion model inference using speculative execution

Abstract

A computer-implemented method, system, and computer program code for generating a data item represented by a plurality of continuous-valued elements, e.g. using a diffusion model. The method generates a draft sequence, of draft representations of the data item, and determines, e.g. in parallel, whether to accept or reject these. If a draft representation is rejected a new representation is determined. In implementations the new representation is computed deterministically as a function of the rejected draft representation.

Description

Attorney Docket No.: MX402430WO ACCELERATING DIFFUSION MODEL INFERENCE USING SPECULATIVE EXECUTION BACKGROUND

[0001] This specification relates generating data items using machine learning models.

[0002] As one example, neural networks are machine learning models that employ one or more layers of nonlinear units to predict an output for a received input. Some neural networks include one or more hidden layers in addition to an output layer. The output of each hidden layer is used as input to another layer in the network, e.g., the next hidden layer or the output layer. Each layer of the network generates an output from a received input in accordance with current values of a respective set of weights. SUMMARY

[0003] This specification describes systems and methods, implemented as computer programs on one or more computers in one or more locations, that generate an output data item using a diffusion model, or other similar model.

[0004] The described techniques can use parallel processing to speed up the generation of an output data item having continuous-valued elements, e.g. using a diffusion model, typically implemented as a diffusion model neural network. More generally the techniques can be used for sampling from any unnormalized distribution that is differentiable.

[0005] In one aspect there is provided a computer-implemented method of generating a data item represented by a plurality of continuous-valued elements. The method involves obtaining an initial representation of the data item, and generating a sequence of representations of the data item representing data items, until obtaining a final representation of the data item. In some implementations, e.g. when applied to a diffusion model, the sequence of representations can represent data items with successively reduced levels of noise. In some other implementations, e.g. when applied to unadjusted Langevin sampling, the level of noise need not reduce and can, e.g., remain constant.

[0006] The method generates a draft sequence of draft representations of the data item from a current representation of the data item in the sequence. The current representation of the data item is the first draft representation in the draft sequence. In some implementations, e.g. when the technique is applied to a diffusion model, each successive draft representation of the data item in the draft sequence represents a data item with a successively reduced level of noise. However this is not an essential feature of the described techniques. For example,Attorney Docket No.: MX402430WO successively reduced levels of noise in the draft sequence need not be present when sampling from an unnormalized distribution, e.g. when using the described techniques to speed up unadjusted Langevin sampling.

[0007] For each successive draft representation of the data item in the draft sequence after the first, e.g. in parallel, the method determines whether to accept or reject the draft representation for inclusion in the sequence of representations of the data item. If the draft representation is rejected a new representation is determined for inclusion in the sequence of representations of the data item.

[0008] In implementations the new representation is computed deterministically as a function of the rejected draft representation.

[0009] The subject matter described in this specification can be implemented in particular embodiments so as to realize one or more of the following advantages.

[0010] In general a diffusion model can generate a data item by performing a reverse diffusion process to gradually de-noise a representation of the data item. This can be done either in a space of the data item, e.g. a pixel space when generating an image, or in a latent space of the data item. However many, e.g. tens or hundreds, of denoising steps can be needed, which is slow and computationally expensive. Similar issues arise more generally when sampling from an unnormalized probability density based on an energy function.

[0011] The described techniques involve speculatively generating a draft sequence comprising draft representations of the data item, and then determining whether to accept or reject successive draft representations of the draft sequence. A difficulty arises when a draft representation is rejected because it is computationally very inefficient to sample from the correct adjusted distribution. Implementations of the techniques obtain a new representation deterministically from the rejected sample (draft representation).

[0012] The draft representations can also be processed in parallel using the target model, i.e. the diffusion model. This can result in much faster generation of the data item; in implementations it can also be done in a way that is computationally less expensive than conventional approaches.

[0013] The draft sequence can be generated by a smaller diffusion model, i.e. one with fewer trained parameters such as weights, or the draft sequence can be generated by a “frozen” version of the target model, optionally with a correction applied. That is, there is no need to train another, draft diffusion model. Obtaining the draft sequence from the frozen target model can significantly reduce the number of model evaluations needed, e.g. to just one per draft sequence or just one for the whole data item generation process.Attorney Docket No.: MX402430WO

[0014] A further advantage can be realized in some particular hardware implementations. For example, when using a hardware device capable of parallel processing, such as a hardware accelerator, loading the diffusion model weights into local, e.g. on-chip memory can be a significant bottleneck, and once loaded it is relatively inexpensive to compute multiple denoising outputs in parallel, e.g. as a batch. The described techniques are able to take advantage of this when processing continuous-values elements using a diffusion model neural network.

[0015] The described techniques can be applied in a space of the data item or in a latent space of the data item. They can be applied to diffusion model neural networks that predict a score (score-based diffusion models that predict the so-called “score”, i.e. the gradient of a log- likelihood function), and to diffusion model neural networks that predict a noise estimate or a corrected representation (which are equivalent to score-based diffusion models). They can be used with existing, pre-trained diffusion models.

[0016] Some implementations of the described techniques can generate exact samples from the target diffusion model. That is, the representations generated by the techniques can be as if they were generated by the target diffusion model, i.e. from the same distribution as that learned by the target diffusion model. Other implementations can relax this constraint for other benefits, e.g. faster data item generation.

[0017] In some implementations generation of the draft representations of the draft sequence can also be parallelized.

[0018] The described techniques are generally orthogonal to other acceleration methods, and can be combined with these.

[0019] The 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

[0020] FIG.1 shows an example system for generating a data item represented by a plurality of continuous-valued elements.

[0021] FIG.2 illustrates a speculative sampling process for continuous-valued data.Attorney Docket No.: MX402430WO

[0022] FIG.3 is a flow diagram of an example process for generating a data item represented by a plurality of continuous-valued elements.

[0023] Like reference numbers and designations in the various drawings indicate like elements. DETAILED DESCRIPTION

[0024] FIG. 1 shows a system 100, implemented as computer programs on one or more computers in one or more locations, for generating a data item represented by a plurality of continuous-valued elements.

[0025] Generally the system 100 includes a target diffusion model neural network 110 for generating the data item. More specifically the target diffusion model neural network 110 is used for generating a sequence of representations of the data item representing data items with successively reduced levels of noise, until obtaining a final representation of the data item.

[0026] The final representation of the data item can be in the output space, i.e., so that the values in the generated data item are the values of a data item of the appropriate type, e.g., values of image pixels, amplitude values of an audio signal, and so on, or can be in a latent space, which can be decoded to the output space.

[0027] The system 100 can optionally also include a draft diffusion model neural network 120 for generating a draft representation of the data item, more specifically a draft sequence comprising draft representations of the data item. In the draft sequence each successive draft representation of the data item in the draft sequence represents a data item with a successively reduced level of noise.

[0028] As described later, the draft diffusion model neural network 120 is optional. For example, the draft sequence can be generated using the target diffusion model neural network 110. The example architectures and implementations described below for the target diffusion model neural network 110 also apply to the draft diffusion model neural network 120.

[0029] In implementations of the described techniques the target diffusion model neural network 110, and the draft diffusion model neural network 120, where present, have been trained. They can have been trained conventionally, i.e. the techniques described herein do not rely on any special training of the diffusion model neural network(s). Nonetheless merely for illustration, some suitable training techniques and training datasets are indicated later.

[0030] The target diffusion model neural network 110 is configured to process a representation 108 of the generated data item to generate a denoising output 112 for obtainingAttorney Docket No.: MX402430WO a subsequent, reduced noise version of the representation of the data item. There are various known ways in which such a denoising output can be generated and used, and some are outlined below.

[0031] In general the techniques described herein are used with stochastic denoising processes, sometimes referred to as stochastic samplers (from a distribution); DDPM (Denoising Diffusion Probabilistic Model) is one example (see, e.g., Ho et al. arXiv:2006:11239). However, the techniques can also be applied to deterministic denoising processes if these are modified to add noise.

[0032] In more detail, the target diffusion model neural network 110 is configured to process each element of the representation 108 to generate denoising output 112 comprising a vector of elements, i.e. a target model value for each element, that can be used to obtain the reduced noise version of the representation of the data item. Later the target diffusion model neural network 110 is also referred to as the target model, ^^. The model used to generate the draft sequence is referred to as the draft model, ^^.

[0033] The trained target diffusion model neural network 110 can also process an index 106, ^^, of the representation to generate the denoising output. The index, n, can also be referred to as a time step t; in some implementations t can run from 1 to (approximately) 0 during the denoising process. In some of the example algorithms given later k also denotes such a time step. Processing the index, n, generally involves processing data specifying the index, e.g. an embedding of the index. There are many different embeddings that can be used, e.g. a sinusoidal embedding (similar to that used in transformers), a learned embedding, and so forth.

[0034] In some implementations the trained target diffusion model neural network 110 (and the draft diffusion model neural network 120, where present) further processes a conditioning input 104 that characterizes one or more desired properties for the data item, i.e., that characterizes one or more properties that the final data item generated by the system should have. For example, for an image frame the content conditioning input may comprise text in a natural or computer language, or features of text, or audio, e.g. speech, or features of audio, that the image frame should represent.

[0035] The target diffusion model neural network 110 can have any appropriate architecture that allows the neural network to map the representation of the data item to a denoising output 112 of the same dimensionality (which is the dimensionality of the final representation of the data item).Attorney Docket No.: MX402430WO

[0036] For example, the target diffusion model neural network 110 can comprise a convolutional neural network, e.g., it can have a U-Net (Ronneberger et al., “U-Net: Convolutional Networks for Biomedical Image Segmentation”, arXiv:1505.04597) or other architecture that maps one input of a given dimensionality to an output of the same dimensionality. As another example, the denoising neural network can comprise a Transformer neural network that processes the representation of the current version of the data item through a set of self-attention layers to generate the denoising output 112. As some particular examples, the denoising neural network can have a U-ViT architecture (Bao, et al., arXiv:2209.12152, 2023; Appendix B of Hoogeboom et al., arXiv:2301.11093, 2023); or may comprise a diffusion transformer (DiT, Peebles and Xie, arXiv:2212.09748).

[0037] The target diffusion model neural network 110 can be conditioned on the conditioning input in any of a variety of ways. For example, it can include one or more cross-attention layers that each cross-attend into the one or more embeddings, and / or can include other types of neural network layer that are conditioned on the embeddings (e.g. a FiLM layer), and / or the conditioning data can be included as one or more additional channels in an input to the denoising neural network.

[0038] Optionally the conditioning input can be obtained by encoding conditioning data, e.g. using a text, audio or image encoder for text, audio, or image data. Example encoders include a Transformer-based language model (e.g. the text encoder of a T5 model, Raffel et al., arXiv:1910.10683. 2023); an audio encoder of an audio language model or of a speech recognition system (e.g. BEST-RQ, Chine et al. arXiv:2202.01855); and a vision Transformer neural network (that generates a set of embeddings that represent the image). As another example a scalar value can be mapped to an embedding using an embedding matrix.

[0039] Generally the denoising process can use so-called classifier-free guidance (e.g., Ho and Salimans, arXiv:2207.12598), in which conditional and unconditional denoising outputs are combined according to a guidance weight; or so-called classifier guidance.

[0040] The draft diffusion model neural network 120 can have the same architecture as the target diffusion model neural network 110 or it can be a cheaper model, i.e. one with fewer learned parameters (e.g. weights).

[0041] As previously described the target diffusion model neural network 110 is configured to generate a denoising output 112 comprising a vector comprising a target model value for each element, that can be used to obtain the subsequent reduced noise version of the representation of the data item. That is, the denoising output at an index, or time stepAttorney Docket No.: MX402430WO (obtained by processing the index or time step) can be used to obtain a reduced noise version of the representation of the data item for the next index, or time step. As used in the example algorithms given later the index, n, increases as denoising progresses; conventionally time decreases towards zero as denoising progresses.

[0042] The described techniques generate a subsequent reduced noise version of the representation of the data item from the target model or from the draft model. There are various ways in which the vector of target model values from the denoising output 112 can be used to obtain the subsequent reduced noise version of the representation of the data item. For example, the denoising output 112 can predict a denoising correction that can be subtracted from or added to the representation to obtain the subsequent representation, or the denoising output 112 can predict a corrected, denoised representation i.e. it can directly predict the subsequent representation, or it can predict a so-called v-parameterization (a linear combination of the noise and a denoised representation of the data item, Salimans and Ho arXiv: 2202.00512, 2022, section 4; Appendix D).

[0043] Implementations of the described techniques use a target model mean vector, i.e. a vector that defines mean values for the elements of the subsequent representation. The target model mean vector can be obtained from the target model values without explicitly determining the subsequent representation (which could be obtained by sampling, e.g., from a multivariate Gaussian distribution according to the target model mean vector).

[0044] For example, where the denoising output 112 predicts a denoising correction, the target model value for each element of the representation can estimate the noise in each element of the representation (to be subtracted from each element), or it can estimate a so- called score for each element of the representation. Generally, the score is the value of a score function that is a gradient of log probability density of the representation with respect to the representation (∇௫log^^^^^^^^^^^^). The score and estimated noise are equivalent up to a scaling factor and a sign (one is the negative of the other and they are related by a scaling factor).

[0045] Determining the target model mean vector (later ^^^; ^^^^ି^ ൫^^^^ି^൯) for arepresentation can involve summing, for each element of the drafta value of the element and the target model value for the element from the (trained) target diffusion model neural network 110.

[0046] As a particular illustrative example, a reverse diffusion process can be represented as: ^^^^௧ ൌ ^^௧^^^௧^^^^^ ^ ^^^^^ି௧^^^^௧ (1)Attorney Docket No.: MX402430WO ^^^~^^^

[0047] The process starts with noisy data, i.e. ^^^~^^^ and ^^^ ^ ^^^0, ^^^^^; time runs from^^ ൌ 1 (most noisy) to ^^ ൌ 0 (least noisy) is a diffusion coefficient controlling anamount of noise injected at time t in the correspondingprocess; and ^^^^௧represents the random noise added at each infinitesimal time step in the forward diffusionprocess (^^௧ is a Wiener process, also known as Brownian motion). There can be, say, ^^ ^1 discrete denoising steps from 0 to ^^, with step size ^^ ൌ 1 / ^^ and ^^^ ൌ ^^^^. Examples ofmodels with which the described techniques can be used are described in Albergo et al., “Stochastic Interpolants: A Unifying Framework for Flows and Diffusions”, arXiv:2303.08797v3, November 2024).

[0048] The term ^^௧^^^௧^ can be approximated by a neural network, e.g. the target model, q, that processes ^^ to generate, in this example, the above-described denoising correction, e.g.௧as ^^^^^^^. In the example implementations described later ^^^௧௧ ௧^^^௧^can be obtained from thetarget diffusion model neural network 110; later, ^^୯^∙^ denotes the trained target diffusion^ౡmodel neural network 110 for time step ^^ .^

[0049] The term ^^^^ corresponds to (is) the standard deviation of the noise in DDPM,^ି௧which is typically derived from a noise schedule; and ^^ is a scaling term that determines the overall strength of the noise (and can be set to 1). During training, the noise added at each step of the forward process is defined by a noise schedule. There are many noise schedules that can be used, e.g. a cosine-based noise schedule; in general the added noise defines a signal to noise ratio that monotonically increases with t (decreases with n).

[0050] The (forward) diffusion process can be described as: ^ଶ^ ^^ ^^^ ^^ |^^ ൌ ^^ ^^ ;^^^^ ,^^ ^^^^ (2)൫൯^ ^ି^ ^ ^ି^ ^ି^^ି^^ ^ ^ ^ where ^^ ^^ ൌ ^^ ^^ ; 0, ^^^^^herein, following convention, a capital ^^ indicates a random variable and a lower case ^^ a realized value or sample). ^

[0051] In the diffusion process ^^ is a mean function that can be obtained from the denoising^^^^ ^^ ^ output 112 directly, or indirectly, e.g. as ^^^^ ൌ ^^ ^ ^^^^^^ .^ ^ ^௧^ଶ

[0052] In this example ^^ ൌ ^^^^^^ . Merely as an the variance ^^ ൌ ^^ , and√^ ^ି௧ೖା^గ^ଶ ^^ ∙^^ is^^^ ^ ^

[0053] In implementations a draft model process, using a draft model p, can be described as:Attorney Docket No.: MX402430WO ^^^^^^^|^^^ି^^ ൌ ^^൫^^^;^^^ି^^^^^:^ି^^,^^^ଶି^^^^^൯ (3)

[0054] This is a reversestates, ^^^:^ି^are used to predict a less noisy state ^^^(where here the notation ^^^:^ି^refers to a sequence from ^^ downto ^^ െ 1). Later the notation ^^^ା^ ൌ ^^^∙ |^^^:^ା^ି^^ for ^^ ∈ ^^ refers to performing a sequenceof L denoising steps (a sequence of length ^^^), generating a draft representation for step ^^ ^^^ from previously obtained representations starting from step n. Note that here (and in the later example algorithms) although n is increasing this corresponds to decreasing t, i.e.moving from ^^ ൌ 1 towards ^^ ൌ 0.

[0055] Implementations of the described techniques use a draft model to generate a sequence of L draft states, i.e. data item representations, with successively reduced levels noise states, using the draft model p. These are processed to determine whether or not to accept each draft state (representation). In general, generating the draft states using the draft model is cheaper, i.e. requires less computation, than generating states (representations) using the target model. The decision to accept or reject a draft state can be made using the target model q. In implementations the drat states are processed in parallel to determine whether they are accepted or rejected, thus facilitating faster sampling than with a serial process.

[0056] If a draft state (representation) is rejected then a new state, i.e. data item representation is obtained. This should be one that appears to have been drawn from the distribution of the target model (at least approximately). However it can be shown that simply re-sampling from the target model leads to an incorrect distribution, broadly because the preceding accept / reject process can lead to over-sampling from the draft model domain. A calculation can be made of an alternative distribution that the new state (representation) should be sampled from, but it turns this is of little or no practical value because of the computational cost of sampling from such a distribution.

[0057] Instead it has been found that the new state (representation) can be computed deterministically as a function of the rejected draft state (representation). This process can use just one (additional) evaluation of the target model to obtain the new state. The exact form of the deterministic computation depends on how the accept / reject decision is made and a few examples are given. Broadly, however, it depends on a combination (sum) of a first term that is the target model mean vector for the previous draft representation (^^^; ^^^^ି^൫^^^^ି^൯), and a second term, dependent on the rejected draft state (representation), that can be viewed as a correction. Some implementations of the model return a sequence ofAttorney Docket No.: MX402430WO exact samples from the target model, but this can be relaxed to obtain a higher probability of accepting proposals from the draft model.

[0058] There are many ways that a draft sequence, of draft states (representations), can be obtained. One way is to use an independent draft model, typically one that is “cheaper” than the target model (and which need not have the same architecture as the target model). Another way is to use a single evaluation of the target model to generate the draft sequence, rather than evaluating the target model at every successive denoising step (as conventional in a reverse diffusion process). In this case the draft model can be viewed as a frozen version of the target model, and the draft states can be sampled in parallel. Optionally a correction for the frozen model can be made to the draft states. A variant of this allows the draft sequence to be computed from the result of processing a previous draft sequence, in which case only one function evaluation of the target model is needed, to initialize the process. In another approach, the predictions of multiple draft models can be combined; and so forth. In implementations distributions, e.g. Gaussian distributions, of samples from the target model and from the draft model at the same index or time step have the same variance (which follows automatically if the draft sequence is obtained from a version of the target model).

[0059] FIG. 2 conceptually illustrates such a speculative sampling process for continuous- valued data. The draft model, p, e.g. a version of the target diffusion model neural network 110 or the draft diffusion model neural network 120, generates a draft sequence 210, that is a sequence of draft representations (states). Starting from an initial, noisy representation 200 of the data item, the draft representations (states) are processed or “verified” 220, here in parallel, to determine whether to accept or reject each draft representation for inclusion in the sequence of representations 230 of the data item. Successive draft sequences are processed to build up the sequence of representations 230 of the data item until a final representation of the data item 240 is obtained. The combined Gaussian distributions pictorially illustrate moving a rejected draft representation to the target model distribution (using reflection- maximal coupling).

[0060] FIG.3 is a flow diagram of an example process for generating a data item represented by a plurality of continuous-valued elements. The process of FIG.3 may be implemented by one or more computers in one or more locations; for convenience the process is described with reference to FIG.1.

[0061] At step 300 the process obtains an initial, noisy representation of the data item, e.g. by sampling values in the representation from a distribution such as a Gaussian distribution.Attorney Docket No.: MX402430WO

[0062] The process generates a sequence of representations of the data item representing data items with successively reduced levels of noise (in a space of the representation), until obtaining a final representation of the data item (step 310). The representation of the data item can be in a space of the data item, i.e. in an output space of the generated data item, or it can be a representation in a latent space, in which case the final representation of the data item can be decoded to obtain the (final) generated data item.

[0063] The process involves generating, from a current representation of the data item in thesequence (^^^), a draft sequence of draft representations of the data item (^^^^, ^^^^ା^, …) (step302). The current representation of the data item can be the first draft representation in the draft sequence. Each successive draft representation of the data item in the draft sequence represents a data item with a successively reduced level of noise compared to a previous representation of the data item in the draft sequence. The draft sequence is generated using the draft model, p, noting that this can be a version of the target model, q (which has the advantage that there is no need to obtain a separate draft model).

[0064] For each successive draft representation of the data item in the draft sequence after the first (^^^^) the process determines whether to accept or reject the draft representation for inclusion in the sequence of representations of the data item, also referred to herein as verifying the drat representation (step 304), until the draft representation is rejected. In a case where the draft representation is rejected the process determines a new representation for inclusion in the sequence of representations of the data item (step 306). In particular the new representation is computed deterministically, as a function of the rejected draft representation.

[0065] The process continues iteratively (308): The new representation becomes the first draft representation of the next draft sequence, until the final representation is obtained (step 310).

[0066] Determining whether to accept or reject a draft representation can involve determining a target model mean vector for the representation. Generally, the target model mean vector is dependent upon the target model value for each element of the draft representation (^^୯^ౡ^^^^^). In implementations determining the respective target model mean vector canweighting the denoising output according to a noise schedule that depends on the index ^^ of ^^^of the draft representation in the sequence of representations of the data items (the weight is a decreasing function of ^^, which runs over the entire sequence not just the draft sequence).Attorney Docket No.: MX402430WO

[0067] In implementations, for each draft representation of the draft sequence, in parallel, the process determines a respective target model mean vector, e.g., as ^^୯୩^y^^ ൌ ^^୩ ^ γ^^୯^ౡ^y^^ for the next step ^^ ^ 1. As previously described, this can berepresentation (^^^) using the trained target diffusion modelthe denoising output 112, a vector ^^୯^ౡ^^^^^ that provides the target model value for each element of the draft representation.

[0068] In implementations offor each successive draft representation of the data item in the draft sequence after the first (e.g., for ^^ ൌ ^^ ^ 1:^^^), determining whetherto accept or reject the draft representation (^^^^) for inclusion in the sequence of representations of the data item is based on i) the previous draft representation (e.g. ^^^^ି^, or ^^^^ for ^^ ൌ ^^ ^1), and ii) the target model mean vector for the previous draft representation (e.g. ^^୯ ^୩ି^ ൫^^^^ି^൯; ^^ ). As previously noted, in implementations the current representation of thedraft representation in the draft sequence (e.g. ^^^^ ← ^^^). The process ofto accept or reject the draft representation is continued until the draft representation (^^^^) is rejected or the end of the draft sequence is reached.

[0069] As described further later, there are various particular ways in which the accept / reject, “verification” decision can be made. For example the decision can be made in a space of the representation, or explicitly in a latent space of the representation, or implicitly in a latent space of the representation.

[0070] In response to the draft representation being rejected the process then determines new representation (for ^^^^), which becomes the next current representation for inclusion in the sequence of representations of the data item, e.g. based on the draft representation and the target model mean vector. The next ^^^^, i.e. the first draft representation of the next draft sequence, can be set to be this new representation (until a final representation has been obtained).

[0071] The parallel determination of the respective target model mean vector for each draft representation of the draft sequence can involve implementing the trained target diffusion model neural network 110, and optionally draft diffusion model neural network 120 (where present), on a hardware device, e.g. a hardware accelerator, capable of parallel processing. The trained parameters, e.g. weights, of the trained target diffusion model neural network 110, and optionally those of the draft diffusion model neural network 120 (where present), can be loaded into a memory of the hardware device. The respective target model mean vector for each draft representation of the draft sequence can then be determined in parallelAttorney Docket No.: MX402430WO using the hardware device, e.g. in a batch. As used herein a hardware accelerator is generally characterized by the ability to handle large matrix computations, such as matrix multiplication, in hardware, e.g. in parallel. Some examples include GPUs (graphics processing units) and TPUs (tensor processing units).

[0072] In some implementations generating the draft sequence of draft representations of the data item involves determining, for each successive draft representation of the data item in the draft sequence after the first, a draft mean vector (^^^). (This is generated for the draft representation, but can be used to obtain the next draft representation, e.g. by sampling from a distribution parameterized by the draft mean vector). In implementations the draft mean vector (^^^) is dependent upon a value, in particular a mean value, for each element of the previous draft representation. Some examples of how the drat mean vector can be determined are given later.

[0073] Determining whether to accept or reject the draft representation (^^^^) based on the previous draft representation and the target model mean vector for the previous draft representation can involve determining whether to accept or reject the draft representation dependent on a difference between the draft mean vector (^^^) for the previous draft representation and the target model mean vector (^^^) for the previous draft representation. The difference can be scaled by (divided by) the standard deviation (^^) of the noise in the draft representation. This can be determined according to a nose schedule that defines a variation of noise or SNR with the index of the draft representation in the sequence. That is implementations of the described techniques determine whether to accept or reject the draft representation dependent on ^^^െ^^^, or on ^^^^െ^^^^ / ^^. In general the target model or draft mean vector for a representation parameterizes a Gaussian that can be sampled to get the next representation.

[0074] In some implementations, e.g. where verification is performed in a space of the draft representation, determining whether to accept or reject the draft representation (^^^^) can involve accepting the draft representation with a probability that depends on, e.g. is a minimum of, 1 and, a ratio of a first multivariate Gaussian probability density and a second multivariate Gaussian probability density (the first multivariate Gaussian distribution divided by the second multivariate Gaussian distribution).

[0075] The second multivariate Gaussian probability density can be evaluated at a valuedependent on the draft representation (^^^^), and the first multivariate Gaussian probabilitydensity can be evaluated at a value dependent on the draft representation (^^^^) with an offsetAttorney Docket No.: MX402430WO dependent on ^^^ െ ^^^ or ^^^^ െ ^^^^ / ^^. Each of these distributions can have zero mean,but this is not necessary for the described techniques; each can have a variance that is the identity matrix or a scaled version of the identity matrix. In general the first and second multivariate Gaussian probability distributions have the same variance.

[0076] For example, the second multivariate Gaussian probability density can be evaluated for a normalized value (^^^ି^) based on the draft representation (^^^^), in particular dependenton the difference between the draft representation and the draft model’s conditional mean (^^^^ ^^ െ ^^^ି^ ൫^^^^:^ି^൯^ / ^^^ି^), which can be referred to as a normalized residual of the draftmultivariate Gaussian probability density can be evaluated for (i.e.vector comprising) a sum, i.e. an elementwise sum, of i) the normalized version (residual), ^^^ି^, of the draft representation and ii) an offset vector, Δ୩ି^, dependent on the difference between the draft mean vector (^^^^ି^) for the previous draft representation (at step ^^ െ 1, and conditioned on the draft sequence ^^: ^^ െ 1) and the target model meanvector (^^^^ି^) for the previous draft representation. That is, this term can be evaluated using the target model, and just the previous draft representation at ^^ െ 1.

[0077] Broadly, the acceptance (verification) test determines whether the sample would more likely under the target model than under the draft model, if so accepting the sample.

[0078] As previously described, if the draft representation is rejected a new representation can be determined based on the draft representation and the target model mean vector, in particular by computing the new representation from a deterministic combination of the draft representation and the target model mean vector for the previous draft representation.

[0079] The previously described normalized value (residual), Z, can be referred to as an offset version of the draft representation, offset by the draft mean vector (^^^), e.g. by subtracting ^^^from the draft representation. Determining the new representation can involve determining this, and also a sum of the target model mean vector (evaluated for the previous draft representation) and the offset version of the draft representation (e.g. scaled by a weight vector dependent on Δ), to obtain the new representation.

[0080] The algorithms below are a particular example of an implementation of the process of FIG.3 for DDPM. The draft sequence is generated sequentially and the evaluation of the target model mean vectors is performed in parallel, according to the first example algorithm. The determinations of whether to accept or reject the draft representations are also performed in parallel, by the ^^^^^^^^^^^^^^^^^^ algorithm. The ^^^^^^^^^^^^^^^^^^ algorithm illustrates one way in which the normalized value (residual), Z, and difference in means (offset vector), Δ, canAttorney Docket No.: MX402430WO be determined. As written in the final line, the ^^^^^^^^^^^^^^^^^^ algorithm outputs a new representation ^^ that has the target distribution, ^^, it is computed deterministically rather than being sampled from the target model (which sampling would not provide a newrepresentation ^^ with the target distribution). At acceptance, ^^ ൌ ^^^ , and the ^^^^^^^^^^^^^^^^^^algorithm also maximizes the probability that the new representation ^^ ൌ ^^^ .A ^^^^^^^^^^^^^^^^^^൫^^,^^,^^^൯ algorithm, that computes a new representation ^^ for two Gaussians with the same covariance is below; some other example ^^^^^^^^^^^^^^^^^^ algorithms are given later.

[0081] In the above example the criterion of having the same covariance implies that the draft and target samplers introduce the same amount of noise at each step. A high acceptance rate is generally desirable, and the acceptance rate is improved when the draft model is a good approximation of the target model. The acceptance rate can be improved by having aAttorney Docket No.: MX402430WO smaller step size, but at the expense of a larger total number of steps. In general the acceptance rate is higher at the start of the denoising process than at the end. The acceptance rate depends on ^^, the overall strength of the noise. Both high and low values of ^^ can decrease the acceptance rate, which can be treated as a hyperparameter to optimize; or ^^ can be set to 1.

[0082] A detailed example algorithm that combines the above two algorithms follows:In this example the target model q can be given by e.g. equation (2) above. The draft model p can be given, e.g., by equation (3); some particular examples are described later.

[0083] In some implementations the determination of whether to accept or reject the draft representation is modified by a temperature parameter. For example, the above described the first and multivariate Gaussian probability densities can each have a variance defined by the identity matrix scaled by a temperature parameter, which can be (but need not be) greater than one. This can be used to tune an acceptance rate of the draft tokens. For example,setting ^^ ^ 1 increases the probability of accepting the draft representation, with thedrawback of no longer sampling from the correct target distribution, instead moving theAttorney Docket No.: MX402430WOdistribution of ^^ towards ^^^ . An example ^^^^^^^^^^^^^^^^^^൫^^,^^,^^^൯ algorithm for two Gaussianswith the same covariance, and that includes a temperature parameter is below:

[0084] In some implementations the acceptance criterion is based on so-call maximalcoupling of ^^ and ^^^ in a latent space (maximal coupling maximizes the probability that ^^ ൌ^^ while ^^~^^ and ^^~^^). This can also increase the acceptance rate compared with the ^^^^^^^^^^^^^^^^^^ algorithm (without temperature), also at the expense of no longer sampling ^^ from the correct target distribution q. As one example a linear operator ^^ can be used to define the draft and target models in a latent space and provide a maximal coupling of ^^^and^^^^ as below:where the target model q and draft model p are each ^^ െdimensional and ^^ is a ^^ ൈ ^^ squarematrix. This acceptance (rejection) algorithm can be used in place of the ^^^^^^^^^^^^^^^^^^ algorithm given previously.

[0085] This algorithm can also be re-written as:Attorney Docket No.: MX402430WOwhere ^^ற is the inverse of ^^ (^^ற ൌ ^^்^^^^^்^ି^, or ^^ற ൌ ^^ି^ if ^^ isinvertible). This implicitly provides a maximal coupling of ^^^and ^^^^given some matrix A.

[0086] A non-linear version of this algorithm with “encoder” and “decoder” mappings ^^^∙^and ^^^∙^ (and ^^൫^^^^^^൯ ^ ^^) is below. This can recover example ^^^^^^^^^^^^^^^^^^ algorithmwith temperature given above (if ^^^^^^ ൌ ^^^^, Δ∗ ൌ ^^൫^^^Δ^൯ ൌ ^^ற^^Δ, Δ∗ ൌ Δ / ^^).

[0087] Draft models.

[0088] As previously noted, there are many ways in which a draft sequence of draft representations can be obtained. In general the draft model can be described mathematicallyas ^^൫^^^^^ା^:^^|^^^൯ ൌ ∏^ୀ^ା^ ^^^^^^|^^^:^ି^^ where ^^^ ൌ min^^^ ^ ^^,^^^, ^^ is the length of themodel, i.e. a neural network ^^^௧, e.g. on the same training dataset as the target model. Such a draft model can be cheaper to evaluate than the target model, i.e. requiring less compute and / or memory to process a model input to generate a model output. Such a draft model can satisfy^^^^^^|^^^:^ି^^ ൌ ^^^^^^|^^^ି^^. The draft mean vector can be obtained from the draft model,e.g., as: ^^^^^^^^^ ൌ ^^^ ^ ^^^^^௧ೖ^^^^^Attorney Docket No.: MX402430WOwhere ^^^ ൌ √^^^^^^ ), where ^^^is cheaper to evaluate than ^^^(and where ^^^^ି௧ೖ ௧ ௧ ^^^^^^couldbe written more as ^^^^^^^^:^^). For example, the targetmayof the draft representation, athe element and the target model value for the element.

[0090] For example, in some implementations generating the draft sequence of draft representations of the data item comprises, for successive draft representations in the draft sequence, processing a current draft representation in the draft sequence (and optionally also an index of the current draft representation) using a trained draft diffusion model neural network to generate a draft model denoising output (^^^௧^^^^^^) for each element of the current draft representation. The next draft representation in the draft sequence can be determined from the draft model denoising output, e.g. by adding the denoising output to the current draft representation and then either adding noise to this or using this to define the mean of an e.g. Gaussian distribution which is sampled from to obtain the next draft representation; or using the denoising output as the prediction of the next draft representation where the diffusion model neural network is trained to output a de-noised representation.

[0091] In some implementations of this approach a plurality of such trained draft diffusion model neural networks is maintained, and the next draft representation in the draft sequence is determined from a convex combination of the draft model denoising outputs from the plurality of trained draft diffusion model neural networks. The weights of the draft model denoising outputs in the convex combination may be defined as hyperparameters, or learned.

[0092] Another way to obtain the draft model is to use a “frozen” version of the target model. That is, in some implementations generating the draft sequence of draft representations from the current representation of the data item in the sequence is done by predicting each draft representation of the draft sequence based on a single prediction from the trained target diffusion model neural network. This does not require training a separate draft model.

[0093] This can be done, e.g., by processing one of the draft representations in the draftsequence, e.g. the first draft representation in the draft sequence (^^^^ ← ^^^) using the trainedtarget diffusion model neural network to generate the denoising output (e.g. ^^୯^^^^^^^). Thesuccessive draft representations in the draft sequence can be obtained bya next draft representation in the draft sequence from a combination (an elementwise combination) of a current draft representation and the denoising output for the first draft representation. In implementations an index-dependent weight (e.g. ^^), defined by a noise schedule for theAttorney Docket No.: MX402430WO diffusion (denoising) process can be used to weight ^^୯^^^^^^^) according to the index ^^. Thiscan be written as ^^^^^^|^^^:^ି^^ ൌ ^^^^^^|^^^, ^^^ି^^.

[0094] In a specific example the draft mean vector can be obtained as: ^^^^^^^^:^^ ൌ ^^^ ^ ^^^^^௧^^^^^^with ^^^ ൌ √^^^^^^ . That ^^^^^^ ^ same) ^^^^ି௧ೖ ^ ௧^^^^^^ for each draftrepresentation, and the target model neural network 110, need only be queried once tosequence. This is computationally inexpensive and allows parallel sampling of the draft sequence. That is, in some implementations of this type of approach the successive draft representations in the draft sequence are determined in parallel.

[0095] Optionally, ^^^௧^^^^^^^ can be replaced by ^^௧^^^^^^^, which enables the draft sequence to be computed any calls to the target model, since ^^^௧^^^^^^^ has already been computed for verification of the final draft representation of the previous draft sequence. Just one call to the target model is needed at the start, for the complete data item generation process, for initialization (specifically, to compute ^^^௧^^^^^^^). Another option is to use ^^^௧^షభ^^^^ି^^inplace of ^^^௧^ ^^^^^^ when ^^^^ ് ^^^.

[0096] Towards the end of the denoising process the draft representations obtained, based on a draft model that is a frozen target model, can differ significantly from the target model distribution. Thus, in some implementations a correction neural network can be maintained, to generate a correction output for correcting the draft representations.

[0097] More specifically applying this correction can involve processing the first draft representation and the current draft representation using a correction neural network to generate a correction output (vector), and then adjusting the denoising output for the first draft representation using the correction output (e.g. by elementwise addition or subtraction) to obtain an adjusted current draft representation. A next draft representation, i.e. following the current draft representation in the draft sequence, in the draft sequence can then be generated based on a combination of the adjusted current draft representation and the denoising output for the first draft representation.

[0098] For example, the draft mean vector can be obtained as ^^^^^^^^:^^ ൌ ^^^ ^^^^^^^௧^ ^^^^^ ^ ^^௧^,௧ೖ^^^^, ^^^^^ where ^^௧^,௧ೖ^^^^,^^^^ denotes the correction neural network. Thecorrection neural network can be trained on a loss that depends on ฮ^^^௧ ^^^ ^ െ ^^ ^ೖ ^ ௧^ ^^^^^ െ^^௧^,௧ೖ^^^^, ^^^^ฮଶ , and therefore predicts ^^^ ^^^ ^ െ ^^^௧ೖ ^ ௧^^^^^^.Attorney Docket No.: MX402430WO

[0099] As yet another example, the draft sequence can be obtained using Picard iterations, e.g. by using a ParaDiGMS sampler as described in Shih et al., arXiv:2305.16317v3, October 2023. Shih et al. also describes a parallel sampling technique that can be combined with the technique described herein (whether or not the ParaDiGMS sampler is used).

[0100] Representations in latent space.

[0101] The final representation of the data item can be in a space of the data item, e.g., the elements of a representation may comprise values of image pixels, or amplitude values of an audio signal. For example when the data item comprises an image the space of the data item can be a pixel space and the representations of the data item can have elements comprising values for each pixel that specify color values, e.g., RGB values or another color encoding scheme. An example of a trained diffusion model neural network that operates in pixel space is the Imagen model.

[0102] Alternatively the final representation of the data item can be in a latent space of the data item. For example the representations of the data item can be representations in a latent space, e.g. a learned, latent space, and the continuous-valued elements can then represent the data item in the latent space. For example when the data items are images the elements of the representations can be latent values rather than color values. The method can then involve processing the final representation of the data item using a decoder to generate the (output) data item.

[0103] The decoder can be a decoder that is associated with the trained target diffusion model neural network, e.g., one that has been pre-trained in an auto-encoder framework. For example the trained target diffusion model neural network can have been trained in a system that uses an encoder neural network, e.g., one that has been pre-trained jointly with the decoder in the auto-encoder framework, to encode target data items in the output space to generate target denoising outputs for the target diffusion neural network in the latent space. A diffusion model neural network that operates in latent space is sometimes referred to as a latent diffusion model.

[0104] Example results.

[0105] In general using a frozen target model as the draft model is more efficient than using an independent draft model in terms of compute, specifically in terms of a number of function evaluations (NFE) of the target model, where a function evaluation is a call to the target model with a batch of data (irrespective of the batch size). In Tables 1 and 2 below a frozen target model was used as the draft model.Attorney Docket No.: MX402430WO

[0106] Table 1 below shows example results for a target diffusion model neural network 110 with a U-Net architecture (4-levels, with a residual 8-head attention block; conditioned on class labels and a sinusoidal time embedding, summed and added to feature maps via adaptive normalization). This was trained on the CIFAR10 image dataset, with various overall noise and temperature parameter values; the best results are shown in bold. The table shows FID (Fréchet Inception Distance) – lower is better, indicating a more realistic image; and IS (Inception Score) – higher is better, evaluating image quality. The column headed “Speculative” is for implementations of the described techniques.

[0107] It can be seen that the described techniques can achieve a x2 to x3 speedup for the same image quality compared to conventional approaches. For example, the described techniques can achieve an FID of 2.34 with 35 calls to the target model, whereas a conventional approach requires 100 calls to the target model to reach an FID of 2.45.

[0108] Table 2 below shows example results for a target diffusion model neural network 110 with an architecture comprising a sequence of MLPs (MultiLayer Perceptrons), trained to perform a robotics task, specifically on the PushT dataset (Chi et al., “Diffusion policy: Visuomotor policy learning via action diffusion” , arXiv:2303.04137v5, March 2024).Attorney Docket No.: MX402430WO

[0109] It can be seen that the described techniques can reduce the number of calls to the target model whilst preserving quality of the actions. For example, with 5 calls to the target model the described techniques can achieve a reward of 0.903, compared with 0.876 with a conventional approach.

[0110] Example data items.

[0111] The data item can, e.g., comprise an image (values of pixels of an image), or audio data representing an audio signal (e.g. defining the audio signal waveform in the time or time- frequency domain), or action selection data for controlling a mechanical agent such as a robot to perform a task, optionally conditioned on a conditioning input that characterizes one or more desired properties for the data item. For example in the case of image or audio data a conditioning input can comprise text that describes the data item, and in the case of action selection data for controlling a mechanical agent to perform a task a conditioning input can comprise an observation of a real-world environment, such as an image observation, in which the mechanical agent performs the actions to accomplish the task (e.g. moving an object or navigation to a destination in the environment).

[0112] As another example, the data item may represent one or more chemical molecules such as one or more proteins or ligands, e.g. as a point cloud. The conditioning input can define one or more characteristics of the chemical molecule(s), e.g. in terms of its / their physical or chemical structure or properties. The data item may be used to determine a 3D structure of the chemical molecule(s), e.g. to identify one or more binding sites of or for a ligand such as a drug. This may be used as part of a screening process to identify one chemical molecule that binds to another. Such a screening process may involve evaluating an interaction of one or more candidate ligands with the structure of a target, e.g. a target protein, and then selecting one or more of the candidate ligands dependent on a result of the evaluation. For example the target may comprise a receptor or enzyme, and the ligand may be an agonist or antagonist of the receptor or enzyme. The ligand may be a drug or a ligand of an industrial enzyme. Such a process may also involve synthesizing a molecule identified by the screening process, e.g. the ligand, and optionally also testing activity, e.g. biological activity, of the molecule, e.g. ligand, in vitro and / or in vivo.

[0113] As another example, the data item may represent the output of a scientific or medical instrument, e.g. the output of an electrocardiograph or of a body scanner such as an MRI machine. The conditioning input can define one or more characteristics of the data item such as whether it represents a signal from a normal body or from a diseased body. The data itemAttorney Docket No.: MX402430WO may then be compared with a corresponding data item obtained from a patient, and a comparison made to identify the likely presence or absence of a disease.

[0114] As a further example the data item may represent text, e.g. where the text is embedded into a continuous space. The conditioning input can define a request, image, or audio, that characterizes the text to be generated, e.g. a response to the request, or a description of the image of objects in the image (e.g. in a captioning task), or text representing the audio (e.g. in a speech recognition task).

[0115] In general the system can be configured to generate any of a variety of data items, optionally conditioned on any of a variety of conditioning inputs. In some implementations the system can use classifier-free guidance to generate a data item conditioned on a conditioning input. Then, a draft representation of the data item can be obtained from a weighted combination of the conditional denoising output and an unconditional denoising output from the trained target diffusion model neural network.

[0116] For example, the system can be configured to generate audio data, e.g., a waveform of audio or a spectrogram, e.g., a mel-spectrogram or a spectrogram where the frequencies are in a different scale, of the audio.

[0117] In this example, the conditioning input can be text or features of text that the audio should represent, i.e., so that the system serves as a text-to-speech machine learning model that converts text or features of the text to audio data for an utterance of the text being spoken.

[0118] As another example, the conditioning input can identify a desired speaker for the audio, i.e., so that the system generates audio data that represents speech by the desired speaker.

[0119] As another example, the conditioning input can characterize properties of a song or other piece of music, e.g., lyrics, genre, and so on, so that the system generates a piece of music that has the properties characterized by the conditioning input.

[0120] As another example, the conditioning input can specify a classification for the audio data into a class from a set of possible classes, so that the system generates audio data that belongs to the class. For example, the classes can represent types of musical instruments or other audio emitting devices, i.e., so that the system generates audio that is emitted by the corresponding class, types of animals, i.e., so that the system generates audio that represent noises generated by the corresponding animal, and so on.

[0121] As another particular example, the data item can be an image. A used herein an “image” can be a still or moving image. The system can perform conditional image generation by generating the intensity values of the pixels of the image.Attorney Docket No.: MX402430WO

[0122] In this particular example, the conditioning input can be a sequence of text and the output data item can be an image that describes the text, i.e., the conditioning input can be a caption for the output image.

[0123] As yet another particular example, the conditioning input can be an object detection input that specifies one or more bounding boxes and, optionally, a respective type of object that should be depicted in each bounding box.

[0124] As yet another particular example, the conditioning input can specify an object class from a plurality of object classes to which an object depicted in the output image should belong.

[0125] As yet another particular example, the conditioning input can specify an image at a first resolution and the output data item can include the image at a second, higher resolution.

[0126] As yet another particular example, the conditioning input can specify an image and the output data item can comprise a de-noised version of the image.

[0127] As yet another particular example, the conditioning input can specify an image including a target entity for detection, e.g. a tumor, and the output data item can comprise the image without the target entity, e.g. to facilitate detection of the target entity by comparing the images.

[0128] As yet another particular example, the conditioning input can be a segmentation that assigns each of a plurality of pixels of the output image to a category from a set of categories, e.g., that assigns to each pixel a respective one of the category.

[0129] More generally, the task can be any task that outputs continuous data conditioned on a conditioning input. For example, the output can be an output of a different sensor, e.g., a lidar point cloud, a radar point cloud, an electrocardiogram reading, and so on, and the conditioning input can represent the type of data that should be measured by the sensor. Where a discrete output is desired this can be obtained, e.g. by thresholding.

[0130] More generally the data item can be a sample from an unnormalized probabilitydensity ^^^^^^ that can be described by an energy function ^^^^^^, as ^^^^^^ ∝ exp൫െ^^^^^^൯. Thetechniques are useful where the energy function can be evaluated pointwise, but where each evaluation is computationally expensive. Applying the described techniques, a cheaper proxyenergy function, ^^^^^^^ ^ ^^^^^^ can be used.

[0131] A standard approach to sample from ^^^^^^ is to use Langevin diffusion, defined by^^^^௧ ൌ െ∇^^^^^௧^^^^^ ^ √2^^^^௧ where ^^௧ is a Brownian motion. This can be implementedusing an iterative, discrete time approach, corresponding to the above described target model:Attorney Docket No.: MX402430WO^^^ା^ ൌ ^^^ െ ^^∇E^^^^^ ൌ ^2^^^^^ for step size ^^ ^ 0 and where ^^^~^^^0, ^^^^^; the draftwith E^^^^^^ E^^^^^. A frozen target model can be used asi.e. setting ∇^E^^^^ା^^ ൌ ∇E^^^^^ for ^^ ൌ 1, … ,^^^.Note that in unadjusted Langevin sampling the level of noise typically does not decrease with time (e.g. so that state space is explored effectively), e.g. it can stay constant.

[0132] For example, the previously described combined algorithm can be written, in this case, as:

[0133] probability densities described by proxy energy functions are used in many areas of engineering and applied science. Some example applications follow; the described techniques may be used with these and other applications.

[0134] In an image generation or processing application the unnormalized probability density data item can define a distribution over real-world images and the energy function can be defined by a neural network that assigns low energies to realistic images, and vice-versa.

[0135] In robotics the unnormalized probability density data item can define a distribution over mechanical robot action trajectories, e.g. a series of actions given respective observations of a real-world environment. The energy function can define, a trajectory cost, e.g. based on distance, or energy use, or collision risk and so forth, thus favoring low cost,Attorney Docket No.: MX402430WO safe trajectories when performing an action on an object or navigating in the real-world environment.

[0136] In neural network training the unnormalized probability density data item can define a distribution over neural network weights in a Bayesian neural network; the energy function can represent a negative log likelihood of the observed data given the parameters, optionally with a negative log prior; to bias the weights towards explaining the data.

[0137] In a drug discovery application, the unnormalized probability density data item can define the conformation of a protein, or ligand, or of a bound complex of the two. The energy function can define a structural or binding energy, and can optionally have other terms such as toxicity, synthesizability, and so forth. This can be used to find a drug for synthesis, by identifying a ligand that binds to a target protein, or vice versa; the drug can then be synthesized.

[0138] There are many other applications, e.g. in structural engineering design, climate modelling, geophysical modelling, and so forth.

[0139] As previously explained, the described techniques do not rely on any particular training techniques or datasets and can be used with a trained diffusion model neural network 110 (or draft model). The diffusion model neural network 110 (or draft diffusion model neural network 120) can be trained using any conventional diffusion model training process using any suitable denoising objective. Nonetheless for completeness a training method is described below.

[0140] An example training method can involve sampling (i) a target data item from a set of training examples, (ii) a corresponding conditioning input for the data item from the set of training examples, (iii) a time step t for the training, e.g., uniformly at random from a time interval, and (iv), and noise ^^ from a noise distribution e.g. a Gaussian noise distribution. A noisy data item xt can be generated by combining the target data item x0 with the sampled noise ^^ in accordance with the sampled time step t, e.g., by setting the noisy data item xt= αtx0+ σt^^. The denoising neural network can process an input that includes the noisy data item xt, data specifying the time step, and the conditioning input to generate the denoising output. An error between the denoising output and the target denoising output, e.g. an MSE, L1 or L2 error, can be used to train the neural network, e.g., by determining and backpropagating gradients of the error and updating the parameters of the neural network using an optimizer such as Adam. Any suitable training dataset can be used, generally one with examples of the data items that are to be generated, e.g. ImageNet or the Open ImagesAttorney Docket No.: MX402430WO Dataset (Google) for images, LibriTTS-R (www.openslr.org / 141 / ) for audio (speech), and so forth.

[0141] In this specification, the term "configured" is used in relation to computing systems and environments, as well as computer program components. A computing system or environment is considered "configured" to perform specific operations or actions when it possesses the necessary software, firmware, hardware, or a combination thereof, enabling it to carry out those operations or actions during operation. For instance, configuring a system might involve installing a software library with specific algorithms, updating firmware with new instructions for handling data, or adding a hardware component for enhanced processing capabilities. Similarly, one or more computer programs are "configured" to perform particular operations or actions when they contain instructions that, upon execution by a computing device or hardware, cause the device to perform those intended operations or actions.

[0142] The embodiments and functional operations described in this specification can be implemented in various forms, including digital electronic circuitry, software, firmware, computer hardware (encompassing the disclosed structures and their structural equivalents), or any combination thereof. The subject matter can be realized as one or more computer programs, essentially modules of computer program instructions encoded on a tangible non- transitory storage medium for execution by or to control the operation of a computing device or hardware. The storage medium can be a storage device such as a hard drive or solid-state drive (SSD), a storage medium, a random or serial access memory device, or a combination of these. Additionally or alternatively, the program instructions can be encoded on a transmitted signal, such as a machine-generated electrical, optical, or electromagnetic signal, designed to carry information for transmission to a receiving device or system for execution by a computing device or hardware. Furthermore, implementations may leverage emerging technologies like quantum computing or neuromorphic computing for specific applications, and may be deployed in distributed or cloud-based environments where components reside on different machines or within a cloud infrastructure.

[0143] The term "computing device or hardware" refers to the physical components involved in data processing and encompasses all types of devices and machines used for this purpose. Examples include processors or processing units, computers, multiple processors or computers working together, graphics processing units (GPUs), tensor processing units (TPUs), and specialized processing hardware such as field-programmable gate arrays (FPGAs) or application-specific integrated circuits (ASICs). In addition to hardware, aAttorney Docket No.: MX402430WO computing device or hardware may also include code that creates an execution environment for computer programs. This code can take the form of processor firmware, a protocol stack, a database management system, an operating system, or a combination of these elements. Embodiments may particularly benefit from utilizing the parallel processing capabilities of GPUs, in a General-Purpose computing on Graphics Processing Units (GPU) context, where code specifically designed for GPU execution, often called kernels or shaders, is employed. Similarly, TPUs excel at running optimized tensor operations crucial for many machine learning algorithms. By leveraging these accelerators and their specialized programming models, the system can achieve significant speedups and efficiency gains for tasks involving artificial intelligence and machine learning, particularly in areas such as computer vision, natural language processing, and robotics.

[0144] A computer program, also referred to as software, an application, a module, a script, code, or simply a program, can be written in any programming language, including compiled or interpreted languages, and declarative or procedural languages. It can be deployed in various forms, such as a standalone program, a module, a component, a subroutine, or any other unit suitable for use within a computing environment. A program may or may not correspond to a single file in a file system and can be stored in various ways. This includes being embedded within a file containing other programs or data (e.g., scripts within a markup language document), residing in a dedicated file, or distributed across multiple coordinated files (e.g., files storing modules, subprograms, or code segments). A computer program can be executed on a single computer or across multiple computers, whether located at a single site or distributed across multiple sites and interconnected through a data communication network. The specific implementation of the computer programs may involve a combination of traditional programming languages and specialized languages or libraries designed for GPGPU programming or TPU utilization, depending on the chosen hardware platform and desired performance characteristics.

[0145] In this specification, the term "engine" broadly refers to a software-based system, subsystem, or process designed to perform one or more specific functions. An engine is typically implemented as one or more software modules or components installed on one or more computers, which can be located at a single site or distributed across multiple locations. In some instances, one or more dedicated computers may be used for a particular engine, while in other cases, multiple engines may operate concurrently on the same one or more computers. Examples of engine functions within the context of AI and machine learning could include data pre-processing and cleaning, feature engineering and extraction, modelAttorney Docket No.: MX402430WO training and optimization, inference and prediction generation, and post-processing of results. The specific design and implementation of engines will depend on the overall architecture and the distribution of computational tasks across various hardware components, including CPUs, GPUs, TPUs, and other specialized processors.

[0146] The processes and logic flows described in this specification can be executed by one or more programmable computers running one or more computer programs to perform functions by operating on input data and generating output. Additionally, graphics processing units (GPUs) and tensor processing units (TPUs) can be utilized to enable concurrent execution of aspects of these processes and logic flows, significantly accelerating performance. This approach offers significant advantages for computationally intensive tasks often found in AI and machine learning applications, such as matrix multiplications, convolutions, and other operations that exhibit a high degree of parallelism. By leveraging the parallel processing capabilities of GPUs and TPUs, significant speedups and efficiency gains compared to relying solely on CPUs can be achieved. Alternatively or in combination with programmable computers and specialized processors, these processes and logic flows can also be implemented using specialized processing hardware, such as field-programmable gate arrays (FPGAs) or application-specific integrated circuits (ASICs), for even greater performance or energy efficiency in specific use cases.

[0147] Computers capable of executing a computer program can be based on general-purpose microprocessors, special-purpose microprocessors, or a combination of both. They can also utilize any other type of central processing unit (CPU). Additionally, graphics processing units (GPUs), tensor processing units (TPUs), and other machine learning accelerators can be employed to enhance performance, particularly for tasks involving artificial intelligence and machine learning. These accelerators often work in conjunction with CPUs, handling specialized computations while the CPU manages overall system operations and other tasks. Typically, a CPU receives instructions and data from read-only memory (ROM), random access memory (RAM), or both. The essential elements of a computer include a CPU for executing instructions and one or more memory devices for storing instructions and data. The specific configuration of processing units and memory will depend on factors like the complexity of the AI model, the volume of data being processed, and the desired performance and latency requirements. Embodiments can be implemented on a wide range of computing platforms, from small embedded devices with limited resources to large-scale data center systems with high-performance computing capabilities. The system may include storage devices like hard drives, SSDs, or flash memory for persistent data storage.Attorney Docket No.: MX402430WO

[0148] Computer-readable media suitable for storing computer program instructions and data encompass all forms of non-volatile memory, media, and memory devices. Examples include semiconductor memory devices such as read-only memory (ROM), solid-state drives (SSDs), and flash memory devices; hard disk drives (HDDs); optical media; and optical discs such as CDs, DVDs, and Blu-ray discs. The specific type of computer-readable media used will depend on factors such as the size of the data, access speed requirements, cost considerations, and the desired level of portability or permanence.

[0149] To facilitate user interaction, embodiments of the subject matter described in this specification can be implemented on a computing device equipped with a display device, such as a liquid crystal display (LCD) or an organic light-emitting diode (OLED) display, for presenting information to the user. Input can be provided by the user through various means, including a keyboard), touchscreens, voice commands, gesture recognition, or other input modalities depending on the specific device and application. Additional input methods can include acoustic, speech, or tactile input, while feedback to the user can take the form of visual, auditory, or tactile feedback. Furthermore, computers can interact with users by exchanging documents with a user's device or application. This can involve sending web content or data in response to requests or sending and receiving text messages or other forms of messages through mobile devices or messaging platforms. The selection of input and output modalities will depend on the specific application and the desired form of user interaction.

[0150] Machine learning models can be implemented and deployed using machine learning frameworks, such as TensorFlow or JAX. These frameworks offer comprehensive tools and libraries that facilitate the development, training, and deployment of machine learning models.

[0151] Embodiments of the subject matter described in this specification can be implemented within a computing system comprising one or more components, depending on the specific application and requirements. These may include a back-end component, such as a back-end server or cloud-based infrastructure; an optional middleware component, such as a middleware server or application programming interface (API), to facilitate communication and data exchange; and a front-end component, such as a client device with a user interface, a web browser, or an app, through which a user can interact with the implemented subject matter. For instance, the described functionality could be implemented solely on a client device (e.g., for on-device machine learning) or deployed as a combination of front-end and back-end components for more complex applications. These components, when present, canAttorney Docket No.: MX402430WO be interconnected using any form or medium of digital data communication, such as a communication network like a local area network (LAN) or a wide area network (WAN) including the Internet. The specific system architecture and choice of components will depend on factors such as the scale of the application, the need for real-time processing, data security requirements, and the desired user experience.

[0152] The computing system can include clients and servers that may be geographically separated and interact through a communication network. The specific type of network, such as a local area network (LAN), a wide area network (WAN), or the Internet, will depend on the reach and scale of the application. The client-server relationship is established through computer programs running on the respective computers and designed to communicate with each other using appropriate protocols. These protocols may include HTTP, TCP / IP, or other specialized protocols depending on the nature of the data being exchanged and the security requirements of the system. In certain embodiments, a server transmits data or instructions to a user's device, such as a computer, smartphone, or tablet, acting as a client. The client device can then process the received information, display results to the user, and potentially send data or feedback back to the server for further processing or storage. This allows for dynamic interactions between the user and the system, enabling a wide range of applications and functionalities.

[0153] While this specification contains many specific implementation details, these should not be construed as limitations on the scope of any invention or on the scope of what may be claimed, but rather as descriptions of features that may be specific to particular embodiments of particular inventions. Certain features that are described in this specification in the context of separate embodiments can also be implemented in combination in a single embodiment. Conversely, various features that are described in the context of a single embodiment can also be implemented in multiple embodiments separately or in any suitable subcombination. Moreover, although features may be described above as acting in certain combinations and even initially be claimed as such, one or more features from a claimed combination can in some cases be excised from the combination, and the claimed combination may be directed to a subcombination or variation of a subcombination.

[0154] Similarly, while operations are depicted in the drawings and recited in the claims in a particular order, this should not be understood as requiring that such operations be performed in the particular order shown or in sequential order, or that all illustrated operations be performed, to achieve desirable results. In certain circumstances, multitasking and parallel processing may be advantageous. Moreover, the separation of various system modules andAttorney Docket No.: MX402430WO 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 can generally be integrated together in a single software product or packaged into multiple software products.

[0155] Particular 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 can be performed in a different order and still achieve desirable results. As one 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.

[0156] What is claimed is:

Claims

Attorney Docket No.: MX402430WO CLAIMS 1. A computer-implemented method of generating a data item represented by a plurality of continuous-valued elements, comprising: obtaining an initial representation of the data item; and generating a sequence of representations of the data item, until obtaining a final representation of the data item, comprising: generating, from a current representation of the data item in the sequence, a draft sequence of draft representations of the data item, wherein the current representation of the data item is the first draft representation in the draft sequence; and for each successive draft representation of the data item in the draft sequence after the first: determining whether to accept or reject the draft representation for inclusion in the sequence of representations of the data item and, in a case where the draft representation is rejected, determining a new representation for inclusion in the sequence of representations of the data item.

2. The method of claim 1, comprising: generating the sequence of representations of the data item representing data items with successively reduced levels of noise; and wherein each successive draft representation of the data item in the draft sequence after the current representation represents a data item with a successively reduced level of noise.

3. The method of claim 1 or 2, wherein determining the new representation comprises computing the new representation as a deterministic function of the rejected draft representation.

4. The method of any of claims 1-3, further comprising: for each draft representation of the draft sequence in parallel determining a respective target model mean vector dependent upon a target model value for each element of the draft representation by processing the draft representation using a trained target diffusion model neural network to generate a denoising output that provides the target model value for each element of the draft representation; and whereinAttorney Docket No.: MX402430WO determining whether to accept or reject the draft representation for inclusion in the sequence of representations of the data item is based on the previous draft representation and the target model mean vector for the previous draft representation.

5. The method of claim 4, comprising determining the new representation for inclusion in the sequence of representations of the data item is based on the draft representation and the target model mean vector.

6. The method of claim 4 or 5, wherein determining, in parallel, the respective target model mean vector for each draft representation of the draft sequence comprises: implementing the trained target diffusion model neural network on a hardware device capable of parallel processing; loading trained parameters of the trained target diffusion model neural network into a memory of the hardware device; and determining the respective target model mean vector for each draft representation of the draft sequence in parallel using the hardware device.

7. The method of any of claims 1-6, wherein generating the draft sequence of draft representations of the data item comprises determining, for each successive draft representation of the data item in the draft sequence after the first, a draft mean vector dependent upon a mean value for each element of the previous draft representation; and wherein determining whether to accept or reject the draft representation based on the previous draft representation and the target model mean vector for the previous draft representation comprises determining whether to accept or reject the draft representation dependent on a difference between the draft mean vector for the previous draft representation and the target model mean vector for the previous draft representation.

8. The method of claim 7, wherein obtaining the initial representation of the data item comprises sampling values of the elements of the representation from a Gaussian distribution; and wherein determining whether to accept or reject the draft representation comprises accepting the draft representation with a probability that depends on a ratio of a first multivariateAttorney Docket No.: MX402430WO Gaussian probability density and a second multivariate Gaussian probability density, wherein the first multivariate Gaussian probability density is evaluated for an elementwise sum of the elements of a normalized version of the draft representation and an offset vector dependent on the difference between the draft mean vector for the previous draft representation and the target model mean vector for the previous draft representation, and the second multivariate Gaussian probability density is evaluated for the elements of the normalized version of the draft representation.

9. The method of claim 8, wherein the first and multivariate Gaussian probability densities each have a variance defined by the identity matrix scaled by a temperature parameter, wherein the temperature parameter is greater than one.

10. The method of any of claims 1-9, comprising determining the new representation from a deterministic combination of the draft representation and a target model mean vector for the previous draft representation.

11. The method of any of claims 1-10, wherein determining the new representation based on the draft representation and the target model mean vector comprises: determining an offset version of the draft representation, wherein the offset version of the draft representation is offset by a draft mean vector for the previous draft representation; and determining a sum of the target model mean vector for the previous draft representation and the offset version of the draft representation to obtain the new representation.

12. The method of any of claims 1-11, wherein determining the respective target model mean vector for the draft representation comprises summing, for each element of the draft representation, a value of the element and the target model value for the element.

13. The method of any of claims 1-12, wherein generating, from the current representation of the data item in the sequence, the draft sequence of draft representations of the data item comprises predicting each draft representation of the draft sequence based on a single prediction from a trained target diffusion model neural network.Attorney Docket No.: MX402430WO 14. The method of any of claims 1-13, wherein generating, from the current representation of the data item in the sequence, the draft sequence of draft representations of the data item comprises: processing the first draft representation in the draft sequence using a trained target diffusion model neural network to generate a denoising output; and, for successive draft representations in the draft sequence: generating a next draft representation in the draft sequence from a combination of a current draft representation and the denoising output for the first draft representation.

15. The method of claim 14, wherein generating the next draft representation in the draft sequence further comprises: processing the first draft representation and the current draft representation using a correction neural network to generate a correction output; adjusting the denoising output for the first draft representation using the correction output to obtain an adjusted current draft representation; and generating the next draft representation in the draft sequence from a combination of the adjusted current draft representation and the denoising output for the first draft representation.

16. The method of any of claims 13-15, comprising determining successive draft representations in the draft sequence in parallel.

17. The method of any of claims 1-12, wherein generating, from the current representation of the data item in the sequence, the draft sequence of draft representations of the data item comprises, for successive draft representations in the draft sequence: processing a current draft representation in the draft sequence using a trained draft diffusion model neural network to generate a draft model denoising output for each element of the current draft representation; and determining the next draft representation in the draft sequence from the draft model denoising output.

18. The method of claim 17, comprising: maintaining a plurality of the trained draft diffusion model neural networks; andAttorney Docket No.: MX402430WO determining the next draft representation in the draft sequence from a convex combination of the draft model denoising outputs from the plurality of trained draft diffusion model neural networks.

19. The method of any of claims 1-18, wherein the representations of the data item are representations in a latent space and the continuous-valued elements represent the data item in the latent space, the method further comprising processing the final representation of the data item using a decoder to generate an output data item.

20. The method of any of claims 1-19, wherein the data item comprises an image, or audio data representing an audio signal, or action selection data for controlling a mechanical agent to perform a task.

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

22. A computer storage medium encoded with instructions that, when executed by one or more computers, cause the one or more computers to perform the operations of the method of any of claims 1-20.

Images