Neural network layers

In short: a module registers its parameters, buffers and child modules in its constructor, which gives every parameter a dotted path such as attn.qkv.weight; optimizers iterate that list and checkpoints are keyed by it. UnaryModule is the one-input, one-output shape that Sequential chains. The layers are Linear, five activation modules, Xavier and Kaiming initializers, Conv2d by patch rows and one matrix product, MaxPool2d, Flatten, Embedding, LayerNorm and RMSNorm, Dropout, RNN and LSTM cells and layers, scaled dot-product and multi-head attention, and TransformerBlock. None of it is on the inference path: the graph executor and the LLM decoder have their own kernels.

Each header links its concept page on /ml-models: neurons and layers, activation functions, initialization, CNNs, embeddings, regularization, layer norm and residuals, RNNs and LSTMs, attention and the Transformer. The losses and optimizers that train these layers are on Losses and optimizers.

Module

Module is non-copyable and holds three ordered lists: parameters, buffers and named children, plus a training_ flag that starts true. A subclass fills them in its constructor through the protected Register* calls (module.h L4–10).

Names come from CollectParameters, which emits this module's own parameters with the prefix, then recurses into each child with prefix + name + "." (L63–73). That order is what CopyParameters compares, and what checkpoints are keyed by: SaveCheckpoint and LoadCheckpoint take an nn::Module (rml/io/checkpoint.h L23–31).

UnaryModule, Sequential and Lambda

UnaryModule adds a pure virtual Forward(const Variable&) and an operator() that calls it (module.h L78–83). Sequential applies its children in order and names them "0", "1", … as they are appended, so its parameters read 0.weight, 2.bias; appending a null layer fails a check (L85–99, module.cc L94–107). Lambda wraps a parameter-free function, such as a reshape, as a module (L101–109). Layers that take something other than one Variable (Embedding's ids, the recurrent cells' state) derive from Module directly.

Linear and activations

Linear(in_features, out_features, gen, {bias}) holds W as [out, in] and computes y = x Wᵀ + b over the last axis, so any leading batch axes pass through (linear.cc L29–35). Both W and b start uniform on [−1/√in, 1/√in] (L16–27), which the header notes gives each output a variance of about Var(x)/3 whatever the width (linear.h L21–23). An input whose last axis is not in_features fails a check.

Activation<kFn> is a template over a function pointer, so each activation module is one alias: Relu, Tanh, Sigmoid, Gelu and Silu forward to the autograd op of the same name, where the derivatives are written (Autograd).

Initialization

init.h states the goal: keep the variance of activations forward and of gradients backward roughly constant from layer to layer. Weights follow the [out, in, …] layout, and the receptive field is the product of the trailing axes.

The uniform forms draw from [−√3·std, √3·std], since a uniform on [−a, a] has variance a²/3 (L40–44). All four draw from the counter-based generator on Random numbers, so the same seed gives the same weights. The layers themselves do not call these: Linear, Conv2d and the recurrent layers use their own uniform bound, and the initializers are there for code that wants a specific scheme.

Worked example. A [128, 256] weight: fan_in = 256, fan_out = 128. Kaiming with ReLU gives std = √2 / 16 ≈ 0.0884; Xavier gives √(2 / 384) ≈ 0.0722; Linear(256, 128)'s own bound is 1/16 = 0.0625, a standard deviation of 0.0625 / √3 ≈ 0.0361.

Convolution and pooling

The header describes convolution as unfold (im2col) followed by one matrix product (conv.h L8–12). The code goes one step further. Conv2dFunctional(x, weight, bias, stride, padding) takes x [N, C, H, W] and a square weight [O, C, k, k] and builds the patches as rows, [N·OH·OW, C·k·k], with a private PatchRowsFunction, written sequentially so it needs neither a zero fill nor a transposed copy (conv.cc L109–176). One product with the reshaped, transposed filters covers the whole batch, the bias is added per output channel, and a reshape and transpose return [N, O, OH, OW] (L239–263). Each output side is (H + 2·padding − k) / stride + 1 (L29–41).

Worked example. A 28 × 28 image through Conv2d(1, 8, 3, gen, {.stride = 2, .padding = 1}): each side is (28 + 2 − 3) / 2 + 1 = 14, so the patch rows are [N·196, 9], the filters [9, 8], and the result [N, 8, 14, 14].

Embedding

Embedding(num_embeddings, embedding_dim, gen, stddev = 1) is a [num, dim] table initialized N(0, stddev²). Forward takes an i32 ids tensor of any shape and returns that shape with dim appended, by autograd::IndexSelect on axis 0 (embedding.cc L16–25). The header's point: it equals a Linear on one-hot inputs without making the one-hots, gathering rows forward and scatter-adding into them backward, so a repeated id accumulates its gradient (embedding.h L4–7).

Normalization

Both are written as autograd ops (Mean, Square, Pow(·, −0.5)), so their gradients come for free and are checked numerically. The graph executor's layer norm is a separate fast kernel (Lowering and fast CPU).

Dropout

Dropout(p, seed = 0) requires p in [0, 1). In training mode it draws a Bernoulli(1 − p) keep mask from its own generator and multiplies by mask / (1 − p), so the expected activation is unchanged and evaluation needs no rescaling; in eval mode, or when p is 0, it returns x unchanged (dropout.cc L15–23). Owning its generator makes a run reproducible from the seed alone.

Recurrent layers

The header gives the equations and the reason for the LSTM (recurrent.h L9–18): an RNN's h′ = tanh(x W_ihᵀ + h W_hhᵀ + b) multiplies one Jacobian per step on the way back, so distant gradients shrink or explode; an LSTM's cell state changes additively, c′ = σ(f)·c + σ(i)·tanh(g), h′ = σ(o)·tanh(c′), so while the forget gate stays near 1 the gradient through c passes almost unchanged.

Every weight and bias starts uniform on [−1/√H, 1/√H] (L26–31), except that the LSTM's forget-gate bias, rows [H, 2H), starts at 1, so the cell begins by keeping its memory (Jozefowicz et al. 2015; L83–91). The same cell, exported from PyTorch, is the graph that runtime fusion handles on Lowering and fast CPU.

Attention

Worked example. CausalMask(2, 4), two new queries after a cached prefix of two keys: query 0 sits at position 2 and sees keys 0–2, query 1 at position 3 sees all four.

[[0, 0, 0, -inf],
 [0, 0, 0,    0]]

The header calls this "the readable training version" and says the inference engine's fused kernel in rml/backend is tested against it (attention.h L9–10). At this commit rml/backend does not exist, so that is planned; the LLM decoder's attention on The decoder is its own code and includes nothing from rml/nn.

TransformerBlock

TransformerBlock(embed_dim, num_heads, gen, {mlp_ratio = 4, causal = true}) is the pre-norm block of GPT-2 and its successors (transformer.h L4–9): normalize before each sublayer and add its output back, which leaves an unnormalized residual path from input to output. The whole forward pass is two lines (transformer.cc L29–32).

Worked example. TransformerBlock(8, 2, gen) reports twelve parameters in registration order: ln1.weight, ln1.bias, attn.qkv.weight, attn.qkv.bias, attn.out.weight, attn.out.bias, ln2.weight, ln2.bias, fc.weight, fc.bias, proj.weight, proj.bias. With d = 8 the count is 4d for the two norms, 4d² + 4d for attention and 8d² + 5d for the MLP: 32 + 288 + 552 = 872 (transformer_test.cc L19–31).

Who uses it

Inside the library, only rml/io/checkpoint depends on rml/nn. The rest of its users are the training demos in examples/: mlp, cnn, rnn, ssm, encoder, char_gpt, vit, vae, gan, diffusion, diffusion_unet, flow_matching, gnn, two_tower, multimodal, rag, distillation, finetune and dpo (Examples). The inference stack does not use it: exported models run on the graph executor's kernels (Executing a graph).

Tests

Every layer with parameters gets a numerical gradient check of its input and parameters through GradCheck, and forward results are compared with the definition written out in the test.