Attention rollout for MViT -- would it make sense visualising it for AR tasks?

I'm curious if anybody tried to do the attention rollout for the MViT model in the AR task.

I played around with the `drop_path` outputs for some clips but nothing spectacular (perhaps I'm doing something terribly wrong). Basically, due to embedding dimension changes I find it troublesome to project differently scaled attention maps onto the original clips -- but for the same-size embedding it is a fairly simple TWH reshape.

However, I'm more interested in actual interpretation -- perhaps somebody has tried some other visualisation techniques that are more suitable for this architecture?