← All Posts
From Scratch · C++ · Transformers · Attention Variants · Part 3 of 13

Causal Attention: The Autoregressive Mask

Introduction

In self-attention, every token can look at every other token. But in autoregressive models (language models generating text one token at a time), a token should never see future tokens—it violates the causality assumption that the model is predicting the next word given only past context.

Causal attention enforces this by applying a triangular mask: set all scores for future positions to $-\infty$ before softmax. This post implements the mask and shows how torch::tril, torch::ones, and masked_fill_ work together in LibTorch.

The Causal Principle

In a sequence $[\text{tok}_0, \text{tok}_1, \text{tok}_2, \text{tok}_3]$:

Token 0 can attend to: [token 0] Token 1 can attend to: [token 0, token 1] Token 2 can attend to: [token 0, token 1, token 2] Token 3 can attend to: [token 0, token 1, token 2, token 3] No token can attend to anything in the future.

In matrix form, the attention score matrix $S$ is $(T \times T)$. We want to zero out the upper triangle (future positions). After softmax, those positions contribute zero weight.

LibTorch Implementation

struct CausalAttentionImpl : public torch::nn::Module { torch::nn::Linear wq{nullptr}, wk{nullptr}, wv{nullptr}; int64_t d_model, d_k; CausalAttentionImpl(int64_t d_model_) : d_model(d_model_), d_k(d_model_) { wq = register_module("wq", torch::nn::Linear( torch::nn::LinearOptions(d_model, d_k).bias(false) ) ); wk = register_module("wk", torch::nn::Linear( torch::nn::LinearOptions(d_model, d_k).bias(false) ) ); wv = register_module("wv", torch::nn::Linear( torch::nn::LinearOptions(d_model, d_k).bias(false) ) ); } torch::Tensor forward(torch::Tensor x) { int64_t T = x.size(1); auto Q = wq(x); // (batch, T, d_k) auto K = wk(x); // (batch, T, d_k) auto V = wv(x); // (batch, T, d_k) auto scores = torch::matmul(Q, K.transpose(-2, -1)); scores = scores / sqrt((float)d_k); // Create causal mask: lower triangular matrix of 1s auto causal_mask = torch::tril(torch::ones( {T, T}, x.options() )); // Set future positions to -inf: where mask == 0 scores.masked_fill_(causal_mask == 0, std::numeric_limits<float>::lowest()); auto attn_weights = torch::softmax(scores, -1); auto output = torch::matmul(attn_weights, V); return output; } }; TORCH_MODULE(CausalAttention);

Mask Construction Step-by-Step

torch::ones({T, T}, x.options()) creates a tensor of all 1s with shape $(T \times T)$, on the same device (CPU/GPU) and dtype as the input.

torch::tril(...) keeps only the lower triangle (including diagonal) and zeros the upper triangle. For a $4 \times 4$ matrix, it produces:

[1 0 0 0] [1 1 0 0] [1 1 1 0] [1 1 1 1]

masked_fill_ replaces positions where the condition is true with a fill value. Here, causal_mask == 0 is true in the upper triangle, so we fill those with $-\infty$ (represented by std::numeric_limits<float>::lowest()).

After StepShapeContent
torch::ones(T, T)All 1s
torch::tril(T, T)Lower triangle = 1, upper triangle = 0
masked_fill_(T, T)Lower triangle = score, upper triangle = $-\infty$
torch::softmax(T, T)Lower triangle = softmax weight, upper triangle ≈ 0

Interactive Mask Visualization

▶ Building the Causal Mask

Watch how the mask is constructed step by step.

Complete Shape Trace

StepVariableShapeNotes
Inputx(batch, T, d_model)Token embeddings
Project Q, K, VQ, K, V(batch, T, d_k)Linear projections
Raw scoresscores(batch, T, T)Q @ K^T / sqrt(d_k)
Create onestemp(T, T)torch::ones
Apply trilcausal_mask(T, T)Lower triangle = 1, upper = 0
Masked fillscores(batch, T, T)Upper triangle set to -∞
Softmaxattn_weights(batch, T, T)Rows normalized, future columns ≈ 0
Outputoutput(batch, T, d_k)Weighted sum of values

Key Insights

The mask is static and cheap. Unlike learned parameters, the causal mask is just a one-time construction of the triangular pattern. It's reused across all batches and training steps.
Broadcasting handles the batch dimension. The mask is (T, T) but scores are (batch, T, T). When we apply masked_fill_, broadcasting automatically extends the mask across the batch.
Causal masking is essential for autoregressive models. Without it, the model would learn to cheat during training by looking at future tokens. With the mask, training and generation use the same attention pattern, ensuring consistency.

What Comes Next

Causal attention applies a static mask to a single source (self-attention). But in encoder-decoder models (translation, summarization), we often need cross-attention: the decoder attends to the encoder output, not itself. Next: Cross-Attention.

Previous Softmax Attention