SGD-Mix: Enhancing Domain-Specific Image Classification with Label-Preserving Data Augmentation

TL;DR

This work addresses domain-specific image classification where data augmentation must balance diversity, faithfulness, and label clarity. It introduces SGD-Mix, a saliency-guided diffusion mix framework that preserves the foreground of a source image, replaces the background with a contextually diverse target image, and refines the result with a domain-specific diffusion model. The method combines saliency-based target selection, saliency-guided mixing via binary masks, and DreamBooth-LoRA/Textual Inversion fine-tuning, with a Strength parameter to trade off faithfulness and diversity. Empirical results across fine-grained, long-tail, few-shot, and background-robustness benchmarks demonstrate consistent improvements over state-of-the-art diffusion-based and traditional augmentation methods, highlighting practical viability for domain-specific tasks.

Abstract

Data augmentation for domain-specific image classification tasks often struggles to simultaneously address diversity, faithfulness, and label clarity of generated data, leading to suboptimal performance in downstream tasks. While existing generative diffusion model-based methods aim to enhance augmentation, they fail to cohesively tackle these three critical aspects and often overlook intrinsic challenges of diffusion models, such as sensitivity to model characteristics and stochasticity under strong transformations. In this paper, we propose a novel framework that explicitly integrates diversity, faithfulness, and label clarity into the augmentation process. Our approach employs saliency-guided mixing and a fine-tuned diffusion model to preserve foreground semantics, enrich background diversity, and ensure label consistency, while mitigating diffusion model limitations. Extensive experiments across fine-grained, long-tail, few-shot, and background robustness tasks demonstrate our method's superior performance over state-of-the-art approaches.

Figures (6)

• Figure 1: Top row: Non-generative mixup-based methods. Bottom row: Generative diffusion-based methods. Inside the dashed box: Methods that mix two different training images and their corresponding labels. Outside the dashed box: Label-preserving methods without label mixing. Note that the translation strength for the bottom row in the figure is consistently set to 0.7.
• Figure 2: Semantic drift in DiffuseMix islam2024diffusemix generated images using the prompt "A transformed version of image into autumn". Stronger transformations reduce semantic fidelity to the source image. More examples in Supplementary Materials 2.
• Figure 3: The pipeline of the proposed SGD-Mix method. The process involves three major stages: (1) saliency-based target selection, (2) saliency-guided mixing, and (3) refining mixed images using our domain-specific fine-tuned diffusion model.
• Figure 4: Visualization of attention maps before and after saliency-guided mixing in SGD-Mix. For a source image $I_i$ and target image $I_j$, the attention maps (bottom row) consistently focus on $I_i$'s foreground region in the mixed image $I_{(i,j)}$, preserving semantic consistency. More examples in Supplementary Materials 4.1.
• Figure 5: Examples of SGD-Mix generated images under varying translation strengths $S \in \{0.1, 0.3, 0.5, 0.7, 0.9\}$. The generated image retains the source image's foreground semantics while the background evolves with increasing $S$, balancing diversity and faithfulness. More examples in Supplementary Materials 4.2.