Diffusion model inference acceleration method based on optimal time step sequence search and knowledge distillation

Abstract

The application discloses a diffusion model reasoning acceleration method based on optimal time step sequence search and knowledge distillation, which constructs a unified search space covering time step sequences and corresponding model architecture configurations without overall fine-tuning of a pre-trained diffusion model, and searches for optimal time step sequences. Based on the optimal diffusion time step sequences obtained through the search, a knowledge distillation training method is designed, a main distillation loss function is constructed within a discrete time step subspace defined by the optimal diffusion time step sequences, and an adjacent time step consistency loss is introduced to relieve the discrete error introduced by a non-uniform time step span. In the reasoning stage, the teacher model is frozen, and the distilled student model is combined with an efficient sampler ACDMS based on UniPC and fusing flow matching dynamic correction to perform limited step sampling on the optimal time step sequence, so that reasoning acceleration is realized under the condition of keeping generation quality.

Description

Technical Field

[0001] This invention relates to the field of generative artificial intelligence and deep generative model optimization technology, and in particular to a method for accelerating diffusion model inference based on optimal time step sequence search and knowledge distillation. Background Technology

[0002] Diffusion models have become one of the core frameworks of current generative models, achieving leading performance in fields such as image generation, video generation, material design, and molecular generation. Traditional diffusion models require hundreds to thousands of inverse diffusion time steps to generate a high-quality sample during the inference phase, making their inference speed much slower than models such as Generative Adversarial Networks (GANs) and Variational Autoencoders (VAEs), which significantly limits their application in real-time generation and edge device inference.

[0003] In recent years, research on accelerating diffusion models has mainly focused on two directions: time-step scheduling optimization, which reduces the number of back-diffusion steps manually or heuristically, but lacks interpretability and global optimum guarantees; and knowledge distillation, which optimizes student models while maintaining the performance of teacher models, enabling them to reconstruct approximate distributions in fewer time steps. However, existing methods typically perform discrete distillation at fixed time steps, failing to utilize time-step search for spatial structure or introduce time-domain consistency constraints, resulting in distribution shifts or decreased generation quality after distillation. Furthermore, while fast sampling methods such as DDIM, UniPC, and DPM-Solver can reduce the number of sampling steps, the generation quality deteriorates significantly with very few sampling steps because they fail to simultaneously balance numerical integration accuracy and distribution consistency.

[0004] Therefore, there is an urgent need for a method that can automatically search for the optimal time step structure, perform explicit teacher-student time domain distillation based on the search structure, and combine a high-order adaptive sampler to achieve fast and high-quality sampling, thereby significantly accelerating the inference stage while maintaining the teacher model generation capability. Summary of the Invention

[0005] The problem to be solved by this invention is to provide a diffusion model inference acceleration method based on optimal time step sequence search and knowledge distillation. By accelerating the process through time step search, explicit multinomial loss distillation training and high-order samplers, the number of diffusion inverse inference steps can be significantly reduced while maintaining or even improving the generation quality.

[0006] This invention adopts the following technical solution: a diffusion model inference acceleration method based on optimal time step sequence search and knowledge distillation, comprising the following steps:

[0007] Step 1: Generative vision task based on diffusion model. Establish a fixed pre-trained diffusion generation model and corresponding sampler. While keeping the parameters of the pre-trained diffusion generation model unchanged, based on the set of discrete time steps of the original diffusion process, introduce the inference model architecture configuration corresponding to each time step, and construct a unified search space covering the time step sequence and the corresponding inference model architecture configuration.

[0008] Step 2: Within the unified search space, with the sampling efficiency and generation quality of the generative image reasoning stage as optimization objectives, under the condition of a fixed length of time step, a neural network architecture is used to search for the optimal time step sequence that satisfies the constraints.

[0009] Step 3: Based on the optimal time step sequence, construct a knowledge distillation training process for the time step subspace. Use a pre-trained diffusion generation model as the teacher model and a diffusion generation model to be trained as the student model. While keeping the teacher model parameters frozen, perform explicit distillation training on the student model. In addition, introduce the consistency loss between adjacent time steps to alleviate the discrete error introduced by the non-uniform time step span.

[0010] Step 4: During the inference phase, freeze the teacher model, use the student diffusion generation model trained by distillation and combine it with the efficient diffusion sampler ACDMS to perform finite-step sampling on the optimal time step sequence, thereby accelerating inference while maintaining generation quality.

[0011] As a preferred embodiment, step 1 includes:

[0012] Step 1.1: Represent the complete set of diffusion time steps of the pre-trained diffusion generation model as follows: The forward diffusion process can be represented in standard parametric form as follows:

[0013] ;

[0014] In the formula, This is the cumulative decay factor; Indicates the original clean sample; Indicates and A Gaussian noise vector with consistent dimensions; Indicates the first The noisy samples obtained at each diffusion time step are used as inputs to the diffusion model during the inference or training phase. It is the identity matrix. Let represent a matrix with a mean of 0 and a covariance of identity. The multidimensional Gaussian distribution;

[0015] Step 1.2: Set the search time step sequence length to a fixed value. The process of selecting time steps in a single diffusion inference process is abstracted into a process by... The time-step sequence encoding composed of discrete decision variables is defined as:

[0016] ;

[0017] In the formula, Represents a time step sequence, by each discrete time step Composition, decision variables at each time step The range of values ​​is the complete set of diffusion time steps. ,Right now To ensure that each candidate time step sequence corresponds to a unique and valid diffusion sampling path, the following constraints are imposed on the time step sequences:

[0018] ;

[0019] Step 1.3: Based on the time step sequence search, introduce the inference model architecture configuration corresponding to each time step, and take the sub-network structure of the noise prediction network at different time steps in the diffusion inference process as the searchable object to obtain a unified search space.

[0020] For each time step in the time step sequence Define the corresponding model architecture and configure it as follows:

[0021] ;

[0022] In the formula, Indicates at time step The selected noise prediction network substructure, Indicates the structure of each layer, superscript This indicates the number of network layers enabled at that time step, and The total number of layers in the pre-trained noise prediction network shall not exceed the total number of layers.

[0023] Under the above definition, the candidate architecture in a single diffusion inference process can be represented as:

[0024] .

[0025] Then, under the condition of satisfying the computational complexity constraint, a joint search is performed on the time step sequence and the corresponding model architecture configuration. The resulting neural network architecture search space is used to obtain the optimal diffusion time step sequence.

[0026] As a preferred embodiment, step 2 includes:

[0027] Step 2.1: Set the complete diffusion time steps used by the pre-trained diffusion generation model. This serves as the initial set of time steps; the length of the time step sequence remains unchanged during the search process, only the values ​​of the time steps and their combinations are optimized.

[0028] Step 2.2: Based on multi-objective indicators For different time step sequences The evaluation criteria include at least the FID (Focus on Generation Quality), the number of sampling steps (Sampling Efficiency), and the inference latency (Inference Speed).

[0029] Step 2.3: Sequence of time steps As structural variables input into the neural network architecture search framework, the optimal time-step sequence satisfying the preset convergence condition is obtained through multi-objective optimization. .

[0030] As a preferred option, the multi-objective index , is represented as:

[0031] ;

[0032] For time step-based subsequences The generation quality index obtained from the generated samples is used to measure the degree of difference between the generated samples and the distribution of the real samples:

[0033] ;

[0034] In the formula, The function represents the Fréchet Inception Distance evaluation function, which measures the statistical distance between the generated sample distribution and the real sample distribution in the feature space; Indicates the time step sequence Under constraints, the first The final generated image sample output by the diffusion generation model during the secondary diffusion inference process; This represents a reference dataset composed of real image data, used as a true distribution benchmark for generating quality assessments;

[0035] To use time step subsequence The total number of sampling steps during the sampling process is used to characterize the sampling efficiency.

[0036] ;

[0037] For time step subsequence The average inference latency for completing one full sampling process is used to characterize inference speed:

[0038] ;

[0039] In the formula, The function represents a time-step sequence on a fixed computing hardware platform. The actual reasoning time required to complete the diffusion reasoning generation; This represents the statistical average of the reasoning delays obtained from multiple independent reasoning processes.

[0040] As a preferred embodiment, step 3 includes:

[0041] Step 3.1: According to the preset time step sampling strategy, sample from the optimal time step sequence. Uniform random sampling of a single time step and the time step As the diffusion time step used in the current training iteration;

[0042] Step 3.2: Under the same noise and time step conditions, noise prediction and intermediate denoising results are obtained through the teacher model and student model respectively. The main distillation loss is constructed to constrain the noise prediction of the student model to approximate the noise prediction of the teacher model, and a gradient-aware regularization term is introduced to suppress high-frequency noise.

[0043] Step 3.3: From the sorted optimal time step sequence In the continuous sampling of two adjacent index time steps Using the same noise conditions, construct the corresponding noisy latent representation and compute the student model at the two time steps. and The reconstruction result is represented as Constructing consistency loss The constraints ensure that the reconstruction results remain consistent across the time domain; , For time steps , The corresponding noisy latent representation;

[0044] Step 3.4: The main distillation loss and the consistency loss are weighted and summed according to preset weights to obtain the total distillation loss of the student model. Backpropagation and parameter updates are performed on the student diffusion generation model.

[0045] During the distillation training process, the parameters of the teacher model are kept frozen, and only the parameters of the student model are updated. This achieves acceleration and quality enhancement of the inference stage without changing the performance of the original pre-trained diffusion generation model.

[0046] Specifically, only in the time step subsequence Sampling of the time step within the defined discrete time step subspace, i.e., the time step random variable from The sampling is performed according to a preset strategy to guide students to focus on learning the time step mapping relationship actually used in the reasoning stage.

[0047] The mean square error term of the main distillation loss is:

[0048] ;

[0049] In the formula, This represents the distillation loss value based on noise prediction consistency, used to measure the degree of difference between the prediction results of the student diffusion model and the prediction results of the teacher diffusion model at a given time step; Indicates time step and noisy random variables The mathematical expectation; This is a noisy latent representation obtained from the diffusion process; and The teacher model and the student model are respectively at the time step Noise prediction.

[0050] To suppress high-frequency noise and enhance structural consistency, a gradient-aware regularization term is further constructed. This is achieved by calculating the gradient difference between the noise prediction feature maps and the input. The gradients are respectively:

[0051] ;

[0052] Gradient-aware regularization Defined as:

[0053] ;

[0054] In the formula, Indicates noisy samples The gradient calculation operator for partial derivatives reflects the sensitivity of the noise prediction result to changes in the input sample; This indicates that the teacher model is at time step The noise prediction result at the output is relative to the noisy input sample. Gradient characteristics; Indicates the student model at time step The noise prediction result at the output is relative to the noisy input sample. The gradient characteristics.

[0055] Therefore, the primary distillation loss is expressed as:

[0056] ;

[0057] In the formula, and These are weighting coefficients used to balance the contributions of different distillation constraint terms. Used to control the proportion of noise prediction consistency constraints in the total loss. Used to control the influence of gradient-aware regularization on high-frequency noise suppression and structural smoothness, satisfying ;

[0058] In addition, to mitigate the time-domain discretization error caused by non-uniform time step intervals, a reconstruction consistency loss based on adjacent time steps is constructed to improve the generation stability over the entire time domain.

[0059] The consistency loss is measured by the student model at time step. and The intermediate denoising results generated above and The difference is implemented to constrain the student model to maintain the continuity and smoothness of the predicted and reconstructed trajectories at adjacent time steps, thereby improving the generation stability over the entire time domain. This constraint is consistent with the core idea of ​​Flow Matching, which uses a smooth vector field to make the state evolve continuously over time. The specific construction method is as follows:

[0060] set up For index-adjacent time step pairs Under the same noise conditions, with the same noise seed and the same Construct the noisy latent representation corresponding to the two steps. :

[0061] ;

[0062] ;

[0063] Reconstruction Estimation Based on Student Model Define the consistency loss, and define the reconstruction estimate at each step as follows:

[0064] ;

[0065] The consistency loss can be expressed as:

[0066] ;

[0067] In the formula, This indicates that the diffusion generation model is at time step The corresponding cumulative noise attenuation factor; Indicates at time step The corresponding cumulative noise attenuation factor; This represents the optimal time step sequence obtained through the time step sequence search phase. Represents the optimal time step sequence The statistical average of the consistency loss of adjacent time steps in the index and under different noise sample conditions.

[0068] Ultimately, the total distillation loss of the student model is the weighted sum of the principal distillation loss and the consistency loss:

[0069] ;

[0070] In the formula, The weight coefficients for consistency loss are used to perform backpropagation and parameter updates on the student model parameters based on this loss function. During the distillation training process, the teacher model parameters remain frozen, and only the parameters of the student diffusion generation model are updated.

[0071] As a preferred embodiment, the ACDMS is a high-order prediction and correction style diffusion sampling method. It uses the high-order integral update of UniPC (a unified numerical solution framework for prediction and correction of diffusion model sampling) as the backbone and introduces a distribution corrector based on flow matching to alleviate the distribution shift under the condition of few steps sampling.

[0072] Specifically, the high-efficiency sampler ACDMS is a high-order diffusion sampler that, for a given optimal time step sequence... The system performs multi-step prediction-correction numerical integration, which achieves an approximate solution to the diffusion process by adaptively updating a small number of time steps, thereby reducing the number of sampling steps and inference overhead while maintaining or improving the generation quality.

[0073] In practical implementation, ACDMS can integrate a UniPC-style high-order predictor-corrector scheme with a flow-matching-based distribution correction term to improve the distribution consistency of small-step sampling.

[0074] Specifically, the time step sequence is sorted from largest to smallest. The diffusion integral update is performed on this sequence, thereby achieving high-quality image generation and inference acceleration under the constraint of a fixed number of sampling steps.

[0075] Let the time step between adjacent times be... The student model at time step The noise prediction is Then, the one-step update process of ACDMS can be represented as:

[0076] ;

[0077] In the formula, For time step The corresponding noisy latent representation, The time step difference between adjacent times. For the student model at time step Noise prediction, For order, The first generation generated for UniPC Order correction term, For FlowMatching, the corrected vector field, To employ the UniPC prediction-correction mechanism at time step The intermediate prediction results obtained from the above are as follows: To correct the strength coefficient.

[0078] Compared with the prior art, the present invention, employing the above technical solution, has the following technical effects:

[0079] 1. This invention uses time step sequences as the core controllable variables. Without retraining or modifying the original diffusion model, it can efficiently accelerate existing diffusion generation models and inference frameworks, and has good engineering transferability and versatility.

[0080] 2. The method of the present invention significantly reduces the number of diffusion inverse inference steps by combining time-step search, explicit multinomial loss distillation training and high-order sampler acceleration. It can still generate high-quality images with very few inference steps, which can greatly improve the efficiency of the diffusion model in the inference stage, while maintaining or even improving the generation quality. Attached Figure Description

[0081] Figure 1 This is a flowchart of the diffusion model reasoning acceleration method based on optimal time step sequence search and knowledge distillation of the present invention;

[0082] Figure 2 This is a schematic diagram of the optimal time step sequence generation process based on neural network architecture search according to the present invention;

[0083] Figure 3 This is a flowchart of the knowledge distillation training process for the teacher-student diffusion model based on the optimal time step sequence of this invention.

[0084] Figure 4 This is a schematic diagram of the structure of the ACDMS high-efficiency cooperative sampler for fast inference of diffusion models according to the present invention;

[0085] Figure 5 This is a comparison chart of the FID performance of the generated samples by the method of this invention and Stable Diffusion v1.4 under the same conditions;

[0086] Figure 6 This is a comparison chart of the IS performance of the generated samples using the method of this invention and Stable Diffusion v1.4 under the same conditions. Detailed Implementation

[0087] To make the objectives, technical solutions, and advantages of this invention clearer, the technical solutions of the application will be further described in detail below with reference to the accompanying drawings. The described embodiments are only a part of the embodiments involved in this invention. All non-innovative embodiments based on these embodiments by other researchers in the art are within the protection scope of this invention. Furthermore, the step numbers in the embodiments of this invention are only set for ease of explanation and do not limit the order of the steps. The execution order of each step in the embodiments can be adaptively adjusted according to the understanding of those skilled in the art.

[0088] In one embodiment of the present invention, a diffusion model inference acceleration method based on optimal time step sequence search and knowledge distillation is described in the following process: Figure 1 As shown, it includes the time step search stage, the distillation training stage, and the fast reasoning stage in sequence.

[0089] During the time step search phase, a pre-trained diffusion generation model and its corresponding sampler are loaded to construct a unified search space consisting of the original discrete time step sequences and their corresponding inference model architectures. A multi-objective evolutionary optimization method is then used to search for candidate time step sequences. During the search process, multi-dimensional performance evaluation is performed on the candidate schemes, taking into account indicators such as generation quality, number of sampling steps, and inference latency, so as to obtain the optimal time step sequence that meets different quality and efficiency trade-off requirements.

[0090] During the distillation training phase, the student model is initialized while keeping the teacher model parameters frozen. Time steps are selected from the time step subspace corresponding to the optimal time step sequence. Random noise and corresponding conditional information are simultaneously input into the teacher model and student model under the same noise conditions for forward inference. Based on the prediction results of the teacher model and student model at each time step, distillation loss and adjacent time step consistency loss are constructed. The student model is trained and updated by combining the loss functions until the distillation process converges.

[0091] In the rapid inference phase, the teacher model is frozen, and only the student diffusion generation model optimized by distillation is used. Under the condition of input random noise and corresponding conditional information, a finite number of back diffusion sampling steps are performed in combination with the optimal time step sequence and efficient sampler to output the final generation result. This accelerates the diffusion model inference process while maintaining stable generation quality.

[0092] In this embodiment, the pre-trained diffusion generation model uses Stable Diffusion v1.4 and adopts a matching sampling framework; throughout the search phase, the parameters of the teacher model (i.e., the pre-trained model) remain frozen and are not updated.

[0093] In the diffusion model, the time step refers to the discrete index of the reverse denoising iteration, used to control the noise intensity range corresponding to each step. Traditional methods often use equal intervals or manual rules to extract time steps, making it difficult to balance quality and speed. This embodiment treats the time step sequence as a searchable discrete structural variable, and obtains a more suitable time step schedule for fewer sampling steps through search without modifying the model parameters.

[0094] Furthermore, to construct the search space of the neural network architecture, a fixed pre-trained diffusion generation model and its corresponding sampler are provided. Under the premise of keeping the parameters of the pre-trained diffusion generation model unchanged, the diffusion inference process is searchably modeled based on the complete set of discrete time steps of the original diffusion process and the corresponding model structure.

[0095] Let the complete set of diffusion time steps of the pre-trained diffusion generation model be . The forward diffusion process is represented in standard parametric form as follows:

[0096] ;

[0097] In the formula, Indicates the original clean sample. This is the cumulative decay factor. , is the decay factor at step s.

[0098] Furthermore, the length of the search time step sequence is set to a fixed value. The scheduling process of time steps in a single diffusion inference process is abstracted into a process by... The encoding of an ordered time-step sequence composed of discrete decision variables is defined as follows: .

[0099] Among them, the decision variables at each time step The range of values ​​is the complete set of diffusion time steps. ,Right now .

[0100] To ensure that each candidate time step sequence corresponds to a unique and valid backdiffusion sampling path, the following constraints are imposed on the time step sequence:

[0101] ;

[0102] The above constraints ensure that the candidate time step sequences are executed in descending order of noise intensity during the actual inference phase, thus maintaining consistency with the reverse denoising process of the diffusion model.

[0103] Furthermore, based on the time-step sequence search, an inference model architecture configuration corresponding to each time step is introduced, and the sub-network structures activated by the noise prediction network at different time steps during the diffusion inference process are taken as searchable objects.

[0104] Specifically, for each time step in the time step sequence Define the corresponding inference model architecture subspace and configure it as follows:

[0105] ;

[0106] In the formula, Indicates at time step The selected noise prediction network substructure, superscript This indicates the number of network layers enabled at that time step, and The total number of layers in the pre-trained noisy prediction network shall not exceed the total number of layers.

[0107] A multi-objective optimization framework is used to evaluate and select candidate sequences, where each candidate is represented as:

[0108] ;

[0109] Under the above definition, the candidate architecture in a single diffusion inference process Represented as:

[0110] .

[0111] In the formula, To satisfy the constraints, the time step sequence Configure the noise prediction network substructure corresponding to each time step.

[0112] Furthermore, such as Figure 2 As shown, under the condition of satisfying computational complexity constraints, a joint search is performed on the time step sequence and the corresponding model architecture configuration. Therefore, a unified search space is constructed based on the time step sequence subspace and the inference model architecture subspace to obtain the optimal diffusion time step sequence, providing a foundation for subsequent distillation training and efficient sampling.

[0113] After completing the construction of the unified search space, multi-objective optimization is performed, and an evolutionary algorithm is used to iteratively search in the candidate time step sequence space.

[0114] The search process begins with initializing the candidate diffusion inference architecture population, followed by updating the population individuals through genetic operators to achieve target awareness. The genetic operators include crossover and mutation operations, which are used to exchange or perturb the time step index and sub-network structure configuration between different candidate time step sequences and their corresponding model architectures, thereby enhancing the diversity and exploration capabilities of the search process.

[0115] To accommodate the multidimensional trade-offs between generation quality, sampling efficiency, and inference speed in the diffusion model, this embodiment defines a multi-objective evaluation vector for each candidate time step subsequence τ to perform multidimensional performance evaluation, including FID evaluation, efficiency evaluation, and latency evaluation.

[0116] ;

[0117] In the formula, To represent the quality target, this embodiment uses FID (Fixed Distance Identifier) ​​for quantification. FID evaluation measures the statistical distance between the generated sample distribution and the real sample distribution in the feature space; the smaller the value, the closer the generated distribution is to the real distribution. The calculation formula is as follows:

[0118] ;

[0119] In the formula, and Let represent the mean and covariance matrices of the true sample feature distribution, respectively. , These represent the mean and covariance matrices of the feature distribution of the generated samples, respectively. The features of the generated samples and the real samples are obtained by a pre-trained feature extraction network.

[0120] The sampling efficiency target is defined as the length of the time step sequence:

[0121] ;

[0122] The inference speed objective, i.e., the Pareto optimal time step sequence, is defined as the sequence of time steps... Average inference latency for completing one full sampling process:

[0123] ;

[0124] In the formula, Indicates the first The time required to perform one complete sampling under sequence τ. This represents the number of times the measurement was repeated.

[0125] Through the aforementioned multi-objective evaluation, the evolutionary algorithm performs non-dominated sorting of candidate sequences and forms a Pareto set (elite archive). Pareto selection and feedback are then performed using a distance metric, ultimately outputting a set of optimal time-step sequences that satisfy different quality-velocity tradeoffs. In practical applications, the final time-step sequence can be selected from the Pareto set according to deployment requirements.

[0126] Obtaining the optimal time step sequence After that, the distillation training phase begins, and the overall process is as follows: Figure 3 As shown. In this stage, based on the obtained optimal time step sequence... The student model is trained by explicit distillation to obtain the optimal student subnetwork. The parameters are updated for the generated samples through loss calculation to achieve performance evaluation.

[0127] In this embodiment, distillation refers to using the output of the teacher model Unet as a supervision signal to guide the student model to learn the denoising behavior of the teacher model at a specific time step; the teacher model parameters remain frozen, while the student model parameters are updated through backpropagation. The key to this embodiment is that distillation training no longer covers the original complete time step space, but only... The training is conducted within the defined discrete time step subspace, allowing the student model to focus on learning the time steps that are actually used in the inference phase, thereby improving training effectiveness and reducing the learning burden caused by irrelevant time steps.

[0128] The forward diffusion process of the diffusion model is used to construct noisy samples, and its standard parameterization form is as follows:

[0129] ;

[0130] In the formula, For the original sample, In time step The noisy representation, This is the cumulative decay factor for the diffusion process.

[0131] Teacher model and student model at time step The above outputs noise predictions respectively and Based on the differences between the two, the primary distillation loss is constructed, and its form is as follows:

[0132] ;

[0133] In the formula, This represents the distillation loss value based on noise prediction consistency, used to measure the degree of difference between the prediction results of the student diffusion model and the prediction results of the teacher diffusion model at a given time step; This indicates sampling time steps from a discrete set of the optimal time step sequence. Indicates time step and noisy random variables The mathematical expectation, and The teacher model and the student model are respectively at the time step Noise prediction.

[0134] In particular, this embodiment preferably employs a uniform random sampling strategy, that is, during training, sampling is not performed according to... Instead of sequential traversal, it performs a process within each batch. Perform uniform random sampling so that the student model is used throughout the process. The teacher model behavior is evenly aligned across the covered time domain, thereby improving training stability and generalization ability.

[0135] because Typically, time step sequences are non-uniformly spaced, with potentially large spans between adjacent time steps. Insufficient sampling can easily introduce discretization errors and lead to unsmooth reconstructed trajectories. Therefore, this embodiment further introduces an adjacent time step consistency loss, assuming... For adjacent time step pairs Apply constraints.

[0136] First, the student model at time step is given. Reconstruction estimate That is, using noise prediction to inversely deduce the estimate of clean samples:

[0137] ;

[0138] Then, under the same noise conditions, i.e., sharing the same noise seed and the same Construct the corresponding time steps for each adjacent time step and And based on the above reconstruction definition of consistency loss:

[0139] ;

[0140] Ultimately, the total distillation loss of the student model is the weighted sum of the principal distillation loss and the consistency loss:

[0141] ;

[0142] In the formula, The weighting coefficients for consistency loss are used to balance between strictly aligned teacher noise predictions and smooth reconstruction at adjacent time steps.

[0143] During the training process Backpropagation and parameter updates are performed on the student model for the objective function, while the teacher model parameters are kept frozen, thus obtaining a student model that is more suitable for few-step inference under the premise of fixed teacher capabilities.

[0144] In the rapid reasoning stage, such as Figure 4 As shown, the teacher model parameters are no longer updated during inference. The actual generation only activates the distilled student model and performs finite-step sampling on the optimal time step sequence.

[0145] To further improve the numerical accuracy and distribution consistency of few-step sampling, this embodiment employs a self-developed ACDMS high-efficiency collaborative sampler. The ACDMS high-efficiency collaborative sampler performs adaptive sampling based on the optimal time-step sequence. During parameter updates at each time step, it comprehensively utilizes the noise prediction results (generated samples) of the student model at the current time step to construct a high-order numerical update form containing higher-order prediction terms and a vector field of the vector corrector. This allows the sampling trajectory to still approximate a continuous diffusion process even over a large time-step span.

[0146] Specifically, in time step First, a higher-order prediction term is constructed based on the noise prediction results output by the student model. A multi-order correction term is then introduced to predict the potential representation for the next time step, thereby improving the accuracy of the back-diffusion numerical integral. Building upon this, a vector corrector field based on the Flow Matching concept is further introduced to correct the prediction results at the distribution level. The vector corrector estimates the flow field direction of the generated distribution between adjacent time steps, guiding the sampling trajectory towards the target generated distribution to alleviate the distribution offset problem that may occur under conditions of few sampling steps.

[0147] Under the synergistic effect of the aforementioned higher-order prediction and distribution correction, the ACDMS sampler achieves better results in the sorted time-step sequence. The up-execution update, whose general step-size update form is represented by up-execution sampling, and whose one-step update form is represented by:

[0148] ;

[0149] In the formula, For time steps The corresponding noisy latent representation, The time step difference between adjacent times. For the student model at time step Noise prediction, For order, The first generation generated for UniPC Order correction term, For FlowMatching, the corrected vector field, To employ the UniPC prediction-correction mechanism at time step The intermediate prediction results obtained from the above are as follows: To correct the strength coefficient.

[0150] By jointly introducing higher-order predictors and vector corrector vector fields in the same update step, the ACDMS sampler can effectively suppress numerical integration errors and distribution shifts while significantly reducing the number of diffusion sampling steps, thereby accelerating the diffusion model inference process while maintaining stable generation quality.

[0151] It should be noted that the diffusion model inference acceleration method based on optimal time step sequence search and knowledge distillation proposed in this invention is applicable to generative image generation and processing tasks with diffusion models as the core, and its application is computer vision systems that generate images by inverse denoising of random noise.

[0152] In practical implementation, the method of this invention can be applied to, but is not limited to, the following typical diffusion model application areas:

[0153] (1) Text-conditional image generation scenario, such as in a text-to-image diffusion generation system, a target image is generated based on a natural language description for use in digital content creation and generative content services;

[0154] (2) Conditional diffusion image generation and control generation scenarios, such as generating structurally controlled images under given semantic segmentation maps, depth maps and other conditional constraints, for image editing and image redrawing tasks;

[0155] (3) Multimodal conditional diffusion generation scenarios, such as jointly using text, images or other structured conditional information to generate images for human-computer interaction or generative aid design systems;

[0156] (4) Deployment scenarios of diffusion models that are sensitive to inference efficiency and latency, such as online generation services or efficient image generation in environments with limited computing power.

[0157] In the above application scenarios, the method of the present invention significantly reduces the amount of inference computation and response latency while maintaining stable generation quality by coordinating the time step scheduling, distillation training strategy and efficient sampling numerical update of the diffusion model in the inference stage. This improves the deployability and engineering application value of the diffusion model in actual systems.

[0158] To verify the effectiveness of the method of this invention, this embodiment uses Stable Diffusion v1.4 as the baseline method and conducts comparative experiments on the COCO2014 dataset. While maintaining consistency in model version, dataset, prompt word generation strategy, and evaluation protocol, each method generates 10K images, and the FID and IS indices are calculated under sampling steps of 4, 6, 8, and 10, respectively.

[0159] Experimental results are as follows Figure 5 and Figure 6 As shown, the method in this embodiment achieves better or more stable quality index performance with fewer sampling steps, namely, lower FID index value and higher IS index value, while significantly reducing inference latency. This verifies the effect of the three-layer collaborative strategy of time step search, distillation training and efficient sampler on accelerating the inference of diffusion model.

[0160] The above description is only a preferred embodiment of the present invention. It should be noted that for those skilled in the art, several improvements and modifications can be made without departing from the principle of the present invention, and these improvements and modifications should also be considered within the scope of protection of the present invention.

Claims

1. A method for accelerating inference in a diffusion model based on optimal time-step sequence search and knowledge distillation, characterized in that, Includes the following steps: Step 1: Generative vision task based on diffusion model. Establish a fixed pre-trained diffusion generation model and corresponding sampler. While keeping the parameters of the pre-trained diffusion generation model unchanged, based on the discrete time step set of the original diffusion process, introduce the corresponding inference model architecture configuration and construct a unified search space covering the time step sequence and the corresponding inference model architecture configuration. Step 2: Within the unified search space, with the sampling efficiency and generation quality of the generative image reasoning stage as optimization objectives, under the condition of a fixed length of time step, a neural network architecture is used to search for the optimal time step sequence that satisfies the constraints. Step 3: Based on the optimal time step sequence, construct a knowledge distillation training process for the time step subspace. Use a pre-trained diffusion generation model as the teacher model and a diffusion generation model to be trained as the student model. While keeping the teacher model parameters frozen, perform explicit distillation training on the student model. In addition, introduce the consistency loss between adjacent time steps to alleviate the discrete error introduced by the non-uniform time step span. Step 4: During the inference phase, freeze the teacher model, use the student diffusion generation model trained by distillation and combine it with the efficient diffusion sampler ACDMS to perform finite-step sampling on the optimal time step sequence, thereby accelerating inference while maintaining generation quality. The diffusion sampler ACDMS uses high-order integral updates of UniPC as its backbone and introduces a flow-matching-based distribution corrector to mitigate distribution shift under low-step sampling conditions. Sort the time step sequences in descending order. ACDMS updates the sorted sequence using diffusion integrals, achieving high-quality image generation and inference acceleration under a fixed sampling step constraint. The ACDMS update process is represented as follows: ； In the formula, For time steps The corresponding noisy latent representation, The time step difference between adjacent times. For the student model at time step Noise prediction, For order, The first generation generated for UniPC Order correction term, For Flow Matching, the correction vector field, To employ the UniPC prediction-correction mechanism at time step The intermediate prediction results obtained from the above are as follows: To correct the strength coefficient.

2. The diffusion model inference acceleration method according to claim 1, characterized in that, Step 1 includes: Step 1.1: Represent the complete set of diffusion time steps of the pre-trained diffusion generation model as follows: in, This represents the total number of discrete time steps in the diffusion process; The forward diffusion process can be represented in standard parametric form as follows: ； In the formula, This is the cumulative decay factor; Indicates the original clean sample; Indicates and A Gaussian noise vector with consistent dimensions; Indicates the first Noisy samples obtained at each diffusion time step; It is the identity matrix. Let represent a matrix with a mean of 0 and a covariance of identity. The multidimensional Gaussian distribution; Step 1.2: Set the search time step sequence length to a fixed value. The process of selecting time steps in a single diffusion inference process is abstracted into a process by... The time-step sequence encoding composed of discrete decision variables is defined as: ； In the formula, Represents a time step sequence, by each discrete time step constitute, ; The following constraints are imposed on the time-step sequence: ； Step 1.3: Based on the time step sequence search, introduce the inference model architecture configuration corresponding to each time step, and take the sub-network structure of the noise prediction network at different time steps in the diffusion inference process as the searchable object to obtain a unified search space. For each time step in the time step sequence Define the corresponding inference model architecture and configure it as follows: ； In the formula, Indicates at time step The selected noise prediction network substructure, Indicates the structure of each layer, superscript This indicates the number of network layers activated at the corresponding time step, and The total number of layers in the pre-trained noise prediction network shall not exceed the total number of layers. The candidate architecture in a single diffusion inference process is represented as: ； In the formula, This represents a candidate architecture in a diffusion inference process.

3. The diffusion model inference acceleration method according to claim 2, characterized in that, Step 2 includes: Step 2.1: Set the complete diffusion time steps used by the pre-trained diffusion generation model. As the initial set of time steps; the length of the time step sequence is not changed during the search process, only the values ​​and combinations of time steps are optimized; Step 2.2: Based on multi-objective indicators For different time step sequences The evaluation criteria include at least the FID (Focus on Generation Quality), the number of sampling steps (Sampling Efficiency), and the inference latency (Inference Speed). Step 2.3: Sequence of time steps As structural variables input into the neural network architecture search framework, the optimal time-step sequence satisfying the preset convergence condition is obtained through multi-objective optimization. .

4. The diffusion model inference acceleration method according to claim 3, characterized in that, The multi-objective indicators , is represented as: ； For time step-based subsequences The generation quality index obtained from the generated samples is used to measure the degree of difference between the generated samples and the distribution of the real samples: ； In the formula, The function represents the Fréchet Inception Distance evaluation function, which measures the statistical distance between the generated sample distribution and the real sample distribution in the feature space; Indicates the time step sequence Under constraints, the first The final generated image sample output by the diffusion generation model during the secondary diffusion inference process; This represents a reference dataset composed of real image data, used as a true distribution benchmark for generating quality assessments; To use time step subsequence The total number of sampling steps during the sampling process is used to characterize the sampling efficiency. ； For time step sequence The average inference latency for completing one full sampling process is used to characterize inference speed: ； In the formula, The function represents a time-step sequence on a fixed computing hardware platform. The actual reasoning time required to complete the diffusion reasoning generation; This represents the statistical average of the reasoning delays obtained from multiple independent reasoning processes.

5. The diffusion model inference acceleration method according to claim 3, characterized in that, Step 3 includes: Step 3.1: According to the preset time step sampling strategy, sample from the optimal time step sequence. Uniform random sampling of a single time step and the time step As the diffusion time step used in the current training iteration; Step 3.2: Under the same noise and time step conditions, noise prediction and intermediate denoising results are obtained through the teacher model and student model respectively. The main distillation loss is constructed to constrain the noise prediction of the student model to approximate the noise prediction of the teacher model, and a gradient-aware regularization term is introduced to suppress high-frequency noise. Step 3.3: From the sorted optimal time step sequence The student model is then continuously sampled at two adjacent time steps with the same index, and the corresponding noisy latent representation is constructed using the same noise conditions. The student model is then computed at these adjacent time steps. and The reconstruction results are used to construct a consistency loss. The constraints ensure that the reconstruction results remain consistent across the time domain; Step 3.4: The main distillation loss and the consistency loss are weighted and summed according to preset weights to obtain the total distillation loss of the student model. Backpropagation and parameter updates are performed on the student diffusion generation model.

6. The diffusion model inference acceleration method according to claim 5, characterized in that, The mean square error term of the main distillation loss is: ； In the formula, This represents the distillation loss value based on noise prediction consistency, used to measure the degree of difference between the prediction results of the student diffusion model and the prediction results of the teacher diffusion model at a given time step; Indicates time step and noisy random variables The mathematical expectation; This is a noisy sample obtained by the diffusion process. and The teacher model and the student model are respectively at the time step Noise prediction.

7. The diffusion model inference acceleration method according to claim 6, characterized in that, The gradient-aware regularization term is implemented by calculating the gradient difference of the noise prediction feature map. In the teacher model and the student model, the noise prediction affects the input. The gradients are respectively: ； Gradient-aware regularization Defined as: ； In the formula, This indicates that for noisy input samples The gradient calculation operator for partial derivatives reflects the sensitivity of the noise prediction result to changes in the input sample; This indicates that the teacher model is at time step The noise prediction result at the output is relative to the noisy input sample. Gradient characteristics; Indicates the student model at time step The noise prediction result at the output is relative to the noisy input sample. The gradient characteristics.

8. The diffusion model inference acceleration method according to claim 7, characterized in that, The consistency loss This is used to improve generation stability across the entire time domain, and the construction method is as follows: set up For index-adjacent time step pairs Under the same noise conditions, construct the corresponding noisy latent representation. : ； ； In the formula, and These represent the time steps. and The corresponding cumulative noise attenuation factor; The reconstruction estimation based on the student model defines consistency loss, and the reconstruction estimation result at each step is defined as follows: : ； The consistency loss can be expressed as: ； In the formula, Represents the optimal time step sequence The statistical average of the consistency loss for adjacent time step pairs in the index and under different noise sample conditions. For the student model at time step The reconstruction results.

9. The diffusion model inference acceleration method according to claim 8, characterized in that, The primary distillation loss is expressed as: ； In the formula, and The weighting coefficients used to balance the contributions of different distillation constraint terms satisfy the following conditions: ; The total distillation loss of the student model is a weighted sum of the principal distillation loss and the consistency loss: ； In the formula, The weights are used for consistency loss. Backpropagation and parameter updates are performed on the student model parameters based on the total distillation loss. During the distillation training process, the teacher model parameters are always frozen, and only the student model parameters are updated.

Images