Efficient diffusion machine learning models
Certain aspects of the present disclosure provide techniques and apparatus for improved machine learning. During a first iteration of processing data using a first denoising backbone of a teacher diffusion machine learning model, a first latent tensor is generated using a lower resolution block of the first denoising backbone. During a first iteration of processing data using a second denoising backbone of a student diffusion machine learning model, a second latent tensor is generated using an adapter block of the second denoising backbone. A loss is generated based on the first and second latent tensors, and one or more parameters of the adapter block are updated based on the loss.
1 . A processing system, comprising:
means for generating, during a first iteration of processing data using a first denoising backbone of a teacher diffusion machine learning model, a first latent tensor using a lower resolution block of the first denoising backbone;
means for generating, during a first iteration of processing data using a second denoising backbone of a student diffusion machine learning model, a second latent tensor using an adapter block of the second denoising backbone;
means for generating a loss based on the first and second latent tensors; and
means for updating one or more parameters of the adapter block based on the loss.
2 . The processing system of claim 1 , further comprising:
means for updating one or more parameters of a higher resolution block of the second denoising backbone based on the loss; and
means for updating one or more parameters of a lower resolution block of the second denoising backbone based on the loss.
3 . The processing system of claim 1 , wherein generating the second latent tensor is performed based further on:
processing an embedding corresponding to the first iteration using the adapter block;
processing an embedding corresponding to an input to the student diffusion machine learning model using the adapter block; and
processing an embedding, generated by a higher resolution block of the second denoising backbone, using the adapter block.
4 . The processing system of claim 1 , wherein generating the second latent tensor is performed based further on processing an embedding corresponding to the first iteration using the adapter block.
5 . The processing system of claim 1 , wherein generating the second latent tensor is performed based further on processing an embedding corresponding to an input to the student diffusion machine learning model using the adapter block.
6 . The processing system of claim 1 , wherein generating the second latent tensor is performed based further on processing an embedding, generated by a higher resolution block of the second denoising backbone, using the adapter block.
7 . The processing system of claim 1 , wherein the adapter block performs one or more convolution operations to generate the second latent tensor.
8 . The processing system of claim 1 , wherein:
the adapter block comprises an encoder and a decoder, and
generating the second latent tensor comprises:
generating a compressed tensor based on processing a third latent tensor using the encoder, and
generating the second latent tensor based on processing the compressed tensor using the decoder.
9 . The processing system of claim 1 , further comprising:
means for generating a third latent tensor based on processing the second latent tensor using the adapter block; and
means for generating, during a second iteration of processing the data using the student diffusion machine learning model, a feature tensor based on processing the third latent tensor using a higher resolution block of the second denoising backbone.
10 . The processing system of claim 1 , further comprising:
means for generating, during a second iteration of processing the data using the student diffusion machine learning model, a third latent tensor using a lower resolution block of the second denoising backbone; and
means for generating, during the second iteration, a feature tensor based on processing the third latent tensor using a higher resolution block of the second denoising backbone.