tinygrad-rust

A small educational Rust project that implements a minimal neural-network training pipeline inspired by tinygrad-style autodiff concepts. The repository focuses on clear, hand-written tensor and matrix operations rather than high-performance abstractions.

What this project does

The crate currently trains a simple multilayer perceptron on the MNIST dataset with:

• a custom Matrix implementation for dense linear algebra
• a basic Linear layer
• ReLU and Softmax activations
• cross-entropy loss
• a dynamic computation graph with reverse-mode vector-Jacobian products
• MNIST data loading, train/validation splitting, and shuffling
• progress logging to stdout and a timestamped log file
• validation accuracy reporting after each epoch

Repository structure

• src/main.rs — binary entry point that loads MNIST data, trains the model, and logs progress
• src/ops.rs — layers and loss functions such as Linear, ReLU, Softmax, and CrossEntropy
• src/linear_algebra.rs — matrix operations used throughout the training code
• src/mnist.rs — MNIST dataset loading, splitting, and shuffling helpers
• src/lib.rs — module exports for the library crate

Requirements

• Rust toolchain (stable)
• Cargo
• The MNIST training data files in the data directory:
  • train-images-idx3-ubyte
  • train-labels-idx1-ubyte

The dataset files are expected to be placed under the repository's data/ folder.

Running the training code

From the repository root:

cargo run --release --bin tinygrad-rust -- --epochs 10

This will:

1. load MNIST data
2. split it into training and validation sets
3. train a small MLP for the requested number of epochs
4. print progress lines to the terminal
5. write the same progress lines to a timestamped log file named like training-YYYYMMDD-HHMM.log

The training binary accepts these command-line flags:

• --epochs <N> to set the number of training epochs, defaulting to 10
• --max-batches <N> to stop each epoch after a fixed number of batches
• --skip-validation to disable validation accuracy reporting after each epoch
• --optimizer <adamw|sgd> to select the optimizer, defaulting to adamw

Logging and metrics

During training, the program logs entries such as:

elapsed=12.34s epoch=1/10 batch=0 loss=2.345678
elapsed=24.56s epoch=1/10 validation_accuracy=12.34%

The log file is created in the repository root and flushed after each logged line so interrupted runs still preserve the latest progress.

Dynamic graph

Each training batch creates a Graph, adds input and target values, and applies operations through Graph::apply. The graph owns intermediate matrices and gradients for that batch; persistent layers own only their parameters. Calling graph.backward(loss) walks the recorded operations in reverse and invokes each operation's vector-Jacobian product. The graph is dropped before parameter updates, then recreated for the next batch.

Operations implement the public Operation trait, so custom operations can participate in the graph by providing compute and vjp methods. Validation calls Operation::compute directly and therefore does not record a graph.

For a short local smoke run, use:

cargo run --bin tinygrad-rust -- --epochs 1 --max-batches 1 --skip-validation

Training on Tiny Shakespeare

Download the character-level corpus:

curl -L https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt \
  -o data/tinyshakespeare.txt

Train the optimized sub-million-parameter Llama-2-style transformer:

cargo run --release --bin train-tiny-shakespeare

Both training binaries use AdamW by default. Select SGD explicitly for comparison:

cargo run --release --bin train-tiny-shakespeare -- --optimizer sgd

The defaults select the best measured training-loss configuration: 8 layers, hidden size 48, intermediate size 96, 12 heads, context length 16, batch size 1, learning rate 0.15, and 640,000 steps. It has 194,897 parameters and completed in 535.59 seconds on the test machine. The binary rejects models at or above one million parameters by default.

Training statistics include both noisy current-batch loss and loss on fixed, evenly spaced training windows. The latter makes runs comparable at the same batch_size * steps sample budget. Held-out loss, prompts, and generated completions are also written to both stdout and tiny-shakespeare-training-YYYYMMDD-HHMMSS.log. Run with --help for all options and see EXPERIMENTS.md for the search methodology and results.

Visualizing a training log

Use the Rust visualizer to generate an SVG plot of training loss and validation accuracy:

cargo run --bin visualize-training-log -- training-YYYYMMDD-HHMM.log

By default, the SVG is written next to the log with the same name and an .svg extension. Use --output to choose another destination:

cargo run --bin visualize-training-log -- training-YYYYMMDD-HHMM.log --output training.svg

Compare two optimizer runs in one plot:

cargo run --release --bin visualize-training-log -- training-sgd.log --compare training-adamw.log --output figures/mnist-sgd-adamw.svg

MNIST Results

The latest 10-epoch AdamW run reached 97.77% validation accuracy in 10.35 seconds. Its validation accuracy rose from 94.92% after the first epoch to 97.77% after the final epoch.

Optimizer	First-epoch validation accuracy	Final validation accuracy	Recorded duration
AdamW	94.92%	97.77%	10.35 seconds
SGD	39.87%	84.63%	9.28 seconds

AdamW improved final validation accuracy by 13.14 percentage points over the fresh SGD run. The durations are single-run measurements and should not be treated as a controlled speed comparison.

The combined plot uses solid lines for training loss and dashed lines for validation accuracy. Blue denotes SGD and red denotes AdamW.

Notes

This repository is intentionally simple and educational. It is not intended to be a production-grade deep learning framework, and the implementation uses explicit loops and matrix-level derivative rules for clarity.

Dataset source

• MNIST dataset
• Tiny Shakespeare dataset