A model training method, an image task processing method, and related apparatus.

Abstract

This invention relates to the field of image task processing technology, and discloses a model training method, an image task processing method, and related apparatus. The invention designs a frequency-aware image patch selection module that adaptively selects image patches with the most valuable frequency information for the current diffusion time step and the current training stage in each training step. This concentrates limited computational resources on the data most helpful to the model's current learning, significantly reducing the training time and computational load of the diffusion model without reducing or even improving the generation quality. This invention upgrades the static computational allocation method of "treating all pixels / image patches equally" in traditional diffusion training to a dynamic adaptive computational allocation method driven by both image frequency content and the current training state. This ensures that computational resources are always tilted towards the frequency content most valuable to the current learning, eliminating a large amount of redundant computation, thereby significantly improving training efficiency while maintaining learning effectiveness.

Description

Technical Field

[0001] This invention relates to the field of image task processing technology, and in particular to a model training method, an image task processing method, and related apparatus. Background Technology

[0002] In image processing tasks such as image synthesis, video generation, text-to-image generation, and semantic segmentation, diffusion models are the most influential deep generative model paradigm in recent years. They have achieved performance that surpasses Generative Adversarial Networks (GANs) in image processing tasks and have been widely used in commercial systems such as DALL-E 2, StableDiffusion, and Imagen.

[0003] However, in image processing tasks, training diffusion models requires significant computational resources and time. For example, training a basic diffusion model (DDPM, Denoising Diffusion Probabilistic Model) using eight NVIDIA V100 GPUs on the FFHQ dataset (64×64 resolution) takes approximately four days; training PatchDiffusion on the LSUN dataset still requires about 64 V100 GPU days. This high training cost severely limits the widespread application of diffusion models in image processing. Therefore, accelerating diffusion model training has become an urgent need in image processing tasks. Summary of the Invention

[0004] This invention provides a model training method, an image task processing method, and related apparatus for accelerating diffusion model training in image task processing scenarios.

[0005] The first aspect of this invention discloses a model training method based on frequency-aware image patch selection, the method comprising: The training model is trained using a preset input image to obtain a trained target diffusion model, which is used to perform the corresponding image processing task. Each parameter update step in the training process of the diffusion model to be trained includes: The input image is divided into multiple image blocks, and each image block is transformed from the spatial domain to the frequency domain to obtain the frequency component corresponding to each image block. Obtain the diffusion time step corresponding to the current parameter update step, and obtain the embedding vector based on the diffusion time step; For each image block: the frequency component corresponding to the image block, the position information of the image block, and the embedding vector are fused to obtain an image block token; Input all the image block tokens into a pre-set frequency-aware gating module. The frequency-aware gating module predicts the diffusion loss corresponding to the current parameter update step for each image block and determines the probability value corresponding to each image block based on the diffusion loss. Based on the probability values ​​corresponding to all the image blocks, select at least one matching image block corresponding to the current parameter update step; Perform forward and backward propagation of the diffusion backbone network on the matched image patch to complete the parameter update corresponding to the current parameter update step.

[0006] As an optional implementation, in a first aspect of the invention, the frequency-aware gating module includes a Transformer encoder with a multi-head self-attention mechanism. Furthermore, the frequency-aware gating module predicts the diffusion loss corresponding to the current parameter update step for each image block, and determines the probability value corresponding to each image block based on the diffusion loss, including: The sequence of image block tokens corresponding to all the image blocks is input into the Transformer encoder to achieve global context-aware image block analysis. The diffusion loss parameter corresponding to each image block is output through a shared linear projection layer. After normalizing all the diffusion loss parameters, the probability value corresponding to each image block is obtained. The diffusion loss parameter is used to measure the diffusion loss corresponding to the image block.

[0007] As an optional implementation, in the first aspect of the present invention, the Transformer encoder is trained using a preset training dataset, the training dataset including multiple training image patch tokens and a training diffusion loss corresponding to each training image patch token, the Transformer encoder being used to predict the corresponding diffusion loss based on the image patch tokens.

[0008] As an optional implementation, in the first aspect of the present invention, each parameter update step in the training process of the diffusion model to be trained further includes: For the current frequency sensing gating module, based on the REINFORCE policy gradient estimation algorithm, the diffusion loss corresponding to the matching image block minus the benchmark value is used as the advantage signal to perform policy gradient update on the control parameters corresponding to the current frequency sensing gating module, so as to update the current frequency sensing gating module. The baseline value is the exponential moving average of the diffusion loss corresponding to all the image blocks.

[0009] As an optional implementation, in the first aspect of the present invention, the frequency-aware gating module includes a 2-layer Transformer encoder and a 4-head self-attention sensor, and the hidden layer dimension of the frequency-aware gating module is 128.

[0010] As an optional implementation, in the first aspect of the present invention, the diffusion model to be trained includes a UNet model and a DiT model.

[0011] As an optional implementation, in the first aspect of the invention, transforming each image block from the spatial domain to the frequency domain to obtain the frequency component corresponding to each image block includes: Perform a two-dimensional Fourier transform on each image block to transform each image block from the spatial domain to the frequency domain, and obtain the frequency components corresponding to each image block; And, for each of the image blocks: the frequency component corresponding to the image block, the position information of the image block, and the embedding vector are fused to obtain an image block token, including: For each image block: the frequency component corresponding to the image block is concatenated with the embedding vector and then the position information of the image block is added to obtain the image block token.

[0012] A second aspect of this invention discloses a model training method based on frequency-aware latent space and pixel space fusion, the method comprising: Obtain a training dataset, which includes multiple training images with different resolutions; Based on the resolution of each training image, the training images are divided into high-resolution images and low-resolution images; The high-resolution image is input into the latent space for training, and the low-resolution image is input into the pixel space for training to obtain a trained target diffusion model, which is used to perform the corresponding image processing task. Both the latent space and the pixel space include a diffusion model to be trained, and the latent space and the pixel space are trained using a model training method based on frequency-aware image patch selection as described in any of the first aspects of the present invention.

[0013] A third aspect of the present invention discloses an image task processing method, the method comprising: The image task to be processed is processed using a target diffusion model trained as described in either the first or second aspect of the present invention. The image task to be processed includes an unconditional image generation task or a text-conditional image generation task.

[0014] A fourth aspect of the present invention discloses an apparatus comprising a memory and a processor, the apparatus comprising: Memory containing executable program code; A processor coupled to the memory; The processor calls the executable program code stored in the memory to execute the model training method based on frequency-aware image patch selection as described in any of the first aspects of the present invention, or the model training method based on frequency-aware latent space and pixel space fusion as described in the second aspect of the present invention.

[0015] Compared with the prior art, the present invention has the following beneficial effects: This invention designs a frequency-aware image patch selection module that adaptively selects image patches with the most valuable frequency information for computation in each training step, for the current diffusion time step and training stage. This concentrates limited computational resources on the data most helpful to the model's current learning, significantly reducing the training time and computational load of the diffusion model without reducing, and even improving, the generation quality. Essentially, this invention upgrades the static computational allocation method of "treating all pixels / image patches equally" in traditional diffusion training to a dynamic adaptive computational allocation method driven by both image frequency content and the current training state. This ensures that computational resources are always tilted towards the frequency content most valuable for current learning, eliminating a large amount of redundant computation and thus significantly improving training efficiency while maintaining learning effectiveness. Attached Figure Description

[0016] To more clearly illustrate the technical solutions in the embodiments of the present invention, the accompanying drawings used in the description of the embodiments will be briefly introduced below. Obviously, the accompanying drawings described below are only some embodiments of the present invention. For those skilled in the art, other drawings can be obtained based on these drawings without creative effort.

[0017] Figure 1(a) is a schematic diagram of the intermediate generated samples at different time steps in the diffusion sampling process and their corresponding two-dimensional Fourier amplitude spectra; Figure 1(b) is a scatter plot showing the changes of intrinsic model indices with training iterations during the diffusion model training process; Figure 2 This is a flowchart illustrating a model training method based on frequency-aware image patch selection disclosed in an embodiment of the present invention. Figure 3 This is an internal structural diagram of the frequency sensing gating module disclosed in an embodiment of the present invention; Figure 4 This is a schematic flowchart of a model training method based on frequency-aware latent space and pixel space fusion disclosed in an embodiment of the present invention. Figure 5This is a schematic diagram of the structure of a device including a memory and a processor disclosed in an embodiment of the present invention. Detailed Implementation

[0018] To enable those skilled in the art to better understand the present invention, the technical solutions of the present invention will be clearly and completely described below with reference to the accompanying drawings of the embodiments of the present invention. Obviously, the described embodiments are only some embodiments of the present invention, and not all embodiments. Based on the embodiments of the present invention, all other embodiments obtained by those skilled in the art without creative effort are within the scope of protection of the present invention.

[0019] The terms "first," "second," etc., used in the specification, claims, and accompanying drawings of this invention are used to distinguish different objects, not to describe a specific order. Furthermore, the terms "comprising" and "having," and any variations thereof, are intended to cover non-exclusive inclusion. For example, a process, method, apparatus, product, or end that includes a series of steps or units is not limited to the listed steps or units, but may optionally include steps or units not listed, or may optionally include other steps or units inherent to these processes, methods, products, or ends.

[0020] In this document, the term "embodiment" means that a particular feature, structure, or characteristic described in connection with an embodiment may be included in at least one embodiment of the invention. The appearance of this phrase in various places throughout the specification does not necessarily refer to the same embodiment, nor is it a separate or alternative embodiment mutually exclusive with other embodiments. It will be explicitly and implicitly understood by those skilled in the art that the embodiments described herein can be combined with other embodiments.

[0021] This invention discloses a model training method, an image task processing method, and related apparatus for accelerating diffusion model training in image task processing scenarios. These are described in detail below.

[0022] Example 1 Deep neural networks exhibit an inherent tendency known as "spectral bias": during training, the network tends to preferentially fit the low-frequency components of the objective function before gradually learning high-frequency details. This phenomenon has been widely verified in various neural networks and tasks. This invention further reveals that in diffusion models, spectral bias exists in a unique "dual" form, which is systematically demonstrated through empirical research: First, observing along the diffusion time step dimension, the generated samples during diffusion sampling gradually evolve from coarse-grained low-frequency structures to fine high-frequency details. Denoising in early time steps (when there is more noise) primarily reconstructs the coarse outline of the image, while later time steps (when there is less noise) gradually restore high-frequency textures and details. Second, observing along the training iteration dimension, after performing principal component analysis (PCA) on samples generated at different training stages, the principal components corresponding to low-frequency components (low-index intrinsic modes) appear much earlier than the principal components corresponding to high-frequency components (high-index intrinsic modes), indicating that the model also preferentially learns low-frequency information during training.

[0023] The empirical results of the embodiments of the present invention are as follows: Figure 1(a) shows the intermediate generated samples at different time steps during the diffusion sampling process and their corresponding two-dimensional Fourier amplitude spectra visualization. As the time steps decrease from large (100) to small (700, 1000), the central region of the spectra corresponds to low-frequency components, while the outer region corresponds to high-frequency components. It can be observed that the spectral energy in the early time steps is highly concentrated in the center, indicating that the reconstruction at this stage is mainly based on low-frequency coarse structures; as the time steps progress, the spectral energy in the outer regions gradually increases, and high-frequency details gradually appear.

[0024] Figure 1(b) is a scatter plot showing the change of intrinsic mode indices (vertical axis, corresponding to the PCA principal components of the generated samples) with training iterations (horizontal axis). The scatter plot marks the training iteration in which each intrinsic mode first appears. The curve shows a monotonically increasing trend, indicating that low-indice (low-frequency) intrinsic modes have appeared in the early stages of training, while high-indice (high-frequency) intrinsic modes gradually appear in the later stages of training, which intuitively shows the spectral bias in the training iteration dimension.

[0025] The discovery of this "dual spectral bias" reveals a significant amount of redundant computation in the training of the diffusion model: at different stages and time steps of training, the frequency information carried by a large number of image patches in the image is not the most useful for the current learning stage, and performing full forward and backward propagation on these invalid image patches results in considerable computational waste.

[0026] Based on the above empirical research results, this invention discloses a model training method based on frequency-aware image patch selection. This method can be integrated into a corresponding model training device, which in turn can be integrated into a corresponding server. Figure 2 As shown, the method includes: Step 101: Train the diffusion model to be trained using a preset input image to obtain a trained target diffusion model. The target diffusion model is used to perform the corresponding image processing task.

[0027] In this embodiment of the invention, the image processing task can be any one of image synthesis, video generation, text-to-image generation, semantic segmentation, etc.

[0028] As will be understood by those skilled in the art, the diffusion model includes a forward process (adding noise) and a backward process (denoising), and is trained during the backward process. The training objective of the diffusion model is to enable the model to learn to restore the image step by step into a clear and realistic image. Therefore, the diffusion model includes multiple parameter update steps step by step.

[0029] Each parameter update step in the training process of the diffusion model to be trained may include: The input image is divided into multiple image blocks, and each image block is transformed from the spatial domain to the frequency domain to obtain the frequency component corresponding to each image block. The transformation of each image block from the spatial domain to the frequency domain can be achieved by using Fourier transform, discrete cosine transform, wavelet transform, Walsh-Hadamard transform, etc.

[0030] Obtain the diffusion time step corresponding to the current parameter update step, and obtain the embedding vector based on the diffusion time step; For each image block: fuse the frequency component corresponding to the image block, the position information of the image block, and the embedding vector to obtain the image block token; Input all image block tokens into the pre-set frequency-aware gating module. The frequency-aware gating module predicts the diffusion loss corresponding to the current parameter update step for each image block, and determines the probability value corresponding to each image block based on the diffusion loss. Based on the probability values ​​corresponding to all image patches, select at least one matching image patch corresponding to the current parameter update step; Perform forward and backward propagation of the diffusion backbone network on the matched image patch to complete the parameter update corresponding to the current parameter update step.

[0031] Specifically, based on the principle of dual spectral bias, this embodiment of the invention designs a lightweight frequency-aware gating module responsible for dynamically evaluating the current learning value of each candidate image patch in each training step and outputting the selection probability. The module's workflow is as follows: First, the input image is segmented into several image patches. Each image patch is transformed from a spatial domain representation to a frequency domain representation to encode its frequency content distribution. Then, the frequency representation is concatenated with the embedding vector of the current diffusion time step to form an image patch token carrying both frequency and time step information. Using the sequence of all image patch tokens as input, the frequency-aware gating module of this embodiment processes the data, enabling the representation of each image patch to be context-modeled with reference to information from other image patches. The frequency-aware gating module predicts the diffusion loss corresponding to the current parameter update step for each image patch and determines the probability value corresponding to each image patch based on this diffusion loss. This design allows the importance evaluation of image patches to not only depend on the frequency content of the image patch itself but also to perceive the frequency composition of the global image and the current diffusion time step, thus naturally aligning with the dynamic evolution of the dual spectral bias.

[0032] As can be seen, this embodiment of the invention designs a frequency-aware image patch selection module, which adaptively selects the image patches with the most valuable frequency information for the current diffusion time step and the current training stage in each training step. This concentrates the limited computational budget on the data most helpful to the model's current learning, significantly compressing the training time and computational load of the diffusion model without reducing or even improving the generation quality. The essence of this embodiment is to upgrade the static computational allocation method of "treating all pixels / image patches equally" in traditional diffusion training to a dynamic adaptive computational allocation method driven by both image frequency content and the current training state. This ensures that computational resources are always tilted towards the frequency content most valuable to the current learning, eliminating a large amount of redundant computation, thereby significantly improving training efficiency while maintaining learning effectiveness.

[0033] In an optional embodiment, the frequency-aware gating module includes a Transformer encoder with a multi-head self-attention mechanism; Furthermore, the frequency-aware gating module predicts the diffusion loss corresponding to the current parameter update step for each image patch, and determines the probability value corresponding to each image patch based on the diffusion loss, including: The sequence of image patch tokens corresponding to all image patches is input into the Transformer encoder to achieve global context-aware image patch analysis. The diffusion loss parameter corresponding to each image patch is output through a shared linear projection layer. After normalizing all diffusion loss parameters, the probability value corresponding to each image patch is obtained. The diffusion loss parameter is used to measure the diffusion loss corresponding to the image patch.

[0034] In this optional embodiment, the core operation of the frequency-aware gating module is to sample from a probability distribution to determine which image patch to select. This is a discrete decision-making process, inherently non-differentiable, and cannot be directly updated using the standard backpropagation algorithm. Therefore, in this optional embodiment, the frequency-aware gating module includes a Transformer encoder with a multi-head self-attention mechanism, which models this discrete decision problem as a stochastic policy network framework in reinforcement learning: the frequency-aware gating module is regarded as a stochastic policy whose "action" is to select an image patch, and the "reward" (a negative value) is the diffusion training loss on the selected image patch, i.e., the diffusion loss, which is measured by a diffusion loss parameter, for example, by its numerical value.

[0035] In this optional embodiment, the Transformer encoder is further optionally trained using a preset training dataset, which includes multiple training image patch tokens and a training diffusion loss corresponding to each training image patch token. The Transformer encoder is used to predict the corresponding diffusion loss based on the image patch tokens.

[0036] However, the frequency-aware gating module also needs to update itself in real time according to the training process of the diffusion model. Therefore, this embodiment of the invention also designs an end-to-end differentiable optimization method based on policy gradient for updating the frequency-aware gating module. In another optional embodiment, each parameter update step in the training process of the diffusion model to be trained may further include: For the current frequency sensing gating module, based on the REINFORCE policy gradient estimation algorithm, the diffusion loss corresponding to the matching image patch minus the benchmark value is used as the advantage signal to perform policy gradient update on the control parameters corresponding to the current frequency sensing gating module, so as to update the current frequency sensing gating module. The baseline value is the exponential moving average of the diffusion loss corresponding to all image patches.

[0037] In this optional embodiment, the frequency-aware gating module is treated as a stochastic policy whose "action" is to select an image patch, and the "reward" (a negative value) is the diffusion training loss on the selected image patch. Based on the REINFORCE policy gradient estimation method, this optional embodiment derives the unbiased gradient estimate of the gating module parameters, thereby achieving joint end-to-end training of the diffusion model parameters and the gating module parameters. This allows the gating module to continuously self-optimize during training, always assigning the selection probability to the image patch most valuable to the current learning stage. To reduce the variance of the policy gradient estimate for stable training, this optional embodiment introduces a baseline subtraction technique based on the exponential moving average of the diffusion loss, effectively reducing the estimation variance without changing the gradient expectation.

[0038] However, introducing a frequency-aware gating module into the diffusion model requires consideration of whether it impairs the standard non-convex convergence behavior of the diffusion model training. This optional embodiment finds that, under standard assumptions (loss function and its gradient satisfying Lipschitz continuity, bounded variance of stochastic gradient estimators, etc.), the synchronous update of the diffusion model parameters and the gating module parameters satisfies the non-convex convergence theorem, indicating that the squared norm of the average gradient decreases at a rate of 1 / N with increasing iterations, converging to a controlled-size neighborhood under the influence of random noise. Therefore, this optional embodiment proves that introducing a frequency-aware gating module does not impair the standard non-convex convergence behavior of the diffusion model training; theoretically, adding the gating module has the same convergence properties as standard training without the gating module.

[0039] In another alternative embodiment, the frequency-aware gating module may include a 2-layer Transformer encoder, a 4-head self-attention sensor, and the hidden layer dimension of the frequency-aware gating module is 128.

[0040] In yet another alternative embodiment, the diffusion model to be trained may include a UNet model and a DiT model.

[0041] In yet another optional embodiment, transforming each image patch from the spatial domain to the frequency domain to obtain the frequency component corresponding to each image patch may include: Perform a two-dimensional Fourier transform on each image block to transform each image block from the spatial domain to the frequency domain, and obtain the frequency components corresponding to each image block; Furthermore, for each image patch: the image patch token is obtained by fusing the frequency component corresponding to the image patch, the location information of the image patch, and the embedding vector, which may include: For each image block: the frequency component corresponding to the image block is concatenated with the embedding vector and then the position information of the image block is added to obtain the image block token.

[0042] The specific process of this optional embodiment can be as follows: Figure 3 The diagram shows a detailed internal structure of the frequency-aware gating module. From left to right: Image block x_k undergoes a 2D DFT (2D Fourier Transform) module to obtain the frequency representation F[x_k]; this is concatenated with the time-step embedding t and then combined with the position embedding to form a complete image block token h_k; K image block tokens form the sequence input Transformer encoder (including multi-head self-attention); the output of the Transformer is projected onto a shared linear layer (Linear) to obtain a scalar score (diffusion loss parameter), which is then normalized by Softmax to obtain the selection probability Prob(x_k). This structure diagram clearly illustrates the four core steps: frequency feature extraction, time-step conditional injection, multi-image block interactive modeling, and probability output.

[0043] Example 2 Based on the same concept, this invention discloses a fast training method for diffusion models applicable to image processing tasks, called FaPS (Frequency-aware Patch Selection). FaPS is based on an image patch-level training framework (such as PatchDiffusion), on which a frequency-aware gating module is superimposed. This replaces the original static strategy of randomly selecting image patches with a dynamic adaptive selection strategy driven by both frequency content and diffusion time steps, thereby significantly reducing the computational cost of training while maintaining or even improving the generation quality.

[0044] The FaPS framework has broad versatility: at the network architecture level, it supports two mainstream backbone networks, UNet (such as EDM-DDPM++, EDM-ADM) and DiT (DiffusionTransformer, such as DiT-XL / 2); at the training space level, it supports training in both the original pixel space and the latent diffusion space; at the task level, it is suitable for unconditional image generation and text-conditional image generation (text-to-image) tasks; and at the resolution level, it is compatible with both low-resolution (such as 64×64) and high-resolution (such as 256×256) training scenarios.

[0045] I. Detailed implementation of the frequency sensing gating module: Given a training image, FaPS first segments it into K candidate image patches. For each image patch, the following operation is performed to construct an image patch token: Two-dimensional Fourier transform: Performing a 2D discrete Fourier transform on an image patch converts it from pixel space to frequency domain representation, obtaining the amplitude and phase information of each spatial frequency component, thereby encoding the distribution characteristics of the image patch in low frequency (smooth structure) and high frequency (texture details).

[0046] Temporal step embedding concatenation: Calculate the embedding vector for the diffusion time step sampled in this training iteration, and concatenate it with the frequency representation of the image patch to form a composite image patch token carrying frequency information and temporal context.

[0047] Location encoding: Learnable location embeddings are added to each image patch token to encode the spatial location information of the image patch in the original image.

[0048] The gating module takes a sequence of tokens from all K image patches as input and processes it through a lightweight Transformer encoder (multi-head self-attention + feedforward network). This allows the representation of each image patch to integrate information from other image patches, achieving a globally context-aware evaluation of image patch importance. Subsequently, a scalar score for each image patch is output through a shared linear projection layer, which is then normalized using Softmax to obtain a selection probability distribution. In each training step, an image patch is sampled from this probability distribution, and forward and backward propagation of the diffusion backbone network is performed only on that image patch to complete the parameter update for that step.

[0049] II. Joint Optimization Algorithm (FaPS Training Process) The complete training process for FaPS is as follows: Sample an image from the training dataset and sample a diffusion time step from a uniform distribution.

[0050] The image is divided into blocks, and a Fourier transform is performed on each block. The transformed blocks are then concatenated with the time step embeddings to construct an image block token sequence.

[0051] The image block token sequence is input into the frequency-aware gating module to obtain the selection probability distribution of each image block; an image block is sampled from this distribution, and diffusion noise is sampled at the same time.

[0052] Diffusion model parameter update: Calculate the diffusion denoising loss for the selected image patch (only the selected image patch is involved, not the whole image), and perform gradient descent update on the diffusion model parameters.

[0053] Gating module parameter update: Based on REINFORCE policy gradient estimation, the diffusion loss on the selected image patch minus the baseline value (exponential moving average of the loss) is used as the advantage signal to perform policy gradient update on the gating module parameters; at the same time, the baseline value is updated with a decay factor of 0.99.

[0054] Repeat the above steps until training is complete.

[0055] In the above process, the parameter updates of the diffusion model and the parameter updates of the gating module are performed synchronously at each step. The two work together to optimize the overall expected diffusion loss, forming an end-to-end joint training.

[0056] III. Coordinate Normalization and Image Patch Management To ensure the consistency and portability of image patch location information, this invention normalizes the image coordinate system to the range [-1, 1] (with the top left corner as (-1, -1) and the bottom right corner as (1, 1)), consistent with the processing method of PatchDiffusion. For an image with resolution R×R and image patch size s, the candidate image patch pool contains candidate patches at (R-s+1)² locations. During training, the gating module adaptively selects the most informative patch from these, rather than traversing all candidate patches, thereby avoiding redundant computation.

[0057] Table 1. Key Hyperparameters:

[0058] Furthermore, existing diffusion model training methods typically include standard diffusion training, patch diffusion, time step redistribution, etc. Table 2 compares the FaPS disclosed in this embodiment with existing diffusion model training methods from the perspectives of training units, block selection strategies, frequency awareness, time step adaptability, end-to-end learnability, convergence theory, maximum speedup, and data efficiency improvement.

[0059] Table 2 Comparison of FaPS and Existing Technologies:

[0060] Advantages and positive effects: 1. Technological advantages Highly versatile and architecture-independent, FaPS supports both UNet (DDPM++, ADM) and DiT (DiT-XL / 2) as mainstream diffusion backbones, is compatible with pixel space and latent diffusion space, is suitable for unconditional generation and text conditional generation, has broad versatility, and can be integrated into existing diffusion training workflows as a plug-and-play module.

[0061] Adaptive dynamic selection precisely matches learning needs. The gating module, in conjunction with the frequency content of image patches and the current diffusion time step, dynamically adjusts the selection probability of each image patch, ensuring that the computation is always focused on the frequency content most valuable for the current training phase, thus eliminating the large amount of redundant computation caused by random selection strategies.

[0062] It is end-to-end learnable and requires no manual rules. The gating module parameters are jointly optimized through policy gradient and diffusion model, eliminating the need for manually predefined fixed-frequency selection rules. It can automatically adjust selection preferences as training progresses, exhibiting excellent adaptability.

[0063] It can be orthogonally superimposed with other acceleration methods. The acceleration mechanism of FaPS (image block-level frequency-aware selection) and the time step redistribution method (such as SpeeD) are independent of each other in principle. The superposition of the two can achieve further additional acceleration. Experiments show that at least 1.4× additional acceleration can be obtained on the basis of SpeeD.

[0064] It effectively improves the generation quality in scenarios with limited data. Since image patch-level training allows a single image to be expanded into hundreds of training patches, FaPS inherently has a data augmentation effect. On small datasets (such as AFHQv2), the FID is further reduced compared to PatchDiffusion, demonstrating its additional value in scenarios with limited data.

[0065] It possesses rigorous theoretical convergence guarantees. The FaPS joint optimization algorithm exhibits non-convex convergence, demonstrating that introducing frequency-aware gating does not impair the theoretical convergence behavior of training, and the method has a solid mathematical foundation.

[0066] 2. Main performance indicators This invention compares the performance of FaPS with several commonly used baseline methods in the field (EDM-DDPM++, PatchDiffusion, LDM-ADM, MDT, SpeeD) on different datasets, backbone networks, and resolutions. Evaluation metrics include: FID (Freche Initial Distance): Measures the quality of the generated image; the lower the better.

[0067] Speedup (training speedup ratio): The speedup factor relative to standard training (1×), the higher the better.

[0068] NFE (Number of Function Evaluations): The number of times the model function is evaluated during sampling, which is fixed at 50 or 250 in this comparison.

[0069] Table 3 shows the FID metric and training acceleration effect of FaPS on standard benchmark datasets (lower FID is better, higher Speedup is better); among them, the standard benchmark datasets selected are CIFAR-10, ImageNet 64×64, LSUN Bedroom, and ImageNet 256×256 datasets commonly used in this field.

[0070] Table 3. FID index and training acceleration effect of FaPS:

[0071] The experimental verification in Table 3 shows that: (1) FaPS is an effective method for accelerating diffusion models. It can stably provide training acceleration of 2×~4.5× on various datasets, backbone networks (UNet, ADM, DiT) and resolutions (64², 256²).

[0072] (2) In most cases, FaPS can maintain or even reduce FID (improve generation quality); in some cases (such as ImageNet 64×64), FID increases slightly, but it is still better than similar acceleration methods such as Patch Diffusion, and the acceleration benefits are significant.

[0073] (3) FaPS has good compatibility and can be combined with other acceleration technologies (such as SpeeD) to achieve a higher speedup ratio (4.5×+) while improving quality.

[0074] Example 3 This invention discloses a method for accelerating low-resolution face image generation (CelebA / FFHQ, UNet backbone), which includes the following steps: Environment configuration: PyTorch is deployed on a server equipped with 8 NVIDIA A100 (80G) GPUs, and the FaPS module is integrated based on the EDM code repository.

[0075] Network configuration: EDM-DDPM++ is used as the diffusion backbone, with an image resolution of 64×64; the frequency-aware gating module adopts a 2-layer Transformer encoder, 4-head self-attention, and a hidden layer dimension of 128. The number of parameters is about 0.3% of the backbone network, and the computational overhead is negligible.

[0076] Image block setting: Divide each 64×64 image into several candidate image blocks of randomly selected size in the range of 8~64 pixels, and normalize the coordinates to [-1,1].

[0077] Training configuration: The AdamW optimizer is used, the learning rate of the diffusion backbone is the same as the original EDM settings, and the learning rate of the gating module is 1×e. -4 ; EMA decay factor beta=0.99; The total number of training iterations is the same as the baseline, but each step only processes a single image patch selected by frequency-aware gating.

[0078] Gating module initialization: In the early stages of training, the gating module output is approximately uniformly distributed. As training progresses, it gradually learns the selection preference that matches the current spectral bias stage. For the principle of the gating module, please refer to Embodiment 1 or Embodiment 2 of this invention.

[0079] Evaluation: Generate 50,000 samples at regular intervals to calculate FID (compare with the full set of real images), plot the curve of FID over training time, and record the training time required to reach the baseline final FID.

[0080] Implementation effect FaPS achieved an FID of 1.65 on CelebA-64×64, outperforming the baseline EDM-DDPM++'s 1.66, with training time approximately 1 / 3 of the baseline (about 3× speedup); on FFHQ-64×64, it achieved an FID of 2.98, with a speedup of approximately 3×.

[0081] Example 4 Background and Objectives: A content platform needs to train a high-quality scene generation model on the LSUN-Bedroom and LSUN-Church datasets (256×256). The original LDM-ADM baseline FID is approximately 4.32 (Bedroom) and 4.66 (Church). The goal is to simultaneously improve generation quality (reduce FID) and significantly improve training efficiency.

[0082] To address this, this invention discloses a method for accelerating high-resolution scene image generation (LSUN-256×256, ADM backbone + latent diffusion), which includes the following steps: Latent space setup: 256×256 images are compressed into a low-dimensional latent representation using a VAE encoder pre-trained with StableDiffusion. FaPS image patch selection and diffusion training are performed in the latent space, with ADM as the diffusion backbone of the latent space.

[0083] Frequency-aware gating: Two-dimensional DFT and image block token construction are performed on the feature map of the latent space. The gating module structure is the same as in Example 1.

[0084] Training and Evaluation: The remaining configurations are set according to the EDM-ADM standard. After training, the VAE decoder is used to restore the generated latent code to a 256×256 image and calculate the FID.

[0085] Implementation effect Compared to the LDM-ADM baseline, FaPS reduced FID from 4.32 to 2.72 on Bedroom (an improvement of approximately 37%) and accelerated training by approximately 3.3×; on Church, it reduced FID from 4.66 to 2.73 (an improvement of approximately 41%) and accelerated training by approximately 3.3×. Compared to the recent strong baseline LPDM-ADM (FID 2.75 / 2.66), FaPS achieved comparable or even better FID, while achieving greater speedup (3.3× vs 2.0×).

[0086] Example 5 This invention discloses a model training method based on frequency-aware latent space and pixel space fusion, the method comprising: Obtain the training dataset, which includes multiple training images with different resolutions; Based on the resolution of each training image, the training images are divided into high-resolution images and low-resolution images; High-resolution images are input into the latent space for training, and low-resolution images are input into the pixel space for training, resulting in a trained target diffusion model. The target diffusion model is then used to perform the corresponding image processing task. The latent space and pixel space both include the diffusion model to be trained, and the latent space and pixel space are trained using any of the diffusion model training methods in Examples 1 to 4.

[0087] In this embodiment of the invention, high-resolution images and low-resolution images can be distinguished by setting a resolution threshold. For example, images with a resolution exceeding the threshold are high-resolution images, and images with a resolution below the threshold are low-resolution images.

[0088] Optional solutions in embodiments of the present invention are as follows: Figure 4 As shown, Figure 4 It includes two training branches: a latent space (upper half, for high-resolution images) and a pixel space (lower half, for low-resolution images). The core component is the central frequency-aware gating module, which receives a sequence of image patch tokens (containing frequency representation and time step embedding) as input, outputs the selection probability of each image patch, and guides subsequent forward and backward propagation of the diffusion backbone network only on the selected image patches.

[0089] Example 6 This invention discloses an image task processing method, which includes: The image task to be processed is processed using any of the pre-trained target diffusion models described in Examples 1 to 5; wherein the image task to be processed may include an unconditional image generation task or a text-conditional image generation task, and the image task to be processed may also include one or more of multiple tasks such as image synthesis, video generation, text-to-image generation, and semantic segmentation.

[0090] Example 7 Please see Figure 5 , Figure 5 This invention discloses an apparatus including a memory and a processor, which may include: Memory 201 storing executable program code; Processor 202 coupled to memory 201; The processor 202 calls the executable program code stored in the memory 201 to execute some or all of the steps in any of the visual language task processing methods based on hierarchical optimal transmission in Embodiment 1 or Embodiment 2 of the present invention.

[0091] Finally, it should be noted that the above embodiments are merely preferred embodiments of the present invention and are only used to illustrate the technical solutions of the present invention, 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. Such 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. A model training method based on frequency-aware image patch selection, characterized in that, The method includes: The training model is trained using a preset input image to obtain a trained target diffusion model, which is used to perform the corresponding image processing task. Each parameter update step in the training process of the diffusion model to be trained includes: The input image is divided into multiple image blocks, and each image block is transformed from the spatial domain to the frequency domain to obtain the frequency component corresponding to each image block. Obtain the diffusion time step corresponding to the current parameter update step, and obtain the embedding vector based on the diffusion time step; For each image block: the frequency component corresponding to the image block, the position information of the image block, and the embedding vector are fused to obtain an image block token; Input all the image block tokens into a pre-set frequency-aware gating module. The frequency-aware gating module predicts the diffusion loss corresponding to the current parameter update step for each image block and determines the probability value corresponding to each image block based on the diffusion loss. Based on the probability values ​​corresponding to all the image blocks, select at least one matching image block corresponding to the current parameter update step; Perform forward and backward propagation of the diffusion backbone network on the matched image patch to complete the parameter update corresponding to the current parameter update step.

2. The model training method based on frequency-aware image patch selection according to claim 1, characterized in that, The frequency-aware gating module includes a Transformer encoder with a multi-head self-attention mechanism; Furthermore, the frequency-aware gating module predicts the diffusion loss corresponding to the current parameter update step for each image block, and determines the probability value corresponding to each image block based on the diffusion loss, including: The sequence of image block tokens corresponding to all the image blocks is input into the Transformer encoder to achieve global context-aware image block analysis. The diffusion loss parameter corresponding to each image block is output through a shared linear projection layer. After normalizing all the diffusion loss parameters, the probability value corresponding to each image block is obtained. The diffusion loss parameter is used to measure the diffusion loss corresponding to the image block.

3. The model training method based on frequency-aware image patch selection according to claim 2, characterized in that, The Transformer encoder is trained using a preset training dataset, which includes multiple training image patch tokens and a training diffusion loss corresponding to each training image patch token. The Transformer encoder is used to predict the corresponding diffusion loss based on the image patch tokens.

4. The model training method based on frequency-aware image patch selection according to claim 2 or 3, characterized in that, Each parameter update step in the training process of the diffusion model to be trained further includes: For the current frequency sensing gating module, based on the REINFORCE policy gradient estimation algorithm, the diffusion loss corresponding to the matching image block minus the benchmark value is used as the advantage signal to perform policy gradient update on the control parameters corresponding to the current frequency sensing gating module, so as to update the current frequency sensing gating module. The baseline value is the exponential moving average of the diffusion loss corresponding to all the image blocks.

5. The model training method based on frequency-aware image patch selection according to claim 1, characterized in that, The frequency-aware gating module includes a 2-layer Transformer encoder and a 4-head self-attention sensor. The hidden layer dimension of the frequency-aware gating module is 128.

6. The model training method based on frequency-aware image patch selection according to claim 1, characterized in that, The diffusion models to be trained include the UNet model and the DiT model.

7. The model training method based on frequency-aware image patch selection according to claim 1, characterized in that, Transforming each image block from the spatial domain to the frequency domain to obtain the frequency components corresponding to each image block includes: Perform a two-dimensional Fourier transform on each image block to transform each image block from the spatial domain to the frequency domain, and obtain the frequency components corresponding to each image block; And, for each of the image blocks: the frequency component corresponding to the image block, the position information of the image block, and the embedding vector are fused to obtain an image block token, including: For each image block: the frequency component corresponding to the image block is concatenated with the embedding vector and then the position information of the image block is added to obtain the image block token.

8. A model training method based on frequency-aware latent space and pixel space fusion, characterized in that, The method includes: Obtain a training dataset, which includes multiple training images with different resolutions; Based on the resolution of each training image, the training images are divided into high-resolution images and low-resolution images; The high-resolution image is input into the latent space for training, and the low-resolution image is input into the pixel space for training to obtain a trained target diffusion model, which is used to perform the corresponding image processing task. Both the latent space and the pixel space include a diffusion model to be trained, and the latent space and the pixel space are trained using the model training method based on frequency-aware image patch selection as described in any one of claims 1-7.

9. An image processing method, characterized in that, The method includes: The image task to be processed is handled using the target diffusion model trained as described in any one of claims 1-8; The image task to be processed includes an unconditional image generation task or a text-conditional image generation task.

10. An apparatus comprising a memory and a processor, characterized in that, The device includes: Memory containing executable program code; A processor coupled to the memory; The processor calls the executable program code stored in the memory to execute the model training method based on frequency-aware image patch selection as described in any one of claims 1-7, or the model training method based on frequency-aware latent space and pixel space fusion as described in claim 8.

Images