OpenAI’s diffusion upgrade: learned noise and cosine schedules

What it does This repository contains the training and sampling code for the paper Improved Denoising Diffusion Probabilistic Models. It lets you train unconditional or class-conditional diffusion models on ImageNet, CIFAR-10, LSUN, or any folder of images, then generate samples including via DDIM. Pre-trained checkpoints are provided for several standard model configurations.

The interesting bit Rather than hiding the paper’s contributions inside library internals, the codebase surfaces them as explicit training flags—learned sigmas, cosine noise schedules, and importance-sampled VLB objectives are all togglable command-line options. That makes the 2021 improvements easy to reproduce or ablate without reimplementing the math.

Key highlights

• Pre-trained checkpoints available for ImageNet-64, CIFAR-10, LSUN bedrooms, and a 256×256 upsampler
• Supports class-conditional training and custom datasets with minimal preprocessing
• DDIM sampling and timestep respacing for faster generation
• Distributed training via MPI with microbatching for memory-limited GPUs
• EMA model saving, with the README noting that EMA checkpoints produce much better samples

Caveats

• The README is essentially a training cookbook; architectural details are delegated to the paper
• Published hyperparameters assume multi-GPU resources, and the README explicitly warns that most runs will not fit on a single GPU without shrinking --microbatch

Verdict A solid starting point if you’re reproducing the 2021 OpenAI diffusion baselines or studying learned variances and cosine schedules in their original implementation. If you want a modern, batteries-included diffusion framework, this is too low-level.