Reference implementations of the algorithms behind modern machine learning, written from scratch and explained. Each page derives the math, implements the algorithm twice, in PyTorch and in JAX behind a language switcher so the two idioms can be read against each other, then covers where the algorithm earns its keep in practice and how the from-scratch version compares with the production libraries (scikit-learn, FAISS, XGBoost, torchvision, diffusers, and friends). GPU-bound submodules additionally carry Triton and CUDA kernels beside the reference, as the softmax page shows.
Featured pageScaled dot-product attention derived as a differentiable dictionary lookup, multi-head attention, causal masking, and a complete transformer block, verified against the fused kernels, with the variants that power modern LLMs covered in their own section below.
Read the page →The pre-deep-learning toolkit, still the right answer surprisingly often. Each implemented in PyTorch and JAX and measured against scikit-learn and its cousins.
Least squares from Gaussian maximum likelihood, the normal equations and why you solve them with QR rather than matrix inversion, gradient descent as the scalable alternative, and ridge regularization, compared against scikit-learn and statsmodels.
Bernoulli likelihood to cross-entropy, the full gradient derivation, convexity, and multinomial extension, with full-batch and minibatch training in both frameworks and the scikit-learn solver landscape (lbfgs, liblinear, saga) mapped out.
Lloyd's algorithm as alternating minimization of the distortion objective, k-means++ initialization and its approximation guarantee, vectorized assignment in both frameworks, and where FAISS's GPU k-means takes over from scikit-learn.
The bias-variance role of k, the curse of dimensionality stated concretely, brute-force vectorized top-k in both frameworks, and the escalation path from scikit-learn's KD-trees to approximate search with FAISS.
The max-margin objective, hinge loss as its unconstrained form, the kernel trick worked on a small example, and sub-gradient training in both frameworks, with the LIBSVM and LIBLINEAR heritage behind scikit-learn's SVC explained.
Bayes rule under the conditional-independence assumption, multinomial and Gaussian variants with Laplace smoothing and log-space arithmetic, a worked spam example, and why it ranks well while calibrating badly.
The mixture likelihood, expectation-maximization derived through responsibilities, log-sum-exp stability, k-means as the hard-assignment limit, and scikit-learn's covariance options and BIC model selection.
Variance maximization and reconstruction error shown to be the same objective, PCA via the SVD of centered data, explained-variance ratios, power iteration for the top component, and when UMAP or t-SNE is the right tool instead.
Recursive partitioning with Gini and entropy worked numerically, greedy split search vectorized over tensors, pruning and depth control, and the honest lesson that trees resist autodiff frameworks, which is exactly why the boosting libraries exist.
Boosting derived as gradient descent in function space, residual fitting, shrinkage and subsampling, second-order boosting as XGBoost's contribution, and the XGBoost, LightGBM, and CatBoost landscape where tabular ML is still won.
From backpropagation written by hand up through the submodules transformers are assembled from. The GPU-bound pages add Triton and CUDA kernels beside the reference.
The chain rule through a hidden layer written out fully, manual forward and backward passes with no autograd in either framework, then the idiomatic nn.Module and JAX pytree versions, and why hand-rolling backprop once makes every framework make sense.
Numerical stability via the subtract-max invariance, the online single-pass variant that FlashAttention is built on, and the backward pass from the Jacobian, implemented in NumPy, PyTorch, Triton, and CUDA.
Convolution as a sliding dot product with weight sharing, output-size arithmetic worked numerically, an im2col implementation from scratch, and a LeNet-style classifier in both frameworks, with torchvision and timm as the production model zoos.
Backpropagation through time, vanishing gradients derived from the repeated Jacobian product, the LSTM cell equations in full, a from-scratch cell in PyTorch and a lax.scan version in JAX where recurrence is naturally idiomatic, and where recurrence still wins today.
The mechanism at the center of modern ML, from the core derivation through the variants every production model ships, with what each one is actually useful for.
Scaled dot-product attention as a soft dictionary lookup, the sqrt(d) scaling, causal masking, multi-head as parallel subspace attention, and one complete transformer block, verified against PyTorch's fused scaled_dot_product_attention. The starting point for everything below.
The KV-cache-efficiency family: the cache-size arithmetic that motivates them, multi-query attention sharing one KV head, grouped-query attention as the interpolation Llama and Mistral ship, and DeepSeek's multi-head latent attention compressing KV into a low-rank latent. One configurable implementation spans MHA to MQA in both frameworks.
The quadratic problem at 128k context, sliding-window and sparse patterns, linear attention's reassociation trick and its honest quality trade-offs, and FlashAttention as exact attention made fast by IO-awareness rather than approximation, with a comparison table of who ships what and why.
Queries from one sequence, keys and values from another: the encoder-decoder original, Stable Diffusion's text conditioning inside the U-Net, and Perceiver-style latent bottlenecks, with a minimal encoder-decoder block in both frameworks.
The model shapes that keep reappearing, each implemented completely and tied to the production codebases that ship them.
The degradation problem and how identity shortcuts solve it, gradient flow through the skip connection, a ResNet-18 in both frameworks with the downsampling subtlety handled, and loading torchvision's pretrained weights into your own implementation as the correctness proof.
The encoder-decoder with skip connections that carry the detail pooling destroys, segmentation losses under class imbalance, a complete readable implementation in both frameworks, and its second life as the denoising backbone inside diffusion models, compared against segmentation-models-pytorch, MONAI, and diffusers.
The ELBO derived step by step, the reparameterization trick and why sampling blocks gradients without it, the closed-form Gaussian KL, JAX's explicit randomness as a feature, and the VAE living inside latent diffusion as Stable Diffusion's autoencoder.
The forward noising process in closed form, the epsilon-prediction objective as a weighted ELBO, beta schedules, the full DDPM sampling loop with a runnable 2-D toy example, and the diffusers library as the production home of the same equations.
What happens between the forward pass and the updated weights, from-scratch optimizers verified step-for-step against torch.optim and optax.
Gradient descent's condition-number convergence rate, stochasticity as noise and regularizer, heavy-ball momentum as an EMA of gradients, and Nesterov's lookahead, with from-scratch optimizers matching torch.optim.SGD and optax.sgd exactly.
Adaptive learning rates from Adagrad through RMSProp to Adam, bias correction derived, and the Loshchilov-Hutter argument for why L2 in the gradient is not weight decay under adaptive methods, verified step-for-step against torch.optim.AdamW and optax.adamw.
Why warmup exists, cosine and warmup-stable-decay shapes with an ASCII plot, the standard LLM recipe, and global-norm gradient clipping implemented and verified against clip_grad_norm_.
Float formats compared bit by bit, why fp16 needs loss scaling and bf16 does not, master weights, which ops stay in fp32, and a manual training step showing exactly what torch.autocast and GradScaler do.
Why tokenization quietly shapes model quality and cost, the BPE, WordPiece, and Unigram families with an original worked merge, why naive encoding is slow and the regex stage dominates, and gigatoken as a case study in making BPE run at gigabytes per second, with the project's reported numbers checked against a benchmark run on this machine and a zero-mismatch parity test against tiktoken and HF tokenizers.
How real projects each use PyTorch their own way, nanoGPT's flat readable loop, torchtitan composing DTensor and FSDP2, Unsloth rewriting the hot path with Triton and manual backward passes, vLLM keeping the model but replacing the execution loop, and diffusers splitting models, schedulers, and pipelines into recombinable parts, closing with a guide to which idiom fits which situation.