Stagewise Learning of Balanced Orthogonal Mixtures by Finite-Width Two-Head Transformers

Understanding how attention learns relationships between tokens is central to interpreting Transformers. We study context–signal correspondence learning in balanced orthogonal mixtures. Prior work introduced a three-stage population gradient-flow protocol for this setting. Each input contains three tokens: a context, its matching classification signal, and a distractor. The prior analysis relies on cancellation identities specific to two mixture components ($K=2$), and its finite-width guarantee uses a sufficient width that grows as the initialization scale decreases. Under the same training setting, we prove high-probability convergence of logistic risk to zero for every fixed number $K\ge 2$ of mixture components in a finite-width two-head softmax Transformer. For fixed $K$, bias, admissible initial readout-training duration, and confidence, the sufficient width is uniform over sufficiently small positive initialization scales and independent of the target risk. Our analysis combines joint Gram dynamics and relative perturbation estimates to control multidimensional attention growth, and handles asynchronous ReLU crossings along the actual random trajectory. We establish feature separation at the end of attention training. Subsequent readout training yields permanent activation of all ReLU gates and exact logistic-risk and max-margin readout asymptotics. For $K\ge 3$, we further quantify attention-probability gaps, including one-sided cubic lower bounds. Finally, successful learning does not require monotone improvement of the bias-augmented feature margin. At fixed admissible parameters, including a positive initialization scale, the probability of successful learning with an initial margin decrease remains bounded away from zero for all sufficiently large widths. With high probability, the margin nevertheless increases throughout the remaining attention-training interval after a delay that vanishes with width at fixed confidence.