Pre-LN vs Post-LN Order

Problem

Compare pre-LN and post-LN transformer sublayer orderings for single-head attention. Pre-LN: h = LN(x), attn = softmax(hWq (hWk)^T / sqrt(d)) (hWv), out = x + attn Wo. Post-LN: compute attention from raw x, add the residual, then normalize: out = LN(x + attn Wo).

The signature is pre_post_ln_order(x, Wq, Wk, Wv, Wo, eps=1e-5). Return [pre_ln_output, post_ln_output].

Hint

Pre-LN normalizes the sublayer input; post-LN normalizes the sublayer output after the residual add.

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]]]

expected: [[[1.8883558715,-0.8883558715],[-0.8883558715,1.8883558715]],[[0.9999888539,-0.9999888539],[-0.9999888539,0.9999888539]]]

case 2

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

expected: [[[0.0000199994,2.9999800006],[2.0000199994,4.9999800006]],[[-0.999995,0.999995],[-0.999995,0.999995]]]

case 3

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

expected: [[[-0.9999600012,2],[1.0000399988,4]],[[0.9999907657,-0.9999907657],[0.99999111,-0.99999111]]]

case 4

in: [[[3,-1]],[[1,0],[0,1]],[[1,0],[0,1]],[[1,0],[0,1]],[[1,0],[0,1]]]

expected: [[[3.99999875,-1.99999875]],[[0.9999996875,-0.9999996875]]]

Starter codePython

import math
def pre_post_ln_order(x, Wq, Wk, Wv, Wo, eps=1e-5):
    # Your code here
    pass

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