Source
Belrose et al. (2023), LIT-569. Read as NOTE-308. It refines the logit lens, LIT-570.
The practice
When asking what a pre-LN transformer “predicts” at an intermediate layer:
- Do not trust the logit lens on a model you have not checked it on. On BLOOM and OPT-125M its top-1 at most layers is the input token. On GPT-Neo it is biased toward some tokens until the last layer. A clean-looking trajectory on GPT-2 says nothing about the next model
- Train a translator per layer:
h ↦ A_ℓ h + b_ℓ, identity-initialized, minimizing KL from the model’s own final distribution on held-out pretraining-like text, then apply the model’s final norm and unembedding. Distilling toward the model, not toward labels, keeps the lens from learning what the model does not know - Optimize it properly. The paper’s later version says Muon reaches much lower KL than its SGD setup, and that its own lenses were undertrained
- Reuse lenses across fine-tunes of the same base model. The cost is at most 0.3 bits per byte in the paper’s one test
Conditions
- Pre-LN transformers
- It is a reading tool, not a detector. The paper’s anomaly-detection use loses to a Mahalanobis baseline on most tasks, and for eliciting hidden knowledge it is not reliably better than the logit lens