Mini GPT: Decoder-Only Language Model

Assemble a small causal language model

This exercise combines token embeddings, learned positions, causal self-attention blocks and a tied vocabulary projection. It computes a randomly initialized forward pass. It does not train the model, calculate a loss or generate tokens.

With T input tokens, vocabulary size V and hidden width D, the flow is:

ids                  (T,)
token embedding      (T, D)
+ learned positions  (T, D)
L causal blocks      (T, D)
final LayerNorm      (T, D)
@ token_embedding.T  (T, V) logits

There is no separate encoder or cross-attention. This is one GPT-style architecture with ReLU feed-forward networks, not an exact specification of every GPT model.

Keep the two embedding initializers distinct

Create one unnamed nn.Embed(V, D) with its default initializer. Its (V, D) parameter supplies both the token lookup and the output weights. Separately declare a root parameter pos_embed of shape (T, D) using nn.initializers.normal(stddev=0.02). Add it to the lookup result before the first block.

The position table in this exercise has exactly T rows. It is not a larger table sliced to the current length. Applying the same variables to a different length requests a different parameter shape and fails. Reinitializing at another length creates a new model; it is not a way to extend a trained checkpoint. A variable-length implementation would need a different position-table contract.

Learned positions supply a vector for each position. Tokens with the same ID can have different outputs because their positions and available context differ. The causal mask also records which positions precede each query. Reordering token rows can therefore change the allowed connections.

Apply each causal block

Let H = num_heads and F = d_ff. Each block has two residual updates:

x1  = x + MHA(LN1(x), mask=lower_triangle)
out = x1 + Dense(D)(relu(Dense(F)(LN2(x1))))

Use separate LayerNorms and nn.MultiHeadDotProductAttention with num_heads=H and qkv_features=D. The (T, T) lower-triangular mask allows a query at position t to read positions 0 through t, including itself. Attention scores have shape (H, T, T).

Instantiate a new GPTBlock for each layer. Reusing one instance would share block parameters, which this exercise does not request. After all L blocks, apply a separate final LayerNorm. It changes the vectors sent to the output head and is part of the required computation.

Project to vocabulary scores

If h is the final (T, D) hidden array and W is the (V, D) token embedding, compute h @ W.T. The result has shape (T, V). No separate output weight, bias or softmax is used. Tying saves V*D weights compared with adding an independent bias-free vocabulary projection.

This example uses supplied hidden vectors to isolate the projection and show how next-token targets would line up in a separate training loop:

import jax.numpy as jnp

ids = jnp.array([0, 2, 1], dtype=jnp.int32)
W = jnp.array([[1., 0.], [0., 1.], [1., 1.]])
h = jnp.array([[1., 2.], [0., -1.], [2., 0.]])
logits = h @ W.T

assert jnp.array_equal(logits, jnp.array([
    [1., 2., 3.], [0., -1., -1.], [2., 0., 2.],
]))
training_logits = logits[:-1]
targets = ids[1:]
assert training_logits.shape == (2, 3)
assert jnp.array_equal(targets, jnp.array([2, 1]))

Row t is paired with token t+1, not the token already visible at t. The last row needs a continuation target that is not present in this sequence. Return all rows for this exercise; do not shift or trim them. Random initialization does not guarantee useful predictions.