In my PyTorch implementation of multi-head attention, i have those in __init__() class MultiHeadAttentionLayer(nn.Module): def __init__(self,d_in,d_out,context_length,dropout,num_heads,use_bias=False): super().__init__() self.d_out=d_out self.num_heads=num_heads # In multi-head attention, the output dimension (d_out) is split across multiple attention heads. # Each head processes a portion of the total output dimensions independently before being concatenated back together. self.head_dim=d_out//num_heads self.query_weight = nn.Linear(d_in, d_out, bias=use_bias) self.key_weight = nn.Linear(d_in, d_out, bias=use_bias) self.value_weight = nn.Linear(d_in, d_out, bias=use_bias) this is the forward method def forward(self,x): batch_size,sequence_length,d_in=x.shape keys=self.key_weight(x) queries=self.query_weight(x) values=self.value_weight(x) # RESHAPING # .view() is a PyTorch tensor method that reshapes a tensor without changing its underlying data. It returns a new tensor with the same data but in a different shape. keys=keys.view(batch_size,sequence_length,self.num_heads,self.head_dim) values=values.view(batch_size,sequence_length,self.num_heads,self.head_dim) queries=queries.view(batch_size,sequence_length,self.num_heads,self.head_dim) I understand that d_out is split across multiple attention heads, but I'm not entirely sure why this reshaping is necessary. How does adding num_heads as a new dimension affect the computation of attention, and what would happen if we skipped this step and kept the shape as "batch_size,sequence_length,d_in" Continue reading...
