StepDrop Stochastic Step Skipper

StepDrop: Stochastic Step Skipping

---

Background

StepDrop began as the final project for ECE 661: Computer Engineering Machine Learning and Deep Neural Nets with my teammates Logan Chu and Wanghley Martins. By the end of the course, we had familiarized ourselves with all kinds of machine learning networks and their architectures, so for the final project, we were given free rein to do whatever we liked. We ended up targeting diffusion models due to their widespread use and the fact that they are miserably slow to sample from: a standard DDPM needs on the order of 1000 sequential U-Net forward passes to turn noise into a picture. On data center GPUs, this isn't too consequential, however, on edge devices, it leads to less cute cat images generated per minute. We were also curious to see if we could exploit a quality of large diffusion models called temporal redundancy, which essentially means that consecutive time steps carry mostly overlapping information, and therefore benefits from importance-aware acceleration. We trained our own tiny (under 10M parameter) U-Net diffusion model on CIFAR-10 from scratch specifically so we'd have full control over the sampling loop, rather than fighting someone else's pretrained pipeline.

◢ Comparison of DDPM, DDIM, and StepDrop sampling methodologies

---

Motivation

The standard fix for slow diffusion sampling is DDIM: instead of walking all 1000 timesteps, you deterministically subsample a fixed stride (every 20th step, say) and get away with it because consecutive timesteps predict similar noise. It works, but it treats every timestep as equally disposable.

That assumption looked wrong to us. Early timesteps (high noise) are laying down the coarse structure of the image: silhouette, color blocks, rough composition. Late timesteps (low noise) are polishing fine texture and edges. The steps in between are mostly just... nudging pixels that are already roughly where they'll end up. If that's true, uniform striding is leaving quality on the table by spending the same effort everywhere instead of protecting the steps that actually matter.

"If mid-trajectory steps contribute the least, why are we paying full price for them?"

So the question we set out to answer: can a sampler that's aware of which steps matter beat one that just skips uniformly, at the same compute budget?

---

Method

StepDrop wraps a standard DDIM update rule with a stochastic (or importance-weighted) decision at every timestep: evaluate the U-Net here, or skip and interpolate through? The skip probability is shaped by a scheduling function P(t) rather than being constant, so the sampler can be told to protect specific regions of the trajectory.

Skip strategies

We implemented several closed-form schedules for P(t), each encoding a different hypothesis about where redundancy lives (t normalized to [0, 1], 1 = pure noise):

Strategy	Formula	Shape
constant	P(t) = p	Flat baseline
linear	P(t) = p · 4t(1-t)	Parabolic peak mid-trajectory
cosine_sq	P(t) = p · sin²(πt)	Smoother middle peak
quadratic	P(t) = p · 16t²(1-t)²	Sharper middle peak
critical_preserve	piecewise	Near-zero skip in [0.3, 0.7], high at extremes

A hard guard rail sits underneath all of them: the first and last 20 timesteps are never skipped, no matter what the schedule says, because that's exactly where structure formation and fine detail live. Ablating this boundary protection alone cost +2.1 FID.

The variant that actually won, though, wasn't stochastic at all. AdaptiveStepDropSampler perturbs the input with a small amount of noise, measures how much the model's noise prediction moves, and uses that prediction-instability as an error proxy: low instability means the model is confident and it's safe to skip; high instability means slow down. And a separate deterministic importance allocator, TargetNFEStepDropSampler, just budgets a fixed step count up front: 30% of the budget to the first decile (t ∈ [900,1000]), 30% to the last decile (t ∈ [0,100]), and the remaining 40% spread across the middle 80%. That deterministic allocation beat every coin-flip strategy we tried by 4-13 FID; planned step placement mattered more than randomness.

◢ Timestep selection barcode comparing uniform DDIM striding, StepDrop's importance allocation, and pure stochastic skipping

The barcode above is the clearest way to see the difference. Uniform DDIM (top) spends its 50-step budget evenly across the whole trajectory. StepDrop's importance allocator (middle) crowds steps toward both ends and thins out the middle. Pure stochastic skipping (bottom) is what you get if you just flip a weighted coin at every step. It's noisier, and it underperforms the planned allocation.

---

Architecture

The model itself is a small residual U-Net (src/modules.py):

• Sinusoidal position embeddings for the timestep, injected into every residual block
• Two down-blocks (base_channels → 2×), a self-attention bottleneck, two up-blocks with skip connections
• GroupNorm + SiLU throughout, trained with the standard ε-prediction objective on a cosine noise schedule

Everything downstream (DDPM, DDIM, and both StepDrop variants) samples from the exact same checkpoint, which is what makes the comparison fair: any quality difference is coming from the sampling schedule, not the model.

---

Results

We benchmarked on CIFAR-10 against FID, Inception Score, and throughput, matched at equal NFE (number of function evaluations, i.e., U-Net forward passes, our actual compute proxy):

Method	NFE	FID ↓	IS ↑	Throughput (img/s) ↑
DDPM-1000	1000	72.25	4.25	7.9
DDIM-50	50	77.20	4.33	57.9
DDIM-25	25	79.46	4.29	73.5
StepDrop (Target-50, importance)	50	76.05	4.33	49.0
StepDrop (Target-25, importance)	25	77.61	4.13	96.8

At 50 NFE, StepDrop beats DDIM-50 by 1.15 FID at identical compute. At 25 NFE it matches DDIM-50's quality while roughly doubling throughput: half the GFLOPs (59 → 29.5 GFLOPs per image) for the same picture. It never beats the full 1000-step DDPM baseline outright, but that's the point: it Pareto-dominates DDIM along the compute axis that actually matters for deployment.

◢ Compute cost vs. FID quality tradeoff scatter plot, StepDrop Pareto-dominating DDIM

◢ Denoising evolution film strip comparing DDIM and StepDrop trajectories from pure noise to final image

◢ Bar chart of total GFLOPs and peak memory across DDPM, DDIM, and StepDrop configurations

Residual analysis against full-length DDIM trajectories confirmed StepDrop stays close to the deterministic path even when skipping through the middle of the trajectory, which is reassuring since the whole method hinges on mid-trajectory error not compounding into something visible.

---

Writeup & Code

This project has a full NeurIPS-formatted report with the derivations, ablations (skip-schedule sweep, boundary-protection ablation, adaptive-vs-fixed comparison), and appendix tables that didn't make it into this page. Code, report, and poster all live in the repo.