Causal Attention

Updated 1st November 2026Started 3rd April 2026

When a model generates text autoregressively it cannot know the tokens that come next. So during training each position may only attend to itself and the positions before it. We enforce that by masking out the future in the attention weights.

This builds on Self Attention.

import torch
import torch.nn as nn
 
inputs = torch.tensor(
    [
        [0.43, 0.15, 0.89],  # Your
        [0.55, 0.87, 0.66],  # journey
        [0.57, 0.85, 0.64],  # starts
        [0.22, 0.58, 0.33],  # with
        [0.77, 0.25, 0.10],  # one
        [0.05, 0.80, 0.55],  # step
    ]
)
 
torch.manual_seed(0)
W_query = nn.Linear(3, 2, bias=False)
W_key = nn.Linear(3, 2, bias=False)
 
queries = W_query(inputs)  # nn.Linear layers are called, not matrix-multiplied
keys = W_key(inputs)
attn_scores = queries @ keys.T
attn_weights = torch.softmax(attn_scores / keys.shape[-1] ** 0.5, dim=-1)

Every row of attn_weights spreads its weight over all six tokens, including the ones in the future.

A simple mask, and why it is not enough

torch.tril gives a lower-triangular matrix of ones:

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

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

Multiplying the softmax output by it zeroes the future positions:

masked_simple = attn_weights * mask_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],
        ...

But the rows no longer sum to 1: the first row sums to 0.18. We would have to renormalise each row:

masked_simple / masked_simple.sum(dim=-1, keepdim=True)

Masking before softmax

The cleaner approach is to mask the scores with $-\infty$ before softmax. Since $e^{-\infty} = 0$ , the masked positions get exactly zero weight and each row still sums to 1, with no second pass:

mask = torch.triu(torch.ones(context_length, context_length), diagonal=1)
masked = attn_scores.masked_fill(mask.bool(), -torch.inf)
attn_weights = torch.softmax(masked / keys.shape[-1] ** 0.5, dim=-1)

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

As a module

Two additions over plain self-attention: the mask is registered as a buffer (so it moves with the module between devices but is not a trainable parameter), and dropout is applied to the attention weights. The mask is sliced to the actual sequence length, so the module works on any input up to context_length tokens, and on batches.

class CausalAttention(nn.Module):
    def __init__(self, d_in, d_out, context_length, dropout, 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, 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)
 
        return attn_weights @ values
 
 
batch = torch.stack((inputs, inputs), dim=0)  # shape (2, 6, 3)
ca = CausalAttention(d_in=3, d_out=2, context_length=6, dropout=0.0)
ca(batch).shape  # torch.Size([2, 6, 2])

Original notebook

Linked from

Self Attention

Revision log

1st November 2026Migrated from ai-notes

Tags: AI