Multi-head attention with Einsum Q = torch.einsum("bnd,di->bni", x, self.W_query) bug

[yujie-jia]

Bug description

ch03/02_bonus_efficient-multihead-attention/mha-implementations.ipynb the implementation in
def forward(self, x):

Q = torch.einsum("bnd,di->bni", x, self.W_query)
K = torch.einsum("bnd,di->bni", x, self.W_key)
V = torch.einsum("bnd,di->bni", x, self.W_value)

should be revised into
Q = torch.einsum("bnd,od->bno", x, self.W_query)
K = torch.einsum("bnd,od->bno", x, self.W_key)
V = torch.einsum("bnd,od->bno", x, self.W_value)
because the initialization is self.W_query = nn.Parameter(torch.randn(d_out, d_in)), first d_out then d_in.

This bug is hidden at first because of the initialization of params
mha_einsum = MHAEinsum(
d_in=embed_dim,
d_out=embed_dim,
...) where d_in == d_out == embed_dim, which hidden the bug. If you change it into d_out = 2 * embed_dim, then there will be an error.

What operating system are you using?

None

Where do you run your code?

None

Environment

[aviralgarg05]

I would like to solve the issue for hacktoberfest

[rasbt]

Yes, feel free to tackle that! That'd be cool.

I think we should then also have unit tests with that fix that check parity with the other implementations. Not just with the in and output dimension but the actual calculations.

Note that this would require initializing with the same random weights (/copying the random weights over).

[aviralgarg05]

Cool, can you mark it under hacktoberfest?

[rasbt]

@aviralgarg05 Sounds good, I added a label. That should be sufficient?

[aviralgarg05]

yup, and on the pr just add the hacktoberfest-accepted!, that would be really helpful for my hacktoberfest progress

[d-kleine]

I think it would be better to use (d_in, d_out)

        self.W_query = nn.Parameter(torch.randn(d_in, d_out))
        self.W_key = nn.Parameter(torch.randn(d_in, d_out))
        self.W_value = nn.Parameter(torch.randn(d_in, d_out))

to align it with the other MHA implementations:

LLMs-from-scratch/ch03/02_bonus_efficient-multihead-attention/mha-implementations.ipynb

Lines 151 to 153 in 7084123

	" self.W_query = nn.Linear(d_in, d_out, bias=qkv_bias)\n",
	" self.W_key = nn.Linear(d_in, d_out, bias=qkv_bias)\n",
	" self.W_value = nn.Linear(d_in, d_out, bias=qkv_bias)\n",

[rasbt]

@d-kleine Do you mean replacing nn.Parameter by nn.Linear? That would be good as nn.Linear comes with good init values by default etc.

[d-kleine]

I mean flipping the weight shape dimensions from d_out, d_in to d_in, d_out