Attention head fusion can need exponentially many entries for exact match

Exact Fusion of Softmax Attention Heads Can Require Exponentially Many Atoms

Machine LearningNeural and Evolutionary Computing

Summary

This paper studies how combining multiple attention heads in AI models can require a huge number of parts to replicate exactly as one head. The authors show that some functions represented by many simple heads need exponentially more components to be written as a single head, even if you allow flexibility in key choices and weights. This means that making heads work exactly as one might want can get very complicated and large in size. The work includes proofs and experiments analyzing these limits and trade-offs.

What this means in practice

  • For machine learning engineers: Understand limitations when compressing multiple attention heads into fewer heads without loss of precision in transformer models.
  • For ai system designers: Design model caching and pruning strategies while knowing when exact attention behavior cannot be preserved under limited parameter budgets.

A theory result. No direct application yet.

Authors

Hiroyasu Usami, Keisuke Hara, Ayato Tsuboi, Naohiko Matsuda

Abstract

Fusing independently normalized attention heads can increase the number of explicit key-value entries, or atoms, needed to preserve their output. We prove a sharp separation for $W_d(q)=\sum_{j=1}^dσ(q_j)$: its $d$ two-atom heads use $2d$ head-atom terms, whereas its minimum single-head representation has exactly $2^d$ atoms, even with arbitrary fixed replacement keys, scalar values, and positive masses. The same lower bound holds for every fixed uniform error below $1/2$ on $\mathbb R^d$, and the threshold is sharp. The proof uses exposed-key normal cones; exact equality on any nonempty open query domain inherits the bound. The resulting atom bound for a fixed head count is tight when that count divides $d$. A normalized family retains its exponential exact minimum while approaching a constant on fixed-radius query balls. For pretrained-model caches, a classical affine-rank bound and a certified instance exclude exact vector preservation at a specified atom count. Learned-cache comparisons on two disjoint banks of 64 contexts show budget-dependent reversals and calibration effects.