Might be cool to try different colors for the different attention heads instead of summing them across all attention heads, and making the backgrounds composed of stacked color layers? So you can see how each attention head attends to tokens individually.