A character-level language diffusion model for text generation. The model is a modified version of the nanochat gpt implementation and is trained on Tiny Shakespeare! It is only 10.7 million parameters, so you can try it out locally! Installation # Clone the repository git clone < repository-url > cd tiny-diffusion # Install dependencies (Python 3.10+) uv sync Quick Start The file training.py puts the weights in weights/diffusion_model.pt . The sample and animation files load the model from this file. Train Your Own Model Currently, the weights are already provided for you! It took me around half an hour to train this model for 20,000 steps on 4xA100s. But if you want to retrain the model again, run: # Train from scratch on Shakespeare uv run training.py # Training will save checkpoints to weights/diffusion_model.pt Generate Text To generate a continuous stream of output (currently 30 context lengths), run: # Generate samples using the pre-trained model uv run sample.py Visualize the Diffusion Process To see the diffusion process as a nice animation, run: # Watch the denoising process step-by-step uv run animations/diffusion-process.py # See Game of Life-inspired sampling (fun little experiment) uv run animations/game-of-life.py Default Config Parameters : 10.7 million : 10.7 million Layers : 6 : 6 Attention Heads : 6 : 6 Embedding Dim : 384 : 384 Sequence Length : 256 characters : 256 characters Diffusion Steps: 128 File Structure
First seen: 2025-11-14 19:52
Last seen: 2025-11-15 11:55