NNX Transformer Encoder Block

Combine attention with a token-wise network

This encoder block applies two transformations. Self-attention lets each token use information from every token in the sequence. A feed-forward network then transforms each token separately using shared weights. Each transformation adds its result to its input through a residual connection.

The exercise uses pre-normalization, or pre-LN. Normalize the input to each branch, while preserving an unnormalized path for the residual:

after_attention = x + attention(layer_norm_1(x))
output = after_attention + feed_forward(layer_norm_2(after_attention))

The second branch uses the updated after_attention value. It does not reuse the original x. Post-LN instead normalizes after the residual sum, which defines a different block. Both arrangements exist; this task specifically requires pre-LN and has no final normalization.

Follow the two branches

Let x have shape (T, D), H be the number of attention heads, and Dh=D//H. LayerNorm normalizes each token across its D features, with separate learned scale and bias for each of the two normalization layers. It does not normalize across tokens or keep running statistics.

Attention projects the normalized input through four D-to-D linear layers: Q, K, V and output. Split each Q/K/V result into H heads of width Dh, compute Q @ K.T / sqrt(Dh) within each head, softmax over key positions, and take a weighted sum of V. Join the heads back in token order before the output projection. This branch is unmasked.

The feed-forward branch is ff2(relu(ff1(z))), where ff1 maps D to F=d_ff and ff2 maps F to D. ReLU makes this branch nonlinear. Without it, the two affine layers could be combined into one affine map. That statement applies to the branch after normalization, not the entire block.

Both branches return (T, D), so their residual additions preserve the shape. F is supplied by the task; it need not be 4*D or even larger than D. No particular width guarantees a hardware-performance balance.

A small example with known weights

Consider two tokens and one head. Set Q/K projections to zero and V, output, ff1 and ff2 to identity matrices. Set all linear biases to zero, and normalization scales/biases to one/zero. Attention then averages the normalized token vectors. The following computes this complete block:

import jax.numpy as jnp

def normalize(x):
    centered = x - x.mean(axis=-1, keepdims=True)
    return centered / jnp.sqrt(jnp.mean(centered ** 2, axis=-1, keepdims=True) + 1e-6)

x = jnp.array([[1.0, 3.0], [2.0, 4.0]])
normalized = normalize(x)
attention = jnp.broadcast_to(normalized.mean(axis=0), x.shape)
after_attention = x + attention
output = after_attention + jnp.maximum(normalize(after_attention), 0.0)
assert jnp.allclose(output, jnp.array([[0.0, 5.0], [1.0, 6.0]]), atol=2e-6)

The attention result is approximately [-1, 1] at both positions. The second branch contributes approximately [0, 1]. Notice that neither residual replaces x with a normalized vector. As another useful check, if both branch outputs are zero, the block returns x exactly.

Scope of this block

This is one encoder-style block, not a complete BERT, vision transformer or other named architecture. It has no input embeddings, positional information, dropout, padding mask, final normalization or task head. Every token is treated as valid. Without position information, permuting the input token rows permutes the output rows in the same way.

The diagram compares attention patterns. The bidirectional panel matches this block: every query can attend to every key. It does not depict the normalization, feed-forward network or residual additions.