Is It Necessary to Transpose Dimensions in Multi-Head Attention or Can We Reshape Directly?

[rohanwinsor]

Is it necessary the implementation should be like this

class MultiHeadAttention(nn.Module):
    def __init__(self, d_in, d_out, context_length, dropout, num_heads, qkv_bias=False):
        super().__init__()
        ...

    def forward(self, x):
        ...
        keys = keys.view(b, num_tokens, self.num_heads, self.head_dim) 
        values = values.view(b, num_tokens, self.num_heads, self.head_dim)
        queries = queries.view(b, num_tokens, self.num_heads, self.head_dim)
        # Transpose: (b, num_tokens, num_heads, head_dim) -> (b, num_heads, num_tokens, head_dim)
        keys = keys.transpose(1, 2)
        queries = queries.transpose(1, 2)
        values = values.transpose(1, 2)
        ...

over casting them directly to the dimension that we want?

class MultiHeadAttention(nn.Module):
    def __init__(self, d_in, d_out, context_length, dropout, num_heads, qkv_bias=False):
        super().__init__()
        ...

    def forward(self, x):
        ...
        keys = keys.view(b, self.num_heads, num_tokens, self.head_dim) 
        values = values.view(b, self.num_heads, num_tokens, self.head_dim)
        queries = queries.view(b, self.num_heads, num_tokens, self.head_dim)
        ...

Got the code from the book as well as from here

[rasbt]

Hey there,

this is a really good question. At first glance, it looks like this should work because the dimensions would be the same. But note that reshaping and transposing are slightly different in terms of how the matrices get arranged for the matrix multiplication that follows. So, no, those are not interchangeable. Actually, there was the same question in #167, where the answer may give a bit more concrete insights.

Anyways, thanks for asking!

[rohanwinsor]

Tks a lot for ur answer. Will refer there.