Course project for HSE/YSDA Lab 4 on diffusion models (2025). The repository contains continuous-time DDPM, Neural Diffusion Models (NDM), and Neural Flow Diffusion Models (NFDM) implementations in PyTorch. It was the final codebase used for the course presentation.
- Alexander Nekrasov
- Artyom Shulyatev
- Elena Grishina
- Nina Krivulets
- DDPM, NDM, and NFDM: Three continuous-time diffusion formulations in one training pipeline
- Trajectory control: Curvature regularization for straighter NFDM trajectories
- Evaluation: NLL estimation with the Hutchinson trace estimator and FID for ODE/SDE sampling
- Multiple architectures: MLP (2D toys), SimpleMLP, CUNet (images)
- Multiple datasets: Checkerboard, moons, circles, swiss roll, MNIST, Fashion-MNIST
- Hydra config: Easy experiment management with CLI overrides
- WandB integration: Optional experiment tracking with visualizations
# Using pixi (recommended)
pixi install
# Or using pip
pip install torch torchvision hydra-core omegaconf wandb scikit-learn tqdm matplotlibFor CUNet architecture, clone the EDM repo:
git clone https://github.com/NVlabs/edm# Train DDPM on checkerboard (default)
python train.py
# Train on MNIST with CUNet
python train.py dataset=mnist architecture=cunet
# Train NDM model
python train.py model=ndm
# Train NFDM model
python train.py model=nfdm
# Enable WandB logging
python train.py logging.wandb=true.
├── train.py # Main training script (Hydra-based)
├── diffusion_model.py # DDPM, NDM, and NFDM implementations
├── metrics.py # NLL estimation
├── compute_fid.py # FID evaluation
├── visualization.py # Dataset-specific visualizations
├── visualize_2d.py # 2D samples and trajectory plots
├── models/
│ ├── mlp_2d.py # MLP for 2D toy datasets
│ ├── simple_mlp.py # Configurable simple MLP
│ ├── cunet.py # Conditional UNet for images
│ └── edm.py # EDM integration wrapper
├── datasets/
│ ├── toy.py # 2D datasets (checkerboard, moons, etc.)
│ └── images.py # Image datasets (MNIST, Fashion-MNIST)
└── configs/
├── config.yaml # Main config with defaults
├── model/ # ddpm.yaml, ndm.yaml, nfdm.yaml
├── architecture/ # mlp_2d.yaml, simple_mlp.yaml, cunet.yaml
├── dataset/ # checkerboard.yaml, mnist.yaml, etc.
└── scheduler/ # cosine.yaml, linear.yaml, step.yaml
# Change hyperparameters
python train.py training.batch_size=512 training.num_epochs=500
# Change optimizer
python train.py optimizer.lr=1e-3 optimizer.name=adam
# Change LR scheduler
python train.py scheduler=linear
# Change sampling settings
python train.py eval.num_steps=200 eval.use_sde=false| Category | Options |
|---|---|
| Models | ddpm, ndm, nfdm |
| Architectures | mlp_2d, simple_mlp, cunet |
| Datasets | checkerboard, moons, circles, swiss_roll, mnist, fashion_mnist |
| Schedulers | cosine, linear, step |
| Optimizers | adamw, adam, sgd |
The final course report used 2,000 generated samples for FID. These values are included to document the submitted experiments, not as a benchmark claim.
| Sampling setup | DDPM | NFDM | NFDM + curvature |
|---|---|---|---|
| FID, ODE, 10 steps | 116.309 | 111.078 | 124.102 |
| FID, ODE, 100 steps | 28.779 | 82.227 | 23.847 |
| FID, SDE, 100 steps | 94.971 | 37.248 | 38.223 |
| FID, SDE, 300 steps | 46.723 | 32.082 | 35.837 |
The report also recorded lower NFDM NLL estimates than DDPM at 100, 300, and 1,000 Euler steps. Because the objectives and discretization behavior differ, the results should be interpreted together with the implementation details.
| Before curvature fine-tuning | After curvature fine-tuning |
|---|---|
 |
 |
The same model family was also evaluated on MNIST with 10 sampling steps:
All authors collaborated across the project. The final course report recorded the main areas of responsibility as follows:
- Model derivations, implementation debugging, and NLL evaluation: Alexander Nekrasov and Artyom Shulyatev
- Experiment infrastructure: Elena Grishina and Nina Krivulets
- Integration of existing architectures into the shared pipeline: Elena Grishina
- FID evaluation: Nina Krivulets
Alexander primarily worked on the mathematical formulation, the 2D/toy-data experiments, and debugging the model implementations. Artyom also developed part of the image pipeline.
The submitted 15-page report documents the mathematical derivations, implementation, experiments, metrics, and contribution split:
The report uses an ICLR-style Typst layout for formatting only. It is a course report, not an ICLR submission or publication; the template's publication header was disabled in this public copy.
Training outputs are saved to outputs/YYYY-MM-DD/HH-MM-SS/:
best_model.pt- Best model checkpointcheckpoint_epoch_N.pt- Periodic checkpointssamples_epoch_N.pt- Generated samplestrain.log- Hydra logs
- Neural Diffusion Models
- Neural Flow Diffusion Models
- Denoising Diffusion Probabilistic Models
- EDM: Elucidating the Design Space of Diffusion-Based Generative Models
- Tiny Diffusion (the 2D MLP was adapted for this project)
The conditional U-Net used for image experiments was adapted from course materials. This cleaned fork intentionally excludes the separate scratch notebook and derivative NFDM reference script used during development.