A method and system for training an image generation model based on representations of injection and internalization

Abstract

The application discloses an image generation model training method and system based on representation injection and internalization. The method comprises the following steps: obtaining training samples including clean images, constructing a noisy input according to the clean images; extracting visual representation from the clean images by using a pre-trained visual encoder, and performing representation space adaptation to obtain target visual representation; in the first training stage, injecting the target visual representation into the network module of the generation network to replace the corresponding original intermediate representation to participate in the denoising calculation; in the second training stage, stopping the replacement and restoring the original calculation, and reducing the difference between the original intermediate representation and the corresponding target representation through internalization constraint; after the training is completed, removing the pre-trained visual encoder to obtain an image generation model including the generation network. The application can make the pre-trained visual representation directly participate in the denoising calculation, and internalize the gain to the model itself, thereby improving the training efficiency and generation quality, and reducing the additional dependence in the inference stage.

Description

Technical Field

[0001] This application relates to the fields of artificial intelligence and computer vision technology, and in particular to a training method and system for image generation models based on representation injection and internalization. Background Technology

[0002] In recent years, diffusion models have become an important technical approach in the field of image generation. Among them, the Diffusion Transformer (DiT), which uses a Transformer as the main denoising network, has good model scalability and generative performance. However, to obtain high-quality generated results, such models usually require a large number of training iterations and consume high computational resources, resulting in long training cycles and high training costs. Training efficiency has become one of the important factors restricting its further application. Especially in application scenarios such as medical image generation, where the cost of acquiring training data is high, model training resources are limited, or high quality requirements are required, improving the training efficiency of generative models and making full use of existing visual representations also has practical application value.

[0003] To improve the training efficiency of the Diffusion Transformer, existing techniques have proposed methods that utilize pre-trained visual encoders to assist model learning. For example, Representation Alignment (REPA) maps the intermediate hidden states of the Diffusion Transformer to the representation space of the pre-trained visual encoder through a trainable projector. By using representation similarity constraints, it aligns the intermediate features of the Diffusion Transformer with the features extracted from the image by the pre-trained visual encoder. This leverages the semantic representation capabilities of the pre-trained visual encoder to assist the Diffusion Transformer in learning semantically structured internal representations.

[0004] However, in the aforementioned representation alignment methods, the representations generated by the pre-trained visual encoder primarily serve as supervisory targets for the auxiliary loss, constraining the internal representations of the diffusion Transformer, and do not directly participate in its denoising calculation process. Therefore, the role of the pre-trained visual representations in the denoising process is mainly achieved through indirect alignment supervision, and the semantic structural information they contain is not directly used for the internal denoising calculation of the diffusion Transformer.

[0005] Therefore, how to make fuller use of the semantic representations provided by the pre-trained visual encoder, so that it can more directly affect the training and denoising process of the diffusion Transformer, and further improve the training efficiency and generation quality of the model, while ensuring the effective denoising process of the image generation model and normal use in the inference stage, remains a problem to be solved in the existing image generation model training technology. Summary of the Invention

[0006] Based on this, the purpose of this application is to address the above-mentioned technical problems by providing a training method and system for image generation models based on representation injection and internalization. This method allows the visual representations extracted by the pre-trained visual encoder to directly participate in the internal denoising calculation of the diffusion Transformer after representation space adaptation. Furthermore, the method internalizes the gains brought by external visual representations into the model's own representation capabilities through a representation internalization mechanism, thereby improving the training efficiency and generation quality of the image generation model.

[0007] To achieve the above-mentioned objectives, the first aspect of this application provides a method for training an image generation model based on representation injection and internalization, comprising: Obtain training samples, including clean images, for training an image generation model, and construct a noisy input for diffusion training based on the clean images; The clean image is input into a pre-trained visual encoder, and the visual representation of the clean image is extracted by the pre-trained visual encoder. The visual representation is then adapted to the representation space to obtain a target visual representation that matches the representation space of the native intermediate representation calculated by the generative network of the image generation model itself. In the first stage of training, the target visual representation is injected into the network module of the generator network, and the target visual representation replaces the corresponding native intermediate representation to participate in the denoising calculation. In the second training phase, the substitution is stopped, the computation of the native intermediate representation by the generative network is resumed, and an internalization constraint is applied to reduce the difference between the native intermediate representation and the target representation calculated by the generative network; wherein, the target representation is the representation obtained by extracting visual representations from the clean image and performing representation space adaptation in the second training phase. After training is completed, the pre-trained visual encoder is removed to obtain an image generation model that includes the trained generative network.

[0008] Preferably, constructing the noisy input for diffusion training based on the clean image includes: Obtain random noise and a time variable, and then perform a weighted combination of the clean image and the random noise to obtain a noisy input sample corresponding to the time variable:

[0009] in, For the clean image, To follow a standard normal distribution random noise, Let be the time variable, and , , ; The basic training objective of the generator network in performing denoising prediction is to minimize the basic denoising loss:

[0010] in, For the generated network, This represents all parameters of the generator network. This refers to the conditional information corresponding to the training samples. The prediction target is determined based on the clean image and the random noise, and is used to predict the target corresponding to the noisy input. It indicates a desire for the expected value.

[0011] Preferably, the generating network is a diffusion generating network based on the Transformer architecture, which includes several network modules, at least one of which is a self-attention module; The step of injecting the target visual representation into the network module of the generator network includes: an intermediate representation calculation process of injecting the target visual representation into the selected self-attention module; in the first training phase, the hidden states passed between the plurality of network modules are obtained by the generator network based on the noisy input.

[0012] Preferably, the visual representation includes key representation and value representation generated by the pre-trained visual encoder; The representation space adaptation is performed through a trainable mapping network, including: mapping the key representation to the key space of a selected self-attention module in the generator network using a trainable key projection head in the mapping network to obtain a mapped key representation; mapping the value representation to the value space of the selected self-attention module using a trainable value projection head in the mapping network to obtain a mapped value representation; the mapped key representation and the mapped value representation constitute the target visual representation; The step of injecting the target visual representation into the network module of the generator network and replacing the corresponding native intermediate representation with the target visual representation to participate in the denoising calculation includes: replacing the native key representation and native value representation obtained by the selected self-attention module with the mapped key representation and the mapped value representation respectively to participate in the attention calculation, and retaining the query representation, output projection and residual connection obtained by the native calculation of the noisy input related internal features.

[0013] Preferably, the key representation and the value representation are extracted from the last self-attention layer of the pre-trained visual encoder; the key projection head and the value projection head are respectively implemented using a multilayer perceptron.

[0014] Preferably, in the first phase of training, the attention output of the selected self-attention module is:

[0015] in, The query representation is obtained by natively calculating the relevant internal features of the noisy input. This is the key representation after the mapping. The mapped value is characterized as follows. The feature dimensions of the query representation and the mapped key representation are defined as follows: the output of the selected self-attention module is obtained by adding the attention output after the output projection to the noisy input-related internal features passed through the residual connection.

[0016] Preferably, in the second training stage, the clean image is continued to be input into the pre-trained visual encoder, and the mapped key representation and the mapped value representation, which serve as internalization targets, are obtained through the key projection head and the value projection head, respectively. The internalization constraint includes the internalization loss:

[0017] in, and These are the native key representation and native value representation, respectively, obtained by the selected self-attention module during the second stage of training. and These are the mapped key representation and the mapped value representation, respectively. This indicates that the gradient operation is stopped. Indicates the number of elements in the corresponding representation, subscript Denotes the Frobenius norm; In the second training phase, the internalization loss and the basic denoising loss are weighted and summed to obtain the total loss, and the generator network is then jointly optimized.

[0018] in, These are the weighting coefficients used to balance the basic denoising loss and the internalization loss.

[0019] Preferably, the first training stage is a training stage from the start of training until the preset training conditions are met. The preset training conditions include the number of training iterations reaching a preset number of steps threshold, and / or the proportion of the number of training iterations to the preset total number of training steps reaching a preset proportion threshold. After the preset training conditions are met, the replacement of the native intermediate representation by the target visual representation is stopped, and the training directly enters the second training stage.

[0020] Preferably, the method further includes: A representation alignment constraint is applied to the intermediate hidden states of the generator network, the representation alignment constraint being used to align the intermediate hidden states of the generator network, mapped by a trainable alignment projector, with the representations extracted from the clean image by the pre-trained visual encoder in the representation space of the pre-trained visual encoder. The network module into which the target visual representation is injected is located before the network layer where the representation alignment constraint is applied, along the forward propagation direction of the generated network.

[0021] To achieve the purpose of the invention, a second aspect of this application provides an image generation model training system based on representation injection and internalization, applying the above-described technical solution for training an image generation model based on representation injection and internalization. The system includes: The data acquisition module is used to acquire training samples, including clean images, for training an image generation model, and to construct a noisy input for diffusion training based on the clean images. A visual representation extraction module is used to input the clean image into a pre-trained visual encoder and extract the visual representation of the clean image through the pre-trained visual encoder. The representation mapping module is used to adapt the visual representation to the representation space to obtain a target visual representation that matches the representation space of the native intermediate representation inside the network module of the generating network. The representation injection module is used to inject the target visual representation into the network module of the generator network during the first stage of training, so that the target visual representation replaces the corresponding native intermediate representation in the denoising calculation. The representation internalization module is used in the second training phase to stop the substitution of the target visual representation for the native intermediate representation, restore the computation of the native intermediate representation by the generative network, and apply internalization constraints to reduce the difference between the native intermediate representation and the target representation calculated by the generative network. The target representation is the representation obtained by extracting visual representations from the clean image and performing representation space adaptation in the second training phase. The model output module is used to remove the pre-trained visual encoder after training is completed and output an image generation model including the generative network.

[0022] Compared with the prior art, the beneficial effects of this invention are: This invention injects visual representations extracted by a pre-trained visual encoder into the network module of the generator network after the representation mapping module adapts the representation space. This replaces the corresponding native intermediate representations in the denoising calculation, allowing the semantic structure information in the pre-trained visual representations to directly affect the denoising process. This helps to accelerate the formation of semantically structured internal representations, thereby improving the training efficiency and convergence performance of the image generation model, and enhancing the semantic consistency and visual detail of the generated images. The target visual representations injected in this invention only serve as temporary representation scaffolds in the first stage of training. Their gains are internalized into the generator network itself through the internalization constraints in the second stage of training. After training is completed, the pre-trained visual encoder and representation mapping module are removed. The resulting image generation model relies solely on its own network to generate images during the inference stage, without adding additional inference computation overhead. Attached Figure Description

[0023] Figure 1 This is a flowchart illustrating an image generation model training method based on representation injection and internalization in one embodiment. Figure 2 A schematic diagram of the structure of the Visual Representation Injection Diffusion Transformer Attention Module in one embodiment; Figure 3 This is a comparison chart of training convergence performance for different model sizes in one embodiment. Figure 4 This is a comparison chart of the image generation effects of different training methods in one embodiment. Detailed Implementation

[0024] To make the objectives, technical solutions, and advantages of this application clearer, the following detailed description is provided in conjunction with the accompanying drawings and embodiments. It should be understood that the specific embodiments described herein are merely illustrative and not intended to limit the scope of the invention. The following embodiments are used to illustrate the invention but are not intended to limit its scope.

[0025] Furthermore, the terms "first" and "second" are used for descriptive purposes only and should not be construed as indicating or implying relative importance or implicitly specifying the number of technical features indicated. Thus, a feature defined as "first" or "second" may explicitly or implicitly include one or more of that feature. In the description of this application, "multiple" means two or more, unless otherwise explicitly specified.

[0026] Example 1 Embodiment 1 of this application provides a training method for an image generation model based on representation injection and internalization, such as... Figure 1 As shown, it includes: S1. Obtain training samples including clean images for training the image generation model, and construct noisy input for diffusion training based on the clean images; S2. Input the clean image into the pre-trained visual encoder, and extract the visual representation of the clean image through the pre-trained visual encoder; perform representation space adaptation on the visual representation to obtain a target visual representation that matches the representation space of the native intermediate representation calculated by the generative network of the image generation model itself. S3. In the first training stage, the target visual representation is injected into the network module of the generator network, and the target visual representation replaces the corresponding native intermediate representation to participate in the denoising calculation. S4. In the second training phase, the substitution is stopped, the computation of the native intermediate representation by the generative network is resumed, and an internalization constraint is applied to reduce the difference between the native intermediate representation and the target representation calculated by the generative network; wherein, the target representation is the representation obtained by extracting visual representations from the clean image and performing representation space adaptation in the second training phase. S5. After training is completed, remove the pre-trained visual encoder to obtain an image generation model that includes the trained generative network.

[0027] This embodiment uses a diffusion image generation model based on the Transformer architecture as an example to further explain the above steps S1 to S5.

[0028] The training method in this embodiment includes processes such as training data construction S1, pre-trained visual representation extraction and mapping S2, visual representation injection S3, representation scaffold removal and representation internalization S4, and training completion model output S5.

[0029] S1. Obtain training data and construct the diffusion model training input: Obtain clean images for training the image generation model and corresponding condition information ,in , Represents the probability distribution and conditional information of real data. This can be either a category label or text. Taking flow-matching-based diffusion training as an example, time variables are randomly sampled during training. and random noise And construct time based on clean image and random noise. Corresponding noisy samples :

[0030] Among them, ,and The generation network employs a diffused Transformer network. Used to predict conditional velocity targets Its basic denoising training objective is to minimize the denoising loss. :

[0031] Step S1 is used to establish the basic denoising learning task of the image generation model and to provide a basic training objective for the subsequent injection and internalization of pre-trained visual representations.

[0032] S2. Extract visual representations and perform representation mapping using a pre-trained visual encoder: The pre-trained visual encoder is a visual feature extraction network pre-trained on large-scale image data, used to extract visual representations containing semantic structure information from input images, such as self-supervised visual encoders like DINOv2; preferably, the network parameters of the pre-trained visual encoder are frozen during training and do not participate in parameter updates, so as to reduce training overhead and ensure the stability of the extracted visual representations.

[0033] Clean images from the training samples Input pre-trained visual encoder Through the pre-trained visual encoder Extracting visual representations containing semantic structural information. In a preferred implementation of this embodiment, the key representations generated by the last self-attention layer of the pre-trained visual encoder are extracted. Sum value representation .

[0034] A parameter-trainable neural network is used as a mapping network to transform the visual representation output by the pre-trained visual encoder to a representation space that matches the internal representation space of the generator network. In this embodiment, the mapping network includes a trainable key projection head. Sum value projection head For the selected self-attention module in the diffusion Transformer, the key projection head is set. Will Mapped to the native key space corresponding to the self-attention module, and through the value projection head Will Mapped to the native value space corresponding to this self-attention module:

[0035] in, and These represent the mapped key and value representations, respectively. The key projection head... Sum value projection head This can be implemented using a multilayer perceptron (MLP). Through the above mapping, the key / value representations generated by the pre-trained visual encoder are adapted to the native key / value representation space of the self-attention module selected by the diffusion Transformer, providing matching key and value representations for subsequent representation injection.

[0036] It should be noted that this embodiment divides the training process into two phases: the first training phase starts from the beginning of training and continues until the preset training conditions are met. In this phase, the target visual representation replaces the corresponding native intermediate representation in the denoising calculation; the second training phase starts from the meeting of the preset training conditions and continues until the training is completed. In this phase, the replacement is stopped and internalization constraints are applied.

[0037] S3. Inject the mapped visual representation into the internal denoising calculation process of the diffusion Transformer.

[0038] like Figure 2 As shown, this embodiment represents the key characterization obtained through mapping. Sum value representation The attention calculation process selected by the attention module in the injection diffusion Transformer. Specifically, in the first phase of training (i.e., the early stage of training), the attention calculation is performed using the aforementioned... and Replace the native key representation and native value representation obtained by the self-attention module, respectively, while retaining the query representation obtained by the native calculation of the noisy input-related internal features. Output projection and residual connections. The attention output after injection can be expressed as:

[0039] in, These are the feature dimensions for query representation and key representation. The injection method described above only changes the source of the key and value representations selected from the attention module, without altering the attention calculation method. This is due to the native query representation... The output projection and residual connections remain unchanged, and information related to noisy inputs can still be represented by native queries. It participates in attention calculation and propagates to subsequent network layers via the residual path. Meanwhile, and The semantic structure information derived from the visual representation of a clean image is introduced into the attention computation in the form of key / value representation, thus directly participating in the denoising process of the diffusion Transformer. Figure 2 The last step before the output of the self-attention module is a residual addition operation, which adds the attention result after the output projection to the noisy input related features passed through the residual path to obtain the output of the self-attention module.

[0040] This embodiment preferably employs key / value representation joint injection. In preliminary diagnostic experiments, different injection locations were also examined, including attention output, all Q / K / V, K / V, K only, and V only. All of these injection methods within the attention module can maintain the information flow path from noisy input to module output without completely overwriting the interlayer hidden states. Within the complete representation injection and internalization training framework, K / V joint injection is the preferred implementation.

[0041] S4. Remove the temporary representation scaffold and perform representation internalization training.

[0042] In the first training phase, the mapped pre-trained visual representations serve as temporary representation scaffolds in the denoising calculation of the diffusion Transformer, enabling the model to first learn to denoise using the semantically structured key / value representations. When training reaches a preset training condition (e.g., the number of training iterations reaches a preset threshold), the temporary representation scaffolds are removed, allowing the selected self-attention module to resume the calculation of its own native key and value representations, and then proceed to the second training phase, the representation internalization phase.

[0043] During the representation internalization stage, the mapped key and value representations no longer replace the original key-value representations in denoising calculations. Instead, they serve as internalization targets to constrain the original representations generated by the model itself. Specifically, in the second training stage, clean images are continued to be processed... Input the pre-trained visual encoder And via the key projection head Sum value projection head The mapping post-key representations as internalization targets are obtained respectively. and the value representation after mapping To further internalize the effective information provided by the representation scaffold into the model itself, an internalization loss is introduced. The native K / V representations calculated by the constraint diffusion Transformer itself approximate the corresponding mapped key representations. and the value representation after mapping :

[0044] in, This indicates that gradient operations are stopped, meaning that the gradient of the internalized target is not propagated back to the pre-trained visual encoder and projector. Indicates the number of elements in the corresponding representation; subscript Denotes the Frobenius norm; and These represent the native key representation and native value representation obtained from the native computation of the attention module selected in the second stage of training, respectively.

[0045] During the representation internalization stage, the internalization loss and the basic diffusion training loss are jointly optimized, and the total loss can be expressed as:

[0046] in, The weighting coefficients are used to balance the denoising training objective and the representation internalization objective. It is robust, and the value can range from 1 to 4, but 2 is commonly used. Alternatively, the value can be determined through conventional experiments based on the magnitudes of the base denoising loss and internalization loss. The value of is determined by the internalization loss. The internalization loss also ensures compatibility between the original K / V representation after scaffold removal and the mapped K / V representation corresponding to the scaffold stage, thus facilitating a smooth transition from external representation-assisted computation to native model computation. This embodiment also examined a linear scheduling transition method, but no further performance improvement was observed, indicating that the internalization loss alone is sufficient to maintain a sufficiently smooth transition. Therefore, this embodiment can adopt the method of directly removing the representation scaffold and restoring native computation.

[0047] S5. Complete the training and output the image generation model.

[0048] After completing the above training, the pre-trained visual encoder and the mapping network (including the key projection head and value projection head) used in the training phase are removed, while the trained diffusion Transformer backbone network is retained. At this point, the image generation model reverts to its native inference structure, requiring no external visual encoder. During the inference phase, no clean image is needed as input for extracting the representation scaffold, nor is an additional pre-trained visual encoder required; the image generation process relies solely on the trained diffusion backbone network. Thus, the gains from external visual representations during the training phase are retained and internalized into the generative model's own representational capabilities, while avoiding continuous reliance on an external pre-trained visual encoder during the inference phase.

[0049] In summary, this embodiment forms a complete training chain of "training image input → visual representation extraction → representation mapping → visual representation injection → diffusion denoising training → representation scaffold removal and representation internalization training → training completed and model output".

[0050] Example 2 Embodiment 2 of this application, based on Embodiment 1, provides a preferred implementation of joint training of representation injection and representation alignment.

[0051] Building upon Example 1, the representation injection and internalization method of this application can also be used in conjunction with the representation alignment method REPA. REPA primarily projects the intermediate hidden states of the diffusion Transformer onto the representation space of the pre-trained visual encoder and supervises this process using representation alignment loss. In contrast, the representation injection (REPI) method described in this application projects the key / value representations generated by the pre-trained visual encoder into the denoising computation process within the diffusion Transformer and internalizes these representations during subsequent training. Since the two methods differ in their representation transfer direction, supervision methods, and mechanisms of action, they can be used in combination.

[0052] During joint training, it is preferable to place the REPI representation injection layer before the REPA representation alignment layer. With this setting, the alignment loss of REPA can continue to supervise the intermediate network layers after the REPI injection layer, keeping their output aligned with the encoder representation. After the representation scaffold is removed, the REPA alignment constraint located after the REPI injection layer can also indirectly constrain the native K / V representation of the injection layer to remain consistent with the corresponding state before the scaffold removal, thus forming a complementary constraint with the internalization loss.

[0053] Example 3 This application, in its third embodiment, provides an image generation model training system based on representation injection and internalization, based on embodiments 1 and 2, for implementing the training method described in embodiment 1 or 2. The system includes a data acquisition module, a visual representation extraction module, a representation mapping module, a representation injection module, a representation internalization module, and a model output module.

[0054] The data acquisition module is used to acquire training images and corresponding training conditions, and to construct noisy samples required for diffusion training. The visual representation extraction module is used to extract visual representations from clean training images using a pre-trained visual encoder. The representation mapping module is used to map the visual representation to the internal representation space corresponding to the selected self-attention module in the diffusion Transformer; preferably, the encoder key representation and value representation are mapped to the native key space and value space of the selected self-attention module, respectively. The representation injection module is used to inject the mapped visual representation into the internal denoising calculation process of the diffusion Transformer during the first stage of training, so that it can participate in the attention calculation of the selected self-attention module and preserve the original information propagation path generated by the noisy input. The representation internalization module is used in the second phase of training after the temporary representation scaffold is removed. It enables the diffusion Transformer to generate native internal representations that are compatible with the corresponding mapped visual representations through internalization constraints. The model output module is used to output the trained diffusion Transformer image generation model after training is completed. The model does not need to rely on the pre-trained visual encoder and the mapping network used in the training phase during the inference phase.

[0055] The modules described above work together according to the training data stream to achieve the training process described in Example 1 or Example 2.

[0056] Example 4 To verify the effectiveness of the training method described in this application, experiments were conducted on the ImageNet 256×256 image generation task with category conditions. The FID evaluation experiment was performed without classifier-free guidance (CFG). A lower FID value indicates that the generated image distribution is closer to the real image distribution, i.e., the better the generation performance.

[0057] As shown in Table 1, the REPI method described in this application exhibits good training performance on diffusion Transformer models of different scales. On the SiT-B / 2 and SiT-L / 2 models, the FID obtained by REPI is lower than that of REPA under the same number of training iterations. On the SiT-XL / 2 model, the FID of REPI after 400K training steps is 7.41, lower than REPA's 7.90 under the same number of training steps, indicating that REPI can achieve better generative performance under the same number of training iterations.

[0058] Furthermore, combining REPA and REPI further improves model performance. On the SiT-XL / 2 model, the FID reaches 6.38 after 400K training steps with REPA+REPI; when the training steps are only 160K, the FID reaches 8.22, which is better than the FID of 8.30 obtained by the original SiT after 7M training steps. Under the above experimental conditions, to achieve FID performance equivalent to or better than the original SiT after 7M training steps, the number of training iterations required by the REPA+REPI joint scheme is about 1 / 43.75 of the original scheme. This shows that this joint training method can significantly reduce the number of training iterations required to achieve comparable generative performance, thereby significantly saving the computational resources required for training.

[0059] like Figure 3 As shown, the convergence curves under different model sizes further demonstrate that, with the same number of training iterations, REPI achieves a lower FID than REPA overall and exhibits a better convergence trend; the combined use of REPI and REPA achieves further improvement, indicating that the two mechanisms of representation injection and representation alignment have complementary effects.

[0060] like Figure 4 As shown, the top row represents samples generated by the iREPA method at different training stages, and the bottom row represents samples generated by the REPI method described in this application at different training stages, from left to right: the results generated after 100K steps, 200K steps, and 400K steps of training. All samples use the same initial noise, the same sampler, and the same number of sampling steps, and none are guided by a classifier-less approach. It can be observed that the method described in this application exhibits better semantic consistency and finer visual detail compared to iREPA, indicating that the representation injection and internalization mechanism described in this application helps improve the quality of the generated images.

[0061] In summary, the experimental results, from multiple aspects such as FID evaluation metrics, model convergence trend, and qualitative quality of generated images, demonstrate that the representation injection and internalization training method described in this application can achieve superior image generation performance while reducing the number of training iterations required to achieve comparable generation performance; furthermore, when used in conjunction with REPA, it can achieve further performance improvements.

[0062] Table 1. FID results for different models and training methods (ImageNet 256×256, without CFG)

[0063] In summary, this application provides a training method and system for image generation models based on representation injection and internalization. This method enables the semantically structured visual representations provided by the pre-trained visual encoder to directly participate in the internal denoising calculation of the diffusion Transformer. Furthermore, through the representation internalization mechanism, the gains brought by external visual representations are internalized into the model's own representation capabilities, thereby improving the training efficiency and generation quality of the image generation model and reducing the dependence on additional models during the inference stage.

[0064] The above description is only a preferred embodiment of the present invention, but the scope of protection of the present invention is not limited thereto. Any changes or substitutions that can be easily conceived by those skilled in the art within the scope of the technology disclosed in the present invention should be included within the scope of protection of the present invention.

Claims

1. A training method for an image generation model based on representation injection and internalization, characterized in that, include: Obtain training samples, including clean images, for training an image generation model, and construct a noisy input for diffusion training based on the clean images; The clean image is input into a pre-trained visual encoder, and the visual representation of the clean image is extracted by the pre-trained visual encoder. The visual representation is adapted to the representation space to obtain a target visual representation that matches the representation space of the native intermediate representation calculated by the generative network of the image generation model itself. In the first stage of training, the target visual representation is injected into the network module of the generator network, and the target visual representation replaces the corresponding native intermediate representation to participate in the denoising calculation. In the second training phase, the substitution is stopped, the computation of the native intermediate representation by the generative network is resumed, and an internalization constraint is applied to reduce the difference between the native intermediate representation and the target representation calculated by the generative network; wherein, the target representation is the representation obtained by extracting visual representations from the clean image and performing representation space adaptation in the second training phase. After training is completed, the pre-trained visual encoder is removed to obtain an image generation model that includes the trained generative network.

2. The image generation model training method based on representation injection and internalization according to claim 1, characterized in that, The construction of the noisy input for diffusion training based on the clean image includes: Obtain random noise and a time variable, and then perform a weighted combination of the clean image and the random noise to obtain a noisy input sample corresponding to the time variable: in, For the clean image, To follow a standard normal distribution random noise, Let be the time variable, and , , ; The basic training objective of the generator network in performing denoising prediction is to minimize the basic denoising loss: in, For the generated network, This represents all parameters of the generator network. This refers to the conditional information corresponding to the training samples. The prediction target is determined based on the clean image and the random noise, and is used to predict the target corresponding to the noisy input. It indicates a desire for the expected value.

3. The image generation model training method based on representation injection and internalization according to claim 1, characterized in that, The generating network is a diffusion generating network based on the Transformer architecture, which includes several network modules, at least one of which is a self-attention module. The step of injecting the target visual representation into the network module of the generator network includes: an intermediate representation calculation process of injecting the target visual representation into the selected self-attention module; In the first stage of training, the hidden states passed between the network modules are obtained by the generator network based on the noisy input.

4. The image generation model training method based on representation injection and internalization according to claim 3, characterized in that, The visual representation includes the key representation and value representation generated by the pre-trained visual encoder; The representation space adaptation is performed through a trainable mapping network, including: mapping the key representation to the key space of a selected self-attention module in the generator network using a trainable key projection head in the mapping network to obtain a mapped key representation; mapping the value representation to the value space of the selected self-attention module using a trainable value projection head in the mapping network to obtain a mapped value representation; the mapped key representation and the mapped value representation constitute the target visual representation; The step of injecting the target visual representation into the network module of the generator network and replacing the corresponding native intermediate representation with the target visual representation to participate in the denoising calculation includes: replacing the native key representation and native value representation obtained by the selected self-attention module with the mapped key representation and the mapped value representation respectively to participate in the attention calculation, and retaining the query representation, output projection and residual connection obtained by the native calculation of the noisy input related internal features.

5. The image generation model training method based on representation injection and internalization according to claim 4, characterized in that, The key representation and the value representation are extracted from the last self-attention layer of the pre-trained visual encoder; the key projection head and the value projection head are respectively implemented using a multilayer perceptron.

6. The image generation model training method based on representation injection and internalization according to claim 4, characterized in that, In the first phase of training, the attention output of the selected self-attention module is: in, The query representation is obtained by natively calculating the relevant internal features of the noisy input. This is the key representation after the mapping. The mapped value is characterized as follows. The feature dimensions of the query representation and the mapped key representation are defined as follows: the output of the selected self-attention module is obtained by adding the attention output after the output projection to the noisy input-related internal features passed through the residual connection.

7. The image generation model training method based on representation injection and internalization according to claim 4, characterized in that, In the second training phase, the clean image is input into the pre-trained visual encoder, and the mapped key representation and the mapped value representation, which serve as internalization targets, are obtained through the key projection head and the value projection head, respectively. The internalization constraint includes the internalization loss: in, and These are the native key representation and native value representation, respectively, obtained by the selected self-attention module during the second stage of training. and These are the mapped key representation and the mapped value representation, respectively. This indicates that the gradient operation is stopped. Indicates the number of elements in the corresponding representation, subscript Denotes the Frobenius norm; In the second training phase, the internalization loss and the basic denoising loss are weighted and summed to obtain the total loss, and the generator network is then jointly optimized. in, These are the weighting coefficients used to balance the basic denoising loss and the internalization loss.

8. The image generation model training method based on representation injection and internalization according to claim 1, characterized in that, The first training phase is the training phase from the start of training until the preset training conditions are met. The preset training conditions include reaching a preset number of training iterations. After the preset training conditions are met, the replacement of the native intermediate representation by the target visual representation is stopped, and the training phase directly begins.

9. The image generation model training method based on representation injection and internalization according to any one of claims 1 to 8, characterized in that, The method further includes: A representation alignment constraint is applied to the intermediate hidden states of the generator network, the representation alignment constraint being used to align the intermediate hidden states of the generator network, mapped by a trainable alignment projector, with the representations extracted from the clean image by the pre-trained visual encoder in the representation space of the pre-trained visual encoder. The network module for injecting the target visual representation is located before the network layer where the representation alignment constraint is applied, along the forward propagation direction of the generated network.

10. A training system for an image generation model based on representation injection and internalization, employing the image generation model training method based on representation injection and internalization as described in any one of claims 1 to 9, characterized in that, The system includes: The data acquisition module is used to acquire training samples, including clean images, for training an image generation model, and to construct a noisy input for diffusion training based on the clean images. A visual representation extraction module is used to input the clean image into a pre-trained visual encoder and extract the visual representation of the clean image through the pre-trained visual encoder. The representation mapping module is used to adapt the visual representation to the representation space to obtain a target visual representation that matches the representation space of the native intermediate representation inside the network module of the generating network. The representation injection module is used to inject the target visual representation into the network module of the generator network during the first stage of training, so that the target visual representation replaces the corresponding native intermediate representation in the denoising calculation. The representation internalization module is used in the second training phase to stop the substitution of the target visual representation for the native intermediate representation, restore the computation of the native intermediate representation by the generative network, and apply internalization constraints to reduce the difference between the native intermediate representation and the target representation calculated by the generative network. The target representation is the representation obtained by extracting visual representations from the clean image and performing representation space adaptation in the second training phase. The model output module is used to remove the pre-trained visual encoder after training is completed and output an image generation model including the generative network.

Images