Multi-step distillation of diffusion model by moment matching

Abstract

Systems and methods for training diffusion models for generating data frames, such as image data frames, implemented as computer programs. The described techniques use a teacher diffusion model to train or "distill" a student diffusion model that may generate data frames using much fewer steps (e.g., less than 20 time steps) than the teacher diffusion model. A trained student diffusion model may be used to generate the data frame.

Description

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 for training a diffusion model for generating data frames, such as image data frames, implemented as computer programs on one or more computers in one or more locations. A method of generating data frames using a trained diffusion model is also described.

[0004] Generally, a diffusion model can generate a data frame by performing a reverse diffusion process to gradually denoise an initial noisy data frame over a series of time steps until reaching a final output data frame. Multiple time steps are typically required. The described techniques train or “distill” a student diffusion model using a teacher diffusion model that can generate a data frame using fewer steps, e.g., fewer than 100, 50, or 20 time steps. The trained student diffusion model can be used to generate data frames.

[0005] The techniques are conceptually based on a variant of a matrix matching process, such as expected matching. More specifically, given noisy data from a sampled trajectory of a diffusion process, implementations of the techniques match a conditional expectation of clean data between a training data distribution and a sampled distribution from a distilled student model. Noisy data is partially denoised data, in particular data obtained by denoising noisy data at points in a forward diffusion process that complements a reverse diffusion process.

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

[0007] In related aspects, a computer-implemented method of training a diffusion model is described that involves training a student diffusion model using a 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 a distribution of predicted frames from the teacher diffusion model.

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

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

[0010] Further described is a system that includes one or more computers and one or more storage devices that are 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 described methods.

[0011] Also described are one or more non-transitory computer storage media storing instructions that, when executed by one or more computers, perform the operations of one or more 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 at a significant computational burden. Implementations of the described technology can significantly alleviate this burden by reducing the number of diffusion model steps required in inference by one or two orders of magnitude or more, while still generating high-quality samples.

[0014] Multi-step diffusion models are generally considered the gold standard for data sample generation. However, contrary to intuition, student diffusion models trained as described herein can even improve the quality of data samples generated by a teacher diffusion model. Broadly speaking, this can be achieved by equalizing the prediction error that arises at each diffusion step of the teacher model. Moreover, the technology described herein can be used with relatively poor teacher diffusion models.

[0015] The described technology can be used with or without conditioned data, and thus can be used to generate data frames with particular characteristics or for particular purposes, or to modify data frames 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 dynamic 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 DRAWINGS

[0018] Figure 1 An example system for training a diffusion model is shown.

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

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

[0021] Figure 4 Operations of a process for training a diffusion model are shown.

[0022] Figure 5 is a flow diagram of an example process for generating data frames using a trained diffusion model.

[0023] Figure 6 An example of an image generated by a diffusion model trained as described herein is shown.

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

[0025] This specification describes computer-implemented methods of training a diffusion model, particularly 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—also a diffusion time in implementations—to generate a teacher predicted frame for denoising the noisy frame. Processing the diffusion time can improve the quality of the teacher predicted frame.

[0027] In general, a predicted frame is a data frame generated by a diffusion model described herein, particularly by a neural network of a diffusion model when the diffusion model processes an input data frame. Typically, the predicted data frame has the same dimensions as the input data frame. References herein to sampling or processing values of a data frame refer to sampling or processing values of data elements defined for the data frame.

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

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

[0030] The techniques described in this paper can be used unchanged, whether operating in the output space (i.e., the space that generates the output data items) or in the latent space (i.e., the space of the latent representation of the output space).

[0031] The diffusion model trained as described in this paper, particularly the student diffusion model, can be used to generate data frames by performing a series of denoising steps. The generated data frames can be data frames in the output space, meaning the values ​​of the output data items are values ​​of appropriate types of data items, such as image pixel values, audio signal amplitude values, etc.; or the generated data frames can be data frames in the latent space, meaning the values ​​in the data frames are values ​​in the latent representation of the output data items in the output space. Typically, this latent space has a lower dimension than the output space.

[0032] When generating data frames in the latent space, the system described herein can generate final data frames in the output space by processing the data frames in the latent space using a decoder neural network (e.g., a decoder neural network already pre-trained in an autoencoder framework). During training, the system can encode training data items in the output space using an encoder neural network (e.g., an encoder neural network already pre-trained in conjunction with the decoder in an autoencoder framework), thereby generating training data items for a diffusion neural network in the latent space.

[0033] In the implementation of the system described, the data frame is a latent representation, and the values ​​in the representation are learned latent values, rather than pixel values, for example, when the data frame represents an image.

[0034] Generally, a denoised version of the input data frame can be generated using predicted data frames from either the teacher diffusion model or the student diffusion model. For example, in some implementations, the teacher prediction frame includes predictions of noise in the noisy frame, which can then be subtracted from the noisy frame to obtain the prediction of the denoised frame, such as a fully denoised frame. In some implementations, the teacher prediction frame includes predictions of the denoised frame, such as a fully denoised frame.

[0035] The diffusion time lies between the initial time and the final time. Generally, the initial time corresponds to the initial noisy frame processed during inference, such as an initial frame containing only noise. Generally, the final time corresponds to the last frame of data generated during inference, such as a supposedly fully denoised frame.

[0036] The method trains a student diffusion model. The student diffusion model is configured to process noisy frames corresponding to diffusion time, and optionally also to process diffusion time, to generate student prediction frames for denoising noisy frames, such as predictions of noise in noisy frames or predictions of denoised frames.

[0037] In inference, when used to generate data frames, the diffusion time processed by the student diffusion model corresponds to the number of discrete (sampling) time steps at the corresponding step time between the initial and final times. k ).For example or ,For example or In the implementation method, Typically, diffusion models perform the denoising process over hundreds or thousands of time steps. The described technique is implemented in a way that generates data frames that are as good as, or (contrary to intuition) better than, those of the teacher, while using only, for example, eight steps.

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

[0039] refer to Figure 1 This figure shows an example system 100 used for training a diffusion model. Figure 1 The system can be implemented as one or more computer programs on one or more computers in one or more locations.

[0040] System 100 includes a teacher diffusion model 110, which is configured to process noisy frames 112 corresponding to diffusion time 114—and in an implementation, diffusion time 114 is also processed—to generate teacher prediction frames 118.

[0041] System 100 also includes a student diffusion model 120, which is configured to process noisy frames 122 corresponding to diffusion time 124—in some implementations, diffusion time 124 is also processed—to generate learning prediction frames 128.

[0042] Some implementations of system 100 also include an auxiliary diffusion model 150, which is configured to process noisy frames 152 corresponding to diffusion time 154—in some implementations, diffusion time 154 is also processed—to generate auxiliary prediction frames 158.

[0043] In the implementation, the teacher diffusion model 110, the student diffusion model 120, and (if present) the auxiliary diffusion model 150 each comprise a teacher diffusion model (denoising) neural network, a student diffusion model (denoising) neural network, and an auxiliary diffusion model (denoising) neural network. Each of these models is configured to process noisy frames corresponding to the diffusion time—and in the implementation, also the diffusion time—to generate teacher prediction frames 118, student prediction frames 128, and auxiliary prediction frames 158, respectively. For example, one of the corresponding models can be used to process noisy training frames to generate denoised prediction frames for the final time.

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

[0045] Teacher diffusion model (denoising) neural networks and student diffusion model (denoising) neural networks can have any suitable architecture that corresponds to the processing of the values ​​of the data frames as input to generate the corresponding set of output values ​​for the data frames. Generally, these neural networks can each include one or more feedforward layers, convolutional layers, attention layers, normalization layers, or other neural network layers.

[0046] For example, 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 thereof, such as the U-ViT architecture (Bao, et al., “All are Worth Words: A ViT Backbone for Diffusion Models”, arXiv:2209.12152, 2023), such as one of the U-Net or U-ViT model configurations described in Appendix B of the following literature: Hoogeboom et al. “simple diffusion: End-to-end diffusion for high-resolution images”, arXiv:2301.11093, 2023. For example, such a neural network can include a U-Net with one or more ResNet blocks and one or more self-attention layers, or can include a diffusion transformer (DiT) or a transformer backbone.

[0047] Using such neural networks to process time typically involves processing data that specifies that time (e.g., an embedding of that time). For example, time can be encoded as an embedding, such as a sine position embedding, and that embedding can be added to or otherwise combined with each processed data element. For example, time can be provided as auxiliary information to one or more layers of the neural network.

[0048] In the implementation, the teacher diffusion model 110 and / or the student diffusion model 120—specifically, the teacher diffusion model neural network and the student diffusion model neural network—are also configured to process conditional data to generate corresponding teacher prediction frames 118 or student prediction frames 128. The neural networks can be configured to focus on the conditional data, for example, by incorporating one or more cross-attention layers, or to process the conditional data in some other way. The conditional data can come from the conditional input ( Figure 1 (Not shown in the image).

[0049] Generally, conditional data, such as conditional input, characterizes the content of teacher prediction frame 118 or student prediction frame 128. That is, conditional data defines one or more properties of the generated data frame (in the latent space or output space). For example, conditional input may include text or audio, or lexical units, or embeddings representing text or audio features, while the prediction frame may include pixel data of a still image or moving image representing text or audio.

[0050] Such conditional inputs may include (but are not limited to) text data and / or image data, each of which may be encoded into an embedding sequence, for example, using a text, image, or multimodal Transformer model, such as an encoder-decoder model, like T5 (Text-to-Text Transfer Transformer, Raffel et al., arXiv:1910.10683, 2023) or a language model.

[0051] In the training methods described later, conditional data can be incorporated into the training of a teacher diffusion model neural network or a student diffusion model neural network, for example, using so-called classifier-free guidance. This may involve randomly masking the conditional data or otherwise removing it from the neural network input so that the neural network can be trained to generate corresponding prediction frames with and without guidance from the conditional data. Classifier-free diffusion model guidance is described, for example, in the following literature: Ho and Salimans, arXiv:2207.12598. During inference, data frames can be generated from a weighted combination of conditional and unconditional prediction frames, for example, from a (1+ wWeighted conditionalized prediction frames and (1+ w Data frames are generated from the weighted differences between unconditional prediction frames, where w It has a positive weight.

[0052] Figure 2 This is a flowchart of an example process for training a diffusion model (such as the student diffusion model 120). Figure 2 The process can be performed by one or more computers in one or more locations; for convenience, refer to Figure 1 Describe the process.

[0053] At step 200, the process initializes the student diffusion model 120. For example, the student diffusion model can be initialized from the teacher diffusion model, for instance, by initializing the learnable parameters (e.g., weights) of the student diffusion model 120 from the 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 necessary. For example, the student diffusion model 120 can be initialized by randomly initializing the learnable parameters of the student diffusion model.

[0054] The student diffusion model 120 is trained using the teacher diffusion model 110. Steps 202 to 214 described below are repeated as many times as appropriate during the training process, for example, until the training loss tends to plateau and reaches an approximate minimum.

[0055] In the implementation, training involves processing the training dataset ( Obtain (e.g., sample) training frames ( (Step 202).

[0056] This process determines the target time ( In some implementations, the target time... By sampling the distribution from the target time ( The result is obtained by randomly sampling from (step 204). Any distribution can be used, such as a uniform distribution. In some implementations, the target time... It is possible to use (evaluate) student diffusion models from a fixed set of discrete time intervals—for example, in inference. k Choose from a set of time steps. Randomly sampling the target time may help improve stability and performance. Generally, the target time is between the initial time (e.g., 1) and the final time (e.g., 0) mentioned above.

[0057] This process determines the training time. Specifically, this refers to the training time over a time step following the target time (step 206). In some implementations, the training time is randomly sampled from a distribution; in others, it can be a fixed time following the target time. The training time can also be called the sampling time. Generally, the sampling time... t Between the initial time (e.g., 1) and the final time (e.g., 0) mentioned above.

[0058] For example, it can be derived from the distribution—for example, in the range. Uniform distribution within — sampling time ,in This refers to the number of discrete (sampling) time steps used during inference. Alternatively, time... It can be fixed. Training time ( ) can be truncated to 1, for example In some implementations, training time... It is continuous, even during reasoning, only when... The student diffusion model is evaluated at each time step. In some implementations, the sampling time is... The sampling is random. These measures may help improve stability and performance, but they are not necessary. For example, training time... It doesn't have to be continuous; and time... It can be fixed to a single value.

[0059] Training frames are processed using noise scheduling. This noise scheduling defines how frame noise levels change over time to determine the training time. t Noisy training frames () (Step 208). This may involve obtaining a noisy frame ( And process the noisy frames and training frames to determine the noisy training frames. ).

[0060] Generally, noise scheduling defines the change in signal-to-noise ratio (SNR) from the final time to the initial time during the denoising process, a change followed by both teacher diffusion and student diffusion models. Therefore, noise scheduling can define the SNR at training time (t) and target time (s). Generally, the SNR monotonically decreases from the initial time to the final time.

[0061] Noise scheduling can define the time-dependent average scaling factor of a frame. ) and the time-dependent variance scaling factor of the frame ( And SNR can be defined as Optionally, but not necessarily, determining one of the mean scaling factor and the variance scaling factor can determine the other, for example, according to... It defines variance-preserving noise (SNR) scheduling. However, using variance-preserving noise scheduling is not mandatory. In some implementations, , , and for example .

[0062] Within this framework, noisy training frames can be identified as ,in and In other words, It can be used as a distribution .

[0063] Noisy training frames were processed using Student diffusion model 120. To generate student prediction frames, which define (e.g.) the final time (student) predicted denoised frames. (Step 210).

[0064] Noisy training frames are processed using noise scheduling. ) and the final time (student) predicted denoised frames ( ), to determine the target time ( Corresponding (partial) noise reduction training frames (Step 212). Noise reduction training frames ( The noise reduction training frame () can be determined randomly, for example, based on another sampled noisy frame, or deterministically, for example, based on previously sampled noisy frames. Generally, the noise reduction training frame () It is still noisy, but better than the noisy training frames. The noise is low.

[0065] As a specific example, the denoised training frame ( ) can be obtained from Sampling; or equivalently from The middle sampling is used to determine, among which

[0066] and and .

[0067] As another specific example, in deterministic scenarios such as DDIM (Song et al., “Denoising Diffusion Implicit Models”, arXiv:2010.02502v4, October 2022), .

[0068] The student model's training objective (e.g., loss) is used to update the learnable parameters (e.g., weights) of the student diffusion model (step 214). Typically, as described herein, updating the model's learnable parameters (e.g., weights) involves backpropagating the gradient of the objective function (e.g., loss) to update the learnable parameters. This can be done using any suitable gradient descent optimization algorithm, such as Adam or another optimization algorithm.

[0069] Generally, the training objective (e.g., loss) of the student model depends on the final (student) predicted denoised frames. ) and noise-reduced training frames ( In particular, the denoising training frame was processed using a teacher diffusion model to generate teacher-predicted denoised frames. ,in This represents a teacher diffusion model with parameters (e.g., weights). ).

[0070] In a general sense, the goal of student model training is to train the student diffusion model so that the conditional expectation of the student diffusion model (i.e., the student's final prediction of the denoised frame) is... Conditional expectation of the teacher diffusion model (teacher predicts denoised frames) Matching. In this implementation, the expectation can be computed using a single sample (Monte Carlo estimation). During training, the matching is actually averaged over the samples from the diffusion process, i.e., at the target time ( s Averaged on top.

[0071] In the implementation method, the noise reduction training frame corresponding to the target time ( The teacher diffusion model has been used to generate teacher-predicted denoised frames. ), and this processing is directly or indirectly related to the final time (student) predicted denoised frames ( Used in combination to determine the loss.

[0072] For example, in some specific implementations, the training objective of the student model depends on the predicted denoised frames at the final time ( ) and the teacher-predicted denoised frames at the final time In some specific implementations, the training objective of the student model depends on the predicted denoised frames at the final time ( ) and in Jacobian matrix of the teacher diffusion model (neural network) for assessment.

[0073] In some implementations, updating the learnable parameters of the student diffusion model involves backpropagating the gradient of the student model's training objective and applying the stopping gradient to the denoised training frames. ), to avoid passing through the noise-reduced training frames ( The gradient is passed to the student diffusion model. This provides a stable objective and promotes convergence when training the student diffusion model.

[0074] The above process can be understood as moment matching, specifically matching one or more statistical moments of the distribution of predicted frames from the student diffusion model with one or more statistical moments of the distribution of predicted frames from the teacher diffusion model. This is similar to simply training the student diffusion model to generate accurate samples. Different, among them Indicates parameters (e.g., weighted) student diffusion models.

[0075] A more detailed description of this will help to understand how the invention is implemented.

[0076] Consider the following algorithm, which demonstrates the use of... k A typical random sampling process for generating data frames from a diffusion model in 1 sampling step (where in each step, Become the next ): Requirement: Denoising model Number of sampling steps k Initialize noisy data

[0077] for {1, ( k -1) / k ,…, 2 / k , 1 / k} do

[0078] use Predicting clean data

[0079] Set a time step s = t - 1 / k

[0080] For the next noisy data point Sampling

[0081] end for

[0082] Return approximate sample

[0083] For conventional diffusion models, this process typically requires at least several hundred sampling steps. This type of algorithm (or deterministic sampler, such as DDIM) can be used to generate data frames using a student diffusion model trained as described herein. However, when the student diffusion model is trained as described, the process can be performed (unmodified) with, for example, fewer than 20 sampling steps, and samples that are as good or better than those of the teacher diffusion model can be obtained. This requires the student diffusion model to take significantly more steps than the teacher diffusion model during the (backward) diffusion process.

[0084] Unlike training a student diffusion model, it enables Approximate to expectation (in q Representing data distribution ), train the student diffusion model to generate from the distribution For approximate samples, in this case, the above algorithm is applicable to any number of sampling steps. k All can be derived from data distribution q This generates accurate samples.

[0085] In other words, the output (conditional expectation) of the student diffusion model is trained to match the output (conditional expectation) of the teacher diffusion model given the same noisy input (and averaged over the diffusion trajectory). In this way, the student diffusion model learns to model the statistical effect of a set of steps, rather than replicating each step, which is a more efficient approach than learning to generate... A more robust training process. Even with fewer steps, the Student diffusion model still generates samples that closely match the expected distribution of the real data.

[0086] Teacher diffusion models may have prediction errors, and these errors may vary at different time steps. The distillation process correlates the predictions from clean data at different diffusion times and helps ensure that these predictions are consistent. Empirically, student diffusion models can generate better samples than teacher diffusion models.

[0087] Conditional expectation of the teacher diffusion model (For example ) can be achieved through and The conditional expectation of the student diffusion model is determined by sampling. By means of , , and Sampling is performed to determine this. The moments between the student diffusion model and the teacher diffusion model can be matched by training the student diffusion model using a moment matching objective (loss), which depends on these expectations (…). and Differences between them, such as L2 distance For both the student diffusion model and the teacher diffusion model, the same noisy input can be used. The loss is determined on the distribution of (denoised training frames). For clarity, the dependence of model predictions on diffusion time is omitted here and thereafter, but it can be taken into account in the implementation.

[0088] In practice, while not strictly necessary, stopping gradients can be used to prevent backpropagation from passing through the target (loss) gradient during backpropagation. Extending back to the student diffusion model, the objective (the expected difference) is considered constant. Then, the gradient of the objective is approximated as... And the objective (loss) can be simplified to (May have additional regularization), where This indicates stopping the gradient (the loss is different from L2 loss, but the gradient of the loss is approximately the same). If This loss is minimized. Using stopping gradients in this way, compared to using full gradients, helps stabilize training.

[0089] depending on This simplified loss depends on the similarity between: i) the predicted denoised frame at the final time, and ii) the difference between the auxiliary predicted denoised frame at the final time and the teacher predicted denoised frame at the final time. Here, similarity is defined by the vector dot product (dot product or normalized dot product, i.e., cosine similarity). It can be understood as making the expected difference... It is zero, making The simplified loss, especially its gradient, is a stochastic approximation of the objective. (It should be noted that...) depending on The distribution of [the data] makes adjusting the prediction [improvement] possible. Also implicitly adjusted ). Differences in expectations In fact, it is a measure of the magnitude of the error in the student diffusion model.

[0090] In some implementations, for example, using the assisted diffusion model 150 based on The loss in the space (the space of the generated data frames) updates the learnable parameters of the student diffusion model. In some implementations, the learnable parameters of the student diffusion model are updated based on the objective, which operates (directly) in the parameter space—that is, in the space of the learnable parameters of the student diffusion model.

[0091] In depending on Of the aforementioned losses, the expected The estimate can be determined from the teacher diffusion model as follows: (Using single-sample estimation). Estimation One approach is to maintain an auxiliary diffusion model. In another approach, the loss is rewritten as a loss in the parameter space, which is essentially determined by analyzing the result of performing a gradient descent step on the loss.

[0092] Therefore, some implementations of this method involve maintaining an auxiliary diffusion model ( ) An auxiliary diffusion model 150 is configured to process noisy frames corresponding to the diffusion time to generate auxiliary prediction frames for denoising the noisy frames. The auxiliary diffusion model is used to process denoising training frames. To generate an auxiliary prediction frame, which is either defined or is an auxiliary prediction denoising frame. (The frame corresponding to the final time).

[0093] The teacher diffusion model can also be used to process the denoised training frames. To generate teacher prediction frames, which are defined or are teacher prediction denoised frames. (Also corresponds to the frame of the final time).

[0094] Then, the learnable parameters (e.g., weights) of the auxiliary diffusion model can be updated using the auxiliary model training objective, which depends on the predicted denoised frames. (corresponding to the final time) and auxiliary prediction denoised frames ( The difference between (corresponding to the final time).

[0095] In some implementations, the assisted diffusion model 150 It can be regularized to match the teacher model. Keep it close. For example, the training objective of the auxiliary model can be regularized by adding a term to the objective that depends on the difference between the auxiliary predicted denoised frame at the final time and the teacher predicted denoised frame at the final time, for example... (Then the auxiliary diffusion model predicts...) and The average value will be lost. Multiply ).

[0096] As mentioned earlier, the learnable parameters of the student diffusion model can be updated using the student model training objective, which depends on the (dot product) similarity between the following two: i) the predicted denoised frame at the final time ( ), and ii) assisted prediction of denoised frames ( ) and teacher-predicted denoised frames ( The difference between the auxiliary prediction denoised frames at the final time and the teacher prediction denoised frames at the final time can be used to stop the gradient. As mentioned earlier, the stopping gradient can be applied to the difference between the auxiliary prediction denoised frames at the final time, i.e., to avoid passing the gradient to the student diffusion model via the difference.

[0097] In some implementations, the learnable parameters of the auxiliary diffusion model and the learnable parameters of the student diffusion model can be updated alternately.

[0098] The following algorithm illustrates an example implementation of these techniques: Requirement: Pre-trained denoising model Generator for distillation g η Assisted denoising model Number of sampling steps k Time sampling distribution Loss weight and dataset .

[0099] for n=0:N do

[0100] Sampling target time Sampling time increment .

[0101] Set sampling time t =minimum ( s + ,1).

[0102] from Sample clean data and perform forward diffusion to produce .

[0103] use Sampled from the distillation generator .

[0104] if n It is an even number then

[0105] about minimize

[0106] else

[0107] about minimize

[0108] end if

[0109] end for

[0110] In the above example implementation, the training objective of the student model directly depends on the predicted denoised frames at the final time. Teacher-predicted denoised frames In another example implementation, the student model training objective indirectly depends on the predicted denoised frames at the final time. Teacher-predicted denoised frames In other words, the training objective for the student model can be in the teacher-parameter space rather than in the teacher-parameter space. x Determined in space. This can be achieved by computationally minimizing the above algorithm. The effect is achieved through this, without the need to actually maintain the auxiliary diffusion model.

[0111] From teacher parameters At the beginning, for The effect of performing a single gradient step can be determined as follows:

[0112] in Indicates in Evaluation gradient vector (Dimension of the vector) p equal to parameter (Quantity). Here, It is a teacher diffusion model The Jacobian matrix, where It is an optional ( (Diagonal) gradient preconditioning matrix; if there is no preconditioning, then It can be an identity matrix. The gradient preconditioning matrix can be, for example, the diagonal gradient preconditioning matrix used in the Adam optimization algorithm (Kingma and Ba, Adam, “A Method for Stochastic Optimization”, arXiv:1412.6980).

[0113] Teacher loss can be defined by weighting. Let's write out this loss in terms of moment matching in the teacher parameter space, where... These are optional loss weights (with zero gradient when sampled from the training distribution). "Instantaneous" loss. The estimate then depends on (e.g., equal to, except for the scaling factor and offset) The expectations, among which ; ;and .

[0114] This "instantaneous" version of moment matching loss can be interpreted as matching the teacher's gradient between the training data and the generated data. If the teacher diffusion model is in the student diffusion model... If the loss gradient is zero when evaluated on the generated data, then the loss... The loss is minimized. That is, the loss is minimized if the teacher diffusion model cannot "see" the difference between the real data and the data generated by the student diffusion model. An instantaneous version of the moment matching loss can be evaluated using independent samples (e.g., mini-batches) from the training data, from which... , and One or more independent samples (referred to as in the algorithm below) and ).

[0115] Figure 3 This is a flowchart illustrating an example implementation of this method, and it is... Figure 2 A specific example of the process. Figure 3 The process can be performed by one or more computers in one or more locations.

[0116] At step 302, the first and second training frames are obtained (e.g., sampled) from the training dataset. The first training frame is processed using noise scheduling to determine the noisy first training frame for training time. (Step 304); and use noise scheduling to process the second training frame to determine the noisy second training frame for the training time. (Step 306).

[0117] The noisy first training frame was processed using the Student diffusion model. To generate a first student prediction frame, which defines a first prediction 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 prediction frame, which defines a second prediction denoising frame. (corresponding to the final time) (step 310).

[0118] Noisy second training frame is processed using noise scheduling. ) and the second predicted denoised frame ( ), to determine the second noise reduction training frame corresponding to the target time ( (Step 310).

[0119] The second noise-reducing training frame was processed using a teacher diffusion model. To generate auxiliary prediction frames (note) Used in (Evaluation at the point of view) (step 314). This defines or corresponds to the auxiliary prediction denoised frame at the final time. .

[0120] The process can then determine the teacher model gradient of the auxiliary objective function (with respect to the learned parameters of the teacher diffusion model), which depends on the second predicted denoised frame corresponding to the final time. ) and (corresponding to the final time) auxiliary prediction denoised frames ( The difference between (step 316).

[0121] The student model training objective is then used to update the learnable parameters of the student diffusion model. In this example implementation, the student model training objective depends on the teacher model gradient. ) and the first predicted denoised frame corresponding to the final time ( (Step 318).

[0122] More specifically, this may involve determining a training objective for the student model, which depends on the first predicted denoised frame at the final time. ) and the gradient of the teacher model ( The product of vectors that are proportional.

[0123] In a general sense, the training objective is minimized when the gradient of the teacher model with respect to its learned parameters is zero. This objective trains the student model so that when the teacher model processes, for example, data obtained from the student model... (via When the gradient of the teacher model is zero, the gradient of the teacher model is zero. As mentioned above, if the teacher model cannot "see" the difference between the real data and the generated data, and its parameters do not change when trained on the data generated by the student diffusion model, then the optimization is successful.

[0124] The specific implementation involves using noise scheduling to process the noisy first training frame at the final time. ) and the first predicted denoised frame ( ), to determine the first noise reduction training frame corresponding to the target time ( In some implementations, the teacher model gradient can then be multiplied by a (predetermined) scaling matrix, such as the gradient preconditioning matrix above. Determine the scaled teacher model gradient .For example (and is a dimension of) p (vector).

[0125] The choice of scaling matrix (if used) is not critical. For example, the scaling matrix could be based on the second moment (variance) of a set of (teacher) gradients, such as a moving average based on the squared gradient, for example, as... ,in and Proportional, of which It is time The gradient at that point, and It is a small constant for numerical stability—as described in the literature by Kingma et al. (ibid.). This scaling matrix can be determined, for example, by running an optimizer over several training steps as part of the initialization process, or by loading checkpoints saved during the training of the teacher diffusion model.

[0126] The implementation of this process can be determined for the first noise reduction training frame ( The Jacobian matrix of the learned parameters of the teacher diffusion model for evaluation This method can then determine the Jacobian matrix (of dimension 1) of the learned parameters of the teacher diffusion model. ,in (The dimension of the flattened output) and the scaled gradient of the teacher model ( The product of ) is used to obtain the training product (dimension ). d The student model training objective can then depend on the first predicted denoised frame at the final time (the vector). (flattened to dimension ) d The product of the vector (the vector) and the training product. Representing the loss using the Jacobian matrix-vector product is convenient for automatic differentiation frameworks such as JAX. In the implementation, The Jacobian matrix does not actually need to be materialized; it only needs to be evaluated implicitly (in the Jacobian matrix vector product in automatic differentiation frameworks such as JAX).

[0127] In this implementation, updating the learnable parameters of the student diffusion model involves applying a stopping gradient to the training product. More specifically, this process may involve updating the learnable parameters of the student diffusion model by backpropagating the gradients of the relevant student model training objective, and applying a stopping gradient such that the learnable parameters of the student diffusion model are not affected by the training product reaching the first denoised training frame. ) and the first predicted denoised frame ( It is updated through backpropagation.

[0128] In some implementations, the method involves obtaining a batch of first training frames from the training dataset and processing each of these first training frames as described above. Similarly, the method involves obtaining a batch of second training frames (independently) from the training dataset and processing each of these second training frames as described above. Specifically, the student model training objective can be determined from a batch of first training frames, and the teacher model gradient can be determined from a batch of second training frames, for example, by averaging the corresponding batches.

[0129] The following algorithm illustrates an example implementation of these techniques, involving moment matching in the parameter space, conceptually based on an instantaneous version of the auxiliary diffusion model (swapped into two independent mini-batches each time the parameters are updated).

[0130] Requirement: Pre-trained denoising model Generator for distillation g η Gradient scaling matrix Number of sampling steps k Time sampling distribution Loss weight and dataset .

[0131] for n=0:N do

[0132] Sampling target time Sampling time increment .

[0133] Set sampling time t =minimum ( s + ,1).

[0134] from Sample two independent batches and perform forward diffusion to produce .

[0135] For two batches, use , Sampled from the distillation generator .

[0136] Assess teacher gradations on a single batch:

[0137] On another batch, regarding minimize

[0138] end for

[0139] Figure 4 It shows Figure 3The above algorithm is a specific example of the process of forward diffusion, which distributes samples from the training data distribution. Move to Then, the student diffusion model was used to... generate .exist and Random interpolation between them is used to obtain and minimize and The gradient of teacher loss.

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

[0141] More specifically, this could involve training a student diffusion model using a moment-matching objective that depends on the distributions of predicted frames from the student diffusion model and the teacher diffusion model, and that the moment-matching objective is minimized when training the teacher diffusion model on the predicted frames from the student diffusion model does not change the learned parameters of the teacher diffusion model. In other words, the student diffusion model can be trained such that the teacher diffusion model cannot "see" the differences between data frames generated by the student diffusion model and, for example, real data frames from the training dataset.

[0142] In some specific implementations, the method may involve obtaining training frames from a training dataset and matching statistical moments defined by a denoising loss (the gradient of which is determined by the difference between a first predicted frame generated by the student diffusion model and a second predicted frame generated by the teacher diffusion model.

[0143] For example, a first predicted frame can be obtained by processing a noisy version of the training frame using a student diffusion model. A second predicted frame can be obtained from the first predicted frame—and, in some implementations, from a noisy version of the training frame—based on noise scheduling.

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

[0145] For example, this could involve determining denoised training frames from prediction frames generated by processing training frames and noisy versions of those training frames. It could also involve processing denoised training frames using an auxiliary diffusion model to generate auxiliary prediction frames. It could further involve processing noisy training frames using a teacher diffusion model to generate teacher prediction frames. Finally, it could involve training an auxiliary diffusion model using an objective that depends on the difference between the auxiliary and teacher prediction frames.

[0146] In a broad sense, one implementation of the training process described above, which operates in the parameter space, can be characterized as first processing the training frames ( Sampling is performed, and then a forward diffusion process is applied to... Modified to Then, the student diffusion model is used to generate... Conditional training frames The (clean) version. Noise-reduced version Through and Obtained through (random) interpolation between [frames]. The student diffusion model is trained with respect to the training frames. The generated version makes in Minimize the gradient norm of the teacher diffusion model at the location.

[0147] For example, this could involve obtaining training frames from a training dataset and applying a forward diffusion process, i.e., a noise addition process, to the training frames to generate a noisy version of the training frames. The Student diffusion model can be used to process a noisy version of the training frames to generate prediction frames, which include denoised prediction frames. The denoised prediction frame is a denoised version of the noisy version of the training frame.

[0148] Specifically, noise scheduling ("random interpolation") can be used to predict frames from denoised frames ( ) and the noisy version of the training frames ( ) Obtain the noisy version of the training frame and the denoised version ( ).

[0149] The student diffusion model can be trained to minimize an objective with respect to learnable parameters of the student diffusion model, which depends on gradients determined from the teacher diffusion model and evaluated, for example, using a denoised version of a noisy version of the training frames. For instance, the objective can depend on the norm of the gradients, such as the L1 or L2 norm.

[0150] Optionally, in the above techniques, the student model training objective and / or the auxiliary model training objective and / or the instantaneous loss... For example, scheduling by weight can be used. Weighting is performed using defined time-related weights. This may help with stability. However, empirically, flat weights... The effect is very good, although the weight scheduling can be adjusted. The gradient magnitude is normalized at different noise levels.

[0151] In some implementations, multiple teacher diffusion models can be used. In this case, multiple moments can be matched, i.e., multiple losses can be used, one loss for each teacher model, as described above. The teacher diffusion models may have been trained on different datasets, such as images with different aspect ratios, crop sizes, etc. More specifically, the described technique is combinatorial and can combine teacher diffusion models covering different training data modalities into student diffusion models that include different modalities. For example, a first teacher diffusion model trained on image data and a second teacher diffusion model trained on dynamic data can be distilled into a student diffusion model that can generate high-quality video data; and so on.

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

[0153] Figure 5 This is a flowchart of an example process for generating data frames using a trained diffusion model, specifically a student diffusion model trained as described above. Figure 5 The process can be performed by one or more computers in one or more locations; for convenience, refer to Figure 1 The system description process.

[0154] The trained diffusion model is configured to process noisy frames corresponding to the diffusion time—and in the implementation, the diffusion time is processed, and optionally conditional data is also processed—to generate prediction frames (predicted noise, or denoised frames) for denoising the noisy frames.

[0155] This process involves initializing the current noisy frame (e.g.) (Step 500). For the discrete number between the initial time and the final time ( k Each of the ) time steps — for example, among them —Execute the example process iteratively.

[0156] Each iteration involves processing the current noisy frame using a diffusion model to generate a prediction frame, which includes either a prediction of the noise in the current noisy frame or a denoised version of the current noisy frame. Prediction of ) (step 502). Denoising version of the current noisy frame ( () can be a predicted frame or determined from a predicted frame.

[0157] Each iteration is based on the denoised version of the current noisy frame. Determine the updated current noisy frame. ), for example, using noise scheduling as described above to determine (step 504). The denoised version of the current noisy frame at the last time step ( This includes the generated data frame (step 506).

[0158] As an illustrative example only, the U-ViT architecture with up to 5 billion parameters can be used for both teacher and student diffusion models simultaneously, and the student diffusion model can be trained for up to 200,000 steps (e.g., using an accelerator TPU). Generally, the selection of hyperparameters can be described as follows: Karras et al, “Elucidating the Design Space of Diffusion-Based Generative Models”, arXiv:2206.00364, 2022.

[0159] Figure 6 Two examples of images generated by a student diffusion model trained (distilled) in this manner are shown, where the conditional input includes the text shown in the corresponding image caption. Figure 6 The example was generated using 8 sampling steps.

[0160] Generally, a diffusion model neural network can be a trained neural network that processes diffusion input, including the current noisy data frame and data at a specified current diffusion time, to generate a diffusion output that defines an estimate of the noise components of the current noisy data frame (given the current diffusion time), i.e., an estimate of the noise that has been added to the original data frame to generate the current noisy data frame; or an estimate of a denoised version of the current noisy data frame. An initial noisy version of the data frame can be sampled from a (multivariate) Gaussian noise distribution.

[0161] As previously described, a data frame may include or represent data in an output space (e.g., an image space), or data in a (low-dimensional) latent space representation of the output space.

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

[0163] Several methods exist for generating dynamic images (i.e., videos) using diffusion models. As an example, the time axis dimension can be considered as an additional spatial dimension. As another example, the rolling diffusion technique described in the following literature can be used: Ruhe et al., arXiv:2402.09470v1.

[0164] The images or training images described herein can be images representing real-world environments. For example, training images can be captured from the real world using a camera or other image sensor. The objects represented in the training images or images generated by a diffusion model trained as described herein can include physically real-world objects.

[0165] In some implementations, data frames do not necessarily represent images, and as in some other examples, can represent audio data, weather or climate data, or fluid dynamics data. In some implementations, image frames can encode another type of data, such as a spectrogram.

[0166] As a concrete example, audio or other data can be represented as a spectrogram. Audio or other data can be processed by performing a time-frequency domain transform on the audio or other signal to generate a frequency domain representation of the audio or other signal within a frequency range, thus producing a spectrogram representing the audio or other data. Many suitable time-frequency domain transforms exist; as an example, the Short-Time Fourier Transform (STFT) is used. By applying the inverse of the time-frequency domain transform, i.e., the inverse time-frequency transform, such as the inverse STFT, audio or other data (signals) can be decoded from the generated spectrogram.

[0167] Generally, a spectrogram can be an image representing a time-frequency domain transformation, for example, a representation with one time axis (e.g., a horizontal axis) and another frequency axis (e.g., a vertical axis). The position of an image pixel along the time axis can represent a time position in audio or other data; the position along the frequency axis can represent the frequency at that time; the pixel value, such as brightness or color value, can represent the audio or other signal component at that time and frequency, such as the amplitude and / or phase of the signal. In some time-frequency domain transformations, the signal components are represented by complex numbers. Generally, a spectrogram can represent the spectrum of an audio or other signal as it changes over time.

[0168] In the case of generating moving images, this may include, as some examples, a continuation of a previous video sequence, an edited version of a video sequence, or a video sequence generated to represent text or audio conditional input.

[0169] As another example, this system can be implemented to generate images representing the predicted state or trajectory of a real-world physical system (such as a robot or vehicle) for use by a control algorithm to control that physical system. For instance, a camera can capture video sequences, which the system then uses to generate still or moving images that predict the future state or configuration of the physical system, optionally conditioned on one or more variables associated with the physical system. This can be used in model predictive control systems to control mechanical agents such as robots to perform specific tasks by processing the predicted images with the control system to generate control signals based on the generated images to control the mechanical agent to perform the task.

[0170] As some other examples, a data frame may include, for example, sensor data, such as data for EEG or ECG traces; or medical imaging data, such as CT or MRI data; or data representing the 3D configuration of a molecule, such as the coordinates of molecular atoms.

[0171] In some implementations, teacher and / or student diffusion models can be trained to perform image processing tasks. The task can be specified by conditional data, which may include the image to be processed.

[0172] In other words, a student diffusion model can be trained to perform image processing tasks that may involve obtaining a reference image to guide an image modification task, and using the reference image to modify (i.e. guide) the predicted frames generated by the student diffusion model (in each iteration of inference), for example, by adjusting the student diffusion model on the representation of the reference image. As some examples, this approach can be used to perform image modification tasks such as image colorization, super-resolution, image inpainting, image extrapolation, interpolation, brush-guided image generation, and denoising.

[0173] In this specification, the term "configured" is used in connection with computing systems and environments, as well as computer program components. A computing system or environment is considered "configured" to perform specific operations or actions when it possesses the necessary software, firmware, hardware, or combinations thereof that enable it to perform specific operations or actions during operation. For example, configuring a system may involve installing a software library with specific algorithms, updating firmware with new instructions for data manipulation, or adding hardware components to gain enhanced processing power. Similarly, one or more computer programs are "configured" to perform those intended operations or actions when they contain instructions that, when executed by a computing device or hardware, cause the device to perform specific operations or actions.

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

[0175] 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, computing devices or hardware may also include code that creates the execution environment for computer programs. This code can take the form of processor firmware, protocol stacks, database management systems, operating systems, or combinations of these elements. In the context of general-purpose computing graphics processing units (GPUs), embodiments can particularly benefit from leveraging the parallel processing capabilities of GPUs, where code specifically designed for GPU execution is employed, often referred to as kernels or shaders. Similarly, TPUs excel at running optimized tensor operations that are crucial to many machine learning algorithms. By utilizing 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 fields such as computer vision, natural language processing, and robotics.

[0176] Computer programs (also known as software, applications, modules, scripts, code, or simply programs) can be written in any programming language, including compiled or interpreted languages, as well as declarative or procedural languages. They can be deployed in various forms, such as standalone programs, modules, components, subroutines, or any other unit suitable for use in 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., a script within a markup language document), residing in a dedicated file, or distributed across multiple coordination files (e.g., files storing modules, subroutines, or code segments). Computer programs can execute on a single computer or across multiple computers, whether located at a single site or distributed across multiple sites and interconnected via data communication networks. The specific implementation of a computer program can 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.

[0177] 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 generally implemented as one or more software modules or components installed on one or more computers, which may be located at a single site or distributed across multiple locations. In some cases, one or more dedicated computers may be used for a particular engine, while in others, multiple engines may operate concurrently on the same one or more computers. Examples of engine functions in the context of AI and machine learning include data preprocessing 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 an engine will depend on the overall architecture and the distribution of computational tasks across various hardware components, including CPUs, GPUs, TPUs, and other specialized processors.

[0178] The processes and logic flows described in this specification can be executed by one or more programmable computers, which run one or more computer programs to perform functions by manipulating input data and generating output. Additionally, graphics processing units (GPUs) and tensor processing units (TPUs) can be used to implement aspects of these processes and logic flows concurrently, significantly accelerating execution. This approach offers significant advantages for computationally intensive tasks common in AI and machine learning applications, such as matrix multiplication, convolution, and other operations exhibiting high parallelism. By leveraging the parallel processing capabilities of GPUs and TPUs, significant speedup and efficiency gains can be achieved compared to relying solely on CPUs. Alternatively, or in combination with programmable computers and dedicated processors, these processes and logic flows can also be implemented using dedicated processing hardware, such as field-programmable gate arrays (FPGAs) or application-specific integrated circuits (ASICs), to achieve even higher performance or energy efficiency in specific use cases.

[0179] Computers capable of executing computer programs 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 typically work in conjunction with the CPU to handle specialized computations, while the CPU manages overall system operation and other tasks. Generally, the CPU receives instructions and data from read-only memory (ROM), random access memory (RAM), or both. The basic components of a computer include a CPU for executing instructions and one or more memory devices for storing instructions and data. The specific configuration of the processing unit and memory will depend on factors such as the complexity of the AI ​​model, the amount of data being processed, and the desired performance and latency requirements. Implementations can be carried out 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. Systems may include storage devices such as hard disks, SSDs, or flash memory for persistent data storage.

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

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

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

[0183] 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 components may include: backend components, such as backend servers or cloud-based infrastructure; optional middleware components, such as middleware servers or application programming interfaces (APIs), for facilitating communication and data exchange; and frontend components, such as client devices having a user interface through which users can interact with the implemented subject matter, a web browser, or an app. For example, the described functionality may be implemented only on the client device (e.g., for on-device machine learning) or deployed as a combination of frontend and backend components for more complex applications. Where present, these components may be interconnected using any form or medium of digital data communication, such as communication networks like local area networks (LANs) or wide area networks (WANs), including the Internet. The specific system architecture and component selection 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.

[0184] Computing systems can include geographically separated clients and servers that interact via communication networks. The specific type of network (such as a local area network (LAN), wide area network (WAN), or the Internet) will depend on the accessibility 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. Depending on the nature of the exchanged data and the system's security requirements, these protocols may include HTTP, TCP / IP, or other specialized protocols. In some embodiments, the server transmits data or instructions to a user device acting as a client, such as a computer, smartphone, or tablet. The client device can then process the received information, display the results to the user, and potentially send data or feedback back to the server for further processing or storage. This allows for dynamic interaction between the user and the system, enabling a wide range of applications and functionalities.

[0185] While this specification contains numerous details of specific implementations, these details should not be construed as limiting the scope of any invention or the scope that may be claimed, but rather as descriptions of features that may be characteristic of particular embodiments of a particular invention. Certain features described in this specification in the context of individual embodiments may also be implemented in combination in a single embodiment. Conversely, various features described in the context of a single embodiment may also be implemented individually or in any suitable sub-combination in multiple embodiments. Furthermore, although features may be described above as operating in certain combinations and even initially claimed in this way, in some cases one or more features from the claimed combination may be removed from the combination, and the claimed combination may involve sub-combinations or variations thereof.

[0186] Similarly, although operations are depicted in the accompanying drawings and described in a specific order in the claims, this should not be construed as requiring such operations to be performed in the specific order shown or in sequential order, or requiring all shown operations to be performed to achieve the desired result. In some contexts, multitasking and parallel processing may be advantageous. Furthermore, the separation of various system modules and components in the above embodiments should not be construed 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.

[0187] Specific embodiments of this subject matter have been described. Other embodiments are within the scope of the appended claims. For example, the actions recited in the claims can be performed in a different order and still achieve the desired result. As an example, the processes depicted in the drawings do not necessarily require a specific order or sequential sequence to achieve the desired result. In some cases, multitasking and parallel processing can be advantageous.

Claims

1. A computer-based method for training a diffusion model, comprising: A teacher diffusion model is obtained, which is configured to process noisy frames corresponding to the diffusion time between the initial time and the final time to generate teacher prediction frames for denoising the noisy frames. Initialize the student diffusion model, which is configured to process noisy frames corresponding to diffusion time to generate student prediction frames for denoising the noisy frames; as well as The student diffusion model is trained using the teacher diffusion model through the following steps: Obtain training frames from the training dataset; Determine the target time and training time, wherein the training time is the duration of a time step following the target time; The training frames are processed using noise scheduling, which defines the change in frame noise level over time, to determine noisy training frames for the training time. The student diffusion model is used to process the noisy training frames to generate student prediction frames, which define the denoised prediction frames for the final time. The noise scheduling is used to process the noisy training frame and the predicted denoised frame at the final time to determine the denoised training frame corresponding to the target time. as well as The learnable parameters of the student diffusion model are updated using a student model training objective, which depends on the denoised training frames and the predicted denoised frames at the final time.

2. The method of claim 1, wherein updating the learnable parameters of the student diffusion model comprises: The gradient of the student model training objective is backpropagated, and a stopping gradient is applied to the denoised training frame to prevent the gradient from being passed to the student diffusion model via the denoised training frame.

3. The method of claim 1 or 2, comprising: The teacher diffusion model is used to process the denoised training frames corresponding to the target time to generate the teacher-predicted denoised frames; as well as The learnable parameters of the student diffusion model are updated using the student model training objective, which depends on the teacher-predicted denoised frame and the predicted denoised frame at the final time.

4. The method according to any one of claims 1 to 3, further comprising: Maintain an auxiliary diffusion model, which is configured to process noisy frames corresponding to diffusion time to generate auxiliary prediction frames for denoising the noisy frames; The auxiliary diffusion model is used to process the denoised training frames to generate auxiliary prediction frames, which define the auxiliary prediction denoised frames at the final time. The teacher diffusion model is used to process the denoised training frames to generate teacher prediction frames, which define the teacher prediction denoised frames at the final time. The learnable parameters of the auxiliary diffusion model are updated using an auxiliary model training objective, the auxiliary model training objective depending on the difference between the predicted denoised frame at the final time and the auxiliary predicted denoised frame at the final time; as well as The student model training objective is used to update the learnable parameters of the student diffusion model, which depends on the similarity between: i) the predicted denoised frame at the final time, and ii) the difference between the auxiliary predicted denoised frame at the final time and the teacher predicted denoised frame at the final time.

5. The method of claim 4, wherein updating the learnable parameters of the student diffusion model comprises: The gradient of the student model training objective is backpropagated, and the stopping gradient is applied to the difference between the auxiliary prediction denoised frame at the final time and the teacher prediction denoised frame at the final time.

6. The method of claim 4 or 5, further comprising: The training objective of the auxiliary model is regularized using the difference between the auxiliary prediction denoised frame at the final time and the teacher prediction denoised frame at the final time.

7. The method according to any one of claims 1 to 3, comprising: Obtain the first training frame and the second training frame from the training dataset; The first training frame is processed using the noise scheduling to determine the noisy first training frame for the training time; The noise scheduling is used to process the second training frame to determine the noisy second training frame for the training time; The student diffusion model is used to process the noisy first training frame to generate a first student prediction frame, which defines the first denoised prediction frame at the final time. The student diffusion model is used to process the noisy second training frame to generate a second student prediction frame, which defines the second prediction denoised frame at the final time. The noise scheduling is used to process the noisy second training frame and the second predicted denoised frame of the final time to determine the second denoised training frame corresponding to the target time. The teacher diffusion model is used to process the second noise-reduced training frame to generate an auxiliary prediction frame; Determine the teacher model gradient of the auxiliary objective function, which depends on the difference between the second predicted denoised frame and the auxiliary predicted denoised frame at the final time; as well as The learnable parameters of the student diffusion model are updated using the student model training objective, which depends on the first predicted denoised frame at the final time and the teacher model gradient.

8. The method of claim 7, comprising: The training objective of the student model is determined, which depends on the product of the first predicted denoised frame at the final time and a vector proportional to the gradient of the teacher model.

9. The method of claim 7 or 8, further comprising: The noise scheduling is used to process the noisy first training frame and the first predicted denoised frame of the final time to determine the first denoised training frame corresponding to the target time. The learned parameters of the teacher diffusion model determine the gradient of the teacher model; The scaled teacher model gradient is determined by multiplying the teacher model gradient by a scaling matrix. Determine the Jacobian matrix of the learned parameters of the teacher diffusion model evaluated for the first denoising training frame; Determine the product of the Jacobian matrix of the learned parameters of the teacher diffusion model and the scaled gradient of the teacher model to obtain the training product; as well as The learnable parameters of the student diffusion model are updated using the student model training objective, which depends on the product of the training product and the first predicted denoised frame at the final time.

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

11. The method according to any one of claims 7 to 10, comprising: Obtain a batch of first training frames from the training dataset; Obtain a batch of second training frames from the training dataset; The training objective of the student model is determined from the first batch of training frames; as well as The gradient of the teacher model is determined from the second batch of training frames.

12. The method of any one of claims 1 to 11, wherein the student prediction frame includes the predicted denoised frame of the final time; and the method comprises: The student diffusion model is used to process the noisy training frames to generate the predicted denoised frames at the final time.

13. The method of any one of claims 1 to 12, wherein the teacher diffusion model and the student diffusion model each include a corresponding denoising neural network configured to process the diffusion time and the noisy frame corresponding to the diffusion time to generate a corresponding teacher prediction frame or student prediction frame.

14. A computer-implemented method for training a diffusion model, comprising: A teacher diffusion model is obtained, which is configured to process noisy frames corresponding to the diffusion time between the initial time and the final time to generate teacher prediction frames for denoising the noisy frames. Initialize the student diffusion model, which is configured to process noisy frames corresponding to diffusion time to generate student prediction frames for denoising the noisy frames; as well as The student diffusion model is trained using the teacher diffusion model by matching one or more statistical moments of the distribution of the predicted frames from the student diffusion model with one or more statistical moments of the distribution of the predicted frames from the teacher diffusion model.

15. The method of claim 14, further comprising training the student diffusion model using a moment matching objective, the moment matching objective depending on the distribution of prediction frames from the student diffusion model and the distribution of prediction frames from the teacher diffusion model, and the moment matching objective being minimized when training the teacher diffusion model on the prediction frames from the student diffusion model does not change the learned parameters of the teacher diffusion model.

16. The method of claim 14 or 15, comprising: Obtain training frames from the training dataset; and Matching the one or more statistical moments includes matching statistical moments defined by a denoising loss, which is determined by the difference between a first prediction frame generated by the student diffusion model and a second prediction frame generated by the teacher diffusion model.

17. The method of claim 16, comprising: The first predicted frame is obtained by processing the noisy version of the training frame using the student diffusion model. The second prediction frame is obtained from the noisy version of the first prediction frame and the training frame according to the noise scheduling, wherein the noise scheduling defines the change of frame noise level over time.

18. The method of claim 17, wherein obtaining the second predicted frame comprises: The noise-reduced training frame is determined from the noisy version of the first prediction frame and the training frame according to the noise scheduling. as well as The teacher diffusion model is used to process the denoised training frames to generate the second prediction frames.

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

20. The method of claim 19, comprising: Denoising training frames are determined from prediction frames from the teacher diffusion model, the prediction frames being generated by processing training frames and noisy versions of the training frames; The denoised training frames are processed using the auxiliary diffusion model to generate auxiliary prediction frames; The teacher diffusion model is used to process the noisy training frames to generate teacher prediction frames; as well as The auxiliary diffusion model is trained using an objective that depends on the difference between the auxiliary prediction frame and the teacher prediction frame.

21. A computer-based method for training a diffusion model, comprising: A teacher diffusion model is obtained, which is configured to process noisy frames corresponding to the diffusion time between the initial time and the final time to generate teacher prediction frames for denoising the noisy frames. Initialize the student diffusion model, which is configured to process noisy frames corresponding to diffusion time to generate student prediction frames for denoising the noisy frames; Obtain training frames from the training dataset; The forward diffusion process is applied to the training frames to generate a noisy version of the training frames; The noisy version of the training frame is processed using the student diffusion model to generate a prediction frame, the prediction frame including a denoised prediction frame, which is a denoised version of the noisy version of the training frame. A denoised version of the training frame is obtained from the denoised prediction frame and the denoised version of the training frame according to the noise scheduling, wherein the noise scheduling defines the change of frame noise level over time. The student diffusion model is trained to minimize an objective with respect to learnable parameters of the student diffusion model, the objective depending on the gradient determined from the teacher diffusion model and evaluated using the denoised version of the noisy version of the training frames.

22. The method of claim 21, wherein the gradient determined from the teacher diffusion model includes the gradient of a loss function that depends on the difference between i) the denoised prediction frame and ii) the teacher denoised prediction frame obtained by processing the noisy version of the training frame using the teacher diffusion model.

23. The method of any one of claims 1 to 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 one of claims 1 to 23, wherein the teacher diffusion model and the student diffusion model respectively comprise a teacher diffusion model neural network and a student diffusion model neural network, the teacher diffusion model neural network and the student diffusion model neural network being respectively configured to process noisy frames corresponding to diffusion time and the diffusion time to generate the teacher prediction frame and the student prediction frame respectively.

25. The method of any one of claims 1 to 24, wherein the teacher diffusion model and the student diffusion model are each further configured to process conditional data to generate a corresponding teacher prediction frame or student prediction frame, wherein the corresponding teacher prediction frame or student prediction frame characterizes the conditional data.

26. A computer-implemented method for generating data frames using a trained diffusion model. The trained diffusion model mentioned above includes a student diffusion model trained using the method described in any one of claims 1 to 25. The trained diffusion model is configured to process noisy frames corresponding to diffusion times between an initial time and a final time to generate predicted frames for denoising the noisy frames, the method comprising: Initialize the current noisy frame; Iteratively, for each of the discrete number of time steps between the initial time and the final time, until the last time step: The diffusion model is used to process the current noisy frame to generate a prediction frame, the prediction frame including a prediction of the noise in the current noisy frame or a prediction of a denoised version of the current noisy frame; The denoised version of the current noisy frame is determined from the predicted frame; The noise scheduling is used to determine the updated current noisy frame from the denoised version of the current noisy frame, the noise scheduling defining the change of frame noise level over time; The denoised version of the current noisy frame at the last time step includes the generated data frame.

27. A computer-implemented method for generating data frames using a trained diffusion model. The trained diffusion model is configured to process noisy frames corresponding to the diffusion time between the initial and final times to generate predicted frames for denoising the noisy frames. The trained diffusion model was trained using prediction frames sampled from the distribution of prediction frames generated by the teacher diffusion model, which was configured to process noisy frames corresponding to the diffusion time between the initial and final times to generate prediction frames for denoising the noisy frames. The training includes matching one or more statistical moments of the distribution of predicted frames from the trained diffusion model with one or more statistical moments of the distribution of predicted frames from the teacher diffusion model, the method comprising: Initialize the current noisy frame; Iteratively, for each of the discrete number of time steps between the initial time and the final time, until the last time step: The diffusion model is used to process the current noisy frame to generate a prediction frame, the prediction frame including a prediction of the noise in the current noisy frame or a prediction of a denoised version of the current noisy frame; The denoised version of the current noisy frame is determined from the predicted frame; The noise scheduling is used to determine the updated current noisy frame from the denoised version of the current noisy frame, the noise scheduling defining the change of frame noise level over time; The denoised version of the current noisy frame at the last time step includes the generated data frame.

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

29. The method of any one of claims 1 to 28, wherein the data frame comprises an image frame that defines pixel values ​​for a static image or a dynamic 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 the operation of the corresponding method as claimed in any one of claims 1 to 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 the operation of the corresponding method as claimed in any one of claims 1 to 29.

Images