Day 49: Week 7 Consolidation: Attention to Logits, plus the Pretraining Data Pipeline

Inspect Reference Solution & Trace:

1. x = [502, 1204, 33, 991], y = [1204, 33, 991, 88]. L supervision pairs (x_t, y_t) need L inputs plus one extra token for the last target.
2. get_batch:

def get_batch(data, batch_size, block_size, device):
    ix = torch.randint(len(data) - block_size, (batch_size,))
    x = torch.stack([torch.from_numpy(data[i:i + block_size].astype(np.int64)) for i in ix])
    y = torch.stack([torch.from_numpy(data[i + 1:i + 1 + block_size].astype(np.int64)) for i in ix])
    return x.to(device), y.to(device)

3. 65,536 tokens * 8 bytes = 524,288 bytes per tensor, so x and y take 1,048,576 bytes (1.0 MiB) as int64 and 524,288 bytes (512 KiB) as int32.
4. loss = F.cross_entropy(logits.view(-1, logits.size(-1)), y.view(-1)). The expected first loss is ln(50257) ≈ 10.8249. A first loss far above this, or NaN, points to broken initialization or a shape bug.
5. With y = x the model only has to copy its input, the loss falls toward zero, and it learns nothing about predicting the next token.

Inspect Reference Solution & Trace:

1. Scaled S = [[1.4142, 1.4142], [1.4142, 0.0]]. Causal weights: row 0 = [1.0, 0.0]; row 1 = softmax([1.4142, 0.0]) = [0.804430, 0.195570]. Out = [[1.0, 0.0], [0.804430, 0.195570]]. Without the mask, row 0 would be [0.5, 0.5], so Out[0] = [0.5, 0.5]: token 0 would read token 1 from the future.
2. d_k = 12 / 3 = 4. Q after view and transpose: (2, 3, 5, 4). Scores: (2, 3, 5, 5). Output after concatenation: (2, 5, 12).
3. Four 12 x 12 matrices = 576 weights, plus 4 * 12 = 48 biases = 624. Without biases: 576.
4. For unit-variance q and k, Var(q . k) = d_k (Day 43); dividing by sqrt(d_k) brings it back to 1, so softmax does not saturate and its gradients do not vanish. exp(15) ≈ 3.27e6 is above the float16 maximum of 65,504, so it overflows to inf; the max-shift keeps every exponent <= 0 (Day 45). Out[0] does not change at all, because the mask gives token 0 zero weight on every later token (Day 44).

Inspect Reference Solution & Trace:

1. LayerNorm: (x - 1) / sqrt(5) = [1.3416, -0.4472, -1.3416, 0.4472]. RMSNorm: x / sqrt(6) = [1.6330, 0.0, -0.8165, 0.8165].
2. x = x + Attn(LN1(x)); x = x + FFN(LN2(x)). Pre-LN leaves the residual path as a clean identity, so gradients reach early layers without passing through a normalization (Day 46).
3. Per block: 12 * 32^2 = 12,288 matrix weights + 2 LayerNorms * 64 = 128, so 12,416. Two blocks: 24,832. Token embedding 100 * 32 = 3,200; positional 16 * 32 = 512; final LayerNorm 64. Total = 28,608. An untied head adds 3,200 (31,808 total); tying saves exactly V * d_model (Day 48).
4. FFN 8 d^2 out of 12 d^2 = 2/3 (Day 47).
5. idx (8, 16) -> token + positional embeddings (8, 16, 32) -> 2 blocks (8, 16, 32) -> final LayerNorm (8, 16, 32) -> logits (8, 16, 100). The next-token logits are logits[:, -1, :] with shape (8, 100).

Inspect Reference Solution & Trace:

1. (A) Divides by d_model instead of sqrt(d_k), which shrinks the scores far more than needed, so the softmax is too flat and attention barely tells tokens apart (Day 43). (B) diagonal=0 also masks each token's own position, so row 0 has every entry at -inf (Days 44-45). (C) Post-LN: the residual sum passes through LayerNorm, so the identity gradient path is lost (Day 46). (D) y = chunk[:-1] is the same as x, so there is no shift (Day 49). Fixes:

scores = q @ k.transpose(-2, -1) / math.sqrt(q.size(-1))               # (A)
mask = torch.triu(torch.ones(T, T, dtype=torch.bool), diagonal=1)      # (B)
x = x + self.attn(self.ln_1(x))                                        # (C)
x, y = chunk[:-1], chunk[1:]                                           # (D)

2. (B): row 0 of the scores is all -inf, so softmax computes exp(-inf - (-inf)) = exp(NaN) on the first forward pass (Day 45).
3. (D): with y = x the model only has to copy its input, so the loss heads toward 0 while generations are useless.
4. ln(100) ≈ 4.6052.