§3.4Self-Attention with Trainable Weights (Scaled Dot-Product Attention)

The simplified mechanism of §3.3 had nothing to learn — its attention weights came straight from fixed embedding similarities. This section adds the one ingredient a real LLM needs: trainable weight matrices, so the model can learn which similarities matter for predicting the next token. The result is the mechanism actual GPTs use, called scaled dot-product attention scaled dot-product attention softmax(QKᵀ/√d_k)·V — self-attention with trainable projections; dividing by √d_k keeps softmax out of its vanishing-gradient regime at large dimensions. defined in ch. 3 — open in glossary . Rung 2 of the chapter’s ladder:

3.4.1 Queries, keys, and values — computed step by step

Three trainable weight matrices enter: $W_q$, $W_k$, $W_v$. Each one projects an embedded input token $x^{(i)}$ into a different role:

$q^{(i)} = x^{(i)} W_q \qquad k^{(i)} = x^{(i)} W_k \qquad v^{(i)} = x^{(i)} W_v$

In GPT-like models the input and output dimensions are usually the same; the book picks different ones (d_in=3, d_out=2) purely so you can follow which matrix produced which vector. Step through the projections exactly as the book’s figure draws them (cell values rounded to 1 decimal, as in the book):

import torch
inputs = torch.tensor(
[[0.43, 0.15, 0.89], # Your     (x^1)
 [0.55, 0.87, 0.66], # journey  (x^2)
 [0.57, 0.85, 0.64], # starts   (x^3)
 [0.22, 0.58, 0.33], # with     (x^4)
 [0.77, 0.25, 0.10], # one      (x^5)
 [0.05, 0.80, 0.55]] # step     (x^6)
)

x_2 = inputs[1] #A
d_in = inputs.shape[1] #B
d_out = 2 #C

torch.manual_seed(123)
W_query = torch.nn.Parameter(torch.rand(d_in, d_out), requires_grad=False) #D
W_key   = torch.nn.Parameter(torch.rand(d_in, d_out), requires_grad=False)
W_value = torch.nn.Parameter(torch.rand(d_in, d_out), requires_grad=False)

query_2 = x_2 @ W_query
key_2   = x_2 @ W_key
value_2 = x_2 @ W_value
print(query_2)
# tensor([0.4306, 1.4551])

Although we only want $z^{(2)}$ here, the query must be matched against — and retrieve from — every token, so all keys and values are computed in one shot:

keys = inputs @ W_key
values = inputs @ W_value
print("keys.shape:", keys.shape)
# keys.shape: torch.Size([6, 2])

From scores to scaled weights

The attention score between query 2 and key $j$, and its scaled softmax normalization:

$\omega_{2j} = q^{(2)} \cdot k^{(j)} \qquad\qquad \alpha_{2j} = \mathrm{softmax}\!\left(\frac{\omega_{2j}}{\sqrt{d_k}}\right)$

Worked example: $\omega_{22} = q^{(2)} \cdot k^{(2)} = 1.8524$, and the full row is $\omega_2 = [1.2705, 1.8524, 1.8111, 1.0795, 0.5577, 1.5440]$. Dividing by $\sqrt{2}$ and applying softmax gives the weights the book prints — verify below, where the exact numbers are reproduced live:

keys_2 = keys[1] #A
attn_score_22 = query_2.dot(keys_2)
print(attn_score_22)
# tensor(1.8524)

attn_scores_2 = query_2 @ keys.T   # all scores for this query
print(attn_scores_2)
# tensor([1.2705, 1.8524, 1.8111, 1.0795, 0.5577, 1.5440])

d_k = keys.shape[-1]
attn_weights_2 = torch.softmax(attn_scores_2 / d_k**0.5, dim=-1) #B
print(attn_weights_2)
# tensor([0.1500, 0.2264, 0.2199, 0.1311, 0.0906, 0.1820])

Feel the effect yourself. The temperature slider below divides the same scores before softmax — small $T$ mimics what unscaled, large dot products do to softmax (one bar takes everything: a step function), large $T$ what over-scaling would do (all bars equal: no signal). $\sqrt{d_k}$ sits in the healthy middle:

The context vector, now built from values

The final step mirrors §3.3, but the weighted sum runs over the value vectors, not the raw inputs:

$z^{(2)} = \sum_{j=1}^{T} \alpha_{2j}\, v^{(j)} = \alpha_2 \cdot \text{values}$

Worked example: $z^{(2)} = 0.1500{\cdot}v^{(1)} + 0.2264{\cdot}v^{(2)} + \dots + 0.1820{\cdot}v^{(6)} = [0.3061, 0.8210]$ — note it is 2-dimensional now ($d_{out}=2$), no longer 3-dimensional like the inputs.

context_vec_2 = attn_weights_2 @ values
print(context_vec_2)
# tensor([0.3061, 0.8210])

3.4.2 A compact self-attention class

Everything above, bundled into a PyTorch module — this is the book’s Listing 3.1:

import torch.nn as nn
class SelfAttention_v1(nn.Module):
  def __init__(self, d_in, d_out):
      super().__init__()
      self.d_out = d_out
      self.W_query = nn.Parameter(torch.rand(d_in, d_out))
      self.W_key   = nn.Parameter(torch.rand(d_in, d_out))
      self.W_value = nn.Parameter(torch.rand(d_in, d_out))

  def forward(self, x):
      keys    = x @ self.W_key
      queries = x @ self.W_query
      values  = x @ self.W_value

      attn_scores = queries @ keys.T  #A
      attn_weights = torch.softmax(
          attn_scores / keys.shape[-1]**0.5, dim=-1)  #B
      context_vec = attn_weights @ values
      return context_vec

torch.manual_seed(123)
sa_v1 = SelfAttention_v1(d_in, d_out)
print(sa_v1(inputs))
# tensor([[0.2996, 0.8053],
#         [0.3061, 0.8210],
#         [0.3058, 0.8203],
#         [0.2948, 0.7939],
#         [0.2927, 0.7891],
#         [0.2990, 0.8040]], grad_fn=<MmBackward0>)

Row 1 of the output is exactly the $z^{(2)} = [0.3061, 0.8210]$ computed by hand — the class runs the identical dataflow for all six queries at once:

Listing 3.2 — the same class with nn.Linear

Using nn.Linear (with bias=False) computes the same projections, but its optimized weight-initialization scheme makes training more stable and effective than raw torch.rand parameters:

class SelfAttention_v2(nn.Module):
  def __init__(self, d_in, d_out, qkv_bias=False):
      super().__init__()
      self.d_out = d_out
      self.W_query = nn.Linear(d_in, d_out, bias=qkv_bias)  #A
      self.W_key   = nn.Linear(d_in, d_out, bias=qkv_bias)
      self.W_value = nn.Linear(d_in, d_out, bias=qkv_bias)

  def forward(self, x):
      keys    = self.W_key(x)
      queries = self.W_query(x)
      values  = self.W_value(x)

      attn_scores = queries @ keys.T
      attn_weights = torch.softmax(
          attn_scores / keys.shape[-1]**0.5, dim=-1)
      context_vec = attn_weights @ values
      return context_vec

torch.manual_seed(789)
sa_v2 = SelfAttention_v2(d_in, d_out)
print(sa_v2(inputs))
# tensor([[-0.0739,  0.0713],
#         [-0.0748,  0.0703],
#         [-0.0749,  0.0702],
#         [-0.0760,  0.0685],
#         [-0.0763,  0.0679],
#         [-0.0754,  0.0693]], grad_fn=<MmBackward0>)