Transformer Encoder Block Forward

Problem

Forward pass of a tiny pre-LN transformer encoder block. With h = LayerNorm(x), compute single-head attention Q = hWq, K = hWk, V = hWv, A = softmax(QK^T / sqrt(d)) V, and the residual a = x + A Wo. Then apply a second LayerNorm, an MLP with ReLU and biases, and the second residual: out = a + relu(LN(a) W1 + b1) W2 + b2.

The signature is encoder_block_forward(x, Wq, Wk, Wv, Wo, W1, b1, W2, b2, eps=1e-5). All projection matrices are square d x d.

Hint

Pre-LN means normalize before attention and before the MLP; both sublayers add back the residual.

Test cases

case 1

in: [[[1,0],[0,1]],[[1,0],[0,1]],[[1,0],[0,1]],[[1,0],[0,1]],[[1,0],[0,1]],[[1,0],[0,1]],[0,0],[[1,0],[0,1]],[0,0]]

expected: [[2.8883532775,-0.8883558715],[-0.8883558715,2.8883532775]]

case 2

in: [[[1,0],[0,1]],[[0,0],[0,0]],[[0,0],[0,0]],[[0,0],[0,0]],[[0,0],[0,0]],[[0,0],[0,0]],[0,0],[[0,0],[0,0]],[0,0]]

expected: [[1,0],[0,1]]

case 3

in: [[[1,2],[3,4]],[[0.5,0],[0,0.5]],[[1,0],[0,1]],[[1,0],[0,1]],[[1,0],[0,1]],[[1,-1],[0,1]],[0.5,-0.5],[[1,0],[0,-1]],[0.1,0.1]]

expected: [[0.1000199994,1.5999844451],[2.1000199994,3.5999844451]]

case 4

in: [[[2,0]],[[1,0],[0,1]],[[1,0],[0,1]],[[1,0],[0,1]],[[1,0],[0,1]],[[1,0],[0,1]],[0,0],[[1,0],[0,1]],[0,0]]

expected: [[3.99999375,-0.999995]]

Starter codePython

import math
def encoder_block_forward(x, Wq, Wk, Wv, Wo, W1, b1, W2, b2, eps=1e-5):
    # Your code here
    pass

4 test cases · run your code against them in the editor