Conditional diffusion neural network distillation training

Abstract

Methods, systems, and apparatus, including computer programs encoded on computer storage media, for performing distillation training of a conditional diffusion neural network.

Description

[0001] CONDITIONAL DIFFUSION NEURAL NETWORK DISTILLATION TRAINING

[0002] CROSS-REFERENCE TO RELATED APPLICATION

[0003] This application claims priority to U.S. Provisional Application No. 63 / 584,852, filed on September 22, 2023. The disclosure of the prior application is considered part of and is incorporated by reference in the disclosure of this application.

[0004] BACKGROUND

[0005] This specification relates to generating images using machine learning models.

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

[0007] SUMMARY

[0008] This specification describes a system implemented as one or more computer programs on one or more computers that trains a conditional diffusion neural network that can be used to generate an output image conditioned on a conditioning input.

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

[0010] Generative diffusion models provide strong priors for image generation, e.g., text-to- image generation and serve as a foundation for conditional generation tasks such as image editing, restoration, super-resolution, and compositing.

[0011] However, one major limitation of diffusion models is their slow sampling time. That is, during the reverse diffusion process, each sampling iteration requires sampling a denoising output from the diffusion model, which is a time and computationally-intensive task. Moreover, generating high quality outputs using existing techniques can require a large number of sampling iterations.

[0012] To address this challenge, this specification describes a conditional distillation method for training a conditional diffusion neural network that, after training, allows for conditional sampling with very few steps. In other words, after being trained using the described techniques, the conditional diffusion neural network can be used to generate high- quality images in a small number of sampling iterations, significantly improving the computational efficiency of the generation process relative to other approaches.

[0013] For example, the techniques can be used to directly distill a pre-trained diffusion neural network in a single stage through j oint-1 earning, largely simplifying previous two- stage procedures that include both distillation and then conditional fine-tuning. This makes the described “single stage7’ distillation approach that both (i) decreases the number of sampling iterations required and (ii) adapts the neural network to a new type of conditioning input significantly more computationally efficient than other approaches that separately “distill’’ and then “adapt’’ (or vice versa).

[0014] Furthermore, this specification describes a new parameter-efficient distillation mechanism that distills an additional conditional image generation task with only a small number of additional parameters combined with a shared frozen unconditional backbone.

[0015] The described approach outperforms existing distillation approaches across multiple tasks, including super-resolution, image editing, and conditioned generation, in terms of output image quality given the same sampling time. Additionally, the described techniques result in trained models that can match or exceed the performance of much slower fine-tuned conditional diffusion models.

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

[0017] BRIEF DESCRIPTION OF THE DRAWINGS

[0018] FIG. 1 is a diagram of an example training system.

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

[0020] FIG. 3 is a flow diagram of an example process for performing a training iteration.

[0021] FIG. 4 shows an example of the operation of the system.

[0022] FIG. 5 shows an example of the performance of the described techniques.

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

[0024] FIG. 1 is a diagram of an example training system 100. The training system 100 is an example of a system implemented as computer programs on one or more computers in one or more locations, in which the systems, components, and techniques described below can be implemented.

[0025] The system 100 trains a conditional diffusion neural network 110 that can be used to generate an output image 112 conditioned on a conditioning input 102.

[0026] That is, the system 100 can generate an output image using the conditional diffusion neural network 110 by performing a reverse diffusion process on a noisy representation of the output image using the conditional diffusion neural network 110 while the conditional diffusion neural network is conditioned on the conditioning input 102.

[0027] For example, the conditioning input 102 can be a lower-resolution image and the output image 112 can be a higher-resolution image, e.g., so that the system 100 performs image super-resolution.

[0028] As another example, the conditioning input 102 can be a first image and the output image 112 can be an edited version of the first image in which a specified edit has been performed on the first image. In some implementations, the specified edit is fixed. In some other implementations, the system can also receive an additional conditioning input 102, e.g., a text or speech input specifying the edit to be performed on the first image, and the output image 112 can be an edited version of the first image in which the specified edit specified in the additional conditioning input has been performed on the first image.

[0029] As another example, the conditioning input 102 can be a non-image input describing the contents of an image, e.g., a text or audio input, and the output image 112 can be an image that is described by the conditioning input.

[0030] In particular, the system 100 trains the conditional diffusion neural network 1 10 to be computationally-effi cient while maintaining output image quality.

[0031] More specifically, performing the reverse diffusion process requires iteratively generating samples from the diffusion neural network 110. By training the conditional diffusion neural network 1 10 as specified in this specification, the conditional diffusion neural network 110 can generate the output image in fewer sampling iterations after training, e.g., relative to a baseline diffusion neural network that is used to initialize at least some of the parameters of the conditional diffusion neural network. That is, in some implementations, the system 100 initializes 120 the training of the conditional diffusional neural network 110 starting from a pre-trained diffusion neural network 130.

[0032] In some cases, the pre-trained diffusion neural network 130 is also a conditional diffusion neural network and the system 100 trains the (new) conditional diffusion neural network 110 to improve the sampling time and, optionally, the output quality of the pretrained neural network 130.

[0033] In some cases, the pre-trained diffusion neural network 130 is an unconditional diffusion neural network or a diffusion neural network that is conditional but does not receive conditioning inputs of the same modality as the (new) conditional diffusion neural network 110. In these cases, the system 100 trains the (new) conditional diffusion neural network 110 to improve the sampling time of the pre-trained neural network 130 and to adapt the pretrained neural network 130 to be able to process conditioning inputs of the new modality.

[0034] When the pre-trained diffusion neural network 130 is conditional but does not receive conditioning inputs of the same modality as the (new) conditional diffusion neural network 110, the new conditional diffusion neural network 1 10 can process multiple conditioning inputs, including those of the new modality and one or more of the conditioning inputs that the pre-trained neural network 130 is already configured to process. For example, the new diffusion neural network 110 can process both text and image inputs while the pre-trained neural network 130 only processes text inputs.

[0035] The diffusion neural network 1 10 can be any appropriate diffusion neural network that is configured to receive an input that includes a current (noisy) representation of an image and a conditioning input and to generate a denoising output.

[0036] In some implementations, the diffusion neural network 110 performs a diffusion process in pixel space, so that the images ('‘representations’’) operated on and generated by the diffusion neural network 110 have values for each pixel that specily color values, e.g., RGB values or another color encoding scheme.

[0037] Examples of such diffusion neural networks include Imagen.

[0038] In some other implementations, the diffusion neural network 110 performs a diffusion process in latent space, e.g., in a latent space that is lower-dimensional than the pixel space. That is, the images (“representations”) operated on by the diffusion neural network 110 are latent images and the values for the pixels of the images are learned, latent values rather than color values.

[0039] Examples of such diffusion neural networks include Stable Diffusion. In these implementations, during training, the diffusion neural network 110 can be associated with an image encoder to encode training images into the latent space and, after training and to generate new target images, a decoder neural network that receives an input that includes a latent representation of an image and decodes the latent representation to reconstruct the image. For example, both the encoder and the decoder neural networks can be convolutional neural netw orks, can be self-attention neural networks, or can include both convolutional and self-attention layers.

[0040] The diffusion neural network 110 can have any appropriate architecture that allows the neural network to map a diffusion input that includes an input representation of a data item and to map the input representation to a denoising output that has the same dimensionality as the input representation.

[0041] For example, when the output data item is an audio signal or an image, the diffusion neural network 110 can be a convolutional neural network, e.g., a U-Net or other architecture that maps one input of a given dimensionality to an output of the same dimensionality7.

[0042] As another example, the diffusion neural network 110 can be a Transformer neural network that processes the diffusion input through a set of self-attention layers to generate the denoising output.

[0043] As yet another example, the diffusion neural network 110 can include both convolutional layers and self-attention layers.

[0044] The neural network 110 can be conditioned on the conditioning input 102 in any of a variety of ways.

[0045] As one example, the system 100 can use an encoder neural network to generate one or more embeddings that represent the conditioning input 102 and the diffusion neural netw ork 110 can include one or more cross-attention layers that each cross-attend into the one or more embeddings.

[0046] An embedding, as used in this specification, is an ordered collection of numerical values, e.g., a vector of floating point values or other types of values.

[0047] For example, when the conditioning input is text, the system can use a text encoder neural network, e.g.. a Transformer neural network, to generate a fixed or variable number of text embeddings that represent the conditioning input.

[0048] When the conditioning input is an image, the system can use an image encoder neural network, e.g., a convolutional neural network or a vision Transformer neural network, to generate a set of embeddings that represent the image. When the conditioning input is audio, the system can use, e.g., an audio encoder neural network, e.g., an audio encoder neural network that has been trained jointly with a decoder neural network as part of a neural audio codec, to generate one or more embeddings that encode the audio.

[0049] When the conditioning input is a scalar value, the system can use, e.g., an embedding matrix to map the scalar value or a one-hot representation of the scalar value to an embedding.

[0050] In some cases, the conditioning input 102 includes multiple different types of inputs, e.g., two or more of text, images, bound values, or context embeddings.

[0051] In some of these cases, the system 100 can generate one or more initial embeddings for each of the different types of inputs, i.e.. using an appropriate encoder neural network as described above, and then process the initial embeddings for all of the different types of inputs using a Transformer encoder neural network to update each of the initial embeddings to generate a set of final embeddings. The one or more cross-attention layers within the diffusion neural network 110 can then cross-attend into the set of final embeddings.

[0052] In others of these cases, different cross-attention layers within the diffusion neural network 110 can cross-attend into embeddings of different types of conditioning inputs.

[0053] In yet others of these cases, the system 100 can concatenate the initial embeddings of the different types of inputs along the sequence dimension and then the one or more crossattention layers can cross-attend into the concatenated set of final embeddings.

[0054] As another example, the diffusion neural network 1 10 can include one or more other ty pes of neural network layers that are conditioned on the one or more embeddings. Examples of such layers include Feature-wise Linear Modulation (FiLM) layers, layers with conditional gated activation functions, and so on.

[0055] As another example, as will be described in more detail below, the output(s) of the encoder(s) when encoding one or more of the conditioning inputs can be combined, e.g., through a weighted sum, with features of the representation of the output image, and the combined features can be processed by the remainder of the diffusion neural network 110.

[0056] The diffusion input at any given updating iteration can also include data defining a noise level for the iteration. Generally, each updating iteration has a corresponding time step t and the noise level for the iteration depends on the time step. For example, the noise level can be a decreasing function of the time step t. Examples of such functions include a linear function, a cosine function, and a sigmoid function. In these cases, data identifying the noise level, the time step, or both can be embedded using an appropriate neural network, e.g., a multi-layer perceptron (MLP) and used to condition the diffusion neural network 110 as described above for the conditioning input.

[0057] FIG. 2 is a flow diagram of an example process 200 for training a diffusion neural network. For convenience, the process 200 will be described as being performed by a system of one or more computers located in one or more locations. For example, a training system, e.g., the training system 100 depicted in FIG. 1. appropriately programmed in accordance with this specification, can perform the process 200.

[0058] The system initializes online values and target values of the network parameters (step 202).

[0059] For example, the system can initialize the online values to be equal to the target values.

[0060] In particular, as described above, in some implementations, the system initializes the training of the diffusion neural network from a pre-trained diffusion neural network. In these implementations, the system uses the pre-trained values of the network parameters of the pretrained diffusion neural network to initialize the online values and the target values.

[0061] The pre-trained diffusion neural network can have been pre-trained using any appropriate objective, e g., a score matching objective or another type of diffusion training objective, on any appropriate data set of images, and can be a conditional diffusion neural network an un-conditional diffusion neural network.

[0062] Generally, in these implementations, the system initializes online values of a first subset of the network parameters to be equal to pre-trained values of corresponding network parameters of the pre-trained diffusion neural network.

[0063] For example, if the diffusion neural network and the pre-trained diffusion neural network have the same architecture, e.g., because they both process the same type(s) of conditioning inputs, the system can initialize the online value of each network parameter of the diffusion neural network to be equal to the pre-trained value of the corresponding network parameter.

[0064] As another example, the first subset can be a proper subset of the network parameters of the diffusion neural network.

[0065] For example, the network parameters of the diffusion neural network can include the first subset of network parameters and a second subset of network parameters that do not have a corresponding network parameter in the pre-trained diffusion neural network.

[0066] This can occur for a variety of reasons. For example, the diffusion neural network can be configured to process conditioning inputs of a first modality, while the pre-trained diffusion neural network is not configured to receive conditioning inputs of the first modality. In this example, the second subset can include the network parameters of the encoder that processes the conditioning inputs of the first modality.

[0067] As a particular example of this, the network parameters can also include a scalar weight value that is assigned to features generated using the second subset of the network parameters in a combination of the features generated using the second subset of the network parameters and features generated using at least some of the first subset of network parameters. In this example, initializing the online values can include setting the scalar weight value to zero, i.e., so that, at the outset of training, the conditioning input of the new modality does not impact the output generated by the diffusion neural network.

[0068] For example, the pre-trained diffusion neural network can include an encoder that encodes a representation of the output image and one or more subsequent layers that process the encoded representation of the input image to generate the output representation of the output image. The new diffusion neural network can include the encoder, the one or more subsequent layers, and a new encoder for the first conditioning input. The scalar value can be used to weight the output of the new encoder with the output of the existing encoder when combining the two outputs for further processing by the subsequent layers.

[0069] As another example, the second subset of the network parameters can include low- rank adaptation (LoRA) parameters. For example, the system can use a low-rank approximation (LoRA) technique when performing the training.

[0070] In this case, for each of one or more weight matrices that are included in the pretrained diffusion neural network, the second set of network parameters include a low-rank factorization of an update weight matrix that can be used to update the weight matrix. The low-rank approximation technique can be performed on multiple different weight matrices to update corresponding different layers of the pre-trained diffusion neural netw ork.

[0071] The system can use the low-rank approximation to approximate an update to the update weight matrix during each training update of the diffusion neural network, e.g.. by optimizing a product of tw o smaller matrices in order to reduce the dimensionality of the calculation required to compute the change in weights required by the update. More specifically , performing a low-rank approximation refers to breaking up the update w eight matrix into a product of two smaller matrices that when multiplied together can recover the values of update weight matrix with high fidelity. In particular, the low-rank decomposition can represent Wo+ AW ~ IVO+ BA, where Wois a weight matrix in the first subset of network parameters, AW is the update weight matrix corresponding to Woand the product BA approximates AW.

[0072] In this case, the rank of a matrix refers to the number of linearly independent vectors, e.g., the sum of columns or rows within the matrix decomposition BA that do not contain correlative data. The rank determined specifies the dimensionality of the update needed by providing a constraint on the dimensions of the tw o smaller matrices. For example, in the case in which B is a matrix of dimension d X r and A has dimension r X k, where r must be the same to enable the matrix multiplication, the rank r can be a value much less than the minimum of d and k. e.g., r « min(d, k).

[0073] Thus, during training, the system learns the weights in matrices B and A instead of directly learning the w eights in A W.

[0074] Generally, the system can initialize the online values of the parameters in the second subset in any of a variety of ways. For example, the system can initialize the values randomly using a parameter initialization technique. As another example, when the component in which the second subset are included has the same architecture as one of the components in the pre-trained diffusion neural network, the system can copy the pre-trained values of the components from the pre-trained diffusion neural network.

[0075] In some of these examples, during the training of the conditional diffusion neural network, the system holds the values of the network parameters in the first subset fixed at the pre-trained values while updating online values of the netw ork parameters not in the first subset. Holding the values of the network parameters in the first subset fixed can improve the computational efficiency of the training while still effectively adapting the neural network.

[0076] The system then repeatedly performs operations to update the online values (step 204). That is, the system repeatedly performs iterations of a training process in order to update the online values using the target values. In some cases, as part of the updating, the system also updates the target values.

[0077] Performing an iteration of the training process wall be described in more detail below' with reference to FIG. 3.

[0078] After training, i.e., after repeatedly performing the operations to update the online values, the system can receive a new conditioning input (step 206) and generate a new image conditioned on the new conditioning input using the trained conditional diffusion neural network and in accordance with the online values of the network parameters (step 208). In particular, the system can perform a reverse diffusion process using the trained conditional diffusion neural network and in accordance with the online values of the network parameters to generate the new image. Because of the way that the system trains the conditional diffusion neural network, the system can perform fewer sampling iterations during the reverse diffusion process than would have been required to generate a high-quality output image using the pre-trained neural network.

[0079] To perform the reverse diffusion process to generate an output image from a given conditioning input, the system initializes a representation of the output image. For example, the system can sample each value in each representation from a noise distribution, e.g., a Gaussian distribution.

[0080] The system then updates the representation at each of a plurality of reverse diffusion steps (also referred to as '‘iterations” or ‘'updating iterations”) using the conditional diffusion neural network. Each reverse diffusion step is associated with a noise level for the iteration. Generally, each updating iteration has a corresponding time step t and the noise level for the iteration depends on the time step. For example, the noise level can be a decreasing function of the time step t. Examples of such functions include a linear function, a cosine function, and a sigmoid function. Thus, early iterations are associated with higher noise levels and later iterations are associated with lower noise levels, resulting in the diffusion neural network gradually “denoising” the representation to generate the final representation.

[0081] As part of the updating at any given step, the system generates a denoising output for the reverse diffusion step.

[0082] The system then updates the representation of the output image using the denoising output for the reverse diffusion step.

[0083] For example, the system can map the denoising output to an initial updated representation and then apply a diffusion sampler, e.g., the DDPM (Denoising Diffusion Probabilistic Model) sampler, the DDIM (Denoising Diffusion Implicit Model) sampler or another appropriate sampler, to the initial updated representation to generate an updated representation.

[0084] Optionally, after the last reverse diffusion iteration, the system can refrain from using the diffusion sampler and can instead use the initial updated representation as the updated representation.

[0085] To generate the denoising output, the system processes a first denoising input for the reverse diffusion step that includes the representation of the output image and the conditioning input using the denoising neural network to generate a first denoising output. In some cases, the first denoising output is the denoising output. In some other cases, the system also generates one or more additional denoising outputs and then combines the additional denoising output(s) with the first denoising output through classifier free guidance, i.e., by computing a weighted sum of the denoising outputs with the weight for each denoising output being determined by a guidance weight for the classifier free guidance.

[0086] After updating the representation at each of the plurality of reverse diffusion steps, the system generates the output image from the representation of the output image. For example, when the diffusion neural network is in the latent space, the system can process the representation using the decoder neural network to generate the output image. When the diffusion neural network is in the output space, the system can use the final representation as the output image.

[0087] FIG. 3 is a flow diagram of an example process 300 for performing an iteration of a training process. For convenience, the process 300 will be described as being performed by a system of one or more computers located in one or more locations. For example, a training system, e.g., the training system 100 depicted in FIG. 1. appropriately programmed in accordance with this specification, can perform the process 300.

[0088] The system receives a set of one or more training examples (step 302). Each training example includes a training image and a training conditioning input that characterizes the training image.

[0089] The system then performs steps 304-318 for each training example.

[0090] The system selects a first time step t from a plurality of time steps that are each associated with a respective noise level (step 304). For example, the time steps can be from a specified range of time steps that ranges from a fixed value, e g., 1, 2, or other appropriate value, to the index of the highest noise updating time step during generation of output images using the pre-trained diffusion neural network.

[0091] For example, the system can sample the first time step from the plurality of time steps, e.g., by sampling from a uniform distribution over the plurality of time steps.

[0092] In some implementations, the system independently samples the first time step for each of the training examples at a given iteration of the process 300. In some other implementations, when there is more than one training example, the system selects a single first time step and uses the same first time step for all of the training examples.

[0093] The system selects, based on the first time step, a second time step .s' that is associated with a lower noise level than the first time step (step 306). In some cases, the system can randomly sample the second time step from the time steps that are associated with a lower noise level than the first time step.

[0094] In some other cases, the second time step can have a predetermined offset within the plurality of time steps relative to the first time step. For example, s can be equal to M.

[0095] The system combines a representation x of the training image in the training example with sampled noise c in accordance with the noise level associated with the first time step to generate a latent representation ztfor the first time step (step 308).

[0096] For example, the system can determine a first weight based on the noise level associated with the first time step, determine a second weight based on the noise level associated with the first time step, and then determine a sum of the representation of the training image in the training example weighted by the first weight and the sampled noise weighted by the second weight to generate a latent representation for the first time step.

[0097] In other words, the system can compute zt— atx + <7t£, where atis the first weight and atis the second weight. As described above, each time step (or '‘updating iteration” or '‘reverse diffusion step”) has pre-computed weights associated with the time step that are based on the noise level for the time step.

[0098] The system processes a first denoising input that includes (i) the latent representation for the first time step, (ii) data identifying the first time step, and (iii) the training conditioning input in the training example using the conditional diffusion neural network and in accordance with the online values of the network parameters to generate a first denoising output (step 310).

[0099] The system can generally generate any appropriate type of denoising output for a given time step that defines a predicted noise component of the current representation as of the given time step, i.e., the noise that has been added to a ground truth representation of the output image to generate the current representation. For example, the denoising output can be (i) an estimate of the ground truth representation of the output image (given the current representation), (ii) an estimate of the noise that has been added to the ground truth representation of the target image to arrive at the current representation, (iii) a v- parameterization of the ground truth representation of the output image and the noise, or (iv) another appropriate ty pe of denoising output.

[0100] As one example, the denoising output can be a velocity model output ve(zt, c) that is a prediction of a combination, in accordance with a noise level associated with the time step. of a noise component of the latent representation and a representation of the image. For example. ve(zt, c) « ats + atx.

[0101] The system determines, from the first denoising output and in accordance with the respective noise level for the first time step, an estimate of the representation of the training image for the first time step (step 312).

[0102] For example, when the denoising output is a velocity model output, the system can determine a sum of the latent representation for the first time step weighted by the first weight and the first denoising output weighted by a negative of the second weight to generate the estimate of the representation of the training image for the first time step. In other words, the estimate of the representation can be: xt= atzt— <jtvgzt, c)

[0103] The system determines, from the first denoising output and in accordance with the respective noise level for the first time step, an estimate of the sampled noise for the first time step (step 314).

[0104] For example, when the denoising output is the velocity model output, the system can determine a sum of the first denoising output for the first time step weighted by the first weight and the latent representation for the first time step weighted by a negative of the second weight to generate the estimate of the sampled noise for the first time step. In other words, the estimate of the sampled noise can be:

[0105] The system combines the estimate of the representation of the training image for the first time step with the sampled noise in accordance with the noise level associated with the second time step to generate a latent representation for the second time step (step 316).

[0106] For example, the system can determine a third w eight based on the noise level associated with the second time step; determine a fourth weight based on the noise level associated with the second time step; and determine a sum of the estimate of the representation of the training image for the first time step weighted by the third w eight and the sampled noise weighted by the fourth weight to generate a latent representation for the second time step. In other words, the system can compute

[0107] Zs= O!sXt+ (Ts£

[0108] Alternatively, the system can determine a third weight based on the noise level associated with the second time step; determine a fourth weight based on the noise level associated with the second time step; and determine a sum of the estimate of the representation of the training image for the first time step weighted by the third weight and the noise estimate for the first time step weighted by the fourth weight to generate the latent representation for the second time step. In other words, the system can compute zs= asxt+ crset

[0109] The system processes a second denoising input that includes (i) the latent representation for the second time step, (ii) data identifying the second time step, and (iii) the conditioning input using the conditional diffusion neural network and in accordance with the target values 0“ of the network parameters to generate a second denoising output (step 318). For example, the system can generate the second denoising output as ve- (zs, c).

[0110] The system determines, from the second denoising output and in accordance with the respective noise level for the second time step, an estimate of the sampled noise for the second time step (step 320).

[0111] For example, the estimate can satisfy7:

[0112] The system then trains the conditional diffusion neural network to update the online values of the network parameters based on gradients of a loss function that includes a first term that measures, for each of the one or more training examples, a difference between (i) the estimate of the sampled noise for the second time step for the training example and (ii) the estimate of the sampled noise for the first time step for the training example (step 322). As an example, the system can apply a gradient-based optimization routine, such as stochastic gradient descent, to the loss function to determine updated values of the online values of the network parameters.

[0113] Optionally, the loss function can also include one or more additional terms.

[0114] As one example, the loss function further can include a second term that measures, for each of the one or more training examples, a difference between (i) the representation of the training example in the training example and (ii) the estimate of the representation of the training image in the training example for the first time step. For example, the difference can be an L2 distance between (i) the representation of the training example in the training example and (ii) the estimate of the representation of the training image in the training example for the first time step.

[0115] For example, the overall loss function can be the sum of the first term and the second term. As another example, the overall loss function can be a weighted sum of the first term and the second term. As an example, the loss function can be given by:

[0116] £(0) = E[A(s, t)de(et, es) + r(x, xt)] where A(s, t) is a weighting, which in some examples is equal to 1, d£is a distance function (e.g., the first loss), such as an L2 loss, 8~ are the target parameters, 9 are the online parameters, e is a noise estimate , xeis an estimate of a representation x of an image from the data distribution, c is a conditioning input, and r is the conditional guidance (e.g.. the second term).

[0117] In some implementations, after updating the online values, the system then updates the target values of the network parameters based on the online values. For example, the system can maintain the target values as an exponential moving average of the online values. In this example, the system can update the target values by updating the exponential moving average using the online values.

[0118] FIG. 4 shows an example 400 of the operation of the system during training.

[0119] In particular, FIG. 4 shows an example of training a “student” neural network that has the online values 0 of the network parameters using a “teacher” neural network that has the target values 0“ of the network parameters.

[0120] As shown in the example 400, the system can sample noise c from a noise distribution and an image x from the data distribution. The system then updates the student neural network using a loss function that is based on an error between xtand x and an error between G and es.

[0121] FIG. 5 shows an example 500 of the performance of the described techniques. In particular, the example 500 shows the performance of two variants of the described techniques (“PE-CoDi” and “CoDi”) relative to a set of baseline techniques on three tasks (super-resolution, inpainting, and text-guided depth-to-image) and for various numbers of sampling steps. As can be seen from the example 500, the described techniques achieve high quality on the tasks, even with a very' small number of sampling steps at inference.

[0122] This specification uses the term “configured” in connection with systems and computer program components. For a system of one or more computers to be configured to perform particular operations or actions means that the system has installed on it software, firmware, hardware, or a combination of them that in operation cause the system to perform the operations or actions. For one or more computer programs to be configured to perform particular operations or actions means that the one or more programs include instructions that, when executed by data processing apparatus, cause the apparatus to perform the operations or actions.

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

[0124] The term “data processing apparatus” refers to data processing hardware and encompasses all kinds of apparatus, devices, and machines for processing data, including by way of example a programmable processor, a computer, or multiple processors or computers. The apparatus can also be. or further include, special purpose logic circuitry, e.g., an FPGA (field programmable gate array) or an ASIC (application specific integrated circuit). The apparatus can optionally include, in addition to hardware, code that creates an execution environment for computer programs, e.g., code that constitutes processor firmware, a protocol stack, a database management system, an operating system, or a combination of one or more of them.

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

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

[0127] Similarly, in this specification the term “engine” is used broadly to refer to a software-based system, subsystem, or process that is programmed to perform one or more specific functions. Generally, an engine will be implemented as one or more software modules or components, installed on one or more computers in one or more locations. In some cases, one or more computers will be dedicated to a particular engine; in other cases, multiple engines can be installed and running on the same computer or computers.

[0128] The processes and logic flows described in this specification can be performed by one or more programmable computers executing one or more computer programs to perform functions by operating on input data and generating output. The processes and logic flows can also be performed by special purpose logic circuitry, e.g., an FPGA or an ASIC, or by a combination of special purpose logic circuitry and one or more programmed computers.

[0129] Computers suitable for the execution of a computer program can be based on general or special purpose microprocessors or both, or any other kind of central processing unit. Generally, a central processing unit will receive instructions and data from a read only memory' or a random access memory' or both. The essential elements of a computer are a central processing unit for performing or executing instructions and one or more memory’ devices for storing instructions and data. The central processing unit and the memory can be supplemented by, or incorporated in, special purpose logic circuitry'. Generally, a computer will also include, or be operatively coupled to receive data from or transfer data to, or both, one or more mass storage devices for storing data, e.g., magnetic, magneto optical disks, or optical disks. However, a computer need not have such devices. Moreover, a computer can be embedded in another device, e.g , a mobile telephone, a personal digital assistant (PDA), a mobile audio or video player, a game console, a Global Positioning System (GPS) receiver, or a portable storage device, e.g., a universal serial bus (USB) flash drive, to name just a few.

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

[0131] To provide for interaction with a user, embodiments of the subject matter described in this specification can be implemented on a computer having a display device, e.g., a CRT (cathode ray tube) or LCD (liquid crystal display) monitor, for displaying information to the user and a keyboard and a pointing device, e.g.. a mouse or a trackball, by which the user can provide input to the computer. Other kinds of devices can be used to provide for interaction with a user as well; for example, feedback provided to the user can be any form of sensory feedback, e.g., visual feedback, auditory feedback, or tactile feedback; and input from the user can be received in any’ form, including acoustic, speech, or tactile input. In addition, a computer can interact with a user by sending documents to and receiving documents from a device that is used by the user; for example, by sending web pages to a web browser on a user’s device in response to requests received from the web browser. Also, a computer can interact with a user by sending text messages or other forms of message to a personal device, e.g., a smartphone that is running a messaging application, and receiving responsive messages from the user in return.

[0132] Data processing apparatus for implementing machine learning models can also include, for example, special-purpose hardware accelerator units for processing common and compute-intensive parts of machine learning training or production, e.g., inference, workloads.

[0133] Machine learning models can be implemented and deployed using a machine learning framework, .e.g., a TensorFlow' framework or a Jax framework.

[0134] Embodiments of the subject matter described in this specification can be implemented in a computing system that includes a back end component, e.g., as a data server, or that includes a middleware component, e.g., an application server, or that includes a front end component, e.g., a client computer having a graphical user interface, a web browser, or an app through which a user can interact with an implementation of the subject matter described in this specification, or any combination of one or more such back end. middleware, or front end components. The components of the system can be interconnected by any form or medium of digital data communication, e.g., a communication netw ork. Examples of communication networks include a local area network (LAN) and a wide area netw ork (WAN), e.g.. the Internet. The computing system can include clients and servers. A client and server are generally remote from each other and typically interact through a communication network. The relationship of client and server arises by virtue of computer programs running on the respective computers and having a client-server relationship to each other. In some embodiments, a server transmits data, e.g., an HTML page, to a user device, e.g., for purposes of displaying data to and receiving user input from a user interacting with the device, which acts as a client. Data generated at the user device, e.g.. a result of the user interaction, can be received at the server from the device.

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

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

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

Claims

CLAIMS1. A method performed by one or more computers and for training a conditional diffusion neural network having a plurality of network parameters and configured to receive a diffusion input that comprises (i) a latent representation of an image, (ii) data identifying a current time step, and (iii) a conditioning input and to process the diffusion input to generate a denoising output that characterizes predicted noise in the latent representation, the method comprising: initializing online values and target values of the network parameters: and repeatedly performing operations to update the online values, the operations comprising: receiving a set of one or more training examples, each training example comprising a training image and a training conditioning input that characterizes the training image; for each training example: selecting a first time step from a plurality of time steps that are each associated with a respective noise level; selecting, based on the first time step, a second time step that is associated with a lower noise level than the first time step; combining a representation of the training image in the training example with sampled noise in accordance with the noise level associated with the first time step to generate a latent representation for the first time step; processing a first denoising input comprising (i) the latent representation for the first time step, (ii) data identifying the first time step, and (iii) the training conditioning input in the training example using the conditional diffusion neural network and in accordance with the online values of the netw ork parameters to generate a first denoising output; determining, from the first denoising output and in accordance with the respective noise level for the first time step, an estimate of the representation of the training image for the first time step; determining, from the first denoising output and in accordance with the respective noise level for the first time step, an estimate of the sampled noise for the first time step; combining the estimate of the representation of the training image for the first time step with the sampled noise in accordance with the noise level associated with thesecond time step to generate a latent representation for the second time step; processing a second denoising input comprising (i) the latent representation for the second time step, (ii) data identifying the second time step, and (iii) the conditioning input using the conditional diffusion neural network and in accordance with target values of the network parameters to generate a second denoising output; and determining, from the second denoising output and in accordance with the respective noise level for the second time step, an estimate of the sampled noise for the second time step; training the conditional diffusion neural network to update the online values of the network parameters based on gradients of a loss function that comprises a first term that measures, for each of the one or more training examples, a difference between (i) the estimate of the sampled noise for the second time step for the training example and (ii) the estimate of the sampled noise for the first time step for the training example.

2. The method of claim 1 , further comprising: after the training: receiving a new conditioning input; and generating a new image conditioning on the new conditioning input using the trained conditional diffusion neural network and in accordance with the online values of the network parameters.

3. The method of claim 1 or claim 2, wherein the loss function further comprises a second term that measures, for each of the one or more training examples, a difference between (i) the representation of the training example in the training example and (ii) the estimate of the representation of the training image in the training example for the first time step.

4. The method of claim 3, wherein the difference between (i) the representation of the training example in the training example and (ii) the estimate of the representation of the training image in the training example for the first time step is an L2 distance between (i) the representation of the training example in the training example and (ii) the estimate of the representation of the training image in the training example for the first time step.

5. The method of any one of claims 1-4, further comprising: after updating the online values, updating the target values of the network parameters based on the online values.

6. The method of claim 5, wherein the target values are maintained as an exponential moving average of the online values and wherein updating the target values comprises updating the exponential moving average using the online values.

7. The method of any preceding claim, wherein the same first and second time step are selected for each of the one or more training examples.

8. The method of any preceding claim, wherein the first time step is sampled from the plurality of time steps.

9. The method of any preceding claim, wherein the second time step has a predetermined offset within the plurality of time steps relative to the first time step.

10. The method of any preceding claim, wherein the denoising output is a velocity model output that is prediction of a combination, in accordance with a noise level associated with the time step, of a noise component of the latent representation and a representation of the image.

11. The method of any preceding claim, wherein initializing the online values and the target values of the network parameters comprises initializing the online values to be equal to the target values.

12. The method of any preceding claim, wherein initializing the online values and the target values of the network parameters comprises: initializing online values of a first subset of the network parameters to be equal to pretrained values of corresponding network parameters of a pre-trained diffusion neural network.

13. The method of claim 12, wherein the conditioning input is of a first modality, and wherein the pre-trained diffusion neural network is not configured to receive conditioninginputs of the first modality.

14. The method of claim 12 or 13, wherein the network parameters comprise the first subset of network parameters, a second subset of network parameters that do not have a corresponding network parameter in the pre-trained diffusion neural network, and a scalar weight value that is assigned to features generated using the second subset of the network parameters in a combination of the features generated using the second subset of the network parameters and features generated using at least some of the first subset of network parameters, and wherein initializing the online values and the target values of the network parameters comprises: setting the scalar weight value to zero.

15. The method of any one of claims 12-14, training the conditional diffusion neural netw ork to update the online values of the netw ork parameters comprises training the conditional diffusion neural network to update the online values of the network parameters comprises holding values of the network parameters in the first subset fixed at the pre-trained values while updating online values of the network parameters not in the first subset.

16. The method of any preceding claim, wherein combining a representation of the training image in the training example with sampled noise in accordance with the noise level associated with the first time step to generate a latent representation for the first time step comprises: determining a first weight based on the noise level associated with the first time step; determining a second weight based on the noise level associated with the first time step; and determining a sum of the representation of the training image in the training example weighted by the first weight and the sampled noise weighted by the second weight to generate a latent representation for the first time step.

17. The method of claim 1 , when dependent on claim 10, wherein determining, from the first denoising output and in accordance with the respective noise level for the first time step, an estimate of the representation of the training image for the first time step comprises: determining a sum of the latent representation for the first time step weighted by the first weight and the first denoising output weighted by a negative of the second weight togenerate the estimate of the representation of the training image for the first time step.

18. The method of claim 16 or claim 17, when dependent on claim 10, wherein determining, from the first denoising output and in accordance with the respective noise level for the first time step, an estimate of the sampled noise for the first time step comprises: determining a sum of the first denoising output for the first time step weighted by the first weight and the latent representation for the first time step weighted by a negative of the second weight to generate the estimate of the sampled noise for the first time step.

19. The method of any preceding claim, wherein combining the estimate of the representation of the training image for the first time step with the sampled noise in accordance with the noise level associated with the second time step to generate a latent representation for the second time step comprises: determining a third weight based on the noise level associated with the second time step; determining a fourth weight based on the noise level associated with the second time step; and determining a sum of the estimate of the representation of the training image for the first time step weighted by the third weight and the sampled noise weighted by the fourth weight to generate a latent representation for the second time step.

20. The method of claim 19, when dependent on claim 10, wherein determining, from the second denoising output and in accordance with the respective noise level for the second time step, an estimate of the sampled noise for the second time step comprises: determining a sum of the second denoising output for the first time step weighted by the third weight and the latent representation for the second time step weighted by a negative of the fourth weight to generate the estimate of the sampled noise for the second time step.

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

22. One or more computer-readable storage media storing instructions that when executed by one or more computers cause the one or more computers to perform the respective operations of the method of any one of claims 1-20.

Images