Diffusion model optimization method and system based on U-Net parameter enhancement

Abstract

The invention discloses a diffusion model optimization method and system based on U-Net parameter enhancement, and relates to the technical field of diffusion model optimization. The method comprises the following steps: acquiring a data set formed by image and text pairs; constructing a diffusion model, and introducing a learnable mask into the diffusion model for sampling the weight from the U-Net structure; current computing resources are obtained, fine tuning strategies are selected according to the computing resources, and the fine tuning strategies comprise the fine tuning strategy based on training and the fine tuning strategy based on a reward model; and pre-training the diffusion model by using the data set, and optimizing the learnable mask based on the selected fine tuning strategy in the pre-training process to complete optimization of the diffusion model. According to the method, the generalization ability of the pre-training model is kept by avoiding updating U-Net parameters.

Description

Technical Field

[0001] The present invention relates to the technical field of diffusion model optimization, and in particular to a diffusion model optimization method and system based on U-Net parameter enhancement. Background Art

[0002] The statements in this section merely provide background information related to the present invention and do not necessarily constitute prior art.

[0003] Diffusion models are a class of generative models based on an iterative denoising process that have attracted much attention as a powerful tool for generating high-quality representations of images, videos, and 3D data. Diffusion models, such as Stable Diffusion (SD), have successfully applied pre-trained U-Net models to downstream tasks, including personalized text-to-image generation, relation inversion, semantic binding, and controllable generation. Diffusion models establish spatial information in the early denoising stage to represent semantic structure, and then expand to regional details of elements in the later stages. Therefore, in different inference steps, diffusion models use the same network parameters (such as U-Net in SD) to force the learning of different information, such as global structure and features, edges and textures, etc.

[0004] However, traditional classification models, such as ResNet, capture image information such as structural and semantic features at different layers. Typically, shallow layers focus on extracting structural information, while deeper layers capture higher-level semantic information. Similarly, in traditional generative models, the first few layers of the generator control the generation of structural information, while deeper layers represent texture and edge details. Traditional classification and generation tasks both utilize different structures to represent the internal properties of samples, thereby reducing the difficulty of network optimization and enhancing its representational capabilities. Unlike the above two types of models, diffusion models use the same parameters to force the learning of different information when generating samples, but this difference has not been fully explored in diffusion U-Net.

[0005] Currently, strategies for enhancing U-Net by updating model parameters based on pre-trained models typically include introducing trainable modules in specific layers to adapt pre-trained weights to new tasks, selectively fine-tuning a subset of existing parameters, or directly updating all parameters. However, these approaches carry the risk of overfitting. Recently, methods such as LoRA and DoRA have been proposed that inject low-rank matrices into pre-trained weights to increase model flexibility and mitigate overfitting. However, these methods still require adjusting the original parameter space, which may affect the generalization ability of the pre-trained model.

[0006] The non-training methods designed to enhance the generation capabilities of U-Net can be roughly divided into three categories. The first method focuses on adjusting the feature scale. For example, FreeU introduces a sample-dependent scaling factor for U-Net features and suppresses skip connection features to redistribute feature weights, thereby improving generation quality. The second method focuses on optimizing latent encoding using various supervision methods, such as attention maps, noise inversion, or reward models, to enhance the generation performance of U-Net. The third method uses text embedding optimization as a starting point. For example, Chen et al. use balanced text embedding loss to eliminate potential problems in the embedding representation of key tags, thereby improving generation quality. However, the above methods cannot fully tap the inherent potential of the model and are not sensitive enough in the selection of effective parameters, resulting in less than ideal U-Net parameter enhancement effects.

[0007] In summary, how to effectively strengthen the U-Net parameters in the diffusion model without affecting the overall training effect of the model has become a technical problem that needs to be solved urgently in existing technologies. Summary of the Invention

[0008] In view of the shortcomings of the existing technology, the purpose of the present invention is to provide a diffusion model optimization method and system based on U-Net parameter enhancement, and propose a MaskUNet strategy to maintain the generalization ability of the pre-trained model by avoiding updating the U-Net parameters.

[0009] In order to achieve the above object, the present invention is implemented through the following technical solutions:

[0010] A first aspect of the present invention provides a diffusion model optimization method based on U-Net parameter enhancement, comprising the following steps:

[0011] Get a dataset of image and text pairs;

[0012] Construct a diffusion model and introduce a learnable mask into the diffusion model to sample weights from the U-Net structure;

[0013] Obtain the current computing resources and select a fine-tuning strategy based on the computing resources. The fine-tuning strategy includes a training-based fine-tuning strategy and a reward model-based fine-tuning strategy (no training required).

[0014] The diffusion model is pre-trained using the dataset. During the pre-training process, the learnable mask is optimized based on the selected fine-tuning strategy to complete the optimization of the diffusion model.

[0015] Furthermore, in the training-based fine-tuning strategy, the learnable mask is generated by a mask generator, and in the reward model-based (no training required) fine-tuning strategy, the learnable mask is generated under the supervision of the reward model.

[0016] Furthermore, the specific steps for selecting a fine-tuning strategy based on computing resources are as follows:

[0017] When computing resources are sufficient, choose a training-based fine-tuning strategy;

[0018] When computing resources are insufficient, a fine-tuning strategy based on a reward model (no training required) is chosen.

[0019] Furthermore, the diffusion model is used to add noise to the dataset and gradually denoise it to perform data recovery to complete pre-training.

[0020] Furthermore, the specific steps of the training-based fine-tuning strategy are:

[0021] The mask generator generates time-step-dependent and sample-dependent masks as the first strategy mask;

[0022] The first strategy mask is used to sample weights from the U-Net structure, and the sampling process is supervised by the diffusion loss.

[0023] Furthermore, the specific steps of fine-tuning the strategy based on the reward model are as follows:

[0024] Generate a mask as the second policy mask under the supervision of the reward model;

[0025] The second policy mask is used to sample weights from the U-Net structure, and the sampling process is supervised by the reward loss.

[0026] A second aspect of the present invention provides a diffusion model optimization system based on U-Net parameter enhancement, comprising:

[0027] A data acquisition module is configured to acquire a data set consisting of image and text pairs;

[0028] A model building module is configured to build a diffusion model, in which a learnable mask is introduced for sampling weights from a U-Net structure;

[0029] a mask optimization module configured to obtain current computing resources and select a fine-tuning strategy based on the computing resources, wherein the fine-tuning strategy includes a training-based fine-tuning strategy and a reward model-based fine-tuning strategy;

[0030] The pre-training module is configured to pre-train the diffusion model using the data set. During the pre-training process, the learnable mask is optimized based on the selected fine-tuning strategy to complete the optimization of the diffusion model.

[0031] A third aspect of the present invention provides a medium having a program stored thereon, which, when executed by a processor, implements the steps in the diffusion model optimization method based on U-Net parameter enhancement as described in the first aspect of the present invention.

[0032] The fourth aspect of the present invention provides a device, including a memory, a processor, and a program stored in the memory and runnable on the processor. When the processor executes the program, the steps in the diffusion model optimization method based on U-Net parameter enhancement as described in the first aspect of the present invention are implemented.

[0033] A fifth aspect of the present invention provides a computer program product, comprising a computer program, which, when executed by a processor, implements the steps in the diffusion model optimization method based on U-Net parameter enhancement as described in the first aspect of the present invention.

[0034] One or more of the above technical solutions have the following beneficial effects:

[0035] To address the defect that current diffusion models use the same parameters to force learning different information when generating samples, which limits their generation adaptability, the present invention discloses a diffusion model optimization method and system based on U-Net parameter enhancement, aiming to learn a time-step-dependent and sample-dependent mask generation model. The model further selects target parameters from the pre-trained U-Net to enhance the pre-trained U-Net in the diffusion model.

[0036] This paper proposes a novel fine-tuning framework for pre-trained diffusion models for text-to-image, called the MaskUNet strategy. In this strategy, training-based methods optimize masks using a diffusion loss, while non-training methods optimize masks using a reward model. The learnable masks enhance the capabilities of the U-Net while maintaining the model's generalization ability. This strategy improves text-to-image generation without updating any parameters of the pre-trained U-Net.

[0037] This paper introduces a strategy for sampling parameters from a pre-trained U-Net using a learnable binary mask, resulting in a timestep-dependent and sample-dependent U-Net, emphasizing the importance of generation-sensitive parameters. To efficiently learn the mask, the paper designs two fine-tuning strategies: a training-based approach and a training-free approach. In the training-based approach, a parameter sampler generates timestep-dependent and sample-dependent masks, supervised by a diffusion loss. On the other hand, the training-free approach directly generates masks under the supervision of a reward model, eliminating the need for a mask generator compared to the training-based approach.

[0038] This paper conducts an in-depth study on the relationship between pre-trained U-Net parameters, samples and time steps, reveals the effectiveness of parameter independence, and provides a new perspective for the efficient use of U-Net parameters.

[0039] Advantages of additional aspects of the present invention will be given in part in the following description and in part will be obvious from the following description, or will be learned through practice of the present invention. BRIEF DESCRIPTION OF THE DRAWINGS

[0040] The accompanying drawings, which constitute a part of the present invention, are used to provide a further understanding of the present invention. The exemplary embodiments of the present invention and their descriptions are used to explain the present invention and do not constitute improper limitations on the present invention.

[0041] Figure 1 This is a flow chart of a diffusion model optimization method based on U-Net parameter enhancement in Example 1 of the present invention;

[0042] Figure 2 This is a schematic diagram comparing the generation results of different models using different prompt words in Example 1 of the present invention;

[0043] Figure 3 This is a schematic diagram comparing the generation results of different models in the image customization task in the first embodiment of the present invention;

[0044] Figure 4 This is a diagram showing the effect of using the MaskUNet strategy in the text reversal task in Example 1 of the present invention;

[0045] Figure 5 This is a schematic diagram comparing the results generated by different methods in the relationship reversal task in Example 1 of the present invention;

[0046] Figure 6 This is a schematic diagram comparing the results of different methods in the text-to-video generation task in Example 1 of the present invention. DETAILED DESCRIPTION

[0047] It should be noted that the following detailed descriptions are exemplary and intended to provide further explanation of the present invention. Unless otherwise specified, all technical and scientific terms used herein have the same meaning as commonly understood by those skilled in the art to which the present invention belongs.

[0048] It should be noted that the terms used herein are only for describing specific embodiments and are not intended to limit the exemplary embodiments according to the present invention. As used herein, unless the context clearly indicates otherwise, the singular form is intended to include the plural form. In addition, it should be understood that when the terms "comprise" and / or "include" are used in this specification, they indicate the presence of features, steps, operations, devices, components and / or combinations thereof;

[0049] Explanation of terms:

[0050] Diffusion model: A diffusion model is a type of generative model that gradually adds noise to the data and learns how to remove the noise to restore the original data. The overall structure consists of two stages: a forward diffusion process and a reverse generation process. In the forward process, noise (usually Gaussian noise) is added to the original image over multiple time steps, eventually converting it into pure noise. In the reverse process, a neural network (usually a U-Net with a time condition) is trained to gradually remove the noise and reconstruct a clear image. During training, the model learns the denoising method at each step by minimizing the gap between the predicted noise and the actual noise (such as the L2 loss). The entire structure can be viewed as a gradual modeling of the data distribution.

[0051] Timestep dependency: The mask generation depends on the current timestep t in the diffusion process. That is, the model adjusts the mask based on the current denoising stage, activating the appropriate U-Net parameters for that stage. For example, early timesteps may focus on recovering structural information, while later timesteps may focus more on texture details, so the mask should also change accordingly.

[0052] Sample dependency: Mask generation also depends on the specific sample being processed, i.e., the input image or latent code z. Each sample contains different feature information, and the model selectively activates the U-Net parameters based on these differences, avoiding the use of all weights.

[0053] CLIP score: The CLIP score is used to assess the degree of match between an image and its description. The Contrastive Language-Image Pretraining model uses contrastive learning to map images and text into the same vector space, thereby measuring their similarity. A higher CLIP score indicates a stronger match between the image and the description.

[0054] BLIP-VQA score: The BLIP-VQA score evaluates the match between an image and its corresponding text description (e.g., question and answer). The Bootstrapping Language-Image Pretraining model is specifically optimized for image question answering. The BLIP-VQA score assesses image quality by checking how well the model correctly answers questions about the image.

[0055] GENEVAL: Generative Evaluation (GENEVAL) is a framework specifically designed to evaluate the quality of text-guided image generation, focusing on whether the image accurately expresses semantic elements such as attributes, position, and quantity in the text.

[0056] Example 1:

[0057] Example 1 of the present invention provides a diffusion model optimization method based on U-Net parameter enhancement. It designs a MaskUNet strategy to enhance generation quality using minimal parameters. This method fully utilizes the time-step- and sample-dependent effective U-Net parameters and provides two fine-tuning strategies: a training-based approach and a training-free approach involving customized networks and optimization functions. This example achieves the best FID score in zero-shot inference on the COCO dataset and further demonstrates its effectiveness in downstream task evaluations.

[0058] like Figure 1 As shown, the specific steps include:

[0059] Step 1: Get a dataset of image and text pairs.

[0060] In this embodiment, image and text pairs are obtained from existing public data to form a data set for training the diffusion model, and image and text pairs are obtained from different public data sources as a test set to test the trained diffusion model.

[0061] Step 2: Build a diffusion model and introduce a learnable mask into the diffusion model to sample weights from the U-Net structure.

[0062] Diffusion models create a Markov chain by adding noise to the data, which approximates a simple prior (usually a Gaussian distribution). A neural network is trained to reverse this process, using the diffusion model to add noise to the dataset and gradually denoise it to complete the pre-training process, learning to extract useful information at each step.

[0063] This embodiment performs a unified preprocessing operation on the training images, mainly including scaling the images to a specified resolution, performing center cropping or random cropping, and then randomly flipping them horizontally. The images are then converted into tensor format and pixel values ​​are normalized to the range of [-1, 1] to improve data consistency and model training effects.

[0064] In the latent diffusion model (LDM), the diffusion process occurs in a low-dimensional latent space rather than in pixel space, which provides significant efficiency improvements. The training objective of LDM can be expressed as minimizing the following loss function:

[0065]

[0066] in, is the diffusion model loss function, represents the mathematical expectation of the joint distribution of the random variable noise ε, sample z and time step t, MSE(·) represents the mean square error, is the noise from the standard Gaussian distribution, z t is the noise latent code at time step t, ε θ (z t ,t) is the noise predicted by the denoising network, and θ represents its parameters.

[0067] In text-to-image diffusion, an additional cue c guides image generation to obtain more controllable outputs, so the training objective is:

[0068]

[0069] in, represents diffusion loss. Denotes the mathematical expectation of the joint distribution of the random variables noise ε, sample z, time step t and additional hint c.

[0070] Step 3: Obtain the current computing resources and select a fine-tuning strategy based on the computing resources. The fine-tuning strategy includes a training-based fine-tuning strategy and a reward model-based fine-tuning strategy.

[0071] Step 3.1: Obtain the current computing resources. The specific steps for selecting a fine-tuning strategy based on the computing resources are as follows:

[0072] When computing resources are sufficient, choose a training-based fine-tuning strategy.

[0073] When computing resources are insufficient, a fine-tuning strategy based on a reward model is chosen.

[0074] Whether computing resources are sufficient is determined based on the size of the task and the matching of GPU computing resources.

[0075] In this embodiment, the presence of at least six A40 graphics cards is considered sufficient computing power. Step 3.2: In this embodiment, in the training-based fine-tuning strategy, the learnable mask is generated by the mask generator. In the reward-based fine-tuning strategy, the learnable mask is generated under the supervision of the reward model.

[0076] 1. Training-based fine-tuning strategy.

[0077] The mask generator generates time-step-dependent and sample-dependent masks as the first strategy mask; the first strategy mask is used to sample weights from the U-Net structure, and the sampling process is supervised by diffusion loss.

[0078] In a specific implementation, in order to exploit the full potential of the model parameters, a learnable mask is introduced to sample weights from the pre-trained U-Net. When computing resources are sufficient, a training-based fine-tuning method is proposed to optimize the mask, such as Figure 1 The mask is trained using the diffusion loss defined in Equation (2).

[0079] The mask is generated by a mask generator, such as Figure 1 As shown in (b) in the figure. Define the flattened input feature map Where B is the batch size, N is the number of slices, and C in Indicates the number of input channels. The mask generator embeds the time steps and latent coding As input, where H and W represent the height and width of z. Flattening refers to the operation of converting the multi-dimensional feature map in the network into a one-dimensional vector. C1 represents the dimension of the time step embedding vector. represents the domain of real numbers, that is, the set of all real numbers. C represents the number of channels.

[0080] Specifically, merge t emb and z are as follows:

[0081] First, perform input feature fusion. The time step embedding and the latent code of the sample are fused into a unified conditional vector. The formula is as follows:

[0082] z ′ =FC(t emb )+GAP(z)(3).

[0083] in, is the merged output, FC(·) is the fully connected layer, and GAP(·) is the global average pooling.

[0084] Sample dependence is reflected in GAP(z), which is the potential encoding of each sample Perform global average pooling to extract global content features. The time step dependency is reflected in t emb , t emb Indicates the current diffusion stage and provides time information for the current moment. This step implements joint modeling, generating different control signals for each sample at different time steps, embodying the core idea of ​​combining time steps with sample dependencies.

[0085] Then the mask logic value is generated by MLP. A small-scale MLP is used to transform the conditional vector z ′ Mapped to masked logical values ​​(logits), which serve as the basis for subsequent discrete sampling. Specifically, a 4-layer MLP is applied, which includes 2 ReLU activations to introduce nonlinearity. The formula is as follows:

[0086]

[0087] in, is the MLP output.

[0088] This step introduces nonlinear modeling capabilities, allowing the model to generate differentiated mask patterns based on different inputs. The MLP weights are learnable, but they are small and independent of the U-Net weights.

[0089] Again, Gumbel-Sigmoid activation is used to generate the probabilistic mask.

[0090] To sample the weights, Treated as a binary mask, the Gumbel-Sigmoid activation function is introduced to generate an approximately discrete binary mask. The formula is as follows:

[0091]

[0092] in, is an intermediate representation generated by the input through a multi-layer perceptron (MLP), which is based on the time step embedding t emb The features generated by the latent code z of the sample are used to further generate the mask m.

[0093] is the output of the Gumbel-Sigmoid activation function σ(·; τ, δ). The temperature coefficient τ∈(0,∞) controls the discreteness of m: as τ→0, m approaches a binomial distribution; as τ→∞, m approaches a uniform distribution. The threshold δ discretizes the probability distribution. C2 represents the number of channels in the mask.

[0094] The temperature coefficient τ∈(0,∞) controls the discreteness of m: as τ→0, m tends to binomial distribution; as τ→∞, m tends to uniform distribution. The threshold δ is used to discretize the probability distribution. The output It is a mask customized for each sample and time step.

[0095] Next, mask reshape is applied to the weight tensor.

[0096] The weight of a certain layer in U-Net that needs to be applied Reshape m into the same shape as it, that is This allows each sample mask m to be ′ It can be element-wise multiplied with the linear layer parameters in the network. reshape is a function that changes the shape. m needs to be reshaped to the same shape as the weight w in order to perform element-wise multiplication.

[0097] The mask Reshape into The dimensions are consistent, and the reshaped Linear layer weights applied to U-Net Then perform element-wise multiplication to obtain the mask weight, the formula is as follows:

[0098]

[0099] in, is the mask weight, ⊙ represents element-wise multiplication, C out is the number of output channels.

[0100] This operation allows the subset of U-Net parameters used by each sample at each time step to change dynamically by selectively retaining or suppressing elements in the weight matrix. Since the mask is sample- and time-step-dependent, the network path activated by each forward propagation is unique. This is essentially a sub-network selection mechanism that can achieve flexible reconstruction of the parameter space while freezing the main network weights (i.e., without updating the U-Net parameters). The sample / time-step dependency is reflected in the mask coming from MLP (z ′ ), and z ′ Embedded by time step t emb It is determined jointly by the potential code z of the sample. Therefore, at different time steps and different samples, the generated mask m is different, thereby activating different weight subsets.

[0101] Finally, forward propagation is performed using the masked weights.

[0102] Combine the input feature z with the masked weight Perform batch matrix-matrix multiplication to obtain the output features and get the output o, as follows:

[0103]

[0104] in, For output features, BMM(·,·) represents batch matrix-matrix multiplication. By introducing a mask generator, the pre-trained U-Net weights are expected to dynamically adapt to different sample features and time-step embeddings. Notably, this design does not update the U-Net parameters; instead, it leverages sample- and time-step-dependent adjustments to enable the model to selectively activate specific U-Net weights specific to each input. This approach enhances the flexibility of the pre-trained U-Net while maintaining the stability of the pre-trained structure.

[0105] It's important to note that masks are learnable, allowing us to optimize U-Net parameters and selectively augment them, implicitly learning weights. Masks are multipliers applied to frozen weights, and all gradient updates are applied only to the mask generator, such as fully connected layers and MLPs. U-Net weights are not propagated through gradients and are therefore not modified.

[0106] 2. Fine-tuning strategy based on reward model.

[0107] To further demonstrate the effectiveness of masks, we propose a training-free algorithm to optimize them. We generate time-step-dependent and sample-dependent masks as a second-strategy mask under the supervision of a reward model. We use this second-strategy mask to sample weights from a U-Net architecture, with the sampling process supervised by a reward loss.

[0108] like Figure 1 As shown in (c) in the figure, this embodiment samples the binary mask from the mask logic value through the Gumbel-Sigmoid function, and the mask is applied to the pre-trained weights in an element-wise multiplication to achieve dynamic sub-network selection. The entire process is executed at each diffusion time step t, and λ mask optimization iterations are performed. In each iteration, the current mask logic value is first Generate binary mask Then apply it to the U-Net weight θ to get the sparsely activated sub-network parameters This subnetwork is used to analyze the noise latent variable z t Perform denoising prediction and get z t-1 ←ε θ′ (z t ,t,c), and then further decode to restore the original latent variable z0 and generate the image Generate images through multiple reward models Evaluate and get the weighted reward loss Guided by reinforcement learning or search mechanisms, the mask logic values ​​are updated to gradually improve the semantic consistency scores of generated images.

[0109] Specifically, given the intermediate state z in the denoising process t , the state is obtained by denoising the representation z of the previous step t+1 And obtained by the prompt c, it can be expressed as:

[0110] z t =ε θ (z t+1 ,t+1,c)(8).

[0111] Among them, ε θ (·,·,·) is a pre-trained U-Net, and θ represents its parameters. In a specific embodiment, similar to the training-based method, the mask m is applied to the parameters θ, i.e., θ ′ ←θ⊙m. The key difference is that m does not depend on the generator. Therefore, formula (8) can be rewritten as:

[0112] z t =ε θ′ (z t+1 ,t+1,c)(9).

[0113] Then z t Decode to pixel space through VAE to get x0 ′ . Use x0 ′ and prompt c, which are fed into the reward model to calculate the loss. The reward loss is then backpropagated to update the mask parameters to improve the consistency between the generated image and the prompt. The reward loss can be expressed as:

[0114]

[0115] where Ψ i (·,·) represents the pre-trained reward model, ω i is the balancing factor. In this work, n=2 is set, where n is the number of reward models. This example uses ImageReward and HPSv2 as reward models.

[0116] ImageReward is a reward model trained through human preference learning that evaluates the consistency between generated images and text descriptions. By comparing human preferences for image-text pairs, it learns an image quality scoring function that reflects both the visual quality and semantic alignment of an image. HPSv2 is a more powerful image-text reward evaluation model, trained on large-scale human rating data, that comprehensively considers the image's realism, detail quality, and semantic alignment with the text description.

[0117] Both reward models use a similar image-text matching architecture, which mainly consists of the following parts:

[0118] Text encoder: Usually encodes the input text description into a semantic vector based on a pre-trained language model.

[0119] Image Encoder: Uses a vision model to extract the input image into a high-dimensional feature representation.

[0120] Fusion module / matching mechanism: fuses image and text features and calculates the similarity or matching score between them.

[0121] Score predictor: Use a small regression module (such as a fully connected layer) to output a reward score that indicates how well the image matches the text or how high its quality is.

[0122] Algorithm 1 shows the complete steps of fine-tuning without training:

[0123] Algorithm 1 fine-tuning method without training:

[0124] Input setting: text prompt c, a pre-trained unet model ε θ ,Reward Model

[0125] The balancing factor ω of the reward model i ,Optimization iteration number λ, mask logic value l, temperature coefficient τ, threshold δ, maximum time step T

[0126] Parameter initialization settings: m = 1.0, τ = 1.0, δ = 0.5,

[0127] for t=T to 0 do

[0128] for k=0 to λ do

[0129] Get binary mask

[0130] Apply the mask to the pre-trained unet g θ ′:

[0131] Predict the noise latent code z t-1 ←ε θ′ (z t ,t,c)

[0132] Predict the original latent code z0←z t-1

[0133] Decoding to image space

[0134] Reward loss

[0135] Update mask logic value

[0136] end for

[0137] end for

[0138] return

[0139] The specific steps include:

[0140] 1. Input and initialization: Given the input required for the initialization algorithm, including: prompt word c, a pre-trained unet model ε θ ,Reward Model The balancing factor ω of the reward model i ,Optimize the number of iterations λ, mask logic value l, temperature coefficient τ, threshold δ, and maximum time step T. The initial noise image is randomly sampled from the standard Gaussian distribution, and random noise is used as the starting point of the diffusion process. All model parameters (including U-Net weights) remain frozen, and only the mask control structure is optimized.

[0141] 2. Time-step reverse diffusion and binary mask generation: For each time step t = 0 ← T, perform λ mask optimization iterations. In the kth iteration, use the Gumbel-Sigmoid activation function to sample from the mask logits to obtain an approximate binary mask:

[0142]

[0143] The mask is a learnable structural control factor. Different time steps and different sample inputs will generate different masks, reflecting time step dependence and sample dependence.

[0144] 3. Masking the pre-trained U-Net weights: Perform element-wise multiplication of the U-Net weights through the mask to obtain a sub-network with only some parameters activated. This is shown in the following formula:

[0145]

[0146] Without modifying θ, this step generates a dynamic activation path through the mask to build a sub-network related to samples and time steps, improve the expressiveness, and ensure the stability of the pre-trained model.

[0147] 4. Denoising prediction and image generation: through z t-1 =ε θ′ (z t ,t,c) predict the potential code of the current time step, use the masked U-Net to perform back diffusion, decode it to the image space, and obtain the image generation result for subsequent reward evaluation. The overall process remains consistent with the standard diffusion, only the sub-network path is replaced.

[0148] 5. Calculate reward loss: Use ImageReward and HPSv2 as reward models for the current generated image The semantic consistency with the text prompt c is evaluated. The reward loss can be expressed as:

[0149]

[0150] The reward function guides the mask structure to adjust towards semantic alignment, avoids explicit supervision loss, is compatible with different types of text constraints, and lays the feedback foundation for the training-free optimization strategy.

[0151] 6. Update the mask logic value and output the final image: according to Backpropagation updates the mask logic value, gradually optimizing the mask through a reward-driven strategy without changing the main model parameters. This allows for lightweight fine-tuning and ultimately results in an image generation path that better conforms to the text constraints. After all time steps and optimization rounds are completed, the image generated in the last step is returned as the final output.

[0152] Step 4: Use the dataset to pre-train the diffusion model. During the pre-training process, optimize the learnable mask based on the selected fine-tuning strategy to complete the optimization of the diffusion model.

[0153] For zero-shot text-to-image generation, this example fine-tunes MaskUNet on the existing Laion-art dataset, which contains 20.1k image and text pairs. To verify the effectiveness of this example's method, 30k images and 5k images were generated for COCO 2014 and COCO 2017, respectively. Image quality was evaluated using the Frechet Inception Distance (FID) and the CLIP score for image-text alignment.

[0154] We evaluated the effectiveness of MaskUNet on two semantic binding datasets, T2I-CompBench and GenEval, using a reward-based fine-tuning strategy. This example uses the BLIP-VQA score to evaluate attribute correspondences and the GENEVAL score to evaluate image correctness.

[0155] For training-based fine-tuning strategies, this example uses SD 1.5, Full Fine-tune, and LoRA as baselines for comparison with MaskUNet. MaskUNet is also applied to downstream tasks such as image customization, relation reversal, and text-to-video generation. For these tasks, Dreambooth, Textual Inversion, ReVersion, and Text2Video-zero are used as baselines. For reward-based fine-tuning strategies, SD 1.5, SD 2.0, SynGen, and Attend-and-excite are selected as baselines for comparison with MaskUNet.

[0156] The above-mentioned SD 1.5, Full Fine-tune, LoRA, SD 2.0, SynGen and Attend-and-excite methods are all existing methods that have been published. This example only uses them as comparative examples to verify the effectiveness of the method of the present invention, so the specific methods are not repeated here.

[0157] For the training-based fine-tuning strategy, this example uses a learning rate (LR) of 1e-5 and an AdamW optimizer with a weight decay coefficient of 1e-2. Training is performed for 12 epochs, with 50 inference steps. The classifier-free guidance (CFG) is set to 7.5, and DDIM is used as the sampler.

[0158] For the training-based fine-tuning strategy, the number of iterations λ was set to 15. The optimizer used was AdamW, with a learning rate of 1e-2. ImageReward and HPSV2 were used as reward models, with balance coefficients set to 1.0 and 5.0, respectively. The number of inference steps was set to 15, and the CFG was set to 7.5. The sampler used was DPM-Solver.

[0159] Table 1. Quantitative results of zero-shot generation on COCO2014 and COCO2017 datasets

[0160]

[0161] Table 1 shows the zero-shot generation performance of our method and the baseline model on the COCO 2014 and COCO 2017 datasets. For the COCO 2014 dataset, MaskUNet reduces the FID by 1.13 compared to SD v1.5 and by 1.10 compared to LoRA. In contrast, the FID value for Full Fine-tune increases by 1.21 compared to SD v1.5, indicating a risk of overfitting. A similar trend is observed on the COCO 2017 dataset. Overall, by leveraging the dynamic mask mechanism, MaskUNet effectively improves the generation performance of the SD model.

[0162] Figure 2 The left panel shows the generation results of different methods for different cues. In the first row, MaskUNet generates realistic and consistent images, while other methods either introduce artifacts (such as LoRA) or display unnecessary background elements due to overfitting (such as Full Fine-tune). In the third row, MaskUNet accurately captures the mouse's clothing and posture, a consistency that other methods find difficult to achieve. Overall, MaskUNet effectively balances image quality and fidelity to the cue, capturing cue-specific details while maintaining visual consistency across different scenes.

[0163] MaskUnet also has the potential to enhance image quality in various downstream tasks, including image customization, relation reversal, and text-to-video generation tasks.

[0164] 1. Image customization. DreamBooth is a cutting-edge method for image customization, which requires full fine-tuning of U-Net. This example compares the performance of full fine-tuning (DreamBooth), LoRA, and MaskUNet. Figure 3As shown, MaskUNet excels at maintaining subject consistency and background diversity, generating high-quality images across diverse cues, while DreamBooth and LoRA exhibit overfitting. For example, for rare cues like "on the moon," DreamBooth fails to generate consistent images, while LoRA retains undesirable elements from the training set, such as background details. Notably, MaskUNet achieves effective personalization without updating U-Net parameters, demonstrating the untapped potential of pre-trained U-Nets.

[0165] Text reversal. Learn text embedding to capture new concepts and further enhance it by introducing MaskUNet. Figure 4 As shown in Figure 2, adding masks can significantly improve the quality of text reversal generation. For example, the results in the first and second columns show enhanced sensitivity to quantity, while the third column better preserves topic features, resulting in more accurate output.

[0166] 2. Reversion. ReVersion is a relation-guided image synthesis method based on SD, which can be enhanced by integrating MaskUNet. Figure 5 As shown in the previous section, adding masks can improve sensitivity to relational embeddings and enhance image fidelity. For example, for a cue like "in...", ReVersion might place the rabbit on the surface or outside of the cup, but using MaskUnet, sensitivity to the "in..." embedding is increased, resulting in images with the correct relational context. Furthermore, for a cue like "cat," adding masks can significantly improve image quality.

[0167] 3. Text-to-video generation. Tex2Video-Zero is a training-based diffusion model for text-to-video generation. By integrating MaskUnet into Tex2Video-Zero, the continuity and consistency of the generated video can be enhanced, such as Figure 6 For example, for a prompt like “A panda plays guitar in Times Square,” adding a mask can generate a complete guitar. This shows that the mask is orthogonal to Tex2Video-Zero, thus helping to generate high-quality content.

[0168] Table 2 shows the quantitative results of MaskUNet on the T2I-Compbench benchmark without training. It can be observed that MaskUNet achieves over 7% improvement on the color, shape, and texture categories compared to SD v1.5. Compared to SD2.0, MaskUNet slightly lags behind in color but surpasses it in the other two categories. To further validate the generalization ability of MaskUNet, it was applied to SynGen, resulting in an improvement of over 4% across all three categories. Similar findings are found in Table 3 of the GenEval benchmark, where SynGen's color attribute score increases by 21% after applying MaskUNet. Overall, MaskUNet demonstrates strong generalization capabilities on semantic binding tasks.

[0169] Figure 2 The right panel compares samples generated by different methods to evaluate the effectiveness of MaskUNet on the semantic binding task. In the first row, adding MaskUNet to SD 1.5 highlights the semantic information of "book," while adding it to SynGen enhances the texture of the vase from blurry to richly detailed. In the second row, MaskUNet improves sensitivity to quantity and shape. In the third row, it enhances texture generation. In the last row, MaskUNet is able to more accurately adhere to the specified color and object combinations. Overall, MaskUNet significantly improves generation quality on the semantic binding task, demonstrating higher fidelity to the semantic requirements of the cue word. The best results are highlighted in bold in Tables 1, 2, and 3.

[0170] Table 2. Semantic binding evaluation results on the T2I-CompBench dataset

[0171]

[0172] Table 3. Semantic binding evaluation results on the GeneVal dataset

[0173]

[0174] Using t-SNE dimensionality reduction, we visualized the distribution of images generated by MaskUNet, images generated by SD, and real images from COCO 2017. The distribution of images generated by MaskUNet is closer to the real image distribution. This reveals why MaskUNet improves SD generalization. As the number of iterations increases, the mask ratio increases while the FID gradually decreases, indicating that the mask is continuously improving its generative ability.

[0175] Table 4 shows an ablation study of the effectiveness of different inputs to the mask generator for training-based fine-tuning strategies. MaskUNet with timestep embeddings and sample inputs achieves the lowest FID scores. Removing the timestep embeddings or sample inputs results in higher FID scores, with SD 1.5 (no mask) performing the worst. The CLIP scores for all experiments are similar, indicating that the mask primarily improves image quality with little impact on semantic alignment.

[0176] Table 4. Ablation experiments on the impact of different inputs on the mask generator on the COCO 2017 dataset

[0177]

[0178] Example 2:

[0179] A second embodiment of the present invention provides a diffusion model optimization system based on U-Net parameter enhancement, including:

[0180] A data acquisition module is configured to acquire a data set consisting of image and text pairs;

[0181] A model building module is configured to build a diffusion model, in which a learnable mask is introduced for sampling weights from a U-Net structure;

[0182] a mask optimization module configured to obtain current computing resources and select a fine-tuning strategy based on the computing resources, wherein the fine-tuning strategy includes a training-based fine-tuning strategy and a reward model-based fine-tuning strategy;

[0183] The pre-training module is configured to pre-train the diffusion model using the data set. During the pre-training process, the learnable mask is optimized based on the selected fine-tuning strategy to complete the optimization of the diffusion model.

[0184] Example 3:

[0185] A third embodiment of the present invention provides a medium having a program stored thereon. When the program is executed by a processor, the steps of the diffusion model optimization method based on U-Net parameter enhancement as described in the first embodiment of the present invention are implemented.

[0186] Example 4:

[0187] Embodiment 4 of the present invention provides a device, including a memory, a processor, and a program stored in the memory and executable on the processor. When the processor executes the program, the steps in the diffusion model optimization method based on U-Net parameter enhancement as described in Embodiment 1 of the present invention are implemented.

[0188] Embodiment 5:

[0189] Embodiment 5 of the present invention provides a computer program product, including a computer program. When the computer program is executed by a processor, the computer program implements the steps in the diffusion model optimization method based on U-Net parameter enhancement as described in Embodiment 1 of the present invention.

[0190] The steps involved in the apparatuses of the above embodiments 2, 3, 4 and 5 correspond to those of the method embodiment 1. For the specific implementation methods, please refer to the relevant description part of the embodiment 1.

[0191] Those skilled in the art will appreciate that the modules or steps of the present invention described above can be implemented using a general-purpose computer device. Alternatively, they can be implemented using program code executable by a computing device, which can then be stored in a storage device and executed by the computing device. Alternatively, they can be fabricated into separate integrated circuit modules, or multiple modules or steps can be fabricated into a single integrated circuit module for implementation. The present invention is not limited to any specific combination of hardware and software.

[0192] Although the above describes the specific embodiments of the present invention in conjunction with the accompanying drawings, it is not intended to limit the scope of protection of the present invention. Those skilled in the art should understand that various modifications or variations that can be made by those skilled in the art on the basis of the technical solution of the present invention without any creative work are still within the scope of protection of the present invention.

Claims

1. A diffusion model optimization method based on U-Net parameter enhancement, characterized in that: The following steps are involved: Get a dataset of image and text pairs; Construct a diffusion model and introduce a learnable mask into the diffusion model to sample weights from the U-Net structure; Obtain the current computing resources and select a fine-tuning strategy based on the computing resources. The fine-tuning strategy includes a training-based fine-tuning strategy and a reward model-based fine-tuning strategy. The diffusion model is pre-trained using the dataset. During the pre-training process, the learnable mask is optimized based on the selected fine-tuning strategy to complete the optimization of the diffusion model.

2. The diffusion model optimization method based on U-Net parameter enhancement according to claim 1, characterized in that In the training-based fine-tuning strategy, the learnable masks are generated by the mask generator, and in the reward model-based fine-tuning strategy, the learnable masks are generated under the supervision of the reward model.

3. The diffusion model optimization method based on U-Net parameter enhancement according to claim 1, characterized in that The specific steps for selecting a fine-tuning strategy based on computing resources are: When computing resources are sufficient, choose a training-based fine-tuning strategy; When computing resources are insufficient, a fine-tuning strategy based on a reward model is chosen.

4. The diffusion model optimization method based on U-Net parameter enhancement according to claim 1, characterized in that The diffusion model is used to add noise to the dataset and gradually denoise it to perform data recovery to complete pre-training.

5. The diffusion model optimization method based on U-Net parameter enhancement according to claim 1, characterized in that The specific steps of the training-based fine-tuning strategy are: The mask generator generates time-step-dependent and sample-dependent masks as the first strategy mask; The first strategy mask is used to sample weights from the U-Net structure, and the sampling process is supervised by the diffusion loss.

6. The diffusion model optimization method based on U-Net parameter enhancement according to claim 1, characterized in that The specific steps of fine-tuning the strategy based on the reward model are: Generate a mask as the second policy mask under the supervision of the reward model; The second policy mask is used to sample weights from the U-Net structure, and the sampling process is supervised by the reward loss.

7. A diffusion model optimization system based on U-Net parameter enhancement, characterized in that: include: A data acquisition module is configured to acquire a data set consisting of image and text pairs; A model building module is configured to build a diffusion model, in which a learnable mask is introduced for sampling weights from a U-Net structure; a mask optimization module configured to obtain current computing resources and select a fine-tuning strategy based on the computing resources, wherein the fine-tuning strategy includes a training-based fine-tuning strategy and a reward model-based fine-tuning strategy; The pre-training module is configured to pre-train the diffusion model using the data set. During the pre-training process, the learnable mask is optimized based on the selected fine-tuning strategy to complete the optimization of the diffusion model.

8. A computer program product comprising a computer program, characterized in that When the computer program is executed by a processor, the diffusion model optimization method based on U-Net parameter enhancement according to any one of claims 1 to 6 is implemented.

9. A computer-readable storage medium, characterized in that Multiple instructions are stored therein, and the instructions are suitable for being loaded by a processor of a terminal device and executing the diffusion model optimization method based on U-Net parameter enhancement according to any one of claims 1 to 6.

10. A terminal device, characterized in that: The invention comprises a processor and a computer-readable storage medium, wherein the processor is used to implement each instruction; the computer-readable storage medium is used to store multiple instructions, wherein the instructions are suitable for being loaded by the processor and executing the diffusion model optimization method based on U-Net parameter enhancement according to any one of claims 1 to 6.

Images