Module 6 Assignment: Attention mechanism analysis

Module 6 Assignment: Attention mechanism analysis#

Theme#

Attention and transformers

Scenario#

A platform team is evaluating whether a transformer block can support document understanding without hiding how token interactions are weighted.

Exercises#

  1. Explain query, key, value, and attention weights using a concrete sequence.

  2. Run the starter attention computation and inspect the attention matrix.

  3. Describe how masking changes autoregressive generation.

  4. Compare self-attention with recurrence for parallelism and dependency modeling.

Evidence Requirements#

  • A short explanation of the data, tensors, objective, and evaluation signal used in the starter experiment.

  • At least one meaningful modification to the starter code, with the changed variable named explicitly.

  • A comparison against the unmodified starter result or another defensible baseline.

  • A limitation statement that separates what the toy experiment demonstrates from what a production model would require.

Submission#

Submit a 600-900 word technical memo plus code, plots, tables, or shape traces needed to support your claims. The memo should read like a review artifact for another AI practitioner: concise, reproducible, and honest about uncertainty.

Rubric Focus#

  • Technical correctness and appropriate neural-network vocabulary.

  • Evidence from the starter experiment or a documented extension.

  • Connection between design choices and data/problem structure.

  • Clear treatment of limitations, failure modes, or next experimental gates.

import torch
import torch.nn.functional as F

torch.manual_seed(6)
Q = torch.randn(1, 4, 8)
K = torch.randn(1, 4, 8)
V = torch.randn(1, 4, 8)
scores = Q @ K.transpose(-2, -1) / (Q.size(-1) ** 0.5)
weights = F.softmax(scores, dim=-1)
context = weights @ V
print("attention weights shape:", tuple(weights.shape))
print(weights[0].round(decimals=3))
print("context shape:", tuple(context.shape))

attention weights shape: (1, 4, 4)
tensor([[0.0160, 0.3920, 0.5170, 0.0760],
        [0.3150, 0.1920, 0.0810, 0.4120],
        [0.4670, 0.1910, 0.1510, 0.1910],
        [0.1360, 0.6080, 0.1970, 0.0590]])
context shape: (1, 4, 8)

Reflection Prompts#

  • What changed when you modified the starter experiment, and why should that change matter?

  • Which result surprised you, and what diagnostic would you run next?

  • What assumption would you document before handing this model to another practitioner?

  • Which failure mode from the module reading is most relevant to your result?