v1.1.1: Latent attention

A new side topic, 2c: latent attention. DeepSeek's multi-head latent attention caches a short latent per token in place of keys and values, projects every head's key and value up from it, and at decoding folds the up-projection into the query, so the cache is never projected up. A rotary key of its own carries the position. Stage 3's 14M model was trained with it and with three kinds of key/value heads, for the 6,000 steps of the reference run.

• Loss. Per token and layer, the cache holds 768 numbers with 6 key/value heads, 256 with stage 2's 2, 128 with 1, and 160 with a latent of 128 and a rotary key of 32. Their validation losses:
  • 6 key/value heads: 1.455
  • 2, stage 2's: 1.465, and 1.461 with another seed
  • 1 key/value head: 1.464
  • latent 128 + 32: 1.464
  • latent 64 + 32: 1.478, and latent 32 + 16: 1.490
• Speed alone. For one sequence, decoding against the latent writes 179 tokens per second, against 509 with stage 2's attention.
• Speed in a batch. With 1,024 sequences at once, it writes 43,800 tokens per second, 1.6 times as many as stage 2's attention. The smallest latent writes 68,300. Projected up at every step, the latent of 128 writes 12,800.
• Why. Stage 2's attention copies shared keys out to every head, so one key/value head is hardly faster than six. The latent path reads its cache once for all heads.

New code: LatentAttention in minllm/model.py, train.py --kv-latent --rope-dim, and bench_attention.py.

The films of every chapter are attached, 27 of them.