nanograd-components

Model components for the nanograd machine learning framework: implementations of ResNet, MLP-Mixer and LLAMA-style Transformer.

Description

Each model comes in two forms:

a strict version on nanograd's eagerly evaluated tensors, which computes values as it goes and is trained with backward!;
a lazy version based on array-morphisms tensors. Its forward pass builds a graph. am-training-step/ssa compiles the forward pass, the backward pass and the parameter gradients of a whole training step into one SSA (static single assignment) program, then replays that program on each new batch.

The lazy models compute the same functions as their strict counterparts and keep their parameters in the same order. am-copy-parameters! can therefore move weights from a strict model to a lazy one.

Model	Strict module	Lazy module	Detailed notes
Transformer	`nanograd-transformer`	`nanograd-am-transformer`	transformer_readme.md
ResNet	`nanograd-resnet`	`nanograd-am-resnet`	resnet_readme.md
MLP-Mixer	`nanograd-mlp-mixer`	`nanograd-am-mlp-mixer`	mlp_mixer_readme.md

The module nanograd-am-common provides the parts shared by the lazy models.

The detailed notes explain each model from the basics, for readers who know Scheme but are new to machine learning. They cover the architecture, API and design of the strict model, and end with a summary of the lazy version. The header comments of the lazy source files (am-transformer.scm, am-resnet.scm, am-mlp-mixer.scm and am-common.scm) describe the lazy models in more detail.

Installation

The lazy examples and the benchmarks in the repository also need the array-morphisms-crunch egg. It is not a build dependency.

Requirements

The test suite needs test.

Common conventions

All models and layers are yasos objects that follow nanograd's layer protocol: forward, parameters, zero-grad-layer!, set-training-mode!, set-eval-mode!, layer-name and layer?.

The dtype keyword selects the element type, 'f32 (the default) or 'f64. The name keyword sets the layer name returned by layer-name.

Strict models take nanograd tensors and are trained with backward! and an optimizer from nanograd-optimizer. Lazy models take lazy tensors and are trained with am-training-step/ssa, or run for inference with am-forward/ssa. Both procedures come from the nanograd-array-morphisms module of nanograd.

Transformer

The architecture follows LLAMA2:

a token embedding;
pre-norm blocks of RMSNorm, multi-head self-attention with rotary position embedding (RoPE), and a SwiGLU feed-forward network;
a final RMSNorm;
an output projection tied to the embedding.

Unlike LLAMA2, the linear projections have bias vectors. Every operation has a backward pass, so the model can be trained.

nanograd-transformer

The strict model processes one sequence at a time as a (seq_len, d_model) matrix. A KV cache supports incremental generation.

make-transformer vocab-size max-seq-len d-model n-layers n-heads #!key (d-ff #f) (use-rope? #t) (rope-base 10000.0) (causal? #t) (dtype 'f32) (name Transformer)procedure

Creates the complete model. d-ff is the hidden size of the feed-forward network and defaults to (* 4 d-model). With causal?, each token can attend only to itself and earlier tokens.

(forward model tokens) takes a list or a tensor of token ids and returns the logits, of shape (seq_len, vocab_size). (forward model tokens kv-cache: cache) appends the tokens after those already cached and returns the logits of the last token only, of shape (vocab_size).

The parameters are in this order: the token embedding, the final norm weight, then for each block its two norm weights, its attention parameters (WQ, WK, WV, WO, each a weight and a bias) and its feed-forward parameters (W1, W2, W3, each a weight and a bias).

make-kv-cache n-layers n-heads max-seq-len head-size #!key (dtype 'f32)procedure
kv-cache-length cacheprocedure
clear-cache! cacheprocedure: make-kv-cache creates a cache for the keys and values of one text. kv-cache-length returns the number of cached positions, and clear-cache! empties the cache. The cache stores keys and values as constants, so use it for generation only, not for training.

make-transformer-block d-model n-heads d-ff #!key (rope #f) (causal? #f) (dtype 'f32) (name TransformerBlock)procedure
make-multihead-attention d-model n-heads #!key (rope #f) (causal? #f) (dtype 'f32) (name MultiHeadAttention)procedure
make-ffn-block d-model d-ff #!key (dtype 'f32) (name FFN)procedure: Building blocks: one pre-norm block, multi-head self-attention, and the SwiGLU feed-forward network. Note that the block and attention constructors default to causal?: #f.

make-rope-embedding max-seq-len head-size #!key (base 10000.0) (dtype 'f32)procedure
rope-rotate-rows x rope start-posprocedure
apply-rope-rotation Q K rope posprocedure: make-rope-embedding precomputes the RoPE tables. rope-rotate-rows rotates all rows of a (seq_len, d) matrix whose first row is at position start-pos. apply-rope-rotation rotates a single query and key of shape (n_heads, head_size).

multihead-attention-core Q K V n-heads #!key (causal? #f)procedure

scaled-dot-product-attention Q K V head-size seq-len #!key (causal? #f)procedure

Attention without the input and output projections: multihead-attention-core works on plain matrices, and scaled-dot-product-attention on (len, n_heads, head_size) tensors.

The module also exports the helpers causal-mask, embedding-lookup, slice-columns, concat-columns, get-frequencies, rope-head-size, cache-kv-pair! and get-cached-kv, and the predicates transformer?, transformer-block?, multihead-attention?, ffn-block?, rope-embedding? and kv-cache?.

nanograd-am-transformer

The lazy model processes a batch of B sequences of length T as one [B*T, vocab] matrix of one-hot rows. A constant matrix product selects each attention head, and a block-diagonal mask keeps the sequences apart. The lazy model has no KV cache.

make-am-transformer vocab-size seq-len d-model n-layers n-heads #!key (d-ff #f) (use-rope? #t) (rope-base 10000.0) (causal? #t) (norm-eps 1e-05) (dtype 'f32) (name AM-Transformer)procedure: Creates the lazy model. Unlike max-seq-len in the strict model, seq-len is the fixed sequence length T.

am-tokens->one-hot sequences vocab-size dtypeprocedure: Returns a one-hot typed vector [B*T, vocab-size] for a list of B token lists of equal length T. Row b*T + t holds token t of sequence b.

am-next-token-targets sequences vocab-size dtypeprocedure

For sequences of length T+1, returns a list (inputs targets) of one-hot typed vectors: the inputs are the first T tokens of each sequence, and the targets the last T.

The module also exports the building blocks make-am-attention, make-am-swiglu-ffn and make-am-transformer-block, and the predicate am-transformer?.

ResNet

nanograd-resnet

The strict models accept a single image (C, H, W) or a batch (N, C, H, W). Batch normalization keeps running statistics for evaluation mode.

make-resnet-18 num-classes #!key (dtype 'f32)procedure
make-resnet-34 num-classes #!key (dtype 'f32)procedure
make-resnet-50 num-classes #!key (dtype 'f32)procedure
make-resnet-101 num-classes #!key (dtype 'f32)procedure
make-resnet-152 num-classes #!key (dtype 'f32)procedure: The ImageNet-style models: a 7x7 stem with max pooling, then four layers of basic blocks (18 and 34) or bottleneck blocks (50, 101 and 152). They reject inputs smaller than 32x32.

make-resnet-cifar stages widths num-classes #!key (in-channels 3) (dtype 'f32) (name ResNet-CIFAR)procedure: A ResNet for small images, such as the 32x32 images of CIFAR-10. It has a 3x3 stem and no max pooling, followed by one layer of basic blocks per entry of stages, with (list-ref widths i) channels. Every layer after the first starts with stride 2.

make-resnet block-fn layers num-classes #!key (dtype 'f32) (name ResNet)procedure

make-residual-layer block-fn in-channels out-channels num-blocks #!key (stride 1) (dtype 'f32) (name Layer)procedure

make-residual-block in-channels out-channels #!key (stride 1) (downsample #f) (dtype 'f32) (name BasicBlock)procedure

make-bottleneck-block in-channels out-channels #!key (stride 1) (downsample #f) (dtype 'f32) (name Bottleneck)procedure

make-downsample in-channels out-channels stride dtypeprocedure

Building blocks. make-resnet builds an ImageNet-style model from a block constructor and a list of four block counts. make-residual-layer gives the first block the stride, and also a downsample shortcut when the stride or the channel count changes. make-downsample is a 1x1 convolution followed by batch normalization; its dtype argument is positional.

The predicates are resnet?, residual-block? and bottleneck-block?.

nanograd-am-resnet

The lazy models take NHWC input [N, H, W, C]. With the default 'cifar stem, they compute the same function as make-resnet-cifar in training mode. Batch normalization always uses the statistics of the current batch, so set-training-mode! and set-eval-mode! have no effect. There is no max pooling, because the SSA engine has no reduce-max gradient.

make-am-resnet stages widths num-classes #!key (in-channels 3) (stem 'cifar) (dtype 'f32) (name AM-ResNet)procedure: Creates the lazy model. stem: 'strided gives the 3x3 stem stride 2 instead of 1.

make-am-resnet-8 num-classes #!key (dtype 'f32)procedure: Stages (1 1 1) with widths 16, 32 and 64.

make-am-resnet-18-narrow num-classes #!key (dtype 'f32)procedure: Stages (2 2 2 2) with widths 16, 32, 64 and 128.

am-nchw->nhwc data n c h w dtypeprocedure

Copies an NCHW typed vector into a new NHWC typed vector.

The module also exports the building blocks make-am-residual-block, make-am-residual-stage and make-am-downsample (whose dtype is a keyword), and the predicate am-resnet?.

MLP-Mixer

An MLP-Mixer is a stack of blocks, each a token-mixing MLP and a channel-mixing MLP, with pre-normalization (RMSNorm or LayerNorm) and skip connections. It is intended for small fixed grids, such as those of the Tiny Recursive Model (TRM). The input is a (seq_len, hidden_dim) matrix: one row per token, one column per feature.

nanograd-mlp-mixer

make-mlp-mixer seq-len hidden-dim n-blocks #!key (token-expansion 4) (channel-expansion 4) (use-rmsnorm #t) (dtype 'f32) nameprocedure: Creates a model of n-blocks blocks for one (seq-len, hidden-dim) sample. The output has the same shape as the input. The token MLP has (* token-expansion seq-len) hidden units, and the channel MLP (* channel-expansion hidden-dim). With use-rmsnorm: #f, the blocks use LayerNorm.

get-blocks modelprocedure
num-blocks modelprocedure
mlp-mixer? objprocedure: The list of blocks of a model, their number, and the model predicate.

make-mlp-mixer-block seq-len hidden-dim #!key (token-expansion 4) (channel-expansion 4) (use-rmsnorm #t) (dtype 'f32) nameprocedure
make-token-mixing-mlp seq-len hidden-dim #!key (expansion-factor 4) (dtype 'f32) (name TokenMixing)procedure
make-channel-mixing-mlp hidden-dim #!key (expansion-factor 4) (dtype 'f32) (name ChannelMixing)procedure
make-mixer-mlp input-dim hidden-dim #!key (dtype 'f32) (name MixerMLP)procedure: Building blocks: one block, the two kinds of mixing MLP, and the underlying two-layer MLP with GELU.

make-rmsnorm-layer normalized-shape #!key (eps 1e-06) (dtype 'f32) (name RMSNorm)procedure
make-layer-norm normalized-shape #!key (eps 1e-06) (dtype 'f32) (name LayerNorm)procedure: Normalization over the last axis. normalized-shape is the list (hidden-dim).

nanograd-am-mlp-mixer

The lazy model works on a whole batch at once. B samples form one token-major [T*B, D] matrix, in which row t*B + b holds token t of sample b.

make-am-mlp-mixer seq-len hidden-dim n-blocks #!key (token-expansion 4) (channel-expansion 4) (use-rmsnorm #t) (norm-eps 1e-06) (num-classes #f) (dtype 'f32) (name AM-MLPMixer)procedure: Creates the lazy model. Without num-classes, the output is the [T*B, D] feature matrix, as in the strict model. With num-classes, the model averages the features over the tokens, and a dense head gives [B, num-classes] logits. The head's parameters come last.

am-mixer-pack-batch samples seq-len hidden-dim dtypeprocedure

am-mixer-unpack-batch data batch seq-len hidden-dim dtypeprocedure

Convert between a list of row-major (T, D) vectors, one per sample, and one token-major [T*B, D] typed vector.

The module also exports the building blocks make-am-token-mixing-mlp, make-am-channel-mixing-mlp and make-am-mlp-mixer-block, and the predicate am-mlp-mixer?.

nanograd-am-common

This module holds the parts shared by the lazy models.

am-copy-parameters! src dstprocedure: Copies the parameter values of model src into model dst, element by element in row-major order. The two parameter lists must have the same length and matching element counts; the shapes may differ. For example, a strict convolution weight [out, in, kh, kw] matches a lazy one [out, in*kh*kw]. The values are copied, not shared.

am-parameter-count modelprocedure
am-parameter-shapes modelprocedure: The total number of parameter elements of model, and the list of parameter shapes.

am-half-mse-loss pred targetprocedure: 0.5 * mean((pred - target)^2), the convention of nanograd's strict mse-loss.

make-am-rmsnorm dim #!key (eps 1e-06) (dtype 'f32) (name AM-RMSNorm)procedure
make-am-layernorm dim #!key (eps 1e-06) (dtype 'f32) (name AM-LayerNorm)procedure
make-am-batch-norm-nhwc channels #!key (eps 1e-05) (dtype 'f32) (name AM-BN)procedure: Normalization layers. The batch normalization works on NHWC input and always uses the statistics of the current batch.

am-softmax-rows sprocedure
am-rmsnorm-rows x g epsprocedure
am-layernorm-rows x g b epsprocedure: Row-wise softmax, RMSNorm and LayerNorm of an [M, D] morph-variable.

am-param data shape dtypeprocedure
am-param-fill shape value dtypeprocedure
am-param-uniform shape scale dtypeprocedure
am-param-normal shape std dtypeprocedure: Constructors of trainable lazy tensors: from a typed vector, filled with one value, uniform on [-scale/2, scale/2), or normal with standard deviation std.

am-const data shape dtypeprocedure
am-const-fill shape dtype fprocedure
am-scalar value dtypeprocedure: Constructors of non-trainable morph-variables. Element i of am-const-fill is (f i). am-scalar has shape (1) and broadcasts against any shape.

mv lazy-tensorprocedure
lazy morph-variableprocedure
am-random-normalprocedure: mv returns the morph-variable of a lazy tensor, and lazy the lazy tensor of a morph-variable. am-random-normal returns a sample of the standard normal distribution.

Examples

Strict Transformer

(import (srfi 4) nanograd-autograd nanograd-layer nanograd-optimizer
        nanograd-transformer)

;; vocab 28, context 32, d_model 64, 2 layers, 4 heads
(define model (make-transformer 28 32 64 2 4 dtype: 'f32))
(define opt (make-adam (parameters model) learning-rate: 0.005))

(define (train-step! inputs targets)        ; lists of token ids
  (let ((loss (cross-entropy-loss
               (forward model inputs)       ; logits (seq_len, vocab)
               (make-tensor32 (list->f32vector (map exact->inexact targets))
                              (list (length targets)) requires-grad?: #f)
               from-logits: #t)))
    (backward! loss)
    (step! opt)
    (zero-grad-layer! model)
    loss))

Lazy ResNet with SSA training

(import (srfi 4) yasos array-morphisms-core array-morphisms-context
        (prefix array-morphisms-grad am:)
        nanograd-autograd nanograd-layer nanograd-optimizer
        nanograd-array-morphisms nanograd-am-common nanograd-am-resnet)

(define (vector->lazy data shape)
  (get-or-make-lazy (am:make-var (make-morphism data shape 'f32) #f)))

(define model (make-am-resnet '(1 1 1) '(16 32 64) 10 dtype: 'f32))
(define opt (make-adam (parameters model) learning-rate: 0.001))
(define ctx (make-morphism-context))      ; one compiled program per context

;; x: NHWC f32vector [N, H, W, 3]; t: one-hot f32vector [N, 10]
(define (train-step! x t n)
  (let ((x-lt (vector->lazy x (list n 32 32 3)))
        (t-lt (vector->lazy t (list n 10))))
    ;; The first call compiles the step; later calls replay it.
    (am-training-step/ssa ctx opt model
                          (lambda (logits) (am-cross-entropy-loss logits t-lt))
                          x-lt t-lt)))

A compiled program is fixed to the shapes of its first call. Use one context per batch shape, and a separate context with am-forward/ssa for inference.

Example scripts

The examples directory of the repository holds complete scripts, run from the repository root with csi -s:

transformer_example.scm: a character-level language model with KV-cache generation, a copy task, training with and without RoPE, KV-cache speed and consistency, and an unstable SGD run with and without gradient clipping.
am_transformer_example.scm: the same tasks with the lazy Transformer, and a check that it agrees with the strict model.
resnet_example.scm: classification of synthetic images, a learning-rate schedule, parameter counts, output statistics and transfer learning.
am_resnet_example.scm: the same with the lazy ResNet, agreement with the strict model, and compiled-inference timing.
mlp-mixer-examples.scm: a forward pass, TRM-style recursive refinement, deep supervision and a training loop.

The training loops monitor gradient norms with the diagnostics of nanograd-diagnostics, through gradient-monitoring.scm.

The lazy examples choose their backend with the AM_EXAMPLE_BACKEND environment variable:

crunch: the crunch kernels (the default); needs array-morphisms-crunch.
blas: the blas egg.
scheme: no backend, only the Scheme code of array-morphisms.

OPENBLAS_NUM_THREADS=8 AM_EXAMPLE_BACKEND=blas csi -s examples/am_resnet_example.scm

Benchmarks

The scripts examples/bench-am-resnet-crunch.scm, examples/bench-am-mlp-mixer-crunch.scm and examples/bench-am-transformer-crunch.scm time one SSA training step of each lazy model against its strict counterpart. Compile them with csc -O3 and run them with OPENBLAS_NUM_THREADS=1. Environment variables, described in examples/bench-am-common.scm, select the backend, configuration and number of steps.

Tests

From the repository root:

csi -s tests/run.scm

The runner runs each test file in its own process. To run only some files, name them, as in csi -s tests/run.scm test-gradcheck.

test-gradcheck: finite-difference gradient checks of the three strict models, and a check of the KV cache.
test-am-models: each lazy model against its strict model. After the weights are copied, each model takes one SGD step on the same batch; the losses and gradients must agree to 1e-12 in f64.
test-transformer, test-resnet, test-resnet-numerical, test-mlp-mixer: unit tests of the strict components.

chickadee » nanograd-components

Identifier search