A diffusion model acceleration method and system based on progressive multi-stage relationship distillation

Abstract

The method of the present application divides the reverse denoising process of the teacher model into a high-noise stage, a medium-noise stage and a low-noise stage according to time steps by introducing a stage-adaptive differential loss function in a progressive distillation framework, and obtains a distillation loss function for each stage; a first-stage progressive distillation is performed to distill the teacher model from N steps to N / 2 steps, and finally an intermediate student model is obtained; when further distillation to a smaller step number is performed, a relationship knowledge transfer loss is introduced, a consistency constraint of sample inter-pair similarity structure is imposed between the student model and the teacher model, and the classifier feature distillation loss is jointly optimized; and finally, the student model and the sample inter-relationship knowledge transfer are output, so that an image with better quality is obtained.

Description

Technical Field

[0001] This invention belongs to the field of deep learning and generative model technology, specifically relating to a diffusion model acceleration method and system based on progressive multi-stage relation distillation. Background Technology

[0002] Diffusion models have become the de facto standard for high-fidelity image generation, consistently outperforming traditional generative models such as Generative Adversarial Networks (GANs) and Variational Autoencoders (VAEs) on a wide range of benchmarks. Their success stems from an iterative denoising process: starting from a random noisy sample, it is progressively refined into a structurally clear image. However, this advantage also introduces a significant sampling bottleneck—generating a high-quality image typically requires hundreds or even thousands of sequential neural network evaluations, making real-time deployment of the model on edge devices extremely impractical.

[0003] To alleviate the sampling efficiency problem, existing work is mainly carried out in two directions: 1. Training-free method: A moderate speedup can be achieved by designing a more efficient numerical integrator for probabilistic flow ordinary differential equations (such as DDIM), but the performance drops sharply at very low step counts (≤8).

[0004] 2. Training-based distillation method: A lightweight student model is trained to mimic the more computationally intensive teacher model, progressively halving the required sampling steps. Progressive distillation (PD) aligns the student model's single-step output with the teacher model's two-step output on the same noisy input, successfully compressing the sampling steps from 1024 to 8. Subsequent classifier-based feature distillation (CFD) further replaces the pixel-level mean squared error with feature-level KL divergence extracted from the pre-trained classifier, improving the generation quality with very few steps (1-4 steps).

[0005] Despite the significant progress made by the aforementioned distillation methods, they still have two fundamental limitations: Limitation 1: Sample-Independent Distillation Paradigm. Both PD and CFD process each training sample independently, forcing the student model to match the teacher model's output (pixel values ​​or feature vectors) on a sample-by-sample basis, without imposing any constraints on how the student model relates different samples. The teacher model encodes rich structural information: two semantically similar inputs should produce similar representations, while inputs with significant differences should remain separate. Traditional knowledge distillation has long demonstrated that distilling relational knowledge—i.e., pairwise similarity or distance between samples—significantly outperforms simple instance-level alignment. However, this important insight has never been systematically explored in the context of diffusion model distillation.

[0006] Limitation 2: Uniform treatment of the denoising stage. The backdiffusion process is not homogeneous. The early denoising steps (high noise level) operate on near-pure Gaussian noise, with the feature distribution exhibiting isotropic and high-entropy characteristics; the mid-stage steps transition from coarse structure to fine detail, exhibiting the most dramatic distribution changes; the late-stage steps (low noise level) refine local textures, making the feature distribution sharper and more class-discriminative. However, existing distillation methods apply the exact same loss function at all time steps, ignoring this inherent stage-specific heterogeneity. This "one-size-fits-all" strategy is suboptimal. Summary of the Invention

[0007] To address the shortcomings of existing technologies, this invention proposes a diffusion model acceleration method based on progressive multi-stage relation distillation. The aim is to significantly improve the quality and stability of generated images with an extremely small number of steps by using a stage-adaptive distillation strategy and knowledge transfer of relationships between samples.

[0008] The objective of this invention is achieved through the following technical measures.

[0009] A diffusion model acceleration method based on progressive multi-stage relational distillation is characterized by the following steps: (1) Obtain a pre-trained teacher model, wherein the teacher model uses the DDIM sampler and the total number of sampling steps is N, where N≥128; (2) The reverse denoising process of the teacher model is divided into three stages according to the time step: high noise stage, medium noise stage and low noise stage, and a differentiated distillation loss function is designed for each stage. (3) Perform the first stage of progressive distillation, distilling the teacher model from N steps to N / 2 steps. During the distillation process, dynamically select the corresponding stage loss function according to the stage to which the current time step belongs to calculate the distillation loss and update the student model parameters. (4) Repeat step (3) until an intermediate student model with 8 steps is obtained; (5) In the process of distilling from 8 steps to a very low number of steps, relational knowledge transfer loss is introduced, and consistency constraints on the pairwise similarity structure between samples are imposed between the student model and the teacher model, and jointly optimized with the classifier feature distillation loss. (6) Based on the time step of the current student model For each high-noise, medium-noise, or low-noise stage, perform progressive distillation steps (3) to (5), dynamically select the corresponding stage-differentiated loss function to update the model parameters, and output the final student model.

[0010] According to the method described in this invention, the three stages in step (2) are divided as follows: the high-noise stage corresponds to the student time step. The middle noise stage corresponds to Low noise stage corresponds to ,in This represents the total number of steps in the current distillation stage of the student model.

[0011] According to the method described in this invention, the differentiated distillation loss function in step (2) is specifically: Mean square error loss is used in the high-noise stage: ; in, This represents the set of batch sample indices corresponding to the high-noise phase. The number of samples in this set. The noise reduction result of the student model's prediction for the i-th sample. The target denoising results provided for the teacher model Represents the L2 norm; In the mid-noise stage, a hybrid loss method combining mean square error and cosine similarity is used: ; in, , This is the set of batch sample indexes corresponding to the medium-noise stage. ;symbol" " represents the vector dot product; The low-noise stage employs a combination of mean square error and L1 loss: ; in , , This is the set of batch sample indexes corresponding to the low-noise phase. This represents the L1 norm.

[0012] According to the method described in this invention, the relational knowledge transfer loss in step (5) is a similarity preservation loss, which is calculated as follows: for B samples in the same batch, the teacher characteristics are calculated separately. and student characteristics Normalization to the similarity matrix .in Similarity Preservation Loss Defined as the mean squared error between two similarity matrices: .

[0013] According to the method described in this invention, the total loss function of the joint optimization in step (5) is: ; in The loss is calculated by distilling features for the classifier (extracting features from a pre-trained classifier and calculating the KL divergence). Loss weights are maintained for similarity.

[0014] According to the method of the present invention, the extremely low number of steps in step (5) includes 4 steps, 2 steps and 1 step, and the distillation process is carried out step by step: first distilling from 8 steps to 4 steps, then distilling to 2 steps, and finally distilling to 1 step.

[0015] According to the method described in this invention, the teacher diffusion model is a U-Net architecture diffusion model pre-trained on the CIFAR-10 or ImageNet 64×64 dataset, with a total sampling step count of 1024 steps.

[0016] A system for executing the diffusion model acceleration method based on progressive multi-stage relational distillation as described in this invention, characterized in that it comprises: The data acquisition module is used to acquire pre-trained teacher diffusion models; The stage division module is used to divide the reverse denoising process into three stages: high noise, medium noise, and low noise, according to time steps. The progressive distillation module is used to halve the number of sampling steps one by one, and select the corresponding differential loss function to update the model according to the stage to which the time step belongs during the distillation process; The relational knowledge transfer module is used to introduce pairwise similarity constraints between samples in the extremely low-step distillation stage, and jointly optimize it with classifier feature distillation; The output module is used to output the final student model with fewer steps.

[0017] This invention significantly improves the quality of diffusion model generation with extremely few steps by introducing a stage-adaptive differentiated loss function and knowledge transfer of inter-sample relationships into a progressive distillation framework. This method overcomes two fundamental limitations of traditional distillation methods, providing a new technical solution for the efficient deployment of diffusion models. Attached Figure Description

[0018] Figure 1 This is a diagram illustrating the overall framework of the Progressive Multistage Relational Distillation (PMRD) method of this invention. Detailed Implementation

[0019] The technical solutions of the embodiments of the present invention will be clearly and completely described below with reference to the accompanying drawings. Obviously, the described embodiments are only a part of the embodiments of the present invention, and not all of them. All other embodiments obtained by those skilled in the art based on the embodiments of the present invention without creative effort are within the scope of protection of the present invention.

[0020] Figure 1 This is a diagram illustrating the overall framework of the Progressive Multistage Relational Distillation (PMRD) method provided by this invention. The following is in conjunction with... Figure 1 This paper provides a detailed description of one specific implementation of the method.

[0021] like Figure 1 As shown, the diffusion model acceleration method based on progressive multi-stage relation distillation of the present invention mainly includes two main stages: the first stage is the progressive multi-stage distillation process, and the second stage is the relation knowledge transfer and fine-tuning process.

[0022] In this embodiment, the following operations need to be performed before the first stage: (1) Use the dataset CIFAR-10 or ImageNet64×64 as the training dataset; Based on this dataset, a pre-trained teacher diffusion model is obtained: the teacher model uses the DDIM sampler with a total of N sampling steps, where N≥128, and preferably N=1024. This teacher model has complete multi-step diffusion generation capability and can serve as a "knowledge source" for subsequent distillation processes, providing supervision signals for the diffusion process under multiple noise levels.

[0023] (2) In the first stage, the pre-trained teacher diffusion model is used as supervision, and the student model is distilled in three stages to compress the teacher model into an intermediate student model. Each stage is divided into denoising stages based on different noise levels (corresponding to diffusion time steps t), and a stage-differentiated distillation loss function is designed: the inverse denoising process of the teacher model is divided into three stages according to time steps—high noise stage, medium noise stage, and low noise stage. Let the total number of steps in the current distillation stage of the student model be... The division method is as follows: 1) High-noise stage: This stage corresponds to the time step interval with high noise levels during the diffusion process. The student model learns the teacher model's basic denoising ability in high-noise scenarios. The network mainly processes pure noise, resulting in high prediction uncertainty. The simplest mean squared error (MSE) loss is used as the distillation loss function, and the loss alignment distribution is as follows: ; in, This represents the set of batch sample indices corresponding to the high-noise phase. The number of samples in this set. The noise reduction result of the student model's prediction for the i-th sample. The target denoising results provided for the teacher model Represents the L2 norm; 2) Medium noise stage: Features change rapidly, and the loss function consists of both MSE loss and cosine similarity loss. Cosine similarity loss is used to constrain the correctness and structural consistency between the student and teacher models, thereby enhancing the semantic structure preservation ability of generated samples. A hybrid loss of MSE and cosine similarity is employed. ; in, , This is the set of batch sample indexes corresponding to the medium-noise stage. ;symbol" " represents the vector dot product; 3) Low noise stage In the low-noise region (near the stage of generating sharp images), the loss function employs a joint form of MSE loss and L1 loss. L1 loss provides stronger pixel-level constraints, helping the student model recover fine textures and edge details in low-noise areas. Since the features are highly structured and subtle differences have a significant impact on the final quality, a combination of MSE and L1 loss is used. ; in , , This is the set of batch sample indexes corresponding to the low-noise phase. This represents the L1 norm.

[0024] (3) Perform the first stage of progressive distillation: Starting with the teacher model, the number of sampling steps is halved successively, from N steps to N / 2 steps of distillation. In each distillation step, the corresponding stage loss function is dynamically selected according to the stage to which the current time step belongs to calculate the distillation loss and update the student model parameters.

[0025] (4) Repeat step (3) until an intermediate student model with 8 steps is obtained.

[0026] (5) After obtaining the 8-step student model, the second stage of distillation is performed. The model is further distilled to a very low number of steps (4 steps, 2 steps, 1 step) through relational knowledge transfer loss. In this process, similarity-preserving loss (SP Loss) is added to the classifier feature distillation (CFD) to force the student model to maintain the pairwise similarity structure between samples in the feature space of the teacher model.

[0027] By introducing classifier feature distillation (CFD) loss and utilizing the intermediate feature layer of the pre-trained image classifier, the consistency of images generated by the student model and the teacher model at the high-level semantic features is supervised, ensuring that the compressed model can still generate images with correct semantic information and avoiding semantic bias caused by step compression.

[0028] By introducing similarity preservation loss (SP Loss), the student model is forced to learn and retain the similarity relationship between the teacher model in the sample space through pairwise similarity constraints between samples. This avoids the loss of sample distribution relationship during compression and ensures the diversity and distribution consistency of the generated images.

[0029] For the B samples within the same batch, calculate the teacher characteristics separately. and student characteristics Normalized pairwise similarity matrix: ; in Similarity Preservation Loss Defined as the mean squared error between two similarity matrices: ; Joint optimization of classifier feature distillation and similarity preservation loss: In the very low-step distillation stage, the total loss function is: ; in The loss is calculated by distilling features for the classifier (extracting features from a pre-trained classifier and calculating the KL divergence). To maintain the loss weight for similarity, a value of 1.0 is preferred.

[0030] (6) Based on the time step of the current student model For each high-noise, medium-noise, or low-noise stage, progressive distillation steps (3) to (5) are performed, and the corresponding stage-differentiated loss function is dynamically selected to update the model parameters, outputting the final student model. This model can generate images of comparable or better quality to the teacher model within 1 to 8 steps.

[0031] The method of this invention significantly improves the efficiency of knowledge distillation by dividing the denoising process into three stages and using differential loss, as well as introducing knowledge transfer of relationships between samples. It achieves generation quality (FID and IS metrics) that is superior to existing methods with extremely few steps (1 step, 2 steps, 4 steps, 8 steps).

[0032] Experimental results show that on the CIFAR-10 dataset, the FID generation efficiency can be reduced to 8.12 in 1 step; on the ImageNet 64×64 dataset, the FID generation efficiency of 8 steps (8.31) even surpasses the 1024-step DDIM benchmark (8.46), achieving a sampling speedup of more than 128 times while obtaining better generation quality.

[0033] The contents not described in detail in this specification are existing technologies known to those skilled in the art.

[0034] It should be noted that the above description of the examples of the present invention is for the purpose of explaining the technical features of the present invention in detail. Various improvements and modifications made without departing from the present invention are also protected by the present invention; therefore, the scope of protection of the present invention should be determined by the content defined in the claims of this application.

Claims

1. A diffusion model acceleration method based on progressive multi-stage relational distillation, characterized in that, Includes the following steps: (1) Obtain a pre-trained teacher model, wherein the teacher model uses the DDIM sampler and the total number of sampling steps is N, where N≥128; (2) The reverse denoising process of the teacher model is divided into three stages according to the time step: high noise stage, medium noise stage and low noise stage, and a differentiated distillation loss function is designed for each stage. (3) Perform the first stage of progressive distillation, distilling the teacher model from N steps to N / 2 steps. During the distillation process, dynamically select the corresponding stage loss function according to the stage to which the current time step belongs to calculate the distillation loss and update the student model parameters. (4) Repeat step (3) until an intermediate student model with 8 steps is obtained; (5) In the process of distilling from 8 steps to a very low number of steps, relational knowledge transfer loss is introduced, and consistency constraints on the pairwise similarity structure between samples are imposed between the student model and the teacher model, and jointly optimized with the classifier feature distillation loss. (6) Based on the time step of the current student model For each of the high-noise, medium-noise, or low-noise stages, the progressive distillation steps (3) to (5) are performed respectively. The corresponding stage-differentiated loss function is dynamically selected to update the model parameters and output the final student model.

2. The method according to claim 1, characterized in that, The three stages described in step (2) are divided as follows: the high-noise stage corresponds to... The middle noise stage corresponds to Low noise stage corresponds to ,in This represents the total number of steps in the current distillation stage of the student model. Provide students with time for their studies.

3. The method according to claim 2, characterized in that, The differentiated distillation loss function mentioned in step (2) is specifically as follows: Mean square error loss is used in the high-noise stage: ； in, This represents the set of batch sample indices corresponding to the high-noise phase. The number of samples in this set. The noise reduction result of the student model's prediction for the i-th sample. The target denoising results provided for the teacher model Represents the L2 norm; In the mid-noise stage, a hybrid loss method combining mean square error and cosine similarity is used: ； in, , This is the set of batch sample indexes corresponding to the medium-noise stage. ;symbol" " represents the vector dot product; The low-noise stage employs a combination of mean square error and L1 loss: ； in , , This is the set of batch sample indexes corresponding to the low-noise phase. This represents the L1 norm.

4. The method according to claim 1, characterized in that, The relational knowledge transfer loss mentioned in step (5) is the similarity preservation loss, which is calculated as follows: for B samples in the same batch, calculate the teacher features respectively. and student characteristics Normalization to the similarity matrix ;in Similarity Preservation Loss Function Defined as the mean square error between two similarity matrices: .

5. The method according to claim 4, characterized in that, The total loss function of the joint optimization in step (5) is: ; in For the classifier feature distillation loss, Loss weights are maintained for similarity.

6. The method according to claim 5, characterized in that, The extremely low number of steps mentioned in step (5) includes 4 steps, 2 steps and 1 step. The distillation process is carried out step by step: first distilling from 8 steps to 4 steps, then distilling to 2 steps, and finally distilling to 1 step.

7. The method according to claim 1, characterized in that, The teacher diffusion model is a U-Net architecture diffusion model pre-trained on the CIFAR-10 or ImageNet 64×64 dataset, with a total sampling step of 1024 steps.

8. A system for performing the diffusion model acceleration method based on progressive multi-stage relational distillation as described in any one of claims 1-7, characterized in that, include: The data acquisition module is used to acquire pre-trained teacher diffusion models; The stage division module is used to divide the reverse denoising process into three stages: high noise, medium noise, and low noise, according to time steps. The progressive distillation module is used to halve the number of sampling steps one by one, and select the corresponding differential loss function to update the model according to the stage to which the time step belongs during the distillation process; The relational knowledge transfer module is used to introduce pairwise similarity constraints between samples in the extremely low-step distillation stage, and jointly optimize it with classifier feature distillation; The output module is used to output the final student model with fewer steps.

Images