Lab 03 - Causal Mask and Attention Scores

Overview

Lab 03: Causal Mask and Attention Scores

Annotated code reading lab. Running code is optional.

Related handbook section

Causal Mask and Attention Scores

This lab maps directly to the handbook section. Read the related handbook section first, then use the lab page and starter file to connect the concept to concrete variables, shapes, APIs, and interview-ready explanations.

Open handbook section Open starter file

Concept Goal

Causal Mask and Attention Scores

Read how a causal mask turns parallel training into left-to-right prediction semantics.

Mental Model

Mechanism to keep in mind

`scores[i][j]` means token i reading token j.
The upper triangle is blocked before softmax.
Training can compute all positions at once while preserving autoregressive visibility.

Annotated Code Preview

Starter preview

Excerpt from code/lab-03-causal-mask-attention/causal_mask.py. The linked starter file is the source of truth.

Open starter file

# Causal Mask and Attention Scores
# Annotated reading material. Running this file is optional.
# Source-of-truth focus: Read how a causal mask turns parallel training into left-to-right prediction semantics.

seq = 5
scores_shape = (seq, seq)
causal_mask = [[j <= i for j in range(seq)] for i in range(seq)]
# False entries are filled with a large negative value before softmax.
visible_to_token_3 = [j for j, allowed in enumerate(causal_mask[3]) if allowed]

# What to explain while reading:
# - scores[i][j] means token i reading token j.
# - The upper triangle is blocked before softmax.
# - Training can compute all positions at once while preserving autoregressive visibility.
#
# Common traps:
# - Causal masking is not an inference-only trick.
# - A future token can be present in the tensor but invisible to the current position.

Line-by-line Explanation

What each block is doing

Setup / contract: `scores[i][j]` means token i reading token j.
Main transition: The upper triangle is blocked before softmax.
Interview hook: Training can compute all positions at once while preserving autoregressive visibility.

What to Notice

Reading checkpoints

The mask changes probability mass, not token order.
Masking happens before softmax.
Padding masks and causal masks solve different problems.

Common Misunderstandings

What this lab prevents

Causal masking is not an inference-only trick.
A future token can be present in the tensor but invisible to the current position.

Interview Explanation

How to say it out loud

Read how a causal mask turns parallel training into left-to-right prediction semantics. Then explain the code by naming the state being transformed, the axis or shape that matters, and the tradeoff that would appear in a real system.

External intuition notes

Additional intuition

Use official docs and papers for API behavior and factual claims; use blogs only to improve the mental picture.
If support matrices, performance behavior or backend choices are version-sensitive, check current docs before repeating them.
A strong interview answer names the state object, the shape or axis it changes, and the tradeoff it creates.

InfraLens