Grouped Query Attention with Linear Transformations

GroupedQueryAttention(GQA) is an optimisation technique that reduces the computational cost and memory usage of traditional Multi-Head Attention (MHA) by sharing key and value projections across multiple attention heads.

Note: I took few lines of code from Sebastian Raschka’s Qwen3 notebook as base reference of implementation.

Key Concept: Parameter Sharing

In traditional MHA(multi-head attention), each attention head has its own unique:

Query projection (W_query)
Key projection (W_key)
Value projection (W_value)

In GQA, we reduce the number of key and value projections by having multiple query heads share(parameter sharing) the same key and value projections.

The Linear Transformations in Detail

Looking at the code from the Qwen3 implementation:

self.W_query = nn.Linear(d_in, self.d_out, bias=False, dtype=dtype)           # Full size
self.W_key = nn.Linear(d_in, num_kv_groups * head_dim, bias=False, dtype=dtype)   # Reduced size
self.W_value = nn.Linear(d_in, num_kv_groups * head_dim, bias=False, dtype=dtype) # Reduced size

Key insight:

->  W_query projects to num_heads * head_dim (full size)
->  W_key and W_value project to num_kv_groups * head_dim (reduced size)

Where num_kv_groups < num_heads and group_size = num_heads // num_kv_groups.

Example Sentence : “The cat sat on the mat”

Let’s trace through the linear transformations step by step.

Setup
Sentence: "The cat sat on the mat" (6 tokens)

Configuration:

num_heads = 8 (8 attention heads)
num_kv_groups = 2 (2 key-value groups)
group_size = 4 (4 heads per group)
head_dim = 64 (64 dimensions per head)
d_in = 512 (input embedding dimension)

Step 1: Input Embeddings

First, our sentence gets tokenized and embedded:

# Input shape: (batch_size=1, num_tokens=6, d_in=512)
x = [
    [0.1, 0.2, ..., 0.5],  # "The" embedding
    [0.3, 0.1, ..., 0.7],  # "cat" embedding  
    [0.2, 0.4, ..., 0.3],  # "sat" embedding
    [0.5, 0.1, ..., 0.2],  # "on" embedding
    [0.1, 0.2, ..., 0.5],  # "the" embedding
    [0.4, 0.3, ..., 0.1]   # "mat" embedding
]
# Shape: (1, 6, 512)

Step 2: Linear Transformations

Now we apply the three linear transformations:

2.1 Query Transformation (Full Size)

queries = self.W_query(x)  # Shape: (1, 6, 8*64) = (1, 6, 512)

Each token gets projected to 512 dimensions (8 heads × 64 dimensions each):

# For token "cat" (position 1):
queries[0, 1] = [0.8, 0.2, 0.5, ..., 0.9]  # 512 dimensions
# This represents 8 different "questions" the model can ask about "cat"

2.2 Key Transformation (Reduced Size)

keys = self.W_key(x)  # Shape: (1, 6, 2*64) = (1, 6, 128)

Each token gets projected to only 128 dimensions (2 groups × 64 dimensions each):

# For token "cat" (position 1):
keys[0, 1] = [0.3, 0.7, 0.1, ..., 0.4]  # 128 dimensions
# This represents 2 different ways to "identify" what "cat" is

2.3 Value Transformation (Reduced Size)

values = self.W_value(x)  # Shape: (1, 6, 2*64) = (1, 6, 128)

Each token gets projected to 128 dimensions:

# For token "cat" (position 1):
values[0, 1] = [0.6, 0.2, 0.8, ..., 0.1]  # 128 dimensions
# This represents 2 different ways to "extract information" from "cat"

Step 3: Reshaping and Transposing

3.1 Reshape Queries

queries = queries.view(1, 6, 8, 64).transpose(1, 2)
# Shape: (1, 8, 6, 64) - 8 heads, each with 6 tokens, 64 dimensions each

3.2 Reshape Keys and Values

keys = keys.view(1, 6, 2, 64).transpose(1, 2)
# Shape: (1, 2, 6, 64) - 2 key groups, each with 6 tokens, 64 dimensions each

values = values.view(1, 6, 2, 64).transpose(1, 2)
# Shape: (1, 2, 6, 64) - 2 value groups, each with 6 tokens, 64 dimensions each

Step 4: The Key Operation — Expanding Keys and Values

This is where the magic happens! We expand the 2 key/value groups to match all 8 heads:

keys = keys.repeat_interleave(4, dim=1)    # 4 = group_size
values = values.repeat_interleave(4, dim=1) # 4 = group_size
# Shape: (1, 8, 6, 64) - Now matches the number of head

What this means in practice:

Get Kushvanth’s stories in your inbox

Join Medium for free to get updates from this writer.

Remember me for faster sign in

Before expansion:

keys[0, 0] = [K1_group1, K1_group1, K1_group1, K1_group1, K1_group1, K1_group1]  # Group 1 keys
keys[0, 1] = [K2_group2, K2_group2, K2_group2, K2_group2, K2_group2, K2_group2]  # Group 2 keys

After expansion:

keys[0, 0] = [K1_group1, K1_group1, K1_group1, K1_group1, K1_group1, K1_group1]  # Head 1 uses Group 1
keys[0, 1] = [K1_group1, K1_group1, K1_group1, K1_group1, K1_group1, K1_group1]  # Head 2 uses Group 1  
keys[0, 2] = [K1_group1, K1_group1, K1_group1, K1_group1, K1_group1, K1_group1]  # Head 3 uses Group 1
keys[0, 3] = [K1_group1, K1_group1, K1_group1, K1_group1, K1_group1, K1_group1]  # Head 4 uses Group 1
keys[0, 4] = [K2_group2, K2_group2, K2_group2, K2_group2, K2_group2, K2_group2]  # Head 5 uses Group 2
keys[0, 5] = [K2_group2, K2_group2, K2_group2, K2_group2, K2_group2, K2_group2]  # Head 6 uses Group 2
keys[0, 6] = [K2_group2, K2_group2, K2_group2, K2_group2, K2_group2, K2_group2]  # Head 7 uses Group 2
keys[0, 7] = [K2_group2, K2_group2, K2_group2, K2_group2, K2_group2, K2_group2]  # Head 8 uses Group 2

Step 5: Attention Computation

Now each head computes attention with its unique query but shared keys/values:

Head 1 (Query Group 1):

# Head 1 has unique query for "cat":
q1_cat = [0.8, 0.2, 0.5, ..., 0.9]  # 64 dimensions

# But uses shared keys from Group 1:
k1_the = [0.3, 0.1, ..., 0.2]   # Group 1 key for "The"
k1_cat = [0.7, 0.4, ..., 0.1]   # Group 1 key for "cat"  
k1_sat = [0.2, 0.8, ..., 0.5]   # Group 1 key for "sat"
# ... etc for all tokens

# Attention score for "cat" attending to "sat":
score = q1_cat @ k1_sat.T  # Dot product

Head 5 (Query Group 2):

# Head 5 has different unique query for "cat":
q5_cat = [0.1, 0.9, 0.3, ..., 0.6]  # Different 64 dimensions

# But uses shared keys from Group 2:
k2_the = [0.5, 0.2, ..., 0.7]   # Group 2 key for "The"
k2_cat = [0.1, 0.8, ..., 0.3]   # Group 2 key for "cat"
k2_sat = [0.9, 0.1, ..., 0.4]   # Group 2 key for "sat"
# ... etc for all tokens

# Different attention score for "cat" attending to "sat":
score = q5_cat @ k2_sat.T  # Different dot product

Step 6: Final Output Projection

After all heads compute attention, we combine them and project back:

# Combine all 8 heads: (1, 6, 8*64) = (1, 6, 512)
context = [head1_output, head2_output, ..., head8_output]

# Final projection back to original dimension
output = self.out_proj(context)  # Shape: (1, 6, 512)

What Each Head Learns

• Heads 1–4 (Group 1): Learn different query patterns but share the same way of identifying and extracting information
• Heads 5–8 (Group 2): Learn different query patterns and share a different way of identifying and extracting information

For our sentence “The cat sat on the mat”:

• Some heads might focus on subject-verb relationships (cat-sat)
• Others might focus on prepositional phrases (on-mat)
• But they share the underlying mechanisms for recognizing what constitutes a subject, verb, or object

This sharing reduces parameters while still allowing diverse attention patterns through the unique query projections!