# Self Attention with a causal mask


<!-- WARNING: THIS FILE WAS AUTOGENERATED! DO NOT EDIT! -->

``` python
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.

``` python
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)
    ]
)
```

``` python
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)
```

``` python
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>)

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

``` python
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.]])

``` python
masked_simple = attn_weights * mask_simple
```

``` python
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>)

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

``` python
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>)

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

``` python
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>)

``` python
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
```

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

``` python
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]]])

``` python
d_in = inputs.shape[1]
d_out = 2

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

``` python
context_vecs = ca(batch)
```

``` python
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>)