You can edit this page using wiki syntax for markup.

Article contents:== nanograd-components Model components for the [[nanograd]] machine learning framework: implementations of ResNet, MLP-Mixer and LLAMA-style Transformer. [[toc:]] === 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. <table> <tr><th>Model</th><th>Strict module</th><th>Lazy module</th><th>Detailed notes</th></tr> <tr><td>Transformer</td><td>{{nanograd-transformer}}</td><td>{{nanograd-am-transformer}}</td><td>[[https://github.com/iraikov/nanograd-components/blob/main/notes/transformer_readme.md|transformer_readme.md]]</td></tr> <tr><td>ResNet</td><td>{{nanograd-resnet}}</td><td>{{nanograd-am-resnet}}</td><td>[[https://github.com/iraikov/nanograd-components/blob/main/notes/resnet_readme.md|resnet_readme.md]]</td></tr> <tr><td>MLP-Mixer</td><td>{{nanograd-mlp-mixer}}</td><td>{{nanograd-am-mlp-mixer}}</td><td>[[https://github.com/iraikov/nanograd-components/blob/main/notes/mlp_mixer_readme.md|mlp_mixer_readme.md]]</td></tr> </table> 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 * [[srfi-1]] * [[srfi-4]] * [[yasos]] * [[blas]] * [[nanograd]] * [[array-morphisms]] 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. <procedure>(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). <procedure>(make-kv-cache n-layers n-heads max-seq-len head-size #!key (dtype 'f32))</procedure> <procedure>(kv-cache-length cache)</procedure> <procedure>(clear-cache! cache)</procedure> {{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. <procedure>(make-transformer-block d-model n-heads d-ff #!key (rope #f) (causal? #f) (dtype 'f32) (name "TransformerBlock"))</procedure> <procedure>(make-multihead-attention d-model n-heads #!key (rope #f) (causal? #f) (dtype 'f32) (name "MultiHeadAttention"))</procedure> <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}}. <procedure>(make-rope-embedding max-seq-len head-size #!key (base 10000.0) (dtype 'f32))</procedure> <procedure>(rope-rotate-rows x rope start-pos)</procedure> <procedure>(apply-rope-rotation Q K rope pos)</procedure> {{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)}}. <procedure>(multihead-attention-core Q K V n-heads #!key (causal? #f))</procedure> <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. <procedure>(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-5) (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. <procedure>(am-tokens->one-hot sequences vocab-size dtype)</procedure> 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. <procedure>(am-next-token-targets sequences vocab-size dtype)</procedure> 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. <procedure>(make-resnet-18 num-classes #!key (dtype 'f32))</procedure> <procedure>(make-resnet-34 num-classes #!key (dtype 'f32))</procedure> <procedure>(make-resnet-50 num-classes #!key (dtype 'f32))</procedure> <procedure>(make-resnet-101 num-classes #!key (dtype 'f32))</procedure> <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. <procedure>(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. <procedure>(make-resnet block-fn layers num-classes #!key (dtype 'f32) (name "ResNet"))</procedure> <procedure>(make-residual-layer block-fn in-channels out-channels num-blocks #!key (stride 1) (dtype 'f32) (name "Layer"))</procedure> <procedure>(make-residual-block in-channels out-channels #!key (stride 1) (downsample #f) (dtype 'f32) (name "BasicBlock"))</procedure> <procedure>(make-bottleneck-block in-channels out-channels #!key (stride 1) (downsample #f) (dtype 'f32) (name "Bottleneck"))</procedure> <procedure>(make-downsample in-channels out-channels stride dtype)</procedure> 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. <procedure>(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. <procedure>(make-am-resnet-8 num-classes #!key (dtype 'f32))</procedure> Stages {{(1 1 1)}} with widths 16, 32 and 64. <procedure>(make-am-resnet-18-narrow num-classes #!key (dtype 'f32))</procedure> Stages {{(2 2 2 2)}} with widths 16, 32, 64 and 128. <procedure>(am-nchw->nhwc data n c h w dtype)</procedure> 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 <procedure>(make-mlp-mixer seq-len hidden-dim n-blocks #!key (token-expansion 4) (channel-expansion 4) (use-rmsnorm #t) (dtype 'f32) name)</procedure> 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. <procedure>(get-blocks model)</procedure> <procedure>(num-blocks model)</procedure> <procedure>(mlp-mixer? obj)</procedure> The list of blocks of a model, their number, and the model predicate. <procedure>(make-mlp-mixer-block seq-len hidden-dim #!key (token-expansion 4) (channel-expansion 4) (use-rmsnorm #t) (dtype 'f32) name)</procedure> <procedure>(make-token-mixing-mlp seq-len hidden-dim #!key (expansion-factor 4) (dtype 'f32) (name "TokenMixing"))</procedure> <procedure>(make-channel-mixing-mlp hidden-dim #!key (expansion-factor 4) (dtype 'f32) (name "ChannelMixing"))</procedure> <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. <procedure>(make-rmsnorm-layer normalized-shape #!key (eps 1e-6) (dtype 'f32) (name "RMSNorm"))</procedure> <procedure>(make-layer-norm normalized-shape #!key (eps 1e-6) (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. <procedure>(make-am-mlp-mixer seq-len hidden-dim n-blocks #!key (token-expansion 4) (channel-expansion 4) (use-rmsnorm #t) (norm-eps 1e-6) (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. <procedure>(am-mixer-pack-batch samples seq-len hidden-dim dtype)</procedure> <procedure>(am-mixer-unpack-batch data batch seq-len hidden-dim dtype)</procedure> 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. <procedure>(am-copy-parameters! src dst)</procedure> 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. <procedure>(am-parameter-count model)</procedure> <procedure>(am-parameter-shapes model)</procedure> The total number of parameter elements of {{model}}, and the list of parameter shapes. <procedure>(am-half-mse-loss pred target)</procedure> {{0.5 * mean((pred - target)^2)}}, the convention of nanograd's strict {{mse-loss}}. <procedure>(make-am-rmsnorm dim #!key (eps 1e-6) (dtype 'f32) (name "AM-RMSNorm"))</procedure> <procedure>(make-am-layernorm dim #!key (eps 1e-6) (dtype 'f32) (name "AM-LayerNorm"))</procedure> <procedure>(make-am-batch-norm-nhwc channels #!key (eps 1e-5) (dtype 'f32) (name "AM-BN"))</procedure> Normalization layers. The batch normalization works on NHWC input and always uses the statistics of the current batch. <procedure>(am-softmax-rows s)</procedure> <procedure>(am-rmsnorm-rows x g eps)</procedure> <procedure>(am-layernorm-rows x g b eps)</procedure> Row-wise softmax, RMSNorm and LayerNorm of an {{[M, D]}} morph-variable. <procedure>(am-param data shape dtype)</procedure> <procedure>(am-param-fill shape value dtype)</procedure> <procedure>(am-param-uniform shape scale dtype)</procedure> <procedure>(am-param-normal shape std dtype)</procedure> 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}}. <procedure>(am-const data shape dtype)</procedure> <procedure>(am-const-fill shape dtype f)</procedure> <procedure>(am-scalar value dtype)</procedure> 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. <procedure>(mv lazy-tensor)</procedure> <procedure>(lazy morph-variable)</procedure> <procedure>(am-random-normal)</procedure> {{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 <enscript highlight="scheme"> (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)) </enscript> ==== Lazy ResNet with SSA training <enscript highlight="scheme"> (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))) </enscript> 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 [[https://github.com/iraikov/nanograd-components/tree/main/examples|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. === Author Ivan Raikov === Repository [[https://github.com/iraikov/nanograd-components]] === License LGPL-3 Description of your changes:

I would like to authenticate

Spam control

What do you get when you multiply 4 by 0?