# Preserve the Hard, Regenerate the Rest: Uncertainty-Guided Synthetic Training Data Augmentation with Diffusion Models

arXiv: https://arxiv.org/abs/2606.31603 · HTML: https://mlsanity.com/paper/2606.31603/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Placed in Medical imaging and segmentation, one of 14 directions in Diffusion models.
- Builds on 34 of its 40 references that are in the map.
- First posted on arXiv 30 Jun 2026.

## Placed in

- Diffusion models / [Medical imaging and segmentation](https://mlsanity.com/diffusion/medical-imaging-and-segmentation/)

## Builds on (most cited first, 30 of 34)

- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2006.11239/) (2020)
- [High-Resolution Image Synthesis with Latent Diffusion Models](https://mlsanity.com/paper/2112.10752/) (2021)
- [DINOv2: Learning Robust Visual Features without Supervision](https://mlsanity.com/paper/2304.07193/) (2023)
- [Scalable Diffusion Models with Transformers](https://mlsanity.com/paper/2212.09748/) (2022)
- [Adding Conditional Control to Text-to-Image Diffusion Models](https://mlsanity.com/paper/2302.05543/) (2023)
- [Dropout as a Bayesian Approximation: Representing Model Uncertainty in Deep Learning](https://mlsanity.com/paper/1506.02142/) (2015)
- [The Cityscapes Dataset for Semantic Urban Scene Understanding](https://mlsanity.com/paper/1604.01685/) (2016)
- [SDXL: Improving Latent Diffusion Models for High-Resolution Image Synthesis](https://mlsanity.com/paper/2307.01952/) (2023)
- [BLIP: Bootstrapping Language-Image Pre-training for Unified Vision-Language Understanding and Generation](https://mlsanity.com/paper/2201.12086/) (2022)
- [SegFormer: Simple and Efficient Design for Semantic Segmentation with Transformers](https://mlsanity.com/paper/2105.15203/) (2021)
- [RePaint: Inpainting using Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2201.09865/) (2022)
- [BDD100K: A Diverse Driving Dataset for Heterogeneous Multitask Learning](https://mlsanity.com/paper/1805.04687/) (2018)
- [Class-Balanced Loss Based on Effective Number of Samples](https://mlsanity.com/paper/1901.05555/) (2019)
- [Deep Bayesian Active Learning with Image Data](https://mlsanity.com/paper/1703.02910/) (2017)
- [Training Region-based Object Detectors with Online Hard Example Mining](https://mlsanity.com/paper/1604.03540/) (2016)
- [Blended Latent Diffusion](https://mlsanity.com/paper/2206.02779/) (2022)
- [Learning Loss for Active Learning](https://mlsanity.com/paper/1905.03677/) (2019)
- [Deep Multi-modal Object Detection and Semantic Segmentation for Autonomous Driving: Datasets, Methods, and Challenges](https://mlsanity.com/paper/1902.07830/) (2019)
- [Effective Data Augmentation With Diffusion Models](https://mlsanity.com/paper/2302.07944/) (2023)
- [Computer Vision for Autonomous Vehicles: Problems, Datasets and State of the Art](https://mlsanity.com/paper/1704.05519/) (2017)
- [DiffuMask: Synthesizing Images with Pixel-level Annotations for Semantic Segmentation Using Diffusion Models](https://mlsanity.com/paper/2303.11681/) (2023)
- [DatasetDM: Synthesizing Data with Perception Annotations Using Diffusion Models](https://mlsanity.com/paper/2308.06160/) (2023)
- [Dataset Diffusion: Diffusion-based Synthetic Dataset Generation for Pixel-Level Semantic Segmentation](https://mlsanity.com/paper/2309.14303/) (2023)
- [DiffuseMix: Label-Preserving Data Augmentation with Diffusion Models](https://mlsanity.com/paper/2405.14881/) (2024)
- [FreeMask: Synthetic Images with Dense Annotations Make Stronger Segmentation Models](https://mlsanity.com/paper/2310.15160/) (2023)
- [CEREALS - Cost-Effective REgion-based Active Learning for Semantic Segmentation](https://mlsanity.com/paper/1810.09726/) (2018)
- [X-Paste: Revisiting Scalable Copy-Paste for Instance Segmentation using CLIP and StableDiffusion](https://mlsanity.com/paper/2212.03863/) (2022)
- [TRACER: Extreme Attention Guided Salient Object Tracing Network](https://mlsanity.com/paper/2112.07380/) (2021)
- [Dataset Enhancement with Instance-Level Augmentations](https://mlsanity.com/paper/2406.08249/) (2024)

## Abstract (arXiv)

Semantic segmentation models struggle with data sparsity and rare or visually diverse regions, e.g., dense regions or small objects in aerial or autonomous mobility data. While synthetic augmentation is an appealing solution, directly generating new labeled data risks misalignment of labels and generated pixels. Existing solutions to this problem often rely on external models, or employ coarse heuristics such as indiscriminately augmenting all foreground objects or entire backgrounds, which wastes capacity on uninformative pixels. To address this, we propose an uncertainty-guided synthetic context augmentation strategy that strictly preserves label validity and efficiently maximizes pixel informativeness per synthetic sample - no external guardrails required. Using a baseline segmenter's predictive entropy, we identify uncertain semantic regions and inpaint only the complementary visual context. When fine-tuning the segmenter on this synthetic data, we compute the loss only over the original pixels, excluding inpainted regions. This focuses learning on the unmodified, uncertain regions while presenting them in novel contexts. We demonstrate substantial mIoU gains on Cityscapes, UAVID, and BDD100K with the largest gains on rare and difficult classes such as buses, trains, or (from the aerial perspective) cars. Our results demonstrate that uncertainty-guided context augmentation is a highly effective lever to improve segmentation performance on complex datasets, with code provided at https://github.com/XITASO/Preserve-the-Hard-Regenerate-the-Rest.