Studying "Build a Large Language Model (From Scratch) by Sebastian Raschka. I was trying to evaluate "Multi Head Simple Implementation" and Alternative Implmentation as mentioned in the book using following code. As per my understanding both method should give me same output. But in my following code values are different (i.e context_vec_Final in Simple implementation and context_vec in Alternative implementation should be same). Can you please help me to correct my mistake. For simplicity purpose, I discarded Causal Attention in this example. #Simple Implementation import torch import torch.nn as nn torch.manual_seed(1223) batch = 1 num_tokens = 2 d_out = 6 num_heads = 3 head_dim = d_out // num_heads out_proj = nn.Linear(d_out*num_heads, d_out*num_heads) print (f'{batch = } {num_tokens = } {d_out = } {num_heads = } {head_dim = }') q = torch.rand(batch,num_tokens,d_out) k = torch.rand(batch,num_tokens,d_out) v = torch.rand(batch,num_tokens,d_out) print(f'{q.shape = } {k.shape = } {v.shape = }') context_vec = [] for _ in range(num_heads) : attn_scores1 = q @ k.transpose(1, 2) context_vec1 = attn_scores1 @ v context_vec.append(context_vec1) context_vec_Final = out_proj(torch.cat (context_vec, dim = -1)) print(f'{context_vec_Final.shape = } {context_vec_Final = }') Alternative Variation import torch torch.manual_seed(1223) batch = 1 num_tokens = 2 d_out = 6 num_heads = 3 head_dim = d_out // num_heads print (f'{batch = } {num_tokens = } {d_out = } {num_heads = } {head_dim = }') q = torch.rand(batch,num_tokens,d_out) k = torch.rand(batch,num_tokens,d_out) v = torch.rand(batch,num_tokens,d_out) print(f'{q.shape = } {k.shape = } {v.shape = }') print(f'{q = }') #The key operation is to split the d_out dimension into num_heads and head_dim using view option #(b, num_tokens, d_out) is reshaped to dimension (b, num_tokens, num_heads, head_dim) #d_out = num_heads * head_dim q = q.view(batch,num_tokens,num_heads,head_dim) k = k.view(batch,num_tokens,num_heads,head_dim) v = v.view(batch,num_tokens,num_heads,head_dim) print(f'{q.shape = } {k.shape = } {v.shape = }') #Now transpose (b, num_tokens, num_heads, head_dim) to (b, num_heads, num_tokens, head_dim) #Logically now we have num_heads of num_tokens x head_dim... #Now we have num_heads (multiple head ) of num_tokens x head_dim q = q.transpose(1, 2) k = k.transpose(1, 2) v = v.transpose(1, 2) print(f'{q.shape = } {k.shape = } {v.shape = }') attn_scores = q @ k.transpose(2, 3) print(f'{attn_scores.shape = } {attn_scores = }') context_vec = (attn_scores @ v) print(f'{context_vec.shape = }') #Now we are transposing (b, num_heads, num_tokens, head_dim) back to (b, num_tokens, num_heads, head_dim) context_vec = context_vec.transpose(1, 2) print(f'{context_vec.shape = }') context_vec = context_vec.contiguous().view(batch, num_tokens, d_out) print(f'{attn_scores.shape = } {context_vec.shape = }') print (f'{context_vec.shape = } {context_vec = }') Continue reading...
