Computes multi-head attention. Supports nested or padded tensors. Args: E_q (int): Size of embedding dim for query E_k (int): Size of embedding dim for key E_v (int): Size of embedding dim for value E_total (int): Total embedding dim of combined heads post i
| 172 | # We can implement this in such a way that it can operate on either padded |
| 173 | # or nested tensors. |
| 174 | class MultiHeadAttention(nn.Module): |
| 175 | """ |
| 176 | Computes multi-head attention. Supports nested or padded tensors. |
| 177 | |
| 178 | Args: |
| 179 | E_q (int): Size of embedding dim for query |
| 180 | E_k (int): Size of embedding dim for key |
| 181 | E_v (int): Size of embedding dim for value |
| 182 | E_total (int): Total embedding dim of combined heads post input projection. Each head |
| 183 | has dim E_total // nheads |
| 184 | nheads (int): Number of heads |
| 185 | dropout_p (float, optional): Dropout probability. Default: 0.0 |
| 186 | """ |
| 187 | def __init__(self, E_q: int, E_k: int, E_v: int, E_total: int, |
| 188 | nheads: int, dropout_p: float = 0.0): |
| 189 | super().__init__() |
| 190 | self.nheads = nheads |
| 191 | self.dropout_p = dropout_p |
| 192 | self.query_proj = nn.Linear(E_q, E_total) |
| 193 | self.key_proj = nn.Linear(E_k, E_total) |
| 194 | self.value_proj = nn.Linear(E_v, E_total) |
| 195 | E_out = E_q |
| 196 | self.out_proj = nn.Linear(E_total, E_out) |
| 197 | assert E_total % nheads == 0, "Embedding dim is not divisible by nheads" |
| 198 | self.E_head = E_total // nheads |
| 199 | |
| 200 | def forward(self, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor) -> torch.Tensor: |
| 201 | """ |
| 202 | Forward pass; runs the following process: |
| 203 | 1. Apply input projection |
| 204 | 2. Split heads and prepare for SDPA |
| 205 | 3. Run SDPA |
| 206 | 4. Apply output projection |
| 207 | |
| 208 | Args: |
| 209 | query (torch.Tensor): query of shape (N, L_t, E_q) |
| 210 | key (torch.Tensor): key of shape (N, L_s, E_k) |
| 211 | value (torch.Tensor): value of shape (N, L_s, E_v) |
| 212 | |
| 213 | Returns: |
| 214 | attn_output (torch.Tensor): output of shape (N, L_t, E_q) |
| 215 | """ |
| 216 | # Step 1. Apply input projection |
| 217 | # TODO: demonstrate packed projection |
| 218 | query = self.query_proj(query) |
| 219 | key = self.key_proj(key) |
| 220 | value = self.value_proj(value) |
| 221 | |
| 222 | # Step 2. Split heads and prepare for SDPA |
| 223 | # reshape query, key, value to separate by head |
| 224 | # (N, L_t, E_total) -> (N, L_t, nheads, E_head) -> (N, nheads, L_t, E_head) |
| 225 | query = query.unflatten(-1, [self.nheads, self.E_head]).transpose(1, 2) |
| 226 | # (N, L_s, E_total) -> (N, L_s, nheads, E_head) -> (N, nheads, L_s, E_head) |
| 227 | key = key.unflatten(-1, [self.nheads, self.E_head]).transpose(1, 2) |
| 228 | # (N, L_s, E_total) -> (N, L_s, nheads, E_head) -> (N, nheads, L_s, E_head) |
| 229 | value = value.unflatten(-1, [self.nheads, self.E_head]).transpose(1, 2) |
| 230 | |
| 231 | # Step 3. Run SDPA |