GPT-2 architecture

[d-kleine]

GPT-2 architecture

[rasbt]

Thanks for the note, but I think my figure is correct. The best source for this would be the original GPT-2 model implementation, which you can find here:

https://github.com/openai/gpt-2/blob/9b63575ef42771a015060c964af2c3da4cf7c8ab/src/model.py#L123-L130

def block(x, scope, *, past, hparams):
    with tf.variable_scope(scope):
        nx = x.shape[-1].value
        a, present = attn(norm(x, 'ln_1'), 'attn', nx, past=past, hparams=hparams)
        x = x + a
        m = mlp(norm(x, 'ln_2'), 'mlp', nx*4, hparams=hparams)
        x = x + m
        return x, present

Unless I am reading this incorrectly, the shortcut path starts before not after LayerNorm.

[d-kleine]

Yes, I think you're right, including about the official code. For Pre-LN, the residual connection starts with the shortcut before the layer normalization step, which is nicely illustrated in your figure here and the official illustration of Pre-LN on the far right side:

Apologies - I have only found a few graphics illustrating the GPT-2 architecture, and it seems like all of them were incorrect about the Pre-LN step (even on Wikipedia, surprisingly).

Thanks again, and issue can be closed then 🙂

[rasbt]

Nice, I am glad that it all makes sense now (and also looks correct haha). Thanks for raising this issue though, it's always good to have multiple sets of eyes on these things!