Pass an intermediate-layer vector through U as if it were the final layer, to read which tokens that layer is currently pointing toward. Apply the final LayerNorm and U directly to the residual vector at layer ℓ to obtain logits. Nothing is trained separately, and the same U is used for every layer. This paper uses this method to read the rank of a specific token at each layer. A single line:
ln_final(resid) @ W_U
Seonglae Cho