NNX Mini-BERT

Build a small bidirectional encoder

This exercise turns token IDs into one feature vector per input position. Its attention lets each position use information from both earlier and later tokens. That makes it useful for studying an encoder's forward pass.

Despite the mini-BERT name, this is a simplified pre-LN encoder with ReLU and a final normalization, not a reproduction of the original BERT architecture or pretraining procedure. Original BERT includes additional embedding components and task heads, and differs in block details. Dropping a GPT mask and head alone is not a general recipe for reproducing BERT.

Follow the arrays

token IDs (T,)
  -> token embedding + position rows (T, D)
  -> L bidirectional encoder blocks (T, D)
  -> final LayerNorm (T, D)

V is vocabulary size, D is feature width and L is the number of blocks. Return TD hidden-state values, not TV vocabulary logits. There is no pooling or vocabulary projection in this task. A hidden feature vector is not a class probability or a predicted token.

Token and position embeddings

nnx.Embed(V, D, rngs=rngs) supplies the token table using its default initializer. Cast the harness's floating-point ID values to int32 before lookup. Valid IDs are integers from zero through V-1; the cast is not a way to validate arbitrary fractional or out-of-range inputs.

Add the first T rows of a separate float32 position parameter shaped (max_T, D). Initialize that table to zeros, exactly as required here. Initially it adds no position signal, although it can learn different rows during training. Set max_T=T in the entrypoint. A retained model only supports lengths within its allocated position table.

Independent pre-LN blocks

Each block makes two sequential residual updates:

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

Norms normalize across features separately at each position. They have independent learned scale and bias. Attention projects Q/K/V from the normalized states using D-to-D linear layers, splits into H heads of width Dh=D//H, scales dot products by sqrt(Dh), and softmaxes over keys. It uses every token as a valid key. Join the heads before the D-to-D output projection.

ff1 maps D to F=d_ff, and ff2 maps F back to D. The FFN is token-wise; attention supplies the interaction between positions. Residuals retain the unnormalized inputs to each branch. The only final norm is ln_f after the whole stack.

Store separately constructed blocks in nnx.List so NNX traverses their parameters. The blocks share a structure, not their parameter objects. Reusing one block instance would define a different, weight-tied stack.

A small bidirectional example

Start after Q/K/V projection. With zero queries and keys, attention averages the values at all positions:

import jax
import jax.numpy as jnp

q = jnp.zeros((1, 3, 1))
k = jnp.zeros((1, 3, 1))
v = jnp.array([[[1.0], [2.0], [6.0]]])
weights = jax.nn.softmax(q @ k.transpose(0, 2, 1), axis=-1)
output = weights @ v
assert jnp.allclose(output, jnp.full((1, 3, 1), 3.0))
changed_values = v.at[0, 2, 0].set(9.0)
assert jnp.allclose(weights @ changed_values, jnp.full((1, 3, 1), 4.0))

Changing the last token's value affects even the first output. That is allowed in this encoder. The full model additionally applies learned projections, normalization and residuals, so its outputs are not generally simple averages.

What this exercise leaves out

Bidirectional context is suitable when the task permits access to the whole input. In masked-token training, constructing corrupted inputs and choosing prediction targets are separate steps; this function does neither. It does not insert special tokens. Position zero is a classification token only if the caller supplies one and trains an appropriate task head.

The model has no segment embeddings, input-embedding normalization, dropout, padding mask or task head. All supplied tokens are valid. At initialization, the zero position table means permuting token rows also permutes output rows. Learned nonzero position rows can break that symmetry.

The diagram's bidirectional panel matches this attention pattern. It does not show normalization, residual connections or a training objective.