Fine-tuning and controlling diffusion models
Systems and methods for fine-tuning diffusion models are described. Embodiments of the present disclosure obtain an input text indicating an element to be included in an image; generate a synthetic image depicting the element based on the input text using a diffusion model trained by comparing synthetic images depicting the element to training images depicting elements similar to the element and updating selected parameters corresponding to an attention layer of the diffusion model based on the comparison.
1 . A method comprising:
obtaining an input text indicating an element to be included in an image, wherein the element belongs to a semantic domain; and
generating a synthetic image depicting the element based on the input text using a diffusion model, wherein the diffusion model is trained by:
performing a first training phase,
determining that the diffusion model generates inaccurate output in the semantic domain after the first training phase,
and performing a second training phase to improve an accuracy of the diffusion model for generating images in the semantic domain by comparing synthetic images depicting the element to training images depicting elements similar to the element and updating selected parameters corresponding to a cross-attention block of an attention layer of the diffusion model based on the comparison while holding additional parameters of the cross-attention block fixed, and wherein the cross-attention block operates based on text features and image features.
2 . The method of claim 1 , wherein:
the selected parameters are limited to K mappings and V mappings of the cross-attention block.
3 . The method of claim 1 , wherein:
the selected parameters are limited to a projection mapping matrix of the cross-attention block.
4 . The method of claim 3 , wherein:
the selected parameters are limited to a restricted set of parameters of the projection mapping matrix.
5 . The method of claim 1 , wherein:
the selected parameters are limited to selected resolution levels of a U-Net architecture.
6 . The method of claim 1 , further comprising:
identifying the semantic domain for improving performance of the diffusion model; and
selecting the training images based on the semantic domain.
7 . The method of claim 1 , further comprising:
encoding the training images to obtain the image features;
adding noise to the image features to obtain noisy image features;
generating de-noised image features from the noisy image features using the diffusion model; and
decoding the de-noised image features to obtain the synthetic images for training.
8 . A method comprising:
identifying a diffusion model trained on a first training set during a first training phase;
determining that the diffusion model generates inaccurate output in a semantic domain after the first training phase;
selecting, based on the determination, a first subset of parameters of the diffusion model corresponding to a cross-attention block of an attention layer and a second subset of parameters of the diffusion model outside of the cross-attention block; and
training, based on the determination, the diffusion model to generate a synthetic image depicting an element in the semantic domain based on text condition features, wherein the training improves an accuracy of the diffusion model for generating images in the semantic domain by fine-tuning the first subset of parameters of the diffusion model based on a second training set that includes data in the semantic domain and is different from the first training set during a second training phase, wherein the first subset of parameters is fine-tuned based on comparing the synthetic image with the data in the semantic domain while holding the second subset of parameters of the cross-attention block fixed during the second training phase, and wherein the cross-attention block operates based on text features and image features.
9 . The method of claim 8 , further comprising:
identifying the semantic domain corresponding to the element, wherein the second training set includes a plurality of images representing the element.
10 . The method of claim 8 , further comprising:
identifying the semantic domain corresponding to a class of elements, wherein the second training set includes a plurality of images representing the class of elements.
11 . The method of claim 8 , further comprising:
identifying an image from the first training set;
encoding the image to obtain an encoded image; and
comparing the encoded image to a target encoding for the second training set, wherein the image is included in the second training set based on the comparison of the encoded image to the target encoding.
12 . The method of claim 8 , further comprising:
identifying an image from the first training set and a caption corresponding to the image;
encoding the caption to obtain an encoded caption; and
comparing the encoded caption to a target encoding for the second training set, wherein the image is included in the second training set based on the comparison of the encoded caption to the target encoding.
13 . The method of claim 8 , further comprising:
obtaining an input text; and
generating an output image based on the input text using the diffusion model during the second training phase, wherein the fine-tuning is based on the output image.
14 . An apparatus comprising:
one or more processors; and
one or more memories including instructions executable by the one or more processors to:
generate text condition features from an input text using a text transformer model, wherein the input text indicates an element to be included in an image, and the element belongs to a semantic domain;
select images related to the input text; and
generate synthetic images based on the text condition features using a diffusion model, wherein the diffusion model is trained by performing a first training phase, determining that the diffusion model generates inaccurate output in the semantic domain after the first training phase, and performing a second training phase to improve an accuracy of the diffusion model for generating images in the semantic domain by comparing synthetic images depicting the element to training images depicting elements similar to the element and updating selected parameters corresponding to a cross-attention block of one or more attention layers of the diffusion model based on the comparison while holding additional parameters of the cross-attention block fixed, and wherein the cross-attention block operates based on text features and image features.
15 . The apparatus of claim 14 , wherein the instructions are further executable to:
add noise to the images and the image features.
16 . The apparatus of claim 14 , wherein the instructions are further executable to:
fine-tune the diffusion model.
17 . The apparatus of claim 14 , wherein:
the diffusion model comprises a latent diffusion model.
18 . The apparatus of claim 14 , wherein:
the diffusion model comprises a U-net architecture.
19 . The apparatus of claim 14 , wherein:
the selected parameters are limited to K mappings and V mappings of the cross-attention block.