Multistep distillation of diffusion models by moment matching

Abstract

Systems and methods, implemented as computer programs, for training a diffusion model for use in generating a frame of data, such as a frame of image data. The described techniques use a teacher diffusion model to train, or "distil", a student diffusion model that can generate a frame of data using many fewer steps than the teacher diffusion model, e.g. less than 20 time steps. The trained student diffusion model can be used to generate the frame of data.

Description

MULTISTEP DISTILLATION OF DIFFUSION MODELS BY MOMENT MATCHING BACKGROUND

[0001] This specification relates to processing data using machine learning models.

[0002] 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 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 a received input in accordance with current values of a respective set of parameters. SUMMARY

[0003] This specification describes systems and methods, implemented as computer programs on one or more computers in one or more locations, for training a diffusion model for use in generating a frame of data, such as a frame of image data. A method of using a trained diffusion model to generate a frame of data is also described.

[0004] In general a diffusion model can generate a frame of data by performing a reverse diffusion process to gradually de-noise an initial, noisy frame of data over a succession of time steps until a final, output frame of data is reached. Many time steps are often needed. The described techniques use a teacher diffusion model to train, or “distil”, a student diffusion model that can generate a frame of data using many fewer steps, e.g. less than 100, 50, or 20 time steps. The trained student diffusion model can be used to generate a frame of data.

[0005] The techniques are conceptually based on variations of a moment matching process, such as expectation matching. More particularly implementations of the techniques match conditional expectations of clean data between the training data distribution and a sampling distribution from the distilled student model, given noisy data from along a sampling trajectory of the diffusion process. The noisy data is partially denoised data, in particular obtained by denoising noisy data from a point along a forward diffusion process that complements the reverse diffusion process.

[0006] In one aspect there is described a computer-implemented method of training a diffusion model that involves processing a noisy training frame and a predicted denoised frame, using a noise schedule, to determine a reduced noise (partially denoised) trainingframe. Learnable parameters of the student diffusion model are updated using a training objective that depends on the predicted denoised frame and the reduced noise training frame.

[0007] In a related aspect there is described a computer-implemented method of training a diffusion model that involves training the student diffusion model using the teacher diffusion model by matching one or more statistical moments of a distribution of predicted frames from the student diffusion model to one or more statistical moments of the distribution of predicted frames from the teacher diffusion model.

[0008] There is also described a computer-implemented method of training a diffusion model that involves training the student diffusion model to minimize (with respect to learnable parameters of the student diffusion model) an objective that depends on a gradient determined from the teacher diffusion model and evaluated using a reduced noise (partially denoised) version of a noisy version of a training frame.

[0009] Also described are methods of using the trained student diffusion model to generate denoised data.

[0010] There is further described a system comprising one or more computers, and one or more storage devices communicatively coupled to the one or more computers. The storage devices store instructions that, when executed by the one or more computers, cause the one or more computers to perform the operations of one or more the described methods.

[0011] There is also described one or more non-transitory computer storage media storing instructions that when executed by one or more computers perform the operations of the described methods.

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

[0013] Diffusion models can generate high quality data samples, but with a significant computational burden. Implementations of the described techniques can substantially reduce this burden, by reducing the number of diffusion model steps required in inference by one or two orders of magnitude or more, whilst still generating high quality samples.

[0014] A many-step diffusion model is often viewed as a gold standard for data sample generation. Counter-intuitively, however, a student diffusion model trained as described herein can even improve on the quality of data samples generated by the teacher diffusion model. Broadly this can be achieved by averaging out prediction errors that arise at each diffusion step of the teacher model. Further, the techniques described herein can be used with a relatively poor teacher diffusion model.

[0015] The described techniques can be used with or without conditioning data, and can thus be used to generate frames of data with particular characteristics or for a particular purpose, or to modify frames of data, such as images.

[0016] The described techniques can be used to distill multiple teacher diffusion models into a single student diffusion model that combines the coverage or characteristics of the teacher diffusion models. For example image and dynamics teacher diffusion model modalities can be combined into a video student diffusion model.

[0017] 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

[0018] FIG.1 shows an example system for training a diffusion model.

[0019] FIG.2 is a flow diagram of an example process for training a diffusion model.

[0020] FIG.3 is a flow diagram of an example implementation of the process of FIG.2.

[0021] FIG.4 illustrates the operation of a process for training a diffusion model.

[0022] FIG. 5 a flow diagram of an example process for using a trained diffusion model to generate a frame of data.

[0023] FIG.6 shows examples of images generated by a diffusion model trained as described herein.

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

[0025] This specification describes computer-implemented methods of training a diffusion model, in particular a student diffusion model, using a teacher diffusion model. Such techniques are generally referred to as distillation.

[0026] The teacher diffusion model is configured to process a noisy frame corresponding to a diffusion time, and in implementations also the diffusion time, to generate a teacher predicted frame for de-noising the noisy frame. Processing the diffusion time can improve a quality of the teacher predicted frame.

[0027] In general a predicted frame is a frame of data that is generated by a diffusion model as described herein, in particular by a neural network of the diffusion model, when thediffusion model processes an input frame of data. Typically the predicted frame of data has the same dimensionality as the input frame of data. References herein to sampling or processing values of a frame of data are to sampling or processing values of data elements defined by the frame of data.

[0028] In some implementations the data elements defined by a frame of data can be, e.g., pixel values of an image frame, audio signal values (e.g. instantaneous amplitude values) of an audio frame, and so forth.

[0029] In some implementations rather than operating in an output space such as a pixel space, the teacher and student models can operate in latent space. That is the teacher diffusion model and student diffusion model can be latent diffusion models. The data elements defined by a frame of data can then be a latent representation of, e.g., pixel values of an image frame, audio signal values of an audio frame, and so forth.

[0030] The techniques described herein can be used unchanged, whether operating in an output space, i.e. a space in which an output data item is generated, or in a latent space, i.e. a space of a latent representation of the output space.

[0031] A diffusion model, in particular a student diffusion model, trained as described herein can be used to generate a frame of data by performing a sequence of de-noising steps. The generated frame of data can be a frame of data in an output space, i.e., so that the values of the output 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 the generated frame of data can be a frame of data in a latent space, i.e., so that the values in the frame of data are values in a latent representation of an output data item in the output space. Typically such a latent space is lower-dimensional than the output space.

[0032] When the frame of data is generated in a latent space, a system as described herein can generate a final frame of data in the output space by processing the frame of data in the latent space using a decoder neural network, e.g., one that has been pre-trained in an auto- encoder framework. During training, the system can use an encoder neural network, e.g., one that has been pre-trained jointly with the decoder in the auto-encoder framework, to encode training data items in the output space to generate training data items for the diffusion neural network in the latent space.

[0033] In implementations of the described system frames of data are latent representations and the values in the representations are learned, latent values rather than, e.g., pixel values when the frames of data represent images.

[0034] In general a predicted frame of data, from the teacher diffusion model or from the student diffusion model, can be used to generate a reduced noise version of the input frame of data. For example, in some implementations the teacher predicted frame comprises a prediction of noise in the noisy frame, which can then be subtracted from the noisy frame to obtain a prediction of a reduced noise frame, e.g. a prediction of a fully de-noised frame. In some implementations the teacher predicted frame comprises a prediction of a reduced noise frame, e.g. a prediction of a fully de-noised frame.

[0035] The diffusion time lies between an initial time and a final time. In general the initial time corresponds to an initial noisy frame that is processed during inference, e.g. an initial frame comprising purely noise. In general the final time corresponds to a final frame of data generated during inference, e.g. to a supposed fully de-noised frame.

[0036] The methods train a student diffusion model. The student diffusion model is configured to process a noisy frame corresponding to a diffusion time, and optionally also the diffusion time, to generate a student predicted frame for de-noising the noisy frame, e.g. a prediction of noise in the noisy frame or a prediction of a de-noised frame.

[0037] In inference, when used to generate a frame of data, the diffusion times processed by the student diffusion model correspond to a discrete number (k) of (sampling) time steps atrespective step times between the initial time and the final time. As some examples ^^ ^ 100or ^^ ^ 20, e.g. 1 ^ ^^ ^ 16 or 1 ^ ^^ ^ 8; in implementations ^^ ^ 2, 3, 4 ^^^^ 8. Typically adiffusion model performs a denoising process over hundreds or thousands of time steps. Implementations of the described techniques enable the generation of frames of data that are as good or (counter-intuitively) better than those of the teacher whilst only using, e.g., 8 steps.

[0038] During training (as described below) the student diffusion model can be trained just at the step times, or over a continuous range of times between the initial time and the final time.

[0039] Referring to FIG.1, this shows an example system 100 for training a diffusion model. The system of FIG.1 can be implemented as computer programs on one or more computers in one or more locations.

[0040] The system 100 comprises a teacher diffusion model 110 configured to process a noisy frame 112 corresponding to a diffusion time 114, and in implementations also the diffusion time 114, to generate a teacher predicted frame 118.

[0041] The system 100 also includes a student diffusion model 120 configured to process a noisy frame 122 corresponding to a diffusion time 124, and in implementations also the diffusion time 124, to generate a student predicted frame 128.

[0042] Some implementations system 100 also include an auxiliary diffusion model 150 configured to process a noisy frame 152 corresponding to a diffusion time 154, and in implementations also the diffusion time 154, to generate an auxiliary predicted frame 158.

[0043] In implementations the teacher diffusion model 110 and the student diffusion model 120 and (where present) the auxiliary diffusion model 150 comprise, respectively, a teacher diffusion model (de-noising) neural network, a student diffusion model (de-noising) neural network, and an auxiliary diffusion model (de-noising) neural network. Each of these is configured to process a noisy frame corresponding to a diffusion time, and in implementations the diffusion time, to generate, respectively, the teacher predicted frame 118, the student predicted frame 128, and the auxiliary predicted frame 158. For example a noisy training frame can be processed using one of the respective models to generate a predicted denoised frame for the final time.

[0044] The system 100 also includes a training dataset 130 storing training frames 132; and a training engine 140 for training the student diffusion model 120 as described later.

[0045] The teacher diffusion model (de-noising) neural network and a student diffusion model (de-noising) neural network can have any suitable architecture consistent with processing values of a frame of data as an input to generate a corresponding set of output values for a frame of data. In general these neural networks may each comprise one or more feedforward, convolutional, attention, normalization, or other neural network layers.

[0046] As some examples, such a neural network can have a U-Net architecture or a variant thereof, or a Transformer architecture or a variant thereof, or a combination of these such as a U-ViT architecture (Bao, et al., “All are Worth Words: A ViT Backbone for Diffusion Models”, arXiv:2209.12152, 2023), e.g. one of the U-Net or U-ViT model configurations described in Appendix B of Hoogeboom et al. “simple diffusion: End-to-end diffusion for high resolution images”, arXiv:2301.11093, 2023. As some further examples such a neural network may comprise a U-Net with one or more ResNet blocks and one or more self- attention layers, or may comprise a diffusion transformer (DiT) or a transformer backbone.

[0047] Processing a time using such a neural network generally involves processing data specifying the time, e.g. an embedding of the time. As one example, the time can be encoded as an embedding, such as a sinusoidal positional embedding, and added to or otherwise combined with each processed data element. As another example, the time can be provided as side information to one or more layers of the neural network.

[0048] In implementations the teacher diffusion model 110 and / or the student diffusion model 120, in particular the teacher diffusion model neural network and the student diffusionmodel neural network, are also configured to process conditioning data to generate the respective teacher predicted frame 118 or student predicted frame 128. A neural network can be configured to process conditioning data, e.g. by incorporating one or more cross-attention layers to attend to the conditioning data, or in some other way. The conditioning data can come from a conditioning input (not shown in FIG.1).

[0049] In general the conditioning data, e.g. the conditioning input, characterizes a content of the teacher predicted frame 118 or student predicted frame 128. That is, the conditioning data defines one or more properties of a generated frame of data (in a latent space or the output space). As an example the conditioning input may comprise text or audio or tokens or an embedding representing features of text or audio, and the predicted frame may comprise data for pixels of a still or moving image that represents the text or audio.

[0050] Such a conditioning input may comprise (but is not limited to) text data and / or image data each of which may be encoded e.g. into a sequence of embeddings using a text, image, or multimodal Transformer model, e.g. an encoder-decoder model such as T5 (Text-to-Text Transfer Transformer, Raffel et al., arXiv:1910.10683, 2023), or a language model.

[0051] In the later described training methods, conditioning data may be incorporated into the training teacher diffusion model neural network or the student diffusion model neural network, e.g., using so-called classifier free guidance. This can involve randomly masking out or otherwise removing the conditioning data from the neural network input so as to train the neural network to generate the respective predicted frame both with and without guidance from the conditioning data. Classifier-free diffusion model guidance is described, e.g., in Ho and Salimans, arXiv:2207.12598. At inference a frame of data can be generated from a weighted combination of predicted frames with and without conditioning, e.g. from a difference between a predicted frame with conditioning weighted by (1+w) and predicted frame without conditioning weighted by (1+w), where w is a positive weight.

[0052] FIG.2 is a flow diagram of an example process for training a diffusion model such as the student diffusion model 120. The process of FIG.2 may be implemented by one or more computers in one or more locations; for convenience the process is described with reference to FIG.1.

[0053] At step 200 the process initializes the student diffusion model 120. As one example the student diffusion model can be initialised from the teacher diffusion model, e.g. by initializing learnable parameters, e.g. weights, of the student diffusion model 120 from corresponding learned parameters of the teacher diffusion model 110. The student diffusion model can have the same architecture as the teacher diffusion model, but this is not essential.As another example the student diffusion model 120 can be initialized by randomly initializing learnable parameters of the student diffusion model.

[0054] The student diffusion model 120 is trained using the teacher diffusion model 110. The steps 202 to 214 described below are repeated during the training process for any suitable number of times, e.g. until a training loss has flattened out to an approximate minimum.

[0055] In implementations the training involves obtaining, e.g. sampling, a training frame (^^) from a training dataset (^^) (step 202).

[0056] The process determines a target time (^^). In some implementations the target time, ^^, is obtained by randomly sampling from a target time sampling distribution (^^^^^^) (step 204). Any distribution can be used, e.g. a uniform distribution. In some implementations the target time, ^^, can be selected from a fixed set of discrete times, e.g. the times of the k time steps at which the student diffusion model is used (evaluated) in inference. Sampling the target time randomly can help stability and performance. In general the target time is between the aforementioned initial time, e.g.1, and final time, e.g.0.

[0057] The process determines a training time (^^), in particular a training time that is within a duration of a time step after the target time (step 206). In some implementations the training time is randomly sampled from a distribution; in some implementations the training time can be a fixed time after the target time. The training time can also be referred to as a sampling time. In general the sampling time, t, is between the aforementioned initial time, e.g.1, and final time, e.g.0.

[0058] As an example, a time ^^௧can be sampled from a distribution such as a uniform distribution in the range ^0,1 / ^^^ where ^^ is the number of discrete (sampling) time steps used during inference. Alternatively time ^^௧can be fixed. The training time (^^), can be truncatedat 1, e.g. ^^ ൌ ^^^^^^^^^^^^^^^^^ ^ ^^௧ , 1^. In some implementations the training time ^^ ∈ ^0,1^)and is continuous, even though at inference the student diffusion model is only evaluated at ^^ time steps. In some implementations the sampled time ^^௧is randomly sampled. These measures can help stability and performance but are not essential. For example the training time ^^ need not be continuous; and the time ^^௧can be fixed to a single value.

[0059] The training frame (^^) is processed using a noise schedule that defines a variation of a frame noise level with time to determine a noisy training frame (^^௧) for the training time (t) (step 208). This can involve obtaining a noise frame (^^) and processing the noise frame and the training frame to determine the noisy training frame (^^௧).

[0060] In general the noise schedule defines a variation in signal to noise ratio (SNR) for a de-noising process from the final time to the initial time that the teacher diffusion model and the student diffusion model both follow. Thus the noise schedule can define the SNR for the training time (t) and for the target time (s). In general the SNR decreases monotonically from the initial time to the final time.

[0061] The noise schedule can define a time-dependent mean scaling factor for a frame (^^௧) and a time-dependent variance scaling factor for a frame (^^௧ଶ) and the SNR can be defined as^^^^^^ ൌ ^^ଶ௧ / ^^ଶ௧. Optionally, but not necessarily, determining one of the mean scaling factorand the variance scaling factor may determine the other, e.g. according to ^^ଶ ଶ௧ ^ ^^௧ ൌ 1,which defines a variance-preserving noise (SNR) schedule. However use of a variancepreserving noise schedule is not essential. In some implementations ^^^ ൌ ^^^ ൌ 0, ^^^ ^ ^^and, e.g., ^^^ ൌ ^^^0, ^^^.

[0062] In this framework the noisy training frame can be determined as ^^௧ ൌ ^^௧^^ ^ ^^௧^^ with^^~^^^0, ^^^ and ^^ ∈ ^0,1^. That is, ^^௧ can be distributed as ^^^^^௧|^^^ ൌ ^^^^^௧|^^௧^^,^^௧^.

[0063] The noisy training frame (^^௧) is processed using the student diffusion model 120 to generate a student predicted frame that defines, e.g. is, a (student) predicted denoised frame (^^^) for the final time (step 210).

[0064] The noisy training frame (^^௧) and the (student) predicted denoised frame (^^^) for the final time are processed, using the noise schedule, to determine a (partially) reduced noise training frame (^^^) corresponding to the target time (^^) (step 212). The reduced noise training frame (^^^) can be determined stochastically, e.g. based on another sampled noise frame, or deterministically, e.g. based on the previously sampled noise frame. In general the reduced noise training frame (^^^) is still noisy, but has less noise than the noisy training frame (^^௧).

[0065] As one particular example, the reduced noise training frame (^^^) can be determinedby sampling from ^^^^^^|^^௧,^^^ ൌ ^^^^^^|^^௧→^^^^௧ ,^^^,^^௧→^^; or equivalently ^^^ ൌ ^^௧→^^^^௧,^^^ ^^^௧→^^^^0, ^^^, where^^

[0066] As another particular example, in a deterministic setting such as DDIM (Song et al., “Denoising Diffusion Implicit Models”, arXiv:2010.02502v4, October 2022), ^^^ൌDDIM^→^^^^, ^^௧^ ൌ ^^^^^ ^ ^ఙೞఙ^ ^ ^^^௧ െ ^^௧^^^.of the student diffusion model are updatede.g. loss (step 214). Typically updating the learnable parameters, e.g. weights, of a model as described herein involves backpropagating gradients of an objective function, e.g. loss, to update the learnable parameters. This can use any appropriate gradient descent optimization algorithm, e.g. Adam or another optimization algorithm.

[0068] In general the student model training objective, e.g. loss is dependent on the (student) predicted denoised frame (^^^) for the final time and on the reduced noise training frame (^^^), in particular as processed using the teacher diffusion model to generate the teacher predicted denoised frame (^^ఏ^^^^^ where ^^ఏ^∙^ denotes the teacher diffusion model with parameters, e.g. weights, ^^).

[0069] Broadly, the student model training objective trains the student diffusion model to match a conditional expectation of the student diffusion model (i.e. the student predicted denoised frame,^ ^^, for the final time) to conditional expectation of the teacher diffusion model (the teacher predicted denoised frame, ^^ఏ^^^^^). In implementations the expectations can be calculated using a single sample (a Monte Carlo estimate). During training the matching is effectively averaged over samples from the diffusion process, i.e. over target times (s).

[0070] In implementations the reduced noise training frame (^^^) corresponding to the target time has been processed by the teacher diffusion model to generate the teacher predicted denoised frame (^^ఏ^^^^^), and this processing is used, directly or indirectly, in conjunction with the (student) predicted denoised frame (^^^) for the final time, to determine the loss.

[0071] For example, in some particular implementations the student model training objective depends on the predicted denoised frame (^^^) for the final time and on the teacher predicted denoised frame, ^^ఏ^^^^^, for the final time. In some particular implementations, the student model training objective depends on the predicted denoised frame (^^^) for the final time and on a Jacobian of the teacher diffusion model (neural network) evaluated at ^^ఏ^^^^^.

[0072] In some implementations updating the learnable parameters of the student diffusion model involves backpropagating gradients of the student model training objective and applying a stop gradient to the reduced noise training frame (^^^) so as not to pass gradientsinto the student diffusion model via the reduced noise training frame (^^^). This can provide a stable target, facilitating convergence when training the student diffusion model.

[0073] The above described process can be understood as moment matching, in particular matching one or more statistical moments of a distribution of predicted frames from the student diffusion model to one or more statistical moments of the distribution of predicted frames from the teacher diffusion model. This is different to simply training the studentdiffusion model to generate accurate samples ^ ^^ ൌ ^^ఎ^^^௧ , ^^^ where ^^ఎ^∙^ denotes the studentdiffusion model with parameters, e.g. weights, ^^.

[0074] It is helpful for understanding implementations of the invention to describe this in more detail.

[0075] Consider the algorithm below, which shows a typical stochastic sampling process for generating a frame of data from a diffusion model using k sampling steps (and where, at each step, ^^^becomes the next ^^௧):

[0076] For a conventional diffusion model this process typically requires at least hundreds of sampling steps. An algorithm of this type (or a deterministic sampler such as DDIM) can be used for generating frames of data using a student diffusion model trained as described herein. However when the student diffusion model is trained as described this process can be performed (unmodified) with, e.g., less than 20 sampling steps, and can obtain samples that are as good or better than from the teacher diffusion model. This entails the student diffusion model making larger steps through the (reverse) diffusion process than does the teacher diffusion model.

[0077] Rather than training the student diffusion model so that ^ ^^ ൌ ^^ఎ^^^௧ , ^^^ approximatesthe expectation ^^^^^^|^^௧^ (where q denotes the data distribution ^^^^^^), the student diffusion model is trained to generate approximate samples from the distribution^ ^^~^^^^^|^^௧^, in which case the above algorithm produces exact samples from the data distribution q for any number of sampling steps k.

[0078] Put differently, the output (conditional expectation) of the student diffusion model is trained to match the output (conditional expectation) of the teacher diffusion model, given thesame noisy input (and averaged over the diffusion trajectory). In this way the student diffusion model learns to model the statistical effect of a block of steps rather than replicatingeach step, which is a more powerful training process than learning to generate ^ ^^ ൌ ^^ఎ^^^௧ , ^^^.Even with fewer steps the student diffusion model is still generating samples that closely match the true data distribution in expectation.

[0079] The teacher diffusion model can have prediction errors and these can be different across different time steps. The distillation process links the predictions of clean data at different diffusion times and helps to ensure that these are mutually consistent. Empirically the student diffusion model can generate better samples than the teacher diffusion model.

[0080] The conditional expectation of the teacher diffusion model ^^^^^^|^^^^ (e.g. ^^ఏ^^^^^) can be determined by sampling ^^~^^^^^^ and ^^^~^^^^^^|^^^). The conditional expectation of thestudent diffusion model ^^^^^^^|^^^^ can be determined by sampling ^^~^^^^^^, ^^௧~^^^^^௧ |^^^, ^ ^^ ൌ^^ఎ^^^௧^, and ^^^~^^^^^^ |^^௧ ,^^^^. Moments between the student diffusion model and the teachercan be matched by training the student diffusion model using a moment(loss) that depends on a difference between these expectations ^^^^^^^|^^^^and ^^ ^^^^|^^ ^), such as an L2 distance ^^ ^^^^|^^ ^ െ ^ ^ଶ ^^ ฮ ^ ^ ^^^ ^^|^^^ ฮ. The loss can be determined over a distribution of the same noisy inputs (reduced noise training frames), ^^^, for the student diffusion model and the teacher diffusion model. Here, and later, dependence of the model predictions on diffusion time has been omitted for clarity, though in implementations it can be included.

[0081] In practice, whilst not essential, a stop gradient can be used so that when backpropagating gradients of the objective (loss) the backpropagation does not extend back into the student diffusion model via ^^^, i.e. the target (the difference in expectations) istreated as fixed. Then the gradient of the objective is approximately ∇ఎ^^^൫^^^^^^^|^^^^ െ^^^^^^|^^^^൯ and the objective (loss) can be simplified to ^^^^^^ ൌ ^ ^^்^^^^൫^^^^^^^|^^^^ െ ^^^^^^|^^^^൯(with possible additional regularization) where ^^^^different to the L2 loss but the gradient of the loss is approximately the same). This loss isminimized if ^^^^^^^|^^^^ ൌ ^^^^^^|^^^^. Empirically using a stop gradient in this way helps tostabilize training, compared with using the full gradient.

[0082] This simplified loss, which depends on ^ ^^்^^^^൫^^^^^^^|^^^^ െ ^^^^^^|^^^^൯ depends on asimilarity between i) the predicted denoised frame for the final time and ii) a difference between the auxiliary predicted denoised frame for the final time and the teacher predicteddenoised frame for the final time. Here the similarity is defined by a vector inner product (dot product or normalized dot product, i.e. cosine, similarity). It can be understood asdriving the difference in expectations ൫^^^^^^^|^^^^ െ ^^^^^^|^^^^൯ to be zero, so that^^^^^^^|^^^^ ൌ ^^^^^^|^^^^. The simplified loss, and in particular its gradient, is a stochasticapproximation of this goal. (Note that ^^^^^^^|^^^^ depends on the distribution of ^ ^^ so thatadjusting the prediction^^^ also, implicitly, adjusts ^^^^^^^|^^^^). The difference in expectations൫^^^^^^^|^^^^ െ ^^^^^^|^^^^൯ is, in effect, a measure of how much the student diffusion model is inIn some implementations the learnable parameters of the student diffusion model are updated based on a loss in ^^-space (a space of the generated frame of data), e.g. using the auxiliary diffusion model 150. In some implementations the learnable parameters of the student diffusion model are updated based on an objective that operates (directly) in parameter space, i.e. in a space of the learnable parameters of the student diffusion model.

[0084] In the above described loss, which depends on ^ ^^்^^^^൫^^^^^^^|^^^^ െ ^^^^^^|^^^^൯, anestimate of the expectation ^^^^^^|^^^^ can be determined from the teacher diffusion model, e.g.as ^^^^^^|^^^^ ൌ ^^ఏ^^^^^ (using a one-sample estimate). One way to estimate ^^^^^^^|^^^^ is bymaintaining an auxiliary diffusion model. In another approach the loss is rewritten as a loss in parameter space, in effect by analytically determining a result of performing a gradient descent step on the loss.

[0085] Thus some implementations of the method involve maintaining an auxiliary diffusion model (^^థ^∙^) auxiliary diffusion model 150, configured to process a noisy frame corresponding to a diffusion time to generate an auxiliary predicted frame for de-noising the noisy frame. The reduced noise training frame^^^^^is processed using the auxiliary diffusion model to generate an auxiliary predicted frame that defines, or is, the auxiliary predicted denoised frame, ^^థ^^^^^ (a frame that corresponds to the final time).

[0086] The reduced noise training frame ^^^^^ can also be processed using the teacher diffusion model to generate a teacher predicted frame that defines, or is, a teacher predicted denoised frame, ^^ఏ^^^^^(a frame that also corresponds to the final time).

[0087] The learnable parameters, e.g. weights, of the auxiliary diffusion model can then be updated using an auxiliary model training objective dependent on a difference between the predicted denoised frame (^^^) (corresponding to the final time) and the auxiliary predicted denoised frame (^^థ^^^^^) (corresponding to the final time).

[0088] In some implementations the auxiliary diffusion model 150 ^^థ^∙^ can be regularized so as to stay close to the teacher model ^^ఏ^∙^. For example the auxiliary model training objective can be regularized by adding a term to the objective thaton the difference between the auxiliary predicted denoised frame for the final time and the teacher predicteddenoised frame for the final time, e.g. ฮ^^ఏ^^^^^ െ ^^థ^^^^^ฮଶ . (The auxiliary diffusion model then predicts the average of ^ ^^ and loss ^^^^^^ by 1 / 2).

[0089] As previously described theof the student diffusion model can be updated using a student model training objective that depends on a (dot product) similarity between i) the predicted denoised frame (^^^) for the final time and ii) a difference between the auxiliary predicted denoised frame (^^థ^^^^^) and the teacher predicted denoised frame (^^ఏ^^^^^). As previously described a stop gradient can be applied to the difference between auxiliary predicted denoised frame for the final time and the teacher predicted denoised for the final time, i.e. so as not to pass gradients into the student diffusion model via the difference.

[0090] In some implementations the learnable parameters of the auxiliary diffusion model and the learnable parameters of the student diffusion model are, conveniently, updated alternately.

[0091] The algorithm below shows one example implementation of these techniques:

[0092] In the above example implementation the student model training objective depends directly on the predicted denoised frame for the final time,^^^, and on the teacher predicted denoised frame ^^ఏ^^^^^. In another example implementation the student model training objective dependson the predicted denoised frame for the final time,^ ^^, and on the teacher predicted denoised frame ^^ఏ^^^^^. That is, the student model training objective can be determined in teacher-parameter space rather than x-space. This can be done by calculatingthe effect of minimizing ^^^^^^ in the above algorithm, without actually needing to maintain the auxiliary diffusion model.

[0093] Starting from the teacher parameters ^^, the effect of taking a single gradient step on ^^^^^^ can be determined as: ^ ^ ^^ ^^^^ ^ ^^^^^ఏ^^^^^^௧^^௧ ൌ ^^ ^^ ^^^^ ^^^^^Λ∇ம^^^^^^|థୀఏ^where ^^ ൌ ^^ denotes (which vector has adimension, p, equal to the number of parameters Here డఏ is the Jacobian of theteacher diffusion model ^^ఏ, where Λ is an optional (^^ ൈ ^^, diagonal) gradient preconditioningmatrix; if there is no preconditioning Λ can be the identity matrix. The gradient preconditioning matrix can be, e.g., the diagonal gradient preconditioning matrix used in the Adam optimization algorithm (Kingma and Ba, Adam, “A Method for Stochastic Optimization”, arXiv:1412.6980).

[0094] This loss can be written in terms of moment matching in teacher parameter space bydefining a (weighted) teacher loss ^^ఏ^^^, ^^^^ ≡ ^^^^^^‖^^ െ ^^ఏ^^^^^‖ଶ, where ^^^^^^ is anoptional loss weight (which has zero gradient when sampling from the training distribution).An estimate of the “instantaneous” loss, ^^^^^^௧^^௧, is then dependent on (e.g. equal to, apartfrom a scale factor and offset) an expectation of ∇ఏ^^ఏ൫^^^, ^^^^^^^^^൯ where ^^௧~^^; ^ ^^ ൌ^^^௧^^^^௧^; and ^^^~^^^^^^|^^௧,^^^^.

[0095] This “instantaneous” version of the moment matching loss can be interpreted as matching teacher gradients between the training data and generated data. The loss^^^^^^௧^^௧^^^^is minimized if the teacher diffusion model has zero loss gradient when evaluatedon data generated by the student diffusion model ^^ఎ. That is the loss is minimized if the teacher diffusion model cannot “see” the difference between real data and data generated by the student diffusion model. The instantaneous version of the moment matching loss can be evaluated using an independent sample, e.g. minibatch, from the training data, from which one or more independent samples of ^^௧,^ ^^, and ^^^are obtained (denoted^ ^^′ and ^^′^in the algorithm below).

[0096] FIG.3 is a flow diagram of one example implementation of this approach, and is a particular example of the process of FIG.2. The process of FIG.3 may be implemented by one or more computers in one or more locations.

[0097] At step 302 first and second training frames (^^) are obtained, e.g. sampled, from the training dataset. The first training frame is processed, using the noise schedule, to determinea noisy first training frame (^^௧) for the training time (step 304); and the second training frame is processed, using the noise schedule, to determine a noisy second training frame (^^′௧) for the training time (step 306).

[0098] The noisy first training frame (^^௧) is processed using the student diffusion model to generate a first student predicted frame that defines a first predicted denoised frame (^^^) (corresponding to the final time) (step 308). The noisy second training frame (^^′௧) is processed using the student diffusion model to generate a second student predicted frame thatdefines a second predicted denoised frame (^^^′ ) (corresponding to the final time) (step 310).

[0099] The noisy second training frame (^^′௧) and the second predicted denoised frame (^^^′ )are processed, using the noise schedule, to determine a second reduced noise training frame (^^′^) corresponding to the target time (step 310).

[0100] The second reduced noise training frame (^^′^) is processed using the teacher diffusionmodel to generate an auxiliary predicted frame (noting that ^^థ^^^′^^ ൌ ^^ఏ^^^′^^ for evaluationat ^^ ൌ ^^) (step 314). This defines, or is, an auxiliary ^^థ^^^′^^corresponding to the final time.

[0101] The process can then determine a teacher model gradient (with respect to learned parameters of the teacher diffusion model) of an auxiliary objective function dependent uponon a difference between the second predicted denoised frame (^^^′ ) corresponding to the finaltime and the auxiliary predicted denoised frame (^^థ^^^′^^ ൌ ^^ఏ^^^′^^) (corresponding to thefinal time) (step 316).

[0102] The learnable parameters of the student diffusion model are then updated using the student model training objective, which in this example implementation depends on the first predicted denoised frame (^^^) corresponding to the final time and the teacher model gradient (^^) (step 318).

[0103] In more detail, this can involve comprising determining the student model training objective dependent on a product of the first predicted denoised frame (^^^) for the final time and a vector proportional to the teacher model gradient (^^).

[0104] In broad terms the training objective is minimized when the teacher model gradient with respect to its learned parameters is zero. The objective trains the student model so thatwhen the teacher model processes, e.g., ^^′^ obtained from the student model (via ^^^′ ) it haszero teacher model gradient. As described above, optimization is successful if the teacher model cannot “see” the difference between real and generated data and would not change its parameters when trained on data generated by the student diffusion model.

[0105] A particular implementation involves processing the noisy first training frame (^^௧) and the first predicted denoised frame (^^^) for the final time using the noise schedule, to determine a first reduced noise training frame (^^^) corresponding to the target time. In some implementations a scaled teacher model gradient, ^^, can then be determined by multiplying the teacher model gradient by a (predetermined) scaling matrix, e.g. the gradientpreconditioning matrix Λ, above. For example ^^ ൌ Λ ∇ఏ^^ఏ^^^^ᇱ, ^^′^ ^ (and is a vector ofdimension p).

[0106] The choice of the scaling matrix, where used, is not critical. As an example, the scaling matrix can be based on the second moment (variance) of a set of (teacher) gradients,e.g. on a moving average of the squared gradients, e.g. as diag^ ^^௩^^ାఢ^ with ^^^௧proportional to ^^௧ଶwhere ^^௧is the gradient at time ^^ and ^^ is a small constant numerical stability -following Kingma et al. (ibid). This scaling matrix can be e.g., by running an optimizer on a few training steps as part of an initialization process, or could be loaded from a checkpoint saved during training of the teacher diffusion model.

[0107] Implementations of the process can determine a Jacobian of the learned parameters of the teacher diffusion model, డ^ഇ^^^ೞ^ డఏ , evaluated for the first reduced noise training frame (^^^). The method can then determine a product of the Jacobian of the learned parameters of theteacher diffusion model (dimension ^^ ൈ ^^ where ^^ is a dimension of the flattened output) andthe scaled teacher model gradient (^^) to obtain a training product (a vector of dimension d). The student model training objective can then depend on a product of the first predicted denoised frame (^^^்) for the final time (flattened into a vector of dimension d) and the training product. Expressing the loss in terms of a Jacobian-vector product is convenient forautomatic differentiation frameworks such as JAX. In implementations the ^^ ൈ ^^ Jacobianneed not actually be materialised, just evaluated implicitly (in a Jacobian vector product in an automatic differentiation framework such as JAX).

[0108] In implementations updating the learnable parameters of the student diffusion model includes applying a stop gradient to the training product. More particularly the process can involve updating the learnable parameters of the student diffusion model by backpropagating gradients of the student model training objective dependent, and applying the stop gradient so that the learnable parameters of the student diffusion model are not updated via backpropagation through the training product to the first reduced noise training frame (^^^) and the first predicted denoised frame (^^^).

[0109] In some implementations the method involves obtaining a batch of first training frames from the training dataset, and processing each of these as described above. Similarly the method can involve (independently) obtaining a batch of second training frames from the training dataset, and processing each of these as described above. In particular, the student model training objective can be determined from the batch of first training frames, and the teacher model gradient can be determined from the batch of second training frames, e.g. by averaging over the respective batches.

[0110] The algorithm below shows one example implementation of these techniques, which involve moment matching in parameter space, conceptually based on an instantaneous version of the auxiliary diffusion model (which is exchanged for two independent minibatches per parameter update).

[0111] FIG.4 illustrates thewhich the above algorithm is a particular example. Forward diffusion moves a sample from the training data distribution, ^^, to ^^௧. Then the student diffusion model is used to generate^ ^^ from ^^௧. Stochastic interpolation between^ ^^ and ^^௧is used to obtain ^^^and the gradient of the teacher loss on^ ^^ and ^^^is minimized.

[0112] As previously described, in general the student diffusion model can be trained using the teacher diffusion model by matching one or more statistical moments of a distribution of predicted frames from the student diffusion model to one or more statistical moments of the distribution of predicted frames from the teacher diffusion model.

[0113] More particularly this can involve training the student diffusion model using a moment matching objective that depends on the distribution of predicted frames from the student diffusion model and on the distribution of predicted frames from the teacher diffusion model and that is minimized when training the teacher diffusion model on predicted frames from the student diffusion model would not change learned parameters of the teacher diffusion model. That is, the student diffusion model can be trained so that the teacherdiffusion model cannot “see” the difference between frames of data generated by the student diffusion model and real frames of data, e.g. from the training dataset.

[0114] In some specific implementations the method can involve obtaining a training frame from a training dataset and matching a statistical moment defined by (a gradient of) a denoising loss determined by a difference between a first predicted frame generated by the student diffusion model and a second predicted frame generated by the teacher diffusion model.

[0115] For example, the first predicted frame can be obtained by processing a noisy version of the training frame using the student diffusion model. The second predicted frame can be obtained from the first predicted frame, and in implementations also the noisy version of the training frame, according to the noise schedule.

[0116] In some specific implementations matching the one or more statistical moments involves matching a statistical moment defined by a denoising loss determined by a difference between a first predicted frame generated by the student diffusion model and a second predicted frame generated using an auxiliary diffusion model. The method can train the auxiliary diffusion model to generate predicted frames that match predicted frames generated by the teacher diffusion model.

[0117] As an example this can involve determining a reduced noise training frame from a predicted frame from the teacher diffusion model generated by processing a training frame and a noisy version of the training frame. This can also involve processing the reduced noise training frame using the auxiliary diffusion model to generate an auxiliary predicted frame. This can also involve processing the noisy training frame using the teacher diffusion model to generate a teacher predicted frame. This can also involve training the auxiliary diffusion model using an objective that depends on a difference between the auxiliary predicted frame and the teacher predicted frame.

[0118] In broad terms one implementation of a training process as described above that operates in the parameter space can be characterized as first sampling a training frame (^^) and then applying a forward diffusion process to modify ^^ to ^^௧. Then a (clean) version of the training frame,^ ^^, conditional on ^^௧, is generated using the student diffusion model. A reduced noise version of ^^௧, ^^^, is obtained by (stochastic) interpolation between^ ^^ and ^^௧. The student diffusion model is trained to minimize the teacher diffusion model gradient norm at ^^^with respect to the generated version of the training frame,^ ^^.

[0119] For example, this can involve obtaining a training frame from a training dataset, and applying a forward diffusion process, i.e. a noise addition process, to the training frame to generate a noisy version of the training frame (^^௧). The noisy version of the training frame can be processed using the student diffusion model to generate a predicted frame comprising a de-noised predicted frame (^^^) that is a de-noised version of the noisy version of the training frame.

[0120] A reduced noise version (^^^) of the noisy version of the training frame can be obtained from the de-noised predicted frame (^^^) and the noisy version of the training frame (^^௧), in particular using the noise schedule (“stochastic interpolation”).

[0121] The student diffusion model can be trained to minimize, with respect to learnable parameters of the student diffusion model, an objective dependent on a gradient determined from of the teacher diffusion model and evaluated, e.g., using the reduced noise version of the noisy version of the training frame. For example the objective can depend on a norm of the gradient such as an L1 or L2 norm.

[0122] Optionally in the above described techniques the student model training objective,and / or the auxiliary model training objective, and / or the instantaneous loss, ^^^^^^௧^^௧, can beweighted using a time-dependent weight, e.g. defined by a weight schedule ^^^^^^. This canhelp stability. Empirically, however, a flat weight ^^^^^^ ൌ 1 works well, although the weightschedule ^^^^^^ can be adapted to normalize gradient magnitudes across different noise levels.

[0123] In some implementations multiple teacher diffusion models can be used. In this case, multiple moment s can be matched, i.e. multiple losses as described above can be used, one per teacher model. The teacher diffusion models can have been trained on different datasets, e.g. using images with different aspect ratios, crop sizes, and so forth. More particularly the describes techniques are compositional and the teacher diffusion models that cover different training data modalities can be combined into a student diffusion model that includes the different modalities. For example a first teacher diffusion model trained on image data and a second teacher diffusion model trained on dynamics data can be distilled into a student diffusion model that can generate high quality video data; and so forth.

[0124] As previously described, classifier-free guidance can be used with the above techniques, in the usual way (straight through gradient approximation can be used when implementing Algorithm 3 and related techniques).

[0125] FIG.5 is a flow diagram of an example process for using a trained diffusion model, in particular a student diffusion model trained as described above, to generate a frame of data.The process of FIG.5 may be implemented by one or more computers in one or more locations; for convenience the process is described with reference to the system of FIG.1.

[0126] The trained diffusion model is configured to process a noisy frame corresponding to a diffusion time, and in implementations the diffusion time, and optionally also conditioning data, to generate a predicted frame (predicting the noise, or a de-noised frame) for de-noising the noisy frame.

[0127] The process involves initializing a current noisy frame (e.g. ^^௧) (step 500). The example process is performed iteratively, for each of a discrete number (k) of time steps, e.g.with ^^ ^ 100, 50, 20, 10, between an initial time and a final time.

[0128] Each iteration involves processing the current noisy frame using the diffusion model to generate a predicted frame comprising either a prediction of noise in the current noisy frame or a prediction of a de-noised version of the current noisy frame (^^^) (step 502). The de-noised version of the current noisy frame (^^^) can be, or be determined from, the predicted frame.

[0129] Each iteration determines an updated current noisy frame (^^^; ^^^) from the de-noisedversion of the current noisy frame (^^^), e.g. using a noise schedule as described above (step 504). The de-noised version of the current noisy frame (^^^) at the final time step comprises the generated frame of data (step 506).

[0130] Merely as an illustrative example, a U-ViT architecture with up to 5B parameters can be used for both the teacher diffusion model and the student diffusion model, and the student diffusion model can be trained for up to 200,000 steps (e.g. using accelerator TPUs). In general hyperparameters can be chosen as described in Karras et al, “Elucidating the Design Space of Diffusion-Based Generative Models”, arXiv:2206.00364, 2022.

[0131] FIG.6 shows two examples of images generated by a student diffusion model trained (distilled) in this way, with a conditioning input comprising the text shown in the respective image captions. The examples of FIG.6 were generated using 8 sampling steps.

[0132] In general a diffusion model neural network can be a neural network that has been trained to process a diffusion input comprising a current noisy frame of data and data specifying a current diffusion time to generate a diffusion output that defines an estimate (given the current diffusion time) of either a noise component of the current noisy frame of data, i.e. an estimate of the noise that has been added to an original frame of data to generate the current noisy frame of data; or of a de-noised version of the current noisy frame of data.An initial noisy version of the frame of data can be sampled from a (multivariate) Gaussian noise distribution.

[0133] As previously described a frame of data may comprise or represent data in an output space, e.g. image space, or in a (lower dimension) latent space representation of the output space.

[0134] As used herein an image can be a still or moving image, in monochrome or color (including in non-visible wavelengths), in 2D or in 3D; and can includes a LIDAR point cloud (a “pixel” may then be a point of the point cloud).

[0135] There are several ways of using a diffusion model to generate a moving image i.e. video. As one example the temporal axis dimension can be treated as an additional spatial dimension. As another example a rolling diffusion technique as described in Ruhe et al., arXiv:2402.09470v1, can be used.

[0136] An image, or training image, as described herein may be an image that represents a real world environment. For example a training image may be captured from the real world by a camera or other image sensor. In a training image, or an image generated by a diffusion model trained as described herein, objects represented in the image may comprise physical real-world objects.

[0137] In some implementations the frames of data need not necessarily represent images and may, as some other examples, represent audio data, weather or climate data, or fluid mechanics data. In some implementations an image frame may encode another type of data, e.g. as a spectrogram.

[0138] As one particular example, audio or other data may be represented as a spectrogram. The audio or other data may be processed to generate a spectrogram representing the audio or other data by performing a time-frequency domain transform on an audio or other signal to generate a frequency domain representation of the audio or other signal for a range of frequencies. There are many suitable time-frequency domain transforms; as one example a short-time Fourier transform, STFT. The audio or other data (signal) may be decoded from a generated spectrogram by applying the inverse of the time-frequency domain transform, i.e. a frequency-time domain transform, e.g. an inverse STFT.

[0139] In general a spectrogram can be an image that represents the time-frequency domain transform, e.g. a representation with time on one axis, e.g. a horizontal axis, and frequency on another axis, e.g. a vertical axis. The location of a pixel of the image along the time axis can represent a time position in the audio or other data; the location along the frequency axis can represent a frequency at that time; the value of the pixel e.g. a luminance or color value, canrepresent a component of the audio or other signal at that time and at that frequency, e.g. a magnitude and / or a phase of the signal. In some time-frequency domain transforms the component of the signal is represented by a complex number. Generally, the spectrogram can represent a changing spectrum of an audio or other signal over time.

[0140] Where a moving image is generated this can comprise, as some examples, a continuation of a previous video sequence, an edited version of a video sequence, or a video sequence generated to represent a text or audio conditioning input.

[0141] As another example, implementations of the system can be used to generate an image representing a predicted state or trajectory of a real-world physical system, such as a robot or vehicle, for use by a control algorithm in controlling the physical system. For example a video sequence may be captured by a camera and then the system used to generate a still or moving image that predicts a future state or configuration of the physical system, optionally conditioned on one or more variables relating to the physical system. This can be used in a model predictive control system to control a mechanical agent such as a robot to perform a particular task, by processing the predicted image using the control system to generate control signals to control the mechanical agent, in accordance with the generated image to perform the task.

[0142] As some other examples, a frame of data may comprise, e.g., sensor data such as data for an EEG or ECG trace; or medical imaging data such as CT or MRI data; or data representing a 3D configuration of a molecule, e.g. as coordinates of atoms of the molecule.

[0143] In some implementations the teacher and / or student diffusion model may be trained to perform an image processing task. The task may be specified by the conditioning data, which may comprise an image to be processed.

[0144] That is, the student diffusion model may be trained to perform an image processing task, which can involve obtaining a reference image to guide the image modification task, and using the reference image to modify, i.e. guide, the predicted frame generated by the student diffusion model (at each iteration, in inference), e.g. by conditioning the student diffusion model on a representation of the reference image. As some examples, such an approach can be used to perform image modification tasks such as image colorization, super- resolution, in-painting, out-painting, interpolation, stroke-guided image generation, and de- noising.

[0145] 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 thenecessary 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.

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

[0147] 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, a 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.

[0148] 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 GPU programming or TPU utilization, depending on the chosen hardware platform and desired performance characteristics.

[0149] 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, model 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.

[0150] 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 functionsby 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.

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

[0152] 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 dependon factors such as the size of the data, access speed requirements, cost considerations, and the desired level of portability or permanence.

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

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

[0155] 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, can 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.

[0156] The computing system can include clients and servers that may be geographically separated and interact through a communication network. The specific type of network, suchas 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.

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

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

[0159] 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, theprocesses 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.

[0160] What is claimed is:

Claims

CLAIMS 1. A computer implemented method of training a diffusion model, comprising: obtaining a teacher diffusion model configured to process a noisy frame corresponding to a diffusion time between an initial time and a final time to generate a teacher predicted frame for de-noising the noisy frame; initializing a student diffusion model, the student diffusion model configured to process a noisy frame corresponding to a diffusion time to generate a student predicted frame for de-noising the noisy frame; and training the student diffusion model using the teacher diffusion model by: obtaining a training frame from a training dataset; determining a target time and a training time that is within a duration of a time step after the target time; processing the training frame using a noise schedule that defines a variation of a frame noise level with time to determine a noisy training frame for the training time; processing the noisy training frame using the student diffusion model to generate a student predicted frame that defines a predicted denoised frame for the final time; processing the noisy training frame and the predicted denoised frame for the final time, using the noise schedule, to determine a reduced noise training frame corresponding to the target time; and updating learnable parameters of the student diffusion model using a student model training objective dependent on the predicted denoised frame for the final time and the reduced noise training frame.

2. The method of claim 1, wherein updating the learnable parameters of the student diffusion model comprises backpropagating gradients of the student model training objective and applying a stop gradient to the reduced noise training frame so as not to pass gradients into the student diffusion model via the reduced noise training frame.

3. The method of claim 1 or 2, comprising: processing the reduced noise training frame corresponding to the target time using the teacher diffusion model to generate the teacher predicted denoised frame; andupdating learnable parameters of the student diffusion model using the student model training objective dependent on the predicted denoised frame for the final time and the teacher predicted denoised frame.

4. The method of any of claims 1-3, further comprising: maintaining an auxiliary diffusion model configured to process a noisy frame corresponding to a diffusion time to generate an auxiliary predicted frame for de-noising the noisy frame; processing the reduced noise training frame using the auxiliary diffusion model to generate an auxiliary predicted frame that defines an auxiliary predicted denoised frame for the final time; processing the reduced noise training frame using the teacher diffusion model to generate a teacher predicted frame that defines a teacher predicted denoised frame for the final time; updating learnable parameters of the auxiliary diffusion model using an auxiliary model training objective dependent on a difference between the predicted denoised frame for the final time and the auxiliary predicted denoised frame for the final time; and updating learnable parameters of the student diffusion model using the student model training objective dependent on a similarity between i) the predicted denoised frame for the final time and ii) a difference between the auxiliary predicted denoised frame for the final time and the teacher predicted denoised frame for the final time.

5. The method of claim 4, wherein updating the learnable parameters of the student diffusion model comprises backpropagating gradients of the student model training objective and applying a stop gradient to the difference between the auxiliary predicted denoised frame for the final time and the teacher predicted denoised frame for the final time.

6. The method of claim 4 or 5, further comprising regularizing the auxiliary model training objective using the difference between the auxiliary predicted denoised frame for the final time and the teacher predicted denoised frame for the final time.

7. The method of any of claims 1-3, comprising: obtaining a first training frame and a second training frame from the training dataset;processing the first training frame using the noise schedule to determine a noisy first training frame for the training time; processing the second training frame using the noise schedule to determine a noisy second training frame for the training time; processing the noisy first training frame using the student diffusion model to generate a first student predicted frame that defines a first predicted denoised frame for the final time; processing the noisy second training frame using the student diffusion model to generate a second student predicted frame that defines a second predicted denoised frame for the final time; processing the noisy second training frame and the second predicted denoised frame for the final time using the noise schedule, to determine a second reduced noise training frame corresponding to the target time; processing the second reduced noise training frame using the teacher diffusion model to generate an auxiliary predicted frame; determining a teacher model gradient of an auxiliary objective function dependent upon on a difference between the second predicted denoised frame for the final time and the auxiliary predicted denoised frame; and updating learnable parameters of the student diffusion model using the student model training objective dependent on the first predicted denoised frame for the final time and the teacher model gradient.

8. The method of claim 7, comprising determining the student model training objective dependent on a product of the first predicted denoised frame for the final time and a vector proportional to the teacher model gradient.

9. The method of claim 7 or 8, further comprising: processing the noisy first training frame and the first predicted denoised frame for the final time using the noise schedule, to determine a first reduced noise training frame corresponding to the target time; determining the teacher model gradient with respect to learned parameters of the teacher diffusion model; determining a scaled teacher model gradient by multiplying the teacher model gradient by a scaling matrix;determining a Jacobian of the learned parameters of the teacher diffusion model evaluated for the first reduced noise training frame; determining a product of the Jacobian of the learned parameters of the teacher diffusion model and the scaled teacher model gradient to obtain a training product; and updating the learnable parameters of the student diffusion model using the student model training objective dependent on a product of the first predicted denoised frame for the final time and the training product.

10. The method of claim 9, wherein updating the learnable parameters of the student diffusion model includes applying a stop gradient to the training product.

11. The method of any of claims 7-10, comprising: obtaining a batch of first training frames from the training dataset; obtaining a batch of second training frames from the training dataset; determining the student model training objective from the batch of first training frames; and determining the teacher model gradient from the batch of second training frames.

12. The method of any of claims 1-11, wherein the student predicted frame comprises a predicted de-noised frame for the final time; and comprising processing the noisy training frame using the student diffusion model to generate the predicted denoised frame for the final time.

13. The method of any of claims 1-12, wherein the teacher diffusion model and the student diffusion model each comprise a respective de-noising neural network configured to process the diffusion time and the noisy frame corresponding to the diffusion time to generate the respective teacher predicted frame or student predicted frame.

14. A computer implemented method of training a diffusion model, comprising: obtaining a teacher diffusion model configured to process a noisy frame corresponding to a diffusion time between an initial time and a final time to generate a teacher predicted frame for de-noising the noisy frame;initializing a student diffusion model, the student diffusion model configured to process a noisy frame corresponding to a diffusion time to generate a student predicted frame for de-noising the noisy frame; and training the student diffusion model using the teacher diffusion model by matching one or more statistical moments of a distribution of predicted frames from the student diffusion model to one or more statistical moments of the distribution of predicted frames from the teacher diffusion model.

15. The method of claim 14, comprising training the student diffusion model using a moment matching objective that depends on the distribution of predicted frames from the student diffusion model and on the distribution of predicted frames from the teacher diffusion model and that is minimized when training the teacher diffusion model on predicted frames from the student diffusion model would not change learned parameters of the teacher diffusion model.

16. The method of claim 14 or 15, comprising: obtaining a training frame from a training dataset; and wherein matching the one or more statistical moments comprises matching a statistical moment defined by a denoising loss determined by a difference between a first predicted frame generated by the student diffusion model and a second predicted frame generated by the teacher diffusion model.

17. The method of claim 16, comprising: obtaining the first predicted frame by processing a noisy version of the training frame using the student diffusion model; obtaining the second predicted frame from the first predicted frame and the noisy version of the training frame according to a noise schedule that defines a variation of a frame noise level with time.

18. The method of claim 17, wherein obtaining the second predicted frame comprises: determining a reduced noise training frame from the first predicted frame and the noisy version of the training frame according to the noise schedule; andprocessing the reduced noise training frame using the teacher diffusion model to generate the second predicted frame.

19. The method of claim 14 or 15, comprising: obtaining a training frame from a training dataset; wherein matching the one or more statistical moments comprises matching a statistical moment defined by a denoising loss determined by a difference between a first predicted frame generated by the student diffusion model and a second predicted frame generated using an auxiliary diffusion model; the method further comprising training the auxiliary diffusion model to generate predicted frames that match predicted frames generated by the teacher diffusion model.

20. The method of claim 19, comprising: determining a reduced noise training frame from a predicted frame from the teacher diffusion model generated by processing a training frame and a noisy version of the training frame; processing the reduced noise training frame using the auxiliary diffusion model to generate an auxiliary predicted frame; processing the noisy training frame using the teacher diffusion model to generate a teacher predicted frame; and training the auxiliary diffusion model using an objective that depends on a difference between the auxiliary predicted frame and the teacher predicted frame.

21. A computer implemented method of training a diffusion model, comprising: obtaining a teacher diffusion model configured to process a noisy frame corresponding to a diffusion time between an initial time and a final time to generate a teacher predicted frame for de-noising the noisy frame; initializing a student diffusion model, the student diffusion model configured to process a noisy frame corresponding to a diffusion time to generate a student predicted frame for de-noising the noisy frame; obtaining a training frame from a training dataset; applying a forward diffusion process to the training frame to generate a noisy version of the training frame;processing the noisy version of the training frame using the student diffusion model to generate a predicted frame comprising a de-noised predicted frame that is a de-noised version of the noisy version of the training frame; obtaining a reduced noise version of the noisy version of the training frame from the de-noised predicted frame and the noisy version of the training frame according to a noise schedule that defines a variation of a frame noise level with time; training the student diffusion model to minimize, with respect to learnable parameters of the student diffusion model, an objective dependent on a gradient determined from the teacher diffusion model and evaluated using the reduced noise version of the noisy version of the training frame.

22. The method of claim 21, wherein the gradient determined from the teacher diffusion model comprises a gradient of a loss function that depends on a difference between i) the de- noised predicted frame and ii) a teacher de-noised predicted frame obtained by processing the reduced noise version of the noisy version of the training frame using the teacher diffusion model.

23. The method of any of claims 1-22, comprising initializing the student diffusion model by initializing learnable parameters of the student diffusion model from corresponding learned parameters of the teacher diffusion model.

24. The method of any of claims 1-23, wherein the teacher diffusion model and the student diffusion model comprise, respectively, a teacher diffusion model neural network and a student diffusion model neural network each configured to process a noisy frame corresponding to a diffusion time, and the diffusion time, to generate, respectively, the teacher predicted frame and the student predicted frame.

25. The method of any of claims 1-24, wherein the teacher diffusion model and the student diffusion model are each further configured to process conditioning data to generate the respective teacher predicted frame or student predicted frame, wherein the respective teacher predicted frame or student predicted frame characterizes the conditioning data.

26. A computer-implemented method of using a trained diffusion model to generate a frame of data,wherein the trained diffusion model comprises a student diffusion model trained using the method of any one of claims 1-25, the trained diffusion model configured to process a noisy frame corresponding to a diffusion time between an initial time and a final time to generate a predicted frame for de- noising the noisy frame, the method comprising: initializing a current noisy frame; iteratively, for each of a discrete number of time steps between the initial time and the final time until a final time step: processing the current noisy frame using the diffusion model to generate a predicted frame comprising either a prediction of noise in the current noisy frame or a prediction of a de-noised version of the current noisy frame; determining the de-noised version of the current noisy frame from the predicted frame; determining an updated current noisy frame from the de-noised version of the current noisy frame, using a noise schedule that defines a variation of a frame noise level with time; wherein the de-noised version of the current noisy frame at the final time step comprises the generated frame of data.

27. A computer-implemented method of using a trained diffusion model to generate a frame of data, the trained diffusion model configured to process a noisy frame corresponding to a diffusion time between an initial time and a final time to generate a predicted frame for de- noising the noisy frame, wherein the trained diffusion model has been trained using predicted frames sampled from a distribution of predicted frames generated by a teacher diffusion model configured to process a noisy frame corresponding to a diffusion time between an initial time and a final time to generate a predicted frame for de-noising the noisy frame, the training comprising matching one or more statistical moments of a distribution of predicted frames from the trained diffusion model to one or more statistical moments of the distribution of predicted frames from the teacher diffusion model, the method comprising: initializing a current noisy frame; iteratively, for each of a discrete number of time steps between the initial time and the final time until a final time step:processing the current noisy frame using the diffusion model to generate a predicted frame comprising either a prediction of noise in the current noisy frame or a prediction of a de-noised version of the current noisy frame; determining the de-noised version of the current noisy frame from the predicted frame; determining an updated current noisy frame from the de-noised version of the current noisy frame, using a noise schedule that defines a variation of a frame noise level with time; wherein the de-noised version of the current noisy frame at the final time step comprises the generated frame of data.

28. The method of claim 26 or 27, wherein the trained diffusion model is further configured to process conditioning data to generate the predicted frame, wherein the predicted frame characterizes the conditioning data.

29. The method of any of claims 1-28, wherein the frame of data comprises an image frame that defines values of pixels of a still or moving image.

30. One or more non-transitory computer storage media storing instructions that when executed by one or more computers cause the one or more computers to perform operations of the respective method of any one of claims 1-29.

31. A system comprising one or more computers and one or more storage devices communicatively coupled to the one or more computers, wherein the one or more storage devices store instructions that, when executed by the one or more computers, cause the one or more computers to perform operations of the respective method of any one of claims 1-29.

Images