Logit Lens presents results in the form of heatmaps. Currently, five visualization modes are supported, and additional ones can be added by expanding the data and/or chart options in generate_and_visualize.
- Current Layer Logit/Prob View: Text shows the current layer’s top-5 tokens; color encodes the logit/probability of the current layer’s top-1 token.
- Last Layer Logit/Prob View: Text shows the current layer’s top-1 token; color encodes the logit/probability of the final layer’s token at the current layer.
- Last Layer Rank View: Both text and color represent the vocabulary rank of the final layer’s token within the current layer.
Features
- Supports inspection of attention, attention residual, MLP, and MLP residual.
- Allows inspection at “every N layers”, while always including the final layer.
- Provides an option to normalize before unembedding.
- Supports batch inference: each input generates one heatmap, trimmed to effective length according to the EOS token.
- Currently does not support teacher-forcing inference, viewing by custom intervals/indices, inspection of prompts or the embedding layer.
Code is implemented based on LLaMA. To adapt to different architectures, modify _register_layer_hooks.