Authors: Wenshuai Zhao, Zhiyuan Li, Yi Zhao, Mohammad Hassan Vali, Martin Trapp, Joni Pajarinen, Juho Kannala, Arno Solin
Organizations: Department of Computer Science, Aalto University, Finland · Department of Electrical Engineering and Automation, Aalto University, Finland.
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, importantly, avoids the training instability commonly observed on small datasets.
Figures & tables
Figure 1: Masking up to 98% of pixels improves training . SSD learns exclusively from unmasked pixels while being encouraged to generalize over the masked regions.
Figure 2: Overview of the proposed method in comparison with standard flow matching (FM).
Figure 3: Effect of masking during training. When η=0 , the model reduces to the baseline flow matching [ 49 ] . Across three datasets, SSD with up to 80% masked pixels can still achieve comparable performance as the baseline. Notably, the baseline in Fig. 3(b) 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. Non-curated samples from models trained on CelebA. Sample #1 shows finer neck wrinkles at η=0.8 . In sample #6, an unrealistic bulge on the head at η=0 gradually disappears as η increases. A similar artifact in sample #9 of the baseline is also removed by masking.
Dataset
Baseline (FM/DiT)
SSD (ours)
CIFAR10
2.28 (0.04)
2.24 (0.05)
LSUN
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
4.98 (–)
4.83 (–)
Table 1: FID ( σ ) ( ↓ ) on different datasets.
Method
Baseline ( η=0.0 )
Ours ( η=0.98 )
FID ( ↓ )
13.72
5.10
ℓ2 distance ( ↑ )
42.32
46.02
Table 2: Masking mitigates memorization. FID ( ↓ ) and average ℓ2 distance ( ↑ ) to the nearest training samples over 10k generated images on CelebA-10K. Larger ℓ2 indicates less memorization.
Figure 5: Average ℓ2 distance to the nearest training samples over training iterations on CelebA-10K. Larger values indicate less memorization.
Figure 6: Masking helps capturing essential contextual information. (a) Distribution of ℓ1 norm of gradients computed per image and, (b) distribution of gradient magnitude per pixel over 10k images for CIFAR10 dataset.
Figure 7: 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 Table 3 .
Method
No. of scatters ( ↓ )
No. of inconsistent images ( ↓ )
Baseline
11,208
3,899
Ours
1,014
792
Table 3: Masking improves spatial consistency . The table shows the number of scattered pixel clusters and number of images containing such clusters over 5k generated images for the illustrative dataset shown in Fig. 7 .
Error type ( ↓ )
Baseline ( η=0 )
SSD ( η=0.8 )
Six fingers
26
8
Four fingers
9
14
Total
35
22
Table 4: Comparison of finger-count errors among 1,000 generated 11kHands samples.
Figure 8: Masking improves estimating the population score function . Visualization of population scores (ground truth) and empirical scores estimated with η=0 (baseline) and η=0.5 (ours), along with a comparison of score estimation errors (right two figures) such that the error is color coded by 0 10.63. The baseline model only accurately estimates scores near observed data points (shown as ), whereas the proposed masking enables more accurate estimation over broader regions (squared areas) and yields lower overall errors.
Metric
Baseline
W/ Masking (ours)
Mean error ( ↓ )
3.65
2.78
Max error ( ↓ )
10.63
6.18
Table 5: Comparison of empirical score estimation errors on a 2D Gaussian data distribution. Lower ( ↓ ) is better.
Appendix figures & tables9 assets
Supplementary material from the paper’s appendix.
Appendix
Parameter
CelebA-50K
CIFAR10
LSUN Bedroom
Training Configuration
Epochs
2000
3000
2000
Effective Batch Size
128
256
1024
Learning Rate
0.0001
0.0001
0.0001
Optimizer Betas
[0.9, 0.95]
[0.9, 0.95]
[0.9, 0.95]
Use EMA
True
True
True
Appendix
Table 6: Key hyperparameters for experiments on three datasets
Method
CelebA-10K FID ( ↓ )
CIFAR10 FID ( ↓ )
MDT
24.72
–
Dropout ( p=0.30 )
13.72
2.28
Dropout ( p=0.98 )
13.38
9.37
SSD (ours, η=0.98 )
5.10
2.20
Appendix
Table 7: Comparison with dropout and MDT baselines. Both dropout and MDT fail to prevent divergence during training. All methods on CelebA-10K are trained for 2000 epochs.
Training data
Training steps
Baseline
SSD ( η=0.8 )
SSD ( η=0.98 )
ImageNet-100K
200K
25.99
18.45
15.58
Full ImageNet
180K
4.98
–
4.83
Appendix
Table 8: ImageNet FID at CFG=1.5 using DiT-XL/2 at 256 × 256.
Masking strategy
Mask ratio η
Final FID ( ↓ )
None (baseline)
0
13.72
Element-wise
0.5
13.29
Element-wise
0.95
6.36
Element-wise
0.98
5.10
Patch-wise (8 × 8)
0.6
7.61
Patch-wise (8 × 8)
0.9
5.14
Appendix
Table 9: Masking-granularity ablation on CelebA-10K. Results are reported at 3,500 epochs.
Figure 9: Comparison of generated samples on the 11kHands dataset. Each panel presents 200 samples generated using the same random seed. Thin red outlines highlight samples exhibiting spatial inconsistencies, specifically an incorrect number of fingers. (a) The baseline model frequently generates hands with six fingers, (b) whereas SSD mitigates this issue but occasionally produces hands with only four fingers. The full evaluation results are reported in Table 4 .
Figure 10: Sample paths from same initial noises with models trained with different η across three datasets.
Figure 11: More samples generated by models trained on CelebA-50K with different mask ratios.
Figure 12: More samples generated by models trained on CIFAR10 with different mask ratios.
Figure 13: More samples generated by models trained on LSUN Bedroom with different mask ratios.
Diffusion models (DMs) excel on dense continuous data, but are not designed for sparse continuous data. They do not model exact zeros that represent the deliberate absence of a signal. As a result, they erase sparsity patterns and perform unnecessary computation on mostly zero entries. With Sparsity-Exploiting Diffusion (SED), we model only non-zero values, preserving sparsity. SED delivers computational savings while maintaining or improving generation quality by skipping zeros during training and inference. Across physics and biology benchmarks, SED matches or surpasses conventional DMs and domain-specific baselines, while vision experiments provide intuitive insights into the limitations of dense DMs and the benefits of SED.
Phil Sidney Ostheimer, Mayank Nagda, Andriy Balinskyy +6
1RPTU University Kaiserslautern-Landau · 2Heidelberg Univer · University of California, Irvine.
Many training-free sparse attention methods are effective for accelerating diffusion models. Recently, several works suggest that making sparse attention trainable can further increase sparsity while preserving generation quality. We study three key questions: (1) when do the two common masking rules, i.e., Top-k and Top-p, fail, and how can we avoid these failures? (2) why can trainable sparse attention reach higher sparsity than training-free methods? (3) what are the limitations of fine-tuning sparse attention using the diffusion loss, and how can we address them? Based on this analysis, we propose SpargeAttention2, a trainable sparse attention method that achieves high sparsity without degrading generation quality. SpargeAttention2 includes (i) a hybrid masking rule that combines Top-k and Top-p for more robust masking at high sparsity, (ii) an efficient trainable sparse attention implementation, and (iii) a distillation-inspired fine-tuning objective to better preserve generation quality during fine-tuning using sparse attention. Experiments on video diffusion models show that SpargeAttention2 reaches 95% attention sparsity and a 16.2x attention speedup while maintaining generation quality, consistently outperforming prior sparse attention methods.
The performance of pre-trained masked diffusion models is often constrained by their sampling procedure, which makes decisions irreversible and struggles in low-step generation regimes. We introduce a novel sampling algorithm that works with pre-trained models and, after a lightweight fine-tuning of a single layer, significantly improves sample quality and efficiency. Our method reformulates the generation process using a star-shaped paradigm, which inherently allows for error correction. To make this process effective, we augment it with a learnable re-masking scheduler that intelligently identifies and revises likely errors. This approach yields a substantial quality boost, particularly when using a small number of sampling steps. We extensively ablate key components of our approach and show its usability in different scenarios. In comprehensive experiments on text, and code generation, our sampling algorithm outperforms or matches existing methods.