Training diffusion neural networks by back-propagating differentiable rewards

Abstract

Methods, systems, and apparatus, including computer programs encoded on a computer storage medium, for training diffusion neural networks using differentiatable reward functions.

Description

The diffusion neural network is trained using backpropagation of differentiable rewards.

[0001] Cross-references to related applications

[0002] This application claims priority to U.S. Provisional Application No. 63 / 541,287, filed September 28, 2023. The disclosure of the earlier application is considered part of the disclosure of this application and is incorporated herein by reference. Background Technology

[0003] This manual relates to using machine learning models to process data.

[0004] As an example, a neural network is a machine learning model that uses one or more layers of non-linear units to predict an output from a received input. In addition to the output layer, some neural networks also include one or more hidden layers. 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 the received input based on the current values ​​of a corresponding set of weights. Summary of the Invention

[0005] This specification describes a system implemented as a computer program on one or more computers at one or more locations, which uses a diffuse neural network to generate output data items conditioned on conditional inputs.

[0006] Typically, conditional inputs characterize one or more desired attributes of the data item; that is, they characterize one or more attributes that the final data item generated by the system should have.

[0007] More specifically, this specification describes how the system can train a diffusion neural network using reinforcement learning with a differentiable reward function, for example, after the diffusion neural network has already been trained on a target without a reward function.

[0008] Specific embodiments of the subject matter described in this specification can be implemented to achieve one or more of the following advantages.

[0009] Diffusion models have revolutionized generative modeling for continuous data, achieving impressive results across modalities including images, videos, and audio. However, for many use cases, accurately modeling the distribution of the training data (e.g., diverse images from a large dataset) does not align with the model's expected behavior after training (e.g., generating aesthetically pleasing outputs).

[0010] To overcome this mismatch, this specification proposes an efficient method for gradient-based reward fine-tuning based on differentiation through a diffusion sampling process. Specifically, the training system can execute the described training method after, for example, "pre-training" the diffusion neural network on a large dataset of diverse data items, to align the diffusion neural network to achieve the desired post-training behavior, i.e., to enable the diffusion neural network to generate data items with the quality measured by the reward function after training.

[0011] In some described methods, the system backpropagates rewards through the complete sampling chain. However, this can lead to excessive memory consumption, computational inefficiency, or both during the training process.

[0012] To maintain low memory and computational costs during training, this specification describes a variety of modifications that the system can employ.

[0013] As an example, this specification describes how the system uses gradient checkpointing to reduce the computational cost of training while minimizing memory overhead.

[0014] As another example, the system can optimize a proper subset of network parameters instead of the complete set of network parameters to improve the computational efficiency of the training process.

[0015] As another example, this specification describes how the system can perform backpropagation to compute gradients only through the last K sampling steps. Specifically, this not only improves the computational efficiency of the training process but also avoids gradient explosion that occurs when backpropagating through the full sampling chain. In other words, with the same number of training steps, using "truncated" gradients performs better than computing the full gradients in both computational efficiency and the quality of the resulting model.

[0016] As yet another example, this specification describes how the system can further improve efficiency by introducing gradient estimates with lower variance computed on a noisy sample set. By using these noisy samples to reduce the variance of the gradients, the system can stabilize the training process, thereby reducing the number of training iterations required (thus improving the computational efficiency of the training process), or obtaining a higher quality trained model, or both.

[0017] Furthermore, because the described techniques utilize gradients, they are far more efficient than fine-tuning baselines based on reinforcement learning.

[0018] Details of one or more embodiments of the subject matter of this specification are set forth in the accompanying drawings and the following description. Other features, aspects, and advantages of the subject matter will become apparent from the description, the drawings, and the claims. Attached Figure Description

[0019] Figure 1 is a diagram of the example training system.

[0020] Figure 2 is a flowchart of an example process for training a diffusion neural network.

[0021] Figure 3 is a flowchart of an example process for determining the variance reduction term of the loss function.

[0022] Figure 4 shows an example of system operation when a diffusion neural network is used to generate images.

[0023] Figure 5 illustrates an example of the performance of the described technique when fine-tuning the reward function that measures aesthetic scores.

[0024] The same reference numerals and names in the various figures indicate the same elements. Detailed Implementation

[0025] This specification describes a system implemented as a computer program on one or more computers at one or more locations, which trains a diffuse neural network to generate output data items conditioned on conditional inputs.

[0026] Typically, conditional inputs characterize one or more desired attributes of the data item; that is, they characterize one or more attributes that the final data item generated by the system should have.

[0027] The system can be configured to generate various output data items based on various input conditions.

[0028] For example, the system can be configured to generate audio data, such as the waveform of the audio or the time spectrum of the audio, such as a Mel time spectrum or a time spectrum with frequencies at different scales.

[0029] In this example, the conditional input can be the text that the audio should represent or the features of the text, i.e., making the system function as a text-to-speech machine learning model that converts the text or the features of the text into audio data of the spoken text.

[0030] As another example, conditional input can identify the expected speaker of the audio, that is, cause the system to generate audio data representing the speech spoken by the expected speaker.

[0031] As another example, conditional inputs can characterize attributes of songs or other musical works, such as lyrics, genres, etc., enabling the system to generate musical works with attributes characterized by the conditional inputs.

[0032] As another example, the conditional input can specify that the audio data is classified into one of a set of possible categories, causing the system to generate audio data belonging to that category. For example, the category could represent the type of musical instrument or other sound-producing device, i.e., causing the system to generate audio emitted by the corresponding category, or the type of animal, i.e., causing the system to generate audio representing the noise produced by the corresponding animal, and so on.

[0033] As another concrete example, the data item can be an image, allowing the system to perform conditional image generation by generating intensity values ​​for the image pixels. Typically, the conditional input can specify one or more properties of the image.

[0034] In this specific example, the conditional input can be a sequence of text, and the output data item can be an image describing that text; that is, the conditional input can be the title of the output image.

[0035] As yet another concrete example, the conditional input can be an object detection input, which specifies one or more bounding boxes and, optionally, the corresponding object type that should be depicted in each bounding box.

[0036] As another concrete example, conditional input can specify a particular object category to which the object depicted in the output image belongs from among multiple object categories.

[0037] As another example, conditional input can specify one or more images.

[0038] For example, the conditional input can specify an image at a first resolution, and the output data items can include a second, higher resolution image.

[0039] For example, the conditional input can specify an image, and the output data items can include a denoised, enhanced, stylized, or otherwise edited version of that image.

[0040] As another concrete example, the conditional input can specify an image that includes the target entity (e.g., a tumor) for detection, and the output data item can include an image that does not contain the target entity, for example, to facilitate the detection of the target entity by comparing images.

[0041] As another concrete example, the conditional input could be a segmentation that assigns each of the multiple pixels in the output image to a category from a set of categories, for example, assigning a corresponding category to each pixel.

[0042] As another example, the conditional input can be a different type of structured input, such as a grid or graph that specifies the attributes of the image to be generated.

[0043] More generally, conditional input can include one or more different types of input in one or more different modalities, such as text only, one or more images only, text and one or more images, and so on.

[0044] As another example, the output data item can be a video. Similarly, conditional inputs can specify one or more characteristics of the video.

[0045] As a specific example, the conditional input may include text, and the output data item may be a video described by that text.

[0046] As another concrete example, the conditional input may include one or more images, and the output data item may be a video that completes the one or more images, for example, a video that begins with the one or more images.

[0047] More generally, the task of generating output data items can be any task that outputs continuous data with conditional inputs. For example, the output can be the output of different sensors, such as lidar point clouds, radar point clouds, electrocardiogram readings, etc., and the conditional inputs can represent the type of data that should be measured by the sensors. When a discrete output is required, it can be obtained, for example, by thresholding the output generated by a diffuse neural network.

[0048] In some applications, output data items can be used in control tasks to control the actions of a mechanical agent acting in a real environment to perform mechanical tasks. For example, the output data items can be processed by the agent's policy neural network to select one or more actions to be performed by the agent as part of a task. The agent can then perform those one or more actions. The output data items (e.g., images) can, for example, characterize the state of the real environment predicted to be obtained by the agent performing those one or more actions. The conditional input can, for example, specify the state of the real environment and those one or more actions. As another example, the conditional input can specify the state of the real environment, and the output data items can be used to select one or more actions to be performed by the mechanical agent to perform a task (i.e., a diffuse neural network can represent an action selection policy).

[0049] In any of the above examples, the output data item generated using the diffusion neural network can be an output data item in the output space, i.e., such that the value in the output data item is the value of a data item of an appropriate type, such as an image pixel value, an audio signal amplitude value, etc.; or an output data item in the latent space, i.e., such that the value in the output data item is the value in the latent representation of the output data item in the output space.

[0050] When output data items are generated in the latent space, the system can generate the final output data items in the output space by processing the output data items in the latent space using a decoder neural network (e.g., a decoder neural network pre-trained in an autoencoder framework). During training, the system can encode the target data items in the output space using an encoder neural network (e.g., an encoder neural network pre-trained jointly with the decoder in an autoencoder framework) to generate the target output for the diffusion neural network in the latent space.

[0051] Figure 1 is a diagram of an example training system 100. The training system 100 is an example of a system implemented as a computer program on one or more computers at one or more locations, wherein the following systems, components and techniques can be implemented.

[0052] System 100 is a system for training a diffusion neural network 110, which is used to generate an output (final) data item 112 given a conditional input 102.

[0053] Specifically, system 100 trains a diffuse neural network 110 such that the output (final) data item 112 has one or more desired properties characterized by the conditional input 102.

[0054] More specifically, system 100 uses reward function 120 to perform "fine-tuning" on diffuse neural network 110, i.e., further training.

[0055] In other words, before training with reward function 120, system 100 or another training system has already trained the diffusion neural network on a different objective (e.g., an objective without reward function 120). Typically, the diffusion neural network can have already been trained in a conventional manner with any diffusion model objective. As an example, the diffusion neural network can have already been trained on a set of training data items with respect to a diffusion score matching objective or a variant thereof.

[0056] The reward function 120 can be any suitable differentiable reward function that maps an input including (i) a data item or (ii) a potential representation of a data item to a reward score 122. Optionally, the reward function input may also include a conditional input 102 or a representation of the conditional input 102.

[0057] For example, reward function 120 may include one or more trained reward machine learning models, such as neural networks.

[0058] As an example, the reward function 120 may include a machine learning model that maps at least a portion of the reward input to a score representing the aesthetic quality of the output data item.

[0059] As a concrete example of a reward function representing aesthetic quality, an aesthetic predictor model can be trained on a dataset comprising multiple data items, each of which has been assigned an aesthetic score that measures the aesthetic quality of that data item. That is, the predictor can be trained, for example, using mean squared error or mean absolute error loss, to predict the aesthetic scores of data items in the dataset. As a more specific example, the aesthetic predictor model can include (i) a pre-trained data item embedding model, such as a model trained through contrastive learning or other representation learning techniques, and (ii) one or more output layers, such as fully connected layers, that process the embeddings generated by the data item embedding model to produce aesthetic scores. In this example, during training, the system can keep the embedding model fixed and only train the output layers.

[0060] As another example, reward function 120 may include a machine learning model that maps at least a portion of the reward input to a score representing the predicted quality of the output data item, such as a human user rating it. For example, the reward function could be a reward model already trained to model human preferences, for instance, on an objective function trained using human preferences between pairs of data items. An example of such a model is the Human Preference Score v2 model, described in the paper “Human preference score v2: A solid benchmark for evaluating human preferences of text-to-image synthesis” by Xiaoshi Wu, Yiming Hao, Keqiang Sun, Yixiong Chen, Feng Zhu, Rui Zhao, and Hongsheng Li.

[0061] As another example, reward function 120 may include a machine learning model trained to perform a data item detection or recognition task, such that the reward function penalizes output data items for including specific categories. For example, when the output data item is an image, the model may be an object detection model, such as an open-vocabulary object detection model. In this example, the system may pass images generated by the diffusion model along with a set of queries Q that should be excluded from the generated images through a pre-trained object detection model. As a reward, the system may use the sum of scores of the located objects corresponding to all queries, the sum of their bounding box areas, or both.

[0062] As another example, reward function 120 may include a reward that causes the diffusion model to generate adversarial examples. That is, the system can fine-tune the diffusion model such that the output data items generated based on hints for category y are classified into distinct categories y′ by a pre-trained classifier for specific types of data items. For example, as a reward, the system may use the negative cross-entropy of the pre-trained classifier for the target category.

[0063] As another example, reward function 120 may include one or more hard-coded differentiable reward functions.

[0064] For example, a hard-coded reward function could include a function that measures the compressibility of the output data items. For instance, a compressibility reward function could pass the output data items through differentiable compression (…). ) and decompression ( The algorithm reconstructs the output data items and then outputs a value as a reward score based on the error between the original image and the reconstructed image (e.g., Euclidean distance), such as a negative value of the error. .

[0065] When there are multiple reward models in the reward function 120, the final reward score 122 can be the sum or weighted sum of the reward scores generated by these models.

[0066] Typically, system 100 uses a set of conditional inputs 102 to perform further training.

[0067] For each conditional input 102, system 100 uses a diffuse neural network 110 to generate the final representation 106 of data item 112.

[0068] The system 100 then uses the final representation 106 to generate a reward input 118 for the reward function 120. For example, the reward input 118 may include a data item 112 generated from the final representation 106, and an optional conditional input 102.

[0069] Then system 100 applies reward function 120 to reward input 118, and generates reward score 122 based on condition input 102.

[0070] The system then uses the reward score 122 to train the diffusion neural network 110 (130), that is, to update the network parameters 150 of the diffusion neural network 110. The parameters of the diffusion neural network 110 typically include the weights of each layer of the diffusion neural network 110, and in some cases, also include the biases of each layer.

[0071] In some implementations, the system 100 updates all network parameters 150 of the diffuse neural network 110 by performing further training.

[0072] In some other implementations, the diffuse neural network 110 has a first set of network parameters and a second set of network parameters, and as part of (further) training of the diffuse neural network 110, the system 100 updates the first set of network parameters while keeping the second set of network parameters fixed.

[0073] For example, before further training, system 100 or another training system may have already trained an instance of the diffuse neural network 110 that does not include the first network parameter set (i.e., only includes the second network parameter set). For example, the training system may have already trained an instance of the diffuse neural network 110 on a score matching target.

[0074] In this example, during further training, system 100 can keep the second network parameter set fixed as pre-trained values ​​determined by, for example, training instances of the diffuse neural network 110 that do not include the first network parameter set with respect to a score matching objective.

[0075] For example, System 100 can use a low-rank approximation (LoRA) technique when performing further training (Hu et al., arXiv:2106.09685, 2021). In this case, for each weight matrix included in one or more weight matrices in the second network parameter set, the first network parameter set includes a low-rank factorization of the updated weight matrix that can be used to update that weight matrix. The low-rank approximation technique can be performed on multiple different weight matrices to update corresponding different layers of the diffuse neural network.

[0076] The system can approximate the update of the weight matrix during each training update of the diffuse neural network using a low-rank approximation. This is achieved, for example, by optimizing the product of two smaller matrices to reduce the computational dimensionality required to compute the weight changes needed for the update. More specifically, performing a low-rank approximation means decomposing the updated weight matrix into the product of two smaller matrices, which, when multiplied, can faithfully recover the values ​​of the updated weight matrix. In particular, the low-rank decomposition can represent… ,in, It is the weight matrix in the second network parameter set. Is with The corresponding updated weight matrix is ​​then used, and the product BA approximates... For example, the second set of network parameters may include the set of parameters of the initial diffuse neural network to be fine-tuned (and remain fixed during fine-tuning), and the first set of parameters may include the set of parameters added to the initial diffuse neural network and adjusted during training (fine-tuning).

[0077] In this case, the rank of the matrix refers to the number of linearly independent vectors, such as in matrix factorization. The sum of columns or rows that do not contain relevant data. The determined rank specifies the dimension to be updated by constraining the dimensions of the two smaller matrices. For example, in B, where the dimension is d... r is a matrix and A has dimension r. In the case of k, r must be the same for matrix multiplication to be performed, and the rank r can be a value much smaller than the minimum of d and k, for example, r min(d, k).

[0078] Therefore, during training, the system learns the weights in matrices B and A, rather than directly learning them. The weights in the equation.

[0079] The diffusion neural network 110 can be any suitable diffusion neural network configured to receive inputs including the current (noisy) representation of the image and conditional inputs, and generate a denoised output.

[0080] In some implementations, the diffusion neural network 110 performs a diffusion process in the output space (e.g., the pixel space when the data item is an image). In this example, when the data item is an image, the data item (“representation”) operated on and generated by the diffusion neural network 110 has a value for each pixel that specifies a color value (e.g., an RGB value or another color encoding scheme).

[0081] Examples of this type of diffusion neural network include Imagen.

[0082] In some other implementations, the diffusion neural network 110 performs the diffusion process in a latent space (e.g., a latent space with a dimension lower than the output space). That is, the data item (“representation”) operated on by the diffusion neural network 110 is a latent representation, and when the data item is an image, the value in the representation is the learned latent value, rather than, for example, a color value.

[0083] Examples of such diffusion neural networks include MobileDiffusion, as described in arxiv:2311.16567.

[0084] In these implementations, during training, the diffusion neural network 110 can be associated with an encoder to encode training data items into a latent space, and after training and in order to generate new output data items, it can be associated with a decoder neural network that receives input including a latent representation of the data items and decodes the latent representation to reconstruct the data items.

[0085] The following section will describe the implementation of further training in more detail.

[0086] After training, system 100 or another inference system can use the diffuse neural network 110 to generate new data items 112 conditioned on the new conditional input 102.

[0087] The diffuse neural network 110 can have any suitable architecture that allows the neural network to map a diffused input, which includes input data items with the same dimension as the output data item 112, to a denoised output that also has the same dimension as the output data item 112.

[0088] For example, when the output data item is an audio signal or an image, the diffuse neural network 110 can be a convolutional neural network, such as U-Net, or other architectures that map an input of a given dimension to an output of the same dimension.

[0089] 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 a denoised output.

[0090] The neural network 110 can be conditional on the conditional input 102 in any of a variety of ways.

[0091] As an example, system 100 may use an encoder neural network to generate one or more embeddings representing conditional input 102, and diffusion neural network 110 may include one or more cross-attention layers, each cross-attention layer cross-attention to the one or more embeddings.

[0092] As used in this specification, an embedding is an ordered collection of numerical values, such as a vector of floating-point values ​​or other types of values.

[0093] For example, when the conditional input is text, the system can use a text encoder neural network, such as a Transformer neural network, to generate a fixed or variable number of text embeddings representing the conditional input.

[0094] When the conditional input is an image, the system can use an image encoder neural network, such as a convolutional neural network or a visual Transformer neural network, to generate a set of embeddings representing the image.

[0095] When the conditional input is audio, the system can use, for example, an audio encoder neural network, such as an audio encoder neural network that has been jointly trained with a decoder neural network as part of a neural audio codec, to generate one or more embeddings that encode the audio.

[0096] When the conditional input is a scalar value, the system can use, for example, an embedding matrix to map the scalar value or its one-hot representation to the embedding.

[0097] In some cases, conditional input 102 may include multiple different types of input, such as two or more types of text, images, boundary values, or contextual embeddings.

[0098] In some of these cases, system 100 can generate one or more initial embeddings for each different type of input; that is, using an appropriate encoder neural network as described above, and then using a Transformer encoder neural network to process the initial embeddings for all different types of input to update each initial embedding, thereby generating a set of final embeddings. The one or more cross-attention layers within the diffusion neural network 110 can then cross-attention to this set of final embeddings.

[0099] In some other cases, different cross-attention layers within the diffuse neural network 110 can cross-focus on embeddings of different types of conditional inputs.

[0100] In some of these cases, system 100 can concatenate initial embeddings of different types of inputs along the sequence dimension, and then the one or more cross-attention layers can cross-focus on a concatenated set of final embeddings.

[0101] As another example, the diffuse neural network 110 may include one or more other types of neural network layers conditioned on the one or more embeddings. Examples of such layers include feature-level linear modulation (FiLM) layers, layers with conditionally gated activation functions, etc.

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

[0103] Figure 2 is a flowchart of an example process 200 for training a diffusion neural network using reward scores. For convenience, process 200 will be described as being executed by a system of one or more computers located in one or more locations. For example, a training system appropriately programmed according to this specification, such as training system 100 depicted in Figure 1, can execute process 200.

[0104] The system initializes the representation of the data item (step 202). For example, the system can initialize the representation by sampling values ​​from the representation from a distribution (e.g., a Gaussian distribution).

[0105] The system receives conditional input c representing one or more desired attributes of a data item (step 204). For example, the system can receive conditional input p from a set of conditional inputs. c Input c for the sampling condition.

[0106] The system then updates the representation of the data item to generate the final representation of the data item (step 206). Specifically, the system generates the final representation over multiple sampling iterations (T iterations in the example algorithm given later).

[0107] In each of the plurality of sampling iterations, the system uses a diffusion neural network to process the diffusion input, which includes representations of data items and representations of conditional inputs, for that sampling iteration to generate a denoised output for that sampling iteration (step 208).

[0108] Typically, the denoised output defines an estimate of the final representation given the current representation.

[0109] In some implementations, the denoised output is an estimate of the noise components of the current representation, which needs to be combined with the final representation (e.g., added or subtracted) to generate the noise of the current representation.

[0110] In some other implementations, the denoised output is an estimate of the final representation given the current representation, i.e., an estimate of the data items generated by removing the noise components from the current representation.

[0111] In other implementations, the system parameterizes the denoised output in different ways, such as using v-parameterization (Salimans and Ho arXiv: 2202.00512, 2022, Section 4; Appendix D) or other appropriate parameterization.

[0112] The system then uses the denoised output to update the representation of the data item (step 210).

[0113] In some implementations, the system uses the denoised output as the final denoised output for that update iteration.

[0114] In some other implementations, the system also generates one or more additional denoised outputs for the sampling iteration.

[0115] For example, the system can utilize classifier-less guidance.

[0116] In this example, the system uses a diffuse neural network to process a second diffuse input, which includes representations of data items but excludes conditional inputs, for the sampling iteration, to generate an unconditionally denoised output for the update iteration. For example, the second diffuse input may include representations of data items and a predetermined representation indicating unconditional sampling.

[0117] The system then updates the representation of the data item using the denoised output and the unconditional denoised output. Specifically, the system can combine the denoised output and the unconditional denoised output to generate a final denoised output based on the guiding weights used for this sampling iteration, and use the final denoised output to update the representation of the data item. For example, the system can set the final denoised output to be equal to (1+w). Denoising output –w Unconditional denoising output. That is, the final denoised output can be determined based on the difference between the first denoised output scaled by (1+w) and the sum of one or more additional denoised outputs scaled by w.

[0118] For example, in each iteration except the last, the system can use the (final) denoised output to generate an estimate of the representation, and then apply a diffusion sampler to the estimate. The system can use any suitable diffusion sampler to update the representation, such as a DDPM (Denoising Diffusion Probabilistic Model) sampler, a DDIM (Denoising Diffusion Implicit Model) sampler, or another suitable sampler for that estimate to generate the updated current data item. DDPM is discussed, for example, in Ho et al.'s arXiv:2006:11239.

[0119] For the last iteration, the estimate can be an updated representation, or the system can use a sampler.

[0120] After the last iteration of these multiple sampling iterations, the system generates the reward input from the final representation (step 212).

[0121] When the representation of a data item is the same as that in the output space, the system can directly include the final representation in the reward input. Optionally, the reward input may also include other information, such as conditional input.

[0122] When the representation of a data item is a latent representation in the latent space, the system can use a decoder neural network to process the final representation to generate the final data item and include the final data item in the reward input.

[0123] The system applies a reward function to the reward input to generate a reward score for the final data item, measuring its quality (step 214). For example, the reward function can be any reward function described above with reference to FIG1, or any other suitable differentiable reward function.

[0124] As a concrete example, the reward score for conditional input 102 (c) is 122 ( ) can be represented as ,in, Indicates the final data item. This represents the initial representation of the data item, and This represents the parameters of the diffusion network (e.g., including LoRA parameters). Later, in the algorithm, this will be represented as... ,in, This represents the final data item.

[0125] The system trains the diffusion neural network using a loss function that includes a first term that measures the reward score of the final data item (step 216). For example, the first term could be... .

[0126] As part of the training, the system can backpropagate the gradient of the first term through the reward function across a subset of sampling iterations to determine the gradient of the first term with respect to the network parameters of the diffusion neural network. For example, the gradient of the first term with respect to the network parameters can be backpropagated across a subset of sampling iterations to determine the gradient of the first term with respect to the network parameters of the diffusion neural network. , (That is, backpropagation through multiple calls of the diffusion model neural network in the sampling chain, similar to backpropagation across time). For example, the system can backpropagate the gradient of the first term through the reward function to determine the gradient of the first term with respect to the reward input, and then, as described above, backpropagate the gradient of the first term with respect to the reward input through a subset of sampling iterations to determine the gradient with respect to the network parameters.

[0127] When the representation of a data item is a latent representation in the latent space, because the system uses a decoder neural network to generate inputs to the reward function, training involves backpropagating the gradient of the first item through the reward function, through the decoder neural network, and through a subset (i.e., partially or entirely) of the sampling iterations.

[0128] For example, the system can perform the above steps in parallel for multiple different conditional inputs, and the loss function can be the sum or average of the corresponding terms for each conditional input.

[0129] As another example, in addition to the first term of the conditional input, the loss function can include additional terms, such as regularization terms.

[0130] As a concrete example of an addendum, the loss function may include an addendum that reduces the variance of the updates, calculated based on the reward score generated for the data item produced from the noisy version of the output data item.

[0131] This additional item is described in more detail below with reference to Figure 3.

[0132] In some cases, the subset of sampling iterations through which the system backpropagates the gradient may include all sampling iterations; that is, the subset is not a so-called "proper subset".

[0133] However, backpropagation throughout all sampling iterations can be computationally expensive.

[0134] The system can use any of a variety of techniques to solve this problem and make training more computationally efficient.

[0135] As an example, a subset of sampling iterations may not include one or more of the earliest sampling iterations. In other words, the subset may include all sampling iterations except for the one or more earliest iterations, or equivalently, exclude one or more of the earliest sampling iterations from the subset.

[0136] As a concrete example, a subset of sampling iterations can include only the most recent iterations among the multiple sampling iterations. That is, a subset of sampling iterations can include only the last K sampling iterations, where K is an integer greater than one.

[0137] As yet another concrete example, a subset of sampling iterations may include only the last of the multiple sampling iterations.

[0138] As another example, the system can use gradient checkpoints during training, except either by not including certain sampling iterations in the subset or by doing so instead.

[0139] In this example, to backpropagate gradients throughout a specific sampling iteration, the system retrieves the representation of the final data item up to that specific sampling iteration from memory. That is, when performing that specific sampling iteration as part of step 206, the system can store the input representation for that specific sampling iteration in memory. However, the system can avoid persistently storing intermediate activations of the diffuse neural network in memory when generating the denoised output for that specific sampling iteration; that is, it can avoid retaining intermediate activations in memory after the sampling iteration has been completed.

[0140] The system can then use the retrieved representation to recompute (rematerialize) the intermediate activations of the diffusion neural network used for that particular sampling iteration, i.e., by reprocessing the diffusion input, which includes the retrieved representation, using the diffusion neural network.

[0141] The system then uses these intermediate activations to compute the gradient of the first term for that particular sampling iteration, for example, by backpropagating the current gradient through a diffusion neural network using the intermediate activations.

[0142] Therefore, by avoiding persistently storing intermediate activations in memory, the system can reduce the amount of memory required to compute gradients. However, by still persistently storing the input representation used for this sampling iteration in memory, the system can ensure that the latency associated with training does not increase excessively, i.e., because the input representation does not need to be recomputed.

[0143] As a result of training (e.g., backpropagation of gradients), the system has a gradient for each parameter in at least one subset of the parameters of the diffuse neural network.

[0144] The system can then apply optimizers, such as SGD, Adam, rmsProp, etc., to these gradients to update the values ​​of the (at least one) subset of parameters.

[0145] By repeatedly performing this training on one or more conditional inputs from different groups, the system effectively “fine-tunes” the diffuse neural network to generate outputs that lead to higher reward scores.

[0146] Figure 3 is a flowchart of an example procedure 300 for determining the variance reduction term of the loss function. For convenience, procedure 300 will be described as being executed by a system of one or more computers located in one or more locations. For example, a training system appropriately programmed according to this specification, such as training system 100 depicted in Figure 1, can execute procedure 300.

[0147] The variance reduction term is also referred to as the "second" term in the loss function in this specification.

[0148] The system can execute process 300 in each of one or more noise iterations.

[0149] The system iteratively samples the noise (step 302). For example, the system can sample values ​​from the noise from the same distribution (e.g., Gaussian distribution) used to initialize the representation of the data items.

[0150] The system applies the noise to the final representation to generate a noisy representation (step 304).

[0151] The system uses a diffuse neural network to process the input, which includes a noisy representation, to iteratively generate a denoised output for the noise (step 306).

[0152] The system then uses the denoised output for the noise iteration to update the noisy representation to generate an updated noisy representation as an estimate of the final representation (step 308).

[0153] The system generates a new reward input from the updated noisy representation (step 310). When the representation of a data item is in the output space, the system can directly include the updated noisy representation in the new reward input. Optionally, the reward input may also include other information, such as conditional input.

[0154] When the representation of a data item is a latent representation in the latent space, the system can use a decoder neural network to process the updated noisy representation to generate a noisy data item and include the noisy data item in the new reward input.

[0155] The system applies a reward function to the new reward input to generate a new reward score (step 312).

[0156] Typically, the loss function described above includes a second term that measures the new reward score for the one or more noisy iterations. When there are multiple noisy iterations, the second term can measure a combination of the new reward scores, such as the average.

[0157] When the loss function includes a second term, training a neural network on the loss function may involve backpropagating the gradient of the second term through the reward function and backpropagating into noisy iterations, but not backpropagating through any sampling iterations. That is, the system may insert a stopping gradient to prevent the computation of the gradient of the second term with respect to the final representation or any representation generated in any previous sampling iteration.

[0158] When the diffusion sampler is DDIM and there are T samples, it is used to train a diffusion neural network with parameter θ. Example techniques are shown in Table 1.

[0159]

[0160] Table 1

[0161] Specifically, in Table 1, DRaFT refers to a version of the described technique where the subset of sampling iterations includes all sampling iterations; DRaFT-K refers to a version of the described technique where the subset of sampling iterations includes only the last K sampling iterations; and DRaFT-LV refers to a version of the described technique where the subset includes only the last sampling iteration and the second term is included in the loss function.

[0162] In Table 1, the “DRaFT-LV then…” part of the algorithm refers to the calculation of the second term mentioned above. However, although the example in Table 1 shows that the second term is used only with DRaFT-LV, in practice, the second term can also be used with either DRaFT or DRaFT-K.

[0163] Figure 4 illustrates an example 400 of the technique described when a diffusion neural network is used to generate an image from a conditional input including a text sequence.

[0164] In example 400, the system receives conditional input (“prompt”) 402 “majestic lion”. The system then initializes the representation 404 x. T The conditional input 402 is processed using a diffuse neural network 110 to update the representation across multiple sampling iterations until the final representation x0406 is reached.

[0165] Specifically, in Example 400, the representation is in the output space, so each representation x is an image, and different representations have different amounts of noise.

[0166] When using DRaFT-LV, the system can then use x0 to generate multiple updated noisy representations 408. For example, the system can add noise to the same final representation x0 n times without regenerating a new final representation, thus avoiding the computational cost of T update steps, i.e., sampling iterations (each step involving, for example, a call to a diffuse neural network).

[0167] The gradients of the reward gradients *g* of these noisy representations—that is, the gradients of the first term of the loss function (which measures the reward score)—can be summed to obtain the gradient of the first term of the loss function used for backpropagation to train the diffusion neural network. This can be a particularly efficient method. For example, using... It is twice as efficient as DRaFT, while using only 10% more computation.

[0168] For all variants, the system then scores the generated representations using a reward function r120. In Example 400, the reward function consists of multiple separate reward functions that measure various properties of these representations, such as aesthetic score, PickScore, compressibility, and object detection, as described above.

[0169] Then, depending on whether the technique used is DRaFT-LV, DRaFT, or DRaFT-K, the system backpropagates the gradients throughout each sampling iteration to determine the final gradients that will be used to update the network parameters.

[0170] Furthermore, in Example 400, the system is using the low-rank adaptation (LoRA) training technique, which updates only a subset of the model's network parameters.

[0171] In some implementations, after training, the system can further modify the values ​​of the network parameters of the neural network before using the diffuse neural network to generate new data items.

[0172] For example, the system can train multiple instances of a diffuse neural network using the LoRA training technique, each instance having a different reward function. In this example, the system can generate new instances of the diffuse neural network by computing a weighted sum of the values ​​of the network parameters across multiple instances for each network parameter in a first subset updated through LoRA training, where the weight of each instance is determined by the strength of one or more attributes, measured by the corresponding reward function, that should be reflected in the new data items generated after training.

[0173] As another example, the system can use the LoRA training technique to train individual instances of a diffusion neural network with different reward functions. In this example, the system can generate new instances of the diffusion neural network by calculating a weighted sum of the fine-tuned values ​​of each network parameter and its pre-trained values ​​for each network parameter in a first subset updated through LoRA training, where the weight of each value is determined by the strength of one or more attributes measured by the reward function that should be reflected in the new data items generated after training.

[0174] Figure 5 illustrates example 500 of the performance of the described technique when fine-tuned on the reward function that measures aesthetic scores. In particular, as can be seen from example 500, the variant of the described technique outperforms two existing techniques (ReFL and DDPO), as well as a pre-trained model (Stable Diffusion) and using cue engineering on a range of reward queries (training condition inputs).

[0175] This specification uses the term "configured to" in conjunction with system and computer program components. For configuring one or more computer systems to perform a specific operation or action, it means that software, firmware, hardware, or a combination thereof are installed on the system to cause the system to perform that operation or action during operation. For one or more computer programs configured to perform a specific operation or action, it means that the one or more programs include instructions that, when executed by a data processing device, cause that device to perform that operation or action.

[0176] Embodiments of the subject matter and functional operation described in this specification may be implemented in digital electronic circuit systems, in tangibly embodied computer software or firmware, in computer hardware (including the structures disclosed in this specification and their equivalents), or in one or more combinations thereof. Embodiments of the subject matter described in this specification may be implemented as one or more computer programs, for example, one or more modules of computer program instructions encoded on a tangible, non-transitory storage medium for execution by a data processing device or for controlling the operation of a data processing device. The computer storage medium may be a machine-readable storage device, a machine-readable storage substrate, a random or serial access memory device, or a combination thereof. Alternatively or additionally, program instructions may be encoded on artificially generated propagated signals (e.g., machine-generated electrical, optical, or electromagnetic signals) generated to encode information for transmission to a suitable receiver device for execution by the data processing device.

[0177] The term "data processing device" refers to data processing hardware and includes all kinds of devices, apparatuses, and machines for processing data, such as programmable processors, computers, or multiple processors or computers. The device may also be or further include special-purpose logic circuit systems, such as FPGAs (Field-Programmable Gate Arrays) or ASICs (Application-Specific Integrated Circuits). In addition to hardware, the device may optionally include code that creates an execution environment for computer programs, such as code constituting processor firmware, protocol stacks, database management systems, operating systems, or combinations thereof.

[0178] A computer program (which may also be referred to or described as a program, software, software application, app, module, software module, script, or code) can be written in any form of programming language, including compiled or interpreted languages ​​or declarative or procedural languages, and can be deployed in any form, including as a standalone program or as a module, component, subroutine, or other unit suitable for use in a computing environment. A program may, but does not necessarily, correspond to a file in a file system. A program may be stored as part 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 coordinating files (e.g., files storing one or more modules, subroutines, or code portions). A computer program may be deployed to execute on a single computer or on multiple computers located at a site or distributed across multiple sites and interconnected via a data communication network.

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

[0180] Similarly, in this specification, the term "engine" is used broadly to refer to a software-based system, subsystem, or process 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 may be installed and run on the same one or more computers.

[0181] The processes and logic flows described in this specification can be executed by one or more programmable computers, which execute one or more computer programs to perform functions by manipulating input data and generating output. The processes and logic flows can also be executed by a dedicated logic circuit system, such as an FPGA or ASIC, or by a combination of a dedicated logic circuit system and one or more programmed computers.

[0182] A computer suitable for executing computer programs may be based on a general-purpose microprocessor or a special-purpose microprocessor or both, or any other type of central processing unit. Generally, the central processing unit receives instructions and data from read-only memory or random access memory or both. The basic components of a computer are the central processing unit for making or executing instructions and one or more memory devices for storing instructions and data. The central processing unit and memory may be supplemented by or incorporated into a special-purpose logic circuit system. Generally, a computer will also include one or more mass storage devices for storing data, such as magnetic disks, magneto-optical disks, or optical disks, or operatively coupled to receive data from or transfer data to one or more mass storage devices or both. However, a computer does not need to have such devices. Furthermore, a computer may be embedded in another device, such as a mobile phone, 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, such as a universal serial bus (USB) flash drive, to name a few.

[0183] Computer-readable media suitable for storing computer program instructions and data include all forms of non-volatile memory, media, and memory devices, including, for 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 ROMs and DVD-ROMs.

[0184] To provide interaction with the user, embodiments of the subject matter described in this specification can be implemented on a computer having a display device for displaying information to the user (e.g., a CRT (cathode ray tube) or LCD (liquid crystal display) monitor) and a keyboard and pointing device (e.g., a mouse or trackball) through which the user can provide input to the computer. Other types of devices can also be used to provide interaction with the user; for example, feedback provided to the user can be any form of sensory feedback, such as visual feedback, auditory feedback, or tactile feedback; and input from the user can be received in any form, including sound, speech, or tactile input. Furthermore, the computer can interact with the user by sending documents to and receiving documents from the device used by the user; for example, by sending a webpage to a web browser in response to a request received from a web browser on the user's device. Additionally, the computer can interact with the user by sending text messages or other forms of messages to a personal device (e.g., a smartphone running a messaging application) and receiving responsive messages from the user in response.

[0185] Data processing devices used to implement machine learning models may also include, for example, dedicated hardware accelerator units for handling the common and computationally intensive parts of machine learning training or production (i.e., inference, workloads).

[0186] Machine learning models can be implemented and deployed using machine learning frameworks (such as TensorFlow or Jax).

[0187] Embodiments of the subject matter described in this specification can be implemented in a computing system that includes back-end components (e.g., as a data server), or middleware components (e.g., an application server), or front-end components (e.g., a client computer having a graphical user interface, web browser, or app that a user can interact with through an implementation of the subject matter described in this specification), or any combination of one or more such back-end components, middleware components, or front-end components. The components of the system can be interconnected via digital data communication (e.g., a communication network) of any form or medium. Examples of communication networks include local area networks (LANs) and wide area networks (WANs), such as the Internet.

[0188] A computing system may include clients and servers. Clients and servers are typically geographically separated and interact via a communication network. The client-server relationship is established by computer programs running on respective computers and having a client-server relationship with each other. In some embodiments, the server transmits data (e.g., HTML pages) to a user device, for example, for the purpose of displaying data to a user interacting with the device acting as a client and receiving user input from that user. Data generated at the user device, such as the result of user interaction, may be received at the server from the device.

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

[0190] Similarly, although operations are depicted in the accompanying drawings and described in a specific order in the claims, this should not be construed as requiring such operations to be performed in the specific order shown or in sequential order, or requiring all shown operations to be performed to achieve the desired result. In some contexts, multitasking and parallel processing may be advantageous. Furthermore, the separation of various system modules and components in the above embodiments should not be construed as requiring such separation in all embodiments, and it should be understood that the described program components and systems can generally be integrated together in a single software product or packaged into multiple software products.

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

Claims

1. A method executed by one or more computers, the method comprising: Initialize the representation of the data items; Receive conditional inputs that characterize one or more desired attributes of the data item; The representation of the data item is updated to generate a final representation of the data item, the update comprising: in each of a plurality of sampling iterations: using a diffusion neural network to process a diffusion input for the sampling iteration, including the representation of the data item and the representation of the conditional input, to generate a denoised output for the sampling iteration; The process includes updating the representation of the data item using the denoised output; generating a reward input from the final representation after the last iteration of the plurality of sampling iterations; applying a reward function to the reward input to generate a reward score for the final data item that measures the quality of the final data item; and training the diffusion neural network using a loss function that includes a first term measuring the reward score of the final data item, the training including backpropagating the gradient of the first term through the reward function across a subset of the sampling iterations.

2. The method as described in claim 1, wherein, The diffusion neural network has been pre-trained with respect to the diffusion model training objective that does not use the reward function.

3. The method as described in any of the preceding claims, wherein, The reward input includes the condition input.

4. The method as described in any of the preceding claims, wherein, The subset of the sampling iterations includes all iterations of the plurality of sampling iterations.

5. The method according to any one of claims 1 to 3, wherein, The subset of the sampling iterations is a proper subset that includes fewer than all iterations in the plurality of iterations.

6. The method of claim 5, wherein, The subset of the sampling iterations does not include one or more of the earliest iterations among the plurality of sampling iterations.

7. The method of claim 5 or claim 6, wherein, The subset of the sampling iterations includes only the latest iterations among the plurality of sampling iterations.

8. The method as described in any of the preceding claims, wherein, Backpropagation of gradients throughout a specific sampling iteration includes: retrieving the representation of the final data item up to the specific sampling iteration from memory; recalculating the intermediate activations of the diffuse neural network for the specific sampling iteration using the retrieved representation; and using the intermediate activations to compute the gradient of the first item for the specific sampling iteration.

9. The method as described in any of the preceding claims, wherein, The diffusion neural network includes a first network parameter set and a second network parameter set, wherein training the diffusion neural network includes updating the first network parameter set while keeping the second network parameter set fixed.

10. The method of claim 9, wherein, Updating the first network parameter set while keeping the second network parameter set fixed includes: keeping the second network parameter set fixed at pre-trained values, which are determined by training instances of the diffuse neural network excluding the first network parameter set with respect to a score matching target.

11. The method of any of the preceding claims, wherein the update further comprises: In each of the plurality of sampling iterations: the second diffuse input of the representation of the data item but excluding the conditional input used for the sampling iteration is processed by the diffuse neural network to generate an unconditionally denoised output for the sampling iteration; And use the denoised output and the conditional denoised output to update the representation of the data item.

12. The method of claim 11, wherein, Updating the representation of the data item using the denoised output and the conditional denoised output includes: generating a final denoised output by combining the denoised output and the unconditional denoised output according to guiding weights used for the sampling iteration; and using the final denoised output to update the representation of the data item.

13. The method of claim 11 or claim 12, wherein, The second diffusion input includes the representation of the data item and a predetermined representation indicating unconditional sampling.

14. The method as claimed in any of the preceding claims, wherein, The representation of the data item is a latent representation in the latent space.

15. The method of claim 14, wherein, Generating a reward input from the final representation after the last sampling iteration of the plurality of sampling iterations includes: processing the reward input using a decoder neural network to generate the final data item; and including the final data item in the reward input, wherein the training includes backpropagating the gradient of the first item through the reward function, through the decoder neural network, and through a subset of the sampling iterations.

16. The method as claimed in any of the preceding claims, wherein, The reward function includes multiple reward models, and the reward score is a combination of corresponding initial reward scores generated by each of the multiple reward models by processing at least a portion of the reward input.

17. The method as claimed in any of the preceding claims, wherein, The data item is an image.

18. The method as claimed in any of the preceding claims, wherein, The data item is audio data representing the audio signal.

19. The method as claimed in any of the preceding claims, wherein, The data item is a video that includes multiple video frames.

20. The method as claimed in any of the preceding claims, wherein, The conditional input includes text prompts.

21. The method as claimed in any of the preceding claims, wherein, The conditional input includes an image.

22. The method as described in any of the preceding claims, wherein, The conditional input includes audio data representing an audio signal.

23. The method according to any of the preceding claims, further comprising: Each noise iteration in one or more noise iterations: sampled noise; The noise is applied to the final representation to generate a noisy representation; The diffuse neural network is used to process the input including the noisy representation to generate a denoised output for the noise iteration; the denoised output for the noise iteration is used to update the noisy representation to generate an updated noisy representation as an estimate of the final representation; A new reward input is generated from the updated noisy representation; and the reward function is applied to the new reward input to generate a new reward score; wherein the loss function includes a second term measuring the new reward score of the one or more noisy iterations.

24. The method of claim 23, wherein, There are multiple noisy iterations, and the second term measures the average of the new reward score.

25. The method of claim 23 or claim 24, wherein the training comprises backpropagating the gradient of the second term through the reward function and backpropagating into the noise iteration, but not backpropagating through any of the sampling iterations.

26. A system comprising one or more computers and one or more storage devices storing instructions that, when executed by the one or more computers, are operable to cause the one or more computers to perform operations as described in any of the preceding claims.

27. One or more computer storage media, the computer storage media being encoded with instructions that, when executed by one or more computers, cause the one or more computers to perform the operation of the corresponding method as described in any one of claims 1-25.

Images