Skip to content

← Flax step 43 of 100

Hard Primitives

Mini BERT: Encoder-Only Hidden States

Jump to the task

Return one hidden vector per token

This exercise builds a small bidirectional Transformer encoder. It looks up token embeddings, adds learned positions and applies a stack of blocks. The output is a (T, D) array of hidden states: one vector for each input token. It is not a vocabulary distribution or a sentence-level prediction.

ids                  (T,)
token embedding      (T, D)
+ learned positions  (T, D)
L encoder blocks     (T, D)
returned states      (T, D)

This is a simplified BERT-inspired model. The original BERT uses GELU and post-LN blocks; this exercise uses ReLU and pre-LN blocks. It also omits token-type embeddings, embedding normalization, dropout and task heads. These choices change the computation, not just its size.

Let attention read both directions

Each block uses unmasked self-attention. Every query position can attend to every key position, including itself and later tokens. Changing a later token can therefore change an earlier token's hidden state.

All positions are treated as valid here. A real padded sequence would usually need a padding mask; bidirectional attention does not require attending to padding.

An attention mask and masked-language-model training are different ideas. An attention mask controls which query-key pairs are allowed. Masked-language modeling hides or replaces selected input tokens and trains a prediction head to recover them from context. This exercise performs neither input corruption nor prediction-head training; it just returns initialized encoder states.

Add positions before the first block

Use one unnamed nn.Embed(V, D) with its default initializer. Declare a separate root parameter pos_embed with shape (T, D), initialized by nn.initializers.normal(stddev=0.02). Add the two arrays directly. Do not change the token embedding initializer to match the position table.

The table has exactly T rows. The same initialized variables cannot be applied to a different sequence length: requesting a different pos_embed shape fails. There is no larger table being sliced here.

This example isolates the addition with known weights:

import jax.numpy as jnp

embedding = jnp.array([[1., 0.], [0., 1.]])
ids = jnp.array([0, 1, 0], dtype=jnp.int32)
positions = jnp.array([[0., .1], [.2, 0.], [.3, -.1]])
x = embedding[ids] + positions

assert x.shape == (3, 2)
assert jnp.allclose(x, jnp.array([[1., .1], [.2, 1.], [1.3, -.1]]))
assert not jnp.array_equal(x[0], x[2])  # Same ID, different position vector.

Without positions, this unmasked stack is permutation-equivariant: reordering token rows reorders the output rows. It does not make distinct token IDs identical. Equal IDs in the same unmasked sequence would have equal states in that position-free version, since they have identical queries and can access the same keys and values.

Apply the pre-LN blocks

Each block uses two separate LayerNorms and two residual updates:

x1  = x + MHA(LN1(x))
out = x1 + Dense(D)(relu(Dense(d_ff)(LN2(x1))))

LayerNorm reduces over the final feature axis independently for each token. Attention uses num_heads=H and qkv_features=D, with scores of shape (H, T, T). The FFN expands to (T, d_ff) and projects back to (T, D). Its weights are shared across token rows, but each row is transformed independently after attention has mixed information across positions.

Create a new BERTBlock instance for each layer so blocks have independent parameters. Return the last block's output directly, with no final LayerNorm, pooling, output projection or softmax.

Problem

Implement mini_bert_forward(seed, token_ids, vocab_size, d_model, num_heads, d_ff, num_layers):

  1. Convert configuration values to Python integers and token IDs to jnp.int32. Define BERTBlock and MiniBERT as compact Linen modules.
  2. Follow the embedding, position and block computations above. Preserve root scopes Embed_0, pos_embed, and BERTBlock_0 through BERTBlock_{L-1}.
  3. Inside each block, use unnamed nn.LayerNorm(), nn.MultiHeadDotProductAttention(num_heads=H, qkv_features=D), another nn.LayerNorm(), nn.Dense(F), nn.relu and nn.Dense(D). Keep the two residual additions in the positions shown above.
  4. Keep layer defaults, including projection biases, LayerNorm epsilon 1e-6 and zero attention dropout. Initialize the complete model directly with jax.random.PRNGKey(int(seed)) and ids, apply those variables to ids, and return .reshape(-1).

token_ids is a one-dimensional float array of length T containing integer-valued IDs in [0, V). T, V, D, H and F are positive, D is divisible by H, and L is nonnegative. With L=0, return the token embeddings plus positions without any block transformations.

Return a one-dimensional array of length T * D, with features varying fastest. Keep seed and all five model configuration arguments static if JIT-compiling the whole function. Preserve module names and initializer choices to reproduce the seeded tests.

The bidirectional panel below matches this encoder. The causal and cross-attention panels are comparisons, not additional model stages.

Loading visualization…