Skip to content

← Flax NNX step 100 of 100

Hard Primitives

NNX Mini Language Model Capstone

Jump to the task

Build a complete forward pass

Combine token embeddings, positional signals, causal attention, RMSNorm and SwiGLU into a small decoder-only model. The output is one vector of vocabulary logits per input position. This exercise does not calculate a loss, update weights, or generate tokens.

The model is freshly initialized. Its logits are not trained predictions. It combines sinusoidal input positions with RoPE, uses biased linear projections, and ties the output weights to the embedding table. These are the exact choices for this exercise, not a claim of compatibility with a named pretrained model.

Tokens and positions

Let V be vocabulary size, D model width, H head count, F feed-forward width, L block count, and T sequence length. Initialize one trainable embedding array E of shape (V, D) with jax.random.normal(rngs.params(), (V, D)) / sqrt(D).

Cast the supplied token IDs to int32. Assume they are integer-valued, have shape (T,), and lie in [0, V). Look up E[ids] and add a fixed sinusoidal position array of shape (T, D):

angle[t, i] = t / 10000 ** (2 * i / D)
position[t, 2*i]     = sin(angle[t, i])
position[t, 2*i + 1] = cos(angle[t, i])

Here t starts at zero and i ranges from zero to D/2 minus one. At position zero, even features are zero and odd features are one. These position values have no trainable parameters.

RMSNorm and residual blocks

RMSNorm scales each row without subtracting its mean: gamma * x / sqrt(mean(x ** 2, axis=-1, keepdims=True) + 1e-6). Each norm owns a trainable length-D gamma initialized to ones, with no bias. Read its array with gamma[...].

Each decoder block applies two pre-normalized residual branches:

  1. x = x + attention(norm1(x)).
  2. x = x + swiglu(norm2(x)).

The skip path is unchanged by the branch's normalization. This describes the architecture, not a guarantee of stable training, nonvanishing gradients, or no need for learning-rate warmup.

Causal attention with adjacent-pair RoPE

Create four separate biased D-to-D Linear layers: query, key, value, then output projection. Reshape each projected array from (T, D) into H heads of width Dh = D // H, preserving feature order. Dh must be even.

For every head, rotate adjacent feature pairs of queries and keys. At position t and pair i, use angle t * 10000 ** (-2 * i / Dh). Rotate a pair (a, b) into (a * cos(angle) - b * sin(angle), a * sin(angle) + b * cos(angle)). Values are not rotated. RoPE's angles use head width Dh, whereas the additive position encoding above uses model width D.

Compute scores as Q @ K.T / sqrt(Dh). Query position t may attend to key positions 0 through t, including itself. Replace scores for later keys with negative infinity, then apply softmax over keys and multiply by values. Concatenate heads in order and apply the output projection. There is no dropout or KV cache.

With finite allowed scores, every row has at least its diagonal available, so the masked softmax has a finite denominator. Negative infinity is only an intermediate mask value; it is not returned as a logit. A finite sentinel such as -1e9 can leak future information if an allowed score is even smaller.

This complete example demonstrates that masking distinction:

import jax
import jax.numpy as jnp
import numpy as np

scores = jnp.array([[-2e9, 0.], [0., 0.]])
allowed = jnp.tril(jnp.ones((2, 2), dtype=bool))
weights = jax.nn.softmax(jnp.where(allowed, scores, -jnp.inf), axis=-1)
np.testing.assert_allclose(weights, [[1., 0.], [0.5, 0.5]])
finite_mask = jax.nn.softmax(jnp.where(allowed, scores, -1e9), axis=-1)
assert float(finite_mask[0, 1]) == 1.0  # Incorrect future attention.

SwiGLU and tied output weights

SwiGLU has three biased Linear layers, constructed as gate (D, F), up (D, F), then down (F, D). Its forward calculation is down(jax.nn.silu(gate(x)) * up(x)). Apply SiLU only to the gate branch; multiply elementwise before the down projection.

After all blocks, apply one final RMSNorm. Then compute logits = x @ E.T, shape (T, V). Reuse the same embedding parameter for lookup and projection; there is no separate head kernel or bias. Tying saves V times D parameters compared with an otherwise identical model with an independent bias-free output matrix. It does not halve the total model size or guarantee better predictions.

The additive positions and RoPE are both required here. Their combination does not guarantee useful behavior beyond training lengths. No training has occurred in this task.

Problem

Implement mini_lm_capstone(seed, token_ids, vocab_size, d_model, num_heads, d_ff, num_layers). Cast scalar configuration values to Python integers and token IDs to int32. Assume T, V, D, H and F are positive, L is nonnegative, D is divisible by H, and D/H is even.

Create one nnx.Rngs(int(seed)) and use it throughout construction:

  1. Draw the embedding table first, using one rngs.params() key and the normal initialization specified above.
  2. Create L distinct DecoderBlock modules in an nnx.List. Within each block, create norm1, attention's query/key/value/output projections, norm2, then SwiGLU's gate/up/down projections.
  3. Create final RMSNorm. Norm gamma arrays are ones and consume no random key. Linear layers use default NNX initializers and biases.

Forward: embedding lookup, additive sinusoidal positions, all decoder blocks, final RMSNorm, then the tied projection. Read parameters through parameter[...] and return logits flattened in row-major order to shape (T * V,). Do not apply softmax to the final logits.

With L=0, still add positions, apply final RMSNorm, and project with E.T. As a causality check, changing later input IDs must not change logits at earlier positions while parameters remain fixed.

The diagram below compares attention connectivity. This model uses its causal pattern; it does not illustrate the position rotations, learned weights, or this model's exact sequence length.

Loading visualization…