Sparsely Supervised Diffusion

Wenshuai Zhao^1,4, Zhiyuan Li¹, Yi Zhao², Mohammad Hassan Vali¹, Martin Trapp³, Joni Pajarinen¹, Juho Kannala^1,5, Arno Solin^1,4

¹Aalto University, Finland ²Imperial College London, UK ³KTH Royal Institute of Technology, Sweden

⁴ELLIS Institute Finland ⁵University of Oulu, Finland

NeurIPS 2026

Paper arXiv Code

Method overview comparing full regression supervision in standard flow matching with masked target supervision in SSD — **SSD applies sparse supervision to diffusion-model training.**

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 excessive correlations learned by the model. This can produce samples that are locally plausible but globally inconsistent. We propose a principled method to mitigate this issue. Our method, sparsely supervised diffusion (SSD), is a simple yet effective masking framework that can be implemented with only a few lines of code. We analytically show that SSD fundamentally alters diffusion model training by modifying the spectrum of the data covariance and that it suppresses correlations in the data covariance matrix. Experiments show that our method more accurately approximates the underlying population score function, reduces memorization on small datasets, and promotes the use of essential contextual information during generation. Moreover, even when up to 98% of pixels are masked, it achieves competitive FID scores across a range of datasets and avoids the training instability commonly observed on small datasets.

Motivation

TL;DR: Diffusion models often produce spatially inconsistent images. We propose sparsely supervised diffusion, a simple, principled method that reduces excessive correlations from limited data without relying on physics or geometry heuristics. It is effective and requires only a few lines of code.

#1#2#3#4#5#6#7#8#9Baseline
η = 0SSD
η = 0.2SSD
η = 0.5SSD
η = 0.8

CelebA samples from the same initial noise. Rows: η = 0, 0.2, 0.5, and 0.8.

Main theoretical analysis

Masking regularizes correlations

For zero-mean Gaussian data with covariance Σ, consider an analytical masked-data model x̃ = m ⊙ x. Each coordinate is retained independently with probability p = 1 − η, and the mask is independent of the data. Writing D = diag(Σ), the masked covariance is

\tilde{Σ} = p^{2} Σ + p (1 - p) D

(1)

Diagonal entries

Σ̃_ii = p Σ_ii

Coordinate variances are scaled by the retention probability.

Off-diagonal entries

Σ̃_ij = p² Σ_ij, i ≠ j

Cross-coordinate covariance is attenuated more strongly.

Masking changes directional variance

Let u_i be an eigenvector of the original covariance Σ, with eigenvalue λ_i. The variance of the masked data along that original direction becomes

{\tilde{λ}}_{i} := u_{i}^{T} \tilde{Σ} u_{i} = p^{2} λ_{i} + p (1 - p) u_{i}^{T} D u_{i}

(2)

Experiments

Generation quality

SSD achieves competitive FID scores on CIFAR-10, CelebA, LSUN Bedroom, and ImageNet.

FID scores (↓). Lower is better.
Dataset	Baseline	SSD
CIFAR-10	2.28 ± 0.04	2.24 ± 0.05
LSUN Bedroom	1.40 ± 0.04	1.39 ± 0.01
CelebA-50K	2.23 ± 0.23	1.85 ± 0.14
CelebA-10K	13.72	5.10
ImageNet (CFG = 1.5)	4.98	4.83

Training on limited data

On CelebA-10K, higher masking ratios mitigate the deterioration in FID that occurs with prolonged training. The training curves compare the baseline with several SSD masking ratios.

CelebA-10K FID versus training epochs for the baseline and SSD models with different masking ratios — FID during training on CelebA-10K. Higher masking ratios improve stability in this small-data setting.

Spatial consistency

For a controlled dataset of triangles and squares, we count scattered clusters in 5,000 generated images. SSD reduces images containing scattered clusters from 3,899 to 792, and the total number of clusters from 11,208 to 1,014.

Training shapes at the top, baseline-generated shapes in the middle, and SSD-generated shapes at the bottom

Finger-count errors on 11kHands

After 400 training epochs, we generate 1,000 samples per model using the same random seed and count four- and six-finger cases. SSD reduces the total from 35 to 22.

#1#2#3#4#5#6#7#8#9#10Baseline
η = 0SSD
η = 0.8

Selected hand pairs from the same initial noise after 400 epochs. Red outlines mark finger-count errors.

Error type	Baseline (η = 0)	SSD (η = 0.8)
Six fingers	26	8
Four fingers	9	14
Total	35	22

BibTeX

Download reference

@inproceedings{zhao2026sparsely,
  title     = {Sparsely Supervised Diffusion},
  author    = {Zhao, Wenshuai and Li, Zhiyuan and Zhao, Yi and Vali, Mohammad Hassan and Trapp, Martin and Pajarinen, Joni and Kannala, Juho and Solin, Arno},
  booktitle = {Advances in Neural Information Processing Systems},
  year      = {2026}
}