Abstract
We analyze transformers from the perspective of iterative inference, seekingto understand how model predictions are refined layer by layer. To do so, wetrain an affine probe for each block in a frozen pretrained model, making itpossible to decode every hidden state into a distribution over the vocabulary.Our method, the \emph{tuned lens}, is a refinement of the earlier ``logitlens'' technique, which yielded useful insights but is often brittle. We test our method on various autoregressive language models with up to 20Bparameters, showing it to be more predictive, reliable and unbiased than thelogit lens. With causal experiments, we show the tuned lens uses similarfeatures to the model itself. We also find the trajectory of latent predictionscan be used to detect malicious inputs with high accuracy. All code needed toreproduce our results can be found athttps://github.com/AlignmentResearch/tuned-lens.