<jupyter_start><jupyter_text>To talk about masking one must have training/inference in mind, the GPT models are autoregressive, which is pretty much just a fancy way of saying, "give it a few words/tokens, let it predict the next one". The huge scaling benefit of this lies in the unsupervised nature (no labels needed, to wit, "training data is also test data") allows autoregressive models to train on any amount of text data available, contents from the whole internet is the limit (and has indeed been attempted: https://commoncrawl.org/).  Masking is just a technicality to ensure that the autoregressive property holds by not allowing the model to look ahead. While it is not a deep concept, its implementation can be a source of confusion. So let's get the questions-party started.There can be other masks for different needs (e.g. BERT). We focus on the autoregressive one here. >   Q1. Why would the model look ahead? This is due to the training process and the very heart of attention (processing many tokens at once). Now picture or refer to a diagram for the attention mechanism. It receives a input sequence of n vectors, each is called a token, representing a word. It outputs n vectors, which eventually will be transformed into n predictions for the next-token, for which the "ground truth" should be each word's next neighhour in the input sequence, i.e. the solution is right in the input. It is obvious that we will reveal the ground truth for attention to overfit on if we let it see the whole sequence, so we must have a way to hide the solution so the model cannot cheat -- since during inference we cannot supply the next tokens.Another way to look at this is that the feature of attention layer that enables it to take in multiple tokens at once (a key advantage over RNNs) means we need something to preserve causality which RNN naturally does, and the trick we use is masking.>   Q2. So what do we do?We add an extra step inside attention so that the i th output token can only receive information from 0 to i th (inclusive) input tokens.>   Q3. What extra step? And why is it positioned right before softmax? And why is the mask a triangular matrix? $-\infty$?Actually, the core logic is almost trivial if we recast this into a translation issue from vector to matrix notion. Masking -- the artifect of matrix notation to capture the vector sums.The logic with summuing vectors to preserve causality is almost trivial, it's only when we want to translate this into matrix notation that it begins to look like an advanced trick, but if we view it this way, the core logic is as trivial as saying "drop any info from future tokens", or "if (model about to cheat): don't. "Let's use a seq. of len. 2 to illustrate the process. Attention is really just about seqeunce to sequence which looks like:[x1, x2] -> attention layer -> [y1, y2]Let's say y1 is the first to come out, it should only see x1, and y2 should be able to see x1, x2 as discussed above. Then we need to mask out info, let's look closerIf we write attention as a sum$$y_1 = w_1 \cdot v_1 \\ \sim \exp(k_1 \cdot q_1)  v_1\\= \exp(k_1 \cdot q_1)  v_1 + 0 \cdot v_2 \\= exp(k_1 \cdot q_1)  v_1 + \exp(-\infty) \cdot v_2 $$$$y_2 = w_3 \cdot v_1+ w_4 \cdot v_2 \\\sim \exp(k_1 \cdot q_2)  v_1 + \exp(k_2 \cdot q_2)  v2 $$where ~ means proportional to since we don't need to other constants to illustrate the point, and $w_i$ are weights for matrices below.At the point, we ask, how do we turn this into matrix operation? In matrix form, it looks like:Y = W*Vwhere,$$ W = \begin{bmatrix}w_1 & 0\\w_3 & w_4\end{bmatrix}$$ $$ Y = \begin{bmatrix} y_1 \\ y_2  \end{bmatrix}$$$$ V = \begin{bmatrix} v_1 \\ v_2  \end{bmatrix}$$Now, we know M is row-wise softmaxed, so it means before softmax it must look like,$$ W_{raw} = \begin{bmatrix}w_{raw1} & -\infty\\w_{raw3} & w_{raw4}\end{bmatrix}$$ This is why we want to apply a "mask" to make the elements above diagonal in $W_{raw}$ to be $-\infty$, which will translate into dropping any vectors derived from future tokens in matrix notation.Now we see clearly where the mask comes in -- right before the softmax step, everything else is business as usual.<jupyter_code>#Let's demonstrate the masking in action
#first we start as usual until the softmax step

import torch
import torch.nn as nn

emb_dim = 4
seq_len = 5

#input tokens:
x = torch.tensor([0,1,2,3,4]) 
embedding = nn.Embedding(20, emb_dim) #vocab size 20, emb dim 4
x = embedding(x)

#attention matrices
M_K, M_Q, M_V = [nn.Linear(emb_dim, emb_dim, bias=False) for _ in range(3)]
K, Q, V = [M(x) for M in [M_K, M_Q, M_V ]]
W_raw = Q@(K.transpose(-1,-2))

# stop before next steps
# W = F.softmax(W_raw, dim=-2)
# Y = W@V
W_raw
seq_len = 5
ones = torch.ones((seq_len, seq_len), dtype=torch.uint8)
mask = torch.triu(ones, diagonal=1)
mask
# in Pytorch we can apply mask with another tensor as boolean 
# i.e. if mask[i][j] == 1: W_raw[i][j] = float('-inf')
W_raw[mask] = float('-inf')
W_raw
import torch.nn.functional as F

# proceed with the next steps to finish attention as usual
W = F.softmax(W_raw, dim=-1)
W
V
Y = W@V
Y
# To summarize, we can write the masked version of attention as:

def self_attention_with_mask(x, emb_dim, seq_len):
    M_K, M_Q, M_V = [nn.Linear(emb_dim, emb_dim, bias=False) for _ in range(3)]
    K, Q, V = [M(x) for M in [M_K, M_Q, M_V ]]
    W_raw = Q@(K.transpose(-1,-2))
    # == masking begins ==
    ones = torch.ones((seq_len, seq_len), dtype=torch.uint8)
    mask = torch.triu(ones, diagonal=1)
    W_raw[mask] = float('-inf')
    # == masking ends ==
    W = F.softmax(W_raw, dim=-1)
    Y = W@V
    return Y
self_attention_with_mask(x, emb_dim, seq_len) 
#result will look different each time due to re initialization of nn.Linear<jupyter_output><ipython-input-47-442fd643b0fd>:10: UserWarning: indexing with dtype torch.uint8 is now deprecated, please use a dtype torch.bool instead. (Triggered internally at  /pytorch/aten/src/ATen/native/IndexingUtils.h:25.)
  W_raw[mask] = float('-inf')
{'fraction_non_alphanumeric': 0.08971589965110484, 'fraction_numerical': 0.015284931051669712, 'mean_word_length': 4.03765690376569, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 10, 'https://': 1, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '1', 'dir': 'jupyter-structured-clean-dedup', 'id': '662489', 'license': 'MIT', 'n_tokens_mistral': 1925, 'n_tokens_neox': 1846, 'n_words': 913, 'path': '3. Masking.ipynb', 'repo_name': 'rayding24/First-Principles-Transformers', 'size': '6019'}