Sparsely Supervised Diffusion

TL;DR

Diffusion models often exhibit spatial inconsistencies due to locality in denoising. Sparsely Supervised Diffusion (SSD) introduces a simple masking mechanism in the regression loss, training on unmasked pixels while encouraging generalization over masked regions. Analytically, SSD alters the data covariance spectrum, changing learning dynamics and reducing memorization on small datasets. Empirically, SSD achieves competitive FID across datasets, improves spatial consistency, strengthens population score estimation, and remains stable under heavy masking, making it a practical, architecture-agnostic improvement for diffusion-based generation.

Abstract

Diffusion models have shown remarkable success across a wide range of generative tasks. However, they often suffer from spatially inconsistent generation, arguably due to the inherent locality of their denoising mechanisms. This can yield samples that are locally plausible but globally inconsistent. To mitigate this issue, we propose sparsely supervised learning for diffusion models, a simple yet effective masking strategy that can be implemented with only a few lines of code. Interestingly, the experiments show that it is safe to mask up to 98\% of pixels during diffusion model training. Our method delivers competitive FID scores across experiments and, most importantly, avoids training instability on small datasets. Moreover, the masking strategy reduces memorization and promotes the use of essential contextual information during generation.

Figures (11)

• Figure 1: Masking up to 98% of pixels improves training. Our model learns exclusively from the unmasked pixels while being encouraged to generalize over the masked regions. The curves show FIDs during training on CelebA-10K. FIDs tend to increase with prolonged training under a small masking ratio $\eta$, whereas a large masking ratio enhances generalization and stabilizes training.
• Figure 2: Overview of the proposed method in comparison with standard flow matching (FM). The proposed SSD differs from FM via sparse learning, which uses masked supervised signals to prevent diffusion models from memorizing training data point, and to encourage the model to leverage contextual information when predicting pixels. An individual image can be randomly masked several times, yielding multiple distinct sparse images.
• Figure 3: Effect of masking during training. When $\eta=0$, the model reduces to the baseline flow matching lipman2024flowmatchingguidecode. Across four datasets, SSD with up to 80% masked pixels can still achieve comparable performance as the baseline. Notably, the baseline in \ref{['fig:celeba_50k']} eventually explodes while SSD remains stable. The shaded regions indicate the 95% confidence interval over four runs.
• Figure 4: Masking leads to more realistic generations. Visualization of non-curated generated samples from models trained on CelebA dataset. Sample $\#1$ shows more detailed neck wrinkles with $\eta=0.8$. In sample $\#6$, an unrealistic bulge on the head is present at $\eta=0$ but gradually disappears as $\eta$ increases. A similar artifact is observed in sample $\#9$ of the baseline, but disappears with masking.
• Figure 5: Demonstration of improved spatial consistency. After training for 5k epochs on a binary image dataset of triangles ( ) and squares ( ) (1st row), our model generates images with less scattered dots (3rd row), i.e. enhanced structural integrity, compared to the standard diffusion model without masking (2nd row). Quantitative results are reported in \ref{['tab:spatial_consistency']}.

Theorems & Definitions (2)

• Remark 4.1
• Remark 4.2