Problems
Loading...
Step 4 of 6
ViT Encoder Block
Medium

Implement one pre-normalized Vision Transformer encoder block. Normalize the tokens, apply unmasked multi-head self-attention, project and add the first residual, then normalize again, apply a two-layer GELU MLP, and add the second residual.

A=softmax⁡(QKTdh)V.A = \operatorname{softmax}\left(\frac{QK^{\mathsf T}}{\sqrt{d_h}}\right)V. y=x+concat⁡(A1,…,AH)Wo.y = x + \operatorname{concat}(A_1,\ldots,A_H)W_o. z=y+GELU⁡(LN⁡(y)W1)W2.z = y + \operatorname{GELU}(\operatorname{LN}(y)W_1)W_2.

Here, xx contains the input tokens, HH is num_heads, dh=D/Hd_h=D/H, DD is the last dimension of x, and the projection matrices are Wq, Wk, Wv, Wo, W1, and W2. LayerNorm uses population variance and ε=10−6\varepsilon=10^{-6}. Return zz as a float64 NumPy array with shape (B,N,D)(B,N,D).

Loading visualization...

Examples

Input: x = [[[1,0],[0,1]]], num_heads = 1, Wq = [[1,0],[0,1]], Wk = [[1,0],[0,1]], Wv = [[1,0],[0,1]], Wo = [[1,0],[0,1]], W1 = [[1,0],[0,1]], W2 = [[1,0],[0,1]]

Output: [[[2.729574,-1.047191],[-1.047191,2.729574]]]

Explanation: Pre-normalized attention and the GELU MLP contribute through separate residual paths.

Input: x = [[[1,0],[0,1],[1,1]]], num_heads = 2, Wq = [[0,1],[1,0]], Wk = [[1,0],[0,1]], Wv = [[1,0],[0,1]], Wo = [[1,0],[0,1]], W1 = [[1,0],[0,1]], W2 = [[1,0],[0,1]]

Output: [[[0.265977,1.416304],[1.416304,0.265977],[1,1]]]

Input: x = [[[1,-1],[0,2]],[[2,0],[-1,1]]], num_heads = 1, Wq = [[1,0],[0,0.5]], Wk = [[0.5,0],[0,1]], Wv = [[1,0.25],[-0.5,1]], Wo = [[1,0],[0.5,1]], W1 = [[1,-0.5],[0.5,1]], W2 = [[0.5,0],[0,1]]

Output: [[[1.857823,-1.557072],[-0.762109,3.856215]],[[2.857823,-0.557072],[-1.762109,2.856215]]]

Hint 1

Derive the model width from x.shape[-1].

Hint 2

Reshape each projection to (B, N, num_heads, head_width).

Hint 3

W1.shape[1] already determines the MLP width.

Requirements

  • Use NumPy.
  • Apply pre-normalized multi-head self-attention and its residual.
  • Apply a pre-normalized two-layer GELU MLP and its residual.
  • Use population variance and epsilon 1e-6 for LayerNorm.
  • Return a float64 NumPy array with shape (B, N, D).

Constraints

  • x has shape (B, N, D) and dtype float64.
  • D is divisible by num_heads.
  • Wq, Wk, Wv, and Wo have shape (D, D) and dtype float64.
  • W1 has shape (D, M), and W2 has shape (M, D), both with dtype float64.
  • B, N, D, num_heads, and M are positive.
Prev4 / 6Next