Mini LM Capstone: A Causal Forward Pass

Assemble a small language model

Combine token embeddings, sinusoidal positions, causal attention, RMSNorm and a SwiGLU feed-forward network. The result is a decoder-only model that returns one vocabulary-score vector per input token.

This exercise computes a randomly initialized forward pass. It does not train the model or generate text. The components are useful building blocks, but this particular combination is not an implementation of a named production model.

token IDs                  (T,)
token embeddings           (T, D)
+ sinusoidal positions     (T, D)
L decoder blocks           (T, D)
final RMSNorm              (T, D)
multiply by embedding.T    (T, V)
flatten                    (T * V,)

Add deterministic positions

Use one unnamed nn.Embed(V, D) with its default initializer. For each position t and pair index i from 0 through D/2 - 1, add:

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

Compute the table in float32 for the current sequence length. It is not a learned parameter. This exercise requires an even D, so each sine coordinate has a cosine partner. Row zero alternates 0, 1, 0, 1, ....

The same initialized model can accept another sequence length because no parameter shape depends on T. Positions still start at zero on each call. This does not guarantee useful predictions at longer lengths, and it is not an incremental decoding cache.

Apply two residual updates per block

Define RMSNorm with eps=1e-6 and a learned vector gamma initialized to ones. It operates independently on each token's final feature axis:

RMSNorm(x) = gamma * x / sqrt(mean(x ** 2, axis=-1, keepdims=True) + eps)

It does not subtract the mean or add a learned bias. Each block creates two separate norm instances and applies:

a   = x + causal_attention(RMSNorm0(x))
out = a + SwiGLUFFN(RMSNorm1(a))

The residual uses the unnormalized input to each branch. Normalization goes before attention and the FFN, not after either addition. These placements define this model; they do not guarantee stable training or remove the need to choose a learning-rate schedule.

Attention uses nn.MultiHeadDotProductAttention(num_heads=H, qkv_features=D) and a lower-triangular (T, T) mask including the diagonal. Position t can read positions 0 through t. Changing later input tokens cannot change earlier outputs. All positions are valid here; no padding mask is added.

Build the SwiGLU branch

Define SwiGLUFFN with three Dense layers, named exactly gate, up and down:

gate = Dense(F, name="gate")(x)
up   = Dense(F, name="up")(x)
h    = silu(gate) * up
out  = Dense(D, name="down")(h)

Both first projections have shape (T, F). Multiply them elementwise after applying SiLU to the gate branch, then project back to (T, D). Keep all three biases. SiLU is z * sigmoid(z); the gate is not a probability and can be negative or greater than one.

Create a new DecoderBlock for each layer so weights are independent. After all blocks, apply one final RMSNorm, then use the same embedding matrix W for hidden @ W.T. There is no separate output matrix, output bias or softmax. Reusing W saves V * D parameters compared with an independent bias-free output matrix of shape (D, V).

Read logits and next-token targets correctly

Logits are unrestricted scores, not probabilities. Applying softmax across the vocabulary would turn a row into a distribution. This function returns the raw scores instead.

In next-token training, row t is compared with the token after the input at position t. Given only one sequence, use logits[:-1] and ids[1:]; the final row needs a next token from outside that sequence. The forward function still returns every row and performs no loss calculation.

This example checks position coordinates, the tied projection and target alignment with known arrays. The hidden states stand in for a model's final normalized states:

import jax.numpy as jnp

T, D = 3, 4
angles = jnp.arange(T, dtype=jnp.float32)[:, None] / (
    10000.0 ** (2 * jnp.arange(D // 2, dtype=jnp.float32)[None, :] / D)
)
positions = jnp.stack((jnp.sin(angles), jnp.cos(angles)), axis=-1).reshape(T, D)
assert jnp.array_equal(positions[0], jnp.array([0., 1., 0., 1.]))

W = jnp.array([[1., 0.], [0., 1.], [1., 1.]])
hidden = jnp.array([[1., 2.], [-1., 0.], [0., 1.]])
ids = jnp.array([0, 2, 1], dtype=jnp.int32)
logits = hidden @ W.T
assert jnp.array_equal(logits, jnp.array([[1., 2., 3.], [-1., 0., -1.], [0., 1., 1.]]))
assert logits[:-1].shape == (2, 3)
assert jnp.array_equal(ids[1:], jnp.array([2, 1]))