An online policy distillation method and device for a diffusion model

Abstract

The present application relates to the technical field of knowledge distillation, and provides an online strategy distillation method and device for a diffusion model, which comprises: based on first teacher prediction features output by a teacher diffusion model, jointly optimizing prediction results of a student diffusion model to obtain an initialized student diffusion model; controlling the initialized student diffusion model to iteratively denoise based on a multi-condition combination guided sampling mode to generate a sampling track; for different multi-dimensional feature positions of the sampling track in the iteration process, obtaining second teacher prediction features output by the teacher diffusion model for the current position, and determining a target dynamic weighting weight by combining an influence degree index of the second teacher prediction features on a generation task and a current fitting ability index of the student diffusion model; based on the target dynamic weighting weight, weighting a difference between the second teacher prediction features and current output features of the student diffusion model to obtain a distillation loss, and updating network parameters of the student diffusion model.

Description

Technical Field

[0001] This invention relates to the field of knowledge distillation technology, and in particular to an online strategy distillation method and apparatus for a diffusion model. Background Technology

[0002] Diffusion models have become the mainstream image generation technique in the field of computer vision, demonstrating excellent generation quality and diversity. However, diffusion models rely on hundreds to thousands of iterative denoising and sampling steps, resulting in extremely slow inference speed and huge computational resource consumption. To deploy high-performance diffusion models on computing-constrained devices, model compression techniques are particularly crucial. Among these, knowledge distillation, as a mainstream method, is widely used to transfer the generation capabilities of large-scale teacher diffusion models to lightweight student diffusion models with fewer parameters.

[0003] In existing knowledge distillation schemes for diffusion models, the most common approach is based on static direct fitting at the model output level. Specifically, existing techniques typically involve the student model directly mimicking the teacher model's predictions under the same time steps and input noise conditions, transferring knowledge by calculating the mean squared error between the two models' feature space outputs. Since the teacher model's guidance is assumed to have equal value at any time and location, existing methods usually apply a uniform global supervision signal to all diffusion time steps and the spatial location of the entire image.

[0004] However, when there are significant differences in parameter count or data distribution between the teacher and student diffusion models, existing static uniform distillation methods face the core technical problem of distillation failure due to distribution shift. Existing methods neglect the severe inconsistency between the sampled states during the training phase and the actual inference phase. Under this static fitting, small errors during training are amplified dramatically during the student model's independent iterative denoising phase, causing the probability flow to deviate from the correct generation trajectory. Therefore, an online policy distillation mechanism is urgently needed, enabling the student model to adaptively select high-value features from the dynamically generated trajectory for learning, thereby eliminating the distribution shift between the training and inference phases and achieving efficient and high-quality knowledge transfer. Summary of the Invention

[0005] This invention provides an online policy distillation method and apparatus for diffusion models, which solves the problems of probability flow distribution deviation caused by differences between teacher and student models in the prior art, as well as the distillation failure and low training efficiency caused by global uniform supervision, and realizes efficient and accurate transfer of generation capabilities of lightweight student models.

[0006] This invention provides an online policy distillation method for diffusion models. The method is applied during the training process of knowledge distillation of a student diffusion model based on a teacher diffusion model, and includes: Based on the first teacher prediction features output by the teacher diffusion model under different guidance conditions, the prediction results of the student diffusion model under the corresponding conditions are jointly optimized to obtain the initialized student diffusion model. The initialized student diffusion model is controlled to perform iterative denoising based on a multi-condition combination guided sampling method to generate sampling trajectories for knowledge distillation; For the sampling trajectory at different multidimensional feature positions during the iteration process, the second teacher prediction feature output by the teacher diffusion model for the current position is obtained, and the target dynamic weighting weight is determined by combining the influence index of the second teacher prediction feature on the generated task and the current fitting ability index of the student diffusion model. Based on the target dynamic weighting, the difference between the second teacher's predicted features and the current output features of the student diffusion model is weighted to obtain the distillation loss, and the network parameters of the student diffusion model are updated according to the distillation loss.

[0007] According to the online policy distillation method for the diffusion model provided by the present invention, the first teacher prediction feature includes a true conditional probability stream and a true basic generation probability stream; The step of jointly optimizing the prediction results of the student diffusion model under corresponding conditions based on the first teacher prediction features output by the teacher diffusion model under different guidance conditions to obtain the initialized student diffusion model includes: Acquire training samples containing the target image and target condition information corresponding to the target image, and perform noise processing on the target image to obtain the noise-added input features at the target diffusion time step; The noisy input features and the target condition information are input into the teacher diffusion model to obtain the corresponding true conditional probability stream; and the noisy input features and the null condition information are input into the teacher diffusion model to obtain the corresponding true basic generation probability stream. Calculate the first loss between the conditional prediction score output by the student diffusion model under the same input and the true conditional probability flow, and the second loss between the unconditional prediction score output by the student diffusion model and the true basic generation probability flow. The initial network parameters of the student diffusion model are updated by combining the first loss and the second loss to obtain the initialized student diffusion model.

[0008] According to the online policy distillation method for the diffusion model provided by the present invention, the student diffusion model that is initialized is controlled to perform iterative denoising based on a multi-condition combined guided sampling method to generate a sampling trajectory for knowledge distillation, including: In the trajectory generation stage of the initialized student diffusion model, the conditional prediction score of the target condition information and the unconditional prediction score of the empty condition information are calculated respectively at the current diffusion time step. The classifier-free guidance strategy is used as the multi-condition combined guided sampling method. The conditional prediction score and the unconditional prediction score are weighted and combined according to the preset guidance scale parameter to obtain the combined prediction score. The combined prediction score is used to perform denoising state updates to obtain the hidden state of the next time step, and the sequence of hidden states generated by multiple consecutive time steps is used as the sampling trajectory to ensure that the student diffusion model maintains a consistent probability flow distribution during the training and actual inference phases.

[0009] According to the online policy distillation method for the diffusion model provided by the present invention, the step of obtaining the second teacher prediction feature output by the teacher diffusion model for the current position at different multi-dimensional feature positions of the sampling trajectory during the iteration process, and determining the target dynamic weighting weight by combining the influence index of the second teacher prediction feature on the generated task and the current fitting ability index of the student diffusion model, includes: For the current diffusion time step and latent space position represented by the multidimensional feature position, evaluate the potential benefit of the second teacher prediction feature to the final generated result, and obtain the influence index. At the current diffusion time step and the latent space position, evaluate whether the difference between the second teacher's predicted features and the current features of the student diffusion model is within a learnable range, and obtain the current fitting ability index. Based on the numerical combination of the influence degree index and the current fitting ability index, the spatiotemporal teachability function value corresponding to each latent space position at the current diffusion time step is calculated, and the spatiotemporal teachability function value is used as the target dynamic weighting weight.

[0010] According to the online policy distillation method for the diffusion model provided by the present invention, the step of evaluating whether the difference between the second teacher predicted features and the current features of the student diffusion model is within a learnable range at the current diffusion time step and the latent space position, and obtaining the current fitting ability index, includes: Calculate the error gradient between the second teacher's predicted feature and the student diffusion model's current output feature at the corresponding latent space location; Obtain the historical optimization state parameters of the student diffusion model, and combine the error gradient and the historical optimization state parameters to calculate the effective fitting probability of the student diffusion model to the current position feature difference; The absorbability value is assigned to the corresponding latent space position based on the effective fitting probability, and the absorbability value is used as the current fitting ability index, wherein the position with the higher effective fitting probability has a larger absorbability value.

[0011] According to the online strategy distillation method of the diffusion model provided by the present invention, the step of evaluating the potential benefit of the second teacher prediction feature to the final generated result and obtaining the influence index includes: calculating the contribution weight of the second teacher prediction feature to the final generated result at the current diffusion time step; Obtain the preset diffusion model signal-to-noise ratio scheduling parameters, and dynamically scale the contribution weights according to the signal-to-noise ratio value corresponding to the current diffusion time step to obtain the corresponding potential benefit value. The potential benefit value is used as the influence index, wherein the larger the potential benefit value of a latent space location, the greater its influence on the final generated result.

[0012] According to the online policy distillation method for the diffusion model provided by the present invention, the step of weighting the difference between the second teacher's predicted features and the current output features of the student diffusion model based on the target dynamic weighting to obtain the distillation loss includes: The target dynamic weighting weights are mapped dimensionally at the corresponding diffusion time step and latent space position to generate a dynamic space weighting tensor with the same size as the latent space feature map. Calculate the initial feature distance between the second teacher's predicted feature and the student diffusion model's current output feature at the corresponding position, as the difference; The initial feature distance is calculated position-by-position using the dynamic spatial weighted tensor to obtain the teachability-aware distillation loss as the distillation loss.

[0013] According to the online policy distillation method for the diffusion model provided by the present invention, the step of using the combined prediction score to perform denoising state update to obtain the hidden state of the next time step, and using the sequence of hidden states generated by multiple consecutive time steps as the sampling trajectory, includes: Based on the preset reverse denoising scheduling rules, the conditional probability distribution of the transition from the current diffusion time step to the next time step is determined; The combined prediction score is substituted into the conditional probability distribution as the mean prediction parameter for random sampling to obtain the hidden state of the next time step. The state update and random sampling process is repeated in descending order of diffusion time steps until the final generation stage is reached, and the hidden state sequence generated by the iterative process is used as the sampling trajectory.

[0014] According to the online policy distillation method for the diffusion model provided by the present invention, updating the network parameters of the student diffusion model based on the distillation loss includes: Based on the backpropagation algorithm, the gradient matrix of the distillation loss relative to the current network parameters of the student diffusion model is calculated; Using a preset optimizer algorithm, the network parameters of the student diffusion model are updated in a single iteration based on the gradient matrix; In the next training iteration, the student diffusion model with updated control parameters regenerates a new sampling trajectory for the next round of online knowledge transfer. This process is repeated until the student diffusion model meets the preset convergence condition.

[0015] According to the online policy distillation method for the diffusion model provided by the present invention, the step of obtaining training samples containing a target image and target condition information corresponding to the target image, and adding noise to the target image to obtain noisy input features at the target diffusion time step includes: obtaining samples containing the target image and corresponding text description as the training samples, and extracting the features corresponding to the text description as the target condition information; The target diffusion time step is obtained by randomly sampling from the preset time step distribution, and the noise variance corresponding to the target diffusion time step is obtained. Standard Gaussian noise is injected into the target image based on the noise variance to generate the noisy input features at the target diffusion time step.

[0016] This invention also provides an online policy distillation apparatus for a diffusion model, the apparatus being used in the training process of knowledge distillation of a student diffusion model based on a teacher diffusion model, comprising: The cold start initialization module is used to jointly optimize the prediction results of the student diffusion model under the corresponding conditions based on the first teacher prediction features output by the teacher diffusion model under different guidance conditions, so as to obtain the initialized student diffusion model. The online trajectory generation module is used to control the initialized student diffusion model to perform iterative denoising based on a multi-condition combination guided sampling method, and generate sampling trajectories for knowledge distillation; The spatiotemporal weight determination module is used to obtain the second teacher prediction feature output by the teacher diffusion model for the current position for different multi-dimensional feature positions of the sampling trajectory during the iteration process, and to determine the target dynamic weighting weight by combining the influence index of the second teacher prediction feature on the generated task and the current fitting ability index of the student diffusion model. The weighted distillation update module is used to weight the difference between the second teacher's predicted features and the current output features of the student diffusion model based on the target dynamic weighting to obtain the distillation loss, and update the network parameters of the student diffusion model according to the distillation loss.

[0017] The present invention also provides an electronic device, including a memory, a processor, and a computer program stored in the memory and executable on the processor, wherein the processor executes the computer program to implement an online strategy distillation method for any of the diffusion models described above.

[0018] The present invention also provides a non-transitory computer-readable storage medium having a computer program stored thereon, which, when executed by a processor, implements an online policy distillation method for any of the diffusion models described above.

[0019] The present invention also provides a computer program product, including a computer program that, when executed by a processor, implements an online policy distillation method for any of the diffusion models described above.

[0020] The online policy distillation method and apparatus for diffusion models provided by this invention jointly optimizes the prediction results of student models under corresponding conditions based on the first teacher prediction features output by the teacher diffusion model under different guidance conditions, effectively reducing the initial probability flow difference between teacher and student models. Then, by controlling the initialized student model to perform iterative denoising based on a multi-condition combined guided sampling method to generate sampling trajectories, the probability distribution offset between the training phase and the actual inference phase is significantly alleviated, avoiding the accumulation and amplification of errors during the iteration process. On this basis, for different multi-dimensional feature positions in the sampling trajectory iteration process, the dynamic weighting weight is determined by combining the influence index of the second teacher prediction feature on the generation task and the current fitting ability index of the student model. Based on this weight, the feature differences are weighted to calculate the distillation loss, so that the student model can prioritize learning teacher knowledge that benefits the generation results. In this way, in scenarios where there are significant distribution differences between teacher and student models, the efficiency and stability of online policy distillation are effectively improved, and the final generation quality of lightweight student models is significantly improved. Attached Figure Description

[0021] To more clearly illustrate the technical solutions in this invention or the prior art, the drawings used in the description of the embodiments or the prior art will be briefly introduced below. Obviously, the drawings described below are some embodiments of this invention. For those skilled in the art, other drawings can be obtained from these drawings without creative effort.

[0022] Figure 1 This is a schematic flowchart of the online strategy distillation method for the diffusion model provided by the present invention.

[0023] Figure 2 This is a schematic diagram of the training process for the online policy distillation method of the diffusion model provided by the present invention.

[0024] Figure 3 This is a schematic diagram of the online strategy distillation apparatus for the diffusion model provided by the present invention.

[0025] Figure 4 This is a schematic diagram of the structure of the electronic device provided by the present invention. Detailed Implementation

[0026] To make the objectives, technical solutions, and advantages of this invention clearer, the technical solutions of this invention will be clearly and completely described below with reference to the accompanying drawings. Obviously, the described embodiments are only some, not all, of the embodiments of this invention. All other embodiments obtained by those skilled in the art based on the embodiments of this invention without creative effort are within the scope of protection of this invention.

[0027] In the opening section of the detailed embodiments, in order to facilitate those skilled in the art to better understand the technical concept of the present invention, the prior art related to this application will be described in more detail first.

[0028] Diffusion models, as a type of generative model, generate images by simulating both the forward and reverse diffusion processes from the data distribution to the prior noise distribution. In the reverse process, the model needs to learn a denoising probability density gradient field to gradually recover a clear image from the noise. However, this process typically requires hundreds or even thousands of iterative sampling steps, resulting in slow inference speeds and severely limiting the deployment and application of diffusion models in computationally limited scenarios such as edge devices and mobile terminals.

[0029] To overcome the aforementioned shortcomings, existing technologies propose using knowledge distillation to compress and accelerate diffusion models. The basic idea is to use a large-scale diffusion model with high computational complexity but excellent generative performance as a teacher model to guide the training of a lightweight student model with fewer parameters and sampling steps, thereby transferring the generative capabilities of the teacher model to the student model. In this type of method, the design of the distillation loss function is the core factor determining the final performance of the student model. Early distillation schemes typically directly calculated the mean square error between the outputs of the teacher and student models in the pixel space or latent space, using this as a supervisory signal to fit the student model. While this method is simple to implement, it only performs point-by-point matching at the output level, failing to fully utilize the rich structured information contained in the teacher model.

[0030] To improve distillation results, other existing techniques focus on knowledge transfer in the feature space. For example, some schemes use pre-trained feature extractors to extract feature maps from images generated by the teacher and student models respectively, then calculate the similarity matrix between these feature maps in the pixel space, and then use KL divergence to constrain the student model to mimic the feature relationship distribution of the teacher model across sample images. This type of method improves the student model's ability to perceive the image spatial structure to some extent, but it is essentially still an offline distillation paradigm. That is, the supervision signal of the teacher model is pre-generated statically, and the student model does not dynamically interact with the generation trajectory of the teacher model during training. More importantly, this type of method assumes that all pixel positions in the feature map contribute equally to the distillation, and therefore uses a uniform weighting method to calculate the loss, without considering the different degrees of influence of different positions on the final generation result.

[0031] Furthermore, existing diffusion model distillation methods generally neglect the issue of inconsistent state distributions between the training and inference phases. Specifically, during training, the student model typically uses noisy intermediate states as input, while in the actual inference phase, the model initially receives pure random Gaussian noise. When the number of sampling steps is large, the impact of this distribution difference is not significant. However, as the distillation process progresses, the number of sampling steps for the student model is gradually compressed to extremely low levels (e.g., 4 steps or even 1 step). The discrepancy in state distribution between the training and inference phases leads to the accumulation of errors in the student model during independent iterative denoising, causing the generated trajectory to deviate from the correct probability flow direction, ultimately resulting in severe degradation of the generated image quality. Existing methods have failed to effectively address this problem from the perspective of probability flow consistency.

[0032] To address the aforementioned problems, this invention provides an online policy distillation method for diffusion models. The core concept is as follows: First, during the cold-start initialization phase, the predicted features output by the teacher model under different guidance conditions are used to jointly optimize the prediction results of the student model under corresponding conditions, thereby reducing the initial probability flow difference between the teacher and student models. Second, during the online distillation phase, the student model is controlled to iteratively denoise based on a multi-condition combined guided sampling method to generate sampling trajectories, ensuring that the state distribution during the training phase remains consistent with the actual inference phase, eliminating error accumulation caused by distribution shifts. Finally, for different multi-dimensional feature positions during the sampling trajectory iteration process, an innovative index is introduced regarding the influence of teacher predicted features on the generation task and the current fitting ability index of the student model. Based on this, a target dynamic weighting weight is determined, and the difference between the teacher predicted features and the student's current output features is weighted based on this weighting weight to calculate the distillation loss, thereby achieving differentiated and adaptive transfer of teacher knowledge at different spatiotemporal locations. The technical solution of this invention will be described in detail below with reference to the accompanying drawings and specific embodiments.

[0033] Before describing the technical solutions of the embodiments of the present invention, the terms and concepts involved in the embodiments of the present invention will be explained illustratively.

[0034] Teacher diffusion model: refers to a large-scale generative diffusion model that has been fully trained. It has a large number of parameters and a complex network structure, and can generate high-quality images or other modal data through iterative denoising. In this invention, the teacher diffusion model serves as a knowledge source to guide the training of a lightweight student model. The teacher diffusion model can be a conditional diffusion model, meaning it can generate corresponding outputs based on input text descriptions, category labels, and other conditional information.

[0035] Student diffusion model: refers to a lightweight generative diffusion model to be trained, whose network structure and parameter count are usually more compact compared to the teacher diffusion model. The student diffusion model learns its generative ability by mimicking the predictive behavior of the teacher model, thereby achieving near-teacher model generative quality while maintaining relatively low computational overhead. At the start of training, the parameters of the student diffusion model can be randomly initialized, or a subset of the parameters from the teacher diffusion model can be used for warm-up initialization; this invention does not impose any limitations on this.

[0036] First Teacher Prediction Feature: During the cold start initialization phase, the prediction results output by the teacher diffusion model after forward propagation of the same noisy input under different guidance conditions include the true conditional probability stream and the true basic generation probability stream. Specifically, when the input contains conditional information, the prediction result output by the teacher model is the conditional prediction score, i.e., the true conditional probability stream; when the input does not contain conditional information or is replaced by empty conditions, the prediction result output by the teacher model is the unconditional prediction score, i.e., the true basic generation probability stream. The first teacher prediction feature is used for joint optimization of the parameters of the student model during the initialization phase.

[0037] Second teacher prediction features refer to the prediction features output by the teacher diffusion model at a specific spatiotemporal location within the student model's sampled trajectory during the online policy distillation phase. Specifically, these features represent the teacher model's predicted velocity field or score at that location, determined by both the diffusion time step and the latent space location. The second teacher prediction features are compared with the student model's current output features at the same location to assess the feature differences and subsequently calculate the distillation loss.

[0038] Impact Index: This measures the potential contribution of the teacher model's correction terms to the quality of the final generated image at a specific spatiotemporal location. The teacher model's correction terms refer to the difference between the teacher's predicted features and the student's current output features. The impact index is obtained by quantifying the sensitivity of the teacher's correction terms to the evaluation function of the final generated result. A higher impact index value indicates a greater impact of the feature correction at that location on the final generated quality, and therefore, it is more worthy of priority learning by the student model.

[0039] Current fit capability metric: This metric measures whether the student model is capable of absorbing the correction information provided by the teacher model at the current diffusion time step and latent space position. It considers the relationship between the magnitude of the teacher correction and the current optimization state of the student model. When the teacher correction is too large and exceeds the student model's current learnable range, forced fitting may lead to training instability. The current fit capability metric is obtained by normalizing the teacher correction based on a noise scale. The closer its value is to 1, the easier it is for the student model to absorb the teacher's knowledge at that position.

[0040] Target dynamic weighting refers to the distillation weight values ​​assigned to different spatiotemporal locations within the sampling trajectory during distillation loss calculation. These weights are jointly determined by the aforementioned influence index and the current fitting ability index, comprehensively reflecting the relative value of teacher knowledge at each spatiotemporal location during the distillation process. Target dynamic weighting exhibits spatiotemporal distribution characteristics, meaning different time steps and latent space locations have different weight values, thus achieving adaptive differentiated supervision.

[0041] Conditional prediction score: This refers to the denoised direction or velocity field predicted by the diffusion model given conditional information as input. It corresponds to the estimated value of the conditional probability density gradient field. The conditional prediction score ensures semantic consistency between the generated samples and the input conditions.

[0042] Unconditional prediction score: This refers to the denoised direction or velocity field predicted by the diffusion model without any given conditional information or with no input conditions. It corresponds to the estimated value of the underlying probability density gradient field. The unconditional prediction score ensures the basic quality and authenticity of the generated samples themselves.

[0043] The hidden state sequence refers to the continuous set of hidden states experienced by the student diffusion model at each diffusion time step during the iterative denoising process based on multi-condition combined guided sampling, from the initial noise state to the final generated result. These hidden states represent the latent space features. This sequence records the complete evolution of the student model along the sampling trajectory and can be used for point-by-point comparison with the predictions of the teacher diffusion model at each corresponding state.

[0044] Current diffusion time step: refers to the current discrete-time index during the inverse denoising sampling process of the diffusion model. This time step determines the noise level of the current state and the scheduling parameters corresponding to the prediction score of the denoising network. The diffusion time step progresses in the direction of decreasing time step number, gradually decreasing from the maximum time step to zero.

[0045] Latent space location: refers to a specific position of the feature map within the diffusion model in a spatial dimension. Diffusion models typically perform denoising in a low-dimensional latent space, where the feature map has two spatial dimensions: height and width. Each spatial location corresponds to a feature vector. The spatiotemporal teachability field in this invention considers both the time dimension and the spatial location dimension to achieve fine-grained distillation supervision assignment.

[0046] Potential benefit value: This refers to the quantitative value of the potential benefit to the final generated result after introducing the teacher model's correction term at the current diffusion time step and latent space position. This value corresponds to the aforementioned impact index and is determined by assessing the sensitivity of the final generated quality evaluation function to perturbations at this position. A higher potential benefit value indicates that distillation correction at this position is more valuable for improving generated quality.

[0047] The signal-to-noise ratio (SNR) scheduling parameter for the diffusion model refers to the ratio of noise variance to signal variance determined based on the diffusion time step during the forward noise addition process of the diffusion model. This scheduling parameter determines the noise intensity in the noise input at different time steps, thus affecting the prediction difficulty and confidence of the model at that time step. In this invention, the SNR scheduling parameter is used to dynamically scale the distillation contribution weights at different time steps.

[0048] Dynamically weighted spatial tensor: This refers to a weighted tensor that represents the dynamically weighted weights of the aforementioned targets according to the spatial dimensions of the latent space feature map at the corresponding diffusion time step. The spatial size of this tensor is related to the latent space features. Figure 1 The tensor, whose height and width correspond to the height and width of the feature map, respectively, has its element value at each location as the dynamic weighted weight of the target at that location. This tensor is used for element-wise multiplication with the distillation loss to achieve spatially differentiated supervised assignment.

[0049] Initial feature distance: This refers to a measure of the difference between the second teacher prediction feature output by the teacher diffusion model and the output feature of the student diffusion model at the current spatiotemporal location, before calculating the distillation loss. This distance can be measured using mean squared error, mean absolute error, or other forms of feature space distance measurement. The initial feature distance reflects the degree of deviation between the student model's current prediction and the teacher's supervision signal at that spatiotemporal location.

[0050] Inverse denoising scheduling rules refer to the rules governing the calculation of conditional probability distributions followed by the diffusion model during the inverse generation process, when transitioning from the current diffusion time step to the next. These rules define how to update the state based on the current hidden state and the model's predicted score, including the calculation methods for the transition mean and variance. Different diffusion model architectures correspond to different inverse denoising scheduling rules.

[0051] Target diffusion time step: refers to a specific time step value randomly sampled from a preset time step distribution during the cold start initialization phase of training. This time step is used to apply noise of corresponding intensity to the original images in the training samples, generating noisy input features to simulate an intermediate state of the diffusion model during the reverse process.

[0052] The online policy distillation method for diffusion models provided in this invention can be executed by any electronic device with data computation and program processing capabilities, such as a server, cloud computing cluster, personal computer, graphics processing unit (GPU) computing device, or embedded artificial intelligence computing device. In an optional implementation, each step of the method can be implemented by a processor integrated in the aforementioned electronic device calling executable instructions stored in memory. In a typical application scenario, the aforementioned electronic device is configured as a deep learning training device, pre-loaded with a large-scale teacher diffusion model that has already been trained, and a lightweight student diffusion model to be trained. This device transfers the generation capability of the teacher model to the student model online by executing the method provided in this invention.

[0053] In one possible deployment method, the teacher model and student model can be deployed on the same electronic device or on different electronic devices. When deployed on different devices, they interact with each other via wired or wireless communication. Those skilled in the art will understand that the embodiments of the present invention do not limit the specific hardware form and deployment method of the executing entity, as long as the functions of each step of the method of the present invention can be achieved.

[0054] Figure 1 This is a flowchart illustrating the online strategy distillation method for a diffusion model provided in this embodiment of the invention. The method includes: Step 101: Based on the first teacher prediction features output by the teacher diffusion model under different guidance conditions, jointly optimize the prediction results of the student diffusion model under the corresponding conditions to obtain the initialized student diffusion model.

[0055] In implementation, considering the potential significant differences between the teacher diffusion model and the student diffusion model in terms of parameter count, network structure, or training data distribution, there is often a noticeable initial deviation between their corresponding probability density gradient fields. If subsequent online policy distillation is performed directly on these differences, the student model will face the problem of failing to effectively fit the teacher probability flow. Therefore, this embodiment of the invention first reduces this initial difference.

[0056] In this step, the electronic device acquires training samples, including the original image and corresponding conditional information. The electronic device adds noise to the original image to obtain noisy input features at a certain diffusion time step. Subsequently, the electronic device inputs the noisy input features into the teacher diffusion model under different guidance conditions. The teacher diffusion model performs forward propagation under each guidance condition, outputting the first teacher prediction feature corresponding to each guidance condition. Simultaneously, the electronic device inputs the noisy input features into the student diffusion model under the same multiple guidance conditions, obtaining the prediction results output by the student diffusion model under each guidance condition. These multiple guidance conditions include different scenarios such as with and without guidance conditions. Differences exist between the prediction features output by the teacher model under different guidance conditions, reflecting the degree of influence of guidance information on the generation process. Based on this, the electronic device optimizes the approximation errors of the teacher diffusion model under multiple guidance conditions by using the prediction results of the student diffusion model under each guidance condition to approximate the first teacher prediction feature of the teacher diffusion model under the corresponding guidance condition. This updates the initial network parameters of the student diffusion model.

[0057] Through the joint optimization process described above, the student model can simultaneously learn the output characteristics of the teacher model under different guidance strategies during the initialization phase. This effectively reduces the difference in probability density gradient fields between the teacher and student models, thereby obtaining a student diffusion model with a better initial state, which lays the foundation for subsequent online policy distillation.

[0058] Step 102: Control the initialized student diffusion model to perform iterative denoising based on a multi-condition combination guided sampling method to generate a sampling trajectory for knowledge distillation.

[0059] After the cold start initialization in step 101, the student diffusion model has acquired a certain ability to approximate the probability flow of the teacher model. However, even after the cold start, if the traditional static supervision method is still used in the subsequent distillation process—that is, inputting pre-noised intermediate states independently into the teacher and student models and calculating the loss in each training round—a significant difference still exists between the state distribution in the training phase and the state distribution in the actual inference phase. Specifically, the input to the student model in the training phase is the noisy intermediate state, while in the actual inference phase, the student model initially receives pure random Gaussian noise, which is then iteratively denoised. When the number of sampling steps is large, the impact of this difference is relatively limited, but at extremely low sampling steps, this distribution shift causes the error to accumulate and amplify continuously during the iteration process, ultimately causing the generated trajectory to deviate from the correct probability flow direction.

[0060] To overcome the aforementioned shortcomings, this embodiment of the invention employs an online strategy distillation method in step 102. Specifically, the electronic device controls the initialized student diffusion model, which uses a multi-condition combined guided sampling method for iterative denoising during its trajectory generation process. The multi-condition combined guided sampling method refers to the student model simultaneously considering the prediction results under multiple guiding conditions at each denoising update step, and fusing the prediction results under multiple guiding conditions according to a preset combination strategy to determine the final denoising update direction.

[0061] In this way, the distribution characteristics of the sampling trajectories generated by the student model during the training phase tend to be consistent with the distribution characteristics of the student model itself during the actual inference phase of iterative denoising, thereby effectively alleviating the probability flow offset problem caused by the inconsistency between training and inference distributions. The electronic device records each intermediate hidden state experienced by the student model during the iterative denoising process from the initial noisy state to the final generated result, forming a sampling trajectory for knowledge distillation. This sampling trajectory contains the sequence of hidden states of the student model at multiple diffusion time steps, providing a set of spatiotemporal locations for distillation supervision in subsequent steps.

[0062] Step 103: For the sampling trajectory at different multi-dimensional feature positions during the iteration process, obtain the second teacher prediction feature output by the teacher diffusion model for the current position, and determine the target dynamic weighting weight by combining the influence index of the second teacher prediction feature on the generated task and the current fitting ability index of the student diffusion model.

[0063] In this step, the electronic device acquires the generated sampling trajectory, which contains multiple different multidimensional feature locations.

[0064] For each multi-dimensional feature location in the sampling trajectory, the electronic device acquires the corresponding second teacher prediction feature, influence index, and current fitting ability index.

[0065] Specifically, for each multi-dimensional feature location in the sampling trajectory, the electronic device acquires the current output feature of the student diffusion model at that location and inputs the corresponding state at that location into the teacher diffusion model to obtain the second teacher prediction feature output by the teacher diffusion model for that location. Simultaneously, the electronic device acquires an indicator of the degree of influence of this second teacher prediction feature on the generation task, and an indicator of the current fitting ability of the student diffusion model for that location. The degree of influence indicator characterizes the contribution of the teacher model's correction at that multi-dimensional feature location to the final generation result, while the current fitting ability indicator characterizes whether the teacher model's correction information at that multi-dimensional feature location is within the current learnable range of the student model.

[0066] The electronic device determines the target dynamic weighting weight for each multidimensional feature location based on the numerical combination of the aforementioned influence index and the current fitting ability index. In this way, the electronic device determines the corresponding target dynamic weighting weight for each multidimensional feature location in the sampling trajectory.

[0067] Step 104: Based on the target dynamic weighting, the difference between the second teacher prediction feature and the current output feature of the student diffusion model is weighted to obtain the distillation loss, and the network parameters of the student diffusion model are updated according to the distillation loss.

[0068] In this step, for each multidimensional feature location in the sampling trajectory, the electronic device acquires the second teacher prediction feature corresponding to that location and the current output feature of the student diffusion model at that location, and calculates the feature difference between the two.

[0069] The electronic device performs a weighted calculation on the target's dynamically weighted weight determined at that location and the feature differences at that location to obtain the weighted feature differences at that location. By summing the weighted feature differences of all multi-dimensional feature locations in the sampled trajectory, the electronic device obtains the distillation loss for this iteration.

[0070] After calculating the distillation loss, the electronic device updates the network parameters of the student diffusion model based on this loss. One possible implementation is that the electronic device uses the backpropagation algorithm to calculate the gradient of the distillation loss relative to the current network parameters of the student model, and adjusts the network parameters of the student model according to the calculated gradient.

[0071] After completing this parameter update, the electronic device will use the updated Student diffusion model for the next round of online distillation iteration, repeating steps 102 to 104 until the Student model meets the preset convergence condition.

[0072] Through steps 101-104, the prediction results of the student model under corresponding conditions are jointly optimized by the first teacher prediction features output by the teacher diffusion model under different guidance conditions, which effectively reduces the initial probability flow difference between the teacher and student models. Then, by controlling the initialized student model to perform iterative denoising based on a multi-condition combined guided sampling method to generate sampling trajectories, the probability distribution offset between the training phase and the actual inference phase is significantly alleviated, and the accumulation and amplification of errors during the iteration process is avoided. On this basis, for different multi-dimensional feature positions in the sampling trajectory iteration process, the target dynamic weighting weight is determined by combining the influence index of the second teacher prediction feature on the generation task and the current fitting ability index of the student model. Based on this weight, the feature differences are weighted to calculate the distillation loss, so that the student model can prioritize learning teacher knowledge that is more beneficial to the generation results. In this way, in scenarios where there are significant distribution differences between the teacher model and the student model, the efficiency and stability of online policy distillation are effectively improved, and the final generation quality of the lightweight student model is significantly improved.

[0073] In an optional implementation, step 101 involves jointly optimizing the student diffusion model based on the first teacher prediction features output by the teacher diffusion model under different guidance conditions, specifically including: First, the electronic device acquires training samples, which include a target image and corresponding target condition information. For example, when performing a text-to-image generation task, the target image is a real image, and the target condition information is the text description corresponding to that real image. The electronic device adds noise to the target image to obtain noisy input features at the target diffusion time step. Specifically, the electronic device randomly samples from a preset time step distribution to determine a target diffusion time step, obtains the noise variance corresponding to that target diffusion time step, and injects standard Gaussian noise into the target image based on this noise variance to generate noisy input features at that target diffusion time step.

[0074] Subsequently, the electronic device combines the noisy input feature with different guiding conditions and inputs it into the teacher diffusion model to obtain the corresponding first teacher prediction feature. Specifically, the electronic device inputs the noisy input feature and target conditional information into the teacher diffusion model, and the teacher diffusion model outputs the true conditional probability stream, denoted as s_c. Simultaneously, the electronic device inputs the noisy input feature and empty conditional information into the teacher diffusion model, and the teacher diffusion model outputs the true basic generation probability stream, denoted as s_u. Empty conditional information refers to conditional inputs that do not contain any valid guiding information, such as all-zero vectors or special placeholder embeddings, used to represent the unconditional guiding case.

[0075] Simultaneously, the electronic device inputs the noisy input feature, along with the target conditional information and the empty conditional information, into the student diffusion model to obtain the conditional prediction score output by the student diffusion model under conditional guidance and the unconditional prediction score output under unconditional guidance. The electronic device calculates the first loss between the conditional prediction score of the student diffusion model and the true conditional probability flow of the teacher diffusion model, and the second loss between the unconditional prediction score of the student diffusion model and the true fundamental generation probability flow of the teacher diffusion model. The first and second losses can be calculated using mean squared error or mean absolute error equidistant metrics.

[0076] Based on this, the electronic device updates the initial network parameters of the student diffusion model with the goal of jointly optimizing the first and second losses. One optional implementation is that the electronic device performs a weighted summation of the first and second losses to obtain the joint optimization loss. Then, it calculates the gradient of this joint optimization loss relative to the student model parameters using the backpropagation algorithm, and updates the student model's network parameters based on the gradient using the optimizer algorithm. Through this method, the student model can learn both the teacher model's predictive ability for denoising direction under conditional guidance and its predictive ability for basic generation quality under unconditional guidance during the initialization phase. This effectively reduces the difference in probability density gradient fields between the teacher and student models, thus obtaining the initialized student diffusion model.

[0077] Since the classifier-free guidance method of the diffusion model in the actual inference process depends on the combination of conditional prediction score and unconditional prediction score, i.e., s=s u +w(s c -s u), where w represents the guidance strength coefficient. Therefore, the capabilities required by the student model during actual inference include both conditional prediction capability and unconditional prediction capability. If only the conditional prediction score is learned while the unconditional prediction score is ignored, the student model cannot accurately reproduce the complete probability flow of the teacher model during the inference phase. Through the joint optimization of the first and second losses in this embodiment, the student model can simultaneously approximate the conditional probability flow and the basic generation probability flow of the teacher model during the initialization phase. This reduces the initial distribution difference between the teacher and student models before the subsequent online policy distillation begins, providing a better initial state for subsequent steps and improving the overall distillation efficiency and final generation quality.

[0078] In one optional implementation, during step 101, when acquiring training samples and adding noise to the target image, the electronic device acquires training samples, which include the target image and a corresponding text description. The electronic device inputs the text description into a pre-trained text encoding model, converts the text description into a corresponding feature vector, and uses this feature vector as target conditional information. The pre-trained text encoding model can be a CLIP model text encoder, a BERT model, or other models capable of encoding natural language text into dense feature representations; this invention does not specifically limit this.

[0079] The electronic device randomly samples a target diffusion time step from a preset time step distribution. This preset time step distribution can be a uniform distribution or other pre-defined probability distribution. The electronic device obtains the noise variance corresponding to this target diffusion time step. This noise variance is predetermined based on the forward noise-adding scheduling strategy of the diffusion model. Different diffusion time steps correspond to different noise variances; generally, the larger the diffusion time step, the larger the corresponding noise variance, and the stronger the noise addition. The electronic device determines the noise intensity based on this noise variance and injects standard Gaussian noise of the corresponding intensity into the target image, generating noisy input features at the target diffusion time step. These noisy input features, together with the target condition information, constitute the input for the subsequent forward propagation of the teacher and student models.

[0080] In this way, the electronic device can randomly sample different diffusion time steps in each training iteration, covering various levels of noise from pure noise to clear images during the training process. This helps the student model learn denoising capabilities under different noise levels during the cold start initialization phase. Combined with the aforementioned implementation method of obtaining conditional prediction scores and unconditional prediction scores using target conditional information and null conditional information respectively, this implementation provides a standardized training data preprocessing procedure for the cold start initialization phase, ensuring the smooth execution of the subsequent joint optimization process.

[0081] In one optional implementation, step 102, controlling the initialization of the student diffusion model, involves iterative denoising based on a multi-condition combined guided sampling method, specifically including: The electronic device acquires the initialized student diffusion model and controls it to iteratively denoise from the initial noisy state. During the trajectory generation phase of the student diffusion model, for each current diffusion time step, the electronic device inputs the hidden state at the current time step, along with the target condition information and the null condition information, into the student diffusion model to obtain the conditional prediction score and the unconditional prediction score output by the student diffusion model.

[0082] The electronic device uses a classifier-free guidance strategy as the aforementioned multi-condition combined guided sampling method. Specifically, the electronic device acquires a preset guidance scale parameter, denoted as w, which is used to control the weight strength of conditional guidance in the denoising direction. Based on this guidance scale parameter, the electronic device performs a weighted combination calculation on the conditional prediction score and the unconditional prediction score to obtain the combined prediction score, i.e.: s=s u +w·(s c -s u ); Where s represents the combined prediction score, s c S represents the conditional prediction score. u represents the unconditional prediction score, and w represents the guiding scaling parameter.

[0083] The electronic device updates the denoised state based on the combined prediction score, calculates the update amount for transitioning from the current hidden state to the hidden state at the next time step according to the reverse denoising scheduling rule, and updates the hidden state accordingly.

[0084] The electronic device repeats the above process in descending order of diffusion time steps, sequentially recording the hidden states generated by multiple consecutive time steps to form a hidden state sequence. This hidden state sequence is then used as the sampling trajectory for knowledge distillation. In this way, the student model employs the same classifier-free guided sampling method for iterative denoising during the training phase as it does during the actual inference phase.

[0085] By using the same sampling method for iterative denoising during the training and inference phases of the student model, the distribution characteristics of the sampling trajectories generated during training are kept consistent with those during the actual inference phase, i.e., p train (z t )≈p inference (z t This effectively reduces the probability flow shift problem caused by the inconsistency between training and inference distributions, avoids the accumulation and amplification of errors in a few iterations, and improves the generation stability of the student model after distillation at extremely low sampling steps.

[0086] In an optional implementation, during step 102, when updating the denoised state based on the combined prediction score to generate the sampling trajectory, the electronic device acquires a preset inverse denoising scheduling rule. This inverse denoising scheduling rule defines the conditional probability distribution that must be followed when transitioning from the hidden state of the current diffusion time step to the hidden state of the next diffusion time step. Different diffusion model architectures correspond to different inverse denoising scheduling rules; for example, DDPM and DDIM have different transition rules. This invention does not specifically limit this, and those skilled in the art can select the appropriate scheduling rule according to the actual model architecture used.

[0087] For the current diffusion time step, the electronic device substitutes the combined prediction score as the mean prediction parameter in the aforementioned conditional probability distribution, and then randomly samples from this conditional probability distribution to obtain the hidden state for the next diffusion time step. This hidden state is the denoising result of the student diffusion model at this time step, containing estimation information of the original data under the current noise level.

[0088] The electronic device repeatedly executes the aforementioned state update and random sampling process in descending order of diffusion time steps. Each step uses the current hidden state as input, calculates the combined prediction score, substitutes it into the conditional probability distribution for random sampling, and obtains the next hidden state, until the preset final generation stage is reached. During this process, the electronic device sequentially records the hidden states generated at each diffusion time step, forming a hidden state sequence, which is then used as a sampling trajectory for knowledge distillation. This sampling trajectory fully records the entire process of the student diffusion model from the initial noisy state to the final generated result, providing a data foundation for the student diffusion model to receive teacher supervision at various spatiotemporal locations in subsequent steps.

[0089] The sampling trajectories generated in the above manner use the same sampling method as the inference stage during the rollout process, which ensures that the distribution characteristics of the sampling trajectories are consistent with the distribution of the generated trajectories in the actual inference stage. This lays a reasonable data foundation for subsequent steps to perform spatiotemporal teachability assessment and distillation loss calculation based on the sampling trajectories.

[0090] In an optional implementation, step 103 determines the target dynamic weighting weight for different multidimensional feature positions of the sampling trajectory during the iteration process, which can be achieved in the following way.

[0091] The electronic device acquires the sampling trajectory generated in step 102. Each multidimensional feature position in the sampling trajectory corresponds to a diffusion time step and a latent space position. For each multidimensional feature position in the sampling trajectory, the electronic device acquires the second teacher prediction feature output by the teacher diffusion model for that position.

[0092] For each multidimensional feature location, the electronic device first evaluates the potential benefit of introducing a teacher model correction at that spatiotemporal location to the final generated result, obtaining an impact index. This impact index reflects the importance of teacher knowledge at that location to the quality of the final generated image. Specifically, the electronic device obtains the difference between the velocity field prediction value of the student diffusion model at that location and the second teacher prediction feature of the teacher diffusion model at that location, uses this difference as a correction term for the teacher model, and determines the magnitude of the potential benefit by evaluating the degree of influence of this correction term on the final generated quality evaluation function.

[0093] Simultaneously, the electronic device also needs to determine whether the teacher's correction information at that location is within the current learnable range of the student model, thus obtaining a current fitting ability index. Specifically, the electronic device acquires the difference between the second teacher prediction feature of the teacher diffusion model at that location and the current output feature of the student diffusion model, and normalizes and evaluates this difference by combining it with the noise scale corresponding to the current diffusion time step. When the difference is relatively small, it indicates that the teacher's knowledge at that location is within the current student model's ability range and is easily absorbed; when the difference is too large, it indicates that forced fitting may exceed the current ability range of the student model, leading to unstable training.

[0094] Next, the electronic device calculates the spatiotemporal teachability function value for each spatiotemporal location based on the numerical combination of the aforementioned influence index and the current fitting ability index. This spatiotemporal teachability function value comprehensively considers the distillation value and learnability of that location. For example, the formula for calculating this spatiotemporal teachability function value can be expressed as: Q(t,h,w)=α·I(t,h,w)+β·A(t,h,w); Where t represents the diffusion time step, (h,w) represents the latent space location, I(t,h,w) represents the influence index, A(t,h,w) represents the current fitting ability index, and α and β are preset weighting coefficients. The influence index I(t,h,w) is obtained by quantifying the sensitivity of the teacher correction term to the final generated result evaluation function, and the current fitting ability index A(t,h,w) is obtained by normalizing the teacher correction term based on the noise scale. The values ​​of both are mapped to the same numerical interval for weighted combination.

[0095] The electronic device uses the calculated spatiotemporal teachability function value as the target dynamic weighting weight for each multidimensional feature location. In this way, the electronic device determines the corresponding target dynamic weighting weight for each multidimensional feature location in the sampling trajectory. The larger the Q(t,h,w) value, the higher the transfer value of the teacher's knowledge at that location during the distillation process.

[0096] Through this implementation, the electronic device can accurately identify which positions in the sampling trajectory have a significant impact on the final generation quality. It can also determine whether the student model is capable of absorbing the teacher correction information at that position during the current training phase. This achieves effective screening of distillation targets in terms of both influence and absorbability, ensuring that each parameter update of the student model focuses on high-value and absorbable teacher knowledge. This avoids training instability caused by forcibly fitting teacher corrections that exceed the model's capabilities, and significantly improves distillation efficiency within a limited computational budget.

[0097] To obtain the impact index, the electronic device can use the following method: For each multi-dimensional feature position in the sampling trajectory, the electronic device calculates the contribution weight of the second teacher's predicted feature to the final generated result at the current diffusion time step. This contribution weight is used to quantify the impact of the teacher's correction at that position on the quality of the final generated image. The electronic device obtains the preset diffusion model signal-to-noise ratio (SNR) scheduling parameters. The diffusion model has different SNR characteristics at different time steps. Generally, the noise proportion is higher in earlier time steps, and the confidence of the model prediction is relatively lower, while the noise proportion is lower in later time steps, and the confidence of the model prediction is relatively higher. The electronic device dynamically scales the above contribution weights according to the SNR value corresponding to the current diffusion time step. For time steps with low SNR, the impact of the correction at that position on the final generated result may be more global, so its weight can be appropriately reduced; for time steps with high SNR, the correction at that position often directly affects the semantic structure or visual details of the generated image, so its weight can be appropriately increased.

[0098] Through the aforementioned dynamic scaling, the electronic device obtains the potential gain value corresponding to that location and uses this potential gain value as an indicator of the degree of influence. Specifically, a larger potential gain value indicates a more significant improvement in the quality of the final generated image after the teacher model's correction is introduced at that location, meaning a higher degree of influence on the generation task. In this way, the electronic device can accurately identify key spatiotemporal locations that significantly impact the final generation result during the distillation process, allowing subsequent distillation supervision to focus on these high-influence regions and avoid dispersing limited computational resources to locations with minimal impact on the final generation result.

[0099] For the current fitting capability index, electronic devices can obtain it in the following ways: For each multidimensional feature location in the sampling trajectory, the electronic device calculates the error gradient between the second teacher-predicted feature output by the teacher diffusion model and the current output feature of the student diffusion model at that location in the corresponding latent space. This error gradient reflects the sensitivity of the current fitting error at that location to the model parameters. The electronic device acquires the historical optimization state parameters of the student diffusion model, which include, but are not limited to, the parameter update amount or gradient norm statistics of the student model in previous iterations. The electronic device comprehensively analyzes the above error gradient and historical optimization state parameters to calculate the effective fitting probability of the student model for the feature difference at the current location. This effective fitting probability characterizes whether the student model is capable of effectively absorbing the correction information of the teacher model at that location and transforming it into its own parameter updates during the current training phase. When the error gradient is within a reasonable range relative to the historical optimization state parameters, it indicates that the teacher correction at that location is within the learning difficulty range of the student model, and the effective fitting probability is high; when the error gradient is much larger than the historical gradient norm, it indicates that the correction exceeds the current learning capacity of the student model, and forced fitting may lead to training instability. The electronic device assigns an absorptivity value to the corresponding latent space location based on the aforementioned effective fitting probability, and uses this absorptivity value as an indicator of the current fitting capability. Locations with higher effective fitting probabilities correspond to larger absorptivity values. The absorptivity value can be calculated as follows: ; Among them, v T v represents the velocity field prediction of the teacher diffusion model at that location. S This represents the velocity field prediction of the student diffusion model at that location, ∥v T -v S ∥2 2 σ represents the square norm of the difference between the two. t τ represents the noise scale corresponding to the current diffusion time step t. A This represents the preset absorption difficulty parameter. The formula means that when the difference between the teacher and student is small relative to the noise scale and the absorption difficulty parameter, the exponential term approaches 1, indicating that the teacher's correction at that position is easily absorbed by the student; when the difference is too large, the exponential term approaches 0, indicating that the correction exceeds the current ability range of the student model. The electronic device uses the calculated absorbability value as the current fitting ability index.

[0100] In this way, the electronic device can dynamically evaluate the student model's ability to absorb teacher corrections at each position during the distillation process, avoiding training oscillations or divergence caused by forcibly fitting teacher corrections that exceed the model's capabilities, thus improving the stability and robustness of the distillation process.

[0101] In an optional implementation, during the calculation of distillation loss based on the target dynamic weighting in step 104, the electronic device acquires the target dynamic weighting at each multidimensional feature position of the sampling trajectory determined in step 103. This target dynamic weighting has a distribution characteristic in the spatiotemporal dimension corresponding to the latent state feature map in the sampling trajectory. The electronic device maps this target dynamic weighting according to the corresponding diffusion time step and latent space position, generating a dynamic spatial weighting tensor with the same size as the latent space feature map. The spatial size of this dynamic spatial weighting tensor is the same as the height and width of the latent space feature map, and the element value at each position is the target dynamic weighting weight corresponding to that position.

[0102] For each multidimensional feature location in the sampling trajectory, the electronic device acquires the second teacher's predicted feature and the current output feature of the student diffusion model at that location, and calculates the feature distance between them as the initial feature distance. This initial feature distance reflects the degree of deviation between the student model's current prediction at that spatiotemporal location and the teacher's supervision signal. The feature distance can be measured using Euclidean distance, cosine distance, or other similarity measures.

[0103] The electronic device uses the aforementioned dynamic spatial weighted tensor to perform position-by-position weighted calculation of the initial feature distance. Specifically, for each spatiotemporal location, the electronic device multiplies the initial feature distance of that location by the weight value of the corresponding location in the dynamic spatial weighted tensor to obtain the weighted feature distance of that location. By summing up the weighted feature distances of all spatiotemporal locations in the sampling trajectory, the teachable perceptual distillation loss is obtained.

[0104] Compared to the traditional uniform weighted loss, this implementation incorporates the spatiotemporal teachability field into the distillation loss in the form of a spatially weighted tensor. This allows the parameter updates of the student diffusion model to automatically focus on those spatiotemporal locations with high influence and strong absorbability in the sampling trajectory, rather than fitting all locations with equal strength. This achieves more efficient distillation within a limited computational budget, while avoiding the training instability problem caused by the student model forcibly fitting teacher corrections that exceed its own capabilities.

[0105] In one alternative implementation, during step 104, when updating the network parameters of the student diffusion model based on the distillation loss, the electronic device calculates the gradient matrix of the distillation loss relative to the current network parameters of the student diffusion model using the backpropagation algorithm. This gradient matrix reflects the sensitivity of the distillation loss to each learnable parameter in the model.

[0106] The electronic device uses a pre-defined optimizer algorithm to perform a single-iteration update of the network parameters of the student diffusion model based on the aforementioned gradient matrix. The pre-defined optimizer algorithm can be stochastic gradient descent, the Adam algorithm, the AdamW algorithm, or other optimizer algorithms suitable for training deep learning models. Through this update, the student model moves in the network parameter space in a direction that reduces distillation loss; that is, its predicted output converges towards the output of the teacher model at those high distillation value locations in the sampling trajectory.

[0107] After each parameter update, the electronic device does not stop training but uses the updated student diffusion model for the next training iteration. Specifically, in the next training iteration, the electronic device re-controls the student diffusion model to generate new sampling trajectories based on a multi-condition combination guided sampling method. It then repeats the steps of acquiring the second teacher's predicted features, determining the target dynamic weighting, calculating the distillation loss, and updating the parameters based on the newly generated sampling trajectories. Through this iterative process, the student diffusion model continuously absorbs teacher knowledge in each training round, and its generated sampling trajectories gradually approximate the generation probability flow of the teacher diffusion model.

[0108] The electronic device repeats the above process until the student diffusion model meets a preset convergence condition. This preset convergence condition may be that the distillation loss decreases below a preset threshold, the parameter update amount of the student diffusion model is less than a preset threshold, or the number of training epochs reaches a preset maximum number of iterations. When the convergence condition is met, training terminates, and the student diffusion model then possesses a generative capability close to that of the teacher diffusion model, realizing the capability transfer from a large-scale teacher diffusion model to a lightweight student diffusion model.

[0109] Unlike static distillation where each training round involves independent sampling, this implementation regenerates sampling trajectories using the updated student diffusion model for each training round. As training progresses, the student diffusion model's generation capability gradually improves, and the distribution of its generated sampling trajectories dynamically changes. This ensures that the spatiotemporal teachability assessment and distillation loss calculation in subsequent rounds are always based on the current actual generation state of the student diffusion model. This achieves a gradual transfer of teacher knowledge from coarse to fine and from easy to difficult, which is beneficial for improving the final distillation effect.

[0110] Figure 2 This is a schematic diagram illustrating the training process of the online policy distillation method for the diffusion model provided in this embodiment of the invention. Figure 2 As shown, the process mainly includes three stages: the dual-scoring field cold start stage, the CFG consistency online strategy distillation stage, and the teacher state selection stage based on the spatiotemporal teachability field.

[0111] During the dual-score field cold start phase, the electronic device acquires a pre-trained large-scale teacher diffusion model T. Teacher model T is a conditional diffusion model, capable of generating corresponding images based on input conditional information c (e.g., text prompts). The electronic device uses teacher model T to generate teacher sampling trajectories, obtaining the hidden state sequence at each diffusion time step. During the generation process, the electronic device obtains the conditional scores s output by teacher model T under the guidance of conditional information c. c And the unconditional score s output under empty conditions u Based on this, electronic devices perform cold-start supervised training on the student model, enabling the student model to simultaneously learn the conditional scores s of the teacher model. c And unconditional score s u Because in actual reasoning, the combination of classifier-free guidance methods is as follows: ; Here, w is the guiding scale parameter. Therefore, only when the student model simultaneously masters both conditional and unconditional scores can it accurately reproduce the complete probability flow of the teacher model during the inference phase. Through the joint learning of the dual score fields described above, the difference in score distribution between the student model and the teacher model is effectively reduced.

[0112] During the CFG consensus online policy distillation phase, the electronic device acquires the student model after cold-start initialization and controls it to perform iterative denoising using a classifier-free guidance strategy to generate student sampling trajectories. Since the classifier-free guidance method used by the student model in trajectory generation is consistent with the combination method represented by the above formula, its probability flow distribution during training remains consistent with the probability flow distribution during actual inference, effectively reducing the probability flow offset problem caused by inconsistencies between training and inference distributions. During the student model rollout process, the electronic device acquires the correction information corresponding to each state in the student sampling state sequence from the teacher model and transfers the teacher model's generation capabilities to the student model through online policy distillation.

[0113] In the teacher state selection phase based on the spatiotemporal teachability field, the electronic device evaluates the influence index and the absorbability index for each multidimensional feature position in the student's sampling trajectory. The influence index I(t,h,w) characterizes the potential benefit to the final generated result after introducing teacher model corrections at that spatiotemporal position. The absorbability index A(t,h,w) characterizes whether the teacher correction information at that position is within the current learnable range of the student model, and its calculation method is as follows: ; Among them, v T v represents the velocity field prediction value of the teacher model. S σ represents the velocity field prediction of the student model. tτ represents the noise scale corresponding to the current diffusion time step. A This indicates the preset absorption difficulty parameter.

[0114] Electronic devices calculate the spatiotemporal teachability function value based on a combination of influence and absorbability indicators: Here, α and β are preset weighting coefficients. The electronic device uses the calculated spatiotemporal teachability function value as the target dynamic weight for that location. Based on the above target dynamic weight, the electronic device calculates the distillation loss and updates the network parameters of the student model, enabling the student model to prioritize learning the teacher knowledge with high influence and absorbability in the sampled trajectory.

[0115] Through the coordinated efforts of the above three stages, the electronic device gradually transfers the generation capability of the teacher model to the student model, ultimately obtaining a lightweight student diffusion model with fewer parameters and fewer sampling steps.

[0116] The online strategy distillation apparatus for the diffusion model provided in the embodiments of the present invention will be described below. The online strategy distillation apparatus for the diffusion model described below can be referred to in correspondence with the online strategy distillation method for the diffusion model described above.

[0117] This invention provides an online strategy distillation apparatus based on a diffusion model, see [link to relevant documentation]. Figure 3 The device is used in the training process of knowledge distillation of the student diffusion model based on the teacher diffusion model, and includes: The cold start initialization module 310 is used to jointly optimize the prediction results of the student diffusion model under corresponding conditions based on the first teacher prediction features output by the teacher diffusion model under different guidance conditions, so as to obtain the initialized student diffusion model. The online trajectory generation module 320 is used to control the initialized student diffusion model to perform iterative denoising based on a multi-condition combination guided sampling method to generate a sampling trajectory for knowledge distillation; The spatiotemporal weight determination module 330 is used to obtain the second teacher prediction feature output by the teacher diffusion model for the current position for different multi-dimensional feature positions of the sampling trajectory during the iteration process, and to determine the target dynamic weighting weight by combining the influence index of the second teacher prediction feature on the generated task and the current fitting ability index of the student diffusion model. The weighted distillation update module 340 is used to weight the difference between the second teacher's predicted features and the current output features of the student diffusion model based on the target dynamic weighting weight to obtain the distillation loss, and update the network parameters of the student diffusion model according to the distillation loss.

[0118] Furthermore, the first teacher prediction feature includes a true conditional probability stream and a true basic generation probability stream; the cold start initialization module 310 is used for: Acquire training samples containing the target image and target condition information corresponding to the target image, and perform noise processing on the target image to obtain the noise-added input features at the target diffusion time step; The noisy input features and the target condition information are input into the teacher diffusion model to obtain the corresponding true conditional probability stream; and the noisy input features and the null condition information are input into the teacher diffusion model to obtain the corresponding true basic generation probability stream. Calculate the first loss between the conditional prediction score output by the student diffusion model under the same input and the true conditional probability flow, and the second loss between the unconditional prediction score output by the student diffusion model and the true basic generation probability flow. The initial network parameters of the student diffusion model are updated by combining the first loss and the second loss to obtain the initialized student diffusion model.

[0119] Furthermore, the online trajectory generation module 320 is used for: In the trajectory generation stage of the initialized student diffusion model, the conditional prediction score of the target condition information and the unconditional prediction score of the empty condition information are calculated respectively at the current diffusion time step. The classifier-free guidance strategy is used as the multi-condition combined guided sampling method. The conditional prediction score and the unconditional prediction score are weighted and combined according to the preset guidance scale parameter to obtain the combined prediction score. The combined prediction score is used to perform denoising state updates to obtain the hidden state of the next time step, and the sequence of hidden states generated by multiple consecutive time steps is used as the sampling trajectory to ensure that the student diffusion model maintains a consistent probability flow distribution during the training and actual inference phases.

[0120] Furthermore, the spatiotemporal weight determination module 330 is used for: For the current diffusion time step and latent space position represented by the multidimensional feature position, evaluate the potential benefit of the second teacher prediction feature to the final generated result, and obtain the influence index. At the current diffusion time step and the latent space position, evaluate whether the difference between the second teacher's predicted features and the current features of the student diffusion model is within a learnable range, and obtain the current fitting ability index. Based on the numerical combination of the influence degree index and the current fitting ability index, the spatiotemporal teachability function value corresponding to each latent space position at the current diffusion time step is calculated, and the spatiotemporal teachability function value is used as the target dynamic weighting weight.

[0121] Furthermore, the weighted distillation update module 340 is used for: The target dynamic weighting weights are mapped dimensionally at the corresponding diffusion time step and latent space position to generate a dynamic space weighting tensor with the same size as the latent space feature map. Calculate the initial feature distance between the second teacher's predicted feature and the student diffusion model's current output feature at the corresponding position, as the difference; The initial feature distance is calculated position-by-position using the dynamic spatial weighted tensor to obtain the teachability-aware distillation loss as the distillation loss.

[0122] Figure 4 An example is a schematic diagram of the physical structure of an electronic device, such as... Figure 4 As shown, the electronic device may include a processor 810, a communications interface 820, a memory 830, and a communication bus 840, wherein the processor 810, the communications interface 820, and the memory 830 communicate with each other via the communication bus 840. The processor 810 can call logical instructions in the memory 830 to execute an online policy distillation method for a diffusion model.

[0123] Furthermore, the logical instructions in the aforementioned memory 830 can be implemented as software functional units and, when sold or used as independent products, can be stored in a computer-readable storage medium. Based on this understanding, the technical solution of the present invention, or the part that contributes to the prior art, or a part of the technical solution, can be embodied in the form of a software product. This computer software product is stored in a storage medium and includes several instructions to cause a computer device (which may be a personal computer, server, or network device, etc.) to execute all or part of the steps of the methods described in the various embodiments of the present invention. The aforementioned storage medium includes various media capable of storing program code, such as USB flash drives, portable hard drives, read-only memory (ROM), random access memory (RAM), magnetic disks, or optical disks.

[0124] On the other hand, the present invention also provides a computer program product, which includes a computer program that can be stored on a non-transitory computer-readable storage medium. When the computer program is executed by a processor, the computer is able to execute the online strategy distillation method of the diffusion model provided by the above methods.

[0125] In another aspect, the present invention also provides a non-transitory computer-readable storage medium having a computer program stored thereon, which, when executed by a processor, implements an online strategy distillation method for the diffusion model provided by the methods described above.

[0126] The device embodiments described above are merely illustrative. The units described as separate components may or may not be physically separate. The components shown as units may or may not be physical units; that is, they may be located in one place or distributed across multiple network units. Some or all of the modules can be selected to achieve the purpose of this embodiment according to actual needs. Those skilled in the art can understand and implement this without any creative effort.

[0127] Through the above description of the embodiments, those skilled in the art can clearly understand that each embodiment can be implemented by means of software plus necessary general-purpose hardware platforms, and of course, it can also be implemented by hardware. Based on this understanding, the above technical solutions, in essence or the part that contributes to the prior art, can be embodied in the form of a software product. This computer software product can be stored in a computer-readable storage medium, such as ROM / RAM, magnetic disk, optical disk, etc., and includes several instructions to cause a computer device (which may be a personal computer, server, or network device, etc.) to execute the methods described in the various embodiments or some parts of the embodiments.

[0128] Finally, it should be noted that the above embodiments are only used to illustrate the technical solutions of the present invention, and not to limit them; although the present invention has been described in detail with reference to the foregoing embodiments, those skilled in the art should understand that modifications can still be made to the technical solutions described in the foregoing embodiments, or equivalent substitutions can be made to some of the technical features; and these modifications or substitutions do not cause the essence of the corresponding technical solutions to deviate from the spirit and scope of the technical solutions of the embodiments of the present invention.

Claims

1. An online strategy distillation method for a diffusion model, characterized in that, The method is applied to the training process of knowledge distillation of the student diffusion model based on the teacher diffusion model, including: Based on the first teacher prediction features output by the teacher diffusion model under different guidance conditions, the prediction results of the student diffusion model under the corresponding conditions are jointly optimized to obtain the initialized student diffusion model. The initialized student diffusion model is controlled to perform iterative denoising based on a multi-condition combination guided sampling method to generate sampling trajectories for knowledge distillation; For the sampling trajectory at different multidimensional feature positions during the iteration process, the second teacher prediction feature output by the teacher diffusion model for the current position is obtained, and the target dynamic weighting weight is determined by combining the influence index of the second teacher prediction feature on the generated task and the current fitting ability index of the student diffusion model. Based on the target dynamic weighting, the difference between the second teacher's predicted features and the current output features of the student diffusion model is weighted to obtain the distillation loss, and the network parameters of the student diffusion model are updated according to the distillation loss.

2. The online strategy distillation method for the diffusion model according to claim 1, characterized in that, The first teacher prediction feature includes the true conditional probability stream and the true basic generation probability stream; The step of jointly optimizing the prediction results of the student diffusion model under corresponding conditions based on the first teacher prediction features output by the teacher diffusion model under different guidance conditions to obtain the initialized student diffusion model includes: Acquire training samples containing the target image and target condition information corresponding to the target image, and perform noise processing on the target image to obtain the noise-added input features at the target diffusion time step; The noisy input features and the target condition information are input into the teacher diffusion model to obtain the corresponding true conditional probability stream; and the noisy input features and the null condition information are input into the teacher diffusion model to obtain the corresponding true basic generation probability stream. Calculate the first loss between the conditional prediction score output by the student diffusion model under the same input and the true conditional probability flow, and the second loss between the unconditional prediction score output by the student diffusion model and the true basic generation probability flow. The initial network parameters of the student diffusion model are updated by combining the first loss and the second loss to obtain the initialized student diffusion model.

3. The online strategy distillation method for the diffusion model according to claim 1, characterized in that, The student diffusion model that controls the initialization performs iterative denoising based on a multi-condition combined guided sampling method to generate sampling trajectories for knowledge distillation, including: In the trajectory generation stage of the initialized student diffusion model, the conditional prediction score of the target condition information and the unconditional prediction score of the empty condition information are calculated respectively at the current diffusion time step. The classifier-free guidance strategy is used as the multi-condition combined guided sampling method. The conditional prediction score and the unconditional prediction score are weighted and combined according to the preset guidance scale parameter to obtain the combined prediction score. The combined prediction score is used to perform denoising state updates to obtain the hidden state of the next time step, and the sequence of hidden states generated by multiple consecutive time steps is used as the sampling trajectory to ensure that the student diffusion model maintains a consistent probability flow distribution during the training and actual inference phases.

4. The online strategy distillation method for the diffusion model according to claim 1, characterized in that, For different multidimensional feature positions of the sampling trajectory during the iteration process, the second teacher prediction feature output by the teacher diffusion model for the current position is obtained. Combined with the influence index of the second teacher prediction feature on the generated task and the current fitting ability index of the student diffusion model, the target dynamic weighting weight is determined, including: For the current diffusion time step and latent space position represented by the multidimensional feature position, evaluate the potential benefit of the second teacher prediction feature to the final generated result, and obtain the influence index. At the current diffusion time step and the latent space position, evaluate whether the difference between the second teacher's predicted features and the current features of the student diffusion model is within a learnable range, and obtain the current fitting ability index. Based on the numerical combination of the influence degree index and the current fitting ability index, the spatiotemporal teachability function value corresponding to each latent space position at the current diffusion time step is calculated, and the spatiotemporal teachability function value is used as the target dynamic weighting weight.

5. The online strategy distillation method for the diffusion model according to claim 4, characterized in that, The step of evaluating whether the difference between the second teacher's predicted features and the current features of the student diffusion model is within a learnable range at the current diffusion time step and the latent space position, and obtaining the current fitting ability index, includes: Calculate the error gradient between the second teacher's predicted feature and the student diffusion model's current output feature at the corresponding latent space location; Obtain the historical optimization state parameters of the student diffusion model, and combine the error gradient and the historical optimization state parameters to calculate the effective fitting probability of the student diffusion model to the current position feature difference; The absorbability value is assigned to the corresponding latent space position based on the effective fitting probability, and the absorbability value is used as the current fitting ability index, wherein the position with the higher effective fitting probability has a larger absorbability value.

6. The online strategy distillation method for the diffusion model according to claim 4, characterized in that, The assessment of the potential benefit of the second teacher predictive feature to the final generated result, and the obtaining of the influence index, includes: Calculate the contribution weight of the second teacher prediction feature to the final generated result at the current diffusion time step; Obtain the preset diffusion model signal-to-noise ratio scheduling parameters, and dynamically scale the contribution weights according to the signal-to-noise ratio value corresponding to the current diffusion time step to obtain the corresponding potential benefit value. The potential benefit value is used as the influence index, wherein the larger the potential benefit value of a latent space location, the greater its influence on the final generated result.

7. The online strategy distillation method for the diffusion model according to claim 1, characterized in that, The step of weighting the difference between the second teacher's predicted features and the current output features of the student diffusion model based on the target dynamic weighting to obtain the distillation loss includes: The target dynamic weighting weights are mapped dimensionally at the corresponding diffusion time step and latent space position to generate a dynamic space weighting tensor with the same size as the latent space feature map. Calculate the initial feature distance between the second teacher's predicted feature and the student diffusion model's current output feature at the corresponding position, as the difference; The initial feature distance is calculated position-by-position using the dynamic spatial weighted tensor to obtain the teachability-aware distillation loss as the distillation loss.

8. The online strategy distillation method for the diffusion model according to claim 3, characterized in that, The step of using the combined prediction score to perform denoising state update to obtain the hidden state of the next time step, and using the sequence of hidden states generated by multiple consecutive time steps as the sampling trajectory, includes: Based on the preset reverse denoising scheduling rules, the conditional probability distribution of the transition from the current diffusion time step to the next time step is determined; The combined prediction score is substituted into the conditional probability distribution as the mean prediction parameter for random sampling to obtain the hidden state of the next time step. The state update and random sampling process is repeated in descending order of diffusion time steps until the final generation stage is reached, and the hidden state sequence generated by the iterative process is used as the sampling trajectory.

9. The online strategy distillation method for the diffusion model according to claim 1, characterized in that, The step of updating the network parameters of the student diffusion model based on the distillation loss includes: Based on the backpropagation algorithm, the gradient matrix of the distillation loss relative to the current network parameters of the student diffusion model is calculated; Using a preset optimizer algorithm, the network parameters of the student diffusion model are updated in a single iteration based on the gradient matrix; In the next training iteration, the student diffusion model with updated control parameters regenerates a new sampling trajectory for the next round of online knowledge transfer. This process is repeated until the student diffusion model meets the preset convergence condition.

10. An online strategy distillation apparatus based on a diffusion model, characterized in that, The device is used in the training process of knowledge distillation of a student diffusion model based on a teacher diffusion model, including: The cold start initialization module is used to jointly optimize the prediction results of the student diffusion model under the corresponding conditions based on the first teacher prediction features output by the teacher diffusion model under different guidance conditions, so as to obtain the initialized student diffusion model. The online trajectory generation module is used to control the initialized student diffusion model to perform iterative denoising based on a multi-condition combination guided sampling method, and generate sampling trajectories for knowledge distillation; The spatiotemporal weight determination module is used to obtain the second teacher prediction feature output by the teacher diffusion model for the current position for different multi-dimensional feature positions of the sampling trajectory during the iteration process, and to determine the target dynamic weighting weight by combining the influence index of the second teacher prediction feature on the generated task and the current fitting ability index of the student diffusion model. The weighted distillation update module is used to weight the difference between the second teacher's predicted features and the current output features of the student diffusion model based on the target dynamic weighting to obtain the distillation loss, and update the network parameters of the student diffusion model according to the distillation loss.

Images