About this notebook
This notebook explains the JiT paper by Li and He. That paper trains a network to predict the clean image, not the noise. The notebook builds the model in PyTorch and uses CIFAR-10 as the real data. It keeps each image as a raw 3072-dimensional pixel vector. It uses no VAE and no tokenizer. The charts use Plotly. The interactive parts use anywidget, traitlets, and the wigglystuff cell tour.
A drag canvas starts the argument. You drag your cursor across the canvas to place noisy points. The clean-image target stays pinned to a 1-D curve. The noise target scatters across the 2-D plane. A button then measures the effective dimension of each target. Later sliders control the classifier- free guidance, the sampler steps, and the input noise level. Each slider change re-runs the next cell on the GPU. One slider drives a gallery of held-out test images. At low noise the model recovers the true image. At high noise the output becomes a plausible hallucination.
The model is a Diffusion Transformer with 7.4M parameters. It predicts the clean image with EDM weighting. Training runs on the GPU for 8000 steps with EMA and classifier-free guidance dropout. The notebook saves a checkpoint and reads it on later runs, so sampling never retrains. The full 50k CIFAR-10 training set stays resident on the GPU as float32. A PCA spectrum over 8000 images measures where the pixels concentrate. The notebook then tests the trained model as a single-pass denoiser on the held-out test set.
