I am trying to make sure I understand the implementation of torch.nn.MultiheadAttention(), as I want to use it for autoregressive sampling in a “decoder-only” image transformer.
In such case, one tries to predict the next pixel by attending on all previous pixels. Say my pixels are just binary (0,1).
In the Multihead attention forward function, the input is supposed to be:
query: [target length, batch size, embed dim]
key: [sequence length, batch size, embed dim]
value: [sequence length, batch size, embed dim]
What is “target length” referring to in this context? If I am just trying to predict the next pixel at each pass, is this just length 1?
Also, is there any way to do batch_first with this module?
Thanks for your help!