Track papers across labs, topics, and trends.
Search papers, labs, and topics across Lattice.
This paper establishes that multi-head scaled dot product attention functions as a parameter identification mechanism, revealing that the ratio of unidentified parameters decreases with an increasing number of heads. The findings indicate that while more heads enhance structural identification, attention mechanisms can never achieve complete identification, and certain bias terms may not influence softmax-based attention layers. Through numerical examples and a rebalancing approach, the authors explore how these insights provide a mathematical rationale for architectural choices in transformer models that lead to improved performance.
More heads in attention layers not only enhance parameter identification but also reveal that some biases can be entirely irrelevant to model performance.
We prove that a multi-head scaled dot product attention can be viewed as a parameter identification strategy. The ratio of unidentified parameters to the total number of parameters scales like the reciprocal of the number of heads ($1/2 \to 1/(2H)$), meaning models with more heads are structurally more identified. A subtle side effect of the mathematics observation that attention can never be fully identified. Similarly we also show that some bias terms can have no effect on softmax-based attention layers in both the single- and multiple-head settings, though this is mostly a curiosity that should have a marginal effect on model size and model training/prediction efficiency. We also touch on modern improvements to transformers including RoPE and GQA from this perspective, illustrating how those as well can improve the ratio of ``meaningful'' parameters to all parameters. Simple numerical examples demonstrate that training can indeed involve updates that overlap model-invariant subspaces that arise from a lack of identification. As part of our experiments we use a ``rebalancing'' approach that can ``fix'' updates that overlap unindentified subspaces but do not try to present evidence this should actually be adopted. Instead we simply view our numerical results as exploring and confirming the theoretical results. As a whole we discuss a purely mathematical/statistical explanation, identification, for why specific architectural choices in transformers may have improved performance.