Training diffusion neural networks using reward value differences

Abstract

Methods, systems, and apparatus, including computer programs encoded on a computer storage medium, performing reward training of a diffusion neural network.

Description

[0001]Attorney Docket No.56113-0561WO1 TRAINING DIFFUSION NEURAL NETWORKS USING REWARD VALUE DIFFERENCES CROSS-REFERENCE TO RELATED APPLICATIONS This application claims priority to U.S. Provisional Application No.63 / 615,775, filed on December 28, 2023. The disclosure of the prior application is considered part of and is incorporated by reference in the disclosure of this application. BACKGROUND This specification relates to processing data using machine learning models. As one example, neural networks are machine learning models that employ one or more layers of nonlinear units to predict an output for a received input. Some neural networks include one or more hidden layers in addition to an output layer. The output of each hidden layer is used as input to another layer in the network, e.g., the next hidden layer or the output layer. Each layer of the network generates an output from a received input in accordance with current values of a respective set of weights. SUMMARY This specification describes a system implemented as computer programs on one or more computers in one or more locations that trains a diffusion neural network that can be used to generate an output data item conditioned on a conditioning input. More specifically, this disclosure describes techniques for “fine-tuning.” i.e., further training, the diffusion neural network using a reward function. The subject matter described in this specification can be implemented in particular embodiments so as to realize one or more of the following advantages. This specification describes techniques for fine-tuning a diffusion neural network using a reward function. While reward finetuning has emerged as a promising approach for aligning various models with user preference, there has only been only limited success in finetuning diffusion models to maximize reward functions. Existing approaches are hard to scale to large-scale training datasets, i.e., datasets that include a large and diverse number of conditioning inputs. By contrast, by predicting the reward difference between pairs of denoising trajectories, the described techniques can be used to achieve stable training on these large-scale data sets, e.g., on a large-scale prompt dataset with over one hundred Attorney Docket No.56113-0561WO1 thousand prompts. As a result, the described techniques achieve remarkably improved generation quality relative to other reinforcement learning (RL) based reward optimization approaches and can also generalize well to unseen prompts after training. The details of one or more embodiments of the subject matter of this specification are set forth in the accompanying drawings and description below. According to a first aspect there is provided a method performed by one or more computers for training a diffusion neural network having one or more network parameter values. The method includes, at each of one or more training iterations, obtaining one or more of conditioning inputs for the training iteration and training the diffusion neural network to update the current values of the network parameters. The method further includes generating, for each conditioning input, one or more denoising trajectories conditioned on the conditioning input and using the diffusion neural network and in accordance with old values of the network parameters. Each denoising trajectory includes a sequence of data items starting from a noisy data item and terminating in a denoised data item. Then, for each denoising trajectory, a reward function is applied to the denoised data items to generate a reward score for the denoising trajectory. Training the diffusion neural network to update current values of the network parameters, includes, at each of one or more training steps, for each of one or more of the conditioning inputs and for each of a set of pairs of denoising trajectories for the conditioning input generating a respective likelihood score for each denoising trajectory in the pair given the conditioning input using the diffusion neural network and in accordance with the current values of the network parameters. Then, the diffusion neural network is trained to update the current values of the network parameters on a loss function that includes a first loss term that measures, for each of the one or more conditioning inputs and for each of the set of pairs of denoising trajectories for the conditioning input, (i) a difference between the respective reward scores for the denoising trajectories in the pair and (ii) a difference between the respective likelihood scores for the denoising trajectories in the pair. In some implementations, the data items are images. In some cases, conditioning inputs comprise text sequences. Further in some cases, at each of the plurality of training iterations and after training the diffusion neural network to update the current values of the network parameters, the old values of the network parameters are updated based on the current values of the network parameters. Attorney Docket No.56113-0561WO1 In some implementations, updating the old values of the network parameters are based on the current values of the network parameters includes setting the reference values equal to the old values of the network parameters. In some implementations, prior to performing the one or more training iterations, the diffusion neural network has been pre-trained to determine reference values of the network parameters. Further in some cases, prior to performing the one or more training iterations, the current values of the network parameters and the old values of the network parameters are set to the reference values of the network parameters. In some implementations, generating a respective likelihood score for each denoising trajectory in the pair given the conditioning input using the diffusion neural network and in accordance with the current values of the network parameters includes, for each denoising trajectory in the pair generating a first likelihood score for the denoising trajectory given the conditioning input using the diffusion neural network and in accordance with the current values of the network parameters, generating a second likelihood score for the denoising trajectory given the conditioning input using the diffusion neural network and in accordance with the reference values of the network parameters and then determining the likelihood score for the denoising trajectory given the conditioning input from the first and second likelihood scores for the denoising trajectory. Further in some implementations, determining the likelihood score for the denoising trajectory given the conditioning input from the first and second likelihood scores for the denoising trajectory includes determining the likelihood score based on a ratio between the first and second likelihood scores. In some cases, the first loss term measures a mean squared error between (i) the difference between the respective likelihood scores for the denoising trajectories in the pair and (ii) the difference between the respective reward scores for the denoising trajectories in the pair divided by a regularization strength hyperparameter. In some cases, the loss function includes a second loss term that measures, for each of the one or more conditioning inputs and for each of the set of pairs of denoising trajectories for the conditioning input, (i) a difference between respective clipped likelihood scores for the denoising trajectories in the pair and (ii) the difference between the respective reward scores for the denoising trajectories in the pair. Further in some cases, the loss function is a maximum of the second loss term and the first loss term. Attorney Docket No.56113-0561WO1 In some implementations, the diffusion neural network has been pre-trained on a score matching objective. In some implementations, the reward function is configured to process a reward input that comprises the conditioning input. According to a second aspect there is provided the methods of the first aspect performed by a system that includes 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 respective method. According to a third aspect there is provided the method of the first aspect performed by 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 respective method. Other features aspects, and advantages of the subject matter will become apparent from the description, drawings, and the claims. BRIEF DESCRIPTION OF THE DRAWINGS FIG.1 shows an example training system. FIG.2 is a diagram that further illustrates the example training system including a diffusion neural network, reward model and loss function. FIG.3 illustrates reward score improvement using the example training system. FIG.4 illustrates the progression of an example output of the training system after one or more training iterations. FIG.5 is a flow diagram of an example process for training the diffusion neural network. FIG.6 is a flow diagram of sub-steps of one of the steps of the process of FIG.4. FIG.7 is a flow diagram of sub-steps of one of the steps of the process of FIG.5. DETAILED DESCRIPTION FIG.1 shows an example training system 100 that includes a diffusion neural network 110. Attorney Docket No.56113-0561WO1 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. The system 100 trains a diffusion neural network 110 that can be used to generate an output data item 112, e.g., an output image, conditioned on a conditioning input 102. That is, the system 100 can generate an output data item 112 using the diffusion neural network 110 by performing a reverse diffusion process on a noisy representation of the output data item 112 using the diffusion neural network 110 while the diffusion neural network 110 is conditioned on the conditioning input 102. Generally, the conditioning input 102 characterizes one or more desired properties for the data item, i.e., characterizes one or more properties that the output data item 112 generated by the system should have. The system 100 can be configured to generate any of a variety of output data items 112 conditioned on any of a variety of conditioning inputs 102. For example, the system 100 can be configured to generate audio data, e.g., a waveform of audio or a spectrogram, e.g., a mel-spectrogram or a spectrogram where the frequencies are in a different scale, of the audio. In this example, the conditioning input 102 can be text or features of text that the audio should represent, i.e., so that the system 100 serves as a text-to-speech machine learning model that converts text or features of the text to audio data for an utterance of the text being spoken. As another example, the conditioning input 102 can identify a desired speaker for the audio, i.e., so that the system 100 generates audio data that represents speech by the desired speaker. As another example, the conditioning input 102 can characterize properties of a song or other piece of music, e.g., lyrics, genre, and so on, so that the system 100 generates a piece of music that has the properties characterized by the conditioning input 102. As another example, the conditioning input 102 can specify a classification for the audio data into a class from a set of possible classes, so that the system 100 generates audio data that belongs to the class. For example, the classes can represent types of musical instruments or other audio emitting devices, i.e., so that the system 100 generates audio that is emitted by the corresponding class, types of animals, i.e., so that the system 100 generates audio that represent noises generated by the corresponding animal, and so on. Attorney Docket No.56113-0561WO1 As another particular example, the output data item 112 can be an image, such that the system 100 can perform conditional image generation by generating the intensity values of the pixels of the image. In this particular example, the conditioning input 102 can include a sequence of text and the output data item 112 can be an image that describes the text, i.e., the conditioning input 102 can be a caption for the output image. As yet another particular example, the conditioning input 102 can be an object detection input that specifies one or more bounding boxes and, optionally, a respective type of object that should be depicted in each bounding box. As yet another particular example, the conditioning input 102 can specify an object class from a plurality of object classes to which an object depicted in the output image should belong. As yet another particular example, the conditioning input 102 can specify a context image and the output data item 112 can be an image that modifies the context image. For example, the output data item 112 can be a higher-resolution version of the context image, an image that depicts the same object as the context image but in a different context, or a de-noised version of the context image. As yet another particular example, the conditioning input 102 can specify an image including a target entity for detection, e.g., a tumor, and the output data item 112 can comprise the image without the target entity, e.g., to facilitate detection of the target entity by comparing the images. As yet another particular example, the conditioning input 102 can be a segmentation that assigns each of a plurality of pixels of the output image 112 to a category from a set of categories, e.g., that assigns to each pixel a respective one of the category. More generally, the task can be any task that outputs continuous data conditioned on a conditioning input 102. For example, the output data item 112 can be an output of a different sensor, e.g., a lidar point cloud, a radar point cloud, an electrocardiogram reading, and so on, and the conditioning input 102 can represent the type of data that should be measured by the sensor, where a discrete output is desired, this can be obtained, e.g., by thresholding. In any of the above examples, the output data item 112 generated using the diffusion neural network can either be an output data item in the output space, i.e., so that the values in the output data item 112 are the values of a data item of the appropriate type, e.g., values of image pixels, amplitude values of an audio signal, and so on, or an output data item in a Attorney Docket No.56113-0561WO1 latent space, i.e., so that the values in the output data item 112 are values in a latent representation of an output data item in the output space. When the output data item is generated in a latent space, the system 100 can generate a final output data item 112 in output space by processing the output data item in the latent space using a decoder neural network, e.g., one that has been pre-trained in an auto-encoder framework. During training, the system 100 can use an encoder neural network, e.g., one that has been pre-trained jointly with the decoder in the auto-encoder framework, to encode target data items in the output space to generate target outputs for the diffusion neural network 110 in the latent space. The diffusion neural network 110 can be any appropriate diffusion neural network that is configured to receive an input that includes a current (noisy) representation of a data item and a conditioning input 102 and to generate a denoising output. In some implementations, the diffusion neural network 110 performs a diffusion process in output space, e.g., pixel space when the data items are images. In this example, when the data items are images, the data items (“representations”) operated on and generated by the diffusion neural network have values for each pixel that specify color values, e.g., RGB values or another color encoding scheme. Examples of such diffusion neural networks include Imagen, details of which can be found in C. Saharia et al, “Photorealistic Text-to-Image Diffusion Models with Deep Language Understanding”, arXiv preprint arXiv:2205.11487 [cs.CV]. 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 output space. That is, the data items (“representations”) operated on by the diffusion neural network are latent representations and the values in the representations are learned, latent values, e.g., rather than color values when the data items are images. Examples of such diffusion neural networks include Stable Diffusion, details of which can found in R. Rombach, et al., “High-Resolution Image Synthesis with Latent Diffusion Models”, arXiv preprint arXiv:2112.10752 [cs.CV]. In these implementations, after training, the diffusion neural network 110 can be associated with a decoder neural network that receives an input that includes a latent representation of the final data item 112 and decodes the latent representation to reconstruct the data item. In some implementations, when the output data item 112 is an audio signal or an image, the diffusion neural network 110 can be a convolutional neural network, e.g., a U- Attorney Docket No.56113-0561WO1 Net or other architecture that maps one input of a given dimensionality to an output of the same dimensionality. 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. As yet another example, the diffusion neural network 110 can include both convolutional layers and self-attention layers. The neural network 110 can be conditioned on the conditioning input 102 in any of a variety of ways. 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 network 110 can include one or more cross-attention layers that each cross-attend into the one or more embeddings. 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. For example, when the conditioning input 102 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 102. When the conditioning input 102 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 102 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. When the conditioning input 102 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. 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. 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 Attorney Docket No.56113-0561WO1 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. 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. 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 cross-attention layers can cross-attend into the concatenated set of final embeddings. As another example, the diffusion neural network 110 can include one or more other types 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. As another example, 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. 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 102. To perform the reverse diffusion process to generate an output image from a given conditioning input 102, 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. 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 Attorney Docket No.56113-0561WO1 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. As part of the updating at any given step, the system generates a denoising output for the reverse diffusion step. The system then updates the representation of the output image using the denoising output for the reverse diffusion step. 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. Details of the DDPM sampler and the DDIM sampler can be found in, for example, Ho et al. “Denoising Diffusion Probabilistic Models” arXiv: 2006.11239 and Song et al. “Denoising Diffusion Implicit Models” arXiv: 2010.02502. 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. That is, the system can keep the initial updated representation from the last reverse diffusion step as the final updated representation and not apply a diffusion sampler to the representation. 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 102 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. 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 Attorney Docket No.56113-0561WO1 diffusion neural network is in the output space, the system can use the final representation as the output image. In particular, the system 100 fine-tunes, e.g., iteratively updates, the diffusion neural network 110 across multiple training iterations to achieve more stable training performance. More specifically, the system 100 fine-tunes the diffusion neural network 110, initialized 120 with a pre-trained diffusion neural network 130, using a reward function. As part of this training, the system 100 uses the diffusion neural network 110 to generate denoising trajectories for a given conditioning input 102 and updates the values of one or more network parameters based on the output of a reward function that compares two of the denoising trajectories. The training process will be described further in more detail below in reference to FIG.2. FIG.2 is a diagram that further illustrates an example training system including a diffusion neural network and a reward function. The training system 200 trains the diffusion neural network 110 to update current values of the network parameters of the diffusion neural network 110. In particular, the system 200 performs a sequence of training steps (“gradient descent steps) and, at each training step 240, updates the current values of the network parameters. In some implementations, the diffusion neural network 110 has been pre-trained to determine reference values of the network parameters. In this particular example, the diffusion neural network then sets the current values of the network parameters and “old” values of the network parameters, e.g., parameters that represent the state of the network 110 after a previous training step, to the reference values before training begins. More specifically, the diffusion neural network 110’s old values of network parameters represent a snapshot of the model from a previous training step. Both the current and old values of the network parameters are initialized with reference values from a pre-trained diffusion neural network, e.g., the pre-trained diffusion neural network 130 of FIG.1. When the diffusion neural network 110 has not been pre-trained, the system 200 can initialize the current, old, and reference values in any appropriate way, e.g., randomly or by setting them to the same pre-determined values. At each training iteration, the system 200 obtains one or more conditioning inputs 202. For example, the system 200 can randomly sample the conditioning inputs 202 from a larger set of conditioning inputs. Attorney Docket No.56113-0561WO1 For each conditioning input 201 from the one or more conditioning inputs 202, the system 200 generates one or more denoising trajectories 212 conditioned on the conditioning input 201 and using the diffusion neural network 110 in accordance with the old values of the network parameters. In the particular example depicted in FIG.2, two denoising trajectories 214 and 216 are generated from a conditioning input 201, e.g., a description of a bunny. Each denoising trajectory 214 and 216 includes a sequence of data items starting from a noisy data item and terminating in a denoised data item. To generate a denoising trajectory, the diffusion neural network 110 performs the reverse diffusion process described in reference to FIG.1, e.g., gradually denoising the data item in small intermediate steps. As described with reference to FIG.1, in some implementations, a Gaussian distribution can be used to model the noise of each of the data items in the denoising trajectory and then sampled to generate the next data item: ^^^^^்^ ൌ ℕ^0, ^^^^^^ ^^ ^^^ ^^^Starting at ^^்sample the data item with a Gaussian distribution ℕ with a mean at 0 and an identity covariance matrix I. Then, at each step t, the model takes the current noisy data item ^^௧ି^and predicts the mean^^ ^^^௧, ^^^ of the next data item, i.e., the initial updated which is thenfrom a Gaussian distribution ℕ^^^௧ି^; ^^^^^^௧, ^^^,^^ଶ௧ ^^^ centered around this predicted meanwith the given variance. the sampler to the initial updatedrepresentation ^^^^^^௧, ^^^ as described above. The full denoising trajectory can be generatedby first sampling the noisy image ^^^^^௧^ and then sampling ^^^^^^௧ି^|^^௧ , ^^^ for t = T,….1until the data item has been denoised. For each denoising trajectory of the one or more denoising trajectories 212 for the conditioning input 201, the system 200 applies the reward function 220 to the denoised data items to generate a reward score for the denoising trajectory. For example, the system 200 applies a reward function 220 to each of the two denoising trajectories 214 and 216 to generate a reward score for each. The reward function 220 can be any appropriate differentiable reward function that maps an input that includes (i) a data item or (ii) a latent representation of a data item to a Attorney Docket No.56113-0561WO1 reward score. Optionally, the reward function input can also include the conditioning input 201 or a representation of the conditioning input 201. For example, the reward function 220 can include one or more trained reward machine learning models, e.g., neural networks. As one example, the reward function 220 can include a machine learning model that maps at least a portion of the reward input to a score that represents the generation quality of the output data item. As another example, the reward function 220 can include a machine learning model that maps at least a portion of the reward input to a score that represents a predicted quality of the output data item. When there are multiple reward models in the reward function 220, the final reward score can be a sum or a weighted sum of the reward scores generated by the models. When the diffusion neural network 110 generates outputs in the latent space, as part of applying the reward function 220 to the output data item, the system can process the output data item using the decoder neural network to generate an output in the output space and then apply the reward function 220 to the output in the output space. For each training step, the system 200 analyzes a pair of denoising trajectories from a set of pairs of denoising trajectories for each conditioning input 201. In this particular example, a pair of denoising trajectories can be the denoising trajectories 212, i.e., the denoising trajectory 214 and denoising trajectory 216 for conditioning input 201 that described a bunny. The system then calculates a reward difference 222 between the reward scores generated from each denoising trajectory in the pair of denoising trajectories 212 for each of a set of pairs of denoising trajectories. For example, the system 200 calculates a reward difference 222 between the reward score of the pair of denoising trajectories 214 and 216. For each of the one or more conditioning inputs 202, i.e., conditioning input 201, and for each of a set of pairs of denoising trajectories for the conditioning input 201, the system 200 also generates a respective likelihood score 218 for each denoising trajectory in the pair 212 given the conditioning input 201 using the diffusion neural network 110 and in accordance with the current values of the network parameters. To generate a respective likelihood score 218 for each denoising trajectory in the pair 212, the system generates a first likelihood score for the denoising trajectory given the conditioning input 201 and in accordance with the current values of the network parameters Attorney Docket No.56113-0561WO1 and, optionally, a second likelihood score for the denoising trajectory given the conditioning input 201 in accordance with the reference values of the network parameters. The system can then determine the final likelihood score for the denoising trajectory from the first and, when computed, the second likelihood scores. The likelihood score of a denoising trajectory quantifies the probability of observing the entire denoising trajectory from the initial noisy image to the final data item given the diffusion model 110’s parameters and the conditioning input 201. By comparing the first likelihood score using the current values and the second likelihood score using the reference values, the diffusion neural network 110 can analyze the updates to the parameters to see if the model is being improved by the updates, e.g., the likelihood is being increased and the model is being optimized. In some implementations, the diffusion neural network 110 can compute a log- likelihood score for the denoising trajectory ^^^^^^:்|^^^ given the conditioning input ^^ and the respective values of the network parameters. The following example depicts the log- likelihood given the current values of the parameters ^: ^^^^^^^^^^^^:்|^^^ ൌ ^^^^^^^^^^^்^ ^ ^் ௧ୀ^൫^^^^^^^^^^^^௧ି^|^^௧ , ^^^൯ൌ noisy image ^^^^^்^ and ^் ௧ୀ^൫^^^^^^^^^^^^௧ି^|^^௧, ^^^൯represents the log-probabilities of each subsequent data item trajectory. The sum is taken over all of the data items in thedenoising trajectory. To calculate the log likelihood for each data item, the rightmost equation is used where ^^௧ି^represents the denoised data item of the previous diffusion step and ^^^represents the predicted mean of the denoised data item at the current step, given the noisy input ^^^௧ , ^^^ where c is the conditioning input. The squared error is computedbetween those two values and normalized by the variance ^^௧ଶto account for any uncertainly. The system can also determine a second likelihood score for the denoising trajectory given the reference values of the parameters, denoted as ref: ் మ ௫ ି ^ ^^^^^^^^^^^^ ^^^^^^^^^^^ ^ ^ ^ ^^^^^^^^^ ^^^ ^^^^் ^ห^షభఓ^^^௫^,^ ห^ ^ Attorney Docket No.56113-0561WO1 The system can then determine the likelihood score for the denoising trajectory from the first likelihood score in accordance with the current values of the network parameters and the second likelihood score in accordance with the reference values of the network parameters, e.g., based on a ratio between the first and second likelihood scores. In some implementations, the likelihood score can determine by the below equation: ^̂^^^^̅^^, ^^^= ^^^^^^గ^^௫̅ೌ|^^ గ^^^^௫̅ೌ|^^ , The likelihood score, also be calculated by taking thelogarithm of the ratio of the first accordance with the current values ^^^^^̅^^|^^^ and the second likelihood score generated in accordance with the reference values ^^^^^^^̅^^|^^^ for a denoising trajectory ^̅^^. The system 200 then trains the diffusion neural network 110 to update the current values of the network parameters 232 on a loss function 230. The loss function 230 includes a first loss term that measures, for each of the one or more conditioning inputs 202, and for each of the set of pairs of denoising trajectories for the conditioning input 201, (i) a difference between the respective reward scores 222 for the denoising trajectories in the pair 212 and (ii) a difference between the respective likelihood scores 218 for the denoising trajectories in the pair 212. That is, the system 200 trains the diffusion neural network 110 on a loss function 230 that includes a first term that operates on the reward difference 222 between the pair of denoising trajectories 212 and the difference between the likelihood scores 218 for the pair of denoising trajectories 212. For example, the loss function 230 can be a sum of or average of respective losses for each of the pairs for each of the conditioning inputs 202. In some implementations, the first loss term in the loss function 230 can be a mean squared error (MSE) between (i) the difference between the respective likelihood scores 218 for the denoising trajectories in the pair 212 and (ii) the difference between the respective reward scores 222 for the denoising trajectories in the pair 212 divided by a regularization strength hyperparameter, as seen below: ℒ^^^^ ൌ ^^௫̅ೌ,௫್̅,^||^^^̂^^^^̅^^, ^̅^^, ^^^ െ ^^^̂^^൫^^^^, ^^^^, ^^൯ / ^^ ||ଶ,where likelihoodscores for the pair of denoising trajectories ^̅^^and ^̅^^, and ^^^̂^^ ^^൫^^^, ^^^ , ^^൯ represents thedifference between the respective reward scores for the pair of denoising trajectories. The regularization strength hyperparameter ^^ discourages the current values of the diffusion Attorney Docket No.56113-0561WO1 model from deviating too far from reference values to preserve overall generation capability and keep the generated images close to the distribution where the reward model is accurate. In some implementations, the loss function 230 includes a second loss term. For example, the second loss term can measure, for each of the one or more conditioning inputs, e.g., conditioning input 201, and for each of the set of pairs of denoising trajectories 212 for the conditioning input 201, (i) a difference between respective clipped likelihood scores 218 for the denoising trajectories in the pair 212 and (ii) the difference between the respective reward scores 222 for the denoising trajectories in the pair 212. The second loss term can be included to ensure that the model is updated in smaller training steps that do not move the current values too far from the old values of the network parameters, e.g., the network parameters from the last training step in the training iteration. The likelihood scores are “clipped” to be within an interval using the old values of the network parameters, i.e., the likelihood score for the current value can only be better than the likelihood score for the old values by a certain amount. The system can do this by clipping a likelihood score ratio to be within a smallinterval ^െ^^ᇱ, ^^ᇱ^, i.e., limit how much the likelihood ratio changes per training step.Generally, the system restricts the likelihood score to be within the interval: ^̂^^^^^^ ^^̅^, ^^^ ൌ ^^^^^^^^൫^̂^^^^̅^, ^^^, ^̂^^^^^^^̅^, ^^^ െ ^^ᇱ, ^̂^^^^^^^̅^, ^^^ ^ ^^ᇱ൯. The second loss term of the loss function would then look like the following: ||^^^̂^^^^^^ ^^̅^^, ^̅^^, ^^^ െ ^^^̂^^൫^^^^ , ^^^^, ^^൯ / ^^ ||ଶ, where the clipped likelihood scores arerepresented by ^^^̂^^^^^^ ^^̅^^, ^̅^^, ^^^. By “clipping” the likelihood scores, the system is able tomaintain training stability by updating the model in small steps. In some implementations, the loss function 230 can be a maximum of the second loss term and the first loss term. In particular, the system 200 can backpropagate gradients of the loss function 230 to compute a respective gradient with respect to each parameter in at least a subset of the parameters of the network parameters of the diffusion neural network 110. The system 200 can then apply an optimizer, e.g., SGD, Adam, rmsProp, and so on, to these gradients to update the values of the (at least a) subset of the parameters. More specifically, the system 200 updates the current values of the network parameters 232 based upon the values computed using the loss function 230 for each pair of denoising trajectories. That is, the current values get updated after each training step, in which a pair of denoising trajectories is analyzed. Attorney Docket No.56113-0561WO1 When criteria are satisfied during training, e.g., after each training step or after every N training steps, where N is greater than one, the system 200 updates the old values of the network parameters based on the current values to take a “snapshot” of the current state of the model before initiating the next training step, starting the above-described process over. That is, the system sets the old values equal to the current values. By repeatedly performing training steps 240 for different sets of one or more conditioning inputs 202, the system 200 effectively “fine-tunes” the diffusion neural network 110 to generate outputs that result in increased reward scores. FIG.3 illustrates reward score improvement using the example training system. Table 302 quantitatively values two reward models, e.g., HPSv2 and PickScore, of differently trained models on a small-scale prompt database, e.g., 45 prompts. In this particular example, the three models compared are a stable diffusion (SD) model 306, a diffusion model fine-tuned using denoising diffusion policy optimization (DDPO) 308, and a diffusion model fine-tuned with the example training system described in this disclosure, denoted as proximal reward difference prediction (PRDP) 310. The evaluation prompts are sampled from a training set and the generated images are evaluated by the two reward models respectively, e.g., HPSv2 and PickScore. As seen in table 302, DDPO 308 and PRDP 310 both significantly outperform SD 306 in terms of maximizing reward scores. For example, with regards to the HPSv2 reward model, the SD model 306 has a reward score of 0.2855, while DDPO 308 has a reward score of 0.3398 and PRDP 310 has a reward score of 0.3471. Moreover, PRDP 310 improves the reward score on top of DDPO 308, e.g., 0.3398 for DDPO versus 0.3471 for PRDP. Generally, both DDPO and PRDP can handle small scale fine-tuning tasks well with PRDP, the training system described in this disclosure, having a small advantage. However, in terms of large-scale fine-tuning tasks, PRDP 310 shows a great advantage over the previous methods. Table 354 presents quantitative reward evaluation for the different approaches on a large-scale prompt dataset, e.g., over 100k prompts. Similar to the above, the evaluation prompts are sampled from a training set and the generated images are evaluated by the two reward models respectively, e.g., HPSv2 and PickScore. However, the prompts were tested both as seen and unseen prompts. The results of the table demonstrate that PRDP 310 can achieve higher generation quality and generalization capability for both reward models over SD 306 and DDPO 310 approaches. For example, for seen HPSv2 prompts, PRDP 310 has a reward score 0.31 Attorney Docket No.56113-0561WO1 while SD 306 has a reward score of 0.27 and DDPO 308 has a reward score of 0.25. This table further illustrates that DDPO 308 has lower reward scores than the SD approach 306 as DDPO performs worse due to training stability since it fails to learn useful knowledge during fine-tuning. As seen in the tables above, PRDP 310, the training process described in this disclosure, maximizes reward scores and improves image generation over previous methods. FIG.4 illustrates the progression of an example output of the training system after one or more training steps. In this particular example, the training system is configured to receive a conditioning input that includes a sequence of text and generate an output data item that includes an image that describes the text, i.e., the conditioning input can be a caption for the output image. For example, the conditioning input 442 could be “a brown squirrel” and the conditioning input 444 could be “a bird.” After the first training iteration 400, described in detail above, the system outputs image 402 from a conditioning input 442 and image 404 from a conditioning input 444. In this particular example, the system outputs a simple image of a detailed brown squirrel as image 402 and an image of a detailed black bird as image 404. The system continues to “fine-tune” the model, as described above, and after the second training iteration 410, the system outputs images 412 and 414 from their respective conditioning inputs 442 and 444. Image 412 is a more detailed image of a squirrel perched on an object with a blurred background composed of leaves and the outline of a tree. The same could be said in regard to image 414, i.e., the bird has more details and is perched on a branch. As illustrated in FIG.4, the quality of the image generated by the diffusion neural network has increased from figure 402 and 412 as well as with figure 404 and 414 by generating images with increased details within the animals as well as increased background details, e.g., colored leaves and the outline of a tree. This pattern continues with every training iteration, ending with images 432 and image 434 after training iteration 430, which are the most detailed and highest quality images generated by the diffusion neural network. As demonstrated in FIG.4, the “fine- tuning” of the diffusion neural network has resulted in iterative image quality improvement with each training iteration given the same conditioning input, e.g., conditioning inputs 442 and 444. As described above, small iterative updates to the current network parameters, allow for improvement in image generation quality, e.g., clearer, and more detailed, while ensuring the model does not stray from its generative capabilities. Attorney Docket No.56113-0561WO1 FIG.5 is a flow diagram of an example process 500 for training a diffusion neural network having a plurality of network parameter values. For convenience, the process 500 will be described as being performed by a system or one or more computers located in one or more locations. For example, a training system 100, e.g., the training system 100 depicted in FIG.2, appropriately programmed with this specification, can perform process 500. 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, e.g., the pre-trained diffusion neural network 130 of FIG.1. 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. Generally, in these implementations, the system initializes current 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. The system also initializes old values, e.g., a snapshot of the values of the parameters from a previous training step to pre- trained values of corresponding network parameters before training begins. 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 current value of each network parameter of the diffusion neural network to be equal to the pre-trained value of the corresponding network parameter. The system obtains a plurality of conditioning inputs for a training iteration (step 502). As described above, conditioning inputs can be multimodal, e.g., audio inputs, text sequences, images, etc. In particular, the conditioning inputs can be text sequences where the output data item generated is an image. The system then performs steps 504-508 for each training iteration. The system then generates a plurality of denoising trajectories for each conditioning input in accordance with the old values of the network parameters (step 504). Each denoising trajectory includes a sequence of data items starting from a noisy data item and terminating in a denoised data item. This is described in further detail with reference to FIGS.1 and 2. Attorney Docket No.56113-0561WO1 The system applies a reward function to the denoised data items for each denoising trajectory and generates a reward score for that denoising trajectory (step 506). In some implementations, the system applies the reward function only to the final, denoised data item in the trajectory. In some other implementations, the system can apply the reward function to each of the denoised data items in the denoising trajectory and calculate an overall reward score. In some implementations, the reward function can be a pretrained reward model, e.g., HPSv2 or PickScore that measures the generation quality. Then, the system trains the diffusion neural network to update current values of network parameters (step 508). The training of the diffusion neural network by updating current values of network parameters is described in further detail below in reference to FIGS.6 and 7. FIG.6 is a flow diagram of substeps of step 508 in FIG.5, further describing training the diffusion neural network to update current values of the network parameters. For each of the one or more conditioning inputs and for each of a set of pairs of denoising trajectories for the conditioning input, the system generates a respected likelihood score for each denoising trajectory in the pair in accordance with current values of network parameters (step 602). The likelihood score quantifies the probability of observing the entire diffusion trajectory from the initial noisy image to the final data item given the diffusion model 110’s current parameters and the conditioning input. For example, for each denoising trajectory, the log-likelihood can be computed as the denoising trajectory’s likelihood score. The log-likelihood equation for a denoising trajectory is described in further detail above with reference to FIG.2. The system then trains the diffusion neural network to update the current values of the network parameters on a loss function (step 604). The loss function can be any appropriate function. For example, the loss function can include a first loss term that measures, for each of the one or more conditioning inputs and for each of the set of pairs of denoising trajectories for the conditioning input, (i) a difference between the respective reward scores for the denoising trajectories in the pair and (ii) a difference between the respective likelihood scores for the denoising trajectories in the pair. That is, the first loss term of a loss function for a pair of denoising trajectories for a conditioning input measures the difference between the respective reward scores and the difference between the respective likelihood scores for the denoising trajectories in the pair. The loss function is described in further detail above with reference to FIG.2. Attorney Docket No.56113-0561WO1 FIG.7 is a flow diagram of sub steps of step 602 in FIG.5. In order to generate a likelihood score for the denoising trajectory, the system generates two likelihood scores, e.g., one corresponding to the current values of the network parameters and one corresponding to the score from the reference values of the network parameters. The system generates a first likelihood score for the denoising trajectory in accordance with the current values of network parameters (step 702). For example, the system can calculate a log-likelihood score for the denoising trajectory that is generated in accordance with the current values of the network parameters, as described above with reference to FIG.2. The system then generates a second likelihood score for the denoising trajectory in accordance with the reference values of the network parameters (step 704). For example, the system can calculate a log-likelihood score for the denoising trajectory that is generated in accordance with the reference values of the network parameters, as described above with reference to FIG.2. The system determines the likelihood score for the denoising trajectory from the first and second likelihood scores for the denoising trajectory (step 706). The system is able to evaluate the updates and the effect on the performance of the model. In some implementations, the system determines the likelihood score from the first and second likelihood scores based on a ratio between the first and second likelihood scores. For example, the likelihood score can be the logarithm of the ratio between the first and second likelihood scores: ^̂^^^^̅^^, ^^^=^^^^^^గ^^௫̅ೌ|^^ , where ^^^^^̅^^|^^^ represents the first in accordance with the current values and ^^^^^^^̅^^|^^^ represents the second likelihood score generated in accordance values for a denoising trajectory ^^ ̅^ . 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 Attorney Docket No.56113-0561WO1 that, when executed by data processing apparatus, cause the apparatus to perform the operations or actions. 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, i.e., 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. 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. 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 Attorney Docket No.56113-0561WO1 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. 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. 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. 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. 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 Attorney Docket No.56113-0561WO1 Positioning System (GPS) receiver, or a portable storage device, e.g., a universal serial bus (USB) flash drive, to name just a few. 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. 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. 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, i.e., inference, workloads. Machine learning models can be implemented and deployed using a machine learning framework, e.g., a TensorFlow framework or a Jax framework. 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 Attorney Docket No.56113-0561WO1 by any form or medium of digital data communication, e.g., a communication network. Examples of communication networks include a local area network (LAN) and a wide area network (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. 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. Similarly, while operations are correspond toed 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. Particular embodiments of the subject matter have been described. Other embodiments are within the scope of the following claims. For example, the actions recited Attorney Docket No.56113-0561WO1 in the claims can be performed in a different order and still achieve desirable results. As one example, the processes correspond toed 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. What is claimed is:

Claims

Attorney Docket No.56113-0561WO1 CLAIMS 1. A method performed by one or more computers and for training a diffusion neural network having a plurality of network parameter values, the method comprising, at each of a plurality of training iterations: obtaining a plurality of conditioning inputs for the training iteration; for each conditioning input: generating a plurality of denoising trajectories conditioned on the conditioning input and using the diffusion neural network and in accordance with old values of the network parameters, wherein each denoising trajectory comprises a sequence of data items starting from a noisy data item and terminating in a denoised data item; and for each denoising trajectory, applying a reward function to the denoised data items to generate a reward score for the denoising trajectory; and training the diffusion neural network to update current values of the network parameters, comprising, at each of one or more training steps: for each of one or more of the conditioning inputs and for each of a set of pairs of denoising trajectories for the conditioning input: generating a respective likelihood score for each denoising trajectory in the pair given the conditioning input using the diffusion neural network and in accordance with the current values of the network parameters; and training the diffusion neural network to update the current values of the network parameters on a loss function that comprises a first loss term that measures, for each of the one or more conditioning inputs and for each of the set of pairs of denoising trajectories for the conditioning input, (i) a difference between the respective reward scores for the denoising trajectories in the pair and (ii) a difference between the respective likelihood scores for the denoising trajectories in the pair.

2. The method of claim 1, wherein the data items are images.

3. The method of any preceding claim, wherein the conditioning inputs comprise text sequences.

4. The method of any preceding claim, further comprising, at each of the plurality of training iterations and after training the diffusion neural network to update the current values of the network parameters, updating the old values of the network parameters based on the current values of the network parameters.Attorney Docket No.56113-0561WO1 5. The method of claim 4, wherein updating the old values of the network parameters based on the current values of the network parameters comprises: setting the reference values equal to the old values of the network parameters.

6. The method of any preceding claim, wherein prior to performing the plurality of training iterations, the diffusion neural network has been pre-trained to determine reference values of the network parameters.

7. The method of claim 6, further comprising, prior to performing the plurality of training iterations: setting the current values of the network parameters and the old values of the network parameters to the reference values of the network parameters.

8. The method of claim 6 or claim 7, wherein generating a respective likelihood score for each denoising trajectory in the pair given the conditioning input using the diffusion neural network and in accordance with the current values of the network parameters comprises, for each denoising trajectory in the pair: generating a first likelihood score for the denoising trajectory given the conditioning input using the diffusion neural network and in accordance with the current values of the network parameters; generating a second likelihood score for the denoising trajectory given the conditioning input using the diffusion neural network and in accordance with the reference values of the network parameters; and determining the likelihood score for the denoising trajectory given the conditioning input from the first and second likelihood scores for the denoising trajectory.

9. The method of claim 8, wherein determining the likelihood score for the denoising trajectory given the conditioning input from the first and second likelihood scores for the denoising trajectory comprises: determining the likelihood score based on a ratio between the first and second likelihood scores.

10. The method of any preceding claim, wherein the first loss term measures a mean squared error between (i) the difference between the respective likelihood scores for the denoising trajectories in the pair and (ii) the difference between the respective reward scores for the denoising trajectories in the pair divided by a regularization strength hyperparameter.Attorney Docket No.56113-0561WO1 11. The method of any preceding claim wherein the loss function comprises a second loss term that measures, for each of the one or more conditioning inputs and for each of the set of pairs of denoising trajectories for the conditioning input, (i) a difference between respective clipped likelihood scores for the denoising trajectories in the pair and (ii) the difference between the respective reward scores for the denoising trajectories in the pair.

12. The method of claim 11, wherein the loss function is a maximum of the second loss term and the first loss term.

13. The method of any preceding claim, wherein the diffusion neural network has been pre-trained on a score matching objective.

14. The method of any preceding claim, wherein the reward function is configured to process a reward input that comprises the conditioning input.

15. The method of any preceding claim when dependent on claim 2, wherein the reward function measures image quality.

16. A system comprising one or more computers and one or more storage devices storing instructions that are operable, when executed by the one or more computers, to cause the one or more computers to perform the operations of the respective method of any preceding claim.

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

Images