Free Lunch of Image-mask Alignment for Anomaly Image Generation and Segmentation

1 Introduction

The objective of anomaly image generation and segmentation is to produce anomalous images, such as industrial product defects or diseased areas in medical imagery, along with their corresponding mask labels in a single process. This task addresses the scarcity of real-world anomaly samples and the associated privacy issues. It is a crucial task in computer vision, characterized by the simultaneous creation of images and masks, which distinguishes it from traditional image generation tasks. By creating these synthetic samples, we can

mance and generalization ability of the downstream segmentation models to better handle various anomaly situations.

In previous research on anomaly image generation and segmentation, prevalent approaches tend to employed a pipeline where the guidance mask was utilized as a condition for the generative model, specifying the location of the anomaly within the synthetic image. Once the synthetic image was produced, the mask was treated as the segmentation label. However, misalignment happens when anomaly regions in the generated images fail to accurately align with the conditional mask. This phenomenon, known as the mask label drift issue, negatively impacts the performance of the segmentation models by introducing false training signals.

* Corresponding author (mjsun@suda.edu.cn).

This drift issue may stem from the generative model's incomplete understanding of the image-mask correspondence during training. To enhance this understanding, we propose a dual-branch training approach for the generative model. Specifically, different from previous methods that focused solely on the generation of anomaly images (Figure. 1(a)), we introduce an innovative step where the binary mask is also been generated during the training of the shared generative model (Figure. 1(b)). This dual-branch framework allows us to assess the discrepancy between the noise predicted for the image and that for the mask by the shared generative model. This discrepancy is termed as the Alignment Regularization loss to improve the generative model's image-mask alignment during training. It is worth noting that only the image generation branch is employed during inference to generate samples to train downstream segmentation models.

We attribute such improvement to the fact that denoising both binary mask and image creates an explicit link between their features, mitigating their distribution gap and establishing a robust internal representation of their relationship at the pixel level. Consequently, the contour characteristics of anomalies can be more rapidly learned through the mask generation branch. These contour characteristics are then transferred to the image generation branch via the Alignment Regularization loss during training, ultimately enhancing the image-mask corresponding of the generative model.

Then, we further delve deeper into the "free lunch" of the aforementioned image-mask correspondence within the welltrained generative model: when the condition mask closely matches the anomaly area within the input noisy image, the generative model can easily reconstruct the image from its noisy counterpart, resulting in a smaller standard denoise loss; conversely, if there is a mismatch between the condition mask and the anomaly area in the input noisy image, the corresponding denoise loss increases substantially as well.

Building on this observation, we propose to integrate the generative model into the training of the segmentation model, to complement the traditional training strategy (Figure. 1(c)). Here is how the process unfolds for a given image: first, this image is fed into the segmentation model. The segmentation model then predicts a logit mask, which serves as the condition for the generative model. Subsequently, we introduce noise to the image, treating the noisy image as the input for the generative model, as seen in Figure. 1(d). The standard denoising loss of the generative model is harnessed to finetune the segmentation model to improve its performance.

The main contributions of this paper are as follows:

• We introduce a dual-branch training approach that enables the simultaneous generation of anomaly images and their masks. In Addition, we incorporate an Alignment Regularization loss between the generated images and masks to enhance their correspondence. Only the image-generation branch is activated when synthesizing samples, with the image-mask drift issue alleviated.
• We propose to integrate the well-trained generative model into the training of downstream segmentation models via Generative Feedback loss. This loss, derived from the original denoising loss of the generative model, can serve as an indicator of the segmentation model's accuracy in identifying anomaly areas for its finetuning.
• Experiments show our method outperforms previous approaches in IoU metrics by 5.03% on the Real-IAD dataset (industrial), 5.68% on the polyp dataset (medical), and 16.63% on the Floor Dirty dataset (indoor).

2 RelatedWork

2.1 Anomaly Image Generation

Since anomaly data is extremely scarce in the real world, researchers want to expand the anomaly dataset in various ways. Previous methods [Zavrtanik et al., 2021; Lin et al., 2021] transferred the existing anomalies or textures to normal images, but they lacked realism and diversity. Then researchers used the GAN-based [Goodfellow et al., 2020] model SDGAN [Niu et al., 2020] and Defect-GAN [Zhang et al., 2021] to generate anomalies on normal samples, but the shortage of anomaly samples limits training. Diffusion models [Ho et al., 2020; Yao et al., 2023; Wu et al., 2025; Qiu et al., 2025] have gradually gained widespread application due to their high-quality generation performance. Recently, AnomalyDiffusion [Hu et al., 2024] has made progress in few-shot anomaly generation, while requiring separate models to generate different types of anomalies.

2.2 Anomaly Image Segmentation

Anomaly Image Segmentation is a crucial computer vision task of localization and segmentation [Sun et al., 2020b; Sun et al., 2021b; Sun et al., 2021a; Sun et al., 2020a; Sun et al., 2024] that aims to identify and segment the abnormal regions within an image. Previous researchers often adopt reconstruction-based methods [Schlegl et al., 2019; Cao et al., 2023], which discover defects by analyzing the residual before and after reconstructing the images. Embedding-based methods [Lee et al., 2022; Cao et al., 2022] can be used, where pre-trained encoder networks are used to extract features, followed by using clustering methods to detect defects.

3 Method

3.1 Overview

The core of this work is the image-mask Alignment Regularization strategy to address the drift issue in the Stable Diffusion (SD) [Rombach et al., 2022] model. The well-learned alignment is then used to assist the segmentation training by high-quality training data augmentation and a novel Generative Feedback mechanism.

Figure 2 shows the framework of our method. For generative training, we finetune an SD model for mask-conditioned image generation as shown in Figure 2 (a). Here, our proposed Alignment Regularization is applied to enforce the model to learn precise image-mask alignment at the pixel level. For segmentation training, we first use the well-tuned SD model to generate highly aligned training samples to augment the training data for segmentation models, as shown in Figure 2 (b) and (c). Then, we transfer the image-mask correspondence from SD model into a Generative Feedback supervision. This feedback mechanism, based on the SD model's noise prediction objective, guides the segmentation model to produce masks that accurately correspond to the image.

3.2 Generative Model Training

Standard Finetuning for Conditional Image Generation

Our primary objective is to generate high-quality imagemask pairs to augment the training data for our segmentation model. To achieve this, we first finetune an SD model on the training dataset for conditional image generation. The SD model comprises four main components: a Variational Autoencoder (VAE) [Kingma and Welling, 2013], a CLIP Text Encoder [Radford et al., 2021], a Scheduler, and a UNet [Ronneberger et al., 2015]. The model is conditioned on both a binary mask m, which specifies the background and the foreground anomalous region, as well as a textual description p, which provides additional contextual details for the content to be generated. Ultimately, m is leveraged as the ground truth label for the generated anomalous image.

During training, the process begins with the VAE encoder E compressing the input image x0 into a latent feature z0. Concurrently, the CLIP Text Encoder τ converts the tokenized prompt p into text embedding vectors τ(p). In the forward diffusion process, the Scheduler progressively adds Gaussian noise ϵ to the latent feature z0 over t timesteps, resulting in zt, which approaches pure Gaussian noise as t increases. During the reverse (denoising) process, the U-Net is trained to predict the added noise ϵ based on the noised latent feature zt, the timestep t, and the text embedding vector τ(p). The finetuning loss Lsd is calculated as the Mean Squared Error (MSE) between the actual noise ϵ and the predicted noise ϵθ, as defined in Equation 1, as follows:

After finetuning, the SD model becomes capable of generating anomalous images from randomly sampled noise, conditioned on a given binary mask. However, a limitation arises due to the model's coarse understanding of the mask, which results in a label drift issue. Specifically, the generated imagemask pairs exhibit misalignment between the binary mask and the corresponding foreground and background regions in the image. The misaligned training pairs can introduce inconsistencies in the training data, potentially producing false supervision signals during subsequent segmentation training.

Enhancing Image-Mask Alignment with Regularization

To address the drift issue observed in conventional finetuning, we propose a novel strategy termed Alignment Regularization to enhance correspondence between generated images and their conditioning masks. Our hypothesis identifies feature misalignment between masks and images as the root cause of this drift phenomenon. Due to their inherent visual differences, a distribution gap can emerge between their encoded features, hindering the model's ability to learn effective mask-to-image correlations and ultimately leading to undesired foreground drift in the generated samples.

To overcome this challenge, we introduce an Alignment Regularization mechanism that treats masks not only as conditions but also as denoising targets. Specifically, we apply identical noise to both images and masks, then enforce consistency between their respective noise predictions through our regularization term. As illustrated in Figure 3, this approach enables the model to adapt to both image and mask features simultaneously, thereby bridging their distribution gap. Furthermore, this mechanism establishes explicit links between the representations of the image and mask, strengthening the mask's effectiveness as a conditioning signal.

We introduce a dual-head training strategy that simultaneously processes the image and mask. Specifically, during the diffusion process, the framework employs the VAE encoder E to compress the mask m into a latent feature h0. The same noise ϵ, added to the latent image, is applied to the latent feature h0 over t timesteps to produce noisy latent feature ht:

In the denoising stage of the generative process, the model conditions on both the mask m and the text embedding τ(p). Using the noise-added feature ht as input to the U-Net model, the framework predicts the noise ϵh for the mask region as:

To enforce structural alignment, the Alignment Regularization term Lal computes the mean squared error (MSE) between the noise ϵz predicted by the input image head and the noise ϵh predicted by the mask head, as shown below:

The overall training objective for generation Lgen combines the conventional finetuning loss Lsd and the proposed structural alignment loss Lal, then weighted by a factor α:

3.3 SegmentationModel Training

Segmentation Training on Augmented Dataset

Building upon our novel finetuning strategy, the SD model is now capable of generating high-quality image-mask pairs. We first leverage this capability to augment our training dataset. Specifically, we generate synthetic training images conditioned on binary masks and combine these with real data pairs. The segmentation model is then trained on this combined dataset. The training process optimizes a crossentropy loss, Lce, which minimizes the discrepancy between the model-predicted logits and the ground truth masks.

Given real images xreal and synthetic images xsyn, the segmentation model f produces corresponding logits lreal and lsyn, respectively. The loss is then computed as:

where y represents the binary masks, and H denotes the cross-entropy calculation. By training on these high-quality and diverse augmented data, we enable the segmentation model to develop more robust representations compared to training solely on real data, yielding better performance.

Improving Segmentation with Generative Feedback

To further leverage the image-mask alignment learned in the SD model, we introduce a novel Generative Feedback loss, Lfb, which utilizes the well-trained SD model to evaluate the output of the segmentation model, providing guidance and iterative refinement. During training, the segmentation model f initially predicts logits l for a given input image x. These logits are then processed through a conversion module T to obtain an approximate mask ˜m, which can be defined as:

The logits l have dimensions of N × C × H × W, where the channel dimension C = 1 corresponds to the prediction scores for the background and foreground. They are then converted by module T defined as in following equations:

It is also illustrated in Figure 4. First, the logits are rescaled by a temperature parameter λ ∈ (0, 1], followed by a softmax normalization along the channel dimension. The temperature parameter λ sharpens the softmax output, emphasizing the foreground scores. As shown in Equation 9, we extract the foreground predicted probabilities by slicing the second channel of the predictions to approximate a binary mask, also shown by the yellow probability map in Figure 4.

The approximate mask ˜m is then used as a condition for the SD model, along with a text embedding vector τ(p). The SD model then processes the noise-added latent image feature zt to generate predicted noise ϵfb. Our feedback loss, Lfb, is then defined as the Mean Squared Error (MSE) between the real noise ϵ added during the forward diffusion process and the noise ϵfb predicted by the SD model, as follows:

The final segmentation training loss, Lseg, is a combination of the cross-entropy loss and the Generative Feedback loss:

This dual loss mechanism allows the learned alignment expertise of the generative model to guide and refine the segmentation process. Specifically, if the predicted mask aligns closely with the real mask, the SD model can effectively reconstruct the image, resulting in a predicted noise close to the real noise and, therefore, reducing the Lfb. Conversely, significant deviations between the predicted and actual masks will hinder the reconstruction, leading to less accurate noise predictions and an increase in Lfb. By integrating Lfb into the overall segmentation loss, the discrepancy between the approximate mask and the actual mask is effectively reduced. By doing so, the segmentation model is encouraged to generate higher-quality mask prediction close to the ground truth.

4 Experiment

4.1 Dataset andMetrics

The evaluation experiments are conducted on multiple scenarios, including medical polyp dataset (ETIS [Silva et al., 2014], CVC-ClinicDB/CVC-612 [Bernal, 2015], CVCColonDB [Tajbakhsh et al., 2015], En-doScene [V´azquez et al., 2017], Kvasir [Jha et al., 2020]), industrial dataset

(Real-IAD [Wang et al., 2024], MVTec-AD[Bergmann et al., 2019]), and Floor Dirty dataset. Frechet Inception Distance (FID) [Heusel et al., 2017], Inception Score (IS) [Salimans et al., 2016] and Learned Perceptual Image Patch Similarity (LPIPS) [Zhang et al., 2018] are adopted to evaluate synthetic image quality. Intersection over Union (IoU) and Dice coefficient are utilized to evaluate the accuracy of downstream segmentation models.

4.2 Implementation Details

Using LoRA [Hu et al., 2022] to add image conditions for Stable Diffusion is adopted as the baseline generative model. Segformer [Xie et al., 2021] is adopted as the baseline segmentation model. AdamW is adopted as the optimizer. Input and output images are constrained to 512 × 512. The learning rate is set as 10-5 . The batch size is set as 4. For the inference of the diffusion model, the classifier free guidance scale is set as 7. We set the factor α for Lal to 0.7. For each dataset, we first finetune SD via real samples. Then, real masks are utilized to guide SD to generate synthetic samples. Ultimately, real samples are combined with augmented synthetic samples to train downstream segmentation models (see details in complementary materials).

4.3 Comparisons with SOTAMethods

Synthetic Data Quality Evaluation

In order to explore the matching degree between the generated image and the mask, we carry out relevant experiments. Figure 5 displays the data distribution of real masks, real images, generated images without alignment loss, and generated images with alignment loss after encoding by the VAE. The figure shows that without the alignment loss, there is a significant distribution gap between the generated images and both the real masks and real images. However, after adding the alignment loss, the distribution gap between the generated images and the real masks is noticeably reduced, and the gap with the real images also becomes relatively smaller.

For qualitative synthetic data assessment, we compare images generated by our approach to those from AnomalyDiffusion [Hu et al., 2024] and ArSDM [Du et al., 2023]. Figure 6 shows the results. In medical anomaly datasets with larger Ground-truth (GT) masks, our model's images closely match the GT, while ArSDM's images show misalignment and irregular shapes. For industrial and floor anomaly datasets with smaller GT masks, our model maintains high consistency between generated and real masks, unlike AnomalyDiffusion's images, which have significant shape inconsistencies.

In terms of the quantitative synthetic data evaluation, we compare our method to previous models using FID, IS, and LPIPS metrics on the medical polyp dataset. Results in Table 2 show our method yields lower FID and LPIPS scores, indicating greater similarity to real images, and a higher IS score, suggesting high-quality and diverse image generation.

Segmentation Accuracy Evaluation

As shown in Table 1, in the medical anomalies scenario, we primarily compare segmentation performance with data generated by ArSDM on the polyp dataset. We train our generative model on a smaller dataset with only 100 image-mask pairs to simulate the scarcity of labelled data in medical scenarios. Compared to using 100 samples to train ArSDM, our method yields much better segmentation performance when training segformer with generated images. ArSDM struggles with limited samples, and even with 1,450 samples, training PVT [Dong et al., 2021] or SANet [Wei et al., 2021] doesn't match our results. Our method excels in datasets like ETIS and EndoScene, which greatly differ from the training set, by better capturing image-mask relationships. ArSDM, however, only generates images similar to the training set and fails to adapt to significant differences. Thus, our improvements are notable on these datasets. On other test sets similar to the training set, where segmentation performance is near saturation, the proposed method shows no significant gains.

Table 3 shows industrial anomaly segmentation performance comparisons on the MvTec-AD and Real-IAD datasets. The Real-IAD dataset, split into easy (10 objects) and hard (20 objects) based on generation difficulty, reveals that segformer trained with our synthetic samples achieves 1.07% and 5.03% higher average mIoU on MvTec-AD and Real-IAD, respectively, than segformer trained on the samples generated by AnomalyDiffusion. Similar accuracy improvement can be observed on the Floor Dirty dataset.

4.4 Ablation Study

Contribution of Main Components

The main components: Alignment Regularization loss Lal and Generative Feedback loss Lfb, play a crucial role in enhancing the performance and stability of the model. Table 4

shows the individual and combined impact of these loss functions. Using Lal and Lfb separately brings gains of 5.9% and 6.07% on mDice and 6.81% and 6.86% on mIoU over baseline, respectively. Combining these two loss functions raises the performance to a new level, leading to improvements of 6.38% on mDice and 7.69% on mIoU. This indicates the effectiveness of the two main loss functions. It is worth noting that the Alignment Regularization impacts the segmentation model indirectly by enhancing its training data, while Generative Feedback mechanism is directly involved in its optimization process. Importantly, these two loss functions have a synergistic effect in improving the segmentation model.

Sample Size on Generative Model Training

Table 5 shows the impact of sample size on finetuning the generative model. As samples increase from 50 to 150, the segmentation model's mIoU improves from 78.80% to 82.22%. With only 20 samples, mIoU drops to 47.27%, likely due to insufficient tuning of the diffusion model, causing misalignments in synthetic data and misleading segmentation model optimization. This highlights the importance of adequate training samples for diffusion model finetuning.

Real Mask Numbers for Generative Model Inference

This study examines how the number of real masks guiding synthetic sample generation affects downstream segmentation performance. Figure 7 shows a non-linear relationship between real mask number and segmentation performance. mIoU fluctuates minimally from 10 to 60 masks, peaking at 80.47% with 20 masks and bottoming at 79.24% with 10 masks. From 60 to 100 masks, mIoU slightly increases to 80.38%. The high performance with few real masks indicates our method's robustness and low sensitivity to mask quality.

5 Conclusion

Our method introduces Alignment Regularization loss during generative model training to improve diffusion's understanding of image-mask relationships and output quality. It also integrates Generative Feedback loss into segmentation model training to optimize performance and accuracy. Extensive experiments validate the proposed approach's effectiveness, supporting their application and value in downstream tasks.

Acknowledgments

This work was supported by Young Scientists Fund of the National Natural Science Foundation of China (Grant No. 62302328), Jiangsu Province Foundation for Young Scientists (Grant No. BK20230482), Suzhou Key Laboratory Open Project (Grant No. 25SZZD07) and Jiangsu Manufacturing Strong Province Construction Special Fund Project (Grant Name: Research and Development and Industrialization of Intelligent Service Robots Integrating Large Model and Multimodal Technology).

Contribution Statement

For the co-authors of this paper, Xiangyue Li and Xiaoyang Wang contributed equally to this work.