Improving Diffusion-Based Augmentation for Data-Scarce Image Classification via Mask-Guided Reference Attention and Repulsive Noise Optimization

Hyun-Jic Oh; GaYeon Koh; Won-Ki Jeong

cs.CVNov 3, 2026

Improving Diffusion-Based Augmentation for Data-Scarce Image Classification via Mask-Guided Reference Attention and Repulsive Noise Optimization

Authors: Hyun-Jic Oh, GaYeon Koh, Won-Ki Jeong

Organizations: Korea University, Republic of Korea · VIENCE Inc., Republic of Korea

Abstract

Diffusion-based data augmentation has emerged as a powerful strategy for addressing data-scarce image classification. Recent pipelines typically adapt a pretrained text-to-image diffusion model to a limited number of class examples and subsequently sample synthetic images using diverse textual prompts. However, in extreme data-scarce regimes such as long-tailed single-sample settings, this approach suffers from a dual failure mode: the generated samples often lack discriminative foreground details crucial for classifier training, and the adapted class prior can bias the early denoising process, causing output images to exhibit similar patterns despite varying text prompts. To overcome these limitations, we propose a sampling-side framework of two complementary components that simultaneously enhance fine details and diversity. First, Mask-Guided Reference (MGR) Attention updates only foreground value tokens with reference features via zero-shot segmentation, limiting reference influence to the target object and mitigating excessive global feature injection. Second, Repulsive Noise Optimization (RNO) encourages initial noise trajectories to diverge from dominant generative modes of the adapted model, enabling target prompts to express broader contextual variation. Extensive evaluations show consistent improvements, with the largest gains in long-tailed single-sample settings of up to 2.50 and 2.94 points on CUB-LT IF100, and further robustness across few-shot and background-shift evaluations.

Figures & tables

Figure 2: Overview of the proposed sampling pipeline. (a) RNO optimizes initial noises before denoising; the target mask $\mathbf{M}$ is extracted at $t_{m}$ and the reference mask $\mathbf{M}_{r}$ is cached offline; MGR Attention refines target foreground values during the remaining steps. (b) RNO optimizes each target noise ( $\mathbf{z}_{T}^{i}$ ) to diverge from adaptation-prompt anchors (blue), which remain fixed throughout optimization. (c) MGR Attention constructs a hybrid bank $\mathbf{K}_{\text{bank}},\mathbf{V}_{\text{bank}}$ from target-background and reference-foreground tokens, updating only foreground values.

Figure 3: Qualitative comparison on CUB-LT and Flower-LT IF100 (prompts: “sitting on a ledge by water” and “against a brick wall”). (a) CL and (b) Diff-II rely on the adapted class prior, limiting foreground details and prompt-conditioned variation. (c) MGR enhances foreground details. (d) MGR+RNO (Ours) increases prompt-conditioned diversity while preserving foreground details, successfully following the given prompts.

	CUB-LT (ResNet-50)				Flower-LT (ResNet-50)				CUB-LT (ViT-B/16)				Flower-LT (ViT-B/16)
Method	IF100	IF20	IF10	$\Delta$ CE	IF100	IF20	IF10	$\Delta$ CE	IF100	IF20	IF10	$\Delta$ CE	IF100	IF20	IF10	$\Delta$ CE
CE	33.65	44.82	58.13	-	80.43	90.87	95.07	-	53.87	68.36	78.82	-	92.84	97.52	99.20	-
Mixup [ 47 ]	37.59	49.15	63.43	+4.52	81.11	92.13	96.16	+1.01	53.64	68.85	79.38	+0.27	92.06	97.79	98.96	-0.25
CutMix [ 46 ]	37.61	49.15	63.27	+4.48	81.96	92.28	95.80	+1.22	54.02	68.45	78.22	-0.12	92.06	97.28	99.00	-0.41
CMO [ 31 ]	32.94	44.08	57.62	-0.65	83.95	91.43	95.19	+1.40	49.22	64.20	79.55	-2.69	92.21	97.57	99.03	-0.25
CMO-DRW [ 2 ]	32.57	46.43	59.25	+0.55	84.07	92.02	95.29	+1.67	51.20	67.29	73.80	-2.92	95.27	98.15	99.17	+1.01

Table 1: LT classification results (%) on CUB-LT and Flower-LT (

\Delta

CE: gain over CE; Diff-II

\dagger

: SYNAuG [ 45 ] -style normalization).

	5-shot					10-shot
Method	CUB	Air.	Pet	Car	Avg	CUB	Air.	Pet	Car	Avg
CE	54.52	36.63	86.98	40.10	54.56	70.26	54.11	90.41	69.64	71.11
Mixup	54.33	36.75	86.36	41.43	54.72	70.00	55.28	89.49	70.28	71.26
CutMix	51.74	34.20	85.41	41.89	53.31	68.91	51.05	89.71	67.71	69.35
Real-Filter	55.03	40.40	87.66	57.60	60.17	67.36	47.21	90.14	69.64	68.59
Real-Guidance	54.67	36.27	87.22	41.60	54.90	69.81	54.77	90.00	70.97	71.40

Table 2: Few-shot (5-shot and 10-shot) classification results.

Method	(L,L)	(W,W)	(L,W)	(W,L)	Avg
CE	38.55	32.15	37.90	30.20	34.70
Real-Filter [ 13 ]	38.87	34.72	38.25	31.64	35.87
Real-Guidance [ 13 ]	39.49	33.64	39.80	29.78	35.68
Real-Mix [ 41 ]	30.29	29.44	30.02	26.95	29.18
Da-Fusion [ 38 ]	44.50	39.04	44.46	33.49	40.37
Diff-AUG [ 41 ]	49.02	40.59	48.71	38.58	44.23

Table 3: Background-shift evaluation on Waterbirds.

Figure 4: Qualitative module comparisons. (left) On CUB-LT IF100 and Pet 5-shot, (a) global K/V replacement (VSP-style [ 17 ] ), (b) attention-masked K/V replacement (MasaCtrl [ 3 ] -style), and (c) SAM3-masked K/V replacement introduce unwanted reference style or disrupt target layout, whereas (d) Ours better preserves target layout while enhancing localized foreground details. (right) On a single-image Flower-LT class with varying prompts, (a) CL collapses to similar outputs, (b) CADS increases variation at the cost of quality degradation, and (c) CNO generates diverse samples but struggles to follow prompt-conditioned variation, whereas (d) RNO produces diverse prompt-conditioned layouts while retaining recognizable foreground structure.

	CUB-LT IF100			Flower-LT IF100
Method	Few	Single	Avg	Few	Single	Avg
CL	42.95	37.79	54.16	77.37	64.39	87.74
MGR	44.70	39.46	55.61	80.67	69.14	89.34
RNO	43.23	37.18	54.49	76.35	61.37	87.16
MGR+RNO	46.08	41.05	56.09	82.18	72.04	90.34

Table 4: Module ablation on CUB-LT and Flower-LT IF100. Single: 1 image per class (CUB: 83 classes, Flower: 37 classes). Few:

<

<

10 images per class, including Single.

		CUB-LT IF100			Flower-LT IF100
	Method	Few	Single	Avg	Few	Single	Avg
(i)	CL	42.95	37.79	54.16	77.37	64.39	87.74
	K/V swap	43.26	37.46	54.71	79.41	66.36	88.57
	K/V swap (Attn. mask)	35.47	27.97	49.28	79.12	66.13	88.66
	K/V swap (SAM3)	51.26	31.67	51.36	79.21	66.01	88.64
	MGR (Attn. mask)	43.70	37.71	54.71	77.95	64.97	87.98

Table 5: Design choice ablation on CUB-LT and Flower-LT IF100. (i) Attention mechanism. (ii) Noise optimization.

	CUB-LT IF100			Flower-LT IF100
Method	MSS $\downarrow$	Vendi $\uparrow$	Ref-FG $\uparrow$	MSS $\downarrow$	Vendi $\uparrow$	Ref-FG $\uparrow$
CL	0.62	7.46	0.45	0.81	3.25	0.59
CADS	0.35	17.51	0.22	0.56	8.89	0.28
CNO	0.58	8.39	0.44	0.76	4.12	0.56
MGR	0.60	7.89	0.46	0.79	3.60	0.61
MGR+CNO	0.63	7.18	0.48	0.77	3.86	0.59

Table 6: Diversity and foreground detail metrics on single-sample classes of CUB-LT and Flower-LT IF100. MSS [ 34 ] and Vendi [ 7 ] measure diversity; Ref-FG Sim measures foreground similarity via DINOv2 [ 30 ] features on SAM3 [ 4 ] -segmented regions. Bold indicates best among MGR-based methods; CADS is excluded as it produces severe quality degradation (see Fig. 4 ).

	CUB-LT IF100			Flower-LT IF100
$t_{m}$	Few	Single	Avg	Few	Single	Avg
40	45.17	39.91	55.92	80.09	68.10	89.10
30 ✓	46.08	41.05	56.09	82.18	72.04	90.34
20	44.72	39.50	55.58	80.52	67.75	89.27
10	43.15	37.99	54.30	79.65	67.40	88.86

Table 7: Mask extraction timestep analysis on CUB-LT and Flower-LT IF100.

	SD v1.5			SD v2.1
Method	Few	Single	Avg	Few	Single	Avg
BOB	79.70	69.58	85.49	80.52	70.27	86.03
BOB+Ours	81.80	72.53	86.80	83.14	75.64	87.83

Table 8: Compatibility with BOB [ 44 ] on Flower-LT IF100. Reproduced following the official repository.

Figure 5: Target value features ( $\hat{\mathbf{V}}$ ) over denoising timesteps and final outputs: (a) CL, (b) foreground-only update ( $\hat{\mathbf{V}}{=}\hat{\mathbf{V}}^{\mathrm{FG}}$ ), and (c) MGR Attention. MGR aligns foreground activations with the reference while keeping target responses outside the foreground.

Figure 6: Layer selection for MGR Attention on CUB-LT IF100. (a) Up-blocks (Ours), (b) Mid-blocks, (c) Down-blocks, (d) CL baseline. Up-blocks best preserve foreground details.

Method	Many	Medium	Few	Single	Avg
Down-blocks	82.89	70.48	42.01	36.28	53.62
Mid-blocks	84.82	71.10	42.95	37.55	54.54
Up-blocks (Ours)	85.54	70.87	44.70	39.46	55.61

Table 9: Layer selection analysis for MGR Attention on CUB-LT IF100 (ResNet-50). Decoder up-blocks yield the best performance.

Figure 7: Feature modification strategy on CUB-LT IF100 and Flower-LT IF100. (a) Reference, (b) CL baseline, (c) key-only ( $\hat{\mathbf{K}}$ ), (d) key-value ( $\hat{\mathbf{K}},\hat{\mathbf{V}}$ ), (e) Ours: foreground value update ( $\hat{\mathbf{V}}$ ).

Figure 8: KV bank composition on CUB-LT and Flower-LT IF100. (a) Reference, (b) CL baseline, (c) foreground-only ( $\mathbf{K}_{r}^{\mathrm{FG}},\mathbf{V}_{r}^{\mathrm{FG}}$ ), (d) reference-full ( $\mathbf{K}_{r},\mathbf{V}_{r}$ ), (e) Ours ( $\mathbf{K}_{\mathrm{bank}}{=}[\mathbf{K}^{\mathrm{BG}};\mathbf{K}_{r}^{\mathrm{FG}}]$ ).

Segmenter	Many	Medium	Few	Single	Avg
CLIPSeg [ 25 ]	84.34	70.19	45.80	40.44	56.01
SAM3 [ 4 ] ✓	85.78	69.57	46.08	41.05	56.09

Table 10: Segmentation backbone analysis on CUB-LT IF100 (MGR+RNO, ResNet-50).

Figure 9: Effect of increasing RNO repulsion strength on CUB-LT and Flower-LT IF100. (a) CL. (b) Ours with increasing RNO repulsion strength from left to right, starting from the default setting (CUB: 3 steps, LR 0.01; Flower: 30 steps, LR 0.08). Prompts: “sitting on a branch of a tree” (CUB) and “with water droplets on it” (Flower). Greater repulsion introduces layout diversity but risks drifting from class-consistent foreground appearance (yellow boxes).

Steps	LR	Many	Medium	Few	Single	Avg
30	0.08	84.34	70.48	44.86	39.22	55.51
10	0.04	85.39	70.66	44.36	39.31	55.67
5	0.04	84.82	71.61	45.14	39.87	56.06
3	0.01✓	85.78	69.57	46.08	41.05	56.09

Table 11: RNO hyperparameter sensitivity on CUB-LT IF100 (MGR+RNO, ResNet-50).

Method	Many	Medium	Few	Single	Avg
CE only	85.54	67.14	18.33	7.38	38.04
Diff-II $\dagger$ [ 40 ]	83.86	70.14	42.01	36.69	53.59
P2P [ 14 ]	84.82	69.91	41.40	35.55	53.21
PHOTOSWAP [ 9 ]	83.37	63.57	19.83	11.62	37.73
MasaCtrl [ 3 ]	84.58	66.74	18.50	7.75	37.95
VSP [ 17 ]	85.30	70.81	42.62	36.77	54.28

Table 12: Comparison with attention-based editing and detection-oriented augmentation on CUB-LT IF100 (ResNet-50). All rows are our runs under the same protocol (main-paper Tab. 1 reports CE from Diff-II [ 40 ] ); attention-editing baselines use their original settings, unlike the in-pipeline K/V variants in main-paper Tab. 5 . Groups (classes): Many

x{>}20

(16), Medium

5{\leq}x{\leq}20

(61), Few

x{<}5

incl. Single (123), Single

x{=}1

(83).

Figure 10: Qualitative comparison on Car 5-shot and Aircraft 5-shot. (a) Global K/V replacement (VSP-style [ 17 ] ), (b) attention-masked K/V replacement (MasaCtrl [ 3 ] -style), (c) SAM3-masked K/V replacement, and (d) Ours (MGR Attention). (a) and (b) introduce global style propagation and blurry object regions; (c) reduces noise but can still produce blurry foreground. (d) better preserves target layout while enhancing localized foreground details.

Diffusion-Based Data Augmentation

See all 20 papers

Mar 9, 2026

Diffusion-Based Data Augmentation for Image Recognition: A Systematic Analysis and Evaluation

Diffusion-based data augmentation (DiffDA) has emerged as a promising approach to improving classification performance under data scarcity. However, existing works vary significantly in task configurations, model choices, and experimental pipelines, making it difficult to fairly compare methods or assess their effectiveness across different scenarios. Moreover, there remains a lack of systematic understanding of the full DiffDA workflow. In this work, we introduce UniDiffDA, a unified analytical framework that decomposes DiffDA methods into three core components: model fine-tuning, sample generation, and sample utilization. This perspective enables us to identify key differences among existing methods and clarify the overall design space. Building on this framework, we develop a comprehensive and fair evaluation protocol, benchmarking representative DiffDA methods across diverse low-data classification tasks. Extensive experiments reveal the relative strengths and limitations of different DiffDA strategies and offer practical insights into method design and deployment. All methods are re-implemented within a unified codebase, with full release of code and configurations to ensure reproducibility and to facilitate future research.

Zekun Li, Yinghuan Shi, Yang Gao +1

State Key Laboratory for Novel Software Technology, Nanjing University, Nanjing, China. · School of Computing and Data Science, The University of Hong Kong, Hong Kong SAR, China.

Diffusion-Based Data Augmentation Data Augmentation Diffusion Model Fine-Tuning Few-Shot Image Classification

Ask PDF

May 4, 2026

Synthetic Data Generation for Long-Tail Medical Image Classification: A Case Study in Skin Lesions

Long-tailed class distributions are pervasive in multi-class medical datasets and pose significant challenges for deep learning models which typically underperform on tail classes with limited samples. This limitation is particularly problematic in medical applications, where rare classes often correspond to severe or high-risk diseases and therefore require high diagnostic accuracy. Existing solutions-including specialized architectures, rebalanced loss functions, and handcrafted data augmentation-offer only marginal improvements and struggle to scale due to their limited and largely deterministic variability. To address these challenges, we introduce a diffusion-model-driven synthetic data augmentation pipeline tailored for medical long-tailed classification. Our approach features a novel inpainting diffusion model combined with an Out-of-Distribution (OOD) post-selection mechanism to ensure diverse, realistic, and clinically meaningful synthetic samples. Evaluated on the ISIC2019 skin lesion classification dataset, one of the largest and most imbalanced medical imaging benchmarks, our method yields substantial improvements in overall performance, with particularly pronounced gains on tail classes with more than

28\%

improvement on the class with the fewest samples. These results demonstrate the effectiveness of diffusion-based augmentation in mitigating long-tail imbalance and enhancing medical classification robustness.

Jiaxiang Jiang, Mahesh Subedar, Omesh Tickoo

Intel · Google

Skin Lesion Classification Synthetic Data Augmentation Long-Tailed Classification Medical Image Classification Diffusion-Based Image Generation

Ask PDF

Jul 14, 2026

Steering Diffusion Models via Class-Contrastive Influence for Few-Shot Medical Classification

When labeled data are scarce, off-the-shelf diffusion models can augment training sets for few-shot medical image classification, but not all generated samples are equally useful for the downstream task. Existing approaches largely improve synthetic data by increasing realism, diversity, or domain adaptation, while overlooking a more fundamental question: how should sample usefulness for classification be measured and optimized? We address this with Class-Contrastive Influence (C2I), a criterion that quantifies a sample's usefulness through its gradient-based influence on the classifier. We find that effective samples exhibit a strong C2I gap: their loss gradients align with validation gradients from the same class and oppose those from other classes. Our analysis further suggests that such high-C2I samples are hard, boundary-proximal examples that help refine the decision boundary and improve robustness. Building on this insight, we fine-tune diffusion models with reinforcement learning using a C2I-based reward to steer generation toward class-informative samples. Across several few-shot medical imaging benchmarks, C2I-guided generation improves downstream accuracy and robustness over diffusion-based augmentation baselines, showing that synthetic augmentation is most effective when guided by task usefulness rather than image quality alone.

Jeeyung Kim, Erfan Esmaeili, Qiang Qiu

Purdue University, West Lafayette, USA

Synthetic Data Augmentation Diffusion Model Fine-Tuning Diffusion-Based Image Generation

Ask PDF

Abstract

Diffusion-Based Data Augmentation

Explore similar work

Diffusion-Based Data Augmentation for Image Recognition: A Systematic Analysis and Evaluation

Synthetic Data Generation for Long-Tail Medical Image Classification: A Case Study in Skin Lesions

Steering Diffusion Models via Class-Contrastive Influence for Few-Shot Medical Classification