Method for obtaining rapid text map diffusion model through low-computing-resource distillation

Abstract

The invention discloses a method for obtaining a rapid text graph model with low computing resources, and the training method comprises the steps: carrying out the UNet optimization of a Stable Diffusion v1.4 model, including module pruning and softmax function approximation, and obtaining a denoising network FastSD of rapid single-step reasoning; stable Diffusion v1.4 is further used as a teacher model, the obtained FastSD is used as a student model, multi-loss combined distillation training is carried out on the FastSD model through combined distillation loss, anti-damage and distribution matching loss, training parameter quantity is reduced through a low-rank decomposition technology, and training loss is continuously iterated until convergence conditions are met.

Description

Technical Field

[0001] This invention belongs to the field of multimodal artificial intelligence generation technology, specifically to a diffusion model optimization method for text-to-image generation, particularly suitable for accelerating model distillation in scenarios with limited computing resources. Furthermore, the invention provides a model compression strategy and a multi-objective joint training method. Background Art

[0002] With the development of artificial intelligence, especially AIGC, the diffusion model for generating images from text has been widely used in social networking, drawing, web pages and other fields, with the characteristics of high creation efficiency and high degree of automation. For example, by simply inputting a piece of text into the above-mentioned text-to-image diffusion model, the model can generate diverse and high-quality images. However, the current text-to-image generation technology based on the diffusion model has the following technical bottlenecks: (1) High inference latency: The traditional U-Net architecture has redundant computing modules, resulting in significant time-consuming single-step inference; (2) High distillation cost: The existing model distillation method requires thousands of iterations on the complete dataset, requiring more than one thousand TFLOPS of computing power; (3) Accuracy-efficiency imbalance: Conventional pruning strategies lead to feature space mismatch, and the generation quality drops by more than 15%. Summary of the Invention

[0003] To address the problem that diffusion model distillation consumes a large amount of computing resources, this paper proposes a diffusion model distillation method with efficient training. By pruning and optimizing key high-latency structures in the model, a compressed diffusion model is obtained. By optimizing the loss function, combining distillation loss, adversarial loss and distribution matching loss, and using low-rank decomposition for stable distillation training, a fast diffusion model with low computing resource consumption is obtained.

[0004] A technical solution proposed by the present invention to achieve the above-mentioned purpose is as follows:

[0005] refer to Figure 1 A diffusion model optimization method with low computational resource consumption is provided, for obtaining a diffusion model with low latency for single-step reasoning. The diffusion model includes a variational autoencoding model, a text information encoding model, and a U-shaped network (UNet). The optimization method includes:

[0006] S1. Take 100 image-text pairs from the dataset, extract the prompt text, input Stable Diffusion v1.4 to perform 100 inferences, and obtain the average inference latency of each UNet module;

[0007] S2. Reference Figure 2 ,According to the analysis results of the inference delay, the two modules with the largest delay are removed, and a diffusion model after module pruning is obtained;

[0008] S3. Use GeLU activation function to replace all ReLU activation functions in Stable Diffusion v1.4 to solve the problem of low computational efficiency of traditional ReLU;

[0009] S4. After the above processing, an optimized diffusion model FastSD is obtained.

[0010] The present invention also proposes an efficient distillation method for training the compressed model, which specifically includes:

[0011] P1. Reference Figure 3 After adding noise to the images in the real dataset, they are fed into the original diffusion model StableDiffusion v1.4 (as the teacher model for distillation) and FastSD (as the student model for distillation), and the following denoising steps P11 and P12 are performed respectively:

[0012] P11, Stable Diffusion v1.4 performs several (usually 25) denoising steps to obtain the teacher model's prediction z 0_T ;

[0013] P12, FastSD performs one-step denoising to obtain the student model's prediction z 0_s ;

[0014] P13, calculating the distillation loss based on the two prediction results;

[0015] P2, the result obtained in step P1 Perform two noise additions to obtain the noisy images. and Add noise to z0 to get z t1 , and perform the following training process:

[0016] P21, z 1_S and z t1 The model is fed into the teacher diffusion model Stable Diffusion v1.4 to perform 25 steps of denoising. The denoised latent space features are fed into a discriminator for discrimination, and the discrimination loss of the two is calculated.

[0017] P22, z t2_S Input the teacher model Stablediffusion v1.4 and the student model FastSD respectively, and calculate the distribution matching loss of the two;

[0018] P3. By using LoRA for FastSD and combining it with the distillation loss, adversarial loss, and distribution loss for training, only TGLOPS of computing power resources are required to obtain the distilled FastSD model.

[0019] The above technical solution proposed by the present invention is a compression-optimized diffusion model FastSD, which can be trained using a small amount of computing resources without significantly reducing the performance of the original Stable Diffusion v1.4 model through the above training process. Other features of the present invention will become easy to understand through the following description. BRIEF DESCRIPTION OF THE DRAWINGS

[0020] Figure 1 This is a flow chart of a method for optimizing the Stable Diffusion v1.4 Wensheng graph model provided by the present invention;

[0021] Figure 2 This is a comparison chart before and after module-level UNet pruning provided by the present invention;

[0022] Figure 3 This is a flow chart of a multi-stage distillation training framework provided by the present invention. DETAILED DESCRIPTION

[0023] The present invention is further described below with reference to the accompanying drawings and specific embodiments.

[0024] The proposed method for optimizing the Stable Diffusion v1.4 text graph model consists of three core components: a variational autoencoder (VAE), a text encoder, and a deeply reconstructed UNet network. In practice, this method addresses the latency bottleneck of the UNet by combining module-level pruning with function approximation. The following detailed explanation is provided in stages:

[0025] refer to Figure 1 The present invention aims to optimize the publicly available Stable Diffusion v1.4, including a variational autoencoder (VAE), a text encoder, and an optimized UNet. The present invention notes that UNet is the latency bottleneck of Stable Diffusion v1.4, so the UNet is pruned and optimized:

[0026] S1. Reference Figure 1 , remove the Transformer module of the outermost layer of UNet (the first layer in the downsampling stage and the last layer in the upsampling stage);

[0027] S2. Use a length-regularized ReLU function to approximate the softamx function in all Transformer blocks in UNet. The length-regularized ReLU function is defined as follows:

[0028] f=L -1 relu

[0029] refer to Figure 3 , further multi-objective joint training of the distilled model:

[0030] 1.1 Sample a sample x0~p(x0) from the data distribution and obtain the latent space sample z0 through the image encoder E;

[0031] 1.2 Sampling a time step t from the time density function π(t), according to x t =α(t)·x0+σ(t)·ε adds Gaussian noise to x0, where ε follows a standard normal distribution and ε~N(0,I);

[0032] 1.3 x t Input the Stable Diffusion v1.4 teacher model and use the teacher model's ODE solver (25-step denoising)

[0033] Generate a teacher model's predictions for the sample

[0034] 1.4 At the same time, x t Input a student model, perform one-step denoising and get a student model prediction for the sample

[0035]

[0036] 1.5 Calculate distillation loss:

[0037] 2. Calculate the adversarial loss through steps 2.1 to 2.5:

[0038] 2.1 Sampling a time t1 from the time density function π(t) again, according to x t=α(t)·x0+σ(t)·ε for x0 and

[0039] Adding Gaussian noise, we get and

[0040] 2.2 At the same time and Input the Stable Diffusion v1.4 teacher model and use the teacher model to predict the corresponding outputs: and

[0041] 2.3 and Input a discriminator, the discriminator will The corresponding prediction is judged to be false. The corresponding prediction is judged to be true; calculate the discriminant loss:

[0042] 2.4 Calculate adversarial loss:

[0043] 3.1-to-1 student model FastSD prediction Again sample a time t2 from the time density function π(t), according to x t =α(t)·x0+σ(t)·ε for x0 and Adding Gaussian noise yields

[0044] 3.2 Input the teacher model Stable Diffusion v1.4 to perform 25-step prediction and obtain the teacher model score At the same time Input student model FastSD to get the score of student model in f is the corresponding Stable Diffusion v1.4 teacher model or FastSD student model;

[0045] 3.3 Calculate distribution matching loss:

[0046] 4. Add the three losses 1, 2, and 3 above in proportion to form the total loss of training:

[0047]

[0048] The preferred weight coefficients are λ1 = 0.7, λ2 = 0.2, and λ3 = 0.1. Experimental verification shows that this method can reduce the UNet inference latency by 62% while maintaining the FID index at 97% of the original model.

[0049] Table 1 Performance and cost comparison of the present invention and other model training methods on the COCO2014-30k evaluation set

[0050]

Claims

1. A method for obtaining a fast cultural graph diffusion model by distillation with low computing resources, characterized in that: The following steps are involved: (a) Pruning the UNet network of the Stable Diffusion v1.4 model: removing the first layer of Transformer modules in the downsampling stage and the last layer of Transformer modules in the upsampling stage, and introducing cross-layer residual connections at the pruned locations; (b) Replace the softmax function in UNet with the sequence-scaled relu function; (c) The pruned model from steps (a) and (b) is denoted as FastSD and used as the student model. StableDiffusionv1.4 is used as the teacher model to construct a multi-stage distillation training framework: (c1) Calculate the hybrid distillation loss by combining the multi-step denoising results of the teacher model with the single-step prediction results of the student model; (c2) Perform adversarial training on the noisy latent variables of the teacher-student model, and use a discriminator containing a spectral normalization convolutional layer to calculate the adversarial loss; (c3) Data manifold alignment is achieved by calculating the fractional matching loss through recursive noise injection and fractional gradient matching; (d) The distillation loss, adversarial loss, and score matching loss described in step (c) are fused according to preset weights to jointly optimize the student model parameters.

2. The method according to claim 1, characterized in that The structure pruning in step (a) specifically includes: After removing the first layer Transformer module in the downsampling stage, the front 3×3 convolution kernel is retained and the number of output channels is reduced by 50%; In the upsampling stage, bilinear interpolation is used at the last layer to replace the original transposed convolution operation.

3. The method according to claim 1, characterized in that The relu function of the (a) sequence scaling is defined as: f=L -1 relu Where L is the length of the sequence.

4. The method according to claim 1, wherein The mixed distillation loss in step (c1) is calculated as follows:

5. The method according to claim 1, wherein The discriminator in step (c2) is a single-layer convolutional neural network with a step size of 2, a convolution kernel size of 4, and a SiLU activation function and group regularization. The adversarial loss function is:

6. The method according to claim 1, characterized in that The fractional gradient matching loss in step (c3) is implemented by implicit differentiation:

7. The method according to claim 1, characterized in that The total loss function in step (d) is: The preferred weight coefficients are λ1 = 0.7, λ2 = 0.2, and λ3 = 0.

1. A phased weight adjustment strategy is adopted during training: the λ1 coefficient is linearly reduced to 0.5 in the first 50% of iterations, and the λ2 coefficient is gradually increased to 0.4 in the last 50% of iterations.

8. The method according to any one of claims 1 to 7, characterized in that: When the optimized model is deployed: Support 8-bit quantized inference in dynamic computation graph mode; Compatible with the original VAE encoder and text encoder, no additional fine-tuning is required; When generating a 512×512 resolution image, the total latency for single-image inference is less than 120ms (NVIDIA RTX4090 GPU).

Images