Skip to content

← Flax step 39 of 100

Hard Primitives

Transformer Encoder Block (Pre-LN)

Jump to the task

Combine attention with a per-token network

This encoder block first lets tokens exchange information through self-attention, then transforms each token with a feed-forward network. Each stage adds an update to the representation it received.

Use two distinct intermediate names to keep the residuals clear:

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

The second residual uses x1, which already includes the attention update. Adding the original x there would discard that update from the skip path. LN1 and LN2 are separate LayerNorm modules with separate learned scales and offsets.

What pre-LN means here

Normalize the input to each sublayer, then add its output to that sublayer's unnormalized input. There is no normalization after either residual addition in this block. Post-LN would instead normalize the residual sum, which is a different computation.

Pre-LN provides a direct identity path through each addition. It does not guarantee stable training for every model and learning rate. Transformer families also differ in their activations, normalization, attention and position handling; this block is one specific recipe, not a complete model.

Follow the shapes

Let x have shape (T, D), H = num_heads and F = d_ff.

Stage Shape What it does
First LayerNorm (T, D) Normalize the final feature axis separately for each token.
Query, key and value projections (T, H, D/H) each Split the projected features into heads.
Attention scores (H, T, T) Compare each query with all keys.
Attention output and first residual (T, D) Project back and add to x.
Second LayerNorm (T, D) Normalize the updated representation x1.
First Dense and ReLU (T, F) Transform each token independently.
Second Dense and final residual (T, D) Restore the width and add to x1.

Use qkv_features=D for the projections in this exercise. The attention module's default output width comes from its query input, so it also returns D features. These are separate settings, even though their widths agree here.

The FFN shares its weights across token rows but does not mix the rows. Attention is bidirectional and unmasked: a token may use both earlier and later positions. No padding mask, position encoding or decoding cache is added by this block.

Keep both residual updates

Suppose one row of the input and its two computed updates are as follows. This example isolates the additions; it does not calculate the attention or FFN updates themselves.

import jax.numpy as jnp

x = jnp.array([[1., 3.]])
attention_update = jnp.array([[.5, -.5]])
x1 = x + attention_update
ffn_update = jnp.array([[2., 1.]])  # Computed from LN2(x1) in the real block.
out = x1 + ffn_update

assert jnp.array_equal(out, jnp.array([[3.5, 3.5]]))
assert not jnp.array_equal(out, x + ffn_update)

Problem

Implement encoder_block_forward(seed, x, num_heads, d_ff):

  1. Define a Linen module with integer num_heads and d_ff fields and a compact call method. Set D = x.shape[-1].
  2. Apply an unnamed nn.LayerNorm() to x, then an unnamed nn.MultiHeadDotProductAttention(num_heads=H, qkv_features=D). Add the attention output to x to obtain x1.
  3. Apply a new unnamed nn.LayerNorm() to x1, followed by nn.Dense(F), nn.relu and nn.Dense(D), with both Dense layers unnamed. Add that result to x1.
  4. Initialize the module directly with jax.random.PRNGKey(int(seed)) and x, apply the returned variables to x, and flatten the output with .reshape(-1).

Keep the layer defaults: LayerNorm epsilon 1e-6, learned scales and offsets, enabled projection biases, and zero attention dropout. Do not add a final LayerNorm or a causal mask.

The default parameter scopes are LayerNorm_0, MultiHeadDotProductAttention_0, LayerNorm_1, Dense_0 and Dense_1. Preserve this structure so seeded initialization matches the tests.

x is float32 with shape (T, D). Dimensions, H and F are positive, and D is divisible by H. Return a one-dimensional array of length T * D, with features varying fastest. Keep seed, num_heads and d_ff static if JIT-compiling the whole function.

The bidirectional panel below shows this block's attention permissions. The other panels are comparisons; they are not extra stages in the block.

Loading visualization…