Scaled dot-product attention implementation in PyTorch · CS336 Language Modeling From Scratch

Deliverableartifact

A PyTorch function or class (e.g., attention.py) implementing single-head causal self-attention with tests that confirm correct masking and output shapes.

Goals

Implement a function that computes single-head causal self-attention and passes basic shape tests.
Show numerically that masked positions receive near-zero attention probability.
Verify gradients flow correctly through your attention implementation using autograd checks.
Compare outputs for cases with and without causal masks to illustrate differences.

Resources

📄stanford-cs336/assignment1-basics 📄Spectual/stanford-cs336-a1 📄wind-wing234/cs336-assignment1-basics

Overview: what you’re building

You’ll build a single-head causal self-attention module in PyTorch, plus a few small tests to prove it works:

It will take a batch of token embeddings.
It will project them into queries, keys, and values.
It will compute scaled dot-product attention with a causal mask.
It will return new, context-aware embeddings of the same shape as the input.

We’ll walk through this in four parts, one for each learning point, and end up with one coherent file (e.g. attention.py plus some test code).

A minimal target interface

We’ll aim for something like:

python

attn = CausalSelfAttention(embed_dim=16)
x = torch.randn(2, 5, 16)  # batch=2, time=5, embedding=16
y = attn(x)                # y shape: (2, 5, 16)

1 / 6