I took those small TinyStories language models and wanted to explore attention mechanism in detail and now you can explore it too- with my app:
show-me-your-attention.streamlit.app
Its interactive.
Actually my main goal is to dive into some sort of interpretability and I would appreciate suggestions, what crazy things i should do- apply regressor on some parts of attention heads etc. 