Self Attention with a causal mask

import math

import torch
import torch.nn as nn

When we are decoding autoregressively, we obviously don’t know what tokens are coming next. So when we are training a model, we need to mask out future tokens.

inputs = torch.tensor(
    [
        [0.43, 0.15, 0.89],  # Your     (x^1)
        [0.55, 0.87, 0.66],  # journey  (x^2)
        [0.57, 0.85, 0.64],  # starts   (x^3)
        [0.22, 0.58, 0.33],  # with     (x^4)
        [0.77, 0.25, 0.10],  # one      (x^5)
        [0.05, 0.80, 0.55],  # step     (x^6)
    ]
)

class SelfAttention(nn.Module):
    def __init__(self, d_in, d_out, qkv_bias=False):
        super().__init__()
        self.W_query = nn.Linear(d_in, d_out, bias=qkv_bias)
        self.W_key = nn.Linear(d_in, d_out, bias=qkv_bias)
        self.W_value = nn.Linear(d_in, d_out, qkv_bias)

    def forward(self, x):
        queries = x @ self.W_query
        keys = x @ self.W_key
        values = x @ self.W_value

        attn_scores = queries @ keys.T
        attn_weights = torch.softmax(attn_scores / keys.shape[-1] ** 0.5, dim=-1)
        context_vec = attn_weights @ values
        return context_vec


torch.manual_seed(0)
self_attention = SelfAttention(d_in=inputs.shape[1], d_out=2)

queries = self_attention.W_query(inputs)
keys = self_attention.W_key(inputs)
attn_scores = queries @ keys.T
attn_weights = torch.softmax(attn_scores / keys.shape[-1] ** 0.5, dim=-1)
print(attn_weights)

tensor([[0.1762, 0.1616, 0.1619, 0.1662, 0.1712, 0.1630],
        [0.1665, 0.1678, 0.1675, 0.1669, 0.1614, 0.1699],
        [0.1664, 0.1679, 0.1676, 0.1670, 0.1612, 0.1700],
        [0.1654, 0.1679, 0.1677, 0.1669, 0.1631, 0.1689],
        [0.1645, 0.1691, 0.1686, 0.1671, 0.1595, 0.1713],
        [0.1666, 0.1671, 0.1670, 0.1668, 0.1648, 0.1678]],
       grad_fn=<SoftmaxBackward0>)

context_length = attn_scores.shape[0]
mask_simple = torch.tril(torch.ones(context_length, context_length))

mask_simple

tensor([[1., 0., 0., 0., 0., 0.],
        [1., 1., 0., 0., 0., 0.],
        [1., 1., 1., 0., 0., 0.],
        [1., 1., 1., 1., 0., 0.],
        [1., 1., 1., 1., 1., 0.],
        [1., 1., 1., 1., 1., 1.]])

masked_simple = attn_weights * mask_simple

masked_simple

tensor([[0.1762, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000],
        [0.1665, 0.1678, 0.0000, 0.0000, 0.0000, 0.0000],
        [0.1664, 0.1679, 0.1676, 0.0000, 0.0000, 0.0000],
        [0.1654, 0.1679, 0.1677, 0.1669, 0.0000, 0.0000],
        [0.1645, 0.1691, 0.1686, 0.1671, 0.1595, 0.0000],
        [0.1666, 0.1671, 0.1670, 0.1668, 0.1648, 0.1678]],
       grad_fn=<MulBackward0>)

mask = torch.triu(torch.ones(context_length, context_length), diagonal=1)
masked = attn_scores.masked_fill(mask.bool(), -torch.inf)

masked

tensor([[-0.0021,    -inf,    -inf,    -inf,    -inf,    -inf],
        [ 0.0191,  0.0297,    -inf,    -inf,    -inf,    -inf],
        [ 0.0196,  0.0319,  0.0293,    -inf,    -inf,    -inf],
        [ 0.0109,  0.0323,  0.0306,  0.0234,    -inf,    -inf],
        [ 0.0231,  0.0619,  0.0585,  0.0451, -0.0203,    -inf],
        [ 0.0068,  0.0113,  0.0104,  0.0086, -0.0085,  0.0175]],
       grad_fn=<MaskedFillBackward0>)

attn_weights = torch.softmax(masked / keys.shape[-1] ** 0.5, dim=1)

attn_weights

tensor([[1.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000],
        [0.4981, 0.5019, 0.0000, 0.0000, 0.0000, 0.0000],
        [0.3316, 0.3345, 0.3339, 0.0000, 0.0000, 0.0000],
        [0.2476, 0.2514, 0.2511, 0.2498, 0.0000, 0.0000],
        [0.1985, 0.2040, 0.2035, 0.2016, 0.1925, 0.0000],
        [0.1666, 0.1671, 0.1670, 0.1668, 0.1648, 0.1678]],
       grad_fn=<SoftmaxBackward0>)

class CausalAttention(nn.Module):
    def __init__(self, d_in, d_out, context_length, dropout, qkv_bias=False):
        super().__init__()
        self.d_out = d_out
        self.W_query = nn.Linear(d_in, d_out, bias=qkv_bias)
        self.W_key = nn.Linear(d_in, d_out, bias=qkv_bias)
        self.W_value = nn.Linear(d_in, d_out, bias=qkv_bias)
        self.dropout = nn.Dropout(dropout)
        self.register_buffer("mask", torch.triu(torch.ones(context_length, context_length), diagonal=1))

    def forward(self, x):
        b, num_tokens, d_in = x.shape
        keys = self.W_key(x)
        queries = self.W_query(x)
        values = self.W_value(x)

        attn_scores = queries @ keys.transpose(1, 2)
        attn_scores.masked_fill_(self.mask.bool()[:num_tokens, :num_tokens], -torch.inf)
        attn_weights = torch.softmax(attn_scores / keys.shape[-1] ** 0.5, dim=-1)
        attn_weights = self.dropout(attn_weights)

        context_vec = attn_weights @ values
        return context_vec

batch = torch.stack((inputs, inputs), dim=0)

batch

tensor([[[0.4300, 0.1500, 0.8900],
         [0.5500, 0.8700, 0.6600],
         [0.5700, 0.8500, 0.6400],
         [0.2200, 0.5800, 0.3300],
         [0.7700, 0.2500, 0.1000],
         [0.0500, 0.8000, 0.5500]],

        [[0.4300, 0.1500, 0.8900],
         [0.5500, 0.8700, 0.6600],
         [0.5700, 0.8500, 0.6400],
         [0.2200, 0.5800, 0.3300],
         [0.7700, 0.2500, 0.1000],
         [0.0500, 0.8000, 0.5500]]])

d_in = inputs.shape[1]
d_out = 2

ca = CausalAttention(d_in, d_out, context_length, 0.0)

context_vecs = ca(batch)

context_vecs

tensor([[[0.2154, 0.3506],
         [0.1528, 0.4540],
         [0.1266, 0.4926],
         [0.1074, 0.4399],
         [0.0555, 0.4571],
         [0.0641, 0.4228]],

        [[0.2154, 0.3506],
         [0.1528, 0.4540],
         [0.1266, 0.4926],
         [0.1074, 0.4399],
         [0.0555, 0.4571],
         [0.0641, 0.4228]]], grad_fn=<UnsafeViewBackward0>)

Other Formats