TINYGPT

A token and its position enter at the left, each is embedded using the embedding tables of the model. The embeddings are then summed to form the source of the residual stream.

The stream passes through the twelve transformer blocks, each block with its own weights, with a token selected after a final normalisation and a tied unembedding into 50,257 logits (the unembedding matrix is the transpose of the token embedding matrix).

Both lenses decode the residual stream after embedding and after each block.

The logit lens projects a hidden state onto the model's vocabulary and shows what the output would be if the model were forced to answer after that block.

The Jacobian lens maps it first through matrix J, the fitted average sensitivity of the final state to that block, before asking the same question — the logit lens is the case where J is the identity.