Token-Specific Attention Heads: Mixing Per-Token and Shared Query, Key, and Value Projections in Transformer Language Models

In standard multi-head attention, every head uses the same query, key, and value (QKV) projectionmatrices for every token. We study a simple variant in which, in every layer, a subset of the heads aretoken-specific: their QKV projections depend on the identity of the input token, while the remaining headskeep ordinary shared projections. Token-specific projections are parameterized as a shared matrix plus aper-token low-rank correction stored in a hashed lookup table, so the extra computation is small, althoughthe parameter count grows. In GPT-style language models trained on FineWeb, with 2 of 4 heads per layertoken-specific, every run had lower validation loss than every standard-attention run, at both of the scaleswe tested. At the larger of two scales (6 layers, width 512, 300M training tokens, 3 seeds per method),validation loss fell from 4.154 ± 0.005 to 4.107 ± 0.008, a 4.6% reduction in perplexity. A controlthat adds the same number of lookup-table parameters to the same heads as per-token vectors, ratherthan per-token matrix corrections, reached only 4.148 ± 0.002, indicating that the gain is not explainedby the additional parameters alone. At the smaller scale, a sweep over the number of token-specificheads suggests that mixing token-specific and shared heads works at least as well as making all headstoken-specific. The main cost is throughput: our unoptimized implementation trains about 30% slowerthan the baseline. These are small-scale results, and we discuss their limitations in detail.