Spaces:
Running on Zero
Running on Zero
| """ | |
| Single-head causal self-attention. | |
| Learns relationships between tokens in the input sequence. | |
| Input Representation (X) | |
| │ | |
| ▼ | |
| Q = X @ Wq | |
| K = X @ Wk | |
| V = X @ Wv | |
| │ | |
| ▼ | |
| Scores = Q @ Kᵀ | |
| │ | |
| ▼ | |
| Scale by √d | |
| │ | |
| ▼ | |
| Apply Causal Mask | |
| │ | |
| ▼ | |
| Softmax | |
| │ | |
| ▼ | |
| Attention Weights | |
| │ | |
| ▼ | |
| Output = Attention Weights @ V | |
| """ | |
| from __future__ import annotations | |
| import math | |
| import torch | |
| import torch.nn as nn | |
| class Attention(nn.Module): | |
| """ | |
| Single-head causal self-attention. | |
| Input: | |
| (batch_size, context_length, d_model) | |
| Output: | |
| (batch_size, context_length, d_model) | |
| """ | |
| def __init__(self, d_model: int, context_length: int,): | |
| super().__init__() | |
| # -------------------------------------------------- | |
| # Learnable projection layers. | |
| # | |
| # Each nn.Linear layer contains a learnable weight | |
| # matrix that projects the input representation into | |
| # a different vector space. | |
| # | |
| # Transformer paper notation: | |
| # | |
| # Query Projection -> Wq | |
| # Key Projection -> Wk | |
| # Value Projection -> Wv | |
| # | |
| # During the forward pass: | |
| # | |
| # Q = X @ Wq | |
| # K = X @ Wk | |
| # V = X @ Wv | |
| # | |
| # In PyTorch, the multiplication is performed by | |
| # calling the corresponding nn.Linear layer. | |
| # -------------------------------------------------- | |
| self.query = nn.Linear( | |
| in_features=d_model, | |
| out_features=d_model, | |
| bias=False, | |
| ) | |
| self.key = nn.Linear( | |
| in_features=d_model, | |
| out_features=d_model, | |
| bias=False, | |
| ) | |
| self.value = nn.Linear( | |
| in_features=d_model, | |
| out_features=d_model, | |
| bias=False, | |
| ) | |
| # -------------------------------------------------- | |
| # Causal Mask | |
| # | |
| # Prevents tokens from attending to future tokens. | |
| # | |
| # The mask is created once during initialization | |
| # and reused during every forward pass. | |
| # -------------------------------------------------- | |
| self.register_buffer( | |
| "mask", | |
| torch.triu( | |
| torch.ones( | |
| context_length, | |
| context_length, | |
| dtype=torch.bool, | |
| ), | |
| diagonal=1, | |
| ), | |
| ) | |
| def forward(self, | |
| x: torch.Tensor, | |
| padding_mask: torch.Tensor | None = None, | |
| ) -> torch.Tensor: | |
| """ | |
| Parameters | |
| ---------- | |
| x | |
| Input representation. | |
| Shape: | |
| (batch_size, context_length, d_model) | |
| padding_mask | |
| Optional mask indicating which positions contain | |
| real tokens. | |
| True -> real token | |
| False -> padding | |
| Shape: | |
| (batch_size, context_length) | |
| """ | |
| # -------------------------------------------------- | |
| # Step 1 | |
| # Create Query (Q), Key (K) and Value (V) | |
| # | |
| # Mathematically: | |
| # | |
| # Q = X @ Wq | |
| # K = X @ Wk | |
| # V = X @ Wv | |
| # | |
| # where X is the input representation and Wq, Wk and | |
| # Wv are the learnable weight matrices contained in | |
| # the projection layers above. | |
| # -------------------------------------------------- | |
| Q = self.query(x) | |
| K = self.key(x) | |
| V = self.value(x) | |
| # -------------------------------------------------- | |
| # Step 2 | |
| # Compute Attention Scores | |
| # | |
| # Mathematically: | |
| # | |
| # Scores = Q @ Kᵀ | |
| # | |
| # Every Query vector is compared with every Key vector | |
| # using the dot product, producing one attention score | |
| # for every Query–Key pair. | |
| # | |
| # scores[i, j] answers: | |
| # | |
| # "How much should token i attend to token j?" | |
| # -------------------------------------------------- | |
| scores = Q @ K.transpose(-2, -1) | |
| # -------------------------------------------------- | |
| # Step 3 | |
| # Scale Attention Scores | |
| # | |
| # Mathematically: | |
| # | |
| # Q @ Kᵀ | |
| # Scores = -------- | |
| # √d | |
| # | |
| # where d is the embedding dimension (d_model). | |
| # | |
| # Scaling keeps the magnitude of the attention | |
| # scores relatively stable regardless of the | |
| # embedding dimension. | |
| # -------------------------------------------------- | |
| scores = scores / math.sqrt(Q.size(-1)) | |
| # -------------------------------------------------- | |
| # Step 4 | |
| # Apply Causal Mask | |
| # | |
| # Prevent each token from attending to future tokens. | |
| # | |
| # The masked positions are assigned negative infinity. | |
| # After applying softmax, they will receive a | |
| # probability of zero. | |
| # -------------------------------------------------- | |
| context_length = scores.size(-1) | |
| # -------------------------------------------------- | |
| # Step 4 | |
| # Apply Causal Mask | |
| # "Don't look into the future." | |
| # -------------------------------------------------- | |
| scores = scores.masked_fill( | |
| self.mask[:context_length, :context_length], | |
| float("-inf"), | |
| ) | |
| # -------------------------------------------------- | |
| # Apply Padding Mask | |
| # | |
| # True -> real token | |
| # False -> PAD | |
| # | |
| # We mask the KEY positions, meaning: | |
| # | |
| # "Do not allow any token to attend | |
| # to a PAD token." | |
| # -------------------------------------------------- | |
| if padding_mask is not None: | |
| scores = scores.masked_fill( | |
| ~padding_mask[:, None, :], | |
| float("-inf"), | |
| ) | |
| # -------------------------------------------------- | |
| # Step 5 | |
| # Apply Softmax | |
| # | |
| # Convert the attention scores into probabilities. | |
| # | |
| # Softmax is applied independently to each row of | |
| # the attention score matrix so that every row sums | |
| # to one. | |
| # | |
| # Each probability represents how much attention a | |
| # Query token pays to each Key token. | |
| # -------------------------------------------------- | |
| attention = torch.softmax( | |
| scores, | |
| dim=-1, | |
| ) | |
| # -------------------------------------------------- | |
| # Step 6 | |
| # Compute the Weighted Sum | |
| # | |
| # Mathematically: | |
| # | |
| # Output = Attention Weights @ V | |
| # | |
| # The attention weights determine how much | |
| # information to retrieve from every Value vector. | |
| # | |
| # Each output vector is a weighted combination of | |
| # all Value vectors visible to the current Query. | |
| # -------------------------------------------------- | |
| output = attention @ V | |
| return output | |