single-head-phase-transitions-eliminated-by-two-heads

IN premise — summaries/2026/08/24/von-oswald-2023-icl-gd-s0-ix-chunk-1.md

Created 2026-08-24T17:11:05+00:00

Single-head self-attention layers exhibit seed-dependent long training plateaus before sudden loss drops; adding a second head eliminates these phase transitions. With limited training data (8192 tasks), grokking-like behavior where test loss spikes then collapses is reproducible.

Summary

Training a transformer with only one attention mechanism often produces long, unpredictable stalls in learning followed by sudden dramatic improvement, and whether that stall even appears can depend on the random initialization — but adding a second attention head smooths that chaotic, phase-transition behavior away. The observation also gives the system a reliable recipe for reproducing "grokking" (a brief worsening of generalization before a sharp recovery) by simply capping training data at around eight thousand tasks.