[D] Visualizing attention

How do you visualize attention especially when there are many tokens (256,512) with multiple layers and multiple heads? Most visualizations and frameworks I’ve tried fail when there are more than 100 tokens.

This is quite important since most nlp and vision models now use multi head attention.

attention machinelearning