Diffusion Model for Lorenz Attractor

A PyTorch implementation of a DDPM (Denoising Diffusion Probabilistic Model) according to the original paper. It is implemented and trained from scratch. In this example, I applied it to learn the spatial distribution of the Lorenz Attractor with specific parameters.

Generated Samples

Evolution of the diffusion process from pure gaussian noise to the butterfly structure.

Model architecture

• MLP-based 512-dim hidden layers with SiLU activations and LayerNorm
• Sinusoidal position embeddings (128-dim)
• 200-step DDPM with linear noise schedule

Installation

pip install -r requirements.txt

Training

Run the training script:

python src/train.py

This will:

• Generate Lorenz attractor data via ODE integration
• Train the diffusion model for 300 epochs. I found that while the loss doesn't decrease much after a while, it is still useful to train longer to better capture the data distribution.
• Save the trained model to models/model.pt

Training Progress

Generating Samples

Generate samples using the trained model:

python src/generate.py

This will:

• Load the trained model from models/model.pt
• Generate novel points via 200-step denoising process
• Create an animated GIF showing the evolution from noise to structure
• Save visualizations to media/

References

Ho, J., Jain, A., & Abbeel, P. (2020). Denoising Diffusion Probabilistic Models.