Michael Brenndoerfer | Data & AI, Private Equity, Technology

Michael Brenndoerfer

Architecture Hyperparameters › Number of Attention Heads

Michael BrenndoerferUpdated June 18, 20259 min read

Language AI HandbookAll chapters

Part 16 of 67Transformer Architectures

Chapter 116Architecture Hyperparameters › Number of Attention Heads

Reading section 4 of 10 · Number of Attention Heads · 9 minSection 4 of 10 · 9 min

Section 4 of 10 · 9 min read

Architecture Hyperparameters › Number of Attention HeadsOne section · Sections in this chapter4 / 10 ↓

Sections in this chapter10

Reading tools Definitions & inline links

Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.

Article links

Make inline references clickable

Multi-head attention works by splitting the model dimension into parallel "heads," each computing its own attention pattern independently. One head might attend to syntactic relationships (subject-verb agreement), another to semantic similarity (synonyms and related concepts), and yet another to positional patterns (nearby words). This parallel structure allows the model to capture multiple types of relationships simultaneously.

But how many heads should you use? The answer involves a fundamental trade-off. More heads mean more parallel attention patterns, enabling the model to track more distinct relationships. Fewer heads mean each head has more dimensions to work with, enabling richer attention computations per head. Since the model dimension must be divided among heads, you can't maximize both simultaneously.

The key constraint is that $d_{\text{model}}$ must be evenly divisible by $n_{\text{heads}}$ . This divisibility requirement determines the per-head dimension:

d_{\text{head}} = \frac{d_{\text{model}}}{n_{\text{heads}}}

where:

$d_{\text{head}}$ : the dimension of QKV representations within each attention head
$d_{\text{model}}$ : the model's hidden dimension
$n_{\text{heads}}$ : the number of parallel attention heads

More heads mean smaller per-head dimensions. With $d_{\text{model}} = 768$ and 12 heads, each head has 64 dimensions. With 24 heads, each would have only 32 dimensions. The trade-off is between having many specialized heads versus having fewer but more expressive heads.

The Head Dimension Sweet Spot

Empirically, head dimensions between 64 and 128 work best. Smaller heads lack the capacity to compute meaningful attention patterns. Larger heads provide diminishing returns while reducing the number of parallel attention mechanisms.

In[10]:

Code

def analyze_head_configurations(d_model, min_d_head=32, max_d_head=256):
    """
    Find valid head configurations for a given model dimension.

    Returns list of (n_heads, d_head) tuples that divide evenly.
    """
    configs = []

    for n_heads in range(1, d_model + 1):
        if d_model % n_heads == 0:
            d_head = d_model // n_heads
            if min_d_head <= d_head <= max_d_head:
                configs.append((n_heads, d_head))

    return configs


# Analyze configurations for common model dimensions
dimensions = [768, 1024, 2048, 4096]
all_configs = {}

for d_model in dimensions:
    configs = analyze_head_configurations(d_model)
    all_configs[d_model] = configs

Out[11]:

Console

Valid head configurations (d_head between 32 and 256):

d_model = 768:
  3h×256d, 4h×192d, 6h×128d, 8h×96d, 12h×64d, 16h×48d, 24h×32d

d_model = 1024:
  4h×256d, 8h×128d, 16h×64d, 32h×32d

d_model = 2048:
  8h×256d, 16h×128d, 32h×64d, 64h×32d

d_model = 4096:
  16h×256d, 32h×128d, 64h×64d, 128h×32d

The number of valid configurations varies with model dimension. Highly composite numbers like 768 (which factors as $2^8 \times 3 = 256 \times 3$ ) offer many choices, while prime-heavy dimensions constrain options. Model designers often choose dimensions specifically for their factorization properties.

Out[12]:

Visualization

Line plot showing head dimension decreasing as number of heads increases for different model dimensions, with preferred zone highlighted. — The inverse relationship between number of heads and per-head dimension for common model sizes. All points along each curve represent valid configurations where the model dimension divides evenly. Shaded region shows the empirically preferred range (64-128 dimensions per head).

The curves illustrate the fundamental constraint: for any model dimension, increasing heads necessarily decreases per-head dimension. The shaded zone shows where most successful architectures operate. Too few heads (right side of curves, high per-head dimension) means limited parallel attention patterns. Too many heads (left side, low per-head dimension) means each head lacks capacity for meaningful computation.

Multi-Query and Grouped-Query AttentionLink Copied

During autoregressive generation, transformers face a memory bottleneck: they must store the key and value vectors for every previous token in the sequence. This "KV cache" grows linearly with sequence length and can dominate GPU memory for long contexts. With 32 heads and 128-dimensional vectors, a 4096-token context requires storing over 33 million floating-point values per layer, just for the cache.

Modern architectures address this through an elegant asymmetry: reduce the number of key-value heads while keeping query heads high. The insight is that queries need fine-grained distinctions (each position asks different questions), but keys and values can be shared across groups of queries without significant quality loss. This design, called grouped-query attention (GQA), dramatically reduces cache size while preserving most of the model's expressiveness.

To understand the parameter and memory implications, let's examine how projections work in standard multi-head attention. Each head has its own QKV projections, with parameter counts computed as the product of the number of heads, the per-head dimension, and the model dimension:

Q projections: $n_{\text{heads}} \times d_{\text{head}} \times d_{\text{model}}$ parameters
K projections: $n_{\text{heads}} \times d_{\text{head}} \times d_{\text{model}}$ parameters
V projections: $n_{\text{heads}} \times d_{\text{head}} \times d_{\text{model}}$ parameters
Output projection: $n_{\text{heads}} \times d_{\text{head}} \times d_{\text{model}}$ parameters

where $n_{\text{heads}}$ is the number of attention heads, $d_{\text{head}}$ is the dimension per head, and $d_{\text{model}}$ is the model dimension. Since $n_{\text{heads}} \times d_{\text{head}} = d_{\text{model}}$ , each projection contributes $d_{\text{model}}^2$ parameters.

In grouped-query attention (GQA) with $n_{\text{kv\_heads}}$ key-value heads, the K and V projections are shared across groups of query heads:

Q projections: unchanged ( $n_{\text{heads}} \times d_{\text{head}} \times d_{\text{model}}$ )
K projections: reduced ( $n_{\text{kv\_heads}} \times d_{\text{head}} \times d_{\text{model}}$ )
V projections: reduced ( $n_{\text{kv\_heads}} \times d_{\text{head}} \times d_{\text{model}}$ )
Output projection: unchanged

Each group of $n_{\text{heads}} / n_{\text{kv\_heads}}$ query heads shares the same K and V projections, reducing parameters while maintaining the full expressiveness of query computations.

In[13]:

Code

def attention_parameters(d_model, n_heads, n_kv_heads=None):
    """
    Calculate attention parameters for different configurations.

    Args:
        d_model: Model dimension
        n_heads: Number of query heads
        n_kv_heads: Number of KV heads (None = standard MHA, 1 = MQA)

    Returns:
        Dictionary with parameter breakdown
    """
    if n_kv_heads is None:
        n_kv_heads = n_heads  # Standard multi-head attention

    d_head = d_model // n_heads

    params = {}
    params["q_proj"] = n_heads * d_head * d_model
    params["k_proj"] = n_kv_heads * d_head * d_model
    params["v_proj"] = n_kv_heads * d_head * d_model
    params["o_proj"] = n_heads * d_head * d_model
    params["total"] = sum(params.values())

    # Also compute KV cache size per token (for inference)
    params["kv_cache_per_token"] = 2 * n_kv_heads * d_head  # K and V vectors

    return params


# Compare different attention configurations
d_model = 4096
n_heads = 32
configs = {
    "MHA (standard)": attention_parameters(d_model, n_heads, n_kv_heads=32),
    "GQA (8 KV heads)": attention_parameters(d_model, n_heads, n_kv_heads=8),
    "GQA (4 KV heads)": attention_parameters(d_model, n_heads, n_kv_heads=4),
    "MQA (1 KV head)": attention_parameters(d_model, n_heads, n_kv_heads=1),
}

Out[14]:

Console

Attention configurations for d_model=4096, n_heads=32:

Configuration            Q params     K params     V params        Total   KV cache/tok
------------------------------------------------------------------------------------------
MHA (standard)             16.78M       16.78M       16.78M       67.11M           8192
GQA (8 KV heads)           16.78M        4.19M        4.19M       41.94M           2048
GQA (4 KV heads)           16.78M        2.10M        2.10M       37.75M           1024
MQA (1 KV head)            16.78M      524.29K      524.29K       34.60M            256

The parameter savings from GQA are modest (the reduction is only in K and V projections), but the real benefit is during inference. The KV cache, which stores key and value vectors for all previous tokens during autoregressive generation, shrinks proportionally with $n_{\text{kv\_heads}}$ . For long sequences, this can be the difference between fitting a model in GPU memory or not.

LLaMA-2 70B uses GQA with 8 KV heads and 64 query heads, reducing KV cache memory by 8x compared to standard MHA while maintaining quality. This design choice enables longer context windows with the same hardware.

Out[15]:

Visualization

Line plot comparing KV cache memory growth for MHA, GQA, and MQA attention variants across sequence lengths. — KV cache memory requirements across sequence lengths for different attention configurations. GQA with 8 KV heads uses 8x less memory than standard MHA, enabling much longer context windows within the same memory budget.

The logarithmic scales reveal the dramatic impact of GQA. At 128K tokens with this LLaMA-70B-like configuration, standard MHA requires roughly 320 GB of KV cache memory across all layers, while GQA with 8 KV heads needs only about 40 GB, an 8x reduction. This difference is why long-context models universally adopt GQA or MQA.