Visualizing Attention Heads in GPT-2

Live Demo: https://llm-attention-viz.streamlit.app

Scientific Context

In the field of Mechanistic Interpretability, understanding why Large Language Models (LLMs) make certain predictions is crucial. While standard metrics (loss, accuracy) measure performance, they treat the model as a black box.

This project aims to peer inside that box by visualizing Self-Attention maps. Specifically, it allows researchers to identify specific "Head Roles", such as:

• Induction Heads: Heads that copy patterns from earlier in the context.
• Previous Token Heads: Heads that consistently attend to the immediate predecessor ($t-1$).

Features

1. Layer-wise Inspection: Navigate through all 12 layers of GPT-2 Small.
2. Entropy Analysis: Calculates the Shannon entropy of attention distributions to measure "focus" vs "diffusion".
3. Prediction Alignment: Correlates internal attention with the final output probability distribution.

Tech Stack

• PyTorch & Hugging Face: Model inference and attention extraction.
• SciPy: Statistical calculations (entropy).
• Streamlit: Interactive UI for rapid experimentation.
• Seaborn: Heatmap visualization.

Theoretical Background

The attention mechanism is defined as: $$ \text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V $$ This tool visualizes the $\text{softmax}(\dots)$ matrix, representing the connection strength between tokens.

How to run locally

pip install -r requirements.txt
streamlit run app.py