NNX Mini-GPT

Assemble a small causal Transformer

This exercise combines token embeddings, a position table, a stack of causal blocks, and a vocabulary projection. It computes one forward pass with freshly initialized parameters. It does not train a language model, generate tokens, or reproduce a particular GPT checkpoint.

The arrays follow this path:

token IDs (T,)
  -> token lookup + position rows (T, D)
  -> L causal blocks (T, D)
  -> final LayerNorm (T, D)
  -> multiply by token table transpose (T, V)

V is vocabulary size, D is model width, and L is the number of blocks. The result contains one score for every vocabulary entry at each position. These are logits, not probabilities. LayerNorm normalizes hidden features; it does not normalize the output logits into a probability distribution.

Token and position tables

Use nnx.Embed(V, D, rngs=rngs) with its default initializer for tokens. It owns a (V, D) parameter table. IDs select rows of that table.

The position table is a separate (max_T, D) float32 nnx.Param, initialized to zeros as part of this task's exact contract. Its first T rows are added to the token vectors. Initially this adds nothing, but the table is trainable and its rows can receive different gradients. A random position initializer would change the required output.

Set max_T=T in the entrypoint. A retained model can handle lengths at most max_T; slicing a too-short table does not create missing rows. Longer inputs need a deliberate capacity change outside this exercise.

Stack independent blocks

Each block has separate parameters but the same structure:

a = x + causal_attention(ln1(x))
output = a + ff2(relu(ff1(ln2(a))))

The attention module has Q, K, V and output projections from D to D. Split the projected features into H heads of width Dh=D//H. Scale dot products by sqrt(Dh), and allow key j only when j<=i for query i. Mask forbidden scores with negative infinity before softmax. The diagonal stays valid, so every row has at least one allowed finite score.

The FFN maps D to F=d_ff, applies ReLU, then maps F back to D. Norms act across features separately per token. The residual paths preserve the unnormalized input to each branch. Add no final norm inside a block; ln_f is applied once after the entire stack.

Store separately constructed blocks in nnx.List. This makes their parameters part of NNX's traversable module state. Do not reuse one block object L times: that would tie all layer weights. An ordinary list's treatment depends on NNX's data/static attribute rules and version; it is not safe to assume that an invalid list silently drops parameters.

Tie the output weights

Read the token table as self.embed.embedding[...] and transpose the resulting array. Multiplying final hidden states by this transpose reuses the input embedding parameter. There is no separate head kernel or bias. Compared with an untied bias-free output projection, this saves V*D parameters. Both uses contribute gradients to the same table.

A small example isolates the final normalization and tied projection, which is also the complete forward pass when L=0 and positions are zero:

import jax.numpy as jnp
from flax import nnx

table = jnp.array([[1.0, 3.0], [2.0, 0.0], [0.0, 1.0]])
ids = jnp.array([0, 1], dtype=jnp.int32)
norm = nnx.LayerNorm(2, epsilon=1e-6, rngs=nnx.Rngs(0))
hidden = norm(table[ids])
logits = hidden @ table.T
scale = 1.0 / jnp.sqrt(1.0 + 1e-6)
expected = scale * jnp.array([[2.0, -2.0, 1.0], [-2.0, 2.0, -1.0]])
assert jnp.allclose(logits, expected)

The vocabulary-wide scores can be negative and do not sum to one. Removing ln_f would produce different dot products.

Causality and scope

With fixed parameters and position rows, changing future token IDs cannot change earlier logits. This is a property of the full stack, not just its attention matrices. Next-token training still needs shifted targets; this function returns logits without constructing targets or a loss.

The diagram's causal panel shows the permitted query-key connections. This model has no cross-attention, KV cache, dropout or padding mask. Treat every supplied token as valid.