Expectation-maximization distillation for diffusion models

Abstract

Methods, systems, and apparatus, including computer programs encoded on computer storage media, for training a student generator neural network to generate output data items using a teacher diffusion model based on Expectation-Maximization distillation. The generator neural network can be configured through training to generate any of a variety of output data items, e.g., image data items, audio data items, or video data items.

Description

[0001]Attorney Docket No.: 45288-0475WO1 EXPECTATION-MAXIMIZATION DISTILLATION FOR DIFFUSION MODELS CROSS-REFERENCE TO RELATED APPLICATIONS This application claims priority to U.S. Provisional Application No. 63 / 650,821, filed on May 22, 2024. The disclosure of the prior application is considered part of and is incorporated by reference in its entirety in the disclosure of this application. BACKGROUND This specification relates to using neural networks to generate data. Neural networks are machine learning models that employ one or more layers of nonlinear units to predict an output for a received input. Some neural networks include one or more hidden layers in addition to an output layer. The output of each hidden layer is used as input to the next layer in the network, i.e., the next hidden layer or the output layer. Each layer of the network generates an output from a received input in accordance with current values of a respective set of parameters. SUMMARY This specification describes a system implemented as computer programs on one or more computers in one or more locations that trains a generator neural network to generate output data items using a teacher diffusion model. In implementations the (trained) teacher diffusion model is configured to generate the output data item as a result of performing a plurality of reverse diffusion steps. A student generator neural network is trained to generate output data items that are generally similar, e.g. from a similar distribution. For example the student generator neural network (g(z)) can be configured to process a latent representation (later, z) to generate an output data item (later, x). In particular the student generator neural network can generate clean output data items in a singlestep (e.g. x0 = g(z) when z is sampled from a noise distribution such as ^^0, ^^^.An implementation of the technique determines the latent representation of an output data item. Some implementations of the technique also determine a noise term (later, ^), e.g. by sampling from a noise distribution, although, as shown in examples later, this is not essential. Attorney Docket No.: 45288-0475WO1 As some examples, the output data item can be a still or moving image or an audio signal. In some implementations the output data item is itself a (different) latent space representation that represents an output space, e.g. a latent space representation of an image or audio signal that can be decoded into the image or audio signal. Implementations of the technique perform a plurality of stochastic sampling steps using the teacher diffusion model to determine an adjusted output data item and an adjusted latent representation of the output data item based on the latent representation and the noise term. As an example, the latent representation can be sampled from an e.g. predetermined, distribution (later, p(z)). In some implementations the output data item can be obtained by processing the latent representation using the student generator neural network (e.g. x ~ p(z)). In some implementations the output data item is obtained from the For example, the output data item (in particular a noisy version of the output data item) can be obtained by adding the noise term to a sample from the student generator neural network. The level of added noise can depend on a noise schedule used by the teacher diffusion model or, equivalently, on a diffusion time step that defines the level of added noise. In implementations the output data item can be a noisy version of the output data item (xt), e.g. obtained by adding noise to the output data item (e.g., but not essentially, according to a noise schedule of a diffusion process modelled by the teacher diffusion model). A distribution of the sampled the latent representation and output data item pair ((x,z)) can be adjusted by the stochastic sampling steps towards a target distribution (q(x)) defined by the teacher diffusion model. In implementations this can be done by an iterative, e.g. MCMC (Markov Chain Monte Carlo) process, in particular one that successively updates the output data item and the latent representation, e.g. in a series of one-step updates, to obtain the adjusted output data item (e.g. ^^^) and the adjusted latent representation (e.g. ^^). A convenient one-step update uses Langevin dynamics to obtain an adjusted output data item at a next iteration from the (adjusted) output data item at a current iteration; and similarly to obtain an adjusted latent representation at a next iteration from the (adjusted) latent representation at a current iteration. In general using Langevin dynamics involves using a so-called score function (a gradient of the log likelihood of the pair, respectively with respect to the current output data item, or with Attorney Docket No.: 45288-0475WO1 respect the current latent representation) to update the current output data item, or the current latent representation (and adding noise). Some implementations of the technique maintain a score neural network to predict a score function for the current output data item from the current output data item (^;^^), i.e. a gradient of the log likelihood of the current output data item with respect to the current outputdata item, e.g. ∇^ log ^^^^. This can be combined with a score function for the current outputdata item obtained from the teacher diffusion model (∇^ log ^^^^) (and with a score from thestudent generator neural network, e.g. ∇^^^^^) to obtain score functions for the pair with respect to the current output data item, and with respect the current latent representation. In some implementations a stability of this process can be improved by a reparameterization of the sampled the latent representation and output data item pair ((x,z)). More specifically in some implementations a deterministic transform is used to transform between a pair comprising the latent representation of the output data item (z) and the noise term (^) and a pair comprising the latent representation of the output data item (z) and the output dataitem (x) (e.g. ^ = ^^^^^ + ^^ or ^^ = ^^^^^^ + ^^^).For example in some implementations the stochastic sampling steps obtain the adjusted output data item (e.g. ^^^) and the adjusted latent representation (e.g. ^^) by employing a reparameterization, more specifically by converting the sampled latent representation and output data item pair ((x,z)) to a pair comprising the sampled latent representation and the noise term(^^, ^^) using the deterministic transform. The stochastic sampling steps can then be performedin ^^, ^^ space and then transformed back to (x,z) space.A first student generator input that comprises the adjusted latent representation can be processed using the student generator neural network to generate a first student estimate of the output data item. The student generator neural network can then be trained, updating values of learnable parameters e.g. weights, of the student generator neural network, based at least in part on optimizing a generator objective function. In implementations the generator objective function evaluates a difference between the adjusted output data item and the first student estimate of the output data item. In general, one innovative aspect of the subject matter described in this specification can be embodied in a method of training a student generator neural network to generate an output data item using a teacher diffusion model, wherein the teacher diffusion model is configured to Attorney Docket No.: 45288-0475WO1 generate the output data item as a result of performing a plurality of reverse diffusion steps, wherein the method comprises: determining a latent representation of an output data item, and a noise term; performing a plurality of stochastic sampling steps using the teacher diffusion model to determine an adjusted output data item and an adjusted latent representation of the output data item based on the latent representation and the noise term; processing a first student generator input that comprises the adjusted latent representation using the student generator neural network to generate a first student estimate of the output data item; and training the student generator neural network to update values of parameters of the student generator neural network based at least in part on optimizing a generator objective function that evaluates a difference between the adjusted output data item and the first student estimate of the output data item. These and other embodiments can each optionally include one or more of the following features. Performing the plurality of stochastic sampling steps may comprise, at each stochastic sampling step: sampling a perturbation for the stochastic sampling step based on a predetermined distribution; generating a teacher score for the output data item by processing a teacher diffusion model input that comprises a noisy representation of the output data item; and determining an updated noise term and an updated latent representation by applying a one-step update based on the perturbation and the teacher score. Determining the updated noise term and the updated latent representation may comprise applying a reparameterization sampling technique that uses the teacher score. Determining the adjusted output data item and the adjusted latent representation of the output data item may comprise, after a last stochastic sampling step: determining the adjusted output data item and the adjusted latent representation based on a final updated noise term and a final updated latent representation generated in the last stochastic sampling step. Determining the adjusted output data item may comprise: determining an initial adjusted output data item based on the final updated latent representation; and removing an accumulated noise term from the initial adjusted output data item. Determining the adjusted latent representation may comprise: using the final updated latent representation as the adjusted latent representation. Attorney Docket No.: 45288-0475WO1 The latent representation may comprise a plurality of latent variables, and wherein determining the latent representation comprises sampling, for each latent variable, a value from a predetermined distribution. Training the student generator neural network may comprise: computing gradients of the generator objective function with respect to the parameters of the student generator neural network; and applying a stochastic gradient ascent technique based on the gradients. The gradients may be weighted by weights computed in accordance with a time- dependent weighting function. The student generator neural network may have a neural network architecture that is dependent on a neural network architecture of the teacher diffusion model. The output data item may comprise one of: image data, audio data, or video data. The method may further comprise: processing a second student generator input that comprises the latent representation using the student generator neural network to generate a second student estimate of the output data item; processing, using a score neural network, a score neural network input that comprises (i) a or the noisy representation of the output data item and (ii) data that identifies a time step to generate a student score estimate of the output data item; and training the score neural network to update values of parameters of the score neural network based at least in part on optimizing a score objective function that evaluates a difference between the student score estimate and a target score that is computed based on the second student estimate of the output data item. Training the score neural network to update the values of the parameters of the score neural network may comprise: computing gradients of the score objective function with respect to the parameters of the score neural network; and applying a stochastic gradient descent technique based on the gradients. The method may further comprise using the score neural network to train the student generator neural network. The student generator neural network may be configured to generate the output data item in a single forward pass through the student generator neural network. The student generator neural network may be configured to generate the output data by performing one or more processing iterations and wherein a total number of processing iterations is less than a total number of reverse diffusion steps. Attorney Docket No.: 45288-0475WO1 Another innovative aspect of the subject matter described in this specification can be embodied in a method comprising: receiving a request for an output data item; and generating the output data item using the student generator neural network that has been trained using the method of the above aspect. A further innovative aspect of the subject matter described in this specification can be embodied in a method comprising: receiving a request for an output data item; and generating the output data item using a student generator neural network that has been trained using a teacher diffusion model, wherein the teacher diffusion model is configured to generate the output data item as a result of performing a plurality of reverse diffusion steps, the student generator neural network having parameters reflecting optimization of a generator objective function that evaluated a difference between an adjusted training data item, and a corresponding student estimate of the training data item derived from an adjusted latent representation using the student generator neural network, wherein the adjusted training data item and the adjusted latent representation have been derived from stochastic sampling steps performed using the teacher diffusion model on a latent representation of a training data item and a noise term. Other embodiments of these aspects include corresponding computer systems, apparatus, and computer programs recorded on one or more computer storage devices, each configured to perform the actions of the methods. A system of one or more computers can be configured to perform particular operations or actions by virtue of software, firmware, hardware, or any combination thereof installed on the system that in operation may cause the system to perform the actions. One or more computer programs can be configured to perform particular operations or actions by virtue of including instructions that, when executed by data processing apparatus, cause the apparatus to perform the actions. The subject matter described in this specification can be implemented in particular embodiments so as to realize one or more of the following advantages. Diffusion models can generate output data items that have a high quality but typically consume a significant amount of computational resources due to the execution of many processing iterations (reverse diffusion steps) to generate each output data item. This makes existing diffusion models impractical to deploy in resource-constrained environments. This also increases the latency of diffusion models when generating output data items. Attorney Docket No.: 45288-0475WO1 The specification describes how a distilled student generator neural network that is easier to deploy than a teacher diffusion model can be trained using the teacher diffusion model that has already been trained. Once trained, the student generator neural network can more easily be deployed to perform on-device inference in a resource-constrained environment such as a mobile or edge device. The student generator neural network requires fewer computational resources such as processor time and power usage to generate an output data item by executing a smaller number of processing iterations than the teacher diffusion model. This improved deployment feasibility is a technical effect of the reduced computational footprint of the student generator neural network, which requires, for example, fewer floating- point operations (FLOPs) and less memory per generated output data item compared to the teacher diffusion model. The reduced number of processing iterations directly translates to a technical effect of lower generation latency, enabling applications such as real-time interactive content generation or closed-loop control systems where timely outputs are critical. For example implementations of the techniques described herein enable an output data item to be generated in a single step or pass through the student generator neural network, rather than the tens of hundreds of reverse diffusion steps that are typically needed. In implementations this is done by distilling the teacher diffusion model into the student generator neural network, which is conditioned on the latent representation (z, which in implementations, does not depend on time. In implementations the latent representation is distilled from the reverse diffusion trajectory modelled by the teacher diffusion model at multiple different time steps or noise levels. This can be done using an Expectation Maximization (EM) approach but in practice sampling the latent representation can be difficult. Implementations of the described techniques address this problem by obtaining an approximate sample of x and z using the student generator neural network and then incrementally (iteratively) improving this towards the desired joint distribution (i.e. as defined by the teacher diffusion model), in particular using stochastic sampling. Some particular implementations of the system described in this specification combine a short-run Markov Chain Monte Carlo (MCMC) sampling scheme with techniques that include reparameterized sampling and noise cancellation to improve the stability and therefore, computational resource efficiency of the distillation training process. Attorney Docket No.: 45288-0475WO1 The reparameterized sampling technique specifically contributes to a technical effect of stabilizing the MCMC updates by reducing the variance of gradient estimates within the sampling process, allowing for more efficient exploration of the target distribution. Noise cancellation further enhances training stability by reducing the variance of the targets used for updating the student generator, leading to a more robust convergence of the distillation process. These features collectively address a technical problem of efficiently and stably training a high- fidelity student model from a complex teacher diffusion model. Leveraging Expectation-Maximization distillation, the system overcomes common limitations faced by existing distillation methods that similarly use a teacher diffusion model, such as performance degradation with very few sampling steps, reliance on training data access, or mode-seeking optimization that may fail to capture the full distribution. The Expectation-Maximization framework, particularly through its MCMC-based E-step, provides a technical effect of enabling the student model to better approximate the full data distribution learned by the teacher model, rather than merely fitting to modes. This results in generated samples with higher diversity and better coverage of the underlying data manifold, which is a measurable improvement in the technical quality of the generative model. Once trained using the described techniques, the student generator neural network can thus generate output data items that are similar in quality, e.g., fidelity, coherence, and realism, to output data items generated by the teacher diffusion model. For example the student generator neural network can generate output data items that have superior quality, e.g., in terms of a Fréchet inception distance (FID) metric or another perceptual metric when the output data items are images, than output data items generated by student generator neural networks trained using existing distillation techniques. This is despite using fewer computational resources than the teacher diffusion model and, in some cases, using comparable or fewer computational resources than alternative generator neural networks trained using existing distillation methods. The details of one or more embodiments of the subject matter described in this specification are set forth in the accompanying drawings and the description below. Other features, aspects, and advantages of the subject matter will become apparent from the description, the drawings, and the claims. Attorney Docket No.: 45288-0475WO1 BRIEF DESCRIPTION OF THE DRAWINGS FIG. 1 shows an example training system and an example inference system. FIG. 2 is a flow diagram of an example process for training a student generator neural network to generate an output data item using a teacher diffusion model. FIG. 3 is a flow diagram of sub-steps of one of the steps of the process of FIG. 2. FIG. 4 is a flow diagram of an example process for training a score neural network. FIG. 5 is a flow diagram of an example process for generating an output data item. FIGS. 6A-C show examples that illustrate the effect of stochastic sampling steps. FIG. 7 shows examples of images generated by a teacher diffusion model and a single- step student generator neural network trained as described herein. Like reference numbers and designations in the various drawings indicate like elements. DETAILED DESCRIPTION FIG. 1 shows an example training system 100 and an example inference system 150. The training system 100 and the inference system 150 are examples 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 training system 100 implements an Expectation-Maximization distillation training process that uses a diffusion model. Some implementations of the training system 100 also include a score neural network 130 as described later. The training system 100 can be used to train a generator neural network for deployment at the inference system 150 to compute inference. That is, once trained, the inference system 150 can use the generator neural network to generate output data items 104 in response to requests 102 for output data items. Because of the computationally intensive nature of the training process, the training system 100 can be implemented on a distributed computing system, e.g., a datacenter, having multiple, e.g., tens or hundreds of, computing devices. In contrast, the inference system 150 can be implemented on a computing device with limited computational resources, power resources, or both. In some implementations, the inference system 150 can be implemented on a smartphone or a tablet computer. Attorney Docket No.: 45288-0475WO1 The generator neural network that is being trained by the training system 100 will be referred to in this specification as a “student generator neural network 110.” The student generator neural network 110 has parameters 116 that have values that are updated during the distillation training process. The diffusion model that is being used by the training system 100 to train the student generator neural network 110 will be referred to in this specification as a “teacher diffusion model 120.” The teacher diffusion model 120 has parameters 126 that have values that have been learned as a result of the previous training, and that generally will not be updated during the distillation training process. In some cases, the student generator neural network 110 generates an output data item in an unconditioned manner, i.e., without conditioning on any conditioning input. In these cases, the output data items 104 generated by the student generator neural network 110 approximate samples of a distribution of training data items that were using during the distillation training process. In other cases, the student generator neural network 110 generates an output data item conditioned on a conditioning input. For example, the inference system 150 can receive a conditioning input as part of, or in association with, the request 102 and generate an output data item 104 conditioned on the conditioning input. The teacher diffusion model 120 and thus, the student generator neural network 110, can be configured to generate any of a variety of output data items, either unconditioned or conditioned on a conditioning input. Generally, the conditioning input characterizes one or more desired properties for the output data item, i.e., characterizes one or more properties that the output data item generated by the teacher diffusion model 120 (or the student generator neural network 110) should have. Conditioning can be performed by any suitable technique, e.g. classifier guided conditioning or classifier free conditioning. The teacher diffusion model 120 can have been configured through previous training to be conditioned on any of a variety of conditioning inputs. During the distillation training process, the training system 100 can likewise train the student generator neural network 110 to generate output data items conditioned on any of the variety of conditioning inputs. Some examples of data items and conditioning inputs that can be processed by the teacher diffusion model 120 now follow. By virtue of the distillation training process, the Attorney Docket No.: 45288-0475WO1 student generator neural network 110 is configured to generate the same types of output data items as the teacher diffusion model 120, either unconditioned or conditioned on a conditioning input. Thus, these examples are also applicable to the student generator neural network 110. For example, the output data item can include 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. As a particular example, the output data item can be an output audio example that includes a sample of an audio wave at each of a sequence of output time steps that span a specified time window, where the output time steps are arranged at regular intervals within the specified time window. The audio sample at a given output time step can be an amplitude value of the audio wave or an amplitude value that has been compressed, companded, or both. For example, the audio sample can be a raw amplitude value or a mu-law companded representation of the amplitude value. In this example, the conditioning input can be text or features of text that the audio should represent, i.e., so that the inference system 150 implementing the student generator neural network 110 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 can identify a desired speaker for the audio, i.e., so that the teacher diffusion model 120 generates audio data that represents speech by the desired speaker. As another example, the conditioning input can characterize properties of a song or other piece of music, e.g., lyrics, genre, and so on, so that the teacher diffusion model 120 generates a piece of music that has the properties characterized by the conditioning input. As another example, the conditioning input can specify a classification for the audio data into a class from a set of possible classes, so that the teacher diffusion model 120 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 teacher diffusion model 120 generates audio that is emitted by the corresponding class, or types of animals, i.e., so that the system generates audio that represents noises generated by the corresponding animal, and so on. As another example, the output data item can be an image (a still or moving image), e.g., such that the teacher diffusion model 120 can perform conditional image generation by Attorney Docket No.: 45288-0475WO1 generating the values, e.g. color or intensity values of the pixels of the image. In general the conditioning input can specify one or more characteristics for the image. In this example, the conditioning input can be a sequence of text and the output data item can be an image that describes the text, i.e., the conditioning input can be a caption for the output image. As yet another example, the conditioning input 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 example, the conditioning input can specify an object class from a plurality of object classes to which an object depicted in the output image should belong. As another example, the conditioning input can specify (e.g. be) one or more images. For example, the conditioning input can specify an image at a first resolution and the output data item can include the image at a second, higher resolution. As another example, the conditioning input can specify an image and the output data item can comprise a compressed, de-noised, restored, enhanced, stylized, or otherwise edited version of the image. As another example, the conditioning input can specify an image including a target entity for detection, e.g. a tumor, and the output data item can comprise the image without the target entity, e.g., to facilitate detection of the target entity by comparing the images. As another example, the conditioning input can be a segmentation that assigns each of a plurality of pixels of the output image to a category from a set of categories, e.g., that assigns to each pixel a respective one of the categories. As yet another example, the conditioning input can be a different type of structured input, e.g., a mesh or a graph that specifies properties of the image to be generated. More generally, the conditioning input can include one or more different types of inputs of one or more different modalities, e.g., only text, only one or more images, both text and one or more images, and so on. As another example, the output data item can be a video. Again the conditioning input can specify one or more characteristics for the video. As a particular example, the conditioning input can include text and the output data item can be a video described by the text. Attorney Docket No.: 45288-0475WO1 As yet another example, the conditioning input can include one or more images and the output data item can be a video that completes the one or images, e.g., video starting from the one or more images. As yet another example, the conditioning input can include one or more reference images and the output data item can be a video that includes a temporal sequence of frames that temporally encompass the one or more reference image, such that the reference images are the middle frames of the video. In some applications, the output data item can be used in a control task to control an action of a mechanical agent acting in a real-world environment to perform a mechanical task. For example, the output data item can be processed by a policy neural network of the agent to select one or more actions to be performed by the agent as part of the task. The agent may then perform the one or more actions. The output data item (e.g., image) can, for example, characterize a state of the real-world environment that is predicted to be obtained by the agent performing the one or more actions. As another example the output data item can specify a movement of a part of the mechanical agent, e.g. a trajectory of an arm or movement of an end effector of a robot. The conditioning input can, e.g., specify a state of the real-world environment and the one or more actions. As another example the conditioning input can specify a state of the real-world environment and the output data item can be used to select one or more actions to be performed by the mechanical agent to perform a task (i.e. the student generator neural network 110 can represent an action selection policy). The teacher diffusion model 120 has been configured through training to generate an output data item across multiple reverse diffusion steps by performing a reverse diffusion process. In particular, the teacher diffusion model 120 can perform the reverse diffusion process in a latent space, e.g., in a latent space that is different, e.g., lower-dimensional, than an output space, e.g., pixel space when the output data items are images. That is, in some implementations the data items (“representations”) operated on by the teacher diffusion model 120 are latent representations (embeddings) and the values in the representations are learned, latent values, e.g., rather than color values when the output data items are images. Note that these latent representations are different to the latent representation, z, referred to later. Attorney Docket No.: 45288-0475WO1 In general a “latent representation” as used in this specification is a vector of variables (“latent variables”) having a pre-determined dimensionality, i.e., having a pre-determined number of variables. Each variable can have a floating point value or another value. The space of possible vectors having the pre-determined dimensionality is referred to as the “latent space.” In these implementations, during training, the teacher diffusion model 120 can be associated with an encoder neural network to encode training data items into the latent space (of the output space). After training, a decoder neural network can receive an input that includes a latent representation and decode the latent representation to reconstruct an output data item in the output space, e.g. image space. The encoder neural network can have been pre-trained jointly with the decoder neural network in an auto-encoder framework, to encode training data items from the output space into a latent representation in a latent space and decode data items from the latent space into the output space. The teacher diffusion model 120 can be any appropriate diffusion neural network that has been trained, e.g., by the training system 100 or another training system, on a set of training data items using a denoising score matching objective to, at any given reverse diffusion step, process a teacher diffusion model input for the reverse diffusion step to generate a teacher diffusion model output for the reverse diffusion step that specifies a teacher score (as of the reverse diffusion step) for the output data item. In score-based diffusion models, a “score” refers to the gradients of the log probability density function of an output data item (e.g. of x, with respect to x). In implementations a score can be represented as a data structure, e.g., a tensor, that has the same dimension as the output data item and encodes the direction in which the log probability density increases most rapidly. Thus, a score may also be referred to as a “score tensor.” For example, a teacher score may also be referred to as a “teacher score tensor.” The teacher diffusion model input includes an intermediate (noisy) representation of the output data item as of the reverse diffusion step. The teacher diffusion model input can also include data that identifies a time step, e.g., an embedding of the time step, which, in turn, identifies the reverse diffusion step in the reverse diffusion process. The teacher diffusion model input can also include a conditioning input, or data derived from the conditioning input, e.g., an embedding of the conditioning input. Attorney Docket No.: 45288-0475WO1 For the very first reverse diffusion step, the intermediate representation is an initial representation that is generated by, e.g., randomly sampling each value in a latent representation of the output data item from a corresponding noise distribution, e.g., a Gaussian distribution, or a different noise distribution. For any subsequent reverse diffusion step, the intermediate representation is the updated intermediate representation that has been generated in the immediately preceding reverse diffusion step. Then, a stochastic differential equation (SDE) solver can be applied to the teacher score to update, e.g., de-noise, the intermediate representation of the output data item to generate the updated intermediate representation of the output data item. The SDE solver can be any appropriate SDE solver. Note that as described herein the teacher diffusion model does not need to be used to generate a denoised data item; in implementation the scores from the teacher model are used for the described distillation process. For example, it can be a SDE solver that employs the Euler- Maruyama method to simulate the reverse-time SDE. As another example, it can be a SDE solver that employs the Jolicoeur-Martineau method to simulate the reverse-time SDE. As another example, it can be a stochastic Runge-Kutta method. As another example, it can be the diffusion probabilistic model (DPM) solver as described in Cheng Lu, et al. Dpm-solver: A fast ode solver for diffusion probabilistic model sampling in around 10 steps. Advances in Neural Information Processing Systems, 35:5775–5787, 2022. The teacher diffusion model 120 can have any appropriate architecture that allows it to perform its described functions, i.e., to generate an output data item across multiple reverse diffusion steps. For example, when the output data item is an audio signal or an image, the teacher diffusion model can be a convolutional neural network, e.g., a U-Net or other convolutional architecture. As another example, the teacher diffusion model can be a Transformer neural network, e.g., a Diffusion Transformer (DiT) architecture, a Universal Vision Transformer (UViT) architecture, or another attention-based architecture. The teacher diffusion model 120 can be conditioned on the conditioning input in any of a variety of ways. Attorney Docket No.: 45288-0475WO1 As one example, the teacher diffusion model 120 can use an encoder neural network to generate one or more embeddings that represent the conditioning input and the teacher diffusion model 120 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 is text, the teacher diffusion model 120 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. When the conditioning input is an image, the teacher diffusion model 120 can use an image encoder neural network, e.g., a convolutional neural network or a vision Transformer neural network, to generate a set of embeddings that represent the image. When the conditioning input is audio, the teacher diffusion model 120 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 is a scalar value, the teacher diffusion model 120 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 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 teacher diffusion model 120 can generate one or more initial embeddings for each of the different types of inputs, i.e., using an appropriate encoder neural network as described above, and then process the initial embeddings for all of the different types of inputs using a Transformer encoder neural network to update each of the initial embeddings to generate a set of final embeddings. The one or more cross-attention layers within the teacher diffusion model 120 can then cross-attend into the set of final embeddings. In others of these cases, different cross-attention layers within the teacher diffusion model 120 can cross-attend into embeddings of different types of conditioning inputs. Attorney Docket No.: 45288-0475WO1 In yet others of these cases, the teacher diffusion model 120 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 teacher diffusion model 120 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. Examples of the architectures of the teacher diffusion model 120, as well as examples of the denoising score matching objectives that can be used to train the teacher diffusion model 120 to learn the trained values of its parameters 126, include those described in Jonathan Ho, et al. Denoising diffusion probabilistic models. Advances in neural information processing systems, 33:6840–6851, 2020; Yang Song, et al. Score-based generative modeling through stochastic differential equations. In International Conference on Learning Representations, 2020; Yang Song, et al. Generative modeling by estimating gradients of the data distribution. Advances in neural information processing systems, 32, 2019. By implementing the Expectation-Maximization distillation training process, the training system 100 trains the student generator neural network 110 to generate an output data item by performing a few-step reverse diffusion process with reduced latency and reduced resource consumption, e.g., reduced processing power and memory consumption, compared to the teacher diffusion model 120, which has been configured to generate an output data item across multiple reverse diffusion steps by performing a reverse diffusion process. For example, the student generator neural network 110 can be configured through the distillation training process to generate the output data item by performing a single-step reverse diffusion process that includes a single reverse diffusion step. That is, the student generator neural network 110 can generate the output data item in a single processing iteration (which corresponds to a single forward pass through the student generator neural network). As another example, the student generator neural network 110 can be configured through the distillation training process to generate the output data by performing one or more processing iterations, where a total number of processing iterations is less than a total number of reverse diffusion steps performed by the teacher diffusion model 120 during the reverse diffusion process. Attorney Docket No.: 45288-0475WO1 For example, when configured as a cascaded generator model, the student generator neural network 110 can generate the output data item across a 2-step reverse diffusion process, by generating an initial version of the data item, e.g., a lower-quality or lower-resolution version, in one reverse diffusion step (which corresponds to one processing iteration through a base model), and then mapping the initial version of the data item to the output data item in another reverse diffusion step (which corresponds to another processing iteration through a refinement model, e.g., a super-resolution model). The student generator neural network 110 can have any appropriate architecture that allows it to perform its described functions, i.e., to generate an output data item by performing a few-step (e.g., single-step) reverse diffusion process, and thereby generate the output data item in a more computationally efficient manner than the teacher diffusion model 120. At each reverse diffusion step in the few-step reverse diffusion process, the student generator neural network 110 is configured to process a student generator input for the reverse diffusion step to generate a student generator output that specifies a student estimate of the output data item. The student estimate of the output data item is a predicted reconstruction of the output data item (as of the reverse diffusion step). The student generator input includes an intermediate (noisy) representation of the output data item as of the reverse diffusion step. In the case of a few-step reverse diffusion process, the student generator input can also include data that identifies a time step, e.g., an embedding of the time step, which, in turn, identifies the reverse diffusion step in the reverse diffusion process. The student generator input can also include a conditioning input, or data derived from the conditioning input, e.g., an embedding of the conditioning input. For the very first reverse diffusion step, the intermediate representation is an initial representation that is generated by, e.g., randomly sampling each value in a latent representation of the output data item from a corresponding noise distribution, e.g., a Gaussian distribution, or a different noise distribution. For any subsequent reverse diffusion step (in the case of a few-step reverse diffusion process), the intermediate representation is the updated intermediate representation that has been generated in the immediately preceding reverse diffusion step. When the student generator neural network 110 is configured to generate the output data item by performing a single-step reverse diffusion process, the intermediate representation is the initial representation generated based on sampling from noise distributions. Attorney Docket No.: 45288-0475WO1 In principle, the student generator neural network 110 can include any appropriate neural network layers (e.g., convolutional layers, fully connected layers, attention layers, conditioning layers, etc.) in any appropriate number, connected in any appropriate configuration (e.g., as a linear sequence of layers), although in some implementations, the student generator neural network 110 has a neural network architecture that is dependent on a neural network architecture of the teacher diffusion model 120. For example, the student generator neural network 110 can have the same architecture, i.e., the same number and type of layers, as the teacher diffusion model 120 (but with different learnable parameter values). In some implementations the teacher diffusion model has a time input; when used for the he student generator neural network this can be selected as a hyperparameters (though should not be set at a value, e.g. 1, that corresponds to zero signal to noise level). As another example, the student generator neural network 110 and the teacher diffusion model 120 can both have the same type of repeating blocks, e.g., attention blocks having one or more attention layers, however the student generator neural network 110 can have a smaller number of such repeating blocks than the teacher diffusion model 120. The training system 100 uses the score neural network 130 together with the teacher diffusion model 120 during the Expectation-Maximization distillation training process to train the student generator neural network 110. The score neural network 130 is used to generate student score estimates that are each an approximation of a student score that would be generated by the student generator neural network 110. The score neural network 130 is configured to process a score neural network input to generate a student score estimate. The score neural network input includes an intermediate (noisy) representation of the output data item as of a reverse diffusion step. The score neural network input can also include data that identifies a time step, e.g., an embedding of the time step, which, in turn, identifies the reverse diffusion step in the reverse diffusion process. The score neural network input can also include a conditioning input, or data derived from the conditioning input, e.g., an embedding of the conditioning input. The score neural network 130 can have any appropriate architecture, i.e., can include any appropriate neural network layers (e.g., convolutional layers, fully connected layers, attention layers, conditioning layers, etc.) in any appropriate number, connected in any appropriate Attorney Docket No.: 45288-0475WO1 configuration (e.g., as a linear sequence of layers), that allows it to perform its described functions, i.e., to generate student score estimates that are each an approximation of a student score that would be generated by the student generator neural network 110. The score neural network 130 has parameters 136 that have values that are updated during the distillation training process. In some implementations, the training system 100 alternates between training the score neural network 130 and training the student generator neural network 110 during the distillation training process. FIG. 2 is a flow diagram of an example process 200 for training a student generator neural network to generate an output data item using a teacher diffusion model. For convenience, the process 200 will be described as being performed by a system of one or more computers located in one or more locations. For example, a training system, e.g., the training system 100 depicted in FIG. 1, appropriately programmed in accordance with this specification, can perform the process 200. The teacher diffusion model is configured to generate an output data item across multiple reverse diffusion steps by performing a reverse diffusion process. That is, in the case of the teacher diffusion model, an output data item is generated as a result of performing multiple reverse diffusion steps. In contrast, the student generator neural network is trained to generate an output data item by performing a few-step reverse diffusion process that includes one or more processing iterations, where a total number of one or more processing iterations is less than a total number of the multiple reverse diffusion steps included in the reverse diffusion process. The system can repeatedly perform iterations of the process 200 on different batches of output data items to update the values of the parameters of the student generator neural network. At each iteration, the system obtains, e.g., through sampling, a batch of one or more output data items from a training dataset that includes a large number of output data items. The output data items can include any of a variety of data items, e.g., image data items, audio data items, or video data items. Optionally, the training dataset stores each output data item in association with a conditioning input. Generally, the conditioning input characterizes one or more desired properties for the output data item. The system can continue performing iterations of the process 200 until termination criteria for the training of the student generator neural network have been satisfied, e.g., until the Attorney Docket No.: 45288-0475WO1 values of the parameters have converged, until a threshold amount of wall clock time has elapsed, or until a threshold number of iterations of the process 200 have been performed. In some implementations, the system trains the student generator neural network by repeatedly performing iterations of each of two training processes: the process 200 that corresponds to a student generator neural network training process and the process 400 that corresponds to a score neural network training process (described with reference to FIG. 4). For example, the system can repeatedly alternate between the student generator neural network training process and the score neural network training process, e.g., by performing one iteration of the process 200 followed by one iteration of the process 400, or by performing two or more iterations of the process 200 followed by one iteration of the process 400, or by performing one iteration of the process 200 followed by two or more iterations of the process 400. During the student generator neural network training process, the system updates the values of the parameters of the student generator neural network, and during the score neural network training process, the system updates the values of the parameters of the score neural network. At step 202 the system determines (i) a time step ^, (ii) a latent representation ^ of an output data item, and (iii) a noise term ^. The time step ^ defines a level (or amount) of noise to be added to the output data item at the time step ^, where the level of noise can be computed in accordance with a time-dependent noise schedule ^^^^. For example, the noise schedule can be a linear schedule, a cosine schedule, a cosine-adjusted noise schedule, or a square root schedule. The noise level for ^ can be any value within a predetermined range, e.g., ^0,1^, or a different range. Determining the noise level for ^ can include sampling a value from the predetermined range with some measure of randomness. As an example, ^^ = ^^^^^^ + ^^^ where, e.g. ^ is sampled from ^^0, ^^^. The noiseschedule can define a time-dependent mean scaling factor, ^^, and a time-dependent variancescaling factor, ^^, and an SNR can be defined as ^^ = ^^ / ^^ ^ ^^ . Optionally, but not necessarily,determining one of the mean scaling factor and the variance scaling factor may determine theother, e.g. according to ^^ ^^ + ^^ = 1 (which defines a variance-preserving noise (SNR)schedule). Attorney Docket No.: 45288-0475WO1 The latent representation ^ of the output data item includes a plurality of latent variables. In some implementations, determining the latent representation ^ includes sampling, for each latent variable, a value from a predetermined distribution (p(z)) that corresponds to the latent variable, e.g., a normal distribution, or a different latent distribution, with some measure of randomness. In some implementations (i.e. when distilling a diffusion process that operates in a latent representation of the output space), determining the latent representation ^ includes processing the output data item using the associated encoder neural network to encode the output data item into the latent representation of the output space. Determining the noise term ^ can include sampling from a noise distribution, e.g., a normal distribution, or a different noise distribution, with some measure of randomness. The system performs a stochastic sampling process that includes a plurality of stochastic sampling steps using the teacher diffusion model to determine an adjusted output data item and an adjusted latent representation of the output data item based on the latent representation ^ of an output data item and the noise term ^ (step 204). That is, the output data item and latent representation of the output data item are jointly updated through the stochastic sampling process. For example, the stochastic sampling process can be a Markov Chain Monte Carlo (MCMC) sampling process, and each stochastic sampling step can be a MCMC sampling step. By updating the output data item and the latent representation, the system adjusts them toward a target data distribution of: where ^^!^ and ^"^!^ represents a data distribution generated by the student generator neural network. For example, when the output data items are images, the stochastic sampling process can remove visual artifacts and improve structure in the output data items. Having generated the adjusted output data item and the adjusted latent representation of the output data item as a result of performing the stochastic sampling process, the system uses the student generator neural network to generate a first student estimate of the output data item ! (step 206). Attorney Docket No.: 45288-0475WO1 The student generator neural network is configured to process a first student generator input that includes the adjusted latent representation ^^and, optionally, a conditioning input, or data derived from the conditioning input, e.g., an embedding of the conditioning input, to generate a first student generator output that specifies the first student estimate of the output data item. The first student estimate is a predicted reconstruction of the output data item. The system trains the student generator neural network to update values of parameters of the student generator neural network based at least in part on optimizing a generator objective function (step 208). The generator objective function evaluates, for each output data item in the batch, a difference between the adjusted output data item !^^and the first student estimate of the output data item. The system can do this by computing gradients of the generator objective function with respect to the parameters of the student generator neural network ^#by backpropagation through the parameters θ of the student generator neural network ^#, and then applying a gradient ascent technique, e.g., a stochastic gradient ascent technique, based on the gradients. In some implementations, the gradients are weighted by weighting factors computed in accordance with a time-dependent weighting function %&^^^. For example, the respective gradients of the parameters of the student generator neural network can be computed as∇ ‖^^ ( ^ ^^‖^" ^ ^^^" ^ ^, optionally weighted. Including sampling, this can be expressed as:where . function is flat In some implementations a noise cancellation technique is used when generating the adjusted output data item, as describe later. FIG. 3 is an example flow diagram of sub-steps 302-306 of the step 204 of the process 200. The system can repeatedly perform an iteration of steps 302-306 at each of K stochastic sampling steps. K can be any positive integer. For example, each stochastic sampling step can be a MCMC sampling step. For each pair (x,z) or (^,z), the system determines a perturbation )*for the stochasticsampling step + in ,1, -. (step 302). For example, the perturbation can be (sampled) noise, Attorney Docket No.: 45288-0475WO1 referred to here as Langevin noise; and determining the perturbation ) can include sampling the perturbation ) (noise) from a predetermined distribution, e.g., a normal distribution, or a different noise distribution. In some implementations, the perturbations for different output data items in the same batch are different, while in other implementations, the perturbations for different output data items in the same batch are the same (although the system generally samples different perturbations across different stochastic sampling steps +). For each output data item ! in the batch, the system uses the teacher diffusion model to generate a teacher score for the output data item (step 304). The teacher diffusion model is configured to processes a teacher diffusion model input to generate a teacher diffusion model output that specifies the teacher score for the output data item. The teacher diffusion model input includes an intermediate (noisy) representation of the output data item that is generated based on adding noise to the output data item in accordance with the time step ^ and the noise term ^. The noise and the time step can be the same for all output data items in the same batch. The teacher diffusion model input can also include data that identifies the time step ^, e.g., an embedding of the time step ^. The teacher diffusion model input can also include a conditioning input, or data derived from the conditioning input, e.g., an embedding of the conditioning input. For each output data item ! in the batch, the system determines (i) an updated noise term and (ii) an updated latent representation of an output data item by applying a one-step update based on the perturbation )*and the teacher score (step 306). In some implementations, to determine the updates, the system computes the gradients ofthe logarithm of a joint probability distribution / #^!, ^^ where ! represents the output data item,^ represents the latent representation of the output data item, and θ represents the parameters of the student generator neural network. In some implementations a re-parameterization from (x,z) to (^,z) can be used. As one example, the system can compute the partial gradients of the logarithm of thejoint probability distribution / #^!, ^^ with respect to the output data item ! as:. Attorney Docket No.: 45288-0475WO1 The first term is the teacher score that is generated by using the teacher diffusion model from processing the teacher diffusion model input, as discussed above. The second term, which is generated by using the score neural network, is a student score estimate that is an approximation of a student score that would be generated by the student generator neural network based on the intermediate (noisy) representation of the output data item. The student score estimate is generated by the score neural network from processing a score neural network input that includes the intermediate (noisy) representation of the output data item that is generated based on adding noise to the output data item in accordance with the time step ^ and the noise term ^. The score neural network input can also include data that identifies the time step ^, e.g., an embedding of the time step ^. The score neural network input can also include a conditioning input, or data derived from the conditioning input, e.g., an embedding of the conditioning input. (In the above equation the t subscript is omitted for clarity). The third term is computed based on using the student generator neural network to process the latent representation ^, where ^#represents the student generator neural network having the parameters θ, and ^ and ^ are signal and noise scaling factors (which are tunable parameters of the system) that represent the relative weighting of output data item and noise term, respectively. For example, the system can compute the partial gradients of the logarithm of the jointprobability distribution / #^!, ^^ with respect to the latent representation ^ of the output data itemas: The the teacher score when determining the updated noise term and the updated latent representation. In particular, the system can leverage a reparameterization of the output data item ! and the latent representation ^ of the output data item to address the challenges in joint MCMC sampling, such as slow convergence and complicated step size tuning. Attorney Docket No.: 45288-0475WO1 For example, by using (with t-dependence omitted) to definea deterministic mapping joint probability distribution is given by: Thus, the logarithm of the joint probability as: . joint probability distribution with respect to the latent representation ^ of the output data item as: . one- step update to the noise term ^*01, and determine the updated latent representation ^*01of the output data item by applying a one-step update to the latent representation ^*. The one-step updates are determined based on the computed partial gradients. For example, the one-step updates can be determined as one-step Langevin updates, and the system can determine the updated noise term ^*01as: ^*01 = ^* + 2∇ *3logρ"^^, z^ + 722) .where 2 is the step sampling process), )*is the perturbation for the stochastic sampling step +, and + represents the sampling step. For example, the system can determine the updated latent representation ^*01of the output data item as: ^*01 = ^* + z^ + 722)* .If the current steps 302-306, in otherwords if + : -, the updated noise term and the updated latent representation will then be used asthe noise term and the latent representation to be updated again in the next iteration (the next stochastic sampling step in the stochastic sampling process). Attorney Docket No.: 45288-0475WO1 Alternatively, if the current iteration is the last iteration in the K iterations of steps 302-306, in other words if + = -, the system uses the final updated noise term ^^ and the finalupdated latent representation ^^that are generated in the last iteration (the last stochastic sampling step in the stochastic sampling process) to generate the adjusted output data item !^^and the adjusted latent representation ^^of the output data item. In some implementations, the system uses the final updated latent representation ^^directly as the adjusted latent representation of the output data item, and uses a noise cancellation technique to generate the adjusted output data item based on the final updated noise term ^^. For example, the above describe Langevin updates involve adding (scaled) noise )*at each step. It can be shown that it is valid to track this noise, more specifically, for one or both of the updated noise terms and the updated latent representations, to accumulate the noise added at each step and then afterwards subtract the accumulated noise from the updated representation(s) (and optionally also a proportion of the initial noise term ^). Cancelling the noise after the stochastic sampling process stabilizes the training of the student generator neural network because doing so reduces variance of the gradients without introducing bias. To illustrate with a one stochastic step example, the system can first generate an initial adjusted output data item !^;based on the final updated latent representation as: , and then generate the adjusted an accumulated noise term from the initial adjusted output data item !^;: . FIG. 4 neural network. For convenience, the process 400 will be described as being performed by a system of one or more computers located in one or more locations. For example, a training system, e.g., the training system 100 depicted in FIG. 1, appropriately programmed in accordance with this specification, can perform the process 400. The system can repeatedly perform iterations of the process 400 on different batches of output data items to update the values of the parameters of the score neural network. At each Attorney Docket No.: 45288-0475WO1 iteration, the system obtains, e.g., through sampling, a batch of one or more output data items from a training dataset that includes a large number of output data items. The output data items can include any of a variety of data items, e.g., image data items, audio data items, or video data items. For example, the system can repeatedly alternate between iterations of the process 200 and iterations of the process 400 using the same or different batches of output data items obtained from the same training dataset. For each output data item ! in the batch, the system uses the student generator neural network to generate a second student estimate of the output data item (step 402). The second student estimate is a predicted reconstruction of the output data item. The student generator neural network is configured to process a second student generator input that includes a latent representation ^ of the output data item and, optionally, a conditioning input, or data derived from the conditioning input, e.g., an embedding of the conditioning input, to generate a second student generator output that specifies the second student estimate of the output data item !. As mentioned previously, the latent representation ^ of the output data item includes a plurality of latent variables. Determining the latent representation ^ can include sampling, for each latent variable, a value from a predetermined distribution that corresponds to the latent variable, e.g., a normal distribution, or a different latent distribution, with some measure of randomness. The system uses the score neural network to generate a student score estimate of the output data item (step 404). The student score estimate is an approximation of a student score that would be generated by the student generator neural network based on an intermediate (noisy) representation of the output data item. The score neural network is configured to process a score neural network input to generate student score estimate. The score neural network input includes an intermediate (noisy) representation of the output data item that is generated based on adding noise to a data item ! (which can be generated by the student generator neural network) in accordance with the time step ^ and the noise term ^. The score neural network input can also include data that identifies the time step ^, e.g., an embedding of the time step ^. The score neural network input can also include a conditioning input, or data derived from the conditioning input, e.g., an embedding of the conditioning input. Attorney Docket No.: 45288-0475WO1 The system trains the score neural network to update values of parameters of the score neural network based at least in part on optimizing a score objective function (step 406). The score objective function evaluates, for each output data item in the batch, a difference between the student score estimate and a target score that is computed based on the second student estimate of the output data item. The system can do this by computing gradients of the score objective function with respect to the parameters of the score neural network <=by backpropagation through the parameters ϕ of the score neural network <=, and then applying a gradient descent technique, e.g., a stochastic gradient descent technique, based on the gradients. For example, the respective gradients of the parameters of the score neural network canbe computed as ∇?@<?^^^, ^^ ( ∇^A log ^^B^^|^"^^^D@^ ^ (optionally weighted), e.g. represented as: . dependent weighting function, <=^!^, ^^ represents the student score estimate of the output dataitem, ∇EAlogp^^!^|^#^z^^ represents the target score that is computed based on the secondstudent estimate of the output data item, and the expectation denotes, e.g. averaging over the distributions in the subscript. An example of the algorithms implementing the student generator neural network training process and the score neural network training process for training the student generator neural network is shown below. Attorney Docket No.: 45288-0475WO1 A student generator neural network trained using the techniques described in this specification can thus have parameters having trained values that reflect optimization of a generator objective function that evaluated, for each output data item from a training dataset, a difference between (i) an adjusted output data item and (ii) a corresponding student estimate of the output data item that was derived from an adjusted latent representation using the student generator neural network. The adjusted output data item and the adjusted latent representation have been derived from a stochastic sampling process that included a plurality of stochastic sampling steps and that was performed using the teacher diffusion model based at least on a latent representation of the output data item and a noise term. Example JAX code to implement these algorithms is below: 1 def eps_corrector ( self , eps , z , lambd , rng ) : Attorney Docket No.: 45288-0475WO1 eps_noises = jax . random . normal ( eps_base_rng , ( ld_steps ,) + eps .shape ) f Attorney Docket No.: 45288-0475WO1 z_next = jax . lax . stop_gradient ( z_next ) Attorney Docket No.: 45288-0475WO1 61 x_hat = jax . lax . stop_gradient ( x_hat ) 62 return x hat , z hat 1 def x_loss_g ( self , eps , z , lambd , rng ) : 2 ### MCMC correction ### FIG. 5 is a flow diagram of an example process 500 for generating an output data item. The process of FIG. 5 can be implemented by one or more computers in one or more locations. At step 502 the process samples a value of the latent representation, z, e.g. from a multivariate Gaussian or other distribution. At step 504 the student generator neural network processes the sampled value of the latent representation to obtain the output data item, optionally conditioned on a conditioning input as previously described. In some implementations the output data item is in a latent space and is decoded into an output space, e.g. an image space (i.e. a space of image pixel values), an audio waveform space, and so forth, by a decoder. FIGS. 6A-C illustrate the effect of stochastic sampling steps on image generation. FIG.6A shows images ^ = ^^^^ without stochastic sampling steps to update the joint distribution(x,z) (left column) and images with 300 stochastic sampling steps as described herein (right column). Attorney Docket No.: 45288-0475WO1 FIG. 6B shows output data items ^ = ^^^^ without stochastic sampling steps to updatethe joint distribution (x,z) (on the left) and output data items with 300 stochastic sampling steps as described herein (on the right). The output data items are in a latent image space (embedding space); FIG. 6C shows the corresponding images. FIG. 7 shows (on the left) four example images generated by a teacher diffusion model that uses 50 denoising steps, (on the right) four example images generated in a single step by a student generator neural network trained as described herein. In this specification, the term “configured” is used in relation to computing systems and environments, as well as computer program components. A computing system or environment is considered "configured" to perform specific operations or actions when it possesses the necessary software, firmware, hardware, or a combination thereof, enabling it to carry out those operations or actions during operation. For instance, configuring a system might involve installing a software library with specific algorithms, updating firmware with new instructions for handling data, or adding a hardware component for enhanced processing capabilities. Similarly, one or more computer programs are “configured” to perform particular operations or actions when they contain instructions that, upon execution by a computing device or hardware, cause the device to perform those intended operations or actions. The embodiments and functional operations described in this specification can be implemented in various forms, including digital electronic circuitry, software, firmware, computer hardware (encompassing the disclosed structures and their structural equivalents), or any combination thereof. The subject matter can be realized as one or more computer programs, essentially modules of computer program instructions encoded on a tangible non-transitory storage medium for execution by or to control the operation of a computing device or hardware. The storage medium can be a storage device such as a hard drive or solid-state drive (SSD), a storage medium, a random or serial access memory device, or a combination of these. Additionally or alternatively, the program instructions can be encoded on a transmitted signal, such as a machine-generated electrical, optical, or electromagnetic signal, designed to carry information for transmission to a receiving device or system for execution by a computing device or hardware. Furthermore, implementations may leverage emerging technologies like quantum computing or neuromorphic computing for specific applications, and may be deployed in Attorney Docket No.: 45288-0475WO1 distributed or cloud-based environments where components reside on different machines or within a cloud infrastructure. The term “computing device or hardware” refers to the physical components involved in data processing and encompasses all types of devices and machines used for this purpose. Examples include processors or processing units, computers, multiple processors or computers working together, graphics processing units (GPUs), tensor processing units (TPUs), and specialized processing hardware such as field-programmable gate arrays (FPGAs) or application- specific integrated circuits (ASICs). In addition to hardware, a computing device or hardware may also include code that creates an execution environment for computer programs. This code can take the form of processor firmware, a protocol stack, a database management system, an operating system, or a combination of these elements. Embodiments may particularly benefit from utilizing the parallel processing capabilities of GPUs, in a General-Purpose computing on Graphics Processing Units (GPGPU) context, where code specifically designed for GPU execution, often called kernels or shaders, is employed. Similarly, TPUs excel at running optimized tensor operations crucial for many machine learning algorithms. By leveraging these accelerators and their specialized programming models, the system can achieve significant speedups and efficiency gains for tasks involving artificial intelligence and machine learning, particularly in areas such as computer vision, natural language processing, and robotics. A computer program, also referred to as software, an application, a module, a script, code, or simply a program, can be written in any programming language, including compiled or interpreted languages, and declarative or procedural languages. It can be deployed in various forms, such as a standalone program, a module, a component, a subroutine, or any other unit suitable for use within a computing environment. A program may or may not correspond to a single file in a file system and can be stored in various ways. This includes being embedded within a file containing other programs or data (e.g., scripts within a markup language document), residing in a dedicated file, or distributed across multiple coordinated files (e.g., files storing modules, subprograms, or code segments). A computer program can be executed on a single computer or across multiple computers, whether located at a single site or distributed across multiple sites and interconnected through a data communication network. The specific implementation of the computer programs may involve a combination of traditional programming languages and specialized languages or libraries designed for GPGPU Attorney Docket No.: 45288-0475WO1 programming or TPU utilization, depending on the chosen hardware platform and desired performance characteristics. In this specification, the term “engine” broadly refers to a software-based system, subsystem, or process designed to perform one or more specific functions. An engine is typically implemented as one or more software modules or components installed on one or more computers, which can be located at a single site or distributed across multiple locations. In some instances, one or more dedicated computers may be used for a particular engine, while in other cases, multiple engines may operate concurrently on the same one or more computers. Examples of engine functions within the context of AI and machine learning could include data pre- processing and cleaning, feature engineering and extraction, model training and optimization, inference and prediction generation, and post-processing of results. The specific design and implementation of engines will depend on the overall architecture and the distribution of computational tasks across various hardware components, including CPUs, GPUs, TPUs, and other specialized processors. The processes and logic flows described in this specification can be executed by one or more programmable computers running one or more computer programs to perform functions by operating on input data and generating output. Additionally, graphics processing units (GPUs) and tensor processing units (TPUs) can be utilized to enable concurrent execution of aspects of these processes and logic flows, significantly accelerating performance. This approach offers significant advantages for computationally intensive tasks often found in AI and machine learning applications, such as matrix multiplications, convolutions, and other operations that exhibit a high degree of parallelism. By leveraging the parallel processing capabilities of GPUs and TPUs, significant speedups and efficiency gains compared to relying solely on CPUs can be achieved. Alternatively or in combination with programmable computers and specialized processors, these processes and logic flows can also be implemented using specialized processing hardware, such as field-programmable gate arrays (FPGAs) or application-specific integrated circuits (ASICs), for even greater performance or energy efficiency in specific use cases. Computers capable of executing a computer program can be based on general-purpose microprocessors, special-purpose microprocessors, or a combination of both. They can also utilize any other type of central processing unit (CPU). Additionally, graphics processing units Attorney Docket No.: 45288-0475WO1 (GPUs), tensor processing units (TPUs), and other machine learning accelerators can be employed to enhance performance, particularly for tasks involving artificial intelligence and machine learning. These accelerators often work in conjunction with CPUs, handling specialized computations while the CPU manages overall system operations and other tasks. Typically, a CPU receives instructions and data from read-only memory (ROM), random access memory (RAM), or both. The elements of a computer include a CPU for executing instructions and one or more memory devices for storing instructions and data. The specific configuration of processing units and memory will depend on factors like the complexity of the AI model, the volume of data being processed, and the desired performance and latency requirements. Embodiments can be implemented on a wide range of computing platforms, from small embedded devices with limited resources to large-scale data center systems with high-performance computing capabilities. The system may include storage devices like hard drives, SSDs, or flash memory for persistent data storage. Computer-readable media suitable for storing computer program instructions and data encompass all forms of non-volatile memory, media, and memory devices. Examples include semiconductor memory devices such as read-only memory (ROM), solid-state drives (SSDs), and flash memory devices; hard disk drives (HDDs); optical media; and optical discs such as CDs, DVDs, and Blu-ray discs. The specific type of computer-readable media used will depend on factors such as the size of the data, access speed requirements, cost considerations, and the desired level of portability or permanence. To facilitate user interaction, embodiments of the subject matter described in this specification can be implemented on a computing device equipped with a display device, such as a liquid crystal display (LCD) or an organic light-emitting diode (OLED) display, for presenting information to the user. Input can be provided by the user through various means, including a keyboard), touchscreens, voice commands, gesture recognition, or other input modalities depending on the specific device and application. Additional input methods can include acoustic, speech, or tactile input, while feedback to the user can take the form of visual, auditory, or tactile feedback. Furthermore, computers can interact with users by exchanging documents with a user's device or application. This can involve sending web content or data in response to requests or sending and receiving text messages or other forms of messages through mobile devices or Attorney Docket No.: 45288-0475WO1 messaging platforms. The selection of input and output modalities will depend on the specific application and the desired form of user interaction. Machine learning models can be implemented and deployed using machine learning frameworks, such as TensorFlow or JAX. These frameworks offer comprehensive tools and libraries that facilitate the development, training, and deployment of machine learning models. Embodiments of the subject matter described in this specification can be implemented within a computing system comprising one or more components, depending on the specific application and requirements. These may include a back-end component, such as a back-end server or cloud-based infrastructure; an optional middleware component, such as a middleware server or application programming interface (API), to facilitate communication and data exchange; and a front-end component, such as a client device with a user interface, a web browser, or an app, through which a user can interact with the implemented subject matter. For instance, the described functionality could be implemented solely on a client device (e.g., for on- device machine learning) or deployed as a combination of front-end and back-end components for more complex applications. These components, when present, can be interconnected using any form or medium of digital data communication, such as a communication network like a local area network (LAN) or a wide area network (WAN) including the Internet. The specific system architecture and choice of components will depend on factors such as the scale of the application, the need for real-time processing, data security requirements, and the desired user experience. The computing system can include clients and servers that may be geographically separated and interact through a communication network. The specific type of network, such as a local area network (LAN), a wide area network (WAN), or the Internet, will depend on the reach and scale of the application. The client-server relationship is established through computer programs running on the respective computers and designed to communicate with each other using appropriate protocols. These protocols may include HTTP, TCP / IP, or other specialized protocols depending on the nature of the data being exchanged and the security requirements of the system. In certain embodiments, a server transmits data or instructions to a user's device, such as a computer, smartphone, or tablet, acting as a client. The client device can then process the received information, display results to the user, and potentially send data or feedback back Attorney Docket No.: 45288-0475WO1 to the server for further processing or storage. This allows for dynamic interactions between the user and the system, enabling a wide range of applications and functionalities. 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 depicted in the drawings and recited in the claims in a particular order, this should not be understood as requiring that such operations be performed in the particular order shown or in sequential order, or that all illustrated operations be performed, to achieve desirable results. In certain circumstances, multitasking and parallel processing may be advantageous. Moreover, the separation of various system modules and components in the embodiments described above should not be understood as requiring such separation in all embodiments, and it should be understood that the described program components and systems can generally be integrated together in a single software product or packaged into multiple software products. Particular embodiments of the subject matter have been described. Other embodiments are within the scope of the following claims. For example, the actions recited in the claims can be performed in a different order and still achieve desirable results. As one example, the processes depicted in the accompanying figures do not necessarily require the particular order shown, or sequential order, to achieve desirable results. In some cases, multitasking and parallel processing may be advantageous. What is claimed is:

Claims

Attorney Docket No.: 45288-0475WO1 CLAIMS1. A method of training a student generator neural network to generate an output data itemusing a teacher diffusion model, wherein the teacher diffusion model is configured to generate the output data item as a result of performing a plurality of reverse diffusion steps, wherein the method comprises: determining a latent representation of an output data item, and a noise term; performing a plurality of stochastic sampling steps using the teacher diffusion model to determine an adjusted output data item and an adjusted latent representation of the output data item based on the latent representation and the noise term; processing a first student generator input that comprises the adjusted latent representation using the student generator neural network to generate a first student estimate of the output data item; and training the student generator neural network to update values of parameters of the student generator neural network based at least in part on optimizing a generator objective function that evaluates a difference between the adjusted output data item and the first student estimate of the output data item.

2. The method of claim 1, wherein performing the plurality of stochastic sampling stepscomprises, at each stochastic sampling step: sampling a perturbation for the stochastic sampling step based on a predetermined distribution; generating a teacher score for the output data item by processing a teacher diffusion model input that comprises a noisy representation of the output data item; and determining an updated noise term and an updated latent representation by applying a one-step update based on the perturbation and the teacher score.

3. The method of claim 2, wherein determining the updated noise term and the updatedlatent representation comprises applying a reparameterization sampling technique that uses the teacher score.Attorney Docket No.: 45288-0475WO1 4. The method of any preceding claim, wherein determining the adjusted output data item and the adjusted latent representation of the output data item comprises, after a last stochastic sampling step: determining the adjusted output data item and the adjusted latent representation based on a final updated noise term and a final updated latent representation generated in the last stochastic sampling step.

5. The method of claim 4, wherein determining the adjusted output data item comprises: determining an initial adjusted output data item based on the final updated latent representation; and removing an accumulated noise term from the initial adjusted output data item.

6. The method of claim 4 or claim 5, wherein determining the adjusted latent representation comprises: using the final updated latent representation as the adjusted latent representation.

7. The method of any one of claims 1-6, wherein the latent representation comprises a plurality of latent variables, and wherein determining the latent representation comprises sampling, for each latent variable, a value from a predetermined distribution.

8. The method of any one of claims 1-7, wherein training the student generator neural network comprises: computing gradients of the generator objective function with respect to the parameters of the student generator neural network; and applying a stochastic gradient ascent technique based on the gradients.

9. The method of claim 8, wherein the gradients are weighted by weighting factors computed in accordance with a time-dependent weighting function.

10. The method of any one of claims 1-9, wherein the student generator neural network has a neural network architecture that is dependent on a neural network architecture of the teacher diffusion model.Attorney Docket No.: 45288-0475WO111. The method of any one of claims 1-10, wherein the output data item comprises one of:image data, audio data, or video data.

12. The method of any one of claims 1-11, wherein the method further comprises:processing a second student generator input that comprises the latent representation using the student generator neural network to generate a second student estimate of the output data item; processing, using a score neural network, a score neural network input that comprises (i) a or the noisy representation of the output data item and (ii) data that identifies a time step to generate a student score estimate of the output data item; and training the score neural network to update values of parameters of the score neural network based at least in part on optimizing a score objective function that evaluates a difference between the student score estimate and a target score that is computed based on the second student estimate of the output data item.

13. The method of claim 12, wherein training the score neural network to update the valuesof the parameters of the score neural network comprises: computing gradients of the score objective function with respect to the parameters of the score neural network; and applying a stochastic gradient descent technique based on the gradients.

14. The method of claim 12 or claim 13, wherein the method further comprises using thescore neural network to train the student generator neural network.

15. The method of any preceding claim, wherein the student generator neural network isconfigured to generate the output data item in a single forward pass through the student generator neural network.

16. The method of any preceding claim, wherein the student generator neural network isconfigured to generate the output data item by performing one or more processing iterations andAttorney Docket No.: 45288-0475WO1 wherein a total number of processing iterations is less than a total number of reverse diffusion steps.

17. A method comprising: receiving a request for an output data item; and generating the output data item using the student generator neural network that has been trained using the method of any one of claims 1-16.

18. A method comprising: receiving a request for an output data item; and generating the output data item using a student generator neural network that has been trained using a teacher diffusion model, wherein the teacher diffusion model is configured to generate the output data item as a result of performing a plurality of reverse diffusion steps, the student generator neural network having parameters reflecting optimization of a generator objective function that evaluated a difference between an adjusted training data item, and a corresponding student estimate of the training data item derived from an adjusted latent representation using the student generator neural network, wherein the adjusted training data item and the adjusted latent representation have been derived from stochastic sampling steps performed using the teacher diffusion model on a latent representation of a training data item and a noise term.

19. 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.

20. 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 preceding claim.

Images