Problems
Loading...

08/08Encoder Block

Medium

Implement one post-normalization Transformer encoder block with multi-head self-attention and a position-wise feed-forward network.

z=LayerNorm⁡(x+MHA⁡(x,x,x))z = \operatorname{LayerNorm}(x + \operatorname{MHA}(x,x,x)) y=LayerNorm⁡(z+FFN⁡(z))y = \operatorname{LayerNorm}(z + \operatorname{FFN}(z))

Here, MHA denotes projected multi-head attention. The feed-forward network is

FFN⁡(z)=ReLU⁡(zW1+b1)W2+b2.\operatorname{FFN}(z)=\operatorname{ReLU}(zW_1+b_1)W_2+b_2.

Each layer normalization computes population variance across the final axis and uses

ε=10−6.\varepsilon=10^{-6}.

Return y as a NumPy float64 array with the same shape as x.

Loading visualization...

Examples

Input: x = [[[1,0]]], W_q = W_k = W_v = W_o = [[1,0],[0,1]], W1 = [[1,0,0,0],[0,1,0,0]], W2 = [[1,0],[0,1],[0,0],[0,0]], b1 = [0,0,0,0], b2 = [0,0], gamma1 = gamma2 = [1,1], beta1 = beta2 = [0,0], num_heads = 1

Output: [[[1.0, -1.0]]]

Explanation: With identity projections and a single token, self-attention returns x itself. After residual + layer norm, FFN, and second residual + layer norm, the output is layer-normalized.

Input: x = [[[1,2],[3,4]]], W_q = W_k = W_v = W_o = I(2), W1 = [[1,0,0,0],[0,1,0,0]], W2 = [[1,0],[0,1],[0,0],[0,0]], b1 = [0,0,0,0], b2 = [0,0], num_heads = 1

Output: [[[-1.0, 1.0], [-1.0, 1.0]]]

Input: x = [[[0,1],[1,0]]], W_q = W_k = W_v = W_o = I(2), W1 = [[1,0,1,0],[0,1,0,1]], W2 = [[1,0],[0,1],[0,0],[0,0]], b1 = [0,0,0,0], b2 = [0,0], num_heads = 1

Output: [[[-1.0, 1.0], [1.0, -1.0]]]

Hint 1

Reshape projected tensors to batch, heads, sequence, head_width.

Hint 2

Normalize x + attention before evaluating the feed-forward network.

Hint 3

Normalize the second residual sum with gamma2 and beta2.

Requirements

  • Use NumPy float64 computations.
  • Split projected Q, K, and V evenly across num_heads.
  • Apply residual addition before each layer normalization.
  • Apply the feed-forward network independently at every position.
  • Return an unrounded NumPy array.

Constraints

  • x has shape (batch, sequence, d_model).
  • d_model is divisible by num_heads.
  • Projection, feed-forward, scale, and shift arrays have compatible shapes.
  • Valid inputs are guaranteed.