Sparsely Supervised Diffusion

TL;DR

Sparsely Supervised Diffusion (SSD) trains a flow-matching / diffusion model with a random binary mask applied to the regression loss, so the network only fits unmasked pixel positions and must inpaint the rest from context. Mask ratios up to 98% are reported safe; FID stays competitive with the unmasked baseline. The headline practical wins are on small datasets: SSD eliminates the long-training divergence that wrecks the baseline on CelebA-10K, cuts memorization (larger nearest-neighbor distances at matched FID), and produces more spatially consistent samples on a toy shapes dataset (11× fewer scattered-pixel clusters). The proposed mechanism: masking attenuates off-diagonal entries of the data covariance, suppressing the spurious correlations that locality of the U-Net would otherwise overfit to on limited data.

Key claims

Masking up to 98% of pixels in the training-loss (not the input) of a flow-matching model retains competitive FID against the unmasked baseline across CIFAR-10, CelebA-50K, LSUN-Bedroom and ImageNet-32 [§5.1, Fig. 3].
On CelebA-10K (small data regime), the unmasked flow-matching baseline diverges at long training horizons (FID 13.72), while SSD at p=0.98 stays stable and reaches FID 5.10 — a 2.7× FID improvement attributed entirely to a training-stability fix [§5.2, Table 1].
SSD mitigates training-set memorization: at matched/lower FID, SSD-generated samples on CelebA-10K sit at average nearest-neighbor L2 distance 46.02 vs the baseline’s 42.32 [§5.3, Table 1].
The masking strategy provably modifies the spectrum of the data covariance: it attenuates off-diagonal entries (inter-pixel correlations) more strongly than diagonal entries, which the authors connect to the “locality” mechanism that drives diffusion-model generalization [§4.2, Eq. 10].
SSD shifts the gradient-sensitivity distribution from narrow (baseline: most pixels exert similar influence on a target pixel) to wide (SSD: a few pixels exert strong, selective influence) — interpreted as the model learning to rely on essential context rather than diffuse averaging [§5.4, Fig. 6].
On a synthetic 500-image binary-shapes dataset, SSD produces 1,014 scattered-pixel clusters across 5,000 generated images vs the baseline’s 11,208 (11× reduction); spatially inconsistent images drop from 3,899 to 792 [§5.5, Table 2].
On 2-D Gaussian toy data with known ground-truth score, masking lowers the average score-estimation error from 3.65 (baseline) to 2.78 (SSD) by inducing an independence prior across data dimensions [§5.6, Fig. 7].
Method is architecture-agnostic and “a few lines of code”: apply a fresh random binary mask to the regression loss per sample, with no MAE-style asymmetric encoder/decoder (contra MDT) [§4.1, Eq. 8].

Method

SSD modifies the flow-matching (equivalently, diffusion) training objective by multiplying the per-pixel squared error by a fresh random binary mask M ∈ {0,1}^d whose entries are i.i.d. Bernoulli with P(M_i = 1) = 1 - p. At p = 0 it recovers vanilla flow matching; at p = 0.98 it supervises only ~2% of pixels per sample. Same U-Net, same Heun-2 / EDM sampling, same NFE — only the loss reduction changes. The paper argues the effect of this seemingly aggressive masking is a covariance reshaping: writing Σ̃ = M ⊙ Σ ⊙ M and decomposing on Σ’s eigenbasis, the ratio of new-to-old variance along eigenvector v_i is λ̃_i / λ_i = ⟨v_i, diag(Σ)⟩ / λ_i, which promotes high-variance-but-low-eigenvalue directions and suppresses dominant cross-pixel correlations. They tie this to Lukoianov et al.’s recent finding that diffusion “locality” is dataset-statistics-driven (not architecture-driven), so suppressing spurious off-diagonal correlations is the direct mechanism by which SSD changes locality.

Results

Headline numbers, all with a single U-Net backbone, Heun-2 / dopri5 sampling, FID over 50k samples:

CIFAR-10 (32×32), ImageNet (32×32), LSUN-Bedroom (32×32), CelebA-50K (64×64): SSD at p=0.80 matches baseline FID across training [§5.1, Fig. 3]. No headline FID number is reported as a SOTA claim — competitive parity is the claim.
CelebA-10K small-data regime: baseline diverges after ~1,500 epochs; SSD at p=0.98 reaches FID 5.10 vs baseline’s 13.72 at the same checkpoint [§5.2, Table 1, Fig. 1].
Memorization: at matched compute, mean NN L2 distance 46.02 (SSD) vs 42.32 (baseline) on CelebA-10K — generated samples sit measurably further from training samples [§5.3, Table 1].
Spatial-consistency toy benchmark (500-image triangles+squares): 5,000 samples evaluated, scattered-pixel cluster count 1,014 (SSD) vs 11,208 (baseline); spatially-inconsistent-image count 792 vs 3,899 [§5.5, Table 2].
Gradient-sensitivity distribution: SSD shows higher peak influence with wider per-pixel variance; baseline shows narrow, uniformly-distributed influence [§5.4, Fig. 6].

The paper does not report ImageNet-256 FID, large-model ablations, or any combination with REPA / DINO / RAE. All experiments use a small convolutional U-Net at ≤64² resolution.

Why it’s interesting

The headline FID parity result on CIFAR/LSUN/ImageNet-32 doesn’t move the dial against the SOTA-on-large-data recipes already on the wiki — RAEv2 (Improved Baselines with Representation Autoencoders), Self-Flow (Self-Flow: Self-Supervised Flow Matching for Scalable Multi-Modal Synthesis), or PixelGen (PixelGen: Improving Pixel Diffusion with Perceptual Supervision) — and the experimental scale (≤64², small U-Net) is far below the regime Luma generates in. Kyle’s read in Slack — not viable on Luma’s large datasets, but possibly interesting for SFT or data-constrained settings — matches the paper’s own pitch: every meaningful win (divergence prevention, memorization, spatial-consistency) is reported on small or toy datasets. Where it could matter: small-dataset SFT / personalization passes where overfitting + memorization on a few hundred to ten thousand images is the actual binding constraint — exactly the regime where the FIT virtual-try-on dataset (FIT: A Large-Scale Dataset for Fit-Aware Virtual Try-On) or any reference-image-conditional fine-tuning lives. The mechanism (covariance off-diagonal suppression) also makes it interesting as a contrast to internal self-distillation recipes (SRA, Self-Flow): SSD argues for reducing the supervision signal to fight memorization; Self-Flow argues for richer self-generated supervision. Both target generalization but pull in opposite directions on what to do with the loss.