search

soniajoseph / ViT-Prisma

ViT Prisma is a mechanistic interpretability library for Vision Transformers (ViTs).
Other
165 stars 18 forks source link

Add colorbar to attention head visualization + key / query labels on the attention head visualization axes #74

Open soniajoseph opened 7 months ago

soniajoseph commented 7 months ago

Currently, the attention head visualization does not have a colorbar so it's hard to get a sense of what the values mean.

The key/query axes are also not labeled.

Code is here: https://github.com/soniajoseph/ViT-Prisma/blob/main/src/vit_prisma/visualization/visualize_attention_js.py See plot_javascript in the above file as the function the user calls for their attention heads.